Upgrade for new spawn (#431)

fuzz/fuzz_targets/snapshot2.rs

@@ -46,7 +46,7 @@ pub struct DummyData {
 }
 
 impl DataSource<u32> for DummyData {
-    fn load_data(&self, id: &u32, offset: u64, length: u64) -> Result<(Bytes, u64), Error> {
+    fn load_data(&self, id: &u32, offset: u64, length: u64) -> Option<(Bytes, u64)> {
         let data = match *id {
             DATA_SOURCE_PROGRAM => &self.program,
             DATA_SOURCE_CONTENT => &self.content,
@@ -59,7 +59,7 @@ impl DataSource<u32> for DummyData {
         } else {
             full_size
         };
-        Ok((data.slice(offset..offset + real_size), full_size as u64))
+        Some((data.slice(offset..offset + real_size), full_size as u64))
     }
 }
 

src/error.rs

@@ -19,8 +19,8 @@ pub enum Error {
     ElfSegmentWritableAndExecutable(u64),
     #[display(fmt = "elf error: segment addr or size is wrong vaddr=0x{:x}", "_0")]
     ElfSegmentAddrOrSizeError(u64),
-    // External error type is for the debugging tool of CKB-VM, it should not be
-    // used in this project.
+    // When users need to implement traits defined in CKB-VM, they can use
+    // this error type to wrap their own errors.
     #[display(fmt = "external error: {}", "_0")]
     External(String),
     #[display(fmt = "invalid syscall {}", "_0")]
@@ -52,8 +52,12 @@ pub enum Error {
     MemWriteOnFreezedPage(u64),
     #[display(fmt = "pause")]
     Pause,
+    #[display(fmt = "snapshot data load error")]
+    SnapshotDataLoadError,
     #[display(fmt = "unexpected error")]
     Unexpected(String),
+    #[display(fmt = "yield")]
+    Yield,
 }
 
 #[derive(Debug, PartialEq, Clone, Eq, Display)]

src/machine/trace.rs

@@ -127,6 +127,10 @@ impl<Inner: SupportMachine> TraceMachine<Inner> {
             .load_program_with_metadata(program, metadata, args)
     }
 
+    pub fn set_max_cycles(&mut self, cycles: u64) {
+        self.machine.inner_mut().set_max_cycles(cycles)
+    }
+
     pub fn run(&mut self) -> Result<i8, Error> {
         let mut decoder = build_decoder::<Inner::REG>(self.isa(), self.version());
         self.run_with_decoder(&mut decoder)

src/snapshot2.rs

@@ -24,7 +24,7 @@ const PAGE_SIZE: u64 = RISCV_PAGESIZE as u64;
 /// 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, u64), Error>;
+    fn load_data(&self, id: &I, offset: u64, length: u64) -> Option<(Bytes, u64)>;
 }
 
 #[derive(Clone, Debug)]
@@ -70,7 +70,7 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
             if address % PAGE_SIZE != 0 {
                 return Err(Error::MemPageUnalignedAccess(*address));
             }
-            let (data, _) = self.data_source().load_data(id, *offset, *length)?;
+            let (data, _) = self.load_data(id, *offset, *length)?;
             if data.len() as u64 % PAGE_SIZE != 0 {
                 return Err(Error::MemPageUnalignedAccess(
                     address.wrapping_add(data.len() as u64),
@@ -104,12 +104,10 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
         Ok(())
     }
 
-    pub fn data_source(&self) -> &D {
-        &self.data_source
-    }
-
-    pub fn data_source_mut(&mut self) -> &mut D {
-        &mut self.data_source
+    pub fn load_data(&mut self, id: &I, offset: u64, length: u64) -> Result<(Bytes, u64), Error> {
+        self.data_source
+            .load_data(id, offset, length)
+            .ok_or(Error::SnapshotDataLoadError)
     }
 
     /// Similar to Memory::store_bytes, but this method also tracks memory
@@ -124,7 +122,7 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
         offset: u64,
         length: u64,
     ) -> Result<(u64, u64), Error> {
-        let (data, full_length) = self.data_source.load_data(id, offset, length)?;
+        let (data, full_length) = self.load_data(id, offset, 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)?;

tests/test_resume2.rs

@@ -328,7 +328,7 @@ 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, u64), Error> {
+    fn load_data(&self, id: &u64, offset: u64, length: u64) -> Option<(Bytes, u64)> {
         match self.0.get(id) {
             Some(data) => {
                 let end = if length > 0 {
@@ -337,12 +337,9 @@ impl DataSource<u64> for TestSource {
                     data.len() as u64
                 };
                 let full_length = end - offset;
-                Ok((data.slice(offset as usize..end as usize), full_length))
+                Some((data.slice(offset as usize..end as usize), full_length))
             }
-            None => Err(Error::Unexpected(format!(
-                "Id {} is missing in source!",
-                id
-            ))),
+            None => None,
         }
     }
 }
@@ -452,7 +449,6 @@ impl Machine {
                 let (program, _) = context
                     .lock()
                     .unwrap()
-                    .data_source()
                     .load_data(&PROGRAM_ID, 0, 0)
                     .unwrap();
                 let metadata = parse_elf::<u64>(&program, inner.machine.version())?;
@@ -469,7 +465,6 @@ impl Machine {
                 let (program, _) = context
                     .lock()
                     .unwrap()
-                    .data_source()
                     .load_data(&PROGRAM_ID, 0, 0)
                     .unwrap();
                 let metadata = parse_elf::<u64>(&program, inner.version())?;
@@ -486,7 +481,6 @@ impl Machine {
                 let (program, _) = context
                     .lock()
                     .unwrap()
-                    .data_source()
                     .load_data(&PROGRAM_ID, 0, 0)
                     .unwrap();
                 let metadata = parse_elf::<u64>(&program, inner.machine.version())?;