feat: multi-space page-table walk + link primitives

Adds SpaceAwareMapping, SpaceReferenceMapping, SpaceId, walk_va_spaces,
and space_aware_map in hyperlight_common::vmem. These primitives let a
caller walk several PT roots together, detect aliased intermediate
tables (the 'kernel-half shared' pattern where multiple process PDs
point at the same PT pages), and rebuild those aliases on the write
side by linking into already-built tables instead of copying.

i686 implements the real walker (depth-1 PT sharing, which is what
Nanvix needs) plus space_aware_map that composes a PDE from the owning
space's rebuilt table. amd64 gets a non-aliasing single-root walker so
the architecture-independent re-export compiles; no current amd64
embedder exercises aliasing. aarch64 remains a TODO with stubs.

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

Author: danbugs

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

@@ -46,6 +46,29 @@ pub unsafe fn virt_to_phys<'a, Op: TableReadOps + 'a>(
     core::iter::empty()
 }
 
+/// Stub — see [`crate::vmem::walk_va_spaces`].
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn walk_va_spaces<Op: TableReadOps>(
+    _op: &Op,
+    _roots: &[Op::TableAddr],
+    _address: u64,
+    _len: u64,
+) -> ::alloc::vec::Vec<(
+    crate::vmem::SpaceId,
+    ::alloc::vec::Vec<crate::vmem::SpaceAwareMapping>,
+)> {
+    ::alloc::vec::Vec::new()
+}
+
+/// Stub — see [`crate::vmem::space_aware_map`].
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn space_aware_map<Op: TableOps>(
+    _op: &Op,
+    _ref_map: crate::vmem::SpaceReferenceMapping,
+    _built_roots: &::alloc::collections::BTreeMap<crate::vmem::SpaceId, Op::TableAddr>,
+) {
+}
+
 pub trait TableMovability<Op: TableReadOps + ?Sized, TableMoveInfo> {}
 impl<Op: TableOps<TableMovability = crate::vmem::MayMoveTable>> TableMovability<Op, Op::TableAddr>
     for crate::vmem::MayMoveTable

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

@@ -265,6 +265,99 @@ unsafe fn map_page<
 /// 2. PDPT (38:30) - allocate PD if needed
 /// 3. PD (29:21) - allocate PT if needed
 /// 4. PT (20:12) - write final PTE with physical address and flags
