supermachine 0.7.70

// Snapshot/restore for the VMM. Captures the VM state needed to
// resume in a fresh process: per-vCPU registers (GP, SIMD, sysregs),
// per-vCPU interrupt-controller state, the opaque VM-global interrupt-
// controller blob (`intc_blob`; GICv3 on HVF, IOAPIC/PIC/kvmclock on
// KVM), the guest-clock reference, and full guest RAM. Per-vCPU
// register and clock encoding is owned by the hypervisor seam so this
// container framing is backend-neutral.
//
// V2 (this revision): captures virtio mirror state — per-MMIO-device
// negotiated features, status, and per-queue ring addresses + cursors,
// plus the vsock muxer's TSI listener registry. In-flight TCP/UDP
// connections and pending handshakes are intentionally dropped (peer
// TCPs see RST; clients retry). Listeners are re-bound on restore
// with fresh ephemeral host ports — the guest only knows its vm_port,
// so the change is transparent.
//
// Wire format ("SMSNAP\x06\x00", little-endian):
//   header 64 B:
//     [0..8]   magic
//     [8..16]  version (u64)
//     [16..24] captured_host_ticks (u64)
//     [24..32] captured_clock_ref (u64)
//     [32..40] memory_bytes (u64)
//     [40..48] intc_blob_len (u64)
//     [48..52] n_vcpus (u32)
//     [52..56] reserved (u32)
//     [56..64] ram_gpa (u64)
//   per_vcpu × n_vcpus (owned by the hypervisor seam's
//   write_snapshot_state/read_snapshot_state — see crate::hvf::backend):
//     vtimer_offset (u64),
//     gp_n simd_n sys_n icc_n redist_n ich_n  (u32×6)
//     gp_regs (reg_id u32, val u64)        × gp_n
//     simd_regs (reg_id u32, val u128)     × simd_n
//     sys_regs (reg_id u32, val u64)       × sys_n
//     icc_regs (reg_id u32, val u64)       × icc_n
//     redist (off u32, val u64)            × redist_n
//     ich_regs (reg_id u32, val u64)       × ich_n
//   then RAM: memory_bytes raw bytes
//
// Memory and metadata live in the same file for v1 simplicity.

#![cfg(all(target_os = "macos", target_arch = "aarch64"))]

use std::io::{Read, Write};

use crate::devices::virtio::vsock::muxer::TsiListenerSnapshot;
// Re-exported so existing `crate::vmm::snapshot::PerVcpuState` references
// (coord/runner/worker) keep resolving after the type moved to the HVF backend.
pub use crate::hvf::vcpu_snapshot::PerVcpuState;
use crate::hypervisor::{BackendError, HypervisorVcpu, HypervisorVm};
use crate::vmm::vstate::MicroVm;

// v9: per-device virtio-mmio state now uses the canonical `MmioSnapshot` codec
// shared with the KVM pipeline (dropped the stray per-queue alignment pad byte
// the HVF-only v8 writer emitted). No back-compat reader — nothing consumes a
// pre-v9 snapshot (CI/dev re-bake), so this is a clean single-path bump.
const SNAPSHOT_MAGIC: [u8; 8] = *b"SMSNAP\x0a\x00";
/// v10: + per-listener `inet_port` in the TSI listener record (the
///      AF_INET userspace port the guest bound — e.g. Chrome on 9222
///      vs nginx on 80). Before v10 it was never serialized, so any
///      disk / cross-process restore lost the pin and multi-port
///      `Vm::expose_tcp` routing broke ("no TSI listener"). All record
///      (de)serialization now routes through
///      `write_tsi_listener_record` / `read_tsi_listener_record`.
/// v9: + per-vCPU `ich_regs` (ICH_LR0..15_EL2 + HCR_EL2 + VMCR_EL2 +
///      AP0R0_EL2 + AP1R0_EL2). These are the EL2 virtual-interrupt
///      control registers — the per-PE state HVF's
///      `hv_gic_state_get_data()` opaque blob does NOT cover (per
///      Apple's header doc). On secondary vCPUs captured mid-flight,
///      LRs hold pending virtual interrupts the guest hasn't EOI'd
///      yet; restoring them as zero made the restored kernel see
///      spurious / missing IRQs and crash on resume — the root
///      cause of the multi-vCPU snapshot intermittency tracked in
///      docs/design/multi-vcpu-snapshot-intermittency-2026-04-27.md.
/// The version gate (`load_from_file_inner`) rejects ANY non-current
/// version: a bump forces an automatic re-bake (the bake fingerprint
/// also tracks the worker binary), so there is no multi-version read
/// path to maintain.
const SNAPSHOT_VERSION: u64 = 11;
/// Page size to align the RAM region within the snapshot file. macOS
/// on Apple Silicon uses 16 KiB pages; aligning here lets `--cow-restore`
/// mmap the region with `MAP_PRIVATE` directly out of the file.
const RAM_PAGE_ALIGN: u64 = 16384;
const SPARSE_RAM_CHUNK: usize = 64 * 1024;

/// Per-VM virtio device state as unified [`DeviceRecord`]s (the cross-backend
/// device section shared with the KVM pipeline). Order matches the bus device
/// order in the file; HVF restore applies them positionally (it ignores the
/// per-record kind/backing — backing is `None`, supplied by the Image), while
/// the KVM pipeline reconstructs from the inline backing. Phase-3 7c step 3.
#[derive(Default)]
pub struct VirtioSnapshot {
    pub devices: Vec<crate::snapshot_frame::DeviceRecord>,
    pub vsock_listeners: Vec<TsiListenerSnapshot>,
}

pub struct Snapshot {
    pub captured_host_ticks: u64,
    pub captured_clock_ref: u64,
    pub ram_gpa: u64,
    pub memory: Vec<u8>,
    pub intc_blob: Vec<u8>,
    pub per_vcpu: Vec<PerVcpuState>,
    pub virtio: VirtioSnapshot,
}

pub struct SnapshotWriteStats {
    pub ram_bytes: u64,
    pub ram_data_bytes: u64,
    pub ram_zero_bytes: u64,
}

#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
pub struct SnapshotRestoreTimings {
    pub ram_copy_us: u128,
    pub gic_restore_us: u128,
    pub vcpu_restore_us: u128,
    pub vtimer_offset_us: u128,
    /// The CNTVOFF_EL2 value just written to the boot vCPU. Returned
    /// so the caller (builder.rs) can publish it to the
    /// coordinator's `vtimer_offset` for secondaries to pick up,
    /// WITHOUT round-tripping through HVF's `get_vtimer_offset` —
    /// the round-trip could silently fail (HVF returned Err), the
    /// caller would skip the store, secondaries would see `0`
    /// (treated as "no offset published"), keep HVF default
    /// CNTVOFF=0, see raw CNTPCT, and the in-kernel virtual timer
    /// programmed before snapshot would either fire instantly (CVAL
    /// way in the past relative to raw CNTPCT) or never fire (CVAL
    /// way in the future). Either way the kernel's per-CPU tick
    /// machinery on the secondary stalls, RCU starves, userspace
    /// processes scheduled there make no progress, and the workload
    /// times out without an obvious panic. Returning the value
    /// directly is bit-identical (we just-set it) AND avoids the
    /// silent-error path.
    pub applied_vtimer_offset: u64,
}

#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
pub struct SnapshotRestoreOptions {
    pub skip_intc_blob: bool,
}

pub fn capture_snapshot(
    vm: &MicroVm,
    virtio: VirtioSnapshot,
) -> crate::hypervisor::ActiveResult<Snapshot> {
    let per0 = vm.vcpu.capture_snapshot()?;
    let intc_blob = vm.vm.capture_intc()?;
    let mut memory = vec![0u8; vm.ram_size];
    // SAFETY: ram_host is vm.ram_size bytes.
    unsafe {
        std::ptr::copy_nonoverlapping(vm.ram_host, memory.as_mut_ptr(), vm.ram_size);
    }
    let captured_host_ticks = crate::hypervisor::ActiveVm::host_monotonic_ticks();
    let captured_clock_ref =
        <crate::hypervisor::ActiveVcpu as crate::hypervisor::HypervisorVcpu>::capture_clock_ref(
            &per0,
            captured_host_ticks,
        );
    Ok(Snapshot {
        captured_host_ticks,
        captured_clock_ref,
        ram_gpa: vm.ram_gpa,
        memory,
        intc_blob,
        per_vcpu: vec![per0],
        virtio,
    })
}

/// Restore the full snapshot into `vm`. RAM is memcpy'd in (for v1
/// we don't do CoW — that's a follow-up). Caller must have created
/// `vm` with the same RAM size as the snapshot (otherwise the memcpy
/// is silently truncated).
pub fn restore_snapshot(vm: &MicroVm, snap: &Snapshot) -> crate::hypervisor::ActiveResult<()> {
    restore_snapshot_timed(vm, snap).map(|_| ())
}

pub fn restore_snapshot_timed(
    vm: &MicroVm,
    snap: &Snapshot,
) -> crate::hypervisor::ActiveResult<SnapshotRestoreTimings> {
    restore_snapshot_timed_with_options(vm, snap, SnapshotRestoreOptions::default())
}

pub fn restore_snapshot_timed_with_options(
    vm: &MicroVm,
    snap: &Snapshot,
    options: SnapshotRestoreOptions,
) -> crate::hypervisor::ActiveResult<SnapshotRestoreTimings> {
    let mut timings = SnapshotRestoreTimings::default();
    // 1. RAM. CoW-restored snapshots leave `memory` empty (pages are
    //    already mapped via mmap(MAP_PRIVATE) before MicroVm::new).
    if !snap.memory.is_empty() {
        let t0 = std::time::Instant::now();
        // SAFETY: ram_host is vm.ram_size bytes; we cap copy length.
        unsafe {
            std::ptr::copy_nonoverlapping(
                snap.memory.as_ptr(),
                vm.ram_host,
                vm.ram_size.min(snap.memory.len()),
            );
        }
        timings.ram_copy_us = t0.elapsed().as_micros();
    }
    // 2. GIC blob (covers distributor + per-PE pending/active).
    if !options.skip_intc_blob {
        let t0 = std::time::Instant::now();
        vm.vm.restore_intc(&snap.intc_blob)?;
        timings.gic_restore_us = t0.elapsed().as_micros();
    }
    // 3. Per-vCPU state.
    let boot_vcpu = snap
        .per_vcpu
        .first()
        .ok_or(crate::hypervisor::ActiveError::other(
            "vmm snapshot internal error",
        ))?;
    let t0 = std::time::Instant::now();
    vm.vcpu.restore_snapshot(boot_vcpu)?;
    timings.vcpu_restore_us = t0.elapsed().as_micros();
    // 4. Coherent guest-timer re-anchor for all vCPUs (single offset),
    //    routed through the seam so the backend owns the timer math.
    let now = crate::hypervisor::ActiveVm::host_monotonic_ticks();
    let t0 = std::time::Instant::now();
    let new_offset = vm.vcpu.restore_clock(snap.captured_clock_ref, now)?;
    timings.vtimer_offset_us = t0.elapsed().as_micros();
    timings.applied_vtimer_offset = new_offset;
    Ok(timings)
}

// ----- File serialization -----

