Use correct memory pages

src/snapshot.rs

@@ -1,10 +1,7 @@
 use crate::instructions::Register;
 use crate::memory::Memory;
 use crate::memory::FLAG_DIRTY;
-use crate::{
-    CoreMachine, Error, RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGES, RISCV_PAGESIZE,
-    RISCV_PAGE_SHIFTS,
-};
+use crate::{CoreMachine, Error, RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGESIZE, RISCV_PAGE_SHIFTS};
 use serde::{Deserialize, Serialize};
 
 // Snapshot provides a mechanism for suspending and resuming a virtual machine.
@@ -47,7 +44,7 @@ pub fn make_snapshot<T: CoreMachine>(machine: &mut T) -> Result<Snapshot, Error>
         snap.registers[i] = v.to_u64();
     }
 
-    for i in 0..RISCV_PAGES {
+    for i in 0..(machine.memory().memory_size() >> RISCV_PAGE_SHIFTS) {
         let flag = machine.memory_mut().fetch_flag(i as u64)?;
         if flag & FLAG_DIRTY != 0 {
             let addr_from = i << RISCV_PAGE_SHIFTS;

src/snapshot2.rs

@@ -3,15 +3,14 @@ use crate::{
     elf::{LoadingAction, ProgramMetadata},
     machine::SupportMachine,
     memory::{Memory, FLAG_DIRTY},
-    Error, Register, RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGES, RISCV_PAGESIZE,
+    Error, Register, RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGESIZE, RISCV_PAGE_SHIFTS,
 };
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use std::cmp::min;
 use std::collections::HashMap;
 
 const PAGE_SIZE: u64 = RISCV_PAGESIZE as u64;
-const PAGES: u64 = RISCV_PAGES as u64;
 
 /// DataSource represents data source that can stay stable and possibly
 /// immutable for the entire lifecycle duration of a VM instance. One example
@@ -147,7 +146,7 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
     /// Create a snapshot for the passed machine.
     pub fn make_snapshot<M: SupportMachine>(&self, machine: &mut M) -> Result<Snapshot2<I>, Error> {
         let mut dirty_pages: Vec<(u64, u8, Vec<u8>)> = vec![];
-        for i in 0..PAGES {
+        for i in 0..(machine.memory().memory_size() as u64 >> RISCV_PAGE_SHIFTS) {
             if self.pages.contains_key(&i) {
                 continue;
             }