From 82bcf07de50499f2ff86ef3bd0bc848dc82c370b Mon Sep 17 00:00:00 2001
From: Mohanson <mohanson@outlook.com>
Date: Thu, 14 Mar 2024 18:22:04 +0800
Subject: [PATCH] Backport #412

---
 src/snapshot2.rs      | 23 +++++++++++++++--------
 tests/test_resume2.rs | 22 ++++++++++++----------
 2 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/src/snapshot2.rs b/src/snapshot2.rs
index dc271048..07bd6110 100644
--- a/src/snapshot2.rs
+++ b/src/snapshot2.rs
@@ -20,9 +20,11 @@ const PAGE_SIZE: u64 = RISCV_PAGESIZE as u64;
 /// we can leverage DataSource for snapshot optimizations: data that is already
 /// locatable in the DataSource will not need to be included in the snapshot
 /// again, all we need is an id to locate it, together with a pair of
-/// offset / length to cut in to the correct slices.
+/// offset / length to cut in to the correct slices. Just like CKB's syscall design,
+/// an extra u64 value is included here to return the remaining full length of data
+/// starting from offset, without considering `length` parameter
 pub trait DataSource<I: Clone + PartialEq> {
-    fn load_data(&self, id: &I, offset: u64, length: u64) -> Result<Bytes, Error>;
+    fn load_data(&self, id: &I, offset: u64, length: u64) -> Result<(Bytes, u64), Error>;
 }
 
 #[derive(Clone, Debug)]
@@ -68,7 +70,7 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
             if address % PAGE_SIZE != 0 {
                 return Err(Error::MemPageUnalignedAccess);
             }
-            let data = self.data_source().load_data(id, *offset, *length)?;
+            let (data, _) = self.data_source().load_data(id, *offset, *length)?;
             if data.len() as u64 % PAGE_SIZE != 0 {
                 return Err(Error::MemPageUnalignedAccess);
             }
@@ -107,7 +109,9 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
     }
 
     /// Similar to Memory::store_bytes, but this method also tracks memory
-    /// pages whose entire content comes from DataSource
+    /// pages whose entire content comes from DataSource. It returns 2 values:
+    /// the actual written bytes, and the full length of data starting from offset,
+    /// but ignoring `length` parameter.
     pub fn store_bytes<M: SupportMachine>(
         &mut self,
         machine: &mut M,
@@ -115,10 +119,11 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
         id: &I,
         offset: u64,
         length: u64,
-    ) -> Result<(), Error> {
-        let data = self.data_source.load_data(id, offset, length)?;
+    ) -> Result<(u64, u64), Error> {
+        let (data, full_length) = self.data_source.load_data(id, offset, length)?;
         machine.memory_mut().store_bytes(addr, &data)?;
-        self.track_pages(machine, addr, data.len() as u64, id, offset)
+        self.track_pages(machine, addr, data.len() as u64, id, offset)?;
+        Ok((data.len() as u64, full_length))
     }
 
     /// Due to the design of ckb-vm right now, load_program function does not
@@ -222,7 +227,9 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
         self.track_pages(machine, start, length, id, offset + action.source.start)
     }
 
-    fn track_pages<M: SupportMachine>(
+    /// This is only made public for advanced usages, but make sure to exercise more
+    /// cautions when calling it!
+    pub fn track_pages<M: SupportMachine>(
         &mut self,
         machine: &mut M,
         start: u64,
diff --git a/tests/test_resume2.rs b/tests/test_resume2.rs
index 96717507..305786ca 100644
--- a/tests/test_resume2.rs
+++ b/tests/test_resume2.rs
@@ -324,15 +324,17 @@ const DATA_ID: u64 = 0x2000;
 struct TestSource(HashMap<u64, Bytes>);
 
 impl DataSource<u64> for TestSource {
-    fn load_data(&self, id: &u64, offset: u64, length: u64) -> Result<Bytes, Error> {
+    fn load_data(&self, id: &u64, offset: u64, length: u64) -> Result<(Bytes, u64), Error> {
         match self.0.get(id) {
-            Some(data) => Ok(data.slice(
-                offset as usize..(if length > 0 {
-                    (offset + length) as usize
+            Some(data) => {
+                let end = if length > 0 {
+                    offset + length
                 } else {
-                    data.len()
-                }),
-            )),
+                    data.len() as u64
+                };
+                let full_length = end - offset;
+                Ok((data.slice(offset as usize..end as usize), full_length))
+            }
             None => Err(Error::Unexpected(format!(
                 "Id {} is missing in source!",
                 id
@@ -443,7 +445,7 @@ impl Machine {
         use Machine::*;
         match self {
             Asm(inner, context) => {
-                let program = context
+                let (program, _) = context
                     .lock()
                     .unwrap()
                     .data_source()
@@ -460,7 +462,7 @@ impl Machine {
                 Ok(bytes)
             }
             Interpreter(inner, context) => {
-                let program = context
+                let (program, _) = context
                     .lock()
                     .unwrap()
                     .data_source()
@@ -477,7 +479,7 @@ impl Machine {
                 Ok(bytes)
             }
             InterpreterWithTrace(inner, context) => {
-                let program = context
+                let (program, _) = context
                     .lock()
                     .unwrap()
                     .data_source()