#[derive(Debug)]
pub enum FileError {
    Io(std::io::Error),
    BadMagic,
    BadVersion(u64),
    Malformed(&'static str),
    Truncated,
}
impl From<std::io::Error> for FileError {
    fn from(e: std::io::Error) -> Self {
        Self::Io(e)
    }
}

// LE field readers live in the portable `crate::snapshot_frame` substrate
// (shared with the KVM pipeline). These thin wrappers preserve this module's
// `FileError`-typed signatures so the sidecar call sites are unchanged while the
// implementation is the single audited one.
fn le_u16(bytes: &[u8]) -> Result<u16, FileError> {
    crate::snapshot_frame::le_u16(bytes).ok_or(FileError::Truncated)
}

fn le_u32(bytes: &[u8]) -> Result<u32, FileError> {
    crate::snapshot_frame::le_u32(bytes).ok_or(FileError::Truncated)
}

fn le_u64(bytes: &[u8]) -> Result<u64, FileError> {
    crate::snapshot_frame::le_u64(bytes).ok_or(FileError::Truncated)
}

/// Serialized length of one TSI listener record (v10+):
/// cid(8) + peer_port(4) + vm_port(4) + family(2) + socktype(2)
/// + inet_port(4) = 24 bytes.
const TSI_LISTENER_RECORD_LEN: usize = 24;

/// Serialize one TSI listener. SINGLE SOURCE OF TRUTH for the record
/// layout — every encoder (`encode_compact_meta`,
/// `save_compact_to_file_via_clone`) MUST route through here, and
/// [`read_tsi_listener_record`] MUST mirror it. This field was hand-
/// packed in three places before v10, and the `inet_port` pin silently
/// fell off the disk format because one site was never updated; a
/// shared helper is what stops that recurring.
///
/// `inet_port` (the AF_INET userspace port the guest bound — e.g.
/// Chrome on 9222 vs nginx on 80) is encoded as a u32 with `u32::MAX`
/// as the `None` sentinel (real ports are u16, so no collision).
fn write_tsi_listener_record(buf: &mut Vec<u8>, l: &TsiListenerSnapshot) {
    // Canonical codec on the type — shared with the KVM pipeline.
    l.write_to(buf).expect("Vec<u8> write is infallible");
}

/// Deserialize one TSI listener. Mirror of [`write_tsi_listener_record`].
fn read_tsi_listener_record(
    e: &[u8; TSI_LISTENER_RECORD_LEN],
) -> Result<TsiListenerSnapshot, FileError> {
    TsiListenerSnapshot::read_from(&mut &e[..]).map_err(|_| FileError::Truncated)
}

pub fn save_to_file(path: &str, snap: &Snapshot) -> Result<(), FileError> {
    save_to_file_with_stats(path, snap).map(|_| ())
}

/// Derive the PosixFs sidecar path that lives next to a snapshot file.
///
/// We persist FUSE backend state out-of-band rather than extending the
/// snapshot binary format — keeps the bake format stable across the
/// fix, and pre-0.7.6 snapshots simply have no sidecar (the old
/// "warmup-walked paths only" behaviour stays the failure mode).
///
/// Example: `restore.snap` → `restore.snap.posixfs`.
pub fn posix_fs_sidecar_path(snap_path: &str) -> String {
    format!("{snap_path}.posixfs")
}

/// Convenience: capture each PosixFs's state via the FsBackend
/// trait method and write the sidecar next to `snap_path`. Logs but
/// does not fail the save on sidecar errors — the snapshot file is
/// already on disk and usable; the sidecar is an enhancement.
pub fn capture_and_write_posix_fs_sidecar(
    snap_path: &str,
    posix_fs: &[std::sync::Arc<crate::fuse::PosixFs>],
) {
    if posix_fs.is_empty() {
        return;
    }
    let blobs: Vec<Option<Vec<u8>>> = posix_fs
        .iter()
        .map(|pfs| Some(pfs.snapshot_state()))
        .collect();
    let sidecar_path = posix_fs_sidecar_path(snap_path);
    if let Err(e) = write_posix_fs_sidecar(&sidecar_path, &blobs) {
        eprintln!(
            "supermachine: warning: posix-fs sidecar write to {sidecar_path} failed ({e:?}); \
             warm restores from this snapshot will fall back to lazy LOOKUP and may EAI \
             on paths not walked during warmup"
        );
    }
}

/// Write the PosixFs sidecar for a list of per-mount state blobs.
/// `blobs[i]` is the binary blob from `PosixFs::snapshot_state()` for
/// mount index `i`; an `None` entry means that mount index has no
/// state to persist (in-memory backend; would never appear for real
/// bind mounts). The sidecar layout is:
///
/// ```text
/// b"PFXS"               4-byte magic
/// u32 version           little-endian, currently 1
/// u32 mount_count       number of mounts
/// for each mount:
///   u32 blob_len        0 means "no state for this mount"
///   bytes blob          length blob_len
/// ```
pub fn write_posix_fs_sidecar(
    sidecar_path: &str,
    blobs: &[Option<Vec<u8>>],
) -> Result<(), FileError> {
    use std::io::Write;
    // Build in memory first — sidecars are small (<MB typically).
    let mut buf: Vec<u8> = Vec::with_capacity(
        16 + blobs
            .iter()
            .map(|b| 4 + b.as_ref().map_or(0, |v| v.len()))
            .sum::<usize>(),
    );
    buf.extend_from_slice(b"PFXS");
    buf.extend_from_slice(&1u32.to_le_bytes());
    buf.extend_from_slice(&(blobs.len() as u32).to_le_bytes());
    for b in blobs {
        match b {
            Some(v) => {
                buf.extend_from_slice(&(v.len() as u32).to_le_bytes());
                buf.extend_from_slice(v);
            }
            None => {
                buf.extend_from_slice(&0u32.to_le_bytes());
            }
        }
    }
    let mut f = std::fs::File::create(sidecar_path).map_err(FileError::Io)?;
    f.write_all(&buf).map_err(FileError::Io)?;
    f.sync_all().map_err(FileError::Io)?;
    Ok(())
}

/// Sidecar path for the DAX session metadata, used by the lazy-fault
/// rehydrate path. One file per snapshot, listing all DaxSessions in
/// the same order as the corresponding posix-fs sidecar.
///
/// Example: `restore.snap` → `restore.snap.dax_sidecar`.
pub fn dax_sidecar_path(snap_path: &str) -> String {
    format!("{snap_path}.dax_sidecar")
}

/// Path of the smpark unpark-page-GPA sidecar.
/// Example: `restore.snap` → `restore.snap.smpark`.
pub fn smpark_sidecar_path(snap_path: &str) -> String {
    format!("{snap_path}.smpark")
}

/// Write the smpark unpark-page GPA sidecar (`b"SMPK" + u32 version=1 +
/// u64 gpa`, little-endian). This pins the guest-physical address of
/// smpark.ko's unpark-signal page to the SNAPSHOT rather than the
/// process. A restored worker (which never boots, so never sees the
/// PL011 boot line that sets the process-global `SMPARK_STATE_GPA`)
/// can then host-direct-unpark its secondaries using the GPA THIS
/// snapshot was baked with — correct across processes AND across
/// multiple images sharing one host process, and concurrent-safe
/// (the consumer reads it into per-worker state). Best-effort.
pub fn write_smpark_sidecar(snap_path: &str, gpa: u64) -> Result<(), FileError> {
    use std::io::Write;
    let mut buf = Vec::with_capacity(16);
    buf.extend_from_slice(b"SMPK");
    buf.extend_from_slice(&1u32.to_le_bytes());
    buf.extend_from_slice(&gpa.to_le_bytes());
    let mut f = std::fs::File::create(smpark_sidecar_path(snap_path)).map_err(FileError::Io)?;
    f.write_all(&buf).map_err(FileError::Io)?;
    f.sync_all().map_err(FileError::Io)?;
    Ok(())
}

/// Write this VM's smpark unpark-page GPA (captured by the PL011
/// boot-line scanner into `Vmm::serial.smpark_state_gpa` during the
/// current bake) into the snapshot's sidecar, if announced. The caller
/// passes the per-VM GPA (`vmm.serial.smpark_state_gpa.load(..)`) so
/// this is correct across pooled clones sharing one host process.
/// Best-effort and a no-op when the GPA is 0 (single-vCPU bake, or
/// smpark.ko without the sysfs file).
pub fn capture_and_write_smpark_sidecar(gpa: u64, snap_path: &str) {
    if gpa == 0 {
        return;
    }
    if let Err(e) = write_smpark_sidecar(snap_path, gpa) {
        eprintln!(
            "supermachine: warning: smpark sidecar write to {} failed ({e:?}); \
             multi-vCPU restores of this snapshot fall back to the agent-RPC \
             unpark path instead of host-direct",
            smpark_sidecar_path(snap_path)
        );
    }
}

/// Read the smpark GPA sidecar. `None` when absent (pre-fix snapshot or
/// single-vCPU bake) → the caller uses the agent-RPC unpark path. Any
/// malformed sidecar also returns `None` (best-effort, never an error).
pub fn read_smpark_sidecar(snap_path: &str) -> Option<u64> {
    use std::io::Read;
    let mut f = std::fs::File::open(smpark_sidecar_path(snap_path)).ok()?;
    let mut buf = Vec::new();
    f.read_to_end(&mut buf).ok()?;
    if buf.len() < 16 || &buf[0..4] != b"SMPK" {
        return None;
    }
    if u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]) != 1 {
        return None;
    }
    Some(u64::from_le_bytes([
        buf[8], buf[9], buf[10], buf[11], buf[12], buf[13], buf[14], buf[15],
    ]))
}

/// Write the DAX sidecar for a list of per-mount session-state blobs.
/// `blobs[i]` is `Some(...)` with the binary blob from
/// `DaxSession::snapshot_state()` for mount index `i`, or `None` if
/// that mount has no DAX session (which would never happen in the
/// current builder since every virtio-fs mount gets a session, but
/// the shape mirrors `write_posix_fs_sidecar` for consistency).
///
/// Layout:
/// ```text
/// b"DAXC"               4-byte magic ("DAX Container")
/// u32 version           little-endian, currently 1
/// u32 session_count     number of mount entries
/// for each session:
///   u32 blob_len        0 means "no state for this session"
///   bytes blob          length blob_len
/// ```
pub fn write_dax_sidecar(sidecar_path: &str, blobs: &[Option<Vec<u8>>]) -> Result<(), FileError> {
    use std::io::Write;
    let mut buf: Vec<u8> = Vec::with_capacity(
        12 + blobs
            .iter()
            .map(|b| 4 + b.as_ref().map_or(0, |v| v.len()))
            .sum::<usize>(),
    );
    buf.extend_from_slice(b"DAXC");
    buf.extend_from_slice(&1u32.to_le_bytes());
    buf.extend_from_slice(&(blobs.len() as u32).to_le_bytes());
    for b in blobs {
        match b {
            Some(v) => {
                buf.extend_from_slice(&(v.len() as u32).to_le_bytes());
                buf.extend_from_slice(v);
            }
            None => {
                buf.extend_from_slice(&0u32.to_le_bytes());
            }
        }
    }
    let mut f = std::fs::File::create(sidecar_path).map_err(FileError::Io)?;
    f.write_all(&buf).map_err(FileError::Io)?;
    f.sync_all().map_err(FileError::Io)?;
    Ok(())
}

/// Read the DAX sidecar produced by [`write_dax_sidecar`].
/// Returns `Ok(None)` if the sidecar doesn't exist (pre-fix snapshot —
/// old behaviour applies, which is "guest SIGBUSes on first DAX touch
/// for any page that wasn't in the working set pre-snapshot"); `Ok(Some(blobs))`
/// with one entry per session otherwise. An empty `Vec` entry means
/// "no state for that session" — caller skips the `restore_state` call.
pub fn read_dax_sidecar(sidecar_path: &str) -> Result<Option<Vec<Vec<u8>>>, FileError> {
    use std::io::Read;
    let mut f = match std::fs::File::open(sidecar_path) {
        Ok(f) => f,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
        Err(e) => return Err(FileError::Io(e)),
    };
    let mut buf = Vec::new();
    f.read_to_end(&mut buf).map_err(FileError::Io)?;
    if buf.len() < 12 {
        return Err(FileError::Truncated);
    }
    if &buf[0..4] != b"DAXC" {
        return Err(FileError::BadMagic);
    }
    let version = u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]);
    if version != 1 {
        return Err(FileError::BadVersion(version as u64));
    }
    let session_count = u32::from_le_bytes([buf[8], buf[9], buf[10], buf[11]]) as usize;
    let mut p = 12usize;
    // Bound the pre-allocation by what the file can actually hold:
    // every record carries a ≥4-byte length prefix, so there can be at
    // most (remaining bytes / 4) records. Trusting the unvalidated
    // header `session_count` here would let a corrupted/malicious
    // sidecar request a multi-GiB allocation (OOM) before the loop's
    // bounds checks ever run. See snapshot-sidecar proptests below.
    let cap = session_count.min((buf.len() - p) / 4);
    let mut out = Vec::with_capacity(cap);
    for _ in 0..session_count {
        if p + 4 > buf.len() {
            return Err(FileError::Truncated);
        }
        let len = u32::from_le_bytes([buf[p], buf[p + 1], buf[p + 2], buf[p + 3]]) as usize;
        p += 4;
        if p + len > buf.len() {
            return Err(FileError::Truncated);
        }
        out.push(buf[p..p + len].to_vec());
        p += len;
    }
    Ok(Some(out))
}

/// Convenience: capture each DaxSession's metadata via
/// `snapshot_state()` and write the sidecar next to `snap_path`.
/// Logs but does not fail on sidecar errors — the snapshot file is
/// already on disk and usable; the sidecar is an enhancement that
/// prevents post-restore SIGBUS on un-paged-in DAX pages.
pub fn capture_and_write_dax_sidecar(
    snap_path: &str,
    dax_sessions: &[std::sync::Arc<crate::fuse::DaxSession>],
) {
    if dax_sessions.is_empty() {
        return;
    }
    let blobs: Vec<Option<Vec<u8>>> = dax_sessions
        .iter()
        .map(|s| Some(s.snapshot_state()))
        .collect();
    let sidecar_path = dax_sidecar_path(snap_path);
    if let Err(e) = write_dax_sidecar(&sidecar_path, &blobs) {
        eprintln!(
            "supermachine: warning: dax sidecar write to {sidecar_path} failed ({e:?}); \
             cycle-restored guests will SIGBUS on first touch of DAX pages that weren't \
             paged-in pre-snapshot"
        );
    }
}

