From b915850217b6db97427ce075860fcd5fe07ce974 Mon Sep 17 00:00:00 2001
From: mohanson <mohanson@outlook.com>
Date: Tue, 28 May 2024 21:44:31 +0800
Subject: [PATCH 1/2] Store size when call store_bytes

---
 src/snapshot2.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/snapshot2.rs b/src/snapshot2.rs
index d7840a54..f4f98164 100644
--- a/src/snapshot2.rs
+++ b/src/snapshot2.rs
@@ -121,8 +121,12 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
         id: &I,
         offset: u64,
         length: u64,
+        size_addr: u64,
     ) -> Result<(u64, u64), Error> {
         let (data, full_length) = self.load_data(id, offset, length)?;
+        machine
+            .memory_mut()
+            .store64(&M::REG::from_u64(size_addr), &M::REG::from_u64(full_length))?;
         self.untrack_pages(machine, addr, data.len() as u64)?;
         machine.memory_mut().store_bytes(addr, &data)?;
         self.track_pages(machine, addr, data.len() as u64, id, offset)?;

From 3b88aa4e10555081f3f8fec13e1b11cba36fa643 Mon Sep 17 00:00:00 2001
From: mohanson <mohanson@outlook.com>
Date: Tue, 28 May 2024 21:53:30 +0800
Subject: [PATCH 2/2] Fix test

---
 tests/test_resume2.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/test_resume2.rs b/tests/test_resume2.rs
index d3360f2f..b02e58ef 100644
--- a/tests/test_resume2.rs
+++ b/tests/test_resume2.rs
@@ -362,7 +362,7 @@ impl<Mac: SupportMachine> Syscalls<Mac> for InsertDataSyscall {
         self.0
             .lock()
             .unwrap()
-            .store_bytes(machine, addr, &DATA_ID, 0, size)?;
+            .store_bytes(machine, addr, &DATA_ID, 0, size, 0)?;
 
         machine.add_cycles_no_checking(100000)?;
 
@@ -634,11 +634,11 @@ pub fn test_store_bytes_twice() {
         Machine::Asm(ref mut inner, ref ctx) => {
             ctx.lock()
                 .unwrap()
-                .store_bytes(&mut inner.machine, 0, &DATA_ID, 2, 29186)
+                .store_bytes(&mut inner.machine, 0, &DATA_ID, 2, 29186, 0)
                 .unwrap();
             ctx.lock()
                 .unwrap()
-                .store_bytes(&mut inner.machine, 0, &DATA_ID, 0, 11008)
+                .store_bytes(&mut inner.machine, 0, &DATA_ID, 0, 11008, 0)
                 .unwrap();
         }
         _ => unimplemented!(),
@@ -667,7 +667,7 @@ pub fn test_mixing_snapshot2_writes_with_machine_raw_writes() {
         Machine::Asm(ref mut inner, ref ctx) => {
             ctx.lock()
                 .unwrap()
-                .store_bytes(&mut inner.machine, 0, &DATA_ID, 0, 29186)
+                .store_bytes(&mut inner.machine, 0, &DATA_ID, 0, 29186, 0)
                 .unwrap();
             inner
                 .machine