4 files changed, 18 insertions, 20 deletions
diff --git a/pkg/sentry/mm/BUILD b/pkg/sentry/mm/BUILD
index 3f396986a..bbdfae247 100644
--- a/pkg/sentry/mm/BUILD
+++ b/pkg/sentry/mm/BUILD
@@ -1,24 +1,7 @@
 package(licenses = ["notice"])  # Apache 2.0
 
 load("//tools/go_generics:defs.bzl", "go_template_instance")
-load("//tools/go_stateify:defs.bzl", "go_library", "go_stateify", "go_test")
-
-go_stateify(
-    name = "mm_state",
-    srcs = [
-        "aio_context.go",
-        "aio_context_state.go",
-        "file_refcount_set.go",
-        "io_list.go",
-        "mm.go",
-        "pma_set.go",
-        "save_restore.go",
-        "special_mappable.go",
-        "vma_set.go",
-    ],
-    out = "mm_state.go",
-    package = "mm",
-)
+load("//tools/go_stateify:defs.bzl", "go_library", "go_test")
 
 go_template_instance(
     name = "file_refcount_set",
@@ -101,7 +84,6 @@ go_library(
         "lifecycle.go",
         "metadata.go",
         "mm.go",
-        "mm_state.go",
         "pma.go",
         "pma_set.go",
         "proc_pid_maps.go",
@@ -131,7 +113,6 @@ go_library(
         "//pkg/sentry/safemem",
         "//pkg/sentry/usage",
         "//pkg/sentry/usermem",
-        "//pkg/state",
         "//pkg/sync",
         "//pkg/syserror",
         "//pkg/tcpip/buffer",
diff --git a/pkg/sentry/mm/aio_context.go b/pkg/sentry/mm/aio_context.go
index 992bde5a5..b42156d45 100644
--- a/pkg/sentry/mm/aio_context.go
+++ b/pkg/sentry/mm/aio_context.go
@@ -28,6 +28,8 @@ import (
 )
 
 // aioManager creates and manages asynchronous I/O contexts.
+//
+// +stateify savable
 type aioManager struct {
 	// mu protects below.
 	mu sync.Mutex `state:"nosave"`
@@ -89,12 +91,16 @@ func (a *aioManager) lookupAIOContext(id uint64) (*AIOContext, bool) {
 }
 
 // ioResult is a completed I/O operation.
+//
+// +stateify savable
 type ioResult struct {
 	data interface{}
 	ioEntry
 }
 
 // AIOContext is a single asynchronous I/O context.
+//
+// +stateify savable
 type AIOContext struct {
 	// done is the notification channel used for all requests.
 	done chan struct{} `state:"nosave"`
@@ -190,6 +196,8 @@ func (ctx *AIOContext) WaitChannel() (chan struct{}, bool) {
 
 // aioMappable implements memmap.MappingIdentity and memmap.Mappable for AIO
 // ring buffers.
+//
+// +stateify savable
 type aioMappable struct {
 	refs.AtomicRefCount
 
diff --git a/pkg/sentry/mm/mm.go b/pkg/sentry/mm/mm.go
index ce8097b7f..3299ae164 100644
--- a/pkg/sentry/mm/mm.go
+++ b/pkg/sentry/mm/mm.go
@@ -46,6 +46,8 @@ import (
 )
 
 // MemoryManager implements a virtual address space.
+//
+// +stateify savable
 type MemoryManager struct {
 	// p is the platform.
 	//
@@ -207,6 +209,8 @@ type MemoryManager struct {
 }
 
 // vma represents a virtual memory area.
+//
+// +stateify savable
 type vma struct {
 	// mappable is the virtual memory object mapped by this vma. If mappable is
 	// nil, the vma represents a private anonymous mapping.
@@ -346,6 +350,8 @@ func (v *vma) loadRealPerms(b int) {
 }
 
 // pma represents a platform mapping area.
+//
+// +stateify savable
 type pma struct {
 	// file is the file mapped by this pma. Only pmas for which file ==
 	// platform.Platform.Memory() may be saved. pmas hold a reference to the
@@ -380,6 +386,7 @@ type pma struct {
 	internalMappings safemem.BlockSeq `state:"nosave"`
 }
 
+// +stateify savable
 type privateRefs struct {
 	mu sync.Mutex `state:"nosave"`
 
diff --git a/pkg/sentry/mm/special_mappable.go b/pkg/sentry/mm/special_mappable.go
index 9d3614034..aa2f87107 100644
--- a/pkg/sentry/mm/special_mappable.go
+++ b/pkg/sentry/mm/special_mappable.go
@@ -28,6 +28,8 @@ import (
 // semantics similar to Linux's mm/mmap.c:_install_special_mapping(), except
 // that SpecialMappable takes ownership of the memory that it represents
 // (_install_special_mapping() does not.)
+//
+// +stateify savable
 type SpecialMappable struct {
 	refs.AtomicRefCount