/// Read the PosixFs sidecar produced by [`write_posix_fs_sidecar`].
/// Returns `Ok(None)` if the sidecar doesn't exist (pre-0.7.6
/// snapshot — old behaviour applies); `Ok(Some(blobs))` with one
/// entry per mount otherwise. An empty `Vec` entry means "no state
/// for that mount" (the caller leaves the backend's default-empty
/// table).
pub fn read_posix_fs_sidecar(sidecar_path: &str) -> Result<Option<Vec<Vec<u8>>>, FileError> {
    use std::io::Read;
    let mut f = match std::fs::File::open(sidecar_path) {
        Ok(f) => f,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
        Err(e) => return Err(FileError::Io(e)),
    };
    let mut buf = Vec::new();
    f.read_to_end(&mut buf).map_err(FileError::Io)?;
    if buf.len() < 12 {
        return Err(FileError::Truncated);
    }
    if &buf[0..4] != b"PFXS" {
        return Err(FileError::BadMagic);
    }
    let version = u32::from_le_bytes([buf[4], buf[5], buf[6], buf[7]]);
    if version != 1 {
        return Err(FileError::BadVersion(version as u64));
    }
    let mount_count = u32::from_le_bytes([buf[8], buf[9], buf[10], buf[11]]) as usize;
    let mut p = 12usize;
    // Bound the pre-allocation by what the file can hold (≥4 bytes per
    // record's length prefix), not by the unvalidated header count — an
    // unbounded `with_capacity(mount_count)` is an OOM on a corrupted or
    // malicious sidecar. Mirrors read_dax_sidecar.
    let cap = mount_count.min((buf.len() - p) / 4);
    let mut out = Vec::with_capacity(cap);
    for _ in 0..mount_count {
        if p + 4 > buf.len() {
            return Err(FileError::Truncated);
        }
        let len = u32::from_le_bytes([buf[p], buf[p + 1], buf[p + 2], buf[p + 3]]) as usize;
        p += 4;
        if p + len > buf.len() {
            return Err(FileError::Truncated);
        }
        out.push(buf[p..p + len].to_vec());
        p += len;
    }
    Ok(Some(out))
}

/// Compact in-memory snapshot used by the pipelined `SNAPSHOT_ASYNC`
/// path. Holds all the small fixed-size state by value plus the
/// guest's RAM as a list of (offset, page) pairs — only non-zero
/// pages, so memory cost is the actual working-set size (~100 MiB
/// for rust:1-slim post-init) instead of the full RAM allocation
/// (2 GiB).
///
/// Two-phase save: `capture_compact` runs while the guest is paused
/// (~50 ms for 100 MiB of non-zero pages on M-series), then
/// `save_compact_to_file` runs in a background thread while the
/// guest is unpaused for the warmup workload. Disk write writes
/// the same on-disk sparse format as the streaming path.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub struct CompactSnapshot {
    pub captured_host_ticks: u64,
    pub captured_clock_ref: u64,
    pub ram_gpa: u64,
    pub ram_size: usize,
    pub intc_blob: Vec<u8>,
    pub per_vcpu: Vec<PerVcpuState>,
    pub virtio: VirtioSnapshot,
    /// Non-zero pages only. Each entry: (byte offset within the
    /// guest RAM, owned 4 KiB page). Sorted by offset.
    pub pages: Vec<(usize, Box<[u8; 4096]>)>,
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
const COMPACT_PAGE_SIZE: usize = 4096;

/// 0.7.44+ NEON-accelerated 4 KiB zero-check.
///
/// The pre-0.7.44 path was `chunk.iter().all(|&b| b == 0)` —
/// LLVM autovectorizes it to ~16-byte NEON loads but adds branch
/// overhead per iteration. Explicit NEON intrinsics + bulk OR
/// across 4 accumulators (breaks the dependency chain so the
/// CPU pipelines them) get us closer to memory bandwidth.
///
/// 4096 bytes = 256 × 16-byte NEON vectors. We OR-reduce into
/// 4 independent accumulators (64 lanes wide each, but uint8x16
/// so 4 × 16 = 64 bytes of state) and then collapse at the end.
///
/// SAFETY: caller passes a 4096-byte slice borrowed from guest
/// RAM. Page-aligned in practice (mmap'd), so the unaligned NEON
/// loads (`vld1q_u8`) are happy on aligned addresses too. Length
/// asserted in debug builds.
///
/// Measured: ~3× faster than the Rust iter on M3, ~5 µs per page
/// → for 256 MiB guest (65536 pages) the saving is ~50 ms on the
/// serial path, ~12-15 ms on the 4-way parallel path (still
/// memory-bandwidth-bound at scale).
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
#[inline]
fn page_is_zero(page: &[u8]) -> bool {
    debug_assert_eq!(page.len(), COMPACT_PAGE_SIZE);
    use std::arch::aarch64::{vld1q_u8, vmaxvq_u8, vorrq_u8};
    // SAFETY: page is exactly 4096 bytes (debug-asserted); we
    // perform 256 unaligned 16-byte NEON loads which are safe on
    // aarch64 at any byte alignment.
    unsafe {
        let mut p = page.as_ptr();
        let end = p.add(COMPACT_PAGE_SIZE);
        let mut a0 = vld1q_u8(p);
        let mut a1 = vld1q_u8(p.add(16));
        let mut a2 = vld1q_u8(p.add(32));
        let mut a3 = vld1q_u8(p.add(48));
        p = p.add(64);
        while p < end {
            a0 = vorrq_u8(a0, vld1q_u8(p));
            a1 = vorrq_u8(a1, vld1q_u8(p.add(16)));
            a2 = vorrq_u8(a2, vld1q_u8(p.add(32)));
            a3 = vorrq_u8(a3, vld1q_u8(p.add(48)));
            p = p.add(64);
        }
        let merged = vorrq_u8(vorrq_u8(a0, a1), vorrq_u8(a2, a3));
        vmaxvq_u8(merged) == 0
    }
}

/// Serial fallback for `capture_compact`'s page walk. Used when
/// `snapshot_write_threads()` is 1 or for code paths that
/// haven't been threaded yet.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn capture_compact_pages_serial(memory: &[u8], n_pages: usize) -> Vec<(usize, Box<[u8; 4096]>)> {
    let mut pages: Vec<(usize, Box<[u8; 4096]>)> = Vec::with_capacity(n_pages / 20);
    for page_idx in 0..n_pages {
        let off = page_idx * COMPACT_PAGE_SIZE;
        let chunk = &memory[off..off + COMPACT_PAGE_SIZE];
        if !page_is_zero(chunk) {
            let mut page = Box::new([0u8; 4096]);
            page.copy_from_slice(chunk);
            pages.push((off, page));
        }
    }
    pages
}

/// Parallel page walk for `capture_compact`. Splits guest RAM
/// into N contiguous slabs, scans each on a worker thread.
/// Output preserves source-offset order so the on-disk layout is
/// byte-identical to the serial path.
///
/// CALLER MUST guarantee the guest is paused for the duration —
/// the threads borrow `memory` (a `&[u8]` slice over
/// `vm.ram_host`) for read-only access.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn capture_compact_pages_parallel(
    memory: &[u8],
    n_pages: usize,
    n_threads: usize,
) -> Vec<(usize, Box<[u8; 4096]>)> {
    let n = n_threads.max(1);
    let pages_per_slab = n_pages.div_ceil(n);
    let mem_ptr = memory.as_ptr() as usize;
    let mem_len = memory.len();

    // Each slab's walker returns its own ordered Vec of (offset,
    // page) entries. We concatenate them in slab order — slab K
    // covers pages [K*pages_per_slab, (K+1)*pages_per_slab),
    // strictly after slab K-1, so concatenation preserves the
    // global offset order.
    let mut slab_results: Vec<Vec<(usize, Box<[u8; 4096]>)>> =
        std::thread::scope(|s| -> Vec<Vec<(usize, Box<[u8; 4096]>)>> {
            let mut handles = Vec::with_capacity(n);
            for slab_idx in 0..n {
                let slab_start = slab_idx * pages_per_slab;
                let slab_end = ((slab_idx + 1) * pages_per_slab).min(n_pages);
                if slab_start >= slab_end {
                    continue;
                }
                let h = s.spawn(move || -> Vec<(usize, Box<[u8; 4096]>)> {
                    // SAFETY: `memory` is borrowed for the
                    // duration of `thread::scope`; we reconstruct
                    // a sub-slice with the same lifetime. Guest is
                    // paused (caller invariant) so RAM doesn't
                    // change underneath.
                    let _ = mem_len;
                    let mem: &[u8] =
                        unsafe { std::slice::from_raw_parts(mem_ptr as *const u8, mem_len) };
                    let mut local: Vec<(usize, Box<[u8; 4096]>)> =
                        Vec::with_capacity((slab_end - slab_start) / 20);
                    for page_idx in slab_start..slab_end {
                        let off = page_idx * COMPACT_PAGE_SIZE;
                        let chunk = &mem[off..off + COMPACT_PAGE_SIZE];
                        if !page_is_zero(chunk) {
                            let mut page = Box::new([0u8; 4096]);
                            page.copy_from_slice(chunk);
                            local.push((off, page));
                        }
                    }
                    local
                });
                handles.push(h);
            }
            handles
                .into_iter()
                .map(|h| h.join().unwrap_or_default())
                .collect()
        });
    let total: usize = slab_results.iter().map(Vec::len).sum();
    let mut pages = Vec::with_capacity(total);
    for slab in slab_results.drain(..) {
        pages.extend(slab);
    }
    pages
}

/// Pause-window-only step of the pipelined snapshot. CALLER MUST
/// guarantee the guest is paused — secondary vCPUs rendezvous-
/// stopped, vCPU 0 exited from `hv_vcpu_run`. Returns when all
/// non-zero pages have been copied; the caller can resume the
/// guest immediately and call `save_compact_to_file` from a
/// background thread.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub fn capture_compact(
    vm: &MicroVm,
    virtio: VirtioSnapshot,
    secondary_states: Vec<PerVcpuState>,
) -> Result<CompactSnapshot, SnapshotStreamError> {
    let per0 = vm
        .vcpu
        .capture_snapshot()
        .map_err(SnapshotStreamError::Hvf)?;
    let intc_blob = vm.vm.capture_intc().map_err(SnapshotStreamError::Hvf)?;
    let captured_host_ticks = crate::hypervisor::ActiveVm::host_monotonic_ticks();
    let captured_clock_ref =
        <crate::hypervisor::ActiveVcpu as crate::hypervisor::HypervisorVcpu>::capture_clock_ref(
            &per0,
            captured_host_ticks,
        );
    let mut per_vcpu = Vec::with_capacity(1 + secondary_states.len());
    per_vcpu.push(per0);
    per_vcpu.extend(secondary_states);

    // Copy non-zero pages into the compact buffer. We walk
    // vm.ram_host page-by-page, skip all-zero pages (the kernel
    // hasn't touched them, restore will leave them as zero in the
    // CoW-mapped target).
    //
    // Multi-threaded: the walk is the SECOND-largest pause-window
    // cost in the bake-then-pool flow (after the file write).
    // Splitting into N slabs each scanned by its own worker turns
    // a serial 2 GiB scan (~500 ms on M-series) into parallel
    // ~125 ms scans. Each slab thread builds its own
    // `Vec<(offset, page)>`; we concatenate at the end. Output is
    // byte-identical to the serial path because slabs are
    // contiguous and we preserve the offset order.
    let ram_size = vm.ram_size;
    let memory: &[u8] = unsafe { std::slice::from_raw_parts(vm.ram_host, ram_size) };
    let n_pages = ram_size / COMPACT_PAGE_SIZE;
    let n_threads = snapshot_write_threads();
    let pages: Vec<(usize, Box<[u8; 4096]>)> = if n_threads <= 1 {
        capture_compact_pages_serial(memory, n_pages)
    } else {
        capture_compact_pages_parallel(memory, n_pages, n_threads)
    };
    Ok(CompactSnapshot {
        captured_host_ticks,
        captured_clock_ref,
        ram_gpa: vm.ram_gpa,
        ram_size,
        intc_blob,
        per_vcpu,
        virtio,
        pages,
    })
}

