4 files changed, 20 insertions, 18 deletions
diff --git a/pkg/sentry/mm/BUILD b/pkg/sentry/mm/BUILD
index bbdfae247..3f396986a 100644
--- a/pkg/sentry/mm/BUILD
+++ b/pkg/sentry/mm/BUILD
@@ -1,7 +1,24 @@
 package(licenses = ["notice"])  # Apache 2.0
 
 load("//tools/go_generics:defs.bzl", "go_template_instance")
-load("//tools/go_stateify:defs.bzl", "go_library", "go_test")
+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",
+)
 
 go_template_instance(
     name = "file_refcount_set",
@@ -84,6 +101,7 @@ go_library(
         "lifecycle.go",
         "metadata.go",
         "mm.go",
+        "mm_state.go",
         "pma.go",
         "pma_set.go",
         "proc_pid_maps.go",
@@ -113,6 +131,7 @@ 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 b42156d45..992bde5a5 100644
--- a/pkg/sentry/mm/aio_context.go
+++ b/pkg/sentry/mm/aio_context.go
@@ -28,8 +28,6 @@ import (
 )
 
 // aioManager creates and manages asynchronous I/O contexts.
-//
-// +stateify savable
 type aioManager struct {
 	// mu protects below.
 	mu sync.Mutex `state:"nosave"`
@@ -91,16 +89,12 @@ 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"`
@@ -196,8 +190,6 @@ 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 3299ae164..ce8097b7f 100644
--- a/pkg/sentry/mm/mm.go
+++ b/pkg/sentry/mm/mm.go
@@ -46,8 +46,6 @@ import (
 )
 
 // MemoryManager implements a virtual address space.
-//
-// +stateify savable
 type MemoryManager struct {
 	// p is the platform.
 	//
@@ -209,8 +207,6 @@ 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.
@@ -350,8 +346,6 @@ 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
@@ -386,7 +380,6 @@ 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 aa2f87107..9d3614034 100644
--- a/pkg/sentry/mm/special_mappable.go
+++ b/pkg/sentry/mm/special_mappable.go
@@ -28,8 +28,6 @@ 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