Upgrade for new spawn (#430)

* Add yield error

* Add index out of bound error

* Add set_max_cycles for TraceMachine

* Refactoring DataSource

fuzz/fuzz_targets/snapshot2.rs

@@ -4,8 +4,8 @@ use ckb_vm::{
     machine::VERSION2,
     memory::{round_page_down, round_page_up, FLAG_EXECUTABLE, FLAG_FREEZED},
     snapshot2::{DataSource, Snapshot2Context},
-    Bytes, CoreMachine, DefaultMachine, DefaultMachineBuilder, Error, Memory, ISA_A, ISA_B,
-    ISA_IMC, ISA_MOP, RISCV_MAX_MEMORY, RISCV_PAGESIZE,
+    Bytes, CoreMachine, DefaultMachine, DefaultMachineBuilder, Memory, ISA_A, ISA_B, ISA_IMC,
+    ISA_MOP, RISCV_MAX_MEMORY, RISCV_PAGESIZE,
 };
 use ckb_vm_definitions::asm::AsmCoreMachine;
 use libfuzzer_sys::fuzz_target;
@@ -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

@@ -16,8 +16,8 @@ pub enum Error {
     ElfSegmentWritableAndExecutable,
     #[display(fmt = "elf error: segment addr or size is wrong")]
     ElfSegmentAddrOrSizeError,
-    // 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")]
@@ -49,10 +49,14 @@ pub enum Error {
     MemWriteOnFreezedPage,
     #[display(fmt = "pause")]
     Pause,
+    #[display(fmt = "snapshot data load error")]
+    SnapshotDataLoadError,
     #[display(fmt = "unexpected error")]
     Unexpected(String),
     #[display(fmt = "unimplemented")]
     Unimplemented,
+    #[display(fmt = "yield")]
+    Yield,
 }
 
 impl std::error::Error for Error {}

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.machine.set_running(true);

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);
             }
-            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);
             }
@@ -100,12 +100,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
@@ -120,7 +118,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())?;