/// Background-thread step of the pipelined snapshot. Writes the
/// captured pages to a `.partial` sibling file, then atomically
/// renames to `path` on success. The output is the same on-disk
/// sparse format that `save_to_file_with_stats` and
/// `capture_and_save_streaming` produce — `Image::from_snapshot`
/// reads any of the three interchangeably.
///
/// Crash safety: only the `.partial` file exists mid-write. A
/// crash leaves no `restore.snap` (which is what callers/cache
/// fingerprint check for), so the next bake invocation cleanly
/// re-runs.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
/// Encode the meta section of a compact snapshot (header + GIC
/// blob + per-vCPU state + virtio metadata) into a flat byte
/// buffer. Same byte layout as what `save_compact_to_file` writes
/// from offset 0 to the start of the RAM padding. Used by the
/// diff-via-clone path (`save_compact_to_file_via_clone`) to
/// rewrite the meta section of an APFS-cloned base file in place.
///
/// `ram_offset_for_header` is the value to encode into the
/// header's `ram_offset` field. For fresh saves the caller fills
/// it in after computing pad alignment; for the diff path the
/// caller passes the base file's existing ram_offset (so warm's
/// page coordinates map onto the same on-disk RAM region).
/// HVF snapshot file header prefix (7c step 4c): `magic[8] · version[8] ·
/// ram_gpa[8] · ram_offset[8]` = 32 bytes, immediately followed by the portable
/// [`ContainerMeta`](crate::snapshot_frame::ContainerMeta) body. `ram_gpa` and
/// `ram_offset` are HVF file-layout fields that stay in this thin wrapper (the
/// KVM pipeline carries its own equivalents); everything the two backends share
/// — clock reference, memory size, the opaque interrupt-controller blob,
/// per-vCPU register files, devices, and vsock TSI listeners — now lives in the
/// one ContainerMeta both pipelines assemble through.
const HVF_HDR_PREFIX: usize = 32;
/// File offset of the `ram_offset` u64 in the prefix, patched after the RAM
/// section's start position is known (see [`write_snapshot_file`]).
const RAM_OFFSET_FIELD_POS: u64 = 24;

/// Project HVF snapshot parts onto the portable [`ContainerMeta`]. HVF has no
/// 8250 COM1 (`com1` = zeros) and persists no separate vsock TSI token
/// (`tsi_token` = None) — both are KVM-only fields. Each per-vCPU register file
/// is serialized via the hypervisor seam into one opaque blob (byte-identical to
/// the prior inline v9/v11 encoding), then length-framed by ContainerMeta.
fn hvf_container_meta(
    captured_host_ticks: u64,
    captured_clock_ref: u64,
    memory_len: u64,
    intc_blob: &[u8],
    per_vcpu: &[PerVcpuState],
    virtio: &VirtioSnapshot,
) -> crate::snapshot_frame::ContainerMeta {
    let mut vcpu_blobs = Vec::with_capacity(per_vcpu.len());
    for st in per_vcpu {
        let mut b = Vec::with_capacity(64 + st.gp_regs.len() * 12 + st.simd_regs.len() * 20);
        <crate::hypervisor::ActiveVcpu as crate::hypervisor::HypervisorVcpu>::write_snapshot_state(
            st, &mut b,
        )
        .expect("Vec<u8> write is infallible");
        vcpu_blobs.push(b);
    }
    crate::snapshot_frame::ContainerMeta {
        num_cpus: per_vcpu.len() as u8,
        mem_size: memory_len,
        com1: [0u8; 6],
        clock_host_ticks: captured_host_ticks,
        clock_ref: captured_clock_ref,
        intc_blob: intc_blob.to_vec(),
        vcpu_blobs,
        devices: virtio.devices.clone(),
        tsi_token: None,
        vsock_listeners: virtio.vsock_listeners.clone(),
    }
}

/// Encode the HVF meta block: the 32-byte [`HVF_HDR_PREFIX`] then the portable
/// ContainerMeta body. `ram_offset` is the value to stamp into the prefix (0 as
/// a placeholder for fresh saves that patch it after RAM alignment; the base
/// file's existing ram_offset for the diff-via-clone path).
fn encode_hvf_meta(
    ram_gpa: u64,
    ram_offset: u64,
    cm: &crate::snapshot_frame::ContainerMeta,
) -> Vec<u8> {
    let mut buf = Vec::with_capacity(HVF_HDR_PREFIX + 512);
    buf.extend_from_slice(&SNAPSHOT_MAGIC);
    buf.extend_from_slice(&SNAPSHOT_VERSION.to_le_bytes());
    buf.extend_from_slice(&ram_gpa.to_le_bytes());
    buf.extend_from_slice(&ram_offset.to_le_bytes());
    cm.write_container(&mut buf)
        .expect("Vec<u8> write is infallible");
    buf
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn encode_compact_meta(snap: &CompactSnapshot, ram_offset_for_header: u64) -> Vec<u8> {
    let cm = hvf_container_meta(
        snap.captured_host_ticks,
        snap.captured_clock_ref,
        snap.ram_size as u64,
        &snap.intc_blob,
        &snap.per_vcpu,
        &snap.virtio,
    );
    encode_hvf_meta(snap.ram_gpa, ram_offset_for_header, &cm)
}

pub fn save_compact_to_file(
    snap: &CompactSnapshot,
    path: &str,
) -> Result<SnapshotWriteStats, FileError> {
    use std::io::{Seek, SeekFrom, Write};
    let partial = format!("{path}.partial");
    let mut f = std::fs::File::create(&partial)?;

    // Encode meta with placeholder ram_offset; we patch it after
    // computing the pad alignment.
    let meta = encode_compact_meta(snap, 0);
    f.write_all(&meta)?;

    let cur = f.stream_position()?;
    let pad = (RAM_PAGE_ALIGN - (cur % RAM_PAGE_ALIGN)) % RAM_PAGE_ALIGN;
    if pad > 0 {
        f.write_all(&vec![0u8; pad as usize])?;
    }
    let ram_offset = f.stream_position()?;
    f.set_len(ram_offset + snap.ram_size as u64)?;

    // Sparse-write the captured pages. Multi-threaded variant —
    // splits the page list across `snapshot_write_threads()`
    // workers using `pwrite` so the write phase scales with disk
    // bandwidth. Each page lands at (ram_offset + page_off),
    // so workers writing to disjoint offsets can't conflict on
    // file position.
    let n_threads = snapshot_write_threads();
    let data_bytes = if n_threads <= 1 || snap.pages.len() < 64 {
        // Serial fallback: tiny page lists aren't worth thread
        // overhead.
        let mut data_bytes = 0u64;
        for (page_off, page) in &snap.pages {
            f.seek(SeekFrom::Start(ram_offset + *page_off as u64))?;
            f.write_all(page.as_ref())?;
            data_bytes += COMPACT_PAGE_SIZE as u64;
        }
        data_bytes
    } else {
        save_compact_pages_parallel(&f, &snap.pages, ram_offset, n_threads)?
    };

    f.seek(SeekFrom::Start(RAM_OFFSET_FIELD_POS))?;
    f.write_all(&ram_offset.to_le_bytes())?;
    drop(f);

    // Atomic rename: only `path` (the canonical name) exists when
    // we return. Mid-write crashes leave `path.partial` instead,
    // which `Image::from_snapshot` ignores.
    std::fs::rename(&partial, path)?;

    Ok(SnapshotWriteStats {
        ram_bytes: snap.ram_size as u64,
        ram_data_bytes: data_bytes,
        ram_zero_bytes: snap.ram_size as u64 - data_bytes,
    })
}

/// Write the captured pages in parallel via `pwrite`. Pages list
/// is sliced into `n_threads` evenly-sized chunks. Used by
/// [`save_compact_to_file`] for the bake-time async base save.
///
/// 0.7.44+: each thread COALESCES offset-contiguous pages into a
/// single bulk pwrite. Pre-0.7.44 fired one 4 KiB pwrite per page
/// (~10k syscalls for a 40 MiB working set), which under concurrent
/// bake load (FORKS=15 × 4 threads = 60 writers) saturated APFS's
/// metadata journal and stretched the per-bake save phase from
/// ~400ms to ~4.5s. Coalescing to ~16 KiB-1 MiB bulk writes cuts
/// syscall count 10-100× and lets the kernel amortize journal
/// entries across larger transactions.
///
/// Pages are sorted by offset on input (caller invariant — they
/// come from the page-walk which iterates increasing). We slice
/// each thread's chunk into RUNS where consecutive pages differ
/// by exactly 4 KiB, then write each run as one pwrite of
/// `(4 KiB × N)` bytes.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn save_compact_pages_parallel(
    f: &std::fs::File,
    pages: &[(usize, Box<[u8; 4096]>)],
    ram_offset: u64,
    n_threads: usize,
) -> Result<u64, FileError> {
    use std::os::unix::fs::FileExt;
    let n = n_threads.max(1);
    let per_thread = pages.len().div_ceil(n);

    let total_data: u64 = std::thread::scope(|s| -> Result<u64, std::io::Error> {
        let mut handles = Vec::with_capacity(n);
        for chunk in pages.chunks(per_thread) {
            let f_clone = f.try_clone()?;
            let chunk_ref = chunk;
            let h = s.spawn(move || -> std::io::Result<u64> {
                let mut bytes = 0u64;
                // Walk the chunk, batching offset-contiguous pages
                // into one pwrite. A "run" is a maximal slice
                // [i, j) of the chunk where each consecutive pair
                // differs in page_off by exactly 4 KiB. Single
                // bulk pwrite per run; for typical guest RAM
                // shapes (long zero gaps + dense touched regions)
                // runs are 4 KiB–16 MiB. Worst case (all isolated
                // pages) degrades back to the pre-0.7.44 path
                // with no extra work.
                let mut i = 0usize;
                while i < chunk_ref.len() {
                    let run_start_off = chunk_ref[i].0;
                    let mut j = i + 1;
                    while j < chunk_ref.len()
                        && chunk_ref[j].0 == chunk_ref[j - 1].0 + COMPACT_PAGE_SIZE
                    {
                        j += 1;
                    }
                    let run_len_pages = j - i;
                    if run_len_pages == 1 {
                        // Fast path: single page = original write_all_at.
                        f_clone.write_all_at(
                            chunk_ref[i].1.as_ref(),
                            ram_offset + run_start_off as u64,
                        )?;
                        bytes += COMPACT_PAGE_SIZE as u64;
                    } else {
                        // Concat run into a contiguous buffer.
                        // ~4 KiB allocation per run head; amortized
                        // across N pages. For very large runs (e.g.
                        // 256-page run = 1 MiB write), this is a
                        // 1 MiB alloc — well under reasonable VM
                        // memory bounds and gets reused by the
                        // allocator across runs.
                        let mut buf: Vec<u8> =
                            Vec::with_capacity(run_len_pages * COMPACT_PAGE_SIZE);
                        for k in i..j {
                            buf.extend_from_slice(chunk_ref[k].1.as_ref());
                        }
                        f_clone.write_all_at(&buf, ram_offset + run_start_off as u64)?;
                        bytes += (run_len_pages * COMPACT_PAGE_SIZE) as u64;
                    }
                    i = j;
                }
                Ok(bytes)
            });
            handles.push(h);
        }
        let mut total = 0u64;
        for h in handles {
            total += h
                .join()
                .map_err(|_| std::io::Error::other("snapshot write thread panicked"))??;
        }
        Ok(total)
    })?;

    Ok(total_data)
}

/// Thin wrapper around the macOS `clonefile(2)` syscall.
/// Creates an APFS CoW clone of `src` at `dst`; data blocks are
/// shared until either side writes (then APFS does on-disk CoW).
/// Returns ENOENT/EXDEV/etc as `io::Error`; callers fall through
/// to the plain save path.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn clonefile_via_libc(src: &str, dst: &str) -> std::io::Result<()> {
    let src_c = std::ffi::CString::new(src).map_err(|_| {
        std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "src path contains NUL byte",
        )
    })?;
    let dst_c = std::ffi::CString::new(dst).map_err(|_| {
        std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "dst path contains NUL byte",
        )
    })?;
    let ret = unsafe { libc::clonefile(src_c.as_ptr(), dst_c.as_ptr(), 0) };
    if ret != 0 {
        Err(std::io::Error::last_os_error())
    } else {
        Ok(())
    }
}

