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 { - fn load_data(&self, id: &I, offset: u64, length: u64) -> Result; + fn load_data(&self, id: &I, offset: u64, length: u64) -> Result<(Bytes, u64), Error>; } #[derive(Clone, Debug)] @@ -68,7 +70,7 @@ impl> Snapshot2Context { 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> Snapshot2Context { } /// 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( &mut self, machine: &mut M, @@ -115,10 +119,11 @@ impl> Snapshot2Context { 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> Snapshot2Context { self.track_pages(machine, start, length, id, offset + action.source.start) } - fn track_pages( + /// This is only made public for advanced usages, but make sure to exercise more + /// cautions when calling it! + pub fn track_pages( &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); impl DataSource for TestSource { - fn load_data(&self, id: &u64, offset: u64, length: u64) -> Result { + 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()