From 312862fee45daef8cb9778f42e17b3bffd13e457 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 15:57:58 +0000
Subject: [PATCH 01/20] Update guest Cargo.lock

Previously, this was inconsistent.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/tests/rust_guests/witguest/Cargo.lock | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/tests/rust_guests/witguest/Cargo.lock b/src/tests/rust_guests/witguest/Cargo.lock
index 3cd3633a1..647cd7438 100644
--- a/src/tests/rust_guests/witguest/Cargo.lock
+++ b/src/tests/rust_guests/witguest/Cargo.lock
@@ -84,9 +84,9 @@ dependencies = [
 
 [[package]]
 name = "cc"
-version = "1.2.53"
+version = "1.2.54"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "755d2fce177175ffca841e9a06afdb2c4ab0f593d53b4dee48147dfaade85932"
+checksum = "6354c81bbfd62d9cfa9cb3c773c2b7b2a3a482d569de977fd0e961f6e7c00583"
 dependencies = [
  "find-msvc-tools",
  "shlex",
@@ -644,9 +644,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
 
 [[package]]
 name = "wasmparser"
-version = "0.243.0"
+version = "0.244.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f6d8db401b0528ec316dfbe579e6ab4152d61739cfe076706d2009127970159d"
+checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe"
 dependencies = [
  "bitflags",
  "hashbrown 0.15.5",

From 58d5506827ac6ed917cbe8e0d185d0462dcca1df Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 17:26:45 +0000
Subject: [PATCH 02/20] [nix] include miri for the nightly toolchain

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 flake.nix | 68 ++++++++++++++++++++++++++++---------------------------
 1 file changed, 35 insertions(+), 33 deletions(-)

diff --git a/flake.nix b/flake.nix
index fd1e67d48..3282a0eff 100644
--- a/flake.nix
+++ b/flake.nix
@@ -17,7 +17,7 @@
             patchRustPkg = pkg: (pkg.overrideAttrs (oA: {
               buildCommand = (builtins.replaceStrings
                 [ "rustc,rustdoc" ]
-                [ "rustc,rustdoc,clippy-driver,cargo-clippy" ]
+                [ "rustc,rustdoc,clippy-driver,cargo-clippy,miri,cargo-miri" ]
                 oA.buildCommand) + (let
                   wrapperPath = pkgs.path + "/pkgs/build-support/bintools-wrapper/ld-wrapper.sh";
                   baseOut = pkgs.clangStdenv.cc.bintools.out;
@@ -68,7 +68,7 @@
                 "x86_64-pc-windows-msvc" "x86_64-unknown-none"
                 "wasm32-wasip1" "wasm32-wasip2" "wasm32-unknown-unknown"
               ];
-              extensions = [ "rust-src" ];
+              extensions = [ "rust-src" ] ++ (if args.channel == "nightly" then [ "miri-preview" ] else []);
             });
 
           # Hyperlight needs a variety of toolchains, since we use Nightly
@@ -96,6 +96,31 @@
             rustc = toolchains.stable.rust;
           };
 
+          # when building a guest with cargo-hyperlight, or when
+          # building a miri sysroot for the main workspace, we need to
+          # include any crates.io dependencies of the standard library
+          # (e.g. rustc-literal-escaper)
+          stdlibLocks = lib.mapAttrsToList (_: toolchain:
+            "${toolchain.rust}/lib/rustlib/src/rust/library/Cargo.lock"
+          ) toolchains;
+          stdlibDeps = builtins.map (lockFile:
+            rust-platform.importCargoLock { inherit lockFile; }) stdlibLocks;
+          withStdlibLock = lockFile:
+            pkgs.symlinkJoin {
+              name = "cargo-deps";
+              paths = stdlibDeps ++ [
+                (rust-platform.importCargoLock {
+                  inherit lockFile;
+                })
+              ];
+            };
+          deps = {
+            "Cargo.toml" = withStdlibLock ./Cargo.lock;
+            "src/tests/rust_guests/dummyguest/Cargo.toml" = withStdlibLock ./src/tests/rust_guests/dummyguest/Cargo.lock;
+            "src/tests/rust_guests/simpleguest/Cargo.toml" = withStdlibLock ./src/tests/rust_guests/simpleguest/Cargo.lock;
+            "src/tests/rust_guests/witguest/Cargo.toml" = withStdlibLock ./src/tests/rust_guests/witguest/Cargo.lock;
+          };
+
           # Script snippet, used in the cargo/rustc wrappers below,
           # which creates a number of .cargo/config.toml files in
           # order to allow using Nix-fetched dependencies (this must
@@ -105,7 +130,7 @@
           # unfortunately that tends not to play well with subcommands
           # like `cargo clippy` and `cargo hyperlight` (see
           # https://github.com/rust-lang/cargo/issues/11031).
-          materialiseDeps = deps: let
+          materialiseDeps = let
             sortedNames =  lib.lists.reverseList (builtins.attrNames deps);
             matchClause = path: ''  */${path}) root="''${manifest%${path}}" ;;'';
             matchClauses = lib.strings.concatStringsSep "\n"
@@ -144,7 +169,7 @@
           # scripts also use `rustup toolchain install` in some cases, in
           # order to work in CI, so we provide a fake rustup that does
           # nothing as well.
-          rustup-like-wrapper = name: deps: pkgs.writeShellScriptBin name
+          rustup-like-wrapper = name: pkgs.writeShellScriptBin name
             (let
               clause = name: toolchain:
                 "+${name}) base=\"${toolchain.rust}\"; shift 1; ;;";
@@ -152,7 +177,7 @@
                 (lib.mapAttrsToList clause toolchains);
             in ''
               base="${toolchains.stable.rust}"
-              ${materialiseDeps deps}
+              ${materialiseDeps}
               case "$1" in
                 ${clauses}
                 install) exit 0; ;;
@@ -160,12 +185,12 @@
               export PATH="$base/bin:$PATH"
               exec "$base/bin/${name}" "$@"
             '');
-          fake-rustup = deps: pkgs.symlinkJoin {
+          fake-rustup = pkgs.symlinkJoin {
             name = "fake-rustup";
             paths = [
               (pkgs.writeShellScriptBin "rustup" "")
-              (rustup-like-wrapper "rustc" deps)
-              (rustup-like-wrapper "cargo" deps)
+              (rustup-like-wrapper "rustc")
+              (rustup-like-wrapper "cargo")
             ];
           };
 
@@ -182,34 +207,11 @@
             };
             cargoHash = "sha256-muiMVrK1TydQiMitihfo7xYidqUIIQ+Hw3BIeo5rLFw=";
           };
-          # when building a guest with cargo-hyperlight, we need to
-          # include any crates.io dependencies of the standard library
-          # (e.g. rustc-literal-escaper)
-          stdlibLocks = lib.mapAttrsToList (_: toolchain:
-            "${toolchain.rust}/lib/rustlib/src/rust/library/Cargo.lock"
-          ) toolchains;
-          stdlibDeps = builtins.map (lockFile:
-            rust-platform.importCargoLock { inherit lockFile; }) stdlibLocks;
-          withStdlibLock = lockFile:
-            pkgs.symlinkJoin {
-              name = "cargo-deps";
-              paths = stdlibDeps ++ [
-                (rust-platform.importCargoLock {
-                  inherit lockFile;
-                })
-              ];
-            };
-          deps = finalRootVendor: {
-            "Cargo.toml" = finalRootVendor;
-            "src/tests/rust_guests/dummyguest/Cargo.toml" = withStdlibLock ./src/tests/rust_guests/dummyguest/Cargo.lock;
-            "src/tests/rust_guests/simpleguest/Cargo.toml" = withStdlibLock ./src/tests/rust_guests/simpleguest/Cargo.lock;
-            "src/tests/rust_guests/witguest/Cargo.toml" = withStdlibLock ./src/tests/rust_guests/witguest/Cargo.lock;
-          };
         in (buildRustPackageClang (mkDerivationAttrs: {
           pname = "hyperlight";
           version = "0.0.0";
           src = lib.cleanSource ./.;
-          cargoLock.lockFile = ./Cargo.lock;
+          cargoDeps = deps."Cargo.toml";
 
           nativeBuildInputs = [
             azure-cli
@@ -246,7 +248,7 @@
           # Set this through shellHook rather than nativeBuildInputs to be
           # really sure that it overrides the real cargo.
           postHook = ''
-            export PATH="${fake-rustup (deps mkDerivationAttrs.cargoDeps)}/bin:$PATH"
+            export PATH="${fake-rustup}/bin:$PATH"
           '';
         })).overrideAttrs(oA: {
           hardeningDisable = [ "all" ];

From 083918421c047f4224614db842e778cf1f4c7051 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 17:54:29 +0000
Subject: [PATCH 03/20] [nix] hack to make libgit2-sys builds more reliable

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 flake.nix | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/flake.nix b/flake.nix
index 3282a0eff..05ddd9561 100644
--- a/flake.nix
+++ b/flake.nix
@@ -152,7 +152,11 @@
             makeClauses = lib.strings.concatStringsSep "\n"
               (lib.mapAttrsToList makeClause deps);
           in ''
-            manifest=$(''${base}/bin/cargo locate-project --message-format plain --workspace)
+            base_cargo() {
+              PATH="$base/bin:$PATH" "$base/bin/cargo" "$@"
+            }
+
+            manifest=$(base_cargo locate-project --message-format plain --workspace)
             case "$manifest" in
               ${matchClauses}
             esac
@@ -160,6 +164,16 @@
               sed -i '/# vendor dependency configuration generated by nix/{N;d;}' $root/.git/info/exclude
               ${makeClauses}
             fi
+
+            # libgit2-sys copies a vendored git2 into the target/
+            # directory somewhere. In certain, rare, cases,
+            # libgit2-sys is rebuilt in the same incremental dep
+            # directory as it was before, and then this copy fails,
+            # because the files, copied from the nix store, already
+            # exist and do not have w permission. Hack around this
+            # issue by making any existing libgit2-sys vendored git2
+            # files writable before a build can be run
+            find "$(base_cargo metadata --format-version 1 | jq -r '.target_directory')" -path '*/build/libgit2-sys-*/out/include' -print0 | xargs -r -0 chmod u+w -R
           '';
 
           # Hyperlight scripts use cargo in a bunch of ways that don't

From 22b433a8f5eb745f7ac7e3b1a546be0d80e884fa Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 16:22:25 +0000
Subject: [PATCH 04/20] [just] Bring `just like-ci` in line with CI changes

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 Justfile  | 72 ++++++++++++++++++++++++++++++++++++++++++-------------
 flake.nix |  1 +
 2 files changed, 56 insertions(+), 17 deletions(-)

diff --git a/Justfile b/Justfile
index e11c16ec3..e25900647 100644
--- a/Justfile
+++ b/Justfile
@@ -99,16 +99,10 @@ test-like-ci config=default-target hypervisor="kvm":
     @# test the tracing related features
     {{ if os() == "linux" { "just test-rust-tracing " + config + " " + if hypervisor == "mshv3" { "mshv3" } else { "kvm" } } else { "" } }}
 
-like-ci config=default-target hypervisor="kvm":
+code-checks-like-ci config=default-target hypervisor="kvm":
     @# Ensure up-to-date Cargo.lock
     cargo fetch --locked
 
-    @# typos
-    typos
-
-    @# check licence headers
-    just check-license-headers
-
     @# fmt
     just fmt-check
 
@@ -122,6 +116,17 @@ like-ci config=default-target hypervisor="kvm":
     @# Verify MSRV
     ./dev/verify-msrv.sh hyperlight-common hyperlight-guest hyperlight-guest-bin hyperlight-host hyperlight-component-util hyperlight-component-macro hyperlight-guest-tracing
 
+    @# Check 32-bit guests
+    {{ if os() == "linux" { "just check-i686 " + config } else { "" } }}
+
+    @# Check cargo features compile
+    just check
+
+    @# Check compilation with no default features
+    just test-compilation-no-default-features debug
+    just test-compilation-no-default-features release
+
+build-guests-like-ci config=default-target hypervisor="kvm":
     @# Build and move Rust guests
     just build-rust-guests {{config}}
     just move-rust-guests {{config}}
@@ -130,29 +135,62 @@ like-ci config=default-target hypervisor="kvm":
     just build-c-guests {{config}}
     just move-c-guests {{config}}
 
+build-test-like-ci config=default-target hypervisor="kvm":
     @# Build
     just build {{config}}
 
+    @# Run Miri tests
+    {{ if os() == "linux" { "just miri-tests" } else { "" } }}
+
     @# Run Rust tests
-    just test-like-ci {{config}} {{hypervisor}}
+    just test {{config}}
+
+    @# Run Rust tests with single driver
+    {{ if os() == "linux" { "just test " + config+ " " + if hypervisor == "mshv3" { "mshv3" } else { "kvm" } } else { "" } }}
 
+    @# Run Rust Gdb tests
+    just test-rust-gdb-debugging {{config}}
+
+    @# Run Rust Crashdump tests
+    just test-rust-crashdump {{config}}
+
+    @# Run Rust Tracing tests
+    {{ if os() == "linux" { "just test-rust-tracing " + config } else { "" } }}
+
+run-examples-like-ci config=default-target hypervisor="kvm":
     @# Run Rust examples - Windows
     {{ if os() == "windows" { "just run-rust-examples " + config } else { "" } }}
 
     @# Run Rust examples - linux
-    {{ if os() == "linux" { "just run-rust-examples-linux " + config + " " + if hypervisor == "mshv3" { "mshv3" } else { "kvm" } } else { "" } }}
+    {{ if os() == "linux" { "just run-rust-examples-linux " + config + " " } else { "" } }}
 
-    @# Run Rust Gdb tests
-    just test-rust-gdb-debugging {{ config }} {{ if hypervisor == "mshv3" { "mshv3" } else { "kvm" } }} 
+benchmarks-like-ci config=default-target hypervisor="$vm":
+    @# Run benchmarks
+    {{ if config == "release" { "just bench-ci main" } else { "" } }}
 
-    @# Run Rust Crashdump tests
-    just test-rust-crashdump {{config}} {{ if hypervisor == "mshv3" { "mshv3" } else { "kvm" } }}
+like-ci config=default-target hypervisor="kvm":
+    @# .github/workflows/dep_code_checks.yml
+    just code-checks-like-ci {{config}} {{hypervisor}}
 
-    @# Run Rust Tracing tests - linux
-    {{ if os() == "linux" { "just test-rust-tracing " + config + " " + if hypervisor == "mshv3" { "mshv3" } else { "kvm" } } else { "" } }}
+    @# .github/workflows/dep_build_guests.yml
+    just build-guests-like-ci {{config}} {{hypervisor}}
 
-    @# Run benchmarks
-    {{ if config == "release" { "just bench-ci main " + if hypervisor == "mshv3" { "mshv3" } else { "kvm" } } else { "" } }}
+    @# .github/workflows/dep_build_test.yml
+    just build-test-like-ci {{config}} {{hypervisor}}
+
+    @# .github/workflows/dep_run_examples.yml
+    just run-examples-like-ci {{config}} {{hypervisor}}
+
+    @# .github/workflows/dep_benchmarks.yml
+    just benchmarks-like-ci {{config}} {{hypervisor}}
+
+    @# can't run fuzzing locally
+
+    @# spelling
+    typos
+
+    @# license-headers
+    just check-license-headers
 
 # runs all tests
 test target=default-target features="": (test-unit target features) (test-isolated target features) (test-integration "rust" target features) (test-integration "c" target features) (test-doc target features)
diff --git a/flake.nix b/flake.nix
index 05ddd9561..6e4c845aa 100644
--- a/flake.nix
+++ b/flake.nix
@@ -67,6 +67,7 @@
                 "x86_64-unknown-linux-gnu"
                 "x86_64-pc-windows-msvc" "x86_64-unknown-none"
                 "wasm32-wasip1" "wasm32-wasip2" "wasm32-unknown-unknown"
+                "i686-unknown-linux-gnu"
               ];
               extensions = [ "rust-src" ] ++ (if args.channel == "nightly" then [ "miri-preview" ] else []);
             });

From 415823f0086044bb723dd3f3a0817b522221f5c1 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 16:34:48 +0000
Subject: [PATCH 05/20] [kvm] Better document the cpuid leaf used to set GPA
 width

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 .../src/hypervisor/virtual_machine/kvm.rs     | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs
index 190cb936c..f2115b451 100644
--- a/src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs
+++ b/src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs
@@ -38,6 +38,24 @@ use crate::mem::memory_region::MemoryRegion;
 #[cfg(feature = "trace_guest")]
 use crate::sandbox::trace::TraceContext as SandboxTraceContext;
 
