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, 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}