diff --git a/cr-restore.c b/cr-restore.c
index 3924e639c..1991edbb4 100644
--- a/cr-restore.c
+++ b/cr-restore.c
@@ -52,7 +52,7 @@
 #include "restorer-blob.h"
 #include "crtools.h"
 #include "namespaces.h"
-#include "shmem.h"
+#include "mem.h"
 #include "mount.h"
 #include "fsnotify.h"
 #include "pstree.h"
@@ -177,7 +177,7 @@ static int root_prepare_shared(void)
 		if (pi->state == TASK_HELPER)
 			continue;
 
-		ret = prepare_shmem_pid(pi->pid.virt);
+		ret = prepare_mm_pid(pi);
 		if (ret < 0)
 			break;
 
diff --git a/include/mem.h b/include/mem.h
index 6cfd2e30f..5269cad02 100644
--- a/include/mem.h
+++ b/include/mem.h
@@ -4,7 +4,9 @@
 struct parasite_ctl;
 struct vm_area_list;
 struct page_pipe;
+struct pstree_item;
 
+extern int prepare_mm_pid(struct pstree_item *i);
 extern int do_task_reset_dirty_track(int pid);
 extern unsigned int dump_pages_args_size(struct vm_area_list *vmas);
 extern int parasite_dump_pages_seized(struct parasite_ctl *ctl,
diff --git a/include/shmem.h b/include/shmem.h
index cdb058907..8eaed56f7 100644
--- a/include/shmem.h
+++ b/include/shmem.h
@@ -21,7 +21,8 @@ struct shmem_info {
 	futex_t		lock;
 };
 
-extern int prepare_shmem_pid(int pid);
+struct _VmaEntry;
+extern int collect_shmem(int pid, struct _VmaEntry *vi);
 extern int prepare_shmem_restore(void);
 extern void show_saved_shmems(void);
 extern int get_shmem_fd(int pid, VmaEntry *vi);
diff --git a/mem.c b/mem.c
index 62097a84e..d3787dd40 100644
--- a/mem.c
+++ b/mem.c
@@ -15,6 +15,8 @@
 #include "kerndat.h"
 #include "stats.h"
 #include "vma.h"
+#include "shmem.h"
+#include "pstree.h"
 
 #include "protobuf.h"
 #include "protobuf/pagemap.pb-c.h"
@@ -338,3 +340,41 @@ int parasite_dump_pages_seized(struct parasite_ctl *ctl,
 	return ret;
 }
 
+int prepare_mm_pid(struct pstree_item *i)
+{
+	pid_t pid = i->pid.virt;
+	int fd, ret = -1;
+	VmaEntry *vi;
+
+	fd = open_image(CR_FD_VMAS, O_RSTR, pid);
+	if (fd < 0) {
+		if (errno == ENOENT)
+			return 0;
+		else
+			return -1;
+	}
+
+	while (1) {
+		ret = pb_read_one_eof(fd, &vi, PB_VMA);
+		if (ret <= 0)
+			break;
+
+		pr_info("vma 0x%"PRIx64" 0x%"PRIx64"\n", vi->start, vi->end);
+
+		if (!vma_entry_is(vi, VMA_ANON_SHARED) ||
+		    vma_entry_is(vi, VMA_AREA_SYSVIPC)) {
+			vma_entry__free_unpacked(vi, NULL);
+			continue;
+		}
+
+		ret = collect_shmem(pid, vi);
+		vma_entry__free_unpacked(vi, NULL);
+
+		if (ret)
+			break;
+	}
+
+	close(fd);
+	return ret;
+}
+
diff --git a/shmem.c b/shmem.c
index 1b753a429..1deb5918c 100644
--- a/shmem.c
+++ b/shmem.c
@@ -38,7 +38,7 @@ static struct shmem_info *find_shmem_by_id(unsigned long id)
 	return find_shmem(si, nr_shmems, id);
 }
 
-static int collect_shmem(int pid, VmaEntry *vi)
+int collect_shmem(int pid, VmaEntry *vi)
 {
 	unsigned long size = vi->pgoff + vi->end - vi->start;
 	struct shmem_info *si;
@@ -85,43 +85,6 @@ static int collect_shmem(int pid, VmaEntry *vi)
 	return 0;
 }
 
-int prepare_shmem_pid(int pid)
-{
-	int fd, ret = -1;
-	VmaEntry *vi;
-
-	fd = open_image(CR_FD_VMAS, O_RSTR, pid);
-	if (fd < 0) {
-		if (errno == ENOENT)
-			return 0;
-		else
-			return -1;
-	}
-
-	while (1) {
-		ret = pb_read_one_eof(fd, &vi, PB_VMA);
-		if (ret <= 0)
-			break;
-
-		pr_info("vma 0x%"PRIx64" 0x%"PRIx64"\n", vi->start, vi->end);
-
-		if (!vma_entry_is(vi, VMA_ANON_SHARED) ||
-		    vma_entry_is(vi, VMA_AREA_SYSVIPC)) {
-			vma_entry__free_unpacked(vi, NULL);
-			continue;
-		}
-
-		ret = collect_shmem(pid, vi);
-		vma_entry__free_unpacked(vi, NULL);
-
-		if (ret)
-			break;
-	}
-
-	close(fd);
-	return ret;
-}
-
 static int shmem_wait_and_open(int pid, struct shmem_info *si)
 {
 	char path[128];