nyx_lite/
snapshot.rs

1use std::{collections::HashMap, sync::Arc};
2
3use vmm::persist::MicrovmState;
4
5use crate::{mem::PAGE_SIZE, vm_continuation_statemachine::VMContinuationState};
6
7#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
8pub enum SnapshotType {
9    Incremental,
10    Base,
11}
12
13pub struct BaseRegionSnapshot {
14    pub start: u64,
15    pub data: Arc<[u8]>,
16}
17pub enum MemorySnapshot {
18    Base(Vec<BaseRegionSnapshot>),
19    Incremental(HashMap<u64, Vec<u8>>),
20}
21
22impl MemorySnapshot {
23    pub fn is_incremental(&self) -> bool {
24        matches!(self, Self::Incremental(_))
25    }
26
27    pub fn get_page(&self, paddr: usize) -> Option<&[u8]> {
28        match self {
29            MemorySnapshot::Base(regions) => {
30                let paddr = paddr as u64;
31                for region in regions {
32                    let region_start = region.start;
33                    let region_len = region.data.len() as u64;
34                    if paddr < region_start || paddr >= region_start + region_len {
35                        continue;
36                    }
37                    let offset = (paddr - region_start) as usize;
38                    let offset_end = offset.checked_add(PAGE_SIZE as usize)?;
39                    if offset_end <= region.data.len() {
40                        return Some(region.data[offset..offset_end].try_into().unwrap());
41                    }
42                }
43                None
44            }
45            MemorySnapshot::Incremental(map) => map.get(&(paddr as u64)).map(|v| v.as_ref()),
46        }
47    }
48}
49
50pub struct NyxSnapshot {
51    pub parent: Option<Arc<NyxSnapshot>>,
52    pub depth: usize, // 0 for root snapshots, n+1 for each child incremental snapshot
53    pub memory: MemorySnapshot,
54    pub state: MicrovmState,
55    pub tsc: u64,
56    pub continuation_state: VMContinuationState,
57}
58
59impl NyxSnapshot {
60    pub fn get_page<Callback>(&self, paddr: usize, mut callback: Callback)
61    where
62        Callback: FnMut(&[u8; PAGE_SIZE as usize]),
63    {
64        if let Some(page) = self.memory.get_page(paddr) {
65            return callback(page.try_into().unwrap());
66        }
67        assert!(self.memory.is_incremental());
68        assert!(self.parent.is_some());
69        self.parent
70            .as_ref()
71            .expect("Incremental Snapshots should always have a parent")
72            .get_page(paddr, callback)
73    }
74
75    pub fn iter_delta(&self) -> impl Iterator<Item = (u64, &Vec<u8>)> {
76        if let MemorySnapshot::Incremental(ref map) = self.memory {
77            return map.iter().map(|(p, s)| (*p, s));
78        }
79        panic!("can't iter delta on a root snapshot");
80    }
81}