From ef3b064aa68583f8fc6f3d76293ba143b711468a Mon Sep 17 00:00:00 2001
From: Mohanson <mohanson@outlook.com>
Date: Tue, 28 May 2024 22:03:29 +0800
Subject: [PATCH] Store size when call store_bytes (#434)

* Store size when call store_bytes

* Fix test
---
 src/snapshot2.rs      | 4 ++++
 tests/test_resume2.rs | 8 ++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

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)?;
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