+/// Multi-space page-table walking on amd64: walks each root
+/// independently and emits all leaves as `ThisSpace`. Aliased
+/// intermediate-table detection is not implemented — no current
+/// embedder exercises that pattern on amd64 (it is an i686/Nanvix-
+/// specific concern).
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn walk_va_spaces<Op: TableReadOps>(
+    op: &Op,
+    roots: &[Op::TableAddr],
+    address: u64,
+    len: u64,
+) -> ::alloc::vec::Vec<(
+    crate::vmem::SpaceId,
+    ::alloc::vec::Vec<crate::vmem::SpaceAwareMapping>,
+)> {
+    use ::alloc::vec::Vec;
+
+    let mut out: Vec<(
+        crate::vmem::SpaceId,
+        Vec<crate::vmem::SpaceAwareMapping>,
+    )> = Vec::with_capacity(roots.len());
+
+    let addr = address & ((1u64 << VA_BITS) - 1);
+    let vmin = addr & !(PAGE_SIZE as u64 - 1);
+    let vmax = core::cmp::min(addr + len, 1u64 << VA_BITS);
+
+    for &root in roots {
+        #[allow(clippy::unnecessary_cast)]
+        let root_id: crate::vmem::SpaceId = Op::to_phys(root) as u64;
+        let mut mappings: Vec<crate::vmem::SpaceAwareMapping> = Vec::new();
+
+        let iter = modify_ptes::<47, 39, Op, _>(MapRequest {
+            table_base: root,
+            vmin,
+            len: vmax.saturating_sub(vmin),
+            update_parent: UpdateParentNone {},
+        })
+        .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, _>);
+
+        for r in iter {
+            let Some(pte) = (unsafe { read_pte_if_present(op, r.entry_ptr) }) else {
+                continue;
+            };
+            let phys_addr = pte & PTE_ADDR_MASK;
+            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 executable = (pte & PAGE_NX) == 0;
+            let avl = pte & PTE_AVL_MASK;
+            let kind = if avl == PAGE_AVL_COW {
+                MappingKind::Cow(CowMapping {
+                    readable: true,
+                    executable,
+                })
+            } else {
+                MappingKind::Basic(BasicMapping {
+                    readable: true,
+                    writable: (pte & PAGE_RW) != 0,
+                    executable,
+                })
+            };
+            mappings.push(crate::vmem::SpaceAwareMapping::ThisSpace(Mapping {
+                phys_base: phys_addr,
+                virt_base: virt_addr,
+                len: PAGE_SIZE as u64,
+                kind,
+                user_accessible: false,
+            }));
+        }
+
+        out.push((root_id, mappings));
+    }
+
+    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.
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn space_aware_map<Op: TableOps>(
+    _op: &Op,
+    _ref_map: crate::vmem::SpaceReferenceMapping,
+    _built_roots: &::alloc::collections::BTreeMap<crate::vmem::SpaceId, Op::TableAddr>,
+) {
+}
+
 #[allow(clippy::missing_safety_doc)]
 pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
     modify_ptes::<47, 39, Op, _>(MapRequest {

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

@@ -22,9 +22,9 @@ limitations under the License.
 //! Entries are 4 bytes wide. There is no NX bit; all pages are executable.
 
 use crate::vmem::{
-    BasicMapping, CowMapping, MapRequest, MapResponse, Mapping, MappingKind, TableMovabilityBase,
-    TableOps, TableReadOps, UpdateParent, UpdateParentNone, modify_ptes, read_pte_if_present,
-    require_pte_exist, write_entry_updating,
+    BasicMapping, CowMapping, MapRequest, MapResponse, Mapping, MappingKind, SpaceAwareMapping,
+    SpaceId, SpaceReferenceMapping, TableMovabilityBase, TableOps, TableReadOps, UpdateParent,
+    UpdateParentNone, modify_ptes, read_pte_if_present, require_pte_exist, write_entry_updating,
 };
 
 pub const PAGE_SIZE: usize = 4096;
@@ -169,6 +169,187 @@ pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
     .for_each(drop);
 }
 
+//==================================================================================================
+// 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, .. })`
+/// entries.
+///
+/// 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).
+///
+/// # Safety
+/// Same invariants as [`virt_to_phys`]. Callers must not mutate the
+/// page tables concurrently.
+#[allow(clippy::missing_safety_doc)]
+pub unsafe fn walk_va_spaces<Op: TableReadOps>(
+    op: &Op,
+    roots: &[Op::TableAddr],
+    address: u64,
+    len: u64,
+) -> ::alloc::vec::Vec<(SpaceId, ::alloc::vec::Vec<SpaceAwareMapping>)> {
+    use ::alloc::vec::Vec;
+
+    // Map: PT PA -> (owner SpaceId, the VA at which the owner used
+    // this PT). Subsequent visits to the same PT PA emit AnotherSpace.
+    let mut seen_pts: ::alloc::collections::BTreeMap<u64, (SpaceId, u64)> =
+        ::alloc::collections::BTreeMap::new();
+    let mut results: Vec<(SpaceId, Vec<SpaceAwareMapping>)> = Vec::with_capacity(roots.len());
+
+    let vmin = address & !(PAGE_SIZE as u64 - 1);
+    let vmax = core::cmp::min(address + len, 1u64 << VA_BITS);
+
+    for &root in roots {
+        #[allow(clippy::unnecessary_cast)]
+        let root_id: SpaceId = Op::to_phys(root) as u64;
+        let mut mappings: Vec<SpaceAwareMapping> = Vec::new();
+
+        // Iterate PDEs covering [vmin, vmax) at the PD level (bits 31:22).
+        let pde_iter = modify_ptes::<31, 22, Op, _>(MapRequest {
+            table_base: root,
+            vmin,
+            len: vmax.saturating_sub(vmin),
+            update_parent: UpdateParentNone {},
+        });
+        for r in pde_iter {
+            let Some(pde) = (unsafe { read_pte_if_present(op, r.entry_ptr) }) else {
+                continue;
+            };
+            let pt_pa: u64 = pde & PTE_ADDR_MASK;
+
+            // Seen this PT via an earlier root? Emit AnotherSpace and
+            // don't descend — the sub-tree is fully described by the
+            // owner's entries.
+            if let Some(&(owner, their_va)) = seen_pts.get(&pt_pa) {
+                if owner != root_id {
+                    mappings.push(SpaceAwareMapping::AnotherSpace(SpaceReferenceMapping {
+                        depth: SHARED_TABLE_DEPTH,
+                        space: owner,
+                        our_va: r.vmin,
+                        their_va,
+                    }));
+                    continue;
+                }
+                // Same space saw this PT before (shouldn't happen with
+                // the virt_to_phys-style per-PDE iteration, but skip
+                // defensively).
+                continue;
+            }
+            seen_pts.insert(pt_pa, (root_id, r.vmin));
+
+            // Descend the PT and emit ThisSpace entries for each live
+            // 4KB leaf, mirroring virt_to_phys's leaf-emission logic.
+            let pt_request = MapRequest {
+                #[allow(clippy::unnecessary_cast)]
+                table_base: Op::from_phys(pt_pa as PhysAddr),
+                vmin: r.vmin,
+                len: r.len,
+                update_parent: UpdateParentNone {},
+            };
+            for leaf in modify_ptes::<21, 12, Op, _>(pt_request) {
+                let Some(pte) = (unsafe { read_pte_if_present(op, leaf.entry_ptr) }) else {
+                    continue;
+                };
+                let phys_addr = pte & PTE_ADDR_MASK;
+                let avl = pte & PTE_AVL_MASK;
+                let kind = if avl == PAGE_AVL_COW {
+                    MappingKind::Cow(CowMapping {
+                        readable: true,
+                        executable: true,
+                    })
+                } else {
+                    MappingKind::Basic(BasicMapping {
+                        readable: true,
+                        writable: (pte & PAGE_RW) != 0,
+                        executable: true,
+                    })
+                };
+                mappings.push(SpaceAwareMapping::ThisSpace(Mapping {
+                    phys_base: phys_addr,
+                    virt_base: leaf.vmin,
+                    len: PAGE_SIZE as u64,
+                    kind,
+                    user_accessible: (pte & PAGE_USER) != 0,
+                }));
+            }
+        }
+
+        results.push((root_id, mappings));
+    }
+
+    results
+}
+
+/// Install the link described by `ref_map` in `op`'s root PT tree:
+/// look up what the owner space's rebuilt root put at `their_va`'s
+/// 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>(
+    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
+    );
+
+    // Their rebuilt root — must have been populated earlier in the
+    // rebuild loop (walk_va_spaces guarantees topological order).
+    let Some(&their_root) = built_roots.get(&ref_map.space) else {
+        // Defensive: we have no linkage target. Skip rather than
+        // panic. A trace print would live here in a debug build.
+        return;
+    };
+
+    // Read their PDE at their_va's index to get the rebuilt PT PA.
+    let their_pdi = (ref_map.their_va >> 22) & 0x3FF;
+    let their_pde_ptr = Op::entry_addr(
+        their_root,
+        their_pdi * core::mem::size_of::<PageTableEntry>() as u64,
+    );
+    let Some(their_pde) = (unsafe { read_pte_if_present(op, their_pde_ptr) }) else {
+        // Owner didn't end up with a PDE here — nothing to link.
+        return;
+    };
+    let their_pt_pa: u64 = their_pde & PTE_ADDR_MASK;
+
+    // Compose our PDE: point at their PT, preserve their PDE's low
+    // bits (PAGE_USER for kernel-accessible, PAGE_RW, etc.) so the
+    // hardware still honours sharing semantics uniformly.
+    let our_pdi = (ref_map.our_va >> 22) & 0x3FF;
+    let our_root = op.root_table();
+    let our_pde_ptr = Op::entry_addr(
+        our_root,
+        our_pdi * core::mem::size_of::<PageTableEntry>() as u64,
+    );
+
+    let new_pde: u64 = their_pt_pa | (their_pde & !PTE_ADDR_MASK) | PAGE_PRESENT;
+    unsafe {
+        write_entry_updating(op, Op::TableMovability::root_update_parent(), our_pde_ptr, new_pde);
+    }
+}
+
 /// Translate a virtual address range to its backing physical pages.
 ///
 /// # Safety