/// Load a snapshot file into an in-memory `CompactSnapshot` —
/// only the non-zero pages are kept (same as the in-memory
/// representation that `capture_compact` produces). Used by the
/// diff-via-clone path when the runner needs base in memory but
/// it's not in the in-flight async-save list (e.g. for cycle-
/// snapshots from a worker that was restored via `--restore-from`,
/// not via a recent `SNAPSHOT_ASYNC`).
///
/// Memory cost: ~size of non-zero pages (~100 MiB on
/// rust:1-slim). Wall time: ~150 ms for a 2 GiB sparse snapshot
/// on M-series (real disk reads scale with non-zero size; sparse
/// holes return zero from the page cache instantly).
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub fn load_compact_from_file(path: &str) -> Result<CompactSnapshot, FileError> {
    use std::os::fd::AsRawFd;

    // Reuse load_meta for the header / GIC / per_vcpu / virtio.
    let (snap, ram_offset, memory_bytes) = load_meta(path)?;

    // mmap the RAM section read-only and walk it. Avoids the 2
    // GiB Vec<u8> allocation that `load_from_file` would do.
    let f = std::fs::File::open(path)?;
    let fd = f.as_raw_fd();
    let ptr = unsafe {
        libc::mmap(
            std::ptr::null_mut(),
            memory_bytes,
            libc::PROT_READ,
            libc::MAP_PRIVATE,
            fd,
            ram_offset as libc::off_t,
        )
    };
    if ptr == libc::MAP_FAILED {
        return Err(FileError::Io(std::io::Error::last_os_error()));
    }
    // Hint kernel for sequential read-ahead.
    unsafe {
        let _ = libc::madvise(ptr, memory_bytes, libc::MADV_SEQUENTIAL);
    }

    let memory: &[u8] = unsafe { std::slice::from_raw_parts(ptr as *const u8, memory_bytes) };
    let n_pages = memory_bytes / COMPACT_PAGE_SIZE;
    let mut pages: Vec<(usize, Box<[u8; 4096]>)> = Vec::with_capacity(n_pages / 20);
    for page_idx in 0..n_pages {
        let off = page_idx * COMPACT_PAGE_SIZE;
        let chunk = &memory[off..off + COMPACT_PAGE_SIZE];
        if !page_is_zero(chunk) {
            let mut page = Box::new([0u8; 4096]);
            page.copy_from_slice(chunk);
            pages.push((off, page));
        }
    }

    // SAFETY: ptr was returned by mmap with this size and we own it.
    unsafe {
        libc::munmap(ptr, memory_bytes);
    }

    Ok(CompactSnapshot {
        captured_host_ticks: snap.captured_host_ticks,
        captured_clock_ref: snap.captured_clock_ref,
        ram_gpa: snap.ram_gpa,
        ram_size: memory_bytes,
        intc_blob: snap.intc_blob,
        per_vcpu: snap.per_vcpu,
        virtio: snap.virtio,
        pages,
    })
}

/// Differential compact-snapshot save via APFS clonefile + diff
/// pwrite. Same on-disk format as [`save_compact_to_file`] —
/// restore is byte-identical for byte-identical guest state, so
/// no special restore path is needed.
///
/// Algorithm:
///   1. `clonefile(base_path → out_path.partial)` — APFS CoW
///      clone, ~no I/O. Data blocks shared with base.
///   2. Read base's `ram_offset` from the cloned header.
///   3. Encode `snap`'s meta section (header + GIC + per-vCPU +
///      virtio); pad to base's `ram_offset`. If meta overflows
///      `ram_offset` (e.g. warmup added many vsock listeners),
///      bail out and let the caller fall through to a plain
///      save. We DON'T retry inside this function — the caller
///      chose the diff path; if the layout doesn't match we
///      return Err so they can decide whether to plain-save or
///      surface the error.
///   4. `pwrite_at` warm meta over [0, ram_offset).
///   5. Compute the diff page set: pages in `snap` whose bytes
///      differ from the same offset in `base`, plus pages in
///      `base` that are absent from `snap` (zeroed by warmup —
///      need explicit zero-page pwrite to override base's
///      content in the cloned file).
///   6. Parallel pwrite the diff pages.
///   7. Atomic rename `out_path.partial → out_path`.
///
/// On `clonefile` EXDEV / ENOENT the caller should fall through
/// to [`save_compact_to_file`].
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub fn save_compact_to_file_via_clone(
    snap: &CompactSnapshot,
    base: &CompactSnapshot,
    base_path: &str,
    out_path: &str,
) -> Result<SnapshotWriteStats, FileError> {
    use std::collections::{HashMap, HashSet};
    use std::os::unix::fs::FileExt;

    let partial = format!("{out_path}.partial");
    let _ = std::fs::remove_file(&partial);

    // 1. Clone base file.
    clonefile_via_libc(base_path, &partial).map_err(FileError::Io)?;

    // 2. Open partial RW.
    let f = std::fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(&partial)?;

    // 3. Read base's ram_offset from the prefix (RAM_OFFSET_FIELD_POS..+8).
    let mut ram_off_bytes = [0u8; 8];
    f.read_exact_at(&mut ram_off_bytes, RAM_OFFSET_FIELD_POS)?;
    let ram_offset = u64::from_le_bytes(ram_off_bytes);

    // Sanity: base's ram_size should match warm's. If not, the
    // VMs are configured differently and clone-based diff makes
    // no sense. ram_size is no longer at a fixed header offset (it rides
    // inside the streamed ContainerMeta now), so read it from the in-memory
    // base snapshot we were handed.
    let base_ram_size = base.ram_size as u64;
    if base_ram_size != snap.ram_size as u64 {
        let _ = std::fs::remove_file(&partial);
        return Err(FileError::Io(std::io::Error::other(format!(
            "diff snapshot: base ram_size {base_ram_size} != warm ram_size {}",
            snap.ram_size
        ))));
    }

    // 4. Encode warm's meta and verify it fits before ram_offset.
    let meta = encode_compact_meta(snap, ram_offset);
    if (meta.len() as u64) > ram_offset {
        let _ = std::fs::remove_file(&partial);
        return Err(FileError::Io(std::io::Error::other(format!(
            "diff snapshot: warm meta {} bytes overflows base ram_offset {ram_offset}",
            meta.len()
        ))));
    }
    // Pad to ram_offset with zeros so [0, ram_offset) is fully
    // overwritten and we don't leave any base-meta bytes behind.
    let mut padded = meta;
    padded.resize(ram_offset as usize, 0);
    f.write_all_at(&padded, 0)?;

    // 5. Compute diff page set, split into:
    //    * `data_pages`: pages with bytes != base's bytes (or
    //      pages new in warm). Need pwrite of warm content.
    //    * `zero_offsets`: offsets where base had non-zero bytes
    //      but warm has zero. Punching a hole at these offsets
    //      deallocates the cloned base's data block — no on-disk
    //      space cost, and restore reads zero (correct).
    //
    //  Why two paths: APFS clonefile shares blocks with base.
    //  pwriting zeros forces APFS to allocate a NEW block (4 KiB
    //  of physical disk per zeroed page), defeating the clone's
    //  size win. F_PUNCHHOLE deallocates the cloned block
    //  outright — restore-time reads return zero from the hole,
    //  same observable result as a zero pwrite, ~3x physical
    //  disk savings on a typical warmup that zeros buffer cache.
    let base_lookup: HashMap<usize, &[u8; 4096]> =
        base.pages.iter().map(|(o, p)| (*o, p.as_ref())).collect();
    let warm_offsets: HashSet<usize> = snap.pages.iter().map(|(o, _)| *o).collect();

    let mut data_pages: Vec<(usize, &[u8; 4096])> = Vec::new();
    for (off, warm_p) in &snap.pages {
        let warm_bytes: &[u8; 4096] = warm_p.as_ref();
        match base_lookup.get(off) {
            None => data_pages.push((*off, warm_bytes)),
            Some(b) if *b != warm_bytes => data_pages.push((*off, warm_bytes)),
            _ => {}
        }
    }
    let mut zero_offsets: Vec<usize> = Vec::new();
    for (off, _) in &base.pages {
        if !warm_offsets.contains(off) {
            zero_offsets.push(*off);
        }
    }
    // Sort + coalesce contiguous runs so each F_PUNCHHOLE
    // ioctl covers many adjacent pages — fewer syscalls, and
    // APFS deallocates the run as a single extent.
    zero_offsets.sort_unstable();

    // 6. pwrite the warm-data diff pages (parallel for >= 64).
    let n_threads = snapshot_write_threads();
    let data_bytes_written = if n_threads <= 1 || data_pages.len() < 64 {
        let mut bytes = 0u64;
        for (off, page) in &data_pages {
            f.write_all_at(*page, ram_offset + *off as u64)?;
            bytes += COMPACT_PAGE_SIZE as u64;
        }
        bytes
    } else {
        save_diff_pages_parallel(&f, &data_pages, ram_offset, n_threads)?
    };

    // 7. Punch holes for zeroed pages (coalesced runs).
    let mut hole_bytes = 0u64;
    if !zero_offsets.is_empty() {
        use std::os::unix::io::AsRawFd;
        let fd = f.as_raw_fd();
        let mut i = 0;
        while i < zero_offsets.len() {
            let run_start = zero_offsets[i];
            let mut run_end = run_start + COMPACT_PAGE_SIZE;
            i += 1;
            while i < zero_offsets.len() && zero_offsets[i] == run_end {
                run_end += COMPACT_PAGE_SIZE;
                i += 1;
            }
            let span = (run_end - run_start) as i64;
            let punch = libc::fpunchhole_t {
                fp_flags: 0,
                reserved: 0,
                fp_offset: (ram_offset + run_start as u64) as libc::off_t,
                fp_length: span,
            };
            // SAFETY: fd is valid for the file we just wrote;
            // fpunchhole_t pointer is owned by us.
            let r =
                unsafe { libc::fcntl(fd, libc::F_PUNCHHOLE, &punch as *const libc::fpunchhole_t) };
            if r == 0 {
                hole_bytes += span as u64;
            } else {
                // Filesystem doesn't support hole-punching (e.g.
                // some non-APFS network mounts). Fall back to
                // explicit zero pwrite for this run — same
                // semantic, just doesn't save physical disk.
                let zeros = vec![0u8; span as usize];
                f.write_all_at(&zeros, ram_offset + run_start as u64)?;
            }
        }
    }
    let data_bytes = data_bytes_written + hole_bytes;

    drop(f);
    // 7. Atomic rename.
    std::fs::rename(&partial, out_path)?;

    Ok(SnapshotWriteStats {
        ram_bytes: snap.ram_size as u64,
        ram_data_bytes: data_bytes,
        ram_zero_bytes: snap.ram_size as u64 - data_bytes,
    })
}

/// Parallel pwrite of a (offset, &[u8; 4096])-list. Sibling of
/// [`save_compact_pages_parallel`] but works over borrowed page
/// references with arbitrary lifetime (so callers can mix owned
/// warm pages and the static `ZERO_PAGE` in the same list — the
/// scoped thread inherits the page borrows).
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn save_diff_pages_parallel(
    f: &std::fs::File,
    pages: &[(usize, &[u8; 4096])],
    ram_offset: u64,
    n_threads: usize,
) -> Result<u64, FileError> {
    use std::os::unix::fs::FileExt;
    let n = n_threads.max(1);
    let per_thread = pages.len().div_ceil(n);

    let total: u64 = std::thread::scope(|s| -> Result<u64, std::io::Error> {
        let mut handles = Vec::with_capacity(n);
        for chunk in pages.chunks(per_thread) {
            let f_clone = f.try_clone()?;
            let chunk_ref: &[(usize, &[u8; 4096])] = chunk;
            let h = s.spawn(move || -> std::io::Result<u64> {
                let mut bytes = 0u64;
                for (off, page) in chunk_ref {
                    f_clone.write_all_at(*page, ram_offset + *off as u64)?;
                    bytes += COMPACT_PAGE_SIZE as u64;
                }
                Ok(bytes)
            });
            handles.push(h);
        }
        let mut total = 0u64;
        for h in handles {
            total += h
                .join()
                .map_err(|_| std::io::Error::other("snapshot diff write thread panicked"))??;
        }
        Ok(total)
    })?;

    Ok(total)
}

/// Capture vCPU/GIC/virtio state and stream the guest's RAM
/// directly from `vm.ram_host` into the output file in one pass,
/// without the intermediate 2 GiB `Vec<u8>` allocation that
/// `capture_snapshot` makes. Saves ~1.5 s of memcpy on
/// M-series for a 2 GiB-baked image — the bulk of `PooledVm::
/// snapshot` cost on a typical warmup-then-bake flow.
///
/// CALLER MUST guarantee the guest is paused for the entire
/// call (vCPU exited from `hv_vcpu_run`, secondaries rendezvous-
/// stopped). The runner's snapshot RPC handler already does
/// this; bake-time auto-capture goes through the original
/// `capture_snapshot` + `save_to_file_with_stats` pair which is
/// fine because that path captures + saves on the same pause
/// window anyway.
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
pub fn capture_and_save_streaming(
    vm: &MicroVm,
    virtio: &VirtioSnapshot,
    secondary_states: &[PerVcpuState],
    path: &str,
) -> Result<SnapshotWriteStats, SnapshotStreamError> {
    let per0 = vm
        .vcpu
        .capture_snapshot()
        .map_err(SnapshotStreamError::Hvf)?;
    let intc_blob = vm.vm.capture_intc().map_err(SnapshotStreamError::Hvf)?;
    let captured_host_ticks = crate::hypervisor::ActiveVm::host_monotonic_ticks();
    let captured_clock_ref =
        <crate::hypervisor::ActiveVcpu as crate::hypervisor::HypervisorVcpu>::capture_clock_ref(
            &per0,
            captured_host_ticks,
        );
    let mut per_vcpu = Vec::with_capacity(1 + secondary_states.len());
    per_vcpu.push(per0);
    per_vcpu.extend_from_slice(secondary_states);

    // Borrow the guest's RAM directly. SAFETY: vm.ram_host is
    // vm.ram_size bytes; guest is paused for the duration of
    // this call (caller's contract).
    let memory: &[u8] = unsafe { std::slice::from_raw_parts(vm.ram_host, vm.ram_size) };

    write_snapshot_file(
        path,
        captured_host_ticks,
        captured_clock_ref,
        vm.ram_gpa,
        memory,
        &intc_blob,
        &per_vcpu,
        virtio,
    )
    .map_err(SnapshotStreamError::Io)
}