+/// On KVM x86-64 only, we have to set this in order to set the guest
+/// physical address width.
+///
+/// The requirement to set this to configure the guest physical
+/// address width for KVM is not well documented, but see e.g. Linux
+/// v6.18.6 arch/x86/kvm/cpuid.c:kvm_vcpu_after_set_cpuid()
+/// (https://elixir.bootlin.com/linux/v6.18.6/source/arch/x86/kvm/cpuid.c#L444)
+/// for how it is processed.
+///
+/// For the architectural definition and format of the system register:
+/// See AMD64 Architecture Programmer's Manual, Volume 3: General-Purpose and
+///                                                       System Instructions
+///     Appendix E: Obtaining Processor Information Via the CPUID Instruction
+///         E.4.7: Function 8000_0008h---Processor Capacity Parameters and
+///                Extended Feature Identification, pp. 627--628
+const CPUID_FUNCTION_PROCESSOR_CAPACITY_PARAMETERS_AND_EXTENDED_FEATURE_IDENTIFICATION: u32 =
+    0x8000_0008;
+
 /// Return `true` if the KVM API is available, version 12, and has UserMemory capability, or `false` otherwise
 #[instrument(skip_all, parent = Span::current(), level = "Trace")]
 pub(crate) fn is_hypervisor_present() -> bool {
@@ -95,7 +113,9 @@ impl KvmVm {
             .get_supported_cpuid(kvm_bindings::KVM_MAX_CPUID_ENTRIES)
             .map_err(|e| CreateVmError::InitializeVm(e.into()))?;
         for entry in kvm_cpuid.as_mut_slice().iter_mut() {
-            if entry.function == 0x8000_0008 {
+            if entry.function
+                == CPUID_FUNCTION_PROCESSOR_CAPACITY_PARAMETERS_AND_EXTENDED_FEATURE_IDENTIFICATION
+            {
                 entry.eax &= !0xff;
                 entry.eax |= hyperlight_common::layout::MAX_GPA.ilog2() + 1;
             }

From 665a86f8c85d862e5788cdc77be79b8bb1be3247 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Wed, 28 Jan 2026 03:08:55 +0000
Subject: [PATCH 06/20] [whp] Remove initial_memory_setup_done, which is no
 longer used

whp can now map any region as long as it contains a file
handle (which, presently, only snapshot and scratch regions do---not
user-provided regions).

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/hypervisor/hyperlight_vm.rs    |  4 ----
 .../src/hypervisor/virtual_machine/mod.rs              |  5 -----
 .../src/hypervisor/virtual_machine/whp.rs              | 10 ----------
 3 files changed, 19 deletions(-)

diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs
index ff9d5c263..44e1c1c75 100644
--- a/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs
+++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs
@@ -367,10 +367,6 @@ impl HyperlightVm {
             };
         }
 
-        // Mark initial setup as complete for Windows - subsequent map_memory calls will fail
-        #[cfg(target_os = "windows")]
-        vm.complete_initial_memory_setup();
-
         #[cfg(feature = "init-paging")]
         vm.set_sregs(&CommonSpecialRegisters::standard_64bit_defaults(_pml4_addr))
             .map_err(VmError::Register)?;
diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs
index f77b09ad2..ced8a7dfc 100644
--- a/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs
+++ b/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs
@@ -324,11 +324,6 @@ pub(crate) trait VirtualMachine: Debug + Send {
     /// Get partition handle
     #[cfg(target_os = "windows")]
     fn partition_handle(&self) -> windows::Win32::System::Hypervisor::WHV_PARTITION_HANDLE;
-
-    /// Mark that initial memory setup is complete. After this, map_memory will fail.
-    /// This is only needed on Windows where dynamic memory mapping is not yet supported.
-    #[cfg(target_os = "windows")]
-    fn complete_initial_memory_setup(&mut self);
 }
 
 #[cfg(test)]
diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs
index 572641361..2bcad8bcd 100644
--- a/src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs
+++ b/src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs
@@ -69,10 +69,6 @@ pub(crate) struct WhpVm {
     partition: WHV_PARTITION_HANDLE,
     // Surrogate process for memory mapping
     surrogate_process: SurrogateProcess,
-    // Track if initial memory setup is complete.
-    // Used to reject later memory mapping since it's not supported  on windows.
-    // TODO remove this flag once memory mapping is supported on windows.
-    initial_memory_setup_done: bool,
 }
 
 // Safety: `WhpVm` is !Send because it holds `SurrogateProcess` which contains a raw pointer
@@ -111,7 +107,6 @@ impl WhpVm {
         Ok(WhpVm {
             partition,
             surrogate_process,
-            initial_memory_setup_done: false,
         })
     }
 
@@ -505,11 +500,6 @@ impl VirtualMachine for WhpVm {
         Ok(xsave_buffer)
     }
 
-    /// Mark that initial memory setup is complete. After this, map_memory will fail.
-    fn complete_initial_memory_setup(&mut self) {
-        self.initial_memory_setup_done = true;
-    }
-
     /// Get the partition handle for this VM
     fn partition_handle(&self) -> WHV_PARTITION_HANDLE {
         self.partition

From 2cc3e55477253593cadac12936bf8bbeb278a2fd Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 15:59:12 +0000
Subject: [PATCH 07/20] Make scratch region writable again

This fixes an error introduced by a merge conflict in a previous PR.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/mem/shared_mem.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/hyperlight_host/src/mem/shared_mem.rs b/src/hyperlight_host/src/mem/shared_mem.rs
index a69394c23..9d068bca0 100644
--- a/src/hyperlight_host/src/mem/shared_mem.rs
+++ b/src/hyperlight_host/src/mem/shared_mem.rs
@@ -655,7 +655,9 @@ impl GuestSharedMemory {
         region_type: MemoryRegionType,
     ) -> MemoryRegion {
         let flags = match region_type {
-            MemoryRegionType::Scratch => MemoryRegionFlags::READ | MemoryRegionFlags::EXECUTE,
+            MemoryRegionType::Scratch => {
+                MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE
+            }
             #[allow(clippy::panic)]
             // In the future, all the host side knowledge about memory
             // region types should collapse down to Snapshot vs

From 7fccad62c7bbb8a167c4efd05af3995aa2546e32 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Wed, 28 Jan 2026 06:52:50 +0000
Subject: [PATCH 08/20] Use MAP_SHARED for shared memory creation again

Previously, in 7ca0a78f, MAP_SHARED was replaced with MAP_PRIVATE,
primarily to bring the Miri and non-Miri code paths closer together,
since it did not appear to have any detrimental effect.  However, this
seems to cause issues with the scratch mapping on mshv on Linux. It is
not entirely clear what mshv is doing here; possibly something like
"mapping a page which has never been written to" causes the problem.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/mem/shared_mem.rs | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/hyperlight_host/src/mem/shared_mem.rs b/src/hyperlight_host/src/mem/shared_mem.rs
index 9d068bca0..f9280b773 100644
--- a/src/hyperlight_host/src/mem/shared_mem.rs
+++ b/src/hyperlight_host/src/mem/shared_mem.rs
@@ -314,12 +314,11 @@ impl ExclusiveSharedMemory {
     #[cfg(target_os = "linux")]
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn new(min_size_bytes: usize) -> Result<Self> {
-        use libc::{
-            MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, PROT_READ, PROT_WRITE, c_int, mmap, off_t,
-            size_t,
-        };
+        #[cfg(miri)]
+        use libc::MAP_PRIVATE;
+        use libc::{MAP_ANONYMOUS, MAP_FAILED, PROT_READ, PROT_WRITE, c_int, mmap, off_t, size_t};
         #[cfg(not(miri))]
-        use libc::{MAP_NORESERVE, PROT_NONE, mprotect};
+        use libc::{MAP_NORESERVE, MAP_SHARED, PROT_NONE, mprotect};
 
         if min_size_bytes == 0 {
             return Err(new_error!("Cannot create shared memory with size 0"));
@@ -347,7 +346,7 @@ impl ExclusiveSharedMemory {
 
         // allocate the memory
         #[cfg(not(miri))]
-        let flags = MAP_ANONYMOUS | MAP_PRIVATE | MAP_NORESERVE;
+        let flags = MAP_ANONYMOUS | MAP_SHARED | MAP_NORESERVE;
         #[cfg(miri)]
         let flags = MAP_ANONYMOUS | MAP_PRIVATE;
 

From a9355b17f31948f06287b5ed820cb43dc0041d1d Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 16:05:39 +0000
Subject: [PATCH 09/20] Modify guest physical page allocator to allocate from
 the scratch region

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_guest/src/arch/amd64/layout.rs | 34 ++++++++
 .../src/arch/amd64/prim_alloc.rs              | 36 ++++++++
 src/hyperlight_guest/src/layout.rs            | 29 +++++++
 src/hyperlight_guest/src/lib.rs               |  2 +
 src/hyperlight_guest/src/prim_alloc.rs        | 36 ++++++++
 src/hyperlight_guest_bin/src/paging.rs        | 25 +-----
 src/hyperlight_host/src/mem/mgr.rs            | 87 +++++++++++++------
 7 files changed, 197 insertions(+), 52 deletions(-)
 create mode 100644 src/hyperlight_guest/src/arch/amd64/layout.rs
 create mode 100644 src/hyperlight_guest/src/arch/amd64/prim_alloc.rs
 create mode 100644 src/hyperlight_guest/src/layout.rs
 create mode 100644 src/hyperlight_guest/src/prim_alloc.rs

diff --git a/src/hyperlight_guest/src/arch/amd64/layout.rs b/src/hyperlight_guest/src/arch/amd64/layout.rs
new file mode 100644
index 000000000..9f83382a4
--- /dev/null
+++ b/src/hyperlight_guest/src/arch/amd64/layout.rs
@@ -0,0 +1,34 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+pub const MAIN_STACK_TOP_GVA: usize = 0xffff_feff_ffff_f000;
+
+pub fn scratch_size() -> u64 {
+    let addr = crate::layout::scratch_size_gva();
+    let x: u64;
+    unsafe {
+        core::arch::asm!("mov {x}, [{addr}]", x = out(reg) x, addr = in(reg) addr);
+    }
+    x
+}
+
+pub fn scratch_base_gpa() -> u64 {
+    hyperlight_common::layout::scratch_base_gpa(scratch_size() as usize)
+}
+
+pub fn scratch_base_gva() -> u64 {
+    hyperlight_common::layout::scratch_base_gva(scratch_size() as usize)
+}
diff --git a/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs b/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs
new file mode 100644
index 000000000..049d92b7f
--- /dev/null
+++ b/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs
@@ -0,0 +1,36 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+// There are no notable architecture-specific safety considerations
+// here, and the general conditions are documented in the
+// architecture-independent re-export in prim_alloc.rs
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn alloc_phys_pages(n: u64) -> u64 {
+    let addr = crate::layout::allocator_gva();
+    let nbytes = n * hyperlight_common::vmem::PAGE_SIZE as u64;
+    let mut x = nbytes;
+    unsafe {
+        core::arch::asm!(
+            "lock xadd qword ptr [{addr}], {x}",
+            addr = in(reg) addr,
+            x = inout(reg) x
+        );
+    }
+    if x.checked_add(nbytes).is_none() {
+        panic!("Out of physical memory!")
+    }
+    x
+}
diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs
new file mode 100644
index 000000000..f28dcb0cf
--- /dev/null
+++ b/src/hyperlight_guest/src/layout.rs
@@ -0,0 +1,29 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+#[cfg_attr(target_arch = "x86_64", path = "arch/amd64/layout.rs")]
+mod arch;
+
+pub use arch::MAIN_STACK_TOP_GVA;
+pub fn scratch_size_gva() -> *mut u64 {
+    use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_SIZE_OFFSET};
+    (MAX_GVA as u64 - SCRATCH_TOP_SIZE_OFFSET + 1) as *mut u64
+}
+pub fn allocator_gva() -> *mut u64 {
+    use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_ALLOCATOR_OFFSET};
+    (MAX_GVA as u64 - SCRATCH_TOP_ALLOCATOR_OFFSET + 1) as *mut u64
+}
+pub use arch::{scratch_base_gpa, scratch_base_gva};
diff --git a/src/hyperlight_guest/src/lib.rs b/src/hyperlight_guest/src/lib.rs
index fce0767c0..951cdedd2 100644
--- a/src/hyperlight_guest/src/lib.rs
+++ b/src/hyperlight_guest/src/lib.rs
@@ -23,6 +23,8 @@ extern crate alloc;
 // Modules
 pub mod error;
 pub mod exit;
+pub mod layout;
+pub mod prim_alloc;
 
 pub mod guest_handle {
     pub mod handle;
diff --git a/src/hyperlight_guest/src/prim_alloc.rs b/src/hyperlight_guest/src/prim_alloc.rs
new file mode 100644
index 000000000..f58a80e13
--- /dev/null
+++ b/src/hyperlight_guest/src/prim_alloc.rs
@@ -0,0 +1,36 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+#[cfg_attr(target_arch = "x86_64", path = "arch/amd64/prim_alloc.rs")]
+mod arch;
+
+/// Allocate n contiguous physical pages and return the physical
+/// addresses of the pages in question.
+/// # Safety
+/// Since this reads and writes specific allocator state addresses, it
+/// is only safe when the allocator has been set up properly. It may
+/// become less safe in the future.
+///
+/// # Panics
+/// This function will panic if memory allocation fails
+///
+/// This is defined in an arch-specific module because it reads and
+/// writes the actual allocator state with inline assembly in order to
+/// access it atomically according to the architecture memory model
+/// rather than the Rust memory model: the stronger constraints of the
+/// latter cannot be perfectly satisfied due to the lack of per-byte
+/// atomic memcpy in the host.
+pub use arch::alloc_phys_pages;
diff --git a/src/hyperlight_guest_bin/src/paging.rs b/src/hyperlight_guest_bin/src/paging.rs
index 09cb85cf3..43adf3e46 100644
--- a/src/hyperlight_guest_bin/src/paging.rs
+++ b/src/hyperlight_guest_bin/src/paging.rs
@@ -14,12 +14,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use alloc::alloc::Layout;
 use core::arch::asm;
 
+use hyperlight_guest::prim_alloc::alloc_phys_pages;
 use tracing::{Span, instrument};
 
-use crate::OS_PAGE_SIZE;
 
 /// Convert a physical address in main memory to a virtual address
 /// through the pysmap
@@ -137,28 +136,6 @@ pub unsafe fn map_region(phys_base: u64, virt_base: *mut u8, len: u64) {
     }
 }
 
-/// Allocate n contiguous physical pages and return the physical
-/// addresses of the pages in question.
-/// # Safety
-/// This function is not inherently unsafe but will likely become so in the future
-/// when a real physical page allocator is implemented.
-/// # Panics
-/// This function will panic if:
-/// - The Layout creation fails
-/// - Memory allocation fails
-pub unsafe fn alloc_phys_pages(n: u64) -> u64 {
-    // Currently, since all of main memory is idmap'd, we can just
-    // allocate any appropriately aligned section of memory.
-    unsafe {
-        let v = alloc::alloc::alloc_zeroed(
-            Layout::from_size_align(n as usize * OS_PAGE_SIZE as usize, OS_PAGE_SIZE as usize)
-                .expect("could not create physical page allocation layout"),
-        );
-        if v.is_null() {
-            panic!("could not allocate a physical page");
-        }
-        v as u64
-    }
 }
 
 pub fn flush_tlb() {
diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs
index 0fedab3aa..d67430edb 100644
--- a/src/hyperlight_host/src/mem/mgr.rs
+++ b/src/hyperlight_host/src/mem/mgr.rs
@@ -95,13 +95,18 @@ impl vmem::TableOps for GuestPageTableBuffer {
             .unwrap_or(0)
     }
 
-    unsafe fn write_entry(&self, addr: (usize, usize), x: PageTableEntry) {
+    unsafe fn write_entry(
+        &self,
+        addr: (usize, usize),
+        entry: PageTableEntry,
+    ) -> Option<(usize, usize)> {
         let mut b = self.buffer.borrow_mut();
         let byte_offset =
             (addr.0 - self.phys_base / PAGE_TABLE_SIZE) * PAGE_TABLE_SIZE + addr.1 * 8;
         if let Some(slice) = b.get_mut(byte_offset..byte_offset + 8) {
-            slice.copy_from_slice(&x.to_ne_bytes());
+            slice.copy_from_slice(&entry.to_ne_bytes());
         }
+        None
     }
 
     fn to_phys(addr: (usize, usize)) -> PhysAddr {
@@ -237,28 +242,28 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
     ) {
         let (hshm, gshm) = self.shared_mem.build();
         let (hscratch, gscratch) = self.scratch_mem.build();
-        (
-            SandboxMemoryManager {
-                shared_mem: hshm,
-                scratch_mem: hscratch,
-                layout: self.layout,
-                load_addr: self.load_addr.clone(),
-                entrypoint_offset: self.entrypoint_offset,
-                mapped_rgns: self.mapped_rgns,
-                stack_cookie: self.stack_cookie,
-                abort_buffer: self.abort_buffer,
-            },
-            SandboxMemoryManager {
-                shared_mem: gshm,
-                scratch_mem: gscratch,
-                layout: self.layout,
-                load_addr: self.load_addr.clone(),
-                entrypoint_offset: self.entrypoint_offset,
-                mapped_rgns: self.mapped_rgns,
-                stack_cookie: self.stack_cookie,
-                abort_buffer: Vec::new(), // Guest doesn't need abort buffer
-            },
-        )
+        let mut host_mgr = SandboxMemoryManager {
+            shared_mem: hshm,
+            scratch_mem: hscratch,
+            layout: self.layout,
+            load_addr: self.load_addr.clone(),
+            entrypoint_offset: self.entrypoint_offset,
+            mapped_rgns: self.mapped_rgns,
+            stack_cookie: self.stack_cookie,
+            abort_buffer: self.abort_buffer,
+        };
+        let guest_mgr = SandboxMemoryManager {
+            shared_mem: gshm,
+            scratch_mem: gscratch,
+            layout: self.layout,
+            load_addr: self.load_addr.clone(),
+            entrypoint_offset: self.entrypoint_offset,
+            mapped_rgns: self.mapped_rgns,
+            stack_cookie: self.stack_cookie,
+            abort_buffer: Vec::new(), // Guest doesn't need abort buffer
+        };
+        host_mgr.update_scratch_bookkeeping();
+        (host_mgr, guest_mgr)
     }
 }
 
@@ -400,9 +405,9 @@ impl SandboxMemoryManager<HostSharedMemory> {
         }
         self.shared_mem.restore_from_snapshot(snapshot)?;
         let new_scratch_size = snapshot.layout().get_scratch_size();
-        if new_scratch_size == self.scratch_mem.mem_size() {
+        let gscratch = if new_scratch_size == self.scratch_mem.mem_size() {
             self.scratch_mem.zero()?;
-            Ok(None)
+            None
         } else {
             let new_scratch_mem = ExclusiveSharedMemory::new(new_scratch_size)?;
             let (hscratch, gscratch) = new_scratch_mem.build();
@@ -413,8 +418,34 @@ impl SandboxMemoryManager<HostSharedMemory> {
             // has been unmapped from the VM.
             self.scratch_mem = hscratch;
 
-            Ok(Some(gscratch))
-        }
+            Some(gscratch)
+        };
+        self.update_scratch_bookkeeping();
+        Ok(gscratch)
+    }
+
+    fn update_scratch_bookkeeping(&mut self) {
+        let scratch_size = self.scratch_mem.mem_size();
+        let size_offset =
+            scratch_size - hyperlight_common::layout::SCRATCH_TOP_SIZE_OFFSET as usize;
+        // The only way that write can fail is if the offset is
+        // outside of the memory, which would be sufficiently much of
+        // an invariant violation that panicking is probably
+        // sensible...
+        #[allow(clippy::unwrap_used)]
+        self.scratch_mem
+            .write::<u64>(size_offset, scratch_size as u64)
+            .unwrap();
+        let alloc_offset =
+            scratch_size - hyperlight_common::layout::SCRATCH_TOP_ALLOCATOR_OFFSET as usize;
+        // See above comment about unwrap() on write
+        #[allow(clippy::unwrap_used)]
+        self.scratch_mem
+            .write::<u64>(
+                alloc_offset,
+                hyperlight_common::layout::scratch_base_gpa(scratch_size),
+            )
+            .unwrap();
     }
 }
 

From b6b839d1402429c521e31a764d5212250dd16e47 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 18:11:08 +0000
Subject: [PATCH 10/20] i686: add layout, vmem, etc stubs

As per the comments preexisting in
src/hyperlight_common/src/layout.rs. This avoids some fairly
cumbersome proliferation of #[cfg(feature = "init-paging")], and may
be developed further in the future.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_common/src/arch/i686/layout.rs | 22 +++++++++++
 src/hyperlight_common/src/arch/i686/vmem.rs   | 38 +++++++++++++++++++
 src/hyperlight_common/src/layout.rs           | 11 +-----
 src/hyperlight_common/src/lib.rs              |  2 +-
 src/hyperlight_common/src/vmem.rs             |  6 ++-
 src/hyperlight_guest/src/arch/i686/layout.rs  | 32 ++++++++++++++++
 .../src/arch/i686/prim_alloc.rs               | 25 ++++++++++++
 src/hyperlight_guest/src/layout.rs            |  1 +
 src/hyperlight_guest/src/prim_alloc.rs        |  1 +
 9 files changed, 127 insertions(+), 11 deletions(-)
 create mode 100644 src/hyperlight_common/src/arch/i686/layout.rs
 create mode 100644 src/hyperlight_common/src/arch/i686/vmem.rs
 create mode 100644 src/hyperlight_guest/src/arch/i686/layout.rs
 create mode 100644 src/hyperlight_guest/src/arch/i686/prim_alloc.rs

diff --git a/src/hyperlight_common/src/arch/i686/layout.rs b/src/hyperlight_common/src/arch/i686/layout.rs
new file mode 100644
index 000000000..c70f38e34
--- /dev/null
+++ b/src/hyperlight_common/src/arch/i686/layout.rs
@@ -0,0 +1,22 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+// This file is just dummy definitions at the moment, in order to
+// allow compiling the guest for real mode boot scenarios.
+
+pub const MAX_GVA: usize = 0xffff_efff;
+pub const SNAPSHOT_PT_GVA: usize = 0xefff_efff;
+pub const MAX_GPA: usize = 0xffff_ffff;
diff --git a/src/hyperlight_common/src/arch/i686/vmem.rs b/src/hyperlight_common/src/arch/i686/vmem.rs
new file mode 100644
index 000000000..1a1e4974c
--- /dev/null
+++ b/src/hyperlight_common/src/arch/i686/vmem.rs
@@ -0,0 +1,38 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+// This file is just dummy definitions at the moment, in order to
+// allow compiling the guest for real mode boot scenarios.
+
+use crate::vmem::{BasicMapping, Mapping, TableOps};
+
+pub const PAGE_SIZE: usize = 4096;
+pub const PAGE_TABLE_SIZE: usize = 4096;
+pub type PageTableEntry = u32;
+pub type VirtAddr = u32;
+pub type PhysAddr = u32;
+
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn map<Op: TableOps>(_op: &Op, _mapping: Mapping) {
+    panic!("vmem::map: i686 guests do not support booting the full hyperlight guest kernel");
+}
+
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn virt_to_phys<Op: TableOps>(_op: &Op, _address: u64) -> Option<(u64, BasicMapping)> {
+    panic!(
+        "vmem::virt_to_phys: i686 guests do not support booting the full hyperlight guest kernel"
+    );
+}
diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs
index 3ba739b73..645f6217e 100644
--- a/src/hyperlight_common/src/layout.rs
+++ b/src/hyperlight_common/src/layout.rs
@@ -14,16 +14,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
  */
 
-// The constraint on the feature is temporary and will be removed when other arch i686 is added
 #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/layout.rs")]
-#[cfg(feature = "init-paging")]
+#[cfg_attr(target_arch = "x86", path = "arch/i686/layout.rs")]
 mod arch;
 
-// The constraint on the feature is temporary and will be removed when other arch i686 is added
-#[cfg(feature = "init-paging")]
-pub use arch::MAX_GPA;
-#[cfg(feature = "init-paging")]
-pub use arch::{MAX_GVA, SNAPSHOT_PT_GVA};
+pub use arch::{MAX_GPA, MAX_GVA, SNAPSHOT_PT_GVA};
 
 // offsets down from the top of scratch memory for various things
 pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08;
@@ -31,11 +26,9 @@ pub const SCRATCH_TOP_USED_OFFSET: u64 = 0x10;
 pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x18;
 pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x20;
 
-#[cfg(feature = "init-paging")]
 pub fn scratch_base_gpa(size: usize) -> u64 {
     (MAX_GPA - size + 1) as u64
 }
-#[cfg(feature = "init-paging")]
 pub fn scratch_base_gva(size: usize) -> u64 {
     (MAX_GVA - size + 1) as u64
 }
diff --git a/src/hyperlight_common/src/lib.rs b/src/hyperlight_common/src/lib.rs
index da9294c19..478aeef8b 100644
--- a/src/hyperlight_common/src/lib.rs
+++ b/src/hyperlight_common/src/lib.rs
@@ -41,6 +41,6 @@ pub mod resource;
 
 /// cbindgen:ignore
 pub mod func;
+
 // cbindgen:ignore
-#[cfg(feature = "init-paging")]
 pub mod vmem;
diff --git a/src/hyperlight_common/src/vmem.rs b/src/hyperlight_common/src/vmem.rs
index 9d2890086..abebaf61c 100644
--- a/src/hyperlight_common/src/vmem.rs
+++ b/src/hyperlight_common/src/vmem.rs
@@ -15,9 +15,13 @@ limitations under the License.
  */
 
 #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/vmem.rs")]
+#[cfg_attr(target_arch = "x86", path = "arch/i686/vmem.rs")]
 pub mod arch;
 
-pub use arch::{PAGE_SIZE, PAGE_TABLE_SIZE, PageTableEntry, PhysAddr, VirtAddr};
+/// This is always the page size that the /guest/ is being compiled
+/// for, which may or may not be the same as the host page size.
+pub use arch::PAGE_SIZE;
+pub use arch::{PAGE_TABLE_SIZE, PageTableEntry, PhysAddr, VirtAddr};
 pub const PAGE_TABLE_ENTRIES_PER_TABLE: usize =
     PAGE_TABLE_SIZE / core::mem::size_of::<PageTableEntry>();
 
diff --git a/src/hyperlight_guest/src/arch/i686/layout.rs b/src/hyperlight_guest/src/arch/i686/layout.rs
new file mode 100644
index 000000000..7ca2f4fcf
--- /dev/null
+++ b/src/hyperlight_guest/src/arch/i686/layout.rs
@@ -0,0 +1,32 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+// This file is just dummy definitions at the moment, in order to
+// allow compiling the guest for real mode boot scenarios.
+
+pub const MAIN_STACK_TOP_GVA: usize = 0xdfff_efff;
+
+pub fn scratch_size() -> u64 {
+    hyperlight_common::vmem::PAGE_SIZE as u64
+}
+
+pub fn scratch_base_gpa() -> u64 {
+    hyperlight_common::layout::scratch_base_gpa(scratch_size() as usize)
+}
+
+pub fn scratch_base_gva() -> u64 {
+    hyperlight_common::layout::scratch_base_gva(scratch_size() as usize)
+}
diff --git a/src/hyperlight_guest/src/arch/i686/prim_alloc.rs b/src/hyperlight_guest/src/arch/i686/prim_alloc.rs
new file mode 100644
index 000000000..5f22b7f56
--- /dev/null
+++ b/src/hyperlight_guest/src/arch/i686/prim_alloc.rs
@@ -0,0 +1,25 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+// This file is just dummy definitions at the moment, in order to
+// allow compiling the guest for real mode boot scenarios.
+
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn alloc_phys_pages(_n: u64) -> u64 {
+    panic!(
+        "prim_alloc::alloc_phys_pages: i686 guests do not support booting the full hyperlight guest kernel"
+    );
+}
diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs
index f28dcb0cf..e7ac0a5e3 100644
--- a/src/hyperlight_guest/src/layout.rs
+++ b/src/hyperlight_guest/src/layout.rs
@@ -15,6 +15,7 @@ limitations under the License.
  */
 
 #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/layout.rs")]
+#[cfg_attr(target_arch = "x86", path = "arch/i686/layout.rs")]
 mod arch;
 
 pub use arch::MAIN_STACK_TOP_GVA;
diff --git a/src/hyperlight_guest/src/prim_alloc.rs b/src/hyperlight_guest/src/prim_alloc.rs
index f58a80e13..3757d1711 100644
--- a/src/hyperlight_guest/src/prim_alloc.rs
+++ b/src/hyperlight_guest/src/prim_alloc.rs
@@ -15,6 +15,7 @@ limitations under the License.
  */
 
 #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/prim_alloc.rs")]
+#[cfg_attr(target_arch = "x86", path = "arch/i686/prim_alloc.rs")]
 mod arch;
 
 /// Allocate n contiguous physical pages and return the physical

From f38272f237786e342cc788e866f90a13584eb568 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 18:14:35 +0000
Subject: [PATCH 11/20] Move guest-written page tables to the scratch region

Instead of updating page tables in place, and using the heap region
memory for new page tables, this changes the hyperlight_guest virtual
memory subsystem to treat snapshot page tables as copy-on-write and to
allocate all new page tables in the scratch region.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_common/src/arch/amd64/vmem.rs | 217 ++++++++++++++-----
 src/hyperlight_common/src/vmem.rs            |  11 +-
 src/hyperlight_guest_bin/src/paging.rs       |  54 +++--
 3 files changed, 211 insertions(+), 71 deletions(-)

diff --git a/src/hyperlight_common/src/arch/amd64/vmem.rs b/src/hyperlight_common/src/arch/amd64/vmem.rs
index 4e104460e..ed0c0611c 100644
--- a/src/hyperlight_common/src/arch/amd64/vmem.rs
+++ b/src/hyperlight_common/src/arch/amd64/vmem.rs
@@ -23,9 +23,9 @@ limitations under the License.
 //! - PT (Page Table) - bits 20:12 - 512 entries, each covering 4KB pages
 //!
 //! The code uses an iterator-based approach to walk the page table hierarchy,
-//! allocating intermediate tables as needed and setting appropriate flags on leaf PTEs.
+//! allocating intermediate tables as needed and setting appropriate flags on leaf PTEs
 
-use crate::vmem::{Mapping, MappingKind, TableOps};
+use crate::vmem::{BasicMapping, Mapping, MappingKind, TableOps};
 
 // Paging Flags
 //
@@ -94,20 +94,114 @@ fn bits<const HIGH_BIT: u8, const LOW_BIT: u8>(x: u64) -> u64 {
     (x & ((1 << (HIGH_BIT + 1)) - 1)) >> LOW_BIT
 }
 
+/// Helper function to generate a page table entry that points to another table
+#[allow(clippy::identity_op)]
+#[allow(clippy::precedence)]
+fn pte_for_table<Op: TableOps>(table_addr: Op::TableAddr) -> u64 {
+    Op::to_phys(table_addr) |
+        PAGE_ACCESSED_CLEAR | // accessed bit cleared (will be set by CPU when page is accessed - but we dont use the access bit for anything at present)
+        PAGE_CACHE_ENABLED | // leave caching enabled
+        PAGE_WRITE_BACK | // use write-back caching
+        PAGE_USER_ACCESS_DISABLED |// dont allow user access (no code runs in user mode for now)
+        PAGE_RW | // R/W - we don't use block-level permissions
+        PAGE_PRESENT // P   - this entry is present
+}
+
+/// Helper function to write a page table entry, updating the whole
+/// chain of tables back to the root if necessary
+unsafe fn write_entry_updating<Op: TableOps, P: UpdateParent<Op>>(
+    op: &Op,
+    parent: P,
+    addr: Op::TableAddr,
+    entry: u64,
+) {
+    if let Some(again) = unsafe { op.write_entry(addr, entry) } {
+        parent.update_parent(op, again);
+    }
+}
+
+/// A helper trait that allows us to move a page table (e.g. from the
+/// snapshot to the scratch region), keeping track of the context that
+/// needs to be updated when that is moved (and potentially
+/// recursively updating, if necessary)
+///
+/// This is done via a trait so that the selected impl knows the exact
+/// nesting depth of tables, in order to assist
+/// inlining/specialisation in generating efficient code.
+trait UpdateParent<Op: TableOps>: Copy {
+    fn update_parent(self, op: &Op, new_ptr: Op::TableAddr);
+}
+
+/// A struct implementing [`UpdateParent`] that keeps track of the
+/// fact that the parent table is itself another table, whose own
+/// ancestors may need to be recursively updated
+struct UpdateParentTable<Op: TableOps, P: UpdateParent<Op>> {
+    parent: P,
+    entry_ptr: Op::TableAddr,
+}
+impl<Op: TableOps, P: UpdateParent<Op>> Clone for UpdateParentTable<Op, P> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+impl<Op: TableOps, P: UpdateParent<Op>> Copy for UpdateParentTable<Op, P> {}
+impl<Op: TableOps, P: UpdateParent<Op>> UpdateParentTable<Op, P> {
+    fn new(parent: P, entry_ptr: Op::TableAddr) -> Self {
+        UpdateParentTable { parent, entry_ptr }
+    }
+}
+impl<Op: TableOps, P: UpdateParent<Op>> UpdateParent<Op> for UpdateParentTable<Op, P> {
+    fn update_parent(self, op: &Op, new_ptr: Op::TableAddr) {
+        let pte = pte_for_table::<Op>(new_ptr);
+        unsafe {
+            write_entry_updating(op, self.parent, self.entry_ptr, pte);
+        }
+    }
+}
+
+/// A struct implementing [`UpdateParent`] that keeps track of the
+/// fact that the parent "table" is actually the CR3 system register.
+#[derive(Copy, Clone)]
+struct UpdateParentCR3 {}
+impl<Op: TableOps> UpdateParent<Op> for UpdateParentCR3 {
+    fn update_parent(self, _op: &Op, new_ptr: Op::TableAddr) {
+        unsafe {
+            core::arch::asm!("mov cr3, {}", in(reg) Op::to_phys(new_ptr));
+        }
+    }
+}
+
+/// A struct implementing [`UpdateParent`] that panics when used, for
+/// the occasional situations in which we should never do an update
+#[derive(Copy, Clone)]
+struct UpdateParentNone {}
+impl<Op: TableOps> UpdateParent<Op> for UpdateParentNone {
+    // This is only used in contexts which should absolutely never try
+    // to update a page table, and so in no case should possibly ever
+    // call an update_parent.  Therefore, this is impossible unless an
+    // extremely significant invariant violation has occurred.
+    #[allow(clippy::panic)]
+    fn update_parent(self, _op: &Op, _new_ptr: Op::TableAddr) {
+        panic!("UpdateParentNone: tried to update parent");
+    }
+}
+
 /// A helper structure indicating a mapping operation that needs to be
 /// performed
-struct MapRequest<T> {
-    table_base: T,
+struct MapRequest<Op: TableOps, P: UpdateParent<Op>> {
+    table_base: Op::TableAddr,
     vmin: VirtAddr,
     len: u64,
+    update_parent: P,
 }
 
 /// A helper structure indicating that a particular PTE needs to be
 /// modified
-struct MapResponse<T> {
-    entry_ptr: T,
+struct MapResponse<Op: TableOps, P: UpdateParent<Op>> {
+    entry_ptr: Op::TableAddr,
     vmin: VirtAddr,
     len: u64,
+    update_parent: P,
 }
 
 /// Iterator that walks through page table entries at a specific level.
@@ -122,14 +216,14 @@ struct MapResponse<T> {
 /// - PDPT: HIGH_BIT=38, LOW_BIT=30 (9 bits = 512 entries, each covering 1GB)
 /// - PD:   HIGH_BIT=29, LOW_BIT=21 (9 bits = 512 entries, each covering 2MB)
 /// - PT:   HIGH_BIT=20, LOW_BIT=12 (9 bits = 512 entries, each covering 4KB)
-struct ModifyPteIterator<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps> {
-    request: MapRequest<Op::TableAddr>,
+struct ModifyPteIterator<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps, P: UpdateParent<Op>> {
+    request: MapRequest<Op, P>,
     n: u64,
 }
-impl<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps> Iterator
-    for ModifyPteIterator<HIGH_BIT, LOW_BIT, Op>
+impl<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps, P: UpdateParent<Op>> Iterator
+    for ModifyPteIterator<HIGH_BIT, LOW_BIT, Op, P>
 {
-    type Item = MapResponse<Op::TableAddr>;
+    type Item = MapResponse<Op, P>;
     fn next(&mut self) -> Option<Self::Item> {
         // Each page table entry at this level covers a region of size (1 << LOW_BIT) bytes.
         // For example, at the PT level (LOW_BIT=12), each entry covers 4KB (0x1000 bytes).
@@ -188,12 +282,13 @@ impl<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps> Iterator
             entry_ptr,
             vmin: next_vmin,
             len: next_len,
+            update_parent: self.request.update_parent,
         })
     }
 }
-fn modify_ptes<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps>(
-    r: MapRequest<Op::TableAddr>,
-) -> ModifyPteIterator<HIGH_BIT, LOW_BIT, Op> {
+fn modify_ptes<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps, P: UpdateParent<Op>>(
+    r: MapRequest<Op, P>,
+) -> ModifyPteIterator<HIGH_BIT, LOW_BIT, Op, P> {
     ModifyPteIterator { request: r, n: 0 }
 }
 
@@ -201,34 +296,31 @@ fn modify_ptes<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps>(
 /// # Safety
 /// This function modifies page table data structures, and should not be called concurrently
 /// with any other operations that modify the page tables.
-unsafe fn alloc_pte_if_needed<Op: TableOps>(
+unsafe fn alloc_pte_if_needed<Op: TableOps, P: UpdateParent<Op>>(
     op: &Op,
-    x: MapResponse<Op::TableAddr>,
-) -> MapRequest<Op::TableAddr> {
+    x: MapResponse<Op, P>,
+) -> MapRequest<Op, UpdateParentTable<Op, P>> {
+    let new_update_parent = UpdateParentTable::new(x.update_parent, x.entry_ptr);
     if let Some(pte) = unsafe { read_pte_if_present(op, x.entry_ptr) } {
         return MapRequest {
             table_base: Op::from_phys(pte & PTE_ADDR_MASK),
             vmin: x.vmin,
             len: x.len,
+            update_parent: new_update_parent,
         };
     }
 
     let page_addr = unsafe { op.alloc_table() };
 
-    #[allow(clippy::identity_op)]
-    #[allow(clippy::precedence)]
-    let pte = Op::to_phys(page_addr) |
-        PAGE_ACCESSED_CLEAR | // accessed bit cleared (will be set by CPU when page is accessed - but we dont use the access bit for anything at present)
-        PAGE_CACHE_ENABLED | // leave caching enabled
-        PAGE_WRITE_BACK | // use write-back caching
-        PAGE_USER_ACCESS_DISABLED |// dont allow user access (no code runs in user mode for now)
-        PAGE_RW | // R/W - we don't use block-level permissions
-        PAGE_PRESENT; // P   - this entry is present
-    unsafe { op.write_entry(x.entry_ptr, pte) };
+    let pte = pte_for_table::<Op>(page_addr);
+    unsafe {
+        write_entry_updating(op, x.update_parent, x.entry_ptr, pte);
+    };
     MapRequest {
         table_base: page_addr,
         vmin: x.vmin,
         len: x.len,
+        update_parent: new_update_parent,
     }
 }
 
@@ -238,7 +330,11 @@ unsafe fn alloc_pte_if_needed<Op: TableOps>(
 /// with any other operations that modify the page tables.
 #[allow(clippy::identity_op)]
 #[allow(clippy::precedence)]
-unsafe fn map_page<Op: TableOps>(op: &Op, mapping: &Mapping, r: MapResponse<Op::TableAddr>) {
+unsafe fn map_page<Op: TableOps, P: UpdateParent<Op>>(
+    op: &Op,
+    mapping: &Mapping,
+    r: MapResponse<Op, P>,
+) {
     let pte = match &mapping.kind {
         MappingKind::BasicMapping(bm) =>
         // TODO: Support not readable
@@ -262,7 +358,7 @@ unsafe fn map_page<Op: TableOps>(op: &Op, mapping: &Mapping, r: MapResponse<Op::
         }
     };
     unsafe {
-        op.write_entry(r.entry_ptr, pte);
+        write_entry_updating(op, r.update_parent, r.entry_ptr, pte);
     }
 }
 
@@ -283,17 +379,18 @@ unsafe fn map_page<Op: TableOps>(op: &Op, mapping: &Mapping, r: MapResponse<Op::
 /// 4. PT (20:12) - write final PTE with physical address and flags
 #[allow(clippy::missing_safety_doc)]
 pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
-    modify_ptes::<47, 39, Op>(MapRequest {
+    modify_ptes::<47, 39, Op, _>(MapRequest {
         table_base: op.root_table(),
         vmin: mapping.virt_base,
         len: mapping.len,
+        update_parent: UpdateParentCR3 {},
     })
     .map(|r| unsafe { alloc_pte_if_needed(op, r) })
-    .flat_map(modify_ptes::<38, 30, Op>)
+    .flat_map(modify_ptes::<38, 30, Op, _>)
     .map(|r| unsafe { alloc_pte_if_needed(op, r) })
-    .flat_map(modify_ptes::<29, 21, Op>)
+    .flat_map(modify_ptes::<29, 21, Op, _>)
     .map(|r| unsafe { alloc_pte_if_needed(op, r) })
-    .flat_map(modify_ptes::<20, 12, Op>)
+    .flat_map(modify_ptes::<20, 12, Op, _>)
     .map(|r| unsafe { map_page(op, &mapping, r) })
     .for_each(drop);
 }
@@ -302,14 +399,15 @@ pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
 /// This function traverses page table data structures, and should not
 /// be called concurrently with any other operations that modify the
 /// page table.
-unsafe fn require_pte_exist<Op: TableOps>(
+unsafe fn require_pte_exist<Op: TableOps, P: UpdateParent<Op>>(
     op: &Op,
-    x: MapResponse<Op::TableAddr>,
-) -> Option<MapRequest<Op::TableAddr>> {
+    x: MapResponse<Op, P>,
+) -> Option<MapRequest<Op, UpdateParentTable<Op, P>>> {
     unsafe { read_pte_if_present(op, x.entry_ptr) }.map(|pte| MapRequest {
         table_base: Op::from_phys(pte & PTE_ADDR_MASK),
         vmin: x.vmin,
         len: x.len,
+        update_parent: UpdateParentTable::new(x.update_parent, x.entry_ptr),
     })
 }
 
@@ -322,20 +420,34 @@ unsafe fn require_pte_exist<Op: TableOps>(
 /// Returns `Some(pte)` containing the leaf page table entry if the address is mapped,
 /// or `None` if any level of the page table hierarchy has a non-present entry.
 #[allow(clippy::missing_safety_doc)]
-pub unsafe fn virt_to_phys<Op: TableOps>(op: &Op, address: u64) -> Option<u64> {
-    modify_ptes::<47, 39, Op>(MapRequest {
+pub unsafe fn virt_to_phys<Op: TableOps>(
+    op: &Op,
+    address: u64,
+) -> Option<(PhysAddr, BasicMapping)> {
+    modify_ptes::<47, 39, Op, _>(MapRequest {
         table_base: op.root_table(),
         vmin: address,
         len: 1,
+        update_parent: UpdateParentNone {},
     })
-    .filter_map(|r| unsafe { require_pte_exist::<Op>(op, r) })
-    .flat_map(modify_ptes::<38, 30, Op>)
-    .filter_map(|r| unsafe { require_pte_exist::<Op>(op, r) })
-    .flat_map(modify_ptes::<29, 21, Op>)
-    .filter_map(|r| unsafe { require_pte_exist::<Op>(op, r) })
-    .flat_map(modify_ptes::<20, 12, Op>)
+    .filter_map(|r| unsafe { require_pte_exist(op, r) })
+    .flat_map(modify_ptes::<38, 30, Op, _>)
+    .filter_map(|r| unsafe { require_pte_exist(op, r) })
+    .flat_map(modify_ptes::<29, 21, Op, _>)
+    .filter_map(|r| unsafe { require_pte_exist(op, r) })
+    .flat_map(modify_ptes::<20, 12, Op, _>)
     .filter_map(|r| unsafe { read_pte_if_present(op, r.entry_ptr) })
     .next()
+    .map(|r| {
+        (
+            r & PTE_ADDR_MASK,
+            BasicMapping {
+                readable: true,
+                writable: (r & PAGE_RW) != 0,
+                executable: (r & PAGE_NX) == 0,
+            },
+        )
+    })
 }
 
 pub const PAGE_SIZE: usize = 4096;
@@ -396,8 +508,9 @@ mod tests {
             self.tables.borrow()[addr.0][addr.1]
         }
 
-        unsafe fn write_entry(&self, addr: Self::TableAddr, entry: u64) {
+        unsafe fn write_entry(&self, addr: Self::TableAddr, entry: u64) -> Option<Self::TableAddr> {
             self.tables.borrow_mut()[addr.0][addr.1] = entry;
+            None
         }
 
         fn to_phys(addr: Self::TableAddr) -> PhysAddr {
@@ -629,7 +742,7 @@ mod tests {
         let result = unsafe { virt_to_phys(&ops, 0x1000) };
         assert!(result.is_some(), "Should find mapped address");
         let pte = result.unwrap();
-        assert_eq!(pte & PTE_ADDR_MASK, 0x1000);
+        assert_eq!(pte.0, 0x1000);
     }
 
     #[test]
@@ -674,9 +787,10 @@ mod tests {
             table_base: ops.root_table(),
             vmin: 0x1000,
             len: PAGE_SIZE as u64,
+            update_parent: UpdateParentNone {},
         };
 
-        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps>(request).collect();
+        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps, _>(request).collect();
         assert_eq!(responses.len(), 1, "Single page should yield one response");
         assert_eq!(responses[0].vmin, 0x1000);
         assert_eq!(responses[0].len, PAGE_SIZE as u64);
@@ -689,9 +803,10 @@ mod tests {
             table_base: ops.root_table(),
             vmin: 0x1000,
             len: 3 * PAGE_SIZE as u64,
+            update_parent: UpdateParentNone {},
         };
 
-        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps>(request).collect();
+        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps, _>(request).collect();
         assert_eq!(responses.len(), 3, "3 pages should yield 3 responses");
     }
 
@@ -702,9 +817,10 @@ mod tests {
             table_base: ops.root_table(),
             vmin: 0x1000,
             len: 0,
+            update_parent: UpdateParentNone {},
         };
 
-        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps>(request).collect();
+        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps, _>(request).collect();
         assert_eq!(responses.len(), 0, "Zero length should yield no responses");
     }
 
@@ -717,9 +833,10 @@ mod tests {
             table_base: ops.root_table(),
             vmin: 0x1800,
             len: 0x1000,
+            update_parent: UpdateParentNone {},
         };
 
-        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps>(request).collect();
+        let responses: Vec<_> = modify_ptes::<20, 12, MockTableOps, _>(request).collect();
         assert_eq!(
             responses.len(),
             2,
diff --git a/src/hyperlight_common/src/vmem.rs b/src/hyperlight_common/src/vmem.rs
index abebaf61c..b85149101 100644
--- a/src/hyperlight_common/src/vmem.rs
+++ b/src/hyperlight_common/src/vmem.rs
@@ -75,7 +75,10 @@ pub trait TableOps {
     unsafe fn read_entry(&self, addr: Self::TableAddr) -> PageTableEntry;
 
     /// Write a u64 to the given address, used to write updated page
-    /// table entries
+    /// table entries. In some cases,the page table in which the entry
+    /// is located may need to be relocated in order for this to
+    /// succeed; if this is the case, the base address of the new
+    /// table is returned.
     ///
     /// # Safety
     /// This writes to the given memory address, and so all the usual
@@ -85,7 +88,11 @@ pub trait TableOps {
     /// invariants. The implementor of the trait should ensure that
     /// nothing else will be reading/writing the address at the same
     /// time as mapping code using the trait.
-    unsafe fn write_entry(&self, addr: Self::TableAddr, entry: PageTableEntry);
+    unsafe fn write_entry(
+        &self,
+        addr: Self::TableAddr,
+        entry: PageTableEntry,
+    ) -> Option<Self::TableAddr>;
 
     /// Convert an abstract table address to a concrete physical address (u64)
     /// which can be e.g. written into a page table entry
diff --git a/src/hyperlight_guest_bin/src/paging.rs b/src/hyperlight_guest_bin/src/paging.rs
index 43adf3e46..3d8b76d1f 100644
--- a/src/hyperlight_guest_bin/src/paging.rs
+++ b/src/hyperlight_guest_bin/src/paging.rs
@@ -16,19 +16,10 @@ limitations under the License.
 
 use core::arch::asm;
 
+use hyperlight_common::vmem;
 use hyperlight_guest::prim_alloc::alloc_phys_pages;
 use tracing::{Span, instrument};
 
-
-/// Convert a physical address in main memory to a virtual address
-/// through the pysmap
-///
-/// This is _not guaranteed_ to work with device memory
-pub fn ptov(x: u64) -> *mut u8 {
-    // Currently, all of main memory is identity mapped
-    x as *mut u8
-}
-
 // TODO: This is not at all thread-safe atm
 // TODO: A lot of code in this file uses inline assembly to load and
 //       store page table entries. It would be nice to use pointer
@@ -45,6 +36,8 @@ static SNAPSHOT_PT_GPA: spin::Once<u64> = spin::Once::new();
 struct GuestMappingOperations {
     snapshot_pt_base_gpa: u64,
     snapshot_pt_base_gva: u64,
+    scratch_base_gpa: u64,
+    scratch_base_gva: u64,
 }
 impl GuestMappingOperations {
     fn new() -> Self {
@@ -57,22 +50,28 @@ impl GuestMappingOperations {
                 snapshot_pt_base_gpa
             }),
             snapshot_pt_base_gva: hyperlight_common::layout::SNAPSHOT_PT_GVA as u64,
+            scratch_base_gpa: hyperlight_guest::layout::scratch_base_gpa(),
+            scratch_base_gva: hyperlight_guest::layout::scratch_base_gva(),
         }
     }
-    fn phys_to_virt(&self, addr: u64) -> u64 {
-        if addr >= self.snapshot_pt_base_gpa {
-            self.snapshot_pt_base_gva + (addr - self.snapshot_pt_base_gpa)
+    fn phys_to_virt(&self, addr: u64) -> *mut u8 {
+        if addr >= self.scratch_base_gpa {
+            (self.scratch_base_gva + (addr - self.scratch_base_gpa)) as *mut u8
+        } else if addr >= self.snapshot_pt_base_gpa {
+            (self.snapshot_pt_base_gva + (addr - self.snapshot_pt_base_gpa)) as *mut u8
         } else {
-            // Assume for now that any of our own PTs are identity mapped.
-            addr
+            panic!("phys_to_virt encountered snapshot non-PT page")
         }
     }
 }
-impl hyperlight_common::vmem::TableOps for GuestMappingOperations {
+impl vmem::TableOps for GuestMappingOperations {
     type TableAddr = u64;
     unsafe fn alloc_table(&self) -> u64 {
         let page_addr = unsafe { alloc_phys_pages(1) };
-        unsafe { ptov(page_addr).write_bytes(0u8, hyperlight_common::vmem::PAGE_TABLE_SIZE) };
+        unsafe {
+            self.phys_to_virt(page_addr)
+                .write_bytes(0u8, vmem::PAGE_TABLE_SIZE)
+        };
         page_addr
     }
     fn entry_addr(addr: u64, offset: u64) -> u64 {
@@ -86,11 +85,27 @@ impl hyperlight_common::vmem::TableOps for GuestMappingOperations {
         }
         ret
     }
-    unsafe fn write_entry(&self, addr: u64, entry: u64) {
+    unsafe fn write_entry(&self, addr: u64, entry: u64) -> Option<u64> {
+        let mut addr = addr;
+        let mut ret = None;
+        if addr >= self.snapshot_pt_base_gpa && addr < self.scratch_base_gpa {
+            // This needs to be CoW'd over to the scratch region
+            unsafe {
+                let new_table = alloc_phys_pages(1);
+                core::ptr::copy(
+                    self.phys_to_virt(addr & !0xfff),
+                    self.phys_to_virt(new_table),
+                    vmem::PAGE_TABLE_SIZE,
+                );
+                addr = new_table | (addr & 0xfff);
+                ret = Some(new_table);
+            }
+        }
         let addr = self.phys_to_virt(addr);
         unsafe {
             asm!("mov qword ptr [{}], {}", in(reg) addr, in(reg) entry);
         }
+        ret
     }
     fn to_phys(addr: u64) -> u64 {
         addr
@@ -118,7 +133,6 @@ impl hyperlight_common::vmem::TableOps for GuestMappingOperations {
 ///   if previously-unmapped ranges are not being mapped, TLB invalidation may need to be performed afterwards.
 #[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub unsafe fn map_region(phys_base: u64, virt_base: *mut u8, len: u64) {
-    use hyperlight_common::vmem;
     unsafe {
         vmem::map(
             &GuestMappingOperations::new(),
@@ -136,6 +150,8 @@ pub unsafe fn map_region(phys_base: u64, virt_base: *mut u8, len: u64) {
     }
 }
 
+pub fn virt_to_phys(gva: u64) -> Option<(u64, vmem::BasicMapping)> {
+    unsafe { vmem::virt_to_phys::<_>(&GuestMappingOperations::new(), gva) }
 }
 
 pub fn flush_tlb() {

From 5bbb3931ad7c70aae5013108df3b70eff2e8b46f Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Tue, 27 Jan 2026 18:17:40 +0000
Subject: [PATCH 12/20] Make the snapshot page table mapping readonly

These were already being CoW'd.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/sandbox/snapshot.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/hyperlight_host/src/sandbox/snapshot.rs b/src/hyperlight_host/src/sandbox/snapshot.rs
index 6f9bd3a04..6a083098a 100644
--- a/src/hyperlight_host/src/sandbox/snapshot.rs
+++ b/src/hyperlight_host/src/sandbox/snapshot.rs
@@ -217,7 +217,7 @@ impl Snapshot {
                     len: (pt_buf.size() - pt_size_mapped) as u64,
                     kind: MappingKind::BasicMapping(BasicMapping {
                         readable: true,
-                        writable: true,
+                        writable: false,
                         executable: false,
                     }),
                 };

From 93faff9590809a612bd0284eb66fb8ec73093906 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Wed, 28 Jan 2026 01:17:34 +0000
Subject: [PATCH 13/20] snapshot: handle the scratch region correctly

Now that some important things are being written to the scratch
region, this makes the snapshot process aware of it. In particular,
rather than just copying the current snapshot region into a new
snapshot, the snapshot process now traverses the guest page tables and
picks all pages that are mapped into guest virtual memory, placing
them into the new snapshot along with new page tables which account
for all physical memory having moved to the snapshot region.

Several tests are temporarily disabled in this commit. One,
`snapshot_restore_handles_remapping_correctly`, will be re-enabled in
the next commit; its failure is due to an issue with state in the
guest page mapping process not being reset on snapshot. Several
others, related to the stack guard page, are also disabled, because
the snapshot process loses the context, kept in the host, of which
physical page is the stack guard page.  These will be re-enabled in an
upcoming patch series, which moves stack management fully into the
guest.

This commit also disables an optimisation around restoring to a
just-taken snapshot in certain cases; see the comment at the head of
`MultiUseSandbox::restore` for details.  Note that the optimisation of
taking a new snapshot immediately after a restore remains intact.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 .../src/arch/amd64/layout.rs                  |   3 +-
 src/hyperlight_common/src/arch/amd64/vmem.rs  |  84 +--
 src/hyperlight_common/src/arch/i686/layout.rs |   3 +-
 src/hyperlight_common/src/arch/i686/vmem.rs   |   9 +-
 src/hyperlight_common/src/layout.rs           |   2 +-
 src/hyperlight_common/src/vmem.rs             |   2 +-
 src/hyperlight_guest_bin/src/paging.rs        |  15 +-
 src/hyperlight_host/src/error.rs              |   7 +-
 .../src/hypervisor/hyperlight_vm.rs           |  99 ++--
 src/hyperlight_host/src/mem/layout.rs         |  19 +-
 src/hyperlight_host/src/mem/memory_region.rs  |   3 +
 src/hyperlight_host/src/mem/mgr.rs            | 137 ++---
 src/hyperlight_host/src/mem/shared_mem.rs     |   3 +-
 .../src/sandbox/initialized_multi_use.rs      |  99 +++-
 src/hyperlight_host/src/sandbox/snapshot.rs   | 481 ++++++++++++++----
 .../src/sandbox/uninitialized_evolve.rs       |   4 +-
 src/hyperlight_host/tests/integration_test.rs |   6 +
 src/tests/rust_guests/simpleguest/src/main.rs |  15 +-
 18 files changed, 701 insertions(+), 290 deletions(-)

diff --git a/src/hyperlight_common/src/arch/amd64/layout.rs b/src/hyperlight_common/src/arch/amd64/layout.rs
index 1c020f4f2..c98d4b06d 100644
--- a/src/hyperlight_common/src/arch/amd64/layout.rs
+++ b/src/hyperlight_common/src/arch/amd64/layout.rs
@@ -18,7 +18,8 @@ limitations under the License.
 /// to make working with start/end ptrs in a few places more
 /// convenient (not needing to worry about overflow)
 pub const MAX_GVA: usize = 0xffff_ffff_ffff_efff;
-pub const SNAPSHOT_PT_GVA: usize = 0xffff_8000_0000_0000;
+pub const SNAPSHOT_PT_GVA_MIN: usize = 0xffff_8000_0000_0000;
+pub const SNAPSHOT_PT_GVA_MAX: usize = 0xffff_80ff_ffff_ffff;
 
 /// We assume 36-bit IPAs for now, since every amd64 processor
 /// supports at least 36 bits.  Almost all of them support at least 40
diff --git a/src/hyperlight_common/src/arch/amd64/vmem.rs b/src/hyperlight_common/src/arch/amd64/vmem.rs
index ed0c0611c..5f75b3488 100644
--- a/src/hyperlight_common/src/arch/amd64/vmem.rs
+++ b/src/hyperlight_common/src/arch/amd64/vmem.rs
@@ -47,14 +47,14 @@ use crate::vmem::{BasicMapping, Mapping, MappingKind, TableOps};
 //
 
 /// Page is Present
-pub const PAGE_PRESENT: u64 = 1;
+const PAGE_PRESENT: u64 = 1;
 /// Page is Read/Write (if not set page is read only so long as the WP bit in CR0 is set to 1 - which it is in Hyperlight)
-pub const PAGE_RW: u64 = 1 << 1;
+const PAGE_RW: u64 = 1 << 1;
 /// Execute Disable (if this bit is set then data in the page cannot be executed)`
-pub const PAGE_NX: u64 = 1 << 63;
+const PAGE_NX: u64 = 1 << 63;
 /// Mask to extract the physical address from a PTE (bits 51:12)
 /// This masks out the lower 12 flag bits AND the upper bits including NX (bit 63)
-pub const PTE_ADDR_MASK: u64 = 0x000F_FFFF_FFFF_F000;
+const PTE_ADDR_MASK: u64 = 0x000F_FFFF_FFFF_F000;
 const PAGE_USER_ACCESS_DISABLED: u64 = 0 << 2; // U/S bit not set - supervisor mode only (no code runs in user mode for now)
 const PAGE_DIRTY_CLEAR: u64 = 0 << 6; // D - dirty bit cleared (set by CPU when written)
 const PAGE_ACCESSED_CLEAR: u64 = 0 << 5; // A - accessed bit cleared (set by CPU when accessed)
@@ -415,41 +415,58 @@ unsafe fn require_pte_exist<Op: TableOps, P: UpdateParent<Op>>(
 // here, and the general conditions are documented in the
 // architecture-independent re-export in vmem.rs
 
-/// Translates a virtual address to its physical address by walking the page tables.
+/// Translates a virtual address range to the physical address pages
+/// that back it by walking the page tables.
 ///
-/// Returns `Some(pte)` containing the leaf page table entry if the address is mapped,
-/// or `None` if any level of the page table hierarchy has a non-present entry.
+/// Returns an iterator with an entry for each mapped page that
+/// intersects the given range.
+///
+/// This takes AsRef<Op> + Copy so that on targets where the
+/// operations have little state (e.g. the guest) the operations state
+/// can be copied into the closure(s) in the iterator, allowing for a
+/// nicer result lifetime.  On targets like the
+/// building-an-original-snapshot portion of the host, where the
+/// operations structure owns a large buffer, a reference can instead
+/// be passed.
 #[allow(clippy::missing_safety_doc)]
-pub unsafe fn virt_to_phys<Op: TableOps>(
-    op: &Op,
+pub unsafe fn virt_to_phys<'a, Op: TableOps + 'a>(
+    op: impl core::convert::AsRef<Op> + Copy + 'a,
     address: u64,
-) -> Option<(PhysAddr, BasicMapping)> {
+    len: u64,
+) -> impl Iterator<Item = (VirtAddr, PhysAddr, BasicMapping)> + 'a {
+    // Undo sign-extension, and mask off any sub-page bits
+    let vmin = (address & ((1u64 << VA_BITS) - 1)) & !(PAGE_SIZE as u64 - 1);
+    let vmax = core::cmp::min(vmin + len, 1u64 << VA_BITS);
     modify_ptes::<47, 39, Op, _>(MapRequest {
-        table_base: op.root_table(),
+        table_base: op.as_ref().root_table(),
         vmin: address,
-        len: 1,
+        len: vmax - vmin,
         update_parent: UpdateParentNone {},
     })
-    .filter_map(|r| unsafe { require_pte_exist(op, r) })
+    .filter_map(move |r| unsafe { require_pte_exist(op.as_ref(), r) })
     .flat_map(modify_ptes::<38, 30, Op, _>)
-    .filter_map(|r| unsafe { require_pte_exist(op, r) })
+    .filter_map(move |r| unsafe { require_pte_exist(op.as_ref(), r) })
     .flat_map(modify_ptes::<29, 21, Op, _>)
-    .filter_map(|r| unsafe { require_pte_exist(op, r) })
+    .filter_map(move |r| unsafe { require_pte_exist(op.as_ref(), r) })
     .flat_map(modify_ptes::<20, 12, Op, _>)
-    .filter_map(|r| unsafe { read_pte_if_present(op, r.entry_ptr) })
-    .next()
-    .map(|r| {
-        (
-            r & PTE_ADDR_MASK,
-            BasicMapping {
-                readable: true,
-                writable: (r & PAGE_RW) != 0,
-                executable: (r & PAGE_NX) == 0,
-            },
-        )
+    .filter_map(move |r| {
+        let pte = unsafe { read_pte_if_present(op.as_ref(), r.entry_ptr) }?;
+        let phys_addr = pte & PTE_ADDR_MASK;
+        // Re-do the sign extension
+        let sgn_bit = r.vmin >> (VA_BITS - 1);
+        let sgn_bits = 0u64.wrapping_sub(sgn_bit) << VA_BITS;
+        let virt_addr = sgn_bits | r.vmin;
+        let perms = BasicMapping {
+            readable: true,
+            writable: (pte & PAGE_RW) != 0,
+            executable: (pte & PAGE_NX) == 0,
+        };
+        Some((virt_addr, phys_addr, perms))
     })
 }
 
+const VA_BITS: usize = 48; // We use 48-bit virtual addresses at the moment.
+
 pub const PAGE_SIZE: usize = 4096;
 pub const PAGE_TABLE_SIZE: usize = 4096;
 pub type PageTableEntry = u64;
@@ -471,6 +488,13 @@ mod tests {
         tables: RefCell<Vec<[u64; PAGE_TABLE_ENTRIES_PER_TABLE]>>,
     }
 
+    // for virt_to_phys
+    impl core::convert::AsRef<MockTableOps> for MockTableOps {
+        fn as_ref(&self) -> &Self {
+            self
+        }
+    }
+
     impl MockTableOps {
         fn new() -> Self {
             // Start with one table (the root/PML4)
@@ -739,10 +763,10 @@ mod tests {
 
         unsafe { map(&ops, mapping) };
 
-        let result = unsafe { virt_to_phys(&ops, 0x1000) };
+        let result = unsafe { virt_to_phys(&ops, 0x1000, 1).next() };
         assert!(result.is_some(), "Should find mapped address");
         let pte = result.unwrap();
-        assert_eq!(pte.0, 0x1000);
+        assert_eq!(pte.1, 0x1000);
     }
 
     #[test]
@@ -750,7 +774,7 @@ mod tests {
         let ops = MockTableOps::new();
         // Don't map anything
 
-        let result = unsafe { virt_to_phys(&ops, 0x1000) };
+        let result = unsafe { virt_to_phys(&ops, 0x1000, 1).next() };
         assert!(result.is_none(), "Should return None for unmapped address");
     }
 
@@ -771,7 +795,7 @@ mod tests {
         unsafe { map(&ops, mapping) };
 
         // Query an address in a different PT entry (unmapped)
-        let result = unsafe { virt_to_phys(&ops, 0x5000) };
+        let result = unsafe { virt_to_phys(&ops, 0x5000, 1).next() };
         assert!(
             result.is_none(),
             "Should return None for unmapped address in same PT"
diff --git a/src/hyperlight_common/src/arch/i686/layout.rs b/src/hyperlight_common/src/arch/i686/layout.rs
index c70f38e34..e46a7c8ff 100644
--- a/src/hyperlight_common/src/arch/i686/layout.rs
+++ b/src/hyperlight_common/src/arch/i686/layout.rs
@@ -18,5 +18,6 @@ limitations under the License.
 // allow compiling the guest for real mode boot scenarios.
 
 pub const MAX_GVA: usize = 0xffff_efff;
-pub const SNAPSHOT_PT_GVA: usize = 0xefff_efff;
+pub const SNAPSHOT_PT_GVA_MIN: usize = 0xef00_0000;
+pub const SNAPSHOT_PT_GVA_MAX: usize = 0xefff_efff;
 pub const MAX_GPA: usize = 0xffff_ffff;
diff --git a/src/hyperlight_common/src/arch/i686/vmem.rs b/src/hyperlight_common/src/arch/i686/vmem.rs
index 1a1e4974c..4b1d16b87 100644
--- a/src/hyperlight_common/src/arch/i686/vmem.rs
+++ b/src/hyperlight_common/src/arch/i686/vmem.rs
@@ -31,8 +31,15 @@ pub unsafe fn map<Op: TableOps>(_op: &Op, _mapping: Mapping) {
 }
 
 #[allow(clippy::missing_safety_doc)]
-pub unsafe fn virt_to_phys<Op: TableOps>(_op: &Op, _address: u64) -> Option<(u64, BasicMapping)> {
+pub unsafe fn virt_to_phys<Op: TableOps>(
+    _op: &Op,
+    _address: u64,
+) -> impl Iterator<Item = (VirtAddr, PhysAddr, BasicMapping)> {
     panic!(
         "vmem::virt_to_phys: i686 guests do not support booting the full hyperlight guest kernel"
     );
+    // necessary to provide a concrete type that impls Iterator as the
+    // return type, even though this will never be executed
+    #[allow(unreachable_code)]
+    core::iter::empty()
 }
diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs
index 645f6217e..1c9d0ed71 100644
--- a/src/hyperlight_common/src/layout.rs
+++ b/src/hyperlight_common/src/layout.rs
@@ -18,7 +18,7 @@ limitations under the License.
 #[cfg_attr(target_arch = "x86", path = "arch/i686/layout.rs")]
 mod arch;
 
-pub use arch::{MAX_GPA, MAX_GVA, SNAPSHOT_PT_GVA};
+pub use arch::{MAX_GPA, MAX_GVA, SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN};
 
 // offsets down from the top of scratch memory for various things
 pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08;
diff --git a/src/hyperlight_common/src/vmem.rs b/src/hyperlight_common/src/vmem.rs
index b85149101..1a83146f0 100644
--- a/src/hyperlight_common/src/vmem.rs
+++ b/src/hyperlight_common/src/vmem.rs
@@ -16,7 +16,7 @@ limitations under the License.
 
 #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/vmem.rs")]
 #[cfg_attr(target_arch = "x86", path = "arch/i686/vmem.rs")]
-pub mod arch;
+mod arch;
 
 /// This is always the page size that the /guest/ is being compiled
 /// for, which may or may not be the same as the host page size.
diff --git a/src/hyperlight_guest_bin/src/paging.rs b/src/hyperlight_guest_bin/src/paging.rs
index 3d8b76d1f..b806bf884 100644
--- a/src/hyperlight_guest_bin/src/paging.rs
+++ b/src/hyperlight_guest_bin/src/paging.rs
@@ -33,6 +33,7 @@ use tracing::{Span, instrument};
 // change the snapshot pt base, we will need to modify this.
 static SNAPSHOT_PT_GPA: spin::Once<u64> = spin::Once::new();
 
+#[derive(Copy, Clone)]
 struct GuestMappingOperations {
     snapshot_pt_base_gpa: u64,
     snapshot_pt_base_gva: u64,
@@ -49,7 +50,7 @@ impl GuestMappingOperations {
                 };
                 snapshot_pt_base_gpa
             }),
-            snapshot_pt_base_gva: hyperlight_common::layout::SNAPSHOT_PT_GVA as u64,
+            snapshot_pt_base_gva: hyperlight_common::layout::SNAPSHOT_PT_GVA_MIN as u64,
             scratch_base_gpa: hyperlight_guest::layout::scratch_base_gpa(),
             scratch_base_gva: hyperlight_guest::layout::scratch_base_gva(),
         }
@@ -64,6 +65,12 @@ impl GuestMappingOperations {
         }
     }
 }
+// for virt_to_phys
+impl core::convert::AsRef<GuestMappingOperations> for GuestMappingOperations {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
 impl vmem::TableOps for GuestMappingOperations {
     type TableAddr = u64;
     unsafe fn alloc_table(&self) -> u64 {
@@ -150,8 +157,10 @@ pub unsafe fn map_region(phys_base: u64, virt_base: *mut u8, len: u64) {
     }
 }
 
-pub fn virt_to_phys(gva: u64) -> Option<(u64, vmem::BasicMapping)> {
-    unsafe { vmem::virt_to_phys::<_>(&GuestMappingOperations::new(), gva) }
+pub fn virt_to_phys(
+    gva: u64,
+) -> impl Iterator<Item = (vmem::VirtAddr, vmem::PhysAddr, vmem::BasicMapping)> {
+    unsafe { vmem::virt_to_phys::<_>(GuestMappingOperations::new(), gva, 1) }
 }
 
 pub fn flush_tlb() {
diff --git a/src/hyperlight_host/src/error.rs b/src/hyperlight_host/src/error.rs
index bf9959d72..837eb3b5a 100644
--- a/src/hyperlight_host/src/error.rs
+++ b/src/hyperlight_host/src/error.rs
@@ -328,6 +328,12 @@ impl HyperlightError {
             | HyperlightError::SnapshotSizeMismatch(_, _)
             | HyperlightError::MemoryRegionSizeMismatch(_, _, _) => true,
 
+            // These errors poison the sandbox because they can leave
+            // it in an inconsistent state due to snapshot restore
+            // failing partway through
+            HyperlightError::HyperlightVmError(HyperlightVmError::UpdateRegion(_))
+            | HyperlightError::HyperlightVmError(HyperlightVmError::AccessPageTable(_)) => true,
+
             // HyperlightVmError::DispatchGuestCall may poison the sandbox
             HyperlightError::HyperlightVmError(HyperlightVmError::DispatchGuestCall(e)) => {
                 e.is_poison_error()
@@ -351,7 +357,6 @@ impl HyperlightError {
             | HyperlightError::HyperlightVmError(HyperlightVmError::Initialize(_))
             | HyperlightError::HyperlightVmError(HyperlightVmError::MapRegion(_))
             | HyperlightError::HyperlightVmError(HyperlightVmError::UnmapRegion(_))
-            | HyperlightError::HyperlightVmError(HyperlightVmError::UpdateScratch(_))
             | HyperlightError::IOError(_)
             | HyperlightError::IntConversionFailure(_)
             | HyperlightError::InvalidFlatBuffer(_)
diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs
index 44e1c1c75..65a4ee6f0 100644
--- a/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs
+++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm.rs
@@ -89,15 +89,20 @@ pub(crate) struct HyperlightVm {
     orig_rsp: GuestPtr,
     interrupt_handle: Arc<dyn InterruptHandleImpl>,
 
-    sandbox_regions: Vec<MemoryRegion>, // Initially mapped regions when sandbox is created
-    mmap_regions: Vec<(u32, MemoryRegion)>, // Later mapped regions (slot number, region)
-    next_slot: u32,                     // Monotonically increasing slot number
-    freed_slots: Vec<u32>,              // Reusable slots from unmapped regions
-    scratch_slot: u32,                  // The slot number used for the scratch region
+    next_slot: u32,        // Monotonically increasing slot number
+    freed_slots: Vec<u32>, // Reusable slots from unmapped regions
+
+    snapshot_slot: u32,
+    // The current snapshot region, used to keep it alive as long as
+    // it is used & when unmapping
+    snapshot_memory: Option<GuestSharedMemory>,
+    scratch_slot: u32, // The slot number used for the scratch region
     // The current scratch region, used to keep it alive as long as it
     // is used & when unmapping
     scratch_memory: Option<GuestSharedMemory>,
 
+    mmap_regions: Vec<(u32, MemoryRegion)>, // Later mapped regions (slot number, region)
+
     #[cfg(gdb)]
     gdb_conn: Option<DebugCommChannel<DebugResponse, DebugMsg>>,
     #[cfg(gdb)]
@@ -251,13 +256,20 @@ pub enum UnmapRegionError {
 
 /// Errors that can occur when updating the scratch mapping
 #[derive(Debug, thiserror::Error)]
-pub enum UpdateScratchError {
+pub enum UpdateRegionError {
     #[error("VM map memory error: {0}")]
     MapMemory(#[from] MapMemoryError),
     #[error("VM unmap memory error: {0}")]
     UnmapMemory(#[from] UnmapMemoryError),
 }
 
+/// Errors that can occur when accessing the root page table state
+#[derive(Debug, thiserror::Error)]
+pub enum AccessPageTableError {
+    #[error("Failed to get/set registers: {0}")]
+    AccessRegs(#[from] RegisterError),
+}
+
 /// Errors that can occur during HyperlightVm creation
 #[derive(Debug, thiserror::Error)]
 pub enum CreateHyperlightVmError {
@@ -274,7 +286,7 @@ pub enum CreateHyperlightVmError {
     #[error("VM operation error: {0}")]
     Vm(#[from] VmError),
     #[error("Set scratch error: {0}")]
-    UpdateScratch(#[from] UpdateScratchError),
+    UpdateRegion(#[from] UpdateRegionError),
 }
 
 /// Errors that can occur during debug exit handling
@@ -324,8 +336,10 @@ pub enum HyperlightVmError {
     MapRegion(#[from] MapRegionError),
     #[error("Unmap region error: {0}")]
     UnmapRegion(#[from] UnmapRegionError),
-    #[error("Update scratch error: {0}")]
-    UpdateScratch(#[from] UpdateScratchError),
+    #[error("Update region error: {0}")]
+    UpdateRegion(#[from] UpdateRegionError),
+    #[error("Access page table error: {0}")]
+    AccessPageTable(#[from] AccessPageTableError),
 }
 
 impl HyperlightVm {
@@ -333,7 +347,7 @@ impl HyperlightVm {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
-        mem_regions: Vec<MemoryRegion>,
+        snapshot_mem: GuestSharedMemory,
         scratch_mem: GuestSharedMemory,
         _pml4_addr: u64,
         entrypoint: u64,
@@ -348,8 +362,7 @@ impl HyperlightVm {
         #[cfg(not(gdb))]
         type VmType = Box<dyn VirtualMachine>;
 
-        #[cfg_attr(not(gdb), allow(unused_mut))]
-        let mut vm: VmType = match get_available_hypervisor() {
+        let vm: VmType = match get_available_hypervisor() {
             #[cfg(kvm)]
             Some(HypervisorType::Kvm) => Box::new(KvmVm::new().map_err(VmError::CreateVm)?),
             #[cfg(mshv3)]
@@ -359,14 +372,6 @@ impl HyperlightVm {
             None => return Err(CreateHyperlightVmError::NoHypervisorFound),
         };
 
-        for (i, region) in mem_regions.iter().enumerate() {
-            // Safety: slots are unique and region points to valid memory since we created the regions
-            unsafe {
-                vm.map_memory((i as u32, region))
-                    .map_err(VmError::MapMemory)?
-            };
-        }
-
         #[cfg(feature = "init-paging")]
         vm.set_sregs(&CommonSpecialRegisters::standard_64bit_defaults(_pml4_addr))
             .map_err(VmError::Register)?;
@@ -407,7 +412,8 @@ impl HyperlightVm {
             }),
         });
 
-        let scratch_slot = mem_regions.len() as u32;
+        let snapshot_slot = 0u32;
+        let scratch_slot = 1u32;
         #[cfg_attr(not(gdb), allow(unused_mut))]
         let mut ret = Self {
             vm,
@@ -417,12 +423,15 @@ impl HyperlightVm {
             page_size: 0, // Will be set in `initialise`
 
             next_slot: scratch_slot + 1,
-            sandbox_regions: mem_regions,
-            mmap_regions: Vec::new(),
             freed_slots: Vec::new(),
+
+            snapshot_slot,
+            snapshot_memory: None,
             scratch_slot,
             scratch_memory: None,
 
+            mmap_regions: Vec::new(),
+
             #[cfg(gdb)]
             gdb_conn,
             #[cfg(gdb)]
@@ -433,6 +442,7 @@ impl HyperlightVm {
             rt_cfg,
         };
 
+        ret.update_snapshot_mapping(snapshot_mem)?;
         ret.update_scratch_mapping(scratch_mem)?;
 
         // Send the interrupt handle to the GDB thread if debugging is enabled
@@ -559,11 +569,28 @@ impl HyperlightVm {
         self.mmap_regions.iter().map(|(_, region)| region)
     }
 
+    /// Update the snapshot mapping to point to a new GuestSharedMemory
+    pub(crate) fn update_snapshot_mapping(
+        &mut self,
+        snapshot: GuestSharedMemory,
+    ) -> Result<(), UpdateRegionError> {
+        let guest_base = crate::mem::layout::SandboxMemoryLayout::BASE_ADDRESS as u64;
+        let rgn = snapshot.mapping_at(guest_base, MemoryRegionType::Snapshot);
+
+        if let Some(old_snapshot) = self.snapshot_memory.replace(snapshot) {
+            let old_rgn = old_snapshot.mapping_at(guest_base, MemoryRegionType::Snapshot);
+            self.vm.unmap_memory((self.snapshot_slot, &old_rgn))?;
+        }
+        unsafe { self.vm.map_memory((self.snapshot_slot, &rgn))? };
+
+        Ok(())
+    }
+
     /// Update the scratch mapping to point to a new GuestSharedMemory
     pub(crate) fn update_scratch_mapping(
         &mut self,
         scratch: GuestSharedMemory,
-    ) -> Result<(), UpdateScratchError> {
+    ) -> Result<(), UpdateRegionError> {
         let guest_base = hyperlight_common::layout::scratch_base_gpa(scratch.mem_size());
         let rgn = scratch.mapping_at(guest_base, MemoryRegionType::Scratch);
 
@@ -577,6 +604,20 @@ impl HyperlightVm {
         Ok(())
     }
 
+    /// Get the current base page table physical address
+    pub(crate) fn get_root_pt(&mut self) -> Result<u64, AccessPageTableError> {
+        let sregs = self.vm.sregs()?;
+        Ok(sregs.cr3)
+    }
+
+    /// Set the current base page table physical address
+    pub(crate) fn set_root_pt(&mut self, addr: u64) -> Result<(), AccessPageTableError> {
+        let mut sregs = self.vm.sregs()?;
+        sregs.cr3 = addr;
+        self.vm.set_sregs(&sregs)?;
+        Ok(())
+    }
+
     /// Dispatch a call from the host to the guest using the given pointer
     /// to the dispatch function _in the guest's address space_.
     ///
@@ -713,7 +754,7 @@ impl HyperlightVm {
                     self.handle_io(mem_mgr, host_funcs, port, data)?;
                 }
                 Ok(VmExit::MmioRead(addr)) => {
-                    let all_regions = self.sandbox_regions.iter().chain(self.get_mapped_regions());
+                    let all_regions = self.get_mapped_regions();
                     match get_memory_access_violation(
                         addr as usize,
                         MemoryRegionFlags::WRITE,
@@ -742,7 +783,7 @@ impl HyperlightVm {
                     }
                 }
                 Ok(VmExit::MmioWrite(addr)) => {
-                    let all_regions = self.sandbox_regions.iter().chain(self.get_mapped_regions());
+                    let all_regions = self.get_mapped_regions();
                     match get_memory_access_violation(
                         addr as usize,
                         MemoryRegionFlags::WRITE,
@@ -1051,9 +1092,9 @@ impl HyperlightVm {
                     .and_then(|name| name.to_os_string().into_string().ok())
             });
 
-            // Include both initial sandbox regions and dynamically mapped regions
-            let mut regions: Vec<MemoryRegion> = self.sandbox_regions.clone();
-            regions.extend(self.get_mapped_regions().cloned());
+            // Include dynamically mapped regions
+            // TODO: include the snapshot and scratch regions
+            let regions: Vec<MemoryRegion> = self.get_mapped_regions().cloned().collect();
             Ok(Some(crashdump::CrashDumpContext::new(
                 regions,
                 regs,
diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs
index e353f25b1..dfa161da4 100644
--- a/src/hyperlight_host/src/mem/layout.rs
+++ b/src/hyperlight_host/src/mem/layout.rs
@@ -81,10 +81,10 @@ use super::memory_region::MemoryRegionType::{
     Code, GuardPage, Heap, InitData, InputData, OutputData, Peb, Stack,
 };
 use super::memory_region::{
-    DEFAULT_GUEST_BLOB_MEM_FLAGS, MemoryRegion, MemoryRegion_, MemoryRegionFlags, MemoryRegionKind,
+    DEFAULT_GUEST_BLOB_MEM_FLAGS, MemoryRegion_, MemoryRegionFlags, MemoryRegionKind,
     MemoryRegionVecBuilder,
 };
-use super::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, SharedMemory};
+use super::shared_mem::{ExclusiveSharedMemory, SharedMemory};
 use crate::error::HyperlightError::{GuestOffsetIsInvalid, MemoryRequestTooBig};
 use crate::sandbox::SandboxConfiguration;
 use crate::{Result, new_error};
@@ -125,6 +125,7 @@ pub(crate) struct SandboxMemoryLayout {
     code_size: usize,
     // The offset in the sandbox memory where the code starts
     guest_code_offset: usize,
+    #[cfg_attr(not(feature = "init-paging"), allow(unused))]
     pub(crate) init_data_permissions: Option<MemoryRegionFlags>,
 
     // The size of the scratch region in physical memory; note that
@@ -330,12 +331,13 @@ impl SandboxMemoryLayout {
     }
 
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
+    #[cfg_attr(not(feature = "init-paging"), allow(unused))]
     pub(super) fn get_guest_stack_size(&self) -> usize {
         self.stack_size
     }
 
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    pub(super) fn get_scratch_size(&self) -> usize {
+    pub(crate) fn get_scratch_size(&self) -> usize {
         self.scratch_size
     }
 
@@ -462,7 +464,6 @@ impl SandboxMemoryLayout {
 
     /// Sets the size of the memory region used for page tables
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    #[cfg(feature = "init-paging")]
     pub(crate) fn set_pt_size(&mut self, size: usize) {
         self.pt_size = Some(size);
     }
@@ -475,17 +476,9 @@ impl SandboxMemoryLayout {
         self.get_top_of_user_stack_offset() + self.stack_size - 0x28
     }
 
-    pub fn get_memory_regions(&self, shared_mem: &GuestSharedMemory) -> Result<Vec<MemoryRegion>> {
-        self.get_memory_regions_(
-            #[cfg(target_os = "windows")]
-            shared_mem.host_region_base(),
-            #[cfg(not(target_os = "windows"))]
-            shared_mem.base_addr(),
-        )
-    }
-
     /// Returns the memory regions associated with this memory layout,
     /// suitable for passing to a hypervisor for mapping into memory
+    #[cfg_attr(not(feature = "init-paging"), allow(unused))]
     pub(crate) fn get_memory_regions_<K: MemoryRegionKind>(
         &self,
         host_base: K::HostBaseType,
diff --git a/src/hyperlight_host/src/mem/memory_region.rs b/src/hyperlight_host/src/mem/memory_region.rs
index 18b044118..36a5b3eb3 100644
--- a/src/hyperlight_host/src/mem/memory_region.rs
+++ b/src/hyperlight_host/src/mem/memory_region.rs
@@ -140,6 +140,8 @@ pub enum MemoryRegionType {
     Stack,
     /// The scratch region
     Scratch,
+    /// The snapshot region
+    Snapshot,
 }
 
 /// A trait that distinguishes between different kinds of memory region representations.
@@ -260,6 +262,7 @@ pub(crate) struct MemoryRegion_<K: MemoryRegionKind> {
 
 pub(crate) type MemoryRegion = MemoryRegion_<HostGuestMemoryRegion>;
 
+#[cfg_attr(not(feature = "init-paging"), allow(unused))]
 pub(crate) struct MemoryRegionVecBuilder<K: MemoryRegionKind> {
     guest_base_phys_addr: usize,
     host_base_virt_addr: K::HostBaseType,
diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs
index d67430edb..20e92c558 100644
--- a/src/hyperlight_host/src/mem/mgr.rs
+++ b/src/hyperlight_host/src/mem/mgr.rs
@@ -22,10 +22,7 @@ use hyperlight_common::flatbuffer_wrappers::function_call::{
 };
 use hyperlight_common::flatbuffer_wrappers::function_types::FunctionCallResult;
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
-#[cfg(feature = "init-paging")]
-use hyperlight_common::vmem;
-#[cfg(feature = "init-paging")]
-use hyperlight_common::vmem::{PAGE_TABLE_SIZE, PageTableEntry, PhysAddr};
+use hyperlight_common::vmem::{self, PAGE_TABLE_SIZE, PageTableEntry, PhysAddr};
 use tracing::{Span, instrument};
 
 use super::layout::SandboxMemoryLayout;
@@ -61,13 +58,11 @@ pub(crate) struct SandboxMemoryManager<S> {
     pub(crate) abort_buffer: Vec<u8>,
 }
 
-#[cfg(feature = "init-paging")]
 pub(crate) struct GuestPageTableBuffer {
     buffer: std::cell::RefCell<Vec<u8>>,
     phys_base: usize,
 }
 
-#[cfg(feature = "init-paging")]
 impl vmem::TableOps for GuestPageTableBuffer {
     type TableAddr = (usize, usize); // (table_index, entry_index)
 
@@ -125,7 +120,6 @@ impl vmem::TableOps for GuestPageTableBuffer {
     }
 }
 
-#[cfg(feature = "init-paging")]
 impl GuestPageTableBuffer {
     pub(crate) fn new(phys_base: usize) -> Self {
         GuestPageTableBuffer {
@@ -134,6 +128,10 @@ impl GuestPageTableBuffer {
         }
     }
 
+    pub(crate) fn phys_base(&self) -> usize {
+        self.phys_base
+    }
+
     pub(crate) fn size(&self) -> usize {
         self.buffer.borrow().len()
     }
@@ -191,13 +189,16 @@ where
         &mut self,
         sandbox_id: u64,
         mapped_regions: Vec<MemoryRegion>,
+        root_pt_gpa: u64,
     ) -> Result<Snapshot> {
         Snapshot::new(
             &mut self.shared_mem,
+            &mut self.scratch_mem,
             sandbox_id,
             self.layout,
             crate::mem::exe::LoadInfo::dummy(),
             mapped_regions,
+            root_pt_gpa,
         )
     }
 }
@@ -395,14 +396,15 @@ impl SandboxMemoryManager<HostSharedMemory> {
     pub(crate) fn restore_snapshot(
         &mut self,
         snapshot: &Snapshot,
-    ) -> Result<Option<GuestSharedMemory>> {
-        if self.shared_mem.mem_size() != snapshot.mem_size() {
-            return Err(new_error!(
-                "Snapshot size does not match current memory size: {} != {}",
-                self.shared_mem.raw_mem_size(),
-                snapshot.mem_size()
-            ));
-        }
+    ) -> Result<(Option<GuestSharedMemory>, Option<GuestSharedMemory>)> {
+        let gsnapshot = if self.shared_mem.mem_size() == snapshot.mem_size() {
+            None
+        } else {
+            let new_snapshot_mem = ExclusiveSharedMemory::new(snapshot.mem_size())?;
+            let (hsnapshot, gsnapshot) = new_snapshot_mem.build();
+            self.shared_mem = hsnapshot;
+            Some(gsnapshot)
+        };
         self.shared_mem.restore_from_snapshot(snapshot)?;
         let new_scratch_size = snapshot.layout().get_scratch_size();
         let gscratch = if new_scratch_size == self.scratch_mem.mem_size() {
@@ -421,7 +423,7 @@ impl SandboxMemoryManager<HostSharedMemory> {
             Some(gscratch)
         };
         self.update_scratch_bookkeeping();
-        Ok(gscratch)
+        Ok((gsnapshot, gscratch))
     }
 
     fn update_scratch_bookkeeping(&mut self) {
@@ -453,7 +455,6 @@ impl SandboxMemoryManager<HostSharedMemory> {
 #[cfg(all(feature = "init-paging", target_arch = "x86_64"))]
 mod tests {
     use hyperlight_common::vmem::PAGE_TABLE_SIZE;
-    use hyperlight_common::vmem::arch::{PAGE_NX, PAGE_PRESENT, PAGE_RW, PTE_ADDR_MASK};
     use hyperlight_testing::sandbox_sizes::{LARGE_HEAP_SIZE, MEDIUM_HEAP_SIZE, SMALL_HEAP_SIZE};
     use hyperlight_testing::simple_guest_as_string;
 
@@ -462,60 +463,6 @@ mod tests {
     use crate::sandbox::SandboxConfiguration;
     use crate::sandbox::snapshot::Snapshot;
 
-    /// Convert MemoryRegionFlags to expected PTE flags.
-    fn expected_pte_flags(flags: MemoryRegionFlags) -> u64 {
-        let mut pte = PAGE_PRESENT;
-
-        // Writable if WRITE or STACK_GUARD (stack guard needs to be writable
-        // so faults go to hypervisor, not guest kernel)
-        if flags.contains(MemoryRegionFlags::WRITE)
-            || flags.contains(MemoryRegionFlags::STACK_GUARD)
-        {
-            pte |= PAGE_RW;
-        }
-
-        // NX unless EXECUTE is set
-        if !flags.contains(MemoryRegionFlags::EXECUTE) {
-            pte |= PAGE_NX;
-        }
-
-        pte
-    }
-
-    /// Walk page tables to get PTE for a virtual address.
-    /// Returns None if address is not mapped.
-    fn get_pte(memory: &[u8], pt_offset: usize, pt_base_gpa: usize, vaddr: u64) -> Option<u64> {
-        let pml4_idx = ((vaddr >> 39) & 0x1ff) as usize;
-        let pdpt_idx = ((vaddr >> 30) & 0x1ff) as usize;
-        let pd_idx = ((vaddr >> 21) & 0x1ff) as usize;
-        let pt_idx = ((vaddr >> 12) & 0x1ff) as usize;
-
-        let read_entry = |table_gpa: usize, index: usize| -> Option<u64> {
-            let entry_gpa = table_gpa + index * 8;
-            let offset = entry_gpa - pt_base_gpa + pt_offset;
-            if offset + 8 > memory.len() {
-                return None;
-            }
-            let entry = u64::from_ne_bytes(memory[offset..offset + 8].try_into().ok()?);
-            if entry & PAGE_PRESENT != 0 {
-                Some(entry)
-            } else {
-                None
-            }
-        };
-
-        let pml4e = read_entry(pt_base_gpa, pml4_idx)?;
-        let pdpt_gpa = (pml4e & PTE_ADDR_MASK) as usize;
-
-        let pdpte = read_entry(pdpt_gpa, pdpt_idx)?;
-        let pd_gpa = (pdpte & PTE_ADDR_MASK) as usize;
-
-        let pde = read_entry(pd_gpa, pd_idx)?;
-        let pt_gpa = (pde & PTE_ADDR_MASK) as usize;
-
-        read_entry(pt_gpa, pt_idx)
-    }
-
     /// Verify page tables for a given configuration.
     /// Creates a Snapshot and verifies every page in every region has correct PTEs.
     fn verify_page_tables(name: &str, config: SandboxConfiguration) {
@@ -523,48 +470,54 @@ mod tests {
         let snapshot = Snapshot::from_env(GuestBinary::FilePath(path), config)
             .unwrap_or_else(|e| panic!("{}: failed to create snapshot: {}", name, e));
 
-        let memory = snapshot.memory();
-        let layout = snapshot.layout();
-        let pt_offset = layout.get_pt_offset();
-        let pt_base_gpa = super::super::layout::SandboxMemoryLayout::BASE_ADDRESS + pt_offset;
         let regions = snapshot.regions();
 
-        // Mask for the flags we care about (PRESENT, RW, NX)
-        const FLAG_MASK: u64 = PAGE_PRESENT | PAGE_RW | PAGE_NX;
-
         // Verify NULL page (0x0) is NOT mapped
         assert!(
-            get_pte(memory, pt_offset, pt_base_gpa, 0).is_none(),
+            unsafe { hyperlight_common::vmem::virt_to_phys(&snapshot, 0, 1) }
+                .next()
+                .is_none(),
             "{}: NULL page (0x0) should NOT be mapped",
             name
         );
 
         // Verify every page in every region
         for region in regions {
-            let expected = expected_pte_flags(region.flags);
             let mut addr = region.guest_region.start as u64;
 
             while addr < region.guest_region.end as u64 {
-                let pte = get_pte(memory, pt_offset, pt_base_gpa, addr).unwrap_or_else(|| {
-                    panic!(
-                        "{}: {:?} region: address 0x{:x} is not mapped",
-                        name, region.region_type, addr
-                    )
-                });
+                let mapping = unsafe { hyperlight_common::vmem::virt_to_phys(&snapshot, addr, 1) }
+                    .next()
+                    .unwrap_or_else(|| {
+                        panic!(
+                            "{}: {:?} region: address 0x{:x} is not mapped",
+                            name, region.region_type, addr
+                        )
+                    });
 
                 // Verify identity mapping (phys == virt for low memory)
-                let phys = pte & PTE_ADDR_MASK;
                 assert_eq!(
-                    phys, addr,
+                    mapping.1, addr,
                     "{}: {:?} region: address 0x{:x} should identity map, got phys 0x{:x}",
-                    name, region.region_type, addr, phys
+                    name, region.region_type, addr, mapping.1
+                );
+
+                // Verify writable
+                let actual = mapping.2.writable;
+                let expected = region.flags.contains(MemoryRegionFlags::WRITE)
+                    || region.flags.contains(MemoryRegionFlags::STACK_GUARD);
+                assert_eq!(
+                    actual, expected,
+                    "{}: {:?} region: address 0x{:x} has writable {}, expected {} (region flags: {:?})",
+                    name, region.region_type, addr, actual, expected, region.flags
                 );
 
-                // Verify flags
-                let actual = pte & FLAG_MASK;
+                // Verify executable
+                let actual = mapping.2.executable;
+                let expected = region.flags.contains(MemoryRegionFlags::EXECUTE);
                 assert_eq!(
                     actual, expected,
-                    "{}: {:?} region: address 0x{:x} has flags 0x{:x}, expected 0x{:x} (region flags: {:?})",
+                    "{}: {:?} region: address 0x{:x} has executable {}, expected {} (region flags: {:?})",
                     name, region.region_type, addr, actual, expected, region.flags
                 );
 
diff --git a/src/hyperlight_host/src/mem/shared_mem.rs b/src/hyperlight_host/src/mem/shared_mem.rs
index f9280b773..cae355aae 100644
--- a/src/hyperlight_host/src/mem/shared_mem.rs
+++ b/src/hyperlight_host/src/mem/shared_mem.rs
@@ -582,6 +582,7 @@ impl ExclusiveSharedMemory {
 
     /// Copy the entire contents of `self` into a `Vec<u8>`, then return it
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    #[cfg(test)]
     pub(crate) fn copy_all_to_vec(&self) -> Result<Vec<u8>> {
         let data = self.as_slice();
         Ok(data.to_vec())
@@ -654,7 +655,7 @@ impl GuestSharedMemory {
         region_type: MemoryRegionType,
     ) -> MemoryRegion {
         let flags = match region_type {
-            MemoryRegionType::Scratch => {
+            MemoryRegionType::Scratch | MemoryRegionType::Snapshot => {
                 MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE
             }
             #[allow(clippy::panic)]
diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs
index 8673c6e9c..0d249aeea 100644
--- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs
+++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs
@@ -170,7 +170,13 @@ impl MultiUseSandbox {
         }
         let mapped_regions_iter = self.vm.get_mapped_regions();
         let mapped_regions_vec: Vec<MemoryRegion> = mapped_regions_iter.cloned().collect();
-        let memory_snapshot = self.mem_mgr.snapshot(self.id, mapped_regions_vec)?;
+        let root_pt_gpa = self
+            .vm
+            .get_root_pt()
+            .map_err(|e| HyperlightError::HyperlightVmError(e.into()))?;
+        let memory_snapshot = self
+            .mem_mgr
+            .snapshot(self.id, mapped_regions_vec, root_pt_gpa)?;
         let snapshot = Arc::new(memory_snapshot);
         self.snapshot = Some(snapshot.clone());
         Ok(snapshot)
@@ -197,7 +203,6 @@ impl MultiUseSandbox {
     /// - **Locked mutexes** - All lock state is reset
     /// - **Partial updates** - Data structures restored to their pre-execution state
     ///
-    ///
     /// # Examples
     ///
     /// ```no_run
@@ -254,22 +259,47 @@ impl MultiUseSandbox {
     /// ```
     #[instrument(err(Debug), skip_all, parent = Span::current())]
     pub fn restore(&mut self, snapshot: Arc<Snapshot>) -> Result<()> {
-        if let Some(snap) = &self.snapshot
-            && snap.as_ref() == snapshot.as_ref()
-        {
-            // If the snapshot is already the current one, no need to restore
-            return Ok(());
-        }
+        // Currently, we do not try to optimise restore to the
+        // most-current snapshot. This is because the most-current
+        // snapshot, while it must have identical virtual memory
+        // layout to the current sandbox, does not necessarily have
+        // the exact same /physical/ memory contents. It is not
+        // entirely inconceivable that this could lead to breakage of
+        // cross-request isolation in some way, although it would
+        // require some /very/ odd code.  For example, suppose that a
+        // service uses Hyperlight to sandbox native code from
+        // clients, and promises cross-request isolation. A tenant
+        // provides a binary that can process two forms of request,
+        // either writing a secret into physical memory, or reading
+        // from arbitrary physical memory, assuming that the two kinds
+        // of requests can never (dangerously) meet in the same
+        // sandbox.
+        //
+        // It is presently unclear whether this is a sensible threat
+        // model, especially since Hyperlight is often used with
+        // managed-code runtimes which do not allow even arbitrary
+        // access to virtual memory, much less physical memory.
+        // However, out of an abundance of caution, the optimisation
+        // is presently disabled.
 
         if self.id != snapshot.sandbox_id() {
             return Err(SnapshotSandboxMismatch);
         }
 
-        if let Some(gscratch) = self.mem_mgr.restore_snapshot(&snapshot)? {
+        let (gsnapshot, gscratch) = self.mem_mgr.restore_snapshot(&snapshot)?;
+        if let Some(gsnapshot) = gsnapshot {
+            self.vm
+                .update_snapshot_mapping(gsnapshot)
+                .map_err(|e| HyperlightError::HyperlightVmError(e.into()))?;
+        }
+        if let Some(gscratch) = gscratch {
             self.vm
                 .update_scratch_mapping(gscratch)
                 .map_err(|e| HyperlightError::HyperlightVmError(e.into()))?;
         }
+        self.vm
+            .set_root_pt(snapshot.root_pt_gpa())
+            .map_err(|e| HyperlightError::HyperlightVmError(e.into()))?;
 
         let current_regions: HashSet<_> = self.vm.get_mapped_regions().cloned().collect();
         let snapshot_regions: HashSet<_> = snapshot.regions().iter().cloned().collect();
@@ -1138,7 +1168,7 @@ mod tests {
         let actual: Vec<u8> = sbox
             .call(
                 "ReadMappedBuffer",
-                (guest_base as u64, expected.len() as u64),
+                (guest_base as u64, expected.len() as u64, true),
             )
             .unwrap();
 
@@ -1232,6 +1262,7 @@ mod tests {
 
     #[test]
     #[cfg(target_os = "linux")]
+    #[ignore] // this test will be re-enabled in the next commit
     fn snapshot_restore_handles_remapping_correctly() {
         let mut sbox: MultiUseSandbox = {
             let path = simple_guest_as_string().unwrap();
@@ -1250,32 +1281,50 @@ mod tests {
 
         unsafe { sbox.map_region(&region).unwrap() };
         assert_eq!(sbox.vm.get_mapped_regions().count(), 1);
+        let orig_read = sbox
+            .call::<Vec<u8>>(
+                "ReadMappedBuffer",
+                (
+                    guest_base as u64,
+                    hyperlight_common::vmem::PAGE_SIZE as u64,
+                    true,
+                ),
+            )
+            .unwrap();
 
         // 3. Take snapshot 2 with 1 region mapped
         let snapshot2 = sbox.snapshot().unwrap();
         assert_eq!(sbox.vm.get_mapped_regions().count(), 1);
 
-        // 4. Restore to snapshot 1 (should unmap the region)
+        // 4. Re(store to snapshot 1 (should unmap the region)
         sbox.restore(snapshot1.clone()).unwrap();
         assert_eq!(sbox.vm.get_mapped_regions().count(), 0);
+        let is_mapped = sbox
+            .call::<bool>("CheckMapped", (guest_base as u64,))
+            .unwrap();
+        assert!(!is_mapped);
 
-        // 5. Restore forward to snapshot 2 (should remap the region)
+        // 5. Restore forward to snapshot 2 (should have folded the
+        //    region into the snapshot)
         sbox.restore(snapshot2.clone()).unwrap();
-        assert_eq!(sbox.vm.get_mapped_regions().count(), 1);
+        assert_eq!(sbox.vm.get_mapped_regions().count(), 0);
+        let is_mapped = sbox
+            .call::<bool>("CheckMapped", (guest_base as u64,))
+            .unwrap();
+        assert!(is_mapped);
 
         // Verify the region is the same
-        let mut restored_regions = sbox.vm.get_mapped_regions();
-        assert_eq!(restored_regions.next().unwrap(), &region);
-        assert!(restored_regions.next().is_none());
-        drop(restored_regions);
-
-        // 6. Try map the region again (should fail since already mapped)
-        let err = unsafe { sbox.map_region(&region) };
-        assert!(
-            err.is_err(),
-            "Expected error when remapping existing region: {:?}",
-            err
-        );
+        let new_read = sbox
+            .call::<Vec<u8>>(
+                "ReadMappedBuffer",
+                (
+                    guest_base as u64,
+                    hyperlight_common::vmem::PAGE_SIZE as u64,
+                    false,
+                ),
+            )
+            .unwrap();
+        assert_eq!(new_read, orig_read);
     }
 
     #[test]
diff --git a/src/hyperlight_host/src/sandbox/snapshot.rs b/src/hyperlight_host/src/sandbox/snapshot.rs
index 6a083098a..839b7b743 100644
--- a/src/hyperlight_host/src/sandbox/snapshot.rs
+++ b/src/hyperlight_host/src/sandbox/snapshot.rs
@@ -16,13 +16,17 @@ limitations under the License.
 
 use std::sync::atomic::{AtomicU64, Ordering};
 
+use hyperlight_common::layout::{scratch_base_gpa, scratch_base_gva};
+use hyperlight_common::vmem::{self, BasicMapping, Mapping, MappingKind, PAGE_SIZE};
 use tracing::{Span, instrument};
 
 use crate::HyperlightError::MemoryRegionSizeMismatch;
 use crate::Result;
 use crate::mem::exe::LoadInfo;
+use crate::mem::layout::SandboxMemoryLayout;
 use crate::mem::memory_region::MemoryRegion;
-use crate::mem::shared_mem::SharedMemory;
+use crate::mem::mgr::GuestPageTableBuffer;
+use crate::mem::shared_mem::{ExclusiveSharedMemory, SharedMemory};
 use crate::sandbox::SandboxConfiguration;
 use crate::sandbox::uninitialized::{GuestBinary, GuestEnvironment};
 pub(super) static SANDBOX_CONFIGURATION_COUNTER: AtomicU64 = AtomicU64::new(0);
@@ -61,6 +65,8 @@ pub struct Snapshot {
     /// It is not a [`blake3::Hash`] because we do not presently
     /// require constant-time equality checking
     hash: [u8; 32],
+    /// The address of the root page table
+    root_pt_gpa: u64,
 
     /// Preinitialisation entry point for snapshots created directly from a
     /// guest binary.
@@ -77,6 +83,59 @@ pub struct Snapshot {
     /// snapshot format and must be preserved.
     preinitialise: Option<u64>,
 }
+impl core::convert::AsRef<Snapshot> for Snapshot {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
+impl hyperlight_common::vmem::TableOps for Snapshot {
+    type TableAddr = u64;
+    // This entire impl is only used in contexts which should
+    // absolutely never try to update a page table, and so in no case
+    // should possibly ever call an update_parent.  Therefore, this is
+    // impossible unless an extremely significant invariant violation
+    // has occurred.
+    #[allow(clippy::panic)]
+    unsafe fn alloc_table(&self) -> u64 {
+        panic!("alloc_table: cannot mutate the page tables inside a snapshot");
+    }
+    fn entry_addr(addr: u64, offset: u64) -> u64 {
+        addr + offset
+    }
+    unsafe fn read_entry(&self, addr: u64) -> u64 {
+        let addr = addr as usize;
+        let Some(pte_bytes) = self.memory.as_slice().get(addr..addr + 8) else {
+            // Attacker-controlled data pointed out-of-bounds. We'll
+            // default to returning 0 in this case, which, for most
+            // architectures (including x86-64 and arm64, the ones we
+            // care about presently) will be a not-present entry.
+            return 0;
+        };
+        // this is statically the correct size, so using unwrap() here
+        // doesn't make this any more panic-y.
+        #[allow(clippy::unwrap_used)]
+        let n: [u8; 8] = pte_bytes.try_into().unwrap();
+        u64::from_ne_bytes(n)
+    }
+    // This entire impl is only used in contexts which should
+    // absolutely never try to update a page table, and so in no case
+    // should possibly ever call an update_parent.  Therefore, this is
+    // impossible unless an extremely significant invariant violation
+    // has occurred.
+    #[allow(clippy::panic)]
+    unsafe fn write_entry(&self, _addr: u64, _entry: u64) -> Option<u64> {
+        panic!("write_entry: cannot mutate the page tables inside a snapshot");
+    }
+    fn to_phys(addr: u64) -> u64 {
+        addr
+    }
+    fn from_phys(addr: u64) -> u64 {
+        addr
+    }
+    fn root_table(&self) -> u64 {
+        self.root_pt_gpa
+    }
+}
 
 /// Compute a deterministic hash of a snapshot.
 ///
@@ -113,6 +172,188 @@ fn hash(memory: &[u8], regions: &[MemoryRegion]) -> Result<[u8; 32]> {
     Ok(hasher.finalize().into())
 }
 
+fn access_gpa<'a>(
+    snap: &'a ExclusiveSharedMemory,
+    scratch: &'a ExclusiveSharedMemory,
+    scratch_size: usize,
+    gpa: u64,
+) -> (&'a ExclusiveSharedMemory, usize) {
+    let scratch_base = scratch_base_gpa(scratch_size);
+    if gpa >= scratch_base {
+        (scratch, (gpa - scratch_base) as usize)
+    } else {
+        (snap, gpa as usize - SandboxMemoryLayout::BASE_ADDRESS)
+    }
+}
+
+pub(crate) struct SharedMemoryPageTableBuffer<'a> {
+    snap: &'a ExclusiveSharedMemory,
+    scratch: &'a ExclusiveSharedMemory,
+    scratch_size: usize,
+    root: u64,
+}
+impl<'a> SharedMemoryPageTableBuffer<'a> {
+    fn new(
+        snap: &'a ExclusiveSharedMemory,
+        scratch: &'a ExclusiveSharedMemory,
+        scratch_size: usize,
+        root: u64,
+    ) -> Self {
+        Self {
+            snap,
+            scratch,
+            scratch_size,
+            root,
+        }
+    }
+}
+impl<'a> hyperlight_common::vmem::TableOps for SharedMemoryPageTableBuffer<'a> {
+    type TableAddr = u64;
+    // This entire impl is only used in contexts which should
+    // absolutely never try to update a page table, and so in no case
+    // should possibly ever call an update_parent.  Therefore, this is
+    // impossible unless an extremely significant invariant violation
+    // has occurred.
+    #[allow(clippy::panic)]
+    unsafe fn alloc_table(&self) -> u64 {
+        panic!("alloc_table: cannot mutate the page tables inside the guest");
+    }
+    fn entry_addr(addr: u64, offset: u64) -> u64 {
+        addr + offset
+    }
+    unsafe fn read_entry(&self, addr: u64) -> u64 {
+        let (mem, off) = access_gpa(self.snap, self.scratch, self.scratch_size, addr);
+        let Some(pte_bytes) = mem.as_slice().get(off..off + 8) else {
+            // Attacker-controlled data pointed out-of-bounds. We'll
+            // default to returning 0 in this case, which, for most
+            // architectures (including x86-64 and arm64, the ones we
+            // care about presently) will be a not-present entry.
+            return 0;
+        };
+        // this is statically the correct size, so using unwrap() here
+        // doesn't make this any more panic-y.
+        #[allow(clippy::unwrap_used)]
+        let n: [u8; 8] = pte_bytes.try_into().unwrap();
+        u64::from_ne_bytes(n)
+    }
+    // This entire impl is only used in contexts which should
+    // absolutely never try to update a page table, and so in no case
+    // should possibly ever call an update_parent.  Therefore, this is
+    // impossible unless an extremely significant invariant violation
+    // has occurred.
+    #[allow(clippy::panic)]
+    unsafe fn write_entry(&self, _addr: u64, _entry: u64) -> Option<u64> {
+        panic!("write_entry: cannot mutate the page tables inside the guest");
+    }
+    fn to_phys(addr: u64) -> u64 {
+        addr
+    }
+    fn from_phys(addr: u64) -> u64 {
+        addr
+    }
+    fn root_table(&self) -> u64 {
+        self.root
+    }
+}
+impl<'a> core::convert::AsRef<SharedMemoryPageTableBuffer<'a>> for SharedMemoryPageTableBuffer<'a> {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
+fn filtered_mappings<'a>(
+    snap: &'a ExclusiveSharedMemory,
+    scratch: &'a ExclusiveSharedMemory,
+    regions: &[MemoryRegion],
+    scratch_size: usize,
+    root_pt: u64,
+) -> Vec<(u64, u64, BasicMapping, &'a [u8])> {
+    let op = SharedMemoryPageTableBuffer::new(snap, scratch, scratch_size, root_pt);
+    unsafe {
+        hyperlight_common::vmem::virt_to_phys(&op, 0, hyperlight_common::layout::MAX_GVA as u64)
+    }
+    .filter_map(move |(gva, gpa, bm)| {
+        // the scratch map doesn't count
+        if gva > scratch_base_gva(scratch_size) {
+            return None;
+        }
+        // neither does the mapping of the snapshot's own page tables
+        if gva >= hyperlight_common::layout::SNAPSHOT_PT_GVA_MIN as u64
+            && gva <= hyperlight_common::layout::SNAPSHOT_PT_GVA_MAX as u64
+        {
+            return None;
+        }
+        // todo: is it useful to warn if we can't resolve this?
+        let contents = unsafe { guest_page(snap, scratch, regions, scratch_size, gpa) }?;
+        Some((gva, gpa, bm, contents))
+    })
+    .collect()
+}
+
+/// Find the contents of the page which starts at gpa in guest physical
+/// memory, taking into account excess host->guest regions
+///
+/// # Safety
+/// The host side of the regions identified by MemoryRegion must be
+/// alive and must not be mutated by any other thread: references to
+/// these regions may be created and live for `'a`.
+unsafe fn guest_page<'a>(
+    snap: &'a ExclusiveSharedMemory,
+    scratch: &'a ExclusiveSharedMemory,
+    regions: &[MemoryRegion],
+    scratch_size: usize,
+    gpa: u64,
+) -> Option<&'a [u8]> {
+    let gpa_u = gpa as usize;
+    for rgn in regions {
+        if gpa_u >= rgn.guest_region.start && gpa_u < rgn.guest_region.end {
+            let off = gpa_u - rgn.guest_region.start;
+            #[allow(clippy::useless_conversion)]
+            let host_region_base: usize = rgn.host_region.start.into();
+            return Some(unsafe {
+                std::slice::from_raw_parts((host_region_base + off) as *const u8, PAGE_SIZE)
+            });
+        }
+    }
+    let (mem, off) = access_gpa(snap, scratch, scratch_size, gpa);
+    if off + PAGE_SIZE <= mem.as_slice().len() {
+        Some(&mem.as_slice()[off..off + PAGE_SIZE])
+    } else {
+        None
+    }
+}
+
+fn map_specials(pt_buf: &GuestPageTableBuffer, scratch_size: usize) {
+    // Map the scratch region
+    let mapping = Mapping {
+        phys_base: scratch_base_gpa(scratch_size),
+        virt_base: scratch_base_gva(scratch_size),
+        len: scratch_size as u64,
+        kind: MappingKind::BasicMapping(BasicMapping {
+            readable: true,
+            writable: true,
+            executable: true,
+        }),
+    };
+    unsafe { vmem::map(pt_buf, mapping) };
+    // Map the page tables themselves, in order to allow the
+    // guest to update them easily
+    let mut pt_size_mapped = 0;
+    while pt_buf.size() > pt_size_mapped {
+        let mapping = Mapping {
+            phys_base: (pt_buf.phys_base() + pt_size_mapped) as u64,
+            virt_base: (hyperlight_common::layout::SNAPSHOT_PT_GVA_MIN + pt_size_mapped) as u64,
+            len: (pt_buf.size() - pt_size_mapped) as u64,
+            kind: MappingKind::BasicMapping(BasicMapping {
+                readable: true,
+                writable: true,
+                executable: false,
+            }),
+        };
+        pt_size_mapped = pt_buf.size();
+        unsafe { vmem::map(pt_buf, mapping) };
+    }
+}
+
 impl Snapshot {
     /// Create a new snapshot from the guest binary identified by `env`. With the configuration
     /// specified in `cfg`.
@@ -159,12 +400,11 @@ impl Snapshot {
             .transpose()?;
 
         #[cfg(feature = "init-paging")]
-        {
+        let pt_base_gpa = {
             // Set up page table entries for the snapshot
             let pt_base_gpa =
                 crate::mem::layout::SandboxMemoryLayout::BASE_ADDRESS + layout.get_pt_offset();
-            let pt_buf = crate::mem::mgr::GuestPageTableBuffer::new(pt_base_gpa);
-            use hyperlight_common::vmem::{self, BasicMapping, Mapping, MappingKind};
+            let pt_buf = GuestPageTableBuffer::new(pt_base_gpa);
 
             use crate::mem::memory_region::{GuestMemoryRegion, MemoryRegionFlags};
 
@@ -193,41 +433,16 @@ impl Snapshot {
                 unsafe { vmem::map(&pt_buf, mapping) };
             }
 
-            // 2. Map the scratch region
-            let scratch_size = cfg.get_scratch_size();
-            let mapping = Mapping {
-                phys_base: hyperlight_common::layout::scratch_base_gpa(scratch_size),
-                virt_base: hyperlight_common::layout::scratch_base_gva(scratch_size),
-                len: scratch_size as u64,
-                kind: MappingKind::BasicMapping(BasicMapping {
-                    readable: true,
-                    writable: true,
-                    executable: true,
-                }),
-            };
-            unsafe { vmem::map(&pt_buf, mapping) };
-
-            // 3. Map the page tables themselves, in order to allow the
-            // guest to update them easily
-            let mut pt_size_mapped = 0;
-            while pt_buf.size() > pt_size_mapped {
-                let mapping = Mapping {
-                    phys_base: (pt_base_gpa + pt_size_mapped) as u64,
-                    virt_base: (hyperlight_common::layout::SNAPSHOT_PT_GVA + pt_size_mapped) as u64,
-                    len: (pt_buf.size() - pt_size_mapped) as u64,
-                    kind: MappingKind::BasicMapping(BasicMapping {
-                        readable: true,
-                        writable: false,
-                        executable: false,
-                    }),
-                };
-                unsafe { vmem::map(&pt_buf, mapping) };
-                pt_size_mapped = pt_buf.size();
-            }
+            // 2. Map the special mappings
+            map_specials(&pt_buf, layout.get_scratch_size());
+
             let pt_bytes = pt_buf.into_bytes();
             layout.set_pt_size(pt_bytes.len());
             memory.extend(&pt_bytes);
-        }
+            pt_base_gpa
+        };
+        #[cfg(not(feature = "init-paging"))]
+        let pt_base_gpa = 0usize;
 
         let extra_regions = Vec::new();
         let hash = hash(&memory, &extra_regions)?;
@@ -239,6 +454,7 @@ impl Snapshot {
             regions: extra_regions,
             load_info,
             hash,
+            root_pt_gpa: pt_base_gpa as u64,
             preinitialise: Some(load_addr + entrypoint_offset),
         })
     }
@@ -248,13 +464,52 @@ impl Snapshot {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(crate) fn new<S: SharedMemory>(
         shared_mem: &mut S,
+        scratch_mem: &mut S,
         sandbox_id: u64,
-        layout: crate::mem::layout::SandboxMemoryLayout,
+        mut layout: SandboxMemoryLayout,
         load_info: LoadInfo,
         regions: Vec<MemoryRegion>,
+        root_pt_gpa: u64,
     ) -> Result<Self> {
-        // TODO: Track dirty pages instead of copying entire memory
-        let memory = shared_mem.with_exclusivity(|e| e.copy_all_to_vec())??;
+        let (new_root_pt_gpa, memory) = shared_mem.with_exclusivity(|snap_e| {
+            scratch_mem.with_exclusivity(|scratch_e| {
+                let scratch_size = layout.get_scratch_size();
+
+                // Pass 1: count how many pages need to live
+                let live_pages =
+                    filtered_mappings(snap_e, scratch_e, &regions, scratch_size, root_pt_gpa);
+
+                // Pass 2: copy them, and map them
+                // TODO: Look for opportunities to hugepage map
+                let pt_base_gpa = SandboxMemoryLayout::BASE_ADDRESS + live_pages.len() * PAGE_SIZE;
+                let pt_buf = GuestPageTableBuffer::new(pt_base_gpa);
+                let mut snapshot_memory: Vec<u8> = Vec::new();
+                for (gva, _, bm, contents) in live_pages {
+                    let new_offset = snapshot_memory.len();
+                    snapshot_memory.extend(contents);
+                    let new_gpa = new_offset + SandboxMemoryLayout::BASE_ADDRESS;
+                    let mapping = Mapping {
+                        phys_base: new_gpa as u64,
+                        virt_base: gva,
+                        len: PAGE_SIZE as u64,
+                        kind: MappingKind::BasicMapping(bm),
+                    };
+                    unsafe { vmem::map(&pt_buf, mapping) };
+                }
+                // Phase 3: Map the special mappings
+                map_specials(&pt_buf, layout.get_scratch_size());
+                let pt_bytes = pt_buf.into_bytes();
+                layout.set_pt_size(pt_bytes.len());
+                snapshot_memory.extend(&pt_bytes);
+                (pt_base_gpa, snapshot_memory)
+            })
+        })??;
+
+        // We do not need the original regions anymore, as any uses of
+        // them in the guest have been incorporated into the snapshot
+        // properly.
+        let regions = Vec::new();
+
         let hash = hash(&memory, &regions)?;
         Ok(Self {
             sandbox_id,
@@ -263,6 +518,7 @@ impl Snapshot {
             regions,
             load_info,
             hash,
+            root_pt_gpa: new_root_pt_gpa as u64,
             preinitialise: None,
         })
     }
@@ -298,6 +554,10 @@ impl Snapshot {
         &self.layout
     }
 
+    pub(crate) fn root_pt_gpa(&self) -> u64 {
+        self.root_pt_gpa
+    }
+
     pub(crate) fn preinitialise(&self) -> Option<u64> {
         self.preinitialise
     }
@@ -311,60 +571,94 @@ impl PartialEq for Snapshot {
 
 #[cfg(test)]
 mod tests {
-    use hyperlight_common::mem::PAGE_SIZE_USIZE;
+    use hyperlight_common::vmem::{self, BasicMapping, Mapping, MappingKind, PAGE_SIZE};
 
     use crate::mem::exe::LoadInfo;
-    use crate::mem::shared_mem::{ExclusiveSharedMemory, SharedMemory};
+    use crate::mem::layout::SandboxMemoryLayout;
+    use crate::mem::mgr::{GuestPageTableBuffer, SandboxMemoryManager};
+    use crate::mem::shared_mem::{ExclusiveSharedMemory, HostSharedMemory, SharedMemory};
+
+    fn make_simple_pt_mems() -> (SandboxMemoryManager<HostSharedMemory>, u64) {
+        let scratch_mem = ExclusiveSharedMemory::new(PAGE_SIZE).unwrap();
+        let pt_base = PAGE_SIZE + SandboxMemoryLayout::BASE_ADDRESS;
+        let pt_buf = GuestPageTableBuffer::new(pt_base);
+        let mapping = Mapping {
+            phys_base: SandboxMemoryLayout::BASE_ADDRESS as u64,
+            virt_base: SandboxMemoryLayout::BASE_ADDRESS as u64,
+            len: PAGE_SIZE as u64,
+            kind: MappingKind::BasicMapping(BasicMapping {
+                readable: true,
+                writable: true,
+                executable: true,
+            }),
+        };
+        unsafe { vmem::map(&pt_buf, mapping) };
+        super::map_specials(&pt_buf, PAGE_SIZE);
+        let pt_bytes = pt_buf.into_bytes();
+
+        let mut snapshot_mem = ExclusiveSharedMemory::new(PAGE_SIZE + pt_bytes.len()).unwrap();
+
+        snapshot_mem.copy_from_slice(&pt_bytes, PAGE_SIZE).unwrap();
+        let cfg = crate::sandbox::SandboxConfiguration::default();
+        let mgr = SandboxMemoryManager::new(
+            SandboxMemoryLayout::new(cfg, 4096, 2048, 4096, 0x3000, 0, None).unwrap(),
+            snapshot_mem,
+            scratch_mem,
+            0.into(),
+            None,
+            [0u8; 16],
+        );
+        let (mgr, _) = mgr.build();
+        (mgr, pt_base as u64)
+    }
 
     #[test]
     fn restore() {
         // Simplified version of the original test
-        let data1 = vec![b'a'; PAGE_SIZE_USIZE];
-        let data2 = vec![b'b'; PAGE_SIZE_USIZE];
+        let data1 = vec![b'a'; PAGE_SIZE];
+        let data2 = vec![b'b'; PAGE_SIZE];
 
-        let mut gm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap();
-        gm.copy_from_slice(&data1, 0).unwrap();
-
-        let cfg = crate::sandbox::SandboxConfiguration::default();
-        let layout =
-            crate::mem::layout::SandboxMemoryLayout::new(cfg, 4096, 2048, 4096, 0x3000, 0, None)
-                .unwrap();
+        let (mut mgr, pt_base) = make_simple_pt_mems();
+        mgr.shared_mem.copy_from_slice(&data1, 0).unwrap();
 
         // Take snapshot of data1
         let snapshot = super::Snapshot::new(
-            &mut gm,
+            &mut mgr.shared_mem,
+            &mut mgr.scratch_mem,
             0,
-            layout,
-            crate::mem::exe::LoadInfo::dummy(),
+            mgr.layout,
+            LoadInfo::dummy(),
             Vec::new(),
+            pt_base,
         )
         .unwrap();
 
         // Modify memory to data2
-        gm.copy_from_slice(&data2, 0).unwrap();
-        assert_eq!(gm.as_slice(), &data2[..]);
+        mgr.shared_mem.copy_from_slice(&data2, 0).unwrap();
+        mgr.shared_mem
+            .with_exclusivity(|e| assert_eq!(&e.as_slice()[0..data2.len()], &data2[..]))
+            .unwrap();
 
         // Restore should bring back data1
-        gm.restore_from_snapshot(&snapshot).unwrap();
-        assert_eq!(gm.as_slice(), &data1[..]);
+        let _ = mgr.restore_snapshot(&snapshot).unwrap();
+        mgr.shared_mem
+            .with_exclusivity(|e| assert_eq!(&e.as_slice()[0..data1.len()], &data1[..]))
+            .unwrap();
     }
 
     #[test]
     fn snapshot_mem_size() {
-        let size = PAGE_SIZE_USIZE * 2;
-        let mut gm = ExclusiveSharedMemory::new(size).unwrap();
-
-        let cfg = crate::sandbox::SandboxConfiguration::default();
-        let layout =
-            crate::mem::layout::SandboxMemoryLayout::new(cfg, 4096, 2048, 4096, 0x3000, 0, None)
-                .unwrap();
+        let (mut mgr, pt_base) = make_simple_pt_mems();
+        let size = mgr.shared_mem.mem_size();
 
         let snapshot = super::Snapshot::new(
-            &mut gm,
+            &mut mgr.shared_mem,
+            &mut mgr.scratch_mem,
             0,
-            layout,
-            crate::mem::exe::LoadInfo::dummy(),
+            mgr.layout,
+            LoadInfo::dummy(),
             Vec::new(),
+            pt_base,
         )
         .unwrap();
         assert_eq!(snapshot.mem_size(), size);
@@ -372,34 +666,49 @@ mod tests {
 
     #[test]
     fn multiple_snapshots_independent() {
-        let mut gm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap();
-
-        let cfg = crate::sandbox::SandboxConfiguration::default();
-        let layout =
-            crate::mem::layout::SandboxMemoryLayout::new(cfg, 4096, 2048, 4096, 0x3000, 0, None)
-                .unwrap();
+        let (mut mgr, pt_base) = make_simple_pt_mems();
 
         // Create first snapshot with pattern A
-        let pattern_a = vec![0xAA; PAGE_SIZE_USIZE];
-        gm.copy_from_slice(&pattern_a, 0).unwrap();
-        let snapshot_a =
-            super::Snapshot::new(&mut gm, 1, layout, LoadInfo::dummy(), Vec::new()).unwrap();
+        let pattern_a = vec![0xAA; PAGE_SIZE];
+        mgr.shared_mem.copy_from_slice(&pattern_a, 0).unwrap();
+        let snapshot_a = super::Snapshot::new(
+            &mut mgr.shared_mem,
+            &mut mgr.scratch_mem,
+            1,
+            mgr.layout,
+            LoadInfo::dummy(),
+            Vec::new(),
+            pt_base,
+        )
+        .unwrap();
 
         // Create second snapshot with pattern B
-        let pattern_b = vec![0xBB; PAGE_SIZE_USIZE];
-        gm.copy_from_slice(&pattern_b, 0).unwrap();
-        let snapshot_b =
-            super::Snapshot::new(&mut gm, 2, layout, LoadInfo::dummy(), Vec::new()).unwrap();
+        let pattern_b = vec![0xBB; PAGE_SIZE];
+        mgr.shared_mem.copy_from_slice(&pattern_b, 0).unwrap();
+        let snapshot_b = super::Snapshot::new(
+            &mut mgr.shared_mem,
+            &mut mgr.scratch_mem,
+            2,
+            mgr.layout,
+            LoadInfo::dummy(),
+            Vec::new(),
+            pt_base,
+        )
+        .unwrap();
 
         // Clear memory
-        gm.copy_from_slice(&[0; PAGE_SIZE_USIZE], 0).unwrap();
+        mgr.shared_mem.copy_from_slice(&[0; PAGE_SIZE], 0).unwrap();
 
         // Restore snapshot A
-        gm.restore_from_snapshot(&snapshot_a).unwrap();
-        assert_eq!(gm.as_slice(), &pattern_a[..]);
+        mgr.restore_snapshot(&snapshot_a).unwrap();
+        mgr.shared_mem
+            .with_exclusivity(|e| assert_eq!(&e.as_slice()[0..pattern_a.len()], &pattern_a[..]))
+            .unwrap();
 
         // Restore snapshot B
-        gm.restore_from_snapshot(&snapshot_b).unwrap();
-        assert_eq!(gm.as_slice(), &pattern_b[..]);
+        mgr.restore_snapshot(&snapshot_b).unwrap();
+        mgr.shared_mem
+            .with_exclusivity(|e| assert_eq!(&e.as_slice()[0..pattern_b.len()], &pattern_b[..]))
+            .unwrap();
     }
 }
diff --git a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs
index c735773d1..aeaff030a 100644
--- a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs
+++ b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs
@@ -113,8 +113,6 @@ pub(crate) fn set_up_hypervisor_partition(
     #[cfg(not(feature = "init-paging"))]
     let rsp_ptr = GuestPtr::try_from(Offset::from(0))?;
 
-    let regions = mgr.layout.get_memory_regions(&mgr.shared_mem)?;
-
     let pml4_ptr = {
         let pml4_offset_u64 = mgr.layout.get_pt_offset() as u64;
         base_ptr + Offset::from(pml4_offset_u64)
@@ -152,7 +150,7 @@ pub(crate) fn set_up_hypervisor_partition(
     let trace_info = MemTraceInfo::new(_load_info.info)?;
 
     Ok(HyperlightVm::new(
-        regions,
+        mgr.shared_mem,
         mgr.scratch_mem,
         pml4_ptr.absolute()?,
         entrypoint_ptr.absolute()?,
diff --git a/src/hyperlight_host/tests/integration_test.rs b/src/hyperlight_host/tests/integration_test.rs
index 9a395ee81..713509a83 100644
--- a/src/hyperlight_host/tests/integration_test.rs
+++ b/src/hyperlight_host/tests/integration_test.rs
@@ -605,6 +605,7 @@ fn guest_panic_no_alloc() {
 
 // Tests libc alloca
 #[test]
+#[ignore] // this test will be re-enabled in the next PR
 fn dynamic_stack_allocate_c_guest() {
     let path = c_simple_guest_as_string().unwrap();
     let guest_path = GuestBinary::FilePath(path);
@@ -631,6 +632,7 @@ fn static_stack_allocate() {
 
 // checks that a huge buffer on stack fails with stackoverflow
 #[test]
+#[ignore] // this test will be re-enabled in the next PR
 fn static_stack_allocate_overflow() {
     let mut sbox1 = new_uninit().unwrap().evolve().unwrap();
     let res = sbox1.call::<i32>("LargeVar", ()).unwrap_err();
@@ -639,6 +641,7 @@ fn static_stack_allocate_overflow() {
 
 // checks that a recursive function with stack allocation works, (that chkstk can be called without overflowing)
 #[test]
+#[ignore] // this test will be re-enabled in the next PR
 fn recursive_stack_allocate() {
     let mut sbox1 = new_uninit().unwrap().evolve().unwrap();
 
@@ -650,6 +653,7 @@ fn recursive_stack_allocate() {
 // checks stack guard page (between guest stack and heap)
 // is properly set up and cannot be written to
 #[test]
+#[ignore] // this test will be re-enabled in the next PR
 fn guard_page_check() {
     // this test is rust-guest only
     let offsets_from_page_guard_start: Vec<i64> = vec![
@@ -683,6 +687,7 @@ fn guard_page_check() {
 }
 
 #[test]
+#[ignore] // this test will be re-enabled in the next PR
 fn guard_page_check_2() {
     // this test is rust-guest only
     let mut sbox1 = new_uninit_rust().unwrap().evolve().unwrap();
@@ -725,6 +730,7 @@ fn execute_on_heap() {
 
 // checks that a recursive function with stack allocation eventually fails with stackoverflow
 #[test]
+#[ignore] // this test will be re-enabled in the next PR
 fn recursive_stack_allocate_overflow() {
     let mut sbox1 = new_uninit().unwrap().evolve().unwrap();
 
diff --git a/src/tests/rust_guests/simpleguest/src/main.rs b/src/tests/rust_guests/simpleguest/src/main.rs
index df7de9469..cb9e99848 100644
--- a/src/tests/rust_guests/simpleguest/src/main.rs
+++ b/src/tests/rust_guests/simpleguest/src/main.rs
@@ -602,17 +602,28 @@ fn read_from_user_memory(num: u64, expected: Vec<u8>) -> Result<Vec<u8>> {
 }
 
 #[guest_function("ReadMappedBuffer")]
-fn read_mapped_buffer(base: u64, len: u64) -> Vec<u8> {
+fn read_mapped_buffer(base: u64, len: u64, do_map: bool) -> Vec<u8> {
     let base = base as usize as *const u8;
     let len = len as usize;
 
-    unsafe { hyperlight_guest_bin::paging::map_region(base as _, base as _, len as u64 + 4096) };
+    if do_map {
+        unsafe {
+            hyperlight_guest_bin::paging::map_region(base as _, base as _, len as u64 + 4096)
+        };
+    }
 
     let data = unsafe { core::slice::from_raw_parts(base, len) };
 
     data.to_vec()
 }
 
+#[guest_function("CheckMapped")]
+fn check_mapped_buffer(base: u64) -> bool {
+    hyperlight_guest_bin::paging::virt_to_phys(base)
+        .next()
+        .is_some()
+}
+
 #[guest_function("WriteMappedBuffer")]
 fn write_mapped_buffer(base: u64, len: u64) -> bool {
     let base = base as usize as *mut u8;

From 350a605f70c9e6f69376dcfa2d3351ca6347c5f4 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Wed, 28 Jan 2026 01:58:46 +0000
Subject: [PATCH 14/20] Explicitly specify the physical addresses of snapshot
 page tables

Previously, the guest tried to infer the addresses of the snapshot
page tables from the values of cr3 that it saw at various times, but
this was brittle and irritating to do without writing to the wrong
address at times.  This replaces that mechanism with an extra word in
the metadata portion of the scratch region that explicitly provides
this information.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_common/src/layout.rs               |  3 ++-
 src/hyperlight_guest/src/layout.rs                |  4 ++++
 src/hyperlight_guest_bin/src/paging.rs            | 15 +++------------
 src/hyperlight_host/src/mem/mgr.rs                | 15 ++++++++++++---
 .../src/sandbox/initialized_multi_use.rs          |  1 -
 5 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs
index 1c9d0ed71..0b0bad9b0 100644
--- a/src/hyperlight_common/src/layout.rs
+++ b/src/hyperlight_common/src/layout.rs
@@ -23,7 +23,8 @@ pub use arch::{MAX_GPA, MAX_GVA, SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN};
 // offsets down from the top of scratch memory for various things
 pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08;
 pub const SCRATCH_TOP_USED_OFFSET: u64 = 0x10;
-pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x18;
+pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x10;
+pub const SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET: u64 = 0x18;
 pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x20;
 
 pub fn scratch_base_gpa(size: usize) -> u64 {
diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs
index e7ac0a5e3..929fe7770 100644
--- a/src/hyperlight_guest/src/layout.rs
+++ b/src/hyperlight_guest/src/layout.rs
@@ -27,4 +27,8 @@ pub fn allocator_gva() -> *mut u64 {
     use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_ALLOCATOR_OFFSET};
     (MAX_GVA as u64 - SCRATCH_TOP_ALLOCATOR_OFFSET + 1) as *mut u64
 }
+pub fn snapshot_pt_gpa_base_gva() -> *mut u64 {
+    use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET};
+    (MAX_GVA as u64 - SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET + 1) as *mut u64
+}
 pub use arch::{scratch_base_gpa, scratch_base_gva};
diff --git a/src/hyperlight_guest_bin/src/paging.rs b/src/hyperlight_guest_bin/src/paging.rs
index b806bf884..85e4fa31b 100644
--- a/src/hyperlight_guest_bin/src/paging.rs
+++ b/src/hyperlight_guest_bin/src/paging.rs
@@ -28,11 +28,6 @@ use tracing::{Span, instrument};
 //       virtual address 0, and Rust raw pointer operations can't be
 //       used to read/write from address 0.
 
-// We get this out of CR3 the first time that we do any mapping
-// operation. In the future, if snapshot/restore changes to be able to
-// change the snapshot pt base, we will need to modify this.
-static SNAPSHOT_PT_GPA: spin::Once<u64> = spin::Once::new();
-
 #[derive(Copy, Clone)]
 struct GuestMappingOperations {
     snapshot_pt_base_gpa: u64,
@@ -43,13 +38,9 @@ struct GuestMappingOperations {
 impl GuestMappingOperations {
     fn new() -> Self {
         Self {
-            snapshot_pt_base_gpa: *SNAPSHOT_PT_GPA.call_once(|| {
-                let snapshot_pt_base_gpa: u64;
-                unsafe {
-                    asm!("mov {}, cr3", out(reg) snapshot_pt_base_gpa);
-                };
-                snapshot_pt_base_gpa
-            }),
+            snapshot_pt_base_gpa: unsafe {
+                hyperlight_guest::layout::snapshot_pt_gpa_base_gva().read_volatile()
+            },
             snapshot_pt_base_gva: hyperlight_common::layout::SNAPSHOT_PT_GVA_MIN as u64,
             scratch_base_gpa: hyperlight_guest::layout::scratch_base_gpa(),
             scratch_base_gva: hyperlight_guest::layout::scratch_base_gva(),
diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs
index 20e92c558..e8a9754ec 100644
--- a/src/hyperlight_host/src/mem/mgr.rs
+++ b/src/hyperlight_host/src/mem/mgr.rs
@@ -263,7 +263,9 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             stack_cookie: self.stack_cookie,
             abort_buffer: Vec::new(), // Guest doesn't need abort buffer
         };
-        host_mgr.update_scratch_bookkeeping();
+        host_mgr.update_scratch_bookkeeping(
+            (SandboxMemoryLayout::BASE_ADDRESS + self.layout.get_pt_offset()) as u64,
+        );
         (host_mgr, guest_mgr)
     }
 }
@@ -422,11 +424,11 @@ impl SandboxMemoryManager<HostSharedMemory> {
 
             Some(gscratch)
         };
-        self.update_scratch_bookkeeping();
+        self.update_scratch_bookkeeping(snapshot.root_pt_gpa());
         Ok((gsnapshot, gscratch))
     }
 
-    fn update_scratch_bookkeeping(&mut self) {
+    fn update_scratch_bookkeeping(&mut self, snapshot_pt_base_gpa: u64) {
         let scratch_size = self.scratch_mem.mem_size();
         let size_offset =
             scratch_size - hyperlight_common::layout::SCRATCH_TOP_SIZE_OFFSET as usize;
@@ -448,6 +450,13 @@ impl SandboxMemoryManager<HostSharedMemory> {
                 hyperlight_common::layout::scratch_base_gpa(scratch_size),
             )
             .unwrap();
+        let snapshot_pt_base_gpa_offset = scratch_size
+            - hyperlight_common::layout::SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET as usize;
+        // See above comment about unwrap() on write
+        #[allow(clippy::unwrap_used)]
+        self.scratch_mem
+            .write::<u64>(snapshot_pt_base_gpa_offset, snapshot_pt_base_gpa)
+            .unwrap();
     }
 }
 
diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs
index 0d249aeea..9d82b22d7 100644
--- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs
+++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs
@@ -1262,7 +1262,6 @@ mod tests {
 
     #[test]
     #[cfg(target_os = "linux")]
-    #[ignore] // this test will be re-enabled in the next commit
     fn snapshot_restore_handles_remapping_correctly() {
         let mut sbox: MultiUseSandbox = {
             let path = simple_guest_as_string().unwrap();

From b1274bd1254d016f50e19b4b411e7b936fb62656 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Wed, 28 Jan 2026 13:25:16 +0000
Subject: [PATCH 15/20] fixup! Explicitly specify the physical addresses of
 snapshot page tables

Remove the unused constant `SCRATCH_TOP_USED_OFFSET`.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_common/src/layout.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs
index 0b0bad9b0..a1a6d2661 100644
--- a/src/hyperlight_common/src/layout.rs
+++ b/src/hyperlight_common/src/layout.rs
@@ -22,7 +22,6 @@ pub use arch::{MAX_GPA, MAX_GVA, SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN};
 
 // offsets down from the top of scratch memory for various things
 pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08;
-pub const SCRATCH_TOP_USED_OFFSET: u64 = 0x10;
 pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x10;
 pub const SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET: u64 = 0x18;
 pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x20;

From 2f5a4fb40178d08de0cd37ac760e2725fed2c713 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Thu, 29 Jan 2026 14:46:26 +0000
Subject: [PATCH 16/20] fixup! snapshot: handle the scratch region correctly

Make virt_to_phys use the aligned and un-canonicalised version of its
argument for traversal.

This does not actually matter, because only certain bits are actually
taken from vmin by the (chained) `modify_ptes` iterator(s). However,
since it doesn't matter, this is probably better for clarity.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_common/src/arch/amd64/vmem.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/hyperlight_common/src/arch/amd64/vmem.rs b/src/hyperlight_common/src/arch/amd64/vmem.rs
index 5f75b3488..cddcabf64 100644
--- a/src/hyperlight_common/src/arch/amd64/vmem.rs
+++ b/src/hyperlight_common/src/arch/amd64/vmem.rs
@@ -439,7 +439,7 @@ pub unsafe fn virt_to_phys<'a, Op: TableOps + 'a>(
     let vmax = core::cmp::min(vmin + len, 1u64 << VA_BITS);
     modify_ptes::<47, 39, Op, _>(MapRequest {
         table_base: op.as_ref().root_table(),
-        vmin: address,
+        vmin,
         len: vmax - vmin,
         update_parent: UpdateParentNone {},
     })

From 638bcb7b498163f6d59f6bbbff137f8a80f7baeb Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Thu, 29 Jan 2026 15:07:59 +0000
Subject: [PATCH 17/20] fixup! snapshot: handle the scratch region correctly

Fix bug that resulted in the first page of scratch memory being
unconditionally copied into the snapshot.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/sandbox/snapshot.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/hyperlight_host/src/sandbox/snapshot.rs b/src/hyperlight_host/src/sandbox/snapshot.rs
index 839b7b743..00afe12d8 100644
--- a/src/hyperlight_host/src/sandbox/snapshot.rs
+++ b/src/hyperlight_host/src/sandbox/snapshot.rs
@@ -273,7 +273,7 @@ fn filtered_mappings<'a>(
     }
     .filter_map(move |(gva, gpa, bm)| {
         // the scratch map doesn't count
-        if gva > scratch_base_gva(scratch_size) {
+        if gva >= scratch_base_gva(scratch_size) {
             return None;
         }
         // neither does the mapping of the snapshot's own page tables

From 2a2b4445087492a73eec6182e7138c575ef29538 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Thu, 29 Jan 2026 15:09:45 +0000
Subject: [PATCH 18/20] fixup! snapshot: handle the scratch region correctly

Be slightly more careful about handling offsets from the guest in
snapshot.rs: previously, a maliciously-constructed sandbox could
possibly cause underflow panics, and another function could be misused
in a way that would result in a slice out-of-bounds panic (although
it was not currently being misused in such a way).

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/sandbox/snapshot.rs | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/hyperlight_host/src/sandbox/snapshot.rs b/src/hyperlight_host/src/sandbox/snapshot.rs
index 00afe12d8..7d3642612 100644
--- a/src/hyperlight_host/src/sandbox/snapshot.rs
+++ b/src/hyperlight_host/src/sandbox/snapshot.rs
@@ -177,12 +177,14 @@ fn access_gpa<'a>(
     scratch: &'a ExclusiveSharedMemory,
     scratch_size: usize,
     gpa: u64,
-) -> (&'a ExclusiveSharedMemory, usize) {
+) -> Option<(&'a ExclusiveSharedMemory, usize)> {
     let scratch_base = scratch_base_gpa(scratch_size);
     if gpa >= scratch_base {
-        (scratch, (gpa - scratch_base) as usize)
+        Some((scratch, (gpa - scratch_base) as usize))
+    } else if gpa >= SandboxMemoryLayout::BASE_ADDRESS as u64 {
+        Some((snap, gpa as usize - SandboxMemoryLayout::BASE_ADDRESS))
     } else {
-        (snap, gpa as usize - SandboxMemoryLayout::BASE_ADDRESS)
+        None
     }
 }
 
@@ -222,8 +224,8 @@ impl<'a> hyperlight_common::vmem::TableOps for SharedMemoryPageTableBuffer<'a> {
         addr + offset
     }
     unsafe fn read_entry(&self, addr: u64) -> u64 {
-        let (mem, off) = access_gpa(self.snap, self.scratch, self.scratch_size, addr);
-        let Some(pte_bytes) = mem.as_slice().get(off..off + 8) else {
+        let memoff = access_gpa(self.snap, self.scratch, self.scratch_size, addr);
+        let Some(pte_bytes) = memoff.and_then(|(mem, off)| mem.as_slice().get(off..off + 8)) else {
             // Attacker-controlled data pointed out-of-bounds. We'll
             // default to returning 0 in this case, which, for most
             // architectures (including x86-64 and arm64, the ones we
@@ -303,9 +305,12 @@ unsafe fn guest_page<'a>(
     scratch_size: usize,
     gpa: u64,
 ) -> Option<&'a [u8]> {
+    if !gpa.is_multiple_of(PAGE_SIZE as u64) {
+        return None;
+    }
     let gpa_u = gpa as usize;
     for rgn in regions {
-        if gpa_u >= rgn.guest_region.start && gpa_u < rgn.guest_region.end {
+        if gpa_u >= rgn.guest_region.start && gpa_u + PAGE_SIZE <= rgn.guest_region.end {
             let off = gpa_u - rgn.guest_region.start;
             #[allow(clippy::useless_conversion)]
             let host_region_base: usize = rgn.host_region.start.into();
@@ -314,7 +319,7 @@ unsafe fn guest_page<'a>(
             });
         }
     }
-    let (mem, off) = access_gpa(snap, scratch, scratch_size, gpa);
+    let (mem, off) = access_gpa(snap, scratch, scratch_size, gpa)?;
     if off + PAGE_SIZE <= mem.as_slice().len() {
         Some(&mem.as_slice()[off..off + PAGE_SIZE])
     } else {

From 54cd0655edb2a8d40525f50ce9a800dccac5cea1 Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Thu, 29 Jan 2026 22:24:31 +0000
Subject: [PATCH 19/20] fixup! Modify guest physical page allocator to allocate
 from the scratch region

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_host/src/hypervisor/gdb/mod.rs |  2 +-
 src/hyperlight_host/src/hypervisor/mod.rs     |  2 +-
 src/hyperlight_host/src/mem/mgr.rs            | 50 +++++++++----------
 src/hyperlight_host/src/sandbox/outb.rs       |  4 +-
 src/hyperlight_host/src/sandbox/snapshot.rs   |  2 +-
 .../src/sandbox/uninitialized_evolve.rs       |  2 +-
 6 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/src/hyperlight_host/src/hypervisor/gdb/mod.rs b/src/hyperlight_host/src/hypervisor/gdb/mod.rs
index 2b6aae2c5..82fe9503f 100644
--- a/src/hyperlight_host/src/hypervisor/gdb/mod.rs
+++ b/src/hyperlight_host/src/hypervisor/gdb/mod.rs
@@ -497,7 +497,7 @@ mod tests {
                 .inspect_err(|_| unsafe {
                     libc::munmap(mapped_mem, size);
                 })?;
-            let (mem_mgr, _) = sandbox.mgr.build();
+            let (mem_mgr, _) = sandbox.mgr.build()?;
 
             // Create the memory access struct
             let mem_access = DebugMemoryAccess {
diff --git a/src/hyperlight_host/src/hypervisor/mod.rs b/src/hyperlight_host/src/hypervisor/mod.rs
index 9838f4ddc..58410aa67 100644
--- a/src/hyperlight_host/src/hypervisor/mod.rs
+++ b/src/hyperlight_host/src/hypervisor/mod.rs
@@ -525,7 +525,7 @@ pub(crate) mod tests {
         let rt_cfg: SandboxRuntimeConfig = Default::default();
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
-        let (mut mem_mgr, gshm) = sandbox.mgr.build();
+        let (mut mem_mgr, gshm) = sandbox.mgr.build().unwrap();
         let mut vm = set_up_hypervisor_partition(
             gshm,
             &config,
diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs
index e8a9754ec..060765bc7 100644
--- a/src/hyperlight_host/src/mem/mgr.rs
+++ b/src/hyperlight_host/src/mem/mgr.rs
@@ -235,12 +235,21 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
     }
 
     /// Wraps ExclusiveSharedMemory::build
+    // Morally, this should not have to be a Result: this operation is
+    // infallible. The source of the Result is
+    // update_scratch_bookkeeping(), which calls functions that can
+    // fail due to bounds checks (which are statically known to be ok
+    // in this situation) or due to failing to take the scratch shared
+    // memory lock, but the scratch shared memory is built in this
+    // function, its lock does not escape before the end of the
+    // function, and the lock is taken by no other code path, so we
+    // know it is not contended.
     pub fn build(
         self,
-    ) -> (
+    ) -> Result<(
         SandboxMemoryManager<HostSharedMemory>,
         SandboxMemoryManager<GuestSharedMemory>,
-    ) {
+    )> {
         let (hshm, gshm) = self.shared_mem.build();
         let (hscratch, gscratch) = self.scratch_mem.build();
         let mut host_mgr = SandboxMemoryManager {
@@ -265,8 +274,8 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
         };
         host_mgr.update_scratch_bookkeeping(
             (SandboxMemoryLayout::BASE_ADDRESS + self.layout.get_pt_offset()) as u64,
-        );
-        (host_mgr, guest_mgr)
+        )?;
+        Ok((host_mgr, guest_mgr))
     }
 }
 
@@ -424,39 +433,30 @@ impl SandboxMemoryManager<HostSharedMemory> {
 
             Some(gscratch)
         };
-        self.update_scratch_bookkeeping(snapshot.root_pt_gpa());
+        self.update_scratch_bookkeeping(snapshot.root_pt_gpa())?;
         Ok((gsnapshot, gscratch))
     }
 
-    fn update_scratch_bookkeeping(&mut self, snapshot_pt_base_gpa: u64) {
+    fn update_scratch_bookkeeping(&mut self, snapshot_pt_base_gpa: u64) -> Result<()> {
         let scratch_size = self.scratch_mem.mem_size();
+
         let size_offset =
             scratch_size - hyperlight_common::layout::SCRATCH_TOP_SIZE_OFFSET as usize;
-        // The only way that write can fail is if the offset is
-        // outside of the memory, which would be sufficiently much of
-        // an invariant violation that panicking is probably
-        // sensible...
-        #[allow(clippy::unwrap_used)]
         self.scratch_mem
-            .write::<u64>(size_offset, scratch_size as u64)
-            .unwrap();
+            .write::<u64>(size_offset, scratch_size as u64)?;
+
         let alloc_offset =
             scratch_size - hyperlight_common::layout::SCRATCH_TOP_ALLOCATOR_OFFSET as usize;
-        // See above comment about unwrap() on write
-        #[allow(clippy::unwrap_used)]
-        self.scratch_mem
-            .write::<u64>(
-                alloc_offset,
-                hyperlight_common::layout::scratch_base_gpa(scratch_size),
-            )
-            .unwrap();
+        self.scratch_mem.write::<u64>(
+            alloc_offset,
+            hyperlight_common::layout::scratch_base_gpa(scratch_size),
+        )?;
+
         let snapshot_pt_base_gpa_offset = scratch_size
             - hyperlight_common::layout::SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET as usize;
-        // See above comment about unwrap() on write
-        #[allow(clippy::unwrap_used)]
         self.scratch_mem
-            .write::<u64>(snapshot_pt_base_gpa_offset, snapshot_pt_base_gpa)
-            .unwrap();
+            .write::<u64>(snapshot_pt_base_gpa_offset, snapshot_pt_base_gpa)?;
+        Ok(())
     }
 }
 
diff --git a/src/hyperlight_host/src/sandbox/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs
index cf97697e7..87c40c744 100644
--- a/src/hyperlight_host/src/sandbox/outb.rs
+++ b/src/hyperlight_host/src/sandbox/outb.rs
@@ -287,7 +287,7 @@ mod tests {
             layout
                 .write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
                 .unwrap();
-            let (hmgr, _) = mgr.build();
+            let (hmgr, _) = mgr.build().unwrap();
             hmgr
         };
         {
@@ -399,7 +399,7 @@ mod tests {
                 layout
                     .write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
                     .unwrap();
-                let (hmgr, _) = mgr.build();
+                let (hmgr, _) = mgr.build().unwrap();
                 hmgr
             };
 
diff --git a/src/hyperlight_host/src/sandbox/snapshot.rs b/src/hyperlight_host/src/sandbox/snapshot.rs
index 7d3642612..52850b3f3 100644
--- a/src/hyperlight_host/src/sandbox/snapshot.rs
+++ b/src/hyperlight_host/src/sandbox/snapshot.rs
@@ -613,7 +613,7 @@ mod tests {
             None,
             [0u8; 16],
         );
-        let (mgr, _) = mgr.build();
+        let (mgr, _) = mgr.build().unwrap();
         (mgr, pt_base as u64)
     }
 
diff --git a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs
index aeaff030a..17395daa7 100644
--- a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs
+++ b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs
@@ -39,7 +39,7 @@ use crate::{MultiUseSandbox, Result, UninitializedSandbox, new_error};
 
 #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
 pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result<MultiUseSandbox> {
-    let (mut hshm, gshm) = u_sbox.mgr.build();
+    let (mut hshm, gshm) = u_sbox.mgr.build()?;
     let mut vm = set_up_hypervisor_partition(
         gshm,
         &u_sbox.config,

From a2ece56b39aa7be1a4e574d9c5c67df54e8d726b Mon Sep 17 00:00:00 2001
From: Lucy Menon <168595099+syntactically@users.noreply.github.com>
Date: Fri, 30 Jan 2026 01:58:04 +0000
Subject: [PATCH 20/20] Neaten up page table operations structures

This splits `hyperlight_common::vmem::TableOps` into `TableReadOps`
and `TableOps`, so that implementations like the one in snapshot.rs
that cannot support writes don't have to provide the (extremely
partial!) write functions.  This is slightly involved, because we need
to keep to track of the right information in the type system about
whether writes are possible, and, if so, whether they can move
existing page tables.

The implementation of of writing to cr3 to update page tables is also
moved into `hyperlight_guest_bin::paging`, where it belongs, which is
a simple change.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>
---
 src/hyperlight_common/src/arch/amd64/vmem.rs | 197 ++++++++++++++-----
 src/hyperlight_common/src/arch/i686/vmem.rs  |   9 +-
 src/hyperlight_common/src/vmem.rs            | 138 ++++++++++---
 src/hyperlight_guest_bin/src/paging.rs       |  47 +++--
 src/hyperlight_host/src/mem/mgr.rs           |  53 ++---
 src/hyperlight_host/src/sandbox/snapshot.rs  |  40 +---
 6 files changed, 320 insertions(+), 164 deletions(-)

diff --git a/src/hyperlight_common/src/arch/amd64/vmem.rs b/src/hyperlight_common/src/arch/amd64/vmem.rs
index cddcabf64..110598a9a 100644
--- a/src/hyperlight_common/src/arch/amd64/vmem.rs
+++ b/src/hyperlight_common/src/arch/amd64/vmem.rs
@@ -25,7 +25,9 @@ limitations under the License.
 //! The code uses an iterator-based approach to walk the page table hierarchy,
 //! allocating intermediate tables as needed and setting appropriate flags on leaf PTEs
 
-use crate::vmem::{BasicMapping, Mapping, MappingKind, TableOps};
+use crate::vmem::{
+    BasicMapping, Mapping, MappingKind, TableMovabilityBase, TableOps, TableReadOps, Void,
+};
 
 // Paging Flags
 //
@@ -78,7 +80,7 @@ const fn page_nx_flag(executable: bool) -> u64 {
 /// # Safety
 /// The caller must ensure that `entry_ptr` points to a valid page table entry.
 #[inline(always)]
-unsafe fn read_pte_if_present<Op: TableOps>(op: &Op, entry_ptr: Op::TableAddr) -> Option<u64> {
+unsafe fn read_pte_if_present<Op: TableReadOps>(op: &Op, entry_ptr: Op::TableAddr) -> Option<u64> {
     let pte = unsafe { op.read_entry(entry_ptr) };
     if (pte & PAGE_PRESENT) != 0 {
         Some(pte)
@@ -107,9 +109,37 @@ fn pte_for_table<Op: TableOps>(table_addr: Op::TableAddr) -> u64 {
         PAGE_PRESENT // P   - this entry is present
 }
 
+/// This trait is used to select appropriate implementations of
+/// [`UpdateParent`] to be used, depending on whether a particular
+/// implementation needs the ability to move tables.
+pub trait TableMovability<Op: TableReadOps + ?Sized, TableMoveInfo> {
+    type RootUpdateParent: UpdateParent<Op, TableMoveInfo = TableMoveInfo>;
+    fn root_update_parent() -> Self::RootUpdateParent;
+}
+impl<Op: TableOps<TableMovability = crate::vmem::MayMoveTable>> TableMovability<Op, Op::TableAddr>
+    for crate::vmem::MayMoveTable
+{
+    type RootUpdateParent = UpdateParentRoot;
+    fn root_update_parent() -> Self::RootUpdateParent {
+        UpdateParentRoot {}
+    }
+}
+impl<Op: TableReadOps> TableMovability<Op, Void> for crate::vmem::MayNotMoveTable {
+    type RootUpdateParent = UpdateParentNone;
+    fn root_update_parent() -> Self::RootUpdateParent {
+        UpdateParentNone {}
+    }
+}
+
 /// Helper function to write a page table entry, updating the whole
 /// chain of tables back to the root if necessary
-unsafe fn write_entry_updating<Op: TableOps, P: UpdateParent<Op>>(
+unsafe fn write_entry_updating<
+    Op: TableOps,
+    P: UpdateParent<
+            Op,
+            TableMoveInfo = <Op::TableMovability as TableMovabilityBase<Op>>::TableMoveInfo,
+        >,
+>(
     op: &Op,
     parent: P,
     addr: Op::TableAddr,
@@ -128,14 +158,26 @@ unsafe fn write_entry_updating<Op: TableOps, P: UpdateParent<Op>>(
 /// This is done via a trait so that the selected impl knows the exact
 /// nesting depth of tables, in order to assist
 /// inlining/specialisation in generating efficient code.
-trait UpdateParent<Op: TableOps>: Copy {
-    fn update_parent(self, op: &Op, new_ptr: Op::TableAddr);
+///
+/// The trait definition only bounds its parameter by
+/// [`TableReadOps`], since [`UpdateParentNone`] does not need to be
+/// able to actually write to the tables.
+pub trait UpdateParent<Op: TableReadOps + ?Sized>: Copy {
+    /// The type of the information about a moved table which is
+    /// needed in order to update its parent.
+    type TableMoveInfo;
+    /// The [`UpdateParent`] type that should be used when going down
+    /// another level in the table, in order to add the current level
+    /// to the chain of ancestors to be updated.
+    type ChildType: UpdateParent<Op, TableMoveInfo = Self::TableMoveInfo>;
+    fn update_parent(self, op: &Op, new_ptr: Self::TableMoveInfo);
+    fn for_child_at_entry(self, entry_ptr: Op::TableAddr) -> Self::ChildType;
 }
 
 /// A struct implementing [`UpdateParent`] that keeps track of the
 /// fact that the parent table is itself another table, whose own
 /// ancestors may need to be recursively updated
-struct UpdateParentTable<Op: TableOps, P: UpdateParent<Op>> {
+pub struct UpdateParentTable<Op: TableOps, P: UpdateParent<Op>> {
     parent: P,
     entry_ptr: Op::TableAddr,
 }
@@ -150,45 +192,64 @@ impl<Op: TableOps, P: UpdateParent<Op>> UpdateParentTable<Op, P> {
         UpdateParentTable { parent, entry_ptr }
     }
 }
-impl<Op: TableOps, P: UpdateParent<Op>> UpdateParent<Op> for UpdateParentTable<Op, P> {
+impl<
+    Op: TableOps<TableMovability = crate::vmem::MayMoveTable>,
+    P: UpdateParent<Op, TableMoveInfo = Op::TableAddr>,
+> UpdateParent<Op> for UpdateParentTable<Op, P>
+{
+    type TableMoveInfo = Op::TableAddr;
+    type ChildType = UpdateParentTable<Op, Self>;
     fn update_parent(self, op: &Op, new_ptr: Op::TableAddr) {
         let pte = pte_for_table::<Op>(new_ptr);
         unsafe {
             write_entry_updating(op, self.parent, self.entry_ptr, pte);
         }
     }
+    fn for_child_at_entry(self, entry_ptr: Op::TableAddr) -> Self::ChildType {
+        Self::ChildType::new(self, entry_ptr)
+    }
 }
 
 /// A struct implementing [`UpdateParent`] that keeps track of the
-/// fact that the parent "table" is actually the CR3 system register.
+/// fact that the parent "table" is actually the root (e.g. the value
+/// of CR3 in the guest)
 #[derive(Copy, Clone)]
-struct UpdateParentCR3 {}
-impl<Op: TableOps> UpdateParent<Op> for UpdateParentCR3 {
-    fn update_parent(self, _op: &Op, new_ptr: Op::TableAddr) {
+pub struct UpdateParentRoot {}
+impl<Op: TableOps<TableMovability = crate::vmem::MayMoveTable>> UpdateParent<Op>
+    for UpdateParentRoot
+{
+    type TableMoveInfo = Op::TableAddr;
+    type ChildType = UpdateParentTable<Op, Self>;
+    fn update_parent(self, op: &Op, new_ptr: Op::TableAddr) {
         unsafe {
-            core::arch::asm!("mov cr3, {}", in(reg) Op::to_phys(new_ptr));
+            op.update_root(new_ptr);
         }
     }
+    fn for_child_at_entry(self, entry_ptr: Op::TableAddr) -> Self::ChildType {
+        Self::ChildType::new(self, entry_ptr)
+    }
 }
 
-/// A struct implementing [`UpdateParent`] that panics when used, for
-/// the occasional situations in which we should never do an update
+/// A struct implementing [`UpdateParent`] that is impossible to use
+/// (since its [`update_parent`] method takes `Void`), used when it is
+/// statically known that a table operation cannot result in a need to
+/// update ancestors.
 #[derive(Copy, Clone)]
-struct UpdateParentNone {}
-impl<Op: TableOps> UpdateParent<Op> for UpdateParentNone {
-    // This is only used in contexts which should absolutely never try
-    // to update a page table, and so in no case should possibly ever
-    // call an update_parent.  Therefore, this is impossible unless an
-    // extremely significant invariant violation has occurred.
-    #[allow(clippy::panic)]
-    fn update_parent(self, _op: &Op, _new_ptr: Op::TableAddr) {
-        panic!("UpdateParentNone: tried to update parent");
+pub struct UpdateParentNone {}
+impl<Op: TableReadOps> UpdateParent<Op> for UpdateParentNone {
+    type TableMoveInfo = Void;
+    type ChildType = Self;
+    fn update_parent(self, _op: &Op, impossible: Void) {
+        match impossible {}
+    }
+    fn for_child_at_entry(self, _entry_ptr: Op::TableAddr) -> Self {
+        self
     }
 }
 
 /// A helper structure indicating a mapping operation that needs to be
 /// performed
-struct MapRequest<Op: TableOps, P: UpdateParent<Op>> {
+struct MapRequest<Op: TableReadOps, P: UpdateParent<Op>> {
     table_base: Op::TableAddr,
     vmin: VirtAddr,
     len: u64,
@@ -197,7 +258,7 @@ struct MapRequest<Op: TableOps, P: UpdateParent<Op>> {
 
 /// A helper structure indicating that a particular PTE needs to be
 /// modified
-struct MapResponse<Op: TableOps, P: UpdateParent<Op>> {
+struct MapResponse<Op: TableReadOps, P: UpdateParent<Op>> {
     entry_ptr: Op::TableAddr,
     vmin: VirtAddr,
     len: u64,
@@ -216,11 +277,16 @@ struct MapResponse<Op: TableOps, P: UpdateParent<Op>> {
 /// - PDPT: HIGH_BIT=38, LOW_BIT=30 (9 bits = 512 entries, each covering 1GB)
 /// - PD:   HIGH_BIT=29, LOW_BIT=21 (9 bits = 512 entries, each covering 2MB)
 /// - PT:   HIGH_BIT=20, LOW_BIT=12 (9 bits = 512 entries, each covering 4KB)
-struct ModifyPteIterator<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps, P: UpdateParent<Op>> {
+struct ModifyPteIterator<
+    const HIGH_BIT: u8,
+    const LOW_BIT: u8,
+    Op: TableReadOps,
+    P: UpdateParent<Op>,
+> {
     request: MapRequest<Op, P>,
     n: u64,
 }
-impl<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps, P: UpdateParent<Op>> Iterator
+impl<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableReadOps, P: UpdateParent<Op>> Iterator
     for ModifyPteIterator<HIGH_BIT, LOW_BIT, Op, P>
 {
     type Item = MapResponse<Op, P>;
@@ -286,7 +352,7 @@ impl<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps, P: UpdateParent<Op>> I
         })
     }
 }
-fn modify_ptes<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps, P: UpdateParent<Op>>(
+fn modify_ptes<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableReadOps, P: UpdateParent<Op>>(
     r: MapRequest<Op, P>,
 ) -> ModifyPteIterator<HIGH_BIT, LOW_BIT, Op, P> {
     ModifyPteIterator { request: r, n: 0 }
@@ -296,11 +362,20 @@ fn modify_ptes<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps, P: UpdatePar
 /// # Safety
 /// This function modifies page table data structures, and should not be called concurrently
 /// with any other operations that modify the page tables.
-unsafe fn alloc_pte_if_needed<Op: TableOps, P: UpdateParent<Op>>(
+unsafe fn alloc_pte_if_needed<
+    Op: TableOps,
+    P: UpdateParent<
+            Op,
+            TableMoveInfo = <Op::TableMovability as TableMovabilityBase<Op>>::TableMoveInfo,
+        >,
+>(
     op: &Op,
     x: MapResponse<Op, P>,
-) -> MapRequest<Op, UpdateParentTable<Op, P>> {
-    let new_update_parent = UpdateParentTable::new(x.update_parent, x.entry_ptr);
+) -> MapRequest<Op, P::ChildType>
+where
+    P::ChildType: UpdateParent<Op>,
+{
+    let new_update_parent = x.update_parent.for_child_at_entry(x.entry_ptr);
     if let Some(pte) = unsafe { read_pte_if_present(op, x.entry_ptr) } {
         return MapRequest {
             table_base: Op::from_phys(pte & PTE_ADDR_MASK),
@@ -330,7 +405,13 @@ unsafe fn alloc_pte_if_needed<Op: TableOps, P: UpdateParent<Op>>(
 /// with any other operations that modify the page tables.
 #[allow(clippy::identity_op)]
 #[allow(clippy::precedence)]
-unsafe fn map_page<Op: TableOps, P: UpdateParent<Op>>(
+unsafe fn map_page<
+    Op: TableOps,
+    P: UpdateParent<
+            Op,
+            TableMoveInfo = <Op::TableMovability as TableMovabilityBase<Op>>::TableMoveInfo,
+        >,
+>(
     op: &Op,
     mapping: &Mapping,
     r: MapResponse<Op, P>,
@@ -383,7 +464,7 @@ pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
         table_base: op.root_table(),
         vmin: mapping.virt_base,
         len: mapping.len,
-        update_parent: UpdateParentCR3 {},
+        update_parent: Op::TableMovability::root_update_parent(),
     })
     .map(|r| unsafe { alloc_pte_if_needed(op, r) })
     .flat_map(modify_ptes::<38, 30, Op, _>)
@@ -399,15 +480,18 @@ pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
 /// This function traverses page table data structures, and should not
 /// be called concurrently with any other operations that modify the
 /// page table.
-unsafe fn require_pte_exist<Op: TableOps, P: UpdateParent<Op>>(
+unsafe fn require_pte_exist<Op: TableReadOps, P: UpdateParent<Op>>(
     op: &Op,
     x: MapResponse<Op, P>,
-) -> Option<MapRequest<Op, UpdateParentTable<Op, P>>> {
+) -> Option<MapRequest<Op, P::ChildType>>
+where
+    P::ChildType: UpdateParent<Op>,
+{
     unsafe { read_pte_if_present(op, x.entry_ptr) }.map(|pte| MapRequest {
         table_base: Op::from_phys(pte & PTE_ADDR_MASK),
         vmin: x.vmin,
         len: x.len,
-        update_parent: UpdateParentTable::new(x.update_parent, x.entry_ptr),
+        update_parent: x.update_parent.for_child_at_entry(x.entry_ptr),
     })
 }
 
@@ -429,7 +513,7 @@ unsafe fn require_pte_exist<Op: TableOps, P: UpdateParent<Op>>(
 /// operations structure owns a large buffer, a reference can instead
 /// be passed.
 #[allow(clippy::missing_safety_doc)]
-pub unsafe fn virt_to_phys<'a, Op: TableOps + 'a>(
+pub unsafe fn virt_to_phys<'a, Op: TableReadOps + 'a>(
     op: impl core::convert::AsRef<Op> + Copy + 'a,
     address: u64,
     len: u64,
@@ -480,7 +564,10 @@ mod tests {
     use core::cell::RefCell;
 
     use super::*;
-    use crate::vmem::{BasicMapping, Mapping, MappingKind, PAGE_TABLE_ENTRIES_PER_TABLE, TableOps};
+    use crate::vmem::{
+        BasicMapping, Mapping, MappingKind, MayNotMoveTable, PAGE_TABLE_ENTRIES_PER_TABLE,
+        TableOps, TableReadOps, Void,
+    };
 
     /// A mock TableOps implementation for testing that stores page tables in memory
     /// needed because the `GuestPageTableBuffer` is in hyperlight_host which would cause a circular dependency
@@ -512,16 +599,9 @@ mod tests {
         }
     }
 
-    impl TableOps for MockTableOps {
+    impl TableReadOps for MockTableOps {
         type TableAddr = (usize, usize); // (table_index, entry_index)
 
-        unsafe fn alloc_table(&self) -> Self::TableAddr {
-            let mut tables = self.tables.borrow_mut();
-            let idx = tables.len();
-            tables.push([0u64; PAGE_TABLE_ENTRIES_PER_TABLE]);
-            (idx, 0)
-        }
-
         fn entry_addr(addr: Self::TableAddr, entry_offset: u64) -> Self::TableAddr {
             // Convert to physical address, add offset, convert back
             let phys = Self::to_phys(addr) + entry_offset;
@@ -532,11 +612,6 @@ mod tests {
             self.tables.borrow()[addr.0][addr.1]
         }
 
-        unsafe fn write_entry(&self, addr: Self::TableAddr, entry: u64) -> Option<Self::TableAddr> {
-            self.tables.borrow_mut()[addr.0][addr.1] = entry;
-            None
-        }
-
         fn to_phys(addr: Self::TableAddr) -> PhysAddr {
             // Each table is 4KB, entries are 8 bytes
             (addr.0 as u64 * PAGE_TABLE_SIZE as u64) + (addr.1 as u64 * 8)
@@ -553,6 +628,26 @@ mod tests {
         }
     }
 
+    impl TableOps for MockTableOps {
+        type TableMovability = MayNotMoveTable;
+
+        unsafe fn alloc_table(&self) -> Self::TableAddr {
+            let mut tables = self.tables.borrow_mut();
+            let idx = tables.len();
+            tables.push([0u64; PAGE_TABLE_ENTRIES_PER_TABLE]);
+            (idx, 0)
+        }
+
+        unsafe fn write_entry(&self, addr: Self::TableAddr, entry: u64) -> Option<Void> {
+            self.tables.borrow_mut()[addr.0][addr.1] = entry;
+            None
+        }
+
+        unsafe fn update_root(&self, impossible: Void) {
+            match impossible {}
+        }
+    }
+
     // ==================== bits() function tests ====================
 
     #[test]
diff --git a/src/hyperlight_common/src/arch/i686/vmem.rs b/src/hyperlight_common/src/arch/i686/vmem.rs
index 4b1d16b87..83c749975 100644
--- a/src/hyperlight_common/src/arch/i686/vmem.rs
+++ b/src/hyperlight_common/src/arch/i686/vmem.rs
@@ -17,7 +17,7 @@ limitations under the License.
 // This file is just dummy definitions at the moment, in order to
 // allow compiling the guest for real mode boot scenarios.
 
-use crate::vmem::{BasicMapping, Mapping, TableOps};
+use crate::vmem::{BasicMapping, Mapping, TableOps, TableReadOps, Void};
 
 pub const PAGE_SIZE: usize = 4096;
 pub const PAGE_TABLE_SIZE: usize = 4096;
@@ -43,3 +43,10 @@ pub unsafe fn virt_to_phys<Op: TableOps>(
     #[allow(unreachable_code)]
     core::iter::empty()
 }
+
+pub trait TableMovability<Op: TableReadOps + ?Sized, TableMoveInfo> {}
+impl<Op: TableOps<TableMovability = crate::vmem::MayMoveTable>> TableMovability<Op, Op::TableAddr>
+    for crate::vmem::MayMoveTable
+{
+}
+impl<Op: TableReadOps> TableMovability<Op, Void> for crate::vmem::MayNotMoveTable {}
diff --git a/src/hyperlight_common/src/vmem.rs b/src/hyperlight_common/src/vmem.rs
index 1a83146f0..6ddf172dd 100644
--- a/src/hyperlight_common/src/vmem.rs
+++ b/src/hyperlight_common/src/vmem.rs
@@ -25,30 +25,15 @@ pub use arch::{PAGE_TABLE_SIZE, PageTableEntry, PhysAddr, VirtAddr};
 pub const PAGE_TABLE_ENTRIES_PER_TABLE: usize =
     PAGE_TABLE_SIZE / core::mem::size_of::<PageTableEntry>();
 
-/// The operations used to actually access the page table structures,
-/// used to allow the same code to be used in the host and the guest
-/// for page table setup
-pub trait TableOps {
+/// The read-only operations used to actually access the page table
+/// structures, used to allow the same code to be used in the host and
+/// the guest for page table setup.  This is distinct from
+/// `TableWriteOps`, since there are some implementations for which
+/// writing does not make sense, and only reading is required.
+pub trait TableReadOps {
     /// The type of table addresses
     type TableAddr: Copy;
 
-    /// Allocate a zeroed table
-    ///
-    /// # Safety
-    /// The current implementations of this function are not
-    /// inherently unsafe, but the guest implementation will likely
-    /// become so in the future when a real physical page allocator is
-    /// implemented.
-    ///
-    /// Currently, callers should take care not to call this on
-    /// multiple threads at the same time.
-    ///
-    /// # Panics
-    /// This function may panic if:
-    /// - The Layout creation fails
-    /// - Memory allocation fails
-    unsafe fn alloc_table(&self) -> Self::TableAddr;
-
     /// Offset the table address by the given offset in bytes.
     ///
     /// # Parameters
@@ -74,6 +59,94 @@ pub trait TableOps {
     /// time as mapping code using the trait.
     unsafe fn read_entry(&self, addr: Self::TableAddr) -> PageTableEntry;
 
+    /// Convert an abstract table address to a concrete physical address (u64)
+    /// which can be e.g. written into a page table entry
+    fn to_phys(addr: Self::TableAddr) -> PhysAddr;
+
+    /// Convert a concrete physical address (u64) which may have been e.g. read
+    /// from a page table entry back into an abstract table address
+    fn from_phys(addr: PhysAddr) -> Self::TableAddr;
+
+    /// Return the address of the root page table
+    fn root_table(&self) -> Self::TableAddr;
+}
+
+/// Our own version of ! until it is stable. Used to avoid needing to
+/// implement [`TableOps::update_root`] for ops that never need
+/// to move a table.
+pub enum Void {}
+
+/// A marker struct, used by an implementation of [`TableOps`] to
+/// indicate that it may need to move existing page tables
+pub struct MayMoveTable {}
+/// A marker struct, used by an implementation of [`TableOps`] to
+/// indicate that it will be able to update existing page tables
+/// in-place, without moving them.
+pub struct MayNotMoveTable {}
+
+mod sealed {
+    use super::{MayMoveTable, MayNotMoveTable, TableReadOps, Void};
+
+    /// A (purposefully-not-exposed) internal implementation detail of the
+    /// logic around whether a [`TableOps`] implementation may or may not
+    /// move page tables.
+    pub trait TableMovabilityBase<Op: TableReadOps + ?Sized> {
+        type TableMoveInfo;
+    }
+    impl<Op: TableReadOps> TableMovabilityBase<Op> for MayMoveTable {
+        type TableMoveInfo = Op::TableAddr;
+    }
+    impl<Op: TableReadOps> TableMovabilityBase<Op> for MayNotMoveTable {
+        type TableMoveInfo = Void;
+    }
+}
+use sealed::*;
+
+/// A sealed trait used to collect some information about the marker structures [`MayMoveTable`] and [`MayNotMoveTable`]
+pub trait TableMovability<Op: TableReadOps + ?Sized>:
+    TableMovabilityBase<Op>
+    + arch::TableMovability<Op, <Self as TableMovabilityBase<Op>>::TableMoveInfo>
+{
+}
+impl<
+    Op: TableReadOps,
+    T: TableMovabilityBase<Op>
+        + arch::TableMovability<Op, <Self as TableMovabilityBase<Op>>::TableMoveInfo>,
+> TableMovability<Op> for T
+{
+}
+
+/// The operations used to actually access the page table structures
+/// that involve writing to them, used to allow the same code to be
+/// used in the host and the guest for page table setup.
+pub trait TableOps: TableReadOps {
+    /// This marker should be either [`MayMoveTable`] or
+    /// [`MayNotMoveTable`], as the case may be.
+    ///
+    /// If this is [`MayMoveTable`], the return type of
+    /// [`Self::write_entry`] and the parameter type of
+    /// [`Self::update_root`] will be `<Self as
+    /// TableReadOps>::TableAddr`. If it is [`MayNotMoveTable`], those
+    /// types will be [`Void`].
+    type TableMovability: TableMovability<Self>;
+
+    /// Allocate a zeroed table
+    ///
+    /// # Safety
+    /// The current implementations of this function are not
+    /// inherently unsafe, but the guest implementation will likely
+    /// become so in the future when a real physical page allocator is
+    /// implemented.
+    ///
+    /// Currently, callers should take care not to call this on
+    /// multiple threads at the same time.
+    ///
+    /// # Panics
+    /// This function may panic if:
+    /// - The Layout creation fails
+    /// - Memory allocation fails
+    unsafe fn alloc_table(&self) -> Self::TableAddr;
+
     /// Write a u64 to the given address, used to write updated page
     /// table entries. In some cases,the page table in which the entry
     /// is located may need to be relocated in order for this to
@@ -92,18 +165,19 @@ pub trait TableOps {
         &self,
         addr: Self::TableAddr,
         entry: PageTableEntry,
-    ) -> Option<Self::TableAddr>;
-
-    /// Convert an abstract table address to a concrete physical address (u64)
-    /// which can be e.g. written into a page table entry
-    fn to_phys(addr: Self::TableAddr) -> PhysAddr;
-
-    /// Convert a concrete physical address (u64) which may have been e.g. read
-    /// from a page table entry back into an abstract table address
-    fn from_phys(addr: PhysAddr) -> Self::TableAddr;
+    ) -> Option<<Self::TableMovability as TableMovabilityBase<Self>>::TableMoveInfo>;
 
-    /// Return the address of the root page table
-    fn root_table(&self) -> Self::TableAddr;
+    /// Change the root page table to one at a different address
+    ///
+    /// # Safety
+    /// This function will directly result in a change to virtual
+    /// memory translation, and so is inherently unsafe w.r.t. the
+    /// Rust memory model.  All the caveats listed on [`map`] apply as
+    /// well.
+    unsafe fn update_root(
+        &self,
+        new_root: <Self::TableMovability as TableMovabilityBase<Self>>::TableMoveInfo,
+    );
 }
 
 #[derive(Debug)]
diff --git a/src/hyperlight_guest_bin/src/paging.rs b/src/hyperlight_guest_bin/src/paging.rs
index 85e4fa31b..2c5d62b23 100644
--- a/src/hyperlight_guest_bin/src/paging.rs
+++ b/src/hyperlight_guest_bin/src/paging.rs
@@ -62,16 +62,8 @@ impl core::convert::AsRef<GuestMappingOperations> for GuestMappingOperations {
         self
     }
 }
-impl vmem::TableOps for GuestMappingOperations {
+impl vmem::TableReadOps for GuestMappingOperations {
     type TableAddr = u64;
-    unsafe fn alloc_table(&self) -> u64 {
-        let page_addr = unsafe { alloc_phys_pages(1) };
-        unsafe {
-            self.phys_to_virt(page_addr)
-                .write_bytes(0u8, vmem::PAGE_TABLE_SIZE)
-        };
-        page_addr
-    }
     fn entry_addr(addr: u64, offset: u64) -> u64 {
         addr + offset
     }
@@ -83,6 +75,31 @@ impl vmem::TableOps for GuestMappingOperations {
         }
         ret
     }
+    fn to_phys(addr: u64) -> u64 {
+        addr
+    }
+    fn from_phys(addr: u64) -> u64 {
+        addr
+    }
+    fn root_table(&self) -> u64 {
+        let pml4_base: u64;
+        unsafe {
+            asm!("mov {}, cr3", out(reg) pml4_base);
+        }
+        pml4_base & !0xfff
+    }
+}
+
+impl vmem::TableOps for GuestMappingOperations {
+    type TableMovability = vmem::MayMoveTable;
+    unsafe fn alloc_table(&self) -> u64 {
+        let page_addr = unsafe { alloc_phys_pages(1) };
+        unsafe {
+            self.phys_to_virt(page_addr)
+                .write_bytes(0u8, vmem::PAGE_TABLE_SIZE)
+        };
+        page_addr
+    }
     unsafe fn write_entry(&self, addr: u64, entry: u64) -> Option<u64> {
         let mut addr = addr;
         let mut ret = None;
@@ -105,18 +122,10 @@ impl vmem::TableOps for GuestMappingOperations {
         }
         ret
     }
-    fn to_phys(addr: u64) -> u64 {
-        addr
-    }
-    fn from_phys(addr: u64) -> u64 {
-        addr
-    }
-    fn root_table(&self) -> u64 {
-        let pml4_base: u64;
+    unsafe fn update_root(&self, new_root: u64) {
         unsafe {
-            asm!("mov {}, cr3", out(reg) pml4_base);
+            core::arch::asm!("mov cr3, {}", in(reg) <Self as vmem::TableReadOps>::to_phys(new_root));
         }
-        pml4_base & !0xfff
     }
 }
 
diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs
index 060765bc7..6d8db7d8c 100644
--- a/src/hyperlight_host/src/mem/mgr.rs
+++ b/src/hyperlight_host/src/mem/mgr.rs
@@ -63,17 +63,9 @@ pub(crate) struct GuestPageTableBuffer {
     phys_base: usize,
 }
 
-impl vmem::TableOps for GuestPageTableBuffer {
+impl vmem::TableReadOps for GuestPageTableBuffer {
     type TableAddr = (usize, usize); // (table_index, entry_index)
 
-    unsafe fn alloc_table(&self) -> (usize, usize) {
-        let mut b = self.buffer.borrow_mut();
-        let table_index = b.len() / PAGE_TABLE_SIZE;
-        let new_len = b.len() + PAGE_TABLE_SIZE;
-        b.resize(new_len, 0);
-        (self.phys_base / PAGE_TABLE_SIZE + table_index, 0)
-    }
-
     fn entry_addr(addr: (usize, usize), offset: u64) -> (usize, usize) {
         // Convert to physical address, add offset, convert back
         let phys = Self::to_phys(addr) + offset;
@@ -90,20 +82,6 @@ impl vmem::TableOps for GuestPageTableBuffer {
             .unwrap_or(0)
     }
 
-    unsafe fn write_entry(
-        &self,
-        addr: (usize, usize),
-        entry: PageTableEntry,
-    ) -> Option<(usize, usize)> {
-        let mut b = self.buffer.borrow_mut();
-        let byte_offset =
-            (addr.0 - self.phys_base / PAGE_TABLE_SIZE) * PAGE_TABLE_SIZE + addr.1 * 8;
-        if let Some(slice) = b.get_mut(byte_offset..byte_offset + 8) {
-            slice.copy_from_slice(&entry.to_ne_bytes());
-        }
-        None
-    }
-
     fn to_phys(addr: (usize, usize)) -> PhysAddr {
         (addr.0 as u64 * PAGE_TABLE_SIZE as u64) + (addr.1 as u64 * 8)
     }
@@ -119,6 +97,35 @@ impl vmem::TableOps for GuestPageTableBuffer {
         (self.phys_base / PAGE_TABLE_SIZE, 0)
     }
 }
+impl vmem::TableOps for GuestPageTableBuffer {
+    type TableMovability = vmem::MayNotMoveTable;
+
+    unsafe fn alloc_table(&self) -> (usize, usize) {
+        let mut b = self.buffer.borrow_mut();
+        let table_index = b.len() / PAGE_TABLE_SIZE;
+        let new_len = b.len() + PAGE_TABLE_SIZE;
+        b.resize(new_len, 0);
+        (self.phys_base / PAGE_TABLE_SIZE + table_index, 0)
+    }
+
+    unsafe fn write_entry(
+        &self,
+        addr: (usize, usize),
+        entry: PageTableEntry,
+    ) -> Option<vmem::Void> {
+        let mut b = self.buffer.borrow_mut();
+        let byte_offset =
+            (addr.0 - self.phys_base / PAGE_TABLE_SIZE) * PAGE_TABLE_SIZE + addr.1 * 8;
+        if let Some(slice) = b.get_mut(byte_offset..byte_offset + 8) {
+            slice.copy_from_slice(&entry.to_ne_bytes());
+        }
+        None
+    }
+
+    unsafe fn update_root(&self, impossible: vmem::Void) {
+        match impossible {}
+    }
+}
 
 impl GuestPageTableBuffer {
     pub(crate) fn new(phys_base: usize) -> Self {
diff --git a/src/hyperlight_host/src/sandbox/snapshot.rs b/src/hyperlight_host/src/sandbox/snapshot.rs
index 52850b3f3..185e4274b 100644
--- a/src/hyperlight_host/src/sandbox/snapshot.rs
+++ b/src/hyperlight_host/src/sandbox/snapshot.rs
@@ -88,17 +88,8 @@ impl core::convert::AsRef<Snapshot> for Snapshot {
         self
     }
 }
-impl hyperlight_common::vmem::TableOps for Snapshot {
+impl hyperlight_common::vmem::TableReadOps for Snapshot {
     type TableAddr = u64;
-    // This entire impl is only used in contexts which should
-    // absolutely never try to update a page table, and so in no case
-    // should possibly ever call an update_parent.  Therefore, this is
-    // impossible unless an extremely significant invariant violation
-    // has occurred.
-    #[allow(clippy::panic)]
-    unsafe fn alloc_table(&self) -> u64 {
-        panic!("alloc_table: cannot mutate the page tables inside a snapshot");
-    }
     fn entry_addr(addr: u64, offset: u64) -> u64 {
         addr + offset
     }
@@ -117,15 +108,6 @@ impl hyperlight_common::vmem::TableOps for Snapshot {
         let n: [u8; 8] = pte_bytes.try_into().unwrap();
         u64::from_ne_bytes(n)
     }
-    // This entire impl is only used in contexts which should
-    // absolutely never try to update a page table, and so in no case
-    // should possibly ever call an update_parent.  Therefore, this is
-    // impossible unless an extremely significant invariant violation
-    // has occurred.
-    #[allow(clippy::panic)]
-    unsafe fn write_entry(&self, _addr: u64, _entry: u64) -> Option<u64> {
-        panic!("write_entry: cannot mutate the page tables inside a snapshot");
-    }
     fn to_phys(addr: u64) -> u64 {
         addr
     }
@@ -209,17 +191,8 @@ impl<'a> SharedMemoryPageTableBuffer<'a> {
         }
     }
 }
-impl<'a> hyperlight_common::vmem::TableOps for SharedMemoryPageTableBuffer<'a> {
+impl<'a> hyperlight_common::vmem::TableReadOps for SharedMemoryPageTableBuffer<'a> {
     type TableAddr = u64;
-    // This entire impl is only used in contexts which should
-    // absolutely never try to update a page table, and so in no case
-    // should possibly ever call an update_parent.  Therefore, this is
-    // impossible unless an extremely significant invariant violation
-    // has occurred.
-    #[allow(clippy::panic)]
-    unsafe fn alloc_table(&self) -> u64 {
-        panic!("alloc_table: cannot mutate the page tables inside the guest");
-    }
     fn entry_addr(addr: u64, offset: u64) -> u64 {
         addr + offset
     }
@@ -238,15 +211,6 @@ impl<'a> hyperlight_common::vmem::TableOps for SharedMemoryPageTableBuffer<'a> {
         let n: [u8; 8] = pte_bytes.try_into().unwrap();
         u64::from_ne_bytes(n)
     }
-    // This entire impl is only used in contexts which should
-    // absolutely never try to update a page table, and so in no case
-    // should possibly ever call an update_parent.  Therefore, this is
-    // impossible unless an extremely significant invariant violation
-    // has occurred.
-    #[allow(clippy::panic)]
-    unsafe fn write_entry(&self, _addr: u64, _entry: u64) -> Option<u64> {
-        panic!("write_entry: cannot mutate the page tables inside the guest");
-    }
     fn to_phys(addr: u64) -> u64 {
         addr
     }