Add new SnapshotManager and Snapshot structs.

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/src/mem/layout.rs

@@ -246,7 +246,7 @@ impl SandboxMemoryLayout {
     pub(crate) const BASE_ADDRESS: usize = 0x0;
 
     // the offset into a sandbox's input/output buffer where the stack starts
-    const STACK_POINTER_SIZE_BYTES: u64 = 8;
+    pub(crate) const STACK_POINTER_SIZE_BYTES: u64 = 8;
 
     /// Create a new `SandboxMemoryLayout` with the given
     /// `SandboxConfiguration`, code size and stack/heap size.
@@ -397,7 +397,7 @@ impl SandboxMemoryLayout {
 
     /// Get the offset in guest memory to the output data pointer.
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    fn get_output_data_pointer_offset(&self) -> usize {
+    pub(super) fn get_output_data_pointer_offset(&self) -> usize {
         // This field is immediately after the output data size field,
         // which is a `u64`.
         self.get_output_data_size_offset() + size_of::<u64>()
@@ -429,7 +429,7 @@ impl SandboxMemoryLayout {
 
     /// Get the offset in guest memory to the input data pointer.
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-    fn get_input_data_pointer_offset(&self) -> usize {
+    pub(super) fn get_input_data_pointer_offset(&self) -> usize {
         // The input data pointer is immediately after the input
         // data size field in the input data `GuestMemoryRegion` struct which is a `u64`.
         self.get_input_data_size_offset() + size_of::<u64>()

src/hyperlight_host/src/mem/mod.rs

@@ -31,6 +31,8 @@ pub mod memory_region;
 /// Functionality that wraps a `SandboxMemoryLayout` and a
 /// `SandboxMemoryConfig` to mutate a sandbox's memory as necessary.
 pub mod mgr;
+/// A compact snapshot representation for memory pages
+pub(crate) mod page_snapshot;
 /// Structures to represent pointers into guest and host memory
 pub mod ptr;
 /// Structures to represent memory address spaces into which pointers
@@ -47,5 +49,7 @@ pub mod shared_mem_snapshot;
 /// Utilities for writing shared memory tests
 #[cfg(test)]
 pub(crate) mod shared_mem_tests;
+/// A wrapper around a `SharedMemory` to manage snapshots of the memory
+pub mod shared_memory_snapshot_manager;
 #[cfg(target_os = "windows")]
 mod windows_dirty_page_tracker;

src/hyperlight_host/src/mem/page_snapshot.rs

@@ -0,0 +1,88 @@
+/*
+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.
+*/
+
+use std::collections::HashMap;
+
+use hyperlight_common::mem::PAGE_SIZE_USIZE;
+
+/// A compact snapshot representation that stores pages in a contiguous buffer
+/// with an index for efficient lookup.
+///
+/// This struct is designed to efficiently store and retrieve memory snapshots
+/// by using a contiguous buffer for all page data combined with a HashMap index
+/// for page lookups. This approach reduces memory overhead
+/// compared to storing pages individually.
+///
+/// # Clone Derivation
+///
+/// This struct derives `Clone` because it's stored in `Vec<PageSnapshot>` within
+/// `SharedMemorySnapshotManager`, which itself derives `Clone`.
+#[derive(Clone)]
+pub(super) struct PageSnapshot {
+    /// Maps page numbers to their offset within the buffer (in page units)
+    page_index: HashMap<usize, usize>, // page_number -> buffer_offset_in_pages
+    /// Contiguous buffer containing all the page data
+    buffer: Vec<u8>,
+    /// How many non-main-RAM regions were mapped when this snapshot was taken?
+    mapped_rgns: u64,
+}
+
+impl PageSnapshot {
+    /// Create a snapshot from a list of page numbers with pre-allocated buffer
+    pub(super) fn with_pages_and_buffer(
+        page_numbers: Vec<usize>,
+        buffer: Vec<u8>,
+        mapped_rgns: u64,
+    ) -> Self {
+        let page_count = page_numbers.len();
+        let mut page_index = HashMap::with_capacity(page_count);
+
+        // Map each page number to its offset in the buffer
+        for (buffer_offset, page_num) in page_numbers.into_iter().enumerate() {
+            page_index.insert(page_num, buffer_offset);
+        }
+
+        Self {
+            page_index,
+            buffer,
+            mapped_rgns,
+        }
+    }
+
+    /// Get page data by page number, returns None if page is not in snapshot
+    pub(super) fn get_page(&self, page_num: usize) -> Option<&[u8]> {
+        self.page_index.get(&page_num).map(|&buffer_offset| {
+            let start = buffer_offset * PAGE_SIZE_USIZE;
+            let end = start + PAGE_SIZE_USIZE;
+            &self.buffer[start..end]
+        })
+    }
+
+    /// Get an iterator over all page numbers in this snapshot
+    pub(super) fn page_numbers(&self) -> impl Iterator<Item = usize> + '_ {
+        self.page_index.keys().copied()
+    }
+
+    /// Get the maximum page number in this snapshot, or None if empty
+    pub(super) fn max_page(&self) -> Option<usize> {
+        self.page_index.keys().max().copied()
+    }
+
+    /// Get the number of mapped regions when this snapshot was taken
+    pub(super) fn mapped_rgns(&self) -> u64 {
+        self.mapped_rgns
+    }
+}

src/hyperlight_host/src/mem/shared_mem.rs

@@ -627,6 +627,15 @@ impl ExclusiveSharedMemory {
         Ok(())
     }
 
+    /// Copies bytes from `self` to `dst` starting at offset
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub fn copy_to_slice(&self, dst: &mut [u8], offset: usize) -> Result<()> {
+        let data = self.as_slice();
+        bounds_check!(offset, dst.len(), data.len());
+        dst.copy_from_slice(&data[offset..offset + dst.len()]);
+        Ok(())
+    }
+
     /// Return the address of memory at an offset to this `SharedMemory` checking
     /// that the memory is within the bounds of the `SharedMemory`.
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
@@ -635,6 +644,15 @@ impl ExclusiveSharedMemory {
         Ok(self.base_addr() + offset)
     }
 
+    /// Fill the memory in the range `[offset, offset + len)` with `value`
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub fn zero_fill(&mut self, offset: usize, len: usize) -> Result<()> {
+        bounds_check!(offset, len, self.mem_size());
+        let data = self.as_mut_slice();
+        data[offset..offset + len].fill(0);
+        Ok(())
+    }
+
     generate_reader!(read_u8, u8);
     generate_reader!(read_i8, i8);
     generate_reader!(read_u16, u16);