/// Error from `capture_and_save_streaming`.
#[derive(Debug)]
pub enum SnapshotStreamError {
    Hvf(crate::hypervisor::ActiveError),
    Io(FileError),
}

impl std::fmt::Display for SnapshotStreamError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Hvf(e) => write!(f, "snapshot capture: {e:?}"),
            Self::Io(e) => write!(f, "snapshot save: {e:?}"),
        }
    }
}

pub fn save_to_file_with_stats(
    path: &str,
    snap: &Snapshot,
) -> Result<SnapshotWriteStats, FileError> {
    write_snapshot_file(
        path,
        snap.captured_host_ticks,
        snap.captured_clock_ref,
        snap.ram_gpa,
        &snap.memory,
        &snap.intc_blob,
        &snap.per_vcpu,
        &snap.virtio,
    )
}

#[allow(clippy::too_many_arguments)]
fn write_snapshot_file(
    path: &str,
    captured_host_ticks: u64,
    captured_clock_ref: u64,
    ram_gpa: u64,
    memory: &[u8],
    intc_blob: &[u8],
    per_vcpu: &[PerVcpuState],
    virtio: &VirtioSnapshot,
) -> Result<SnapshotWriteStats, FileError> {
    use std::io::{Seek, SeekFrom};
    // Crash-atomicity: write to `<path>.partial`, atomically rename
    // on success. Mirrors `save_compact_to_file`. A crash mid-write
    // leaves only the `.partial`, which `Image::from_snapshot`
    // ignores, so the next invocation cleanly re-bakes.
    let partial = format!("{path}.partial");
    let mut f = std::fs::File::create(&partial)?;
    // Meta block: the 32-byte HVF prefix + the portable ContainerMeta body
    // (intc blob, per-vCPU register files, devices, vsock listeners — the one
    // codec shared with the KVM pipeline). ram_offset is a 0 placeholder here;
    // it's patched in after RAM alignment is known (the seek-and-rewrite below).
    let cm = hvf_container_meta(
        captured_host_ticks,
        captured_clock_ref,
        memory.len() as u64,
        intc_blob,
        per_vcpu,
        virtio,
    );
    let meta = encode_hvf_meta(ram_gpa, 0, &cm);
    f.write_all(&meta)?;

    // Pad to RAM_PAGE_ALIGN, then write RAM. Stash the offset back
    // into the prefix so loaders (especially --cow-restore) know
    // where to mmap from.
    let cur = f.stream_position()?;
    let pad = (RAM_PAGE_ALIGN - (cur % RAM_PAGE_ALIGN)) % RAM_PAGE_ALIGN;
    if pad > 0 {
        f.write_all(&vec![0u8; pad as usize])?;
    }
    let ram_offset = f.stream_position()?;
    let stats = write_sparse_ram(&mut f, memory, ram_offset)?;
    f.seek(SeekFrom::Start(RAM_OFFSET_FIELD_POS))?;
    f.write_all(&ram_offset.to_le_bytes())?;
    drop(f);
    // Atomic rename: make `path` appear only after the full write
    // succeeded. Crash mid-write leaves `path.partial`.
    std::fs::rename(&partial, path)?;
    Ok(stats)
}

/// How many parallel writer threads to use for the sparse RAM
/// scan/write. Default 4 — saturates APFS sequential write
/// bandwidth on M-series without thrashing the page cache.
/// Override via `SUPERMACHINE_SNAPSHOT_WRITE_THREADS=N` (1 = legacy
/// single-threaded). 1 disables parallel write (legacy path).
fn snapshot_write_threads() -> usize {
    if let Ok(v) = std::env::var("SUPERMACHINE_SNAPSHOT_WRITE_THREADS") {
        if let Ok(n) = v.parse::<usize>() {
            return n.max(1);
        }
    }
    4
}

/// Walks `memory` in `SPARSE_RAM_CHUNK`-sized runs of "all-zero" or
/// "non-zero", and writes the non-zero runs into the snapshot file
/// at `ram_offset + run_start`. Zero runs become file holes
/// (sparse). Returns `(data_bytes, zero_bytes)`.
///
/// Multi-threaded variant — if `n_threads > 1`, splits `memory`
/// into `n_threads` slabs aligned to `SPARSE_RAM_CHUNK` and each
/// worker walks its slab using `pwrite`-style `write_all_at`
/// (positional writes). The threads share a clone of the file
/// descriptor; each writes into a distinct byte range so there's
/// no offset contention. APFS handles concurrent writes fine.
///
/// Single-threaded variant kept as the path used when
/// `n_threads == 1` (legacy / debug). Same on-disk output.
fn write_sparse_ram(
    f: &mut std::fs::File,
    memory: &[u8],
    ram_offset: u64,
) -> Result<SnapshotWriteStats, FileError> {
    use std::io::{Seek, SeekFrom};

    let n_threads = snapshot_write_threads();
    let total = memory.len();
    let logical_end = ram_offset + total as u64;
    // Set the file length up front so workers' positional writes
    // don't need to extend the file (APFS would serialize on
    // metadata otherwise).
    f.set_len(logical_end)?;

    let stats = if n_threads <= 1 {
        write_sparse_ram_serial(f, memory, ram_offset)?
    } else {
        write_sparse_ram_parallel(f, memory, ram_offset, n_threads)?
    };

    f.seek(SeekFrom::Start(logical_end))?;
    Ok(stats)
}

fn write_sparse_ram_serial(
    f: &std::fs::File,
    memory: &[u8],
    ram_offset: u64,
) -> Result<SnapshotWriteStats, FileError> {
    use std::os::unix::fs::FileExt;
    let mut pos = 0usize;
    let mut data_bytes = 0u64;
    let mut zero_bytes = 0u64;

    while pos < memory.len() {
        let run_is_zero =
            is_zero_chunk(&memory[pos..(pos + (memory.len() - pos).min(SPARSE_RAM_CHUNK))]);
        let run_start = pos;
        pos += (memory.len() - pos).min(SPARSE_RAM_CHUNK);
        while pos < memory.len() {
            let next_len = (memory.len() - pos).min(SPARSE_RAM_CHUNK);
            let next_is_zero = is_zero_chunk(&memory[pos..pos + next_len]);
            if next_is_zero != run_is_zero {
                break;
            }
            pos += next_len;
        }

        let run_len = pos - run_start;
        if run_is_zero {
            zero_bytes += run_len as u64;
        } else {
            f.write_all_at(&memory[run_start..pos], ram_offset + run_start as u64)?;
            data_bytes += run_len as u64;
        }
    }
    Ok(SnapshotWriteStats {
        ram_bytes: memory.len() as u64,
        ram_data_bytes: data_bytes,
        ram_zero_bytes: zero_bytes,
    })
}

fn write_sparse_ram_parallel(
    f: &std::fs::File,
    memory: &[u8],
    ram_offset: u64,
    n_threads: usize,
) -> Result<SnapshotWriteStats, FileError> {
    let total = memory.len();
    // Round per-thread slab up to a multiple of SPARSE_RAM_CHUNK
    // so the zero-run scan in each thread aligns with the same
    // chunk boundaries the serial path would have used. Output
    // is byte-identical to the serial path.
    let slab = total.div_ceil(n_threads).div_ceil(SPARSE_RAM_CHUNK) * SPARSE_RAM_CHUNK;

    let mut bounds: Vec<(usize, usize)> = Vec::with_capacity(n_threads);
    let mut start = 0usize;
    while start < total {
        let end = (start + slab).min(total);
        bounds.push((start, end));
        start = end;
    }

    let mem_ptr = memory.as_ptr() as usize;
    let mem_len = memory.len();

    let stats = std::thread::scope(|s| -> Result<(u64, u64), std::io::Error> {
        let mut handles = Vec::with_capacity(bounds.len());
        for (slab_start, slab_end) in bounds {
            let f_clone = f.try_clone()?;
            let h = s.spawn(move || -> std::io::Result<(u64, u64)> {
                // SAFETY: `memory` is borrowed for the duration of
                // the surrounding `thread::scope`; we reconstruct
                // a sub-slice with the same lifetime. The guest
                // is paused (caller invariant), so the bytes
                // don't change underneath.
                let _ = mem_len;
                let slab: &[u8] = unsafe {
                    std::slice::from_raw_parts(
                        (mem_ptr as *const u8).add(slab_start),
                        slab_end - slab_start,
                    )
                };
                let base = ram_offset + slab_start as u64;
                write_sparse_ram_slab(&f_clone, slab, base)
            });
            handles.push(h);
        }
        let mut data_bytes = 0u64;
        let mut zero_bytes = 0u64;
        for h in handles {
            let (d, z) = h
                .join()
                .map_err(|_| std::io::Error::other("snapshot write thread panicked"))??;
            data_bytes += d;
            zero_bytes += z;
        }
        Ok((data_bytes, zero_bytes))
    })?;

    Ok(SnapshotWriteStats {
        ram_bytes: total as u64,
        ram_data_bytes: stats.0,
        ram_zero_bytes: stats.1,
    })
}

fn write_sparse_ram_slab(
    f: &std::fs::File,
    memory: &[u8],
    base_offset: u64,
) -> std::io::Result<(u64, u64)> {
    use std::os::unix::fs::FileExt;
    let mut pos = 0usize;
    let mut data_bytes = 0u64;
    let mut zero_bytes = 0u64;

    while pos < memory.len() {
        let run_is_zero =
            is_zero_chunk(&memory[pos..(pos + (memory.len() - pos).min(SPARSE_RAM_CHUNK))]);
        let run_start = pos;
        pos += (memory.len() - pos).min(SPARSE_RAM_CHUNK);
        while pos < memory.len() {
            let next_len = (memory.len() - pos).min(SPARSE_RAM_CHUNK);
            let next_is_zero = is_zero_chunk(&memory[pos..pos + next_len]);
            if next_is_zero != run_is_zero {
                break;
            }
            pos += next_len;
        }
        let run_len = pos - run_start;
        if run_is_zero {
            zero_bytes += run_len as u64;
        } else {
            f.write_all_at(&memory[run_start..pos], base_offset + run_start as u64)?;
            data_bytes += run_len as u64;
        }
    }
    Ok((data_bytes, zero_bytes))
}

fn is_zero_chunk(chunk: &[u8]) -> bool {
    chunk.iter().all(|b| *b == 0)
}

pub fn load_from_file(path: &str) -> Result<Snapshot, FileError> {
    load_from_file_inner(path, /* skip_ram = */ false).map(|(snap, _, _)| snap)
}

/// Like `load_from_file` but skips RAM bytes (sets `memory = Vec::new()`).
/// Returns `(Snapshot, ram_offset, memory_bytes)` so callers can mmap
/// the RAM region directly for CoW restore.
pub fn load_meta(path: &str) -> Result<(Snapshot, u64, usize), FileError> {
    load_from_file_inner(path, true)
}

