feat(vmem): fold space_aware_map dispatch

Fixes: #1409

Signed-off-by: Tomasz Andrzejak <andreiltd@gmail.com>

Author: andreiltd

src/hyperlight_common/src/arch/aarch64/vmem.rs

@@ -60,13 +60,18 @@ pub unsafe fn walk_va_spaces<Op: TableReadOps>(
     ::alloc::vec::Vec::new()
 }
 
-/// Stub — see [`crate::vmem::space_aware_map`].
+/// aarch64 never emits `AnotherSpace` entries, so [`SpaceReferenceMapping`]
+/// is [`crate::vmem::Void`] making the `AnotherSpace` match arm statically unreachable.
+pub type SpaceReferenceMapping = crate::vmem::Void;
+
+/// Statically unreachable: amd64 never emits `AnotherSpace`.
 #[allow(clippy::missing_safety_doc)]
-pub unsafe fn space_aware_map<Op: TableOps>(
+pub(super) unsafe fn link_shared_table<Op: TableOps>(
     _op: &Op,
-    _ref_map: crate::vmem::SpaceReferenceMapping,
+    impossible: crate::vmem::Void,
     _built_roots: &::alloc::collections::BTreeMap<crate::vmem::SpaceId, Op::TableAddr>,
 ) {
+    match impossible {}
 }
 
 pub trait TableMovability<Op: TableReadOps + ?Sized, TableMoveInfo> {}

src/hyperlight_common/src/arch/amd64/vmem.rs

@@ -412,15 +412,18 @@ pub unsafe fn walk_va_spaces<Op: TableReadOps>(
     out
 }
 
-/// See [`walk_va_spaces`]: amd64 never emits `AnotherSpace`, so this
-/// is unreachable in practice. It silently no-ops (rather than
-/// panicking) to keep the architecture-independent re-export usable.
+/// aarch64 never emits `AnotherSpace` entries, so [`SpaceReferenceMapping`]
+/// is [`crate::vmem::Void`] making the `AnotherSpace` match arm statically unreachable.
+pub type SpaceReferenceMapping = crate::vmem::Void;
+
+/// Statically unreachable: amd64 never emits `AnotherSpace`.
 #[allow(clippy::missing_safety_doc)]
-pub unsafe fn space_aware_map<Op: TableOps>(
+pub(super) unsafe fn link_shared_table<Op: TableOps>(
     _op: &Op,
-    _ref_map: crate::vmem::SpaceReferenceMapping,
+    impossible: crate::vmem::Void,
     _built_roots: &::alloc::collections::BTreeMap<crate::vmem::SpaceId, Op::TableAddr>,
 ) {
+    match impossible {}
 }
 
 #[allow(clippy::missing_safety_doc)]

src/hyperlight_common/src/arch/i686/vmem.rs

@@ -23,8 +23,8 @@ limitations under the License.
 
 use crate::vmem::{
     BasicMapping, CowMapping, MapRequest, MapResponse, Mapping, MappingKind, SpaceAwareMapping,
-    SpaceId, SpaceReferenceMapping, TableMovabilityBase, TableOps, TableReadOps, UpdateParent,
-    UpdateParentNone, modify_ptes, write_entry_updating,
+    SpaceId, TableMovabilityBase, TableOps, TableReadOps, UpdateParent, UpdateParentNone,
+    modify_ptes, write_entry_updating,
 };
 
 pub const PAGE_SIZE: usize = 4096;
@@ -44,6 +44,40 @@ const PAGE_AVL_COW: u64 = 1 << 9;
 
 const VA_BITS: usize = 32;
 
+/// The level of the page-table hierarchy at which sharing occurs.
+///
+/// On i686 the only shareable level is the leaf page table: the table a PDE points at.
+#[derive(Debug, Clone, Copy)]
+pub enum SharedTableLevel {
+    /// A leaf page table
+    Pt,
+}
+
+/// A reference from one address space to an intermediate page table
+/// that lives in a different space. Produced by [`walk_va_spaces`]
+/// when the walker encounters an intermediate table whose physical
+/// address was already seen via an earlier root.
+///
+/// On i686 the only shareable level is [`SharedTableLevel::Pt`] (one
+/// level below the root PD).
+///
+/// `our_va` and `their_va` may be any virtual address inside the
+/// aliased region but callers must not assume they are region-aligned.
+#[derive(Debug, Clone, Copy)]
+pub struct SpaceReferenceMapping {
+    /// The level at which the alias starts.
+    pub level: SharedTableLevel,
+    /// The "owning" space is the first root that visited this
+    /// intermediate PA during [`walk_va_spaces`].
+    pub space: SpaceId,
+    /// Any VA inside the aliased sub-tree in OUR space.
+    pub our_va: u64,
+    /// Any VA inside the aliased sub-tree in the owning space.
+    /// Usually equal to `our_va` (kernel mappings at the same VA
+    /// across processes) but the design permits different VAs.
+    pub their_va: u64,
+}
+
 pub trait TableMovability<Op: TableReadOps + ?Sized, TableMoveInfo> {
     type RootUpdateParent: UpdateParent<Op, TableMoveInfo = TableMoveInfo>;
     fn root_update_parent() -> Self::RootUpdateParent;
@@ -213,17 +247,16 @@ pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
 // Multi-space walk / link (shared intermediate tables)
 //==================================================================================================
 
-/// i686 has two levels (PD -> PT). The only sharable thing is a PT,
-/// at depth 1 (one level below the root PD).
-const SHARED_TABLE_DEPTH: usize = 1;
-
 /// Walk multiple root PDs together, detecting PDEs that point at the
 /// same PT PA across roots (i.e. aliased PTs — the standard
 /// "kernel-half shared" trick on x86 without KPTI). The first root to
 /// visit a given PT PA becomes the "owner"; later roots that alias it
-/// receive `AnotherSpace(SpaceReferenceMapping { depth: 1, .. })`
+/// receive `AnotherSpace(SpaceReferenceMapping { level: SharedTableLevel::Pt, .. })`
 /// entries.
 ///
+/// i686 has two levels (PD -> PT). The only sharable thing is a PT,
+/// one level below the root PD ([`SharedTableLevel::Pt`]).
+///
 /// Generic over `TableAddr` so it works with both the in-guest
 /// implementation (`TableAddr = u32`, backed by raw pointers) and the
 /// host-side snapshot buffer (`TableAddr = u64`, byte offsets).
@@ -273,7 +306,7 @@ pub unsafe fn walk_va_spaces<Op: TableReadOps>(
             if let Some(&(owner, their_va)) = seen_pts.get(&pt_pa) {
                 if owner != root_id {
                     mappings.push(SpaceAwareMapping::AnotherSpace(SpaceReferenceMapping {
-                        depth: SHARED_TABLE_DEPTH,
+                        level: SharedTableLevel::Pt,
                         space: owner,
                         our_va: r.vmin,
                         their_va,
@@ -335,24 +368,18 @@ pub unsafe fn walk_va_spaces<Op: TableReadOps>(
 /// PDE slot, and write that PA into our root's PDE slot for
 /// `our_va`. The owner's rebuilt root is found via `built_roots`.
 ///
-/// On i686 `ref_map.depth` must be 1 (PT-level sharing). Other depths
-/// are rejected defensively.
-///
 /// # Safety
 /// Same invariants as [`map`]: caller owns the concurrency story and
 /// must invalidate the TLB if the page tables are live.
 #[allow(clippy::missing_safety_doc)]
-pub unsafe fn space_aware_map<Op: TableOps>(
+pub(super) unsafe fn link_shared_table<Op: TableOps>(
     op: &Op,
     ref_map: SpaceReferenceMapping,
     built_roots: &::alloc::collections::BTreeMap<SpaceId, Op::TableAddr>,
 ) {
-    assert!(
-        ref_map.depth == SHARED_TABLE_DEPTH,
-        "i686 only supports depth={} sharing; got depth={}",
-        SHARED_TABLE_DEPTH,
-        ref_map.depth
-    );
+    match ref_map.level {
+        SharedTableLevel::Pt => {}
+    }
 
     // Their rebuilt root — must have been populated earlier in the
     // rebuild loop (walk_va_spaces guarantees topological order).

src/hyperlight_common/src/vmem.rs

@@ -289,7 +289,10 @@ pub trait TableReadOps {
 
 /// 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.
+/// to move a table.  Also used as the `SpaceReferenceMapping` type on
+/// architectures that never emit `AnotherSpace` entries, making the
+/// corresponding match arm statically unreachable.
+#[derive(Debug, Clone, Copy)]
 pub enum Void {}
 
 /// A marker struct, used by an implementation of [`TableOps`] to
@@ -464,34 +467,11 @@ pub use arch::virt_to_phys;
 pub type SpaceId = u64;
 
 /// A reference from one address space to an intermediate page table
-/// that lives in a different space. Produced by [`walk_va_spaces`] when
-/// the walker encounters an intermediate table (at some `depth` below
-/// the root) whose physical address was already seen via an earlier
-/// root — i.e. the two spaces alias that sub-tree.
-///
-/// Semantics: the level-`depth` block in **our** space that contains
-/// VAs starting at `our_va` is aliased to the level-`depth` block in
-/// `space` that contains VAs starting at `their_va`. Everything below
-/// that sub-tree — PDEs, PTEs, leaf mappings — is shared wholesale.
-///
-/// `depth` is counted from the root:
-/// - `depth = 1` on i686: the shared thing is a leaf PT (the thing a
-///   PDE points to).
-/// - `depth = 1, 2, 3` on amd64: PDPT, PD, or PT respectively.
-#[derive(Debug, Clone, Copy)]
-pub struct SpaceReferenceMapping {
-    /// Depth from the root at which the alias starts (1-based).
-    pub depth: usize,
-    /// The "owning" space — the first root that visited this
-    /// intermediate PA during [`walk_va_spaces`].
-    pub space: SpaceId,
-    /// Start VA of the aliased sub-tree in OUR space.
-    pub our_va: u64,
-    /// Start VA of the aliased sub-tree in the owning space. Usually
-    /// equal to `our_va` (kernel mappings at the same VA across
-    /// processes) but the design permits different VAs.
-    pub their_va: u64,
-}
+/// that lives in a different space. Defined per-arch: on
+/// architectures that never emit `AnotherSpace` entries (amd64,
+/// aarch64) it is [`Void`], making the `AnotherSpace` variant of
+/// [`SpaceAwareMapping`] statically unreachable.
+pub use arch::SpaceReferenceMapping;
 
 /// Either a normal leaf mapping in the current space, or a reference
 /// to an intermediate table in another space. The compaction loop in
@@ -501,19 +481,21 @@ pub struct SpaceReferenceMapping {
 ///   backing page is compacted into the new snapshot blob, the PTE is
 ///   written, and intermediate tables are allocated on demand.
 /// - `AnotherSpace(r)` is rebuilt by *linking*: the entry in our
-///   rebuilt root at depth `r.depth - 1` for `r.our_va` is made to
-///   point at whatever table the owning space ended up with at
-///   `r.their_va`. See [`space_aware_map`].
+///   rebuilt root for `r.our_va` is made to point at whatever table
+///   the owning space ended up with at `r.their_va`.
+///   See [`space_aware_map`].
+///
+/// On architectures where `SpaceReferenceMapping` is [`Void`], the
+/// `AnotherSpace` variant is statically unreachable so match arms
+/// for it should use `match impossible {}`.
 #[derive(Debug)]
 pub enum SpaceAwareMapping {
     ThisSpace(Mapping),
     AnotherSpace(SpaceReferenceMapping),
 }
 
-/// Counterpart of [`walk_va_spaces`]'s `AnotherSpace` entries on the
-/// write side: installs a link in `op`'s root PT tree at `ref_map.our_va`
-/// that points at whatever intermediate table the owning space ended
-/// up with at `ref_map.their_va` (in `built_roots[ref_map.space]`).
+/// Install a [`SpaceAwareMapping`] into the page tables managed by
+/// `op`.
 ///
 /// Callers must ensure that `built_roots` contains populated page
 /// tables for any other space referenced by the mapping.
@@ -522,7 +504,19 @@ pub enum SpaceAwareMapping {
 /// Same invariants as [`map`]: the caller owns the concurrency story
 /// around the page tables being written, and must invalidate TLBs
 /// afterwards if they were live.
-pub use arch::space_aware_map;
+pub unsafe fn space_aware_map<Op: TableOps>(
+    op: &Op,
+    sam: SpaceAwareMapping,
+    built_roots: &::alloc::collections::BTreeMap<SpaceId, Op::TableAddr>,
+) {
+    match sam {
+        SpaceAwareMapping::ThisSpace(mapping) => unsafe { arch::map(op, mapping) },
+        SpaceAwareMapping::AnotherSpace(ref_map) => unsafe {
+            arch::link_shared_table(op, ref_map, built_roots)
+        },
+    }
+}
+
 /// Walk multiple page-table roots together, emitting either a normal
 /// leaf mapping (`ThisSpace`) or a reference to an alias that was
 /// already seen via an earlier root (`AnotherSpace`).

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -552,18 +552,25 @@ impl Snapshot {
                                     kind,
                                     user_accessible: mapping.user_accessible,
                                 };
-                                unsafe { vmem::map(&pt_buf, compacted) };
-                            }
-                            SpaceAwareMapping::AnotherSpace(ref_map) => {
-                                // Link to the owning space's already-
-                                // rebuilt intermediate table — this
-                                // is what preserves Nanvix's
-                                // kernel-half-shared invariant across
-                                // process PDs after relocation.
+
                                 unsafe {
-                                    vmem::space_aware_map(&pt_buf, ref_map, &built_roots);
-                                }
+                                    vmem::space_aware_map(
+                                        &pt_buf,
+                                        SpaceAwareMapping::ThisSpace(compacted),
+                                        &built_roots,
+                                    )
+                                };
                             }
+                            // Preserve Nanvix's kernel-half-shared
+                            // invariant by linking to the owning
+                            // space's rebuilt intermediate table.
+                            SpaceAwareMapping::AnotherSpace(ref_map) => unsafe {
+                                vmem::space_aware_map(
+                                    &pt_buf,
+                                    SpaceAwareMapping::AnotherSpace(ref_map),
+                                    &built_roots,
+                                )
+                            },
                         }
                     }
                 }