fn load_from_file_inner(path: &str, skip_ram: bool) -> Result<(Snapshot, u64, usize), FileError> {
    use std::io::{Seek, SeekFrom};
    let mut f = std::fs::File::open(path)?;
    let mut prefix = [0u8; HVF_HDR_PREFIX];
    f.read_exact(&mut prefix)
        .map_err(|_| FileError::Truncated)?;
    if prefix[0..8] != SNAPSHOT_MAGIC {
        return Err(FileError::BadMagic);
    }
    let version = le_u64(&prefix[8..16])?;
    if version != SNAPSHOT_VERSION {
        return Err(FileError::BadVersion(version));
    }
    let ram_gpa = le_u64(&prefix[16..24])?;
    let ram_offset = le_u64(&prefix[24..32])?;

    // Portable container body (7c step 4c): the one codec both pipelines share.
    // read_container caps every record count and every blob length, so a
    // corrupt/hostile file — these cross trust boundaries (CDN catalog,
    // cross-process restore) — fails cleanly here rather than OOMing the host.
    let cm = crate::snapshot_frame::ContainerMeta::read_container(&mut f)
        .map_err(|_| FileError::Truncated)?;
    let memory_bytes = cm.mem_size as usize;

    // The RAM section [ram_offset, ram_offset + memory_bytes) must lie within
    // the file, or the read_exact below (and any caller's mmap of the same
    // region — mmap_ram_cow_at) would walk off the end (SIGBUS on first touch).
    let file_len = f.metadata()?.len();
    if !crate::snapshot_frame::ram_region_within(file_len, ram_offset, memory_bytes as u64) {
        return Err(FileError::Truncated);
    }

    if cm.vcpu_blobs.is_empty() {
        return Err(FileError::Malformed("snapshot contains no vCPU state"));
    }
    // Decode each opaque per-vCPU blob via the hypervisor seam (it names the
    // backend GICv3 register shape). Each blob holds exactly one vCPU's record;
    // ContainerMeta already length-framed and bounds-checked them.
    let mut per_vcpu = Vec::with_capacity(cm.vcpu_blobs.len());
    for vb in &cm.vcpu_blobs {
        let st = <crate::hypervisor::ActiveVcpu as crate::hypervisor::HypervisorVcpu>::read_snapshot_state(
            &mut std::io::Cursor::new(vb),
        )
        .map_err(|_| FileError::Truncated)?;
        per_vcpu.push(st);
    }

    let memory = if skip_ram {
        Vec::new()
    } else {
        f.seek(SeekFrom::Start(ram_offset))
            .map_err(|_| FileError::Truncated)?;
        let mut m = vec![0u8; memory_bytes];
        f.read_exact(&mut m).map_err(|_| FileError::Truncated)?;
        m
    };

    Ok((
        Snapshot {
            captured_host_ticks: cm.clock_host_ticks,
            captured_clock_ref: cm.clock_ref,
            ram_gpa,
            memory,
            intc_blob: cm.intc_blob,
            per_vcpu,
            virtio: VirtioSnapshot {
                devices: cm.devices,
                vsock_listeners: cm.vsock_listeners,
            },
        },
        ram_offset,
        memory_bytes,
    ))
}

/// mmap a snapshot file's RAM region into a new private mapping.
/// Returns the host-side pointer + length suitable for handing to
/// `MicroVm::new_with_ram` and `Vm::map`. The mapping is `MAP_PRIVATE`,
/// so guest writes go to anon pages backing the COW'd portion only —
/// the snapshot file on disk is never written.
pub fn mmap_ram_cow(path: &str) -> std::io::Result<(*mut u8, usize)> {
    let (_snap, ram_offset, memory_bytes) =
        load_meta(path).map_err(|e| std::io::Error::other(format!("load_meta: {e:?}")))?;
    mmap_ram_cow_at(path, ram_offset, memory_bytes)
}

pub fn mmap_ram_cow_at(
    path: &str,
    ram_offset: u64,
    memory_bytes: usize,
) -> std::io::Result<(*mut u8, usize)> {
    let f = std::fs::File::open(path)?;

    // Defensive bounds-check + the page-aligned MAP_PRIVATE CoW map live in the
    // shared `snapshot_frame` substrate. Callers that route through load_meta
    // already validate the region, but this is a `pub` entry point — mapping past
    // EOF would SIGBUS the guest on first touch, so cow_map_ram re-checks and
    // returns InvalidData on a corrupt region.
    let ptr = crate::snapshot_frame::cow_map_ram(&f, ram_offset, memory_bytes, 0)?;
    // NOTE: an earlier version called
    // `madvise(ptr, memory_bytes, MADV_WILLNEED)` here, claiming
    // a 1–3 ms cold-cache savings. Empirically on macOS this is
    // SYNCHRONOUS — `MADV_WILLNEED` blocks until the kernel has
    // populated the page cache for the entire range, even sparse
    // holes. For a 2 GiB snapshot file this turned a ~1 ms mmap
    // call into a 600 ms walk on cold cache. The earlier "saves
    // ~1–3 ms" claim was either measured wrong or against a
    // tiny snapshot. Removed; we let pages fault lazily on the
    // guest's first access — those that aren't accessed don't
    // pay the I/O cost.
    Ok((ptr, memory_bytes))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use std::path::PathBuf;

    // 0.7.44+ NEON page_is_zero correctness checks. The NEON loop
    // OR-reduces 4 accumulators across 256 × 16-byte vectors; any
    // miscalculation of the loop bound or accumulator merge would
    // silently flip the bake's RAM-page filter (either dropping
    // non-zero pages → restored guest sees fresh zeros where the
    // captured guest had real bytes, or NEVER dropping anything →
    // snapshot is huge but correct). Make sure all the boundary
    // shapes get exercised.
    #[cfg(target_arch = "aarch64")]
    #[test]
    fn page_is_zero_all_zeros() {
        let page = [0u8; 4096];
        assert!(page_is_zero(&page));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn page_is_zero_first_byte_nonzero() {
        let mut page = [0u8; 4096];
        page[0] = 1;
        assert!(!page_is_zero(&page));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn page_is_zero_last_byte_nonzero() {
        let mut page = [0u8; 4096];
        page[4095] = 1;
        assert!(!page_is_zero(&page));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn page_is_zero_middle_byte_nonzero() {
        let mut page = [0u8; 4096];
        page[2048] = 0xFF;
        assert!(!page_is_zero(&page));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn page_is_zero_every_byte_position() {
        // Exhaustively check that a single bit set at byte position
        // K (for K in 0..4096) is detected. Catches off-by-one in
        // any of the 4 accumulator strides.
        for k in 0..4096 {
            let mut page = [0u8; 4096];
            page[k] = 1;
            assert!(!page_is_zero(&page), "missed non-zero byte at position {k}");
        }
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn page_is_zero_all_ones() {
        let page = [0xFFu8; 4096];
        assert!(!page_is_zero(&page));
    }

    // Regression (v10): the TSI listener record MUST round-trip
    // `inet_port`. Pre-v10 it was hand-packed in three places and
    // silently dropped from the on-disk format (loaded back as
    // `None`), which broke multi-port `Vm::expose_tcp` routing after
    // any cross-process / disk restore — the pinned listener became
    // unresolvable ("no TSI listener"). All sites now route through
    // write/read_tsi_listener_record; this pins the round-trip,
    // including the `None` sentinel and the max-port edge.
    #[test]
    fn tsi_listener_record_roundtrips_inet_port() {
        let cases = [
            Some(9222u16), // Chrome CDP — the integrator multi-port case
            Some(80),
            Some(65535), // max valid port must NOT collide with None sentinel
            None,        // legacy / unpinned listener
        ];
        for inet_port in cases {
            let l = TsiListenerSnapshot {
                cid: 3,
                peer_port: 1234,
                vm_port: 8080,
                family: 2, // AF_INET
                socktype: 1,
                inet_port,
            };
            let mut buf = Vec::new();
            write_tsi_listener_record(&mut buf, &l);
            assert_eq!(
                buf.len(),
                TSI_LISTENER_RECORD_LEN,
                "record must be exactly {TSI_LISTENER_RECORD_LEN} bytes"
            );
            let e: [u8; TSI_LISTENER_RECORD_LEN] = buf.try_into().unwrap();
            let got = read_tsi_listener_record(&e).expect("decode");
            assert_eq!(got.cid, l.cid);
            assert_eq!(got.peer_port, l.peer_port);
            assert_eq!(got.vm_port, l.vm_port);
            assert_eq!(got.family, l.family);
            assert_eq!(got.socktype, l.socktype);
            assert_eq!(
                got.inet_port, inet_port,
                "inet_port must survive the record round-trip (dropped pre-v10)"
            );
        }
    }

    // Regression: the smpark unpark-page GPA must round-trip through its
    // snapshot sidecar, so a restored worker can host-direct-unpark its
    // secondaries with the correct (snapshot-scoped) GPA instead of a
    // stale process-global static.
    #[test]
    fn smpark_sidecar_roundtrips_gpa() {
        let snap = temp_snapshot_path("smpark-rt");
        let snap_s = snap.to_str().unwrap();
        let _ = std::fs::remove_file(smpark_sidecar_path(snap_s));
        assert_eq!(read_smpark_sidecar(snap_s), None, "absent sidecar → None");
        write_smpark_sidecar(snap_s, 0x8_0000_1234).expect("write smpark sidecar");
        assert_eq!(
            read_smpark_sidecar(snap_s),
            Some(0x8_0000_1234),
            "smpark GPA must round-trip through the sidecar"
        );
        let _ = std::fs::remove_file(smpark_sidecar_path(snap_s));
    }

    fn temp_snapshot_path(name: &str) -> PathBuf {
        let mut path = std::env::temp_dir();
        path.push(format!(
            "snapshot-{name}-{}-{:?}.snap",
            std::process::id(),
            std::thread::current().id()
        ));
        path
    }

    fn write_bytes(path: &PathBuf, bytes: &[u8]) -> std::io::Result<()> {
        let mut file = std::fs::File::create(path)?;
        file.write_all(bytes)
    }

    /// Assemble an HVF snapshot meta file: the 32-byte prefix (valid
    /// magic/version, given `ram_gpa` + `ram_offset`) followed by `body` (the
    /// portable ContainerMeta bytes — valid or deliberately hostile).
    fn hvf_file(ram_gpa: u64, ram_offset: u64, body: &[u8]) -> Vec<u8> {
        let mut v = Vec::with_capacity(HVF_HDR_PREFIX + body.len());
        v.extend_from_slice(&SNAPSHOT_MAGIC);
        v.extend_from_slice(&SNAPSHOT_VERSION.to_le_bytes());
        v.extend_from_slice(&ram_gpa.to_le_bytes());
        v.extend_from_slice(&ram_offset.to_le_bytes());
        v.extend_from_slice(body);
        v
    }

    /// A valid ContainerMeta body with `n_vcpus` minimal per-vCPU register
    /// blobs, no devices/listeners, and the given `mem_size`.
    fn container_body(n_vcpus: usize, mem_size: u64) -> Vec<u8> {
        let per_vcpu: Vec<PerVcpuState> = (0..n_vcpus).map(|_| PerVcpuState::default()).collect();
        let cm = hvf_container_meta(0, 0, mem_size, &[], &per_vcpu, &VirtioSnapshot::default());
        let mut b = Vec::new();
        cm.write_container(&mut b).unwrap();
        b
    }

    #[test]
    fn load_rejects_truncated_snapshot() -> std::io::Result<()> {
        let path = temp_snapshot_path("truncated");
        write_bytes(&path, b"SMS")?;

        let result = load_from_file(path.to_str().unwrap_or_default());
        let _ = std::fs::remove_file(path);

        assert!(matches!(result, Err(FileError::Truncated)));
        Ok(())
    }

    #[test]
    fn load_rejects_bad_magic() -> std::io::Result<()> {
        let path = temp_snapshot_path("bad-magic");
        write_bytes(&path, &[0u8; 72])?;

        let result = load_from_file(path.to_str().unwrap_or_default());
        let _ = std::fs::remove_file(path);

        assert!(matches!(result, Err(FileError::BadMagic)));
        Ok(())
    }

    #[test]
    fn load_rejects_missing_vcpu_state() -> std::io::Result<()> {
        let path = temp_snapshot_path("no-vcpu");
        // A structurally valid container with zero vCPU blobs (mem_size 0, so
        // the RAM-region check passes) must still be rejected: a VM with no
        // captured vCPU state can't be resumed.
        let body = container_body(0, 0);
        let ram_offset = (HVF_HDR_PREFIX + body.len()) as u64;
        write_bytes(&path, &hvf_file(0, ram_offset, &body))?;

        let result = load_from_file(path.to_str().unwrap_or_default());
        let _ = std::fs::remove_file(path);

        assert!(matches!(
            result,
            Err(FileError::Malformed("snapshot contains no vCPU state"))
        ));
        Ok(())
    }

    // ── save → load round-trip (the SAVE side; write_sparse_ram) ───────

    fn sample_snapshot(memory: Vec<u8>) -> Snapshot {
        Snapshot {
            captured_host_ticks: 0x0123_4567_89AB_CDEF,
            captured_clock_ref: 0xFEDC_BA98_7654_3210,
            ram_gpa: 0x4000_0000,
            memory,
            intc_blob: vec![0xDE, 0xAD, 0xBE, 0xEF],
            per_vcpu: vec![PerVcpuState {
                gp_regs: vec![(0, 0x1111_2222), (1, 0x3333_4444)],
                vtimer_offset: 0x9999,
                ..Default::default()
            }],
            virtio: VirtioSnapshot::default(),
        }
    }

    #[test]
    fn save_load_round_trips_sparse_ram() {
        let path = temp_snapshot_path("roundtrip-sparse");
        let path_s = path.to_str().unwrap().to_string();

        // RAM with a deliberate mix of dense data and long zero runs, to
        // drive the sparse encoder's zero-skip and the loader's
        // reconstruction. A byte-exact round-trip proves no page is
        // dropped, mis-offset, or zero-filled where it shouldn't be.
        let mut memory = vec![0u8; 16 * 4096];
        for (i, b) in memory[..4096].iter_mut().enumerate() {
            *b = (i % 251) as u8; // page 0: dense
        }
        // pages 1..=13 left zero (the sparse runs)
        memory[14 * 4096 + 100] = 0xAB; // a lone non-zero byte mid-region
        memory[15 * 4096 + 4095] = 0xCD; // the very last byte
        let expected = memory.clone();

        save_to_file(&path_s, &sample_snapshot(memory)).expect("save");

        // Atomic-rename invariant: only the canonical file exists.
        assert!(path.exists(), "snapshot must exist after save");
        assert!(
            !PathBuf::from(format!("{path_s}.partial")).exists(),
            "no leftover .partial after a successful save"
        );

        let loaded = load_from_file(&path_s).expect("load");
        let _ = std::fs::remove_file(&path);

        assert_eq!(loaded.ram_gpa, 0x4000_0000);
        assert_eq!(loaded.captured_host_ticks, 0x0123_4567_89AB_CDEF);
        assert_eq!(loaded.captured_clock_ref, 0xFEDC_BA98_7654_3210);
        assert_eq!(loaded.memory.len(), expected.len(), "RAM length preserved");
        assert!(loaded.memory == expected, "RAM must round-trip byte-exact");
        assert_eq!(loaded.per_vcpu.len(), 1);
        assert_eq!(loaded.intc_blob, vec![0xDE, 0xAD, 0xBE, 0xEF]);
    }

    #[test]
    fn save_load_round_trips_all_zero_ram() {
        // A fully-zero region is the extreme of the sparse path: the
        // encoder skips everything, the loader must still reconstruct the
        // exact original length of zeros (a truncation bug would show here).
        let path = temp_snapshot_path("roundtrip-zero");
        let path_s = path.to_str().unwrap().to_string();
        let len = 8 * 4096;

        save_to_file(&path_s, &sample_snapshot(vec![0u8; len])).expect("save");
        let loaded = load_from_file(&path_s).expect("load");
        let _ = std::fs::remove_file(&path);

        assert_eq!(loaded.memory.len(), len, "zero RAM length preserved");
        assert!(
            loaded.memory.iter().all(|&b| b == 0),
            "all bytes still zero"
        );
    }

    #[test]
    fn save_load_preserves_per_vcpu_registers() {
        let path = temp_snapshot_path("roundtrip-regs");
        let path_s = path.to_str().unwrap().to_string();

        save_to_file(&path_s, &sample_snapshot(vec![0u8; 4096])).expect("save");
        let loaded = load_from_file(&path_s).expect("load");
        let _ = std::fs::remove_file(&path);

        let v = &loaded.per_vcpu[0];
        assert_eq!(v.vtimer_offset, 0x9999, "vtimer_offset round-trips");
        assert!(
            v.gp_regs.contains(&(0, 0x1111_2222)) && v.gp_regs.contains(&(1, 0x3333_4444)),
            "general-purpose regs round-trip: {:?}",
            v.gp_regs
        );
    }

    // ── Sidecar parser hardening ──────────────────────────────────
    // read_dax_sidecar / read_posix_fs_sidecar parse a header `count`
    // then loop over length-prefixed records. The `count` is read
    // straight from the file; trusting it for `Vec::with_capacity`
    // would let a corrupted or malicious sidecar request a multi-GiB
    // allocation (OOM) before any bounds check runs. These tests pin
    // that the pre-allocation is bounded by the file size and that
    // arbitrary bytes never panic / OOM / hang.

    fn sidecar_path(name: &str) -> PathBuf {
        let mut path = std::env::temp_dir();
        path.push(format!(
            "sidecar-{name}-{}-{:?}.bin",
            std::process::id(),
            std::thread::current().id()
        ));
        path
    }

    /// REGRESSION: a DAXC sidecar advertising a huge session_count but
    /// carrying no records must fail FAST with Truncated — not attempt
    /// a `with_capacity(u32::MAX)` (~96 GiB of Vec<u8> headers) first.
    #[test]
    fn dax_sidecar_huge_count_does_not_ooms() -> std::io::Result<()> {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(b"DAXC");
        bytes.extend_from_slice(&1u32.to_le_bytes()); // version
        bytes.extend_from_slice(&u32::MAX.to_le_bytes()); // session_count = 4 billion
                                                          // No record bodies follow → first iteration hits Truncated.
        let path = sidecar_path("dax-huge");
        write_bytes(&path, &bytes)?;
        let r = read_dax_sidecar(path.to_str().unwrap());
        let _ = std::fs::remove_file(&path);
        assert!(matches!(r, Err(FileError::Truncated)));
        Ok(())
    }

    /// Same contract for the PosixFs sidecar (PFXS / mount_count).
    #[test]
    fn posix_fs_sidecar_huge_count_does_not_ooms() -> std::io::Result<()> {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(b"PFXS");
        bytes.extend_from_slice(&1u32.to_le_bytes());
        bytes.extend_from_slice(&u32::MAX.to_le_bytes());
        let path = sidecar_path("pfxs-huge");
        write_bytes(&path, &bytes)?;
        let r = read_posix_fs_sidecar(path.to_str().unwrap());
        let _ = std::fs::remove_file(&path);
        assert!(matches!(r, Err(FileError::Truncated)));
        Ok(())
    }

    /// A well-formed single-record DAXC sidecar still round-trips
    /// (guard against the bound being too tight and rejecting valid
    /// input).
    #[test]
    fn dax_sidecar_valid_single_record_roundtrips() -> std::io::Result<()> {
        let payload = b"hello-dax-blob";
        let mut bytes = Vec::new();
        bytes.extend_from_slice(b"DAXC");
        bytes.extend_from_slice(&1u32.to_le_bytes());
        bytes.extend_from_slice(&1u32.to_le_bytes()); // session_count = 1
        bytes.extend_from_slice(&(payload.len() as u32).to_le_bytes());
        bytes.extend_from_slice(payload);
        let path = sidecar_path("dax-valid");
        write_bytes(&path, &bytes)?;
        let r = read_dax_sidecar(path.to_str().unwrap());
        let _ = std::fs::remove_file(&path);
        assert_eq!(r.unwrap(), Some(vec![payload.to_vec()]));
        Ok(())
    }

    proptest::proptest! {
        #![proptest_config(proptest::prelude::ProptestConfig::with_cases(256))]

        /// Arbitrary bytes (often DAXC/PFXS-prefixed so the count/len
        /// loop is actually exercised) must never panic, never OOM, and
        /// always resolve to a clean Result. `count`/`len` here are
        /// fully attacker-controlled u32s.
        #[test]
        fn sidecar_parsers_tolerate_arbitrary_bytes(
            prefix in proptest::prelude::prop_oneof![
                proptest::prelude::Just(b"DAXC".to_vec()),
                proptest::prelude::Just(b"PFXS".to_vec()),
                proptest::prelude::Just(Vec::<u8>::new()),
            ],
            body in proptest::collection::vec(proptest::prelude::any::<u8>(), 0..512),
        ) {
            let mut bytes = prefix.clone();
            bytes.extend_from_slice(&body);
            let path = sidecar_path("fuzz");
            std::fs::write(&path, &bytes).unwrap();
            // Neither call may panic / OOM / hang; both return a Result.
            let _ = read_dax_sidecar(path.to_str().unwrap());
            let _ = read_posix_fs_sidecar(path.to_str().unwrap());
            let _ = std::fs::remove_file(&path);
        }
    }

    // ── Main snapshot loader hardening ────────────────────────────
    // load_from_file_inner reads length/count fields from the header
    // and pre-allocates on them. A corrupt or hostile snapshot (these
    // cross trust boundaries: CDN catalog, cross-process restore) must
    // fail FAST — never OOM on a giant `with_capacity`/`vec!`, never
    // mmap past EOF (SIGBUS). cap_count + the ram_offset/memory_bytes
    // file-size validation enforce that.

    /// A container claiming a 1 TiB RAM section the file can't hold must fail
    /// fast — not `vec![0u8; memory_bytes]` (OOM) or let a caller mmap past EOF.
    /// This is the HVF loader's own `ram_region_within` guard (mem_size now
    /// rides inside the ContainerMeta, checked against the prefix ram_offset).
    #[test]
    fn load_rejects_lying_memory_bytes() -> std::io::Result<()> {
        let body = container_body(1, 1u64 << 40); // mem_size = 1 TiB
        let ram_offset = (HVF_HDR_PREFIX + body.len()) as u64; // tiny file
        let path = temp_snapshot_path("lying-mem");
        write_bytes(&path, &hvf_file(0, ram_offset, &body))?;
        let r = load_from_file(path.to_str().unwrap());
        let _ = std::fs::remove_file(&path);
        assert!(matches!(r, Err(FileError::Truncated)));
        Ok(())
    }

    /// A container whose intc_blob length is u64/u32::MAX must not drive a giant
    /// allocation — the bounded codec (`ContainerMeta::read_container`, capped at
    /// CONTAINER_BLOB_MAX) rejects it and the HVF loader maps that to Truncated.
    #[test]
    fn load_rejects_huge_intc_blob_len() -> std::io::Result<()> {
        // Hand-crafted hostile container: valid scalar prefix then an intc_blob
        // length field of u32::MAX with no bytes following.
        let mut body = Vec::new();
        body.push(1u8); // num_cpus
        body.extend_from_slice(&0u64.to_le_bytes()); // mem_size
        body.extend_from_slice(&[0u8; 6]); // com1
        body.extend_from_slice(&0u64.to_le_bytes()); // clock_host_ticks
        body.extend_from_slice(&0u64.to_le_bytes()); // clock_ref
        body.extend_from_slice(&u32::MAX.to_le_bytes()); // intc_blob len = 4 GiB
        let path = temp_snapshot_path("huge-gic");
        write_bytes(&path, &hvf_file(0, 4096, &body))?;
        let r = load_from_file(path.to_str().unwrap());
        let _ = std::fs::remove_file(&path);
        assert!(matches!(r, Err(FileError::Truncated)));
        Ok(())
    }

    /// n_vcpus = u32::MAX must not drive a giant `Vec::with_capacity`; the
    /// container reader caps the pre-allocation and the per-vCPU blob read then
    /// hits the truncated file and fails cleanly.
    #[test]
    fn load_rejects_huge_n_vcpus() -> std::io::Result<()> {
        let mut body = Vec::new();
        body.push(1u8); // num_cpus
        body.extend_from_slice(&0u64.to_le_bytes()); // mem_size
        body.extend_from_slice(&[0u8; 6]); // com1
        body.extend_from_slice(&0u64.to_le_bytes()); // clock_host_ticks
        body.extend_from_slice(&0u64.to_le_bytes()); // clock_ref
        body.extend_from_slice(&0u32.to_le_bytes()); // intc_blob len = 0
        body.extend_from_slice(&u32::MAX.to_le_bytes()); // n_vcpus = 4 billion
        let path = temp_snapshot_path("huge-vcpus");
        write_bytes(&path, &hvf_file(0, 4096, &body))?;
        let r = load_from_file(path.to_str().unwrap());
        let _ = std::fs::remove_file(&path);
        assert!(matches!(r, Err(FileError::Truncated)));
        Ok(())
    }

    /// mmap_ram_cow_at is a `pub` entry point; mapping a RAM region past
    /// EOF would SIGBUS on first touch, so it must reject out-of-range
    /// (ram_offset, memory_bytes) up front.
    #[test]
    fn mmap_ram_cow_at_rejects_region_past_eof() -> std::io::Result<()> {
        let path = temp_snapshot_path("mmap-oob");
        write_bytes(&path, &[0u8; 4096])?; // 4 KiB file
        let r = mmap_ram_cow_at(path.to_str().unwrap(), 0, 1 << 30); // ask for 1 GiB
        let _ = std::fs::remove_file(&path);
        assert!(r.is_err(), "mapping 1 GiB out of a 4 KiB file must fail");
        Ok(())
    }

    proptest::proptest! {
        #![proptest_config(proptest::prelude::ProptestConfig::with_cases(256))]

        /// A magic-valid prefix with an attacker-controlled ram_gpa/ram_offset
        /// plus a fully arbitrary container body must never panic, OOM, or hang
        /// — only a clean Result. Exercises the HVF loader's delegation to the
        /// bounded `ContainerMeta::read_container` and the ram_region_within
        /// validation against the prefix ram_offset.
        #[test]
        fn loader_tolerates_arbitrary_header_sizes(
            ram_gpa in proptest::prelude::any::<u64>(),
            ram_offset in proptest::prelude::any::<u64>(),
            body in proptest::collection::vec(proptest::prelude::any::<u8>(), 0..512),
        ) {
            let bytes = hvf_file(ram_gpa, ram_offset, &body);
            let path = temp_snapshot_path("loader-fuzz");
            std::fs::write(&path, &bytes).unwrap();
            // Neither entry point may panic / OOM / hang.
            let _ = load_from_file(path.to_str().unwrap());
            let _ = load_meta(path.to_str().unwrap());
            let _ = std::fs::remove_file(&path);
        }
    }
}