Let snapshots stores the load_reservation_address (#407)

* Let snapshots stores the load_reservation_address * Add testcase
nervosnetwork · Feb 28, 2024 · a83a0cc · a83a0cc
1 parent b7da390
commit a83a0cc
Show file tree

Hide file tree

Showing 8 changed files with 100 additions and 8 deletions.
diff --git a/src/snapshot.rs b/src/snapshot.rs
@@ -32,12 +32,14 @@ pub struct Snapshot {
     pub page_indices: Vec<u64>,
     pub page_flags: Vec<u8>,
     pub pages: Vec<Vec<u8>>,
+    pub load_reservation_address: u64,
 }
 
 pub fn make_snapshot<T: CoreMachine>(machine: &mut T) -> Result<Snapshot, Error> {
     let mut snap = Snapshot {
         version: machine.version(),
         pc: machine.pc().to_u64(),
+        load_reservation_address: machine.memory().lr().to_u64(),
         ..Default::default()
     };
     for (i, v) in machine.registers().iter().enumerate() {
@@ -92,6 +94,8 @@ pub fn resume<T: CoreMachine>(machine: &mut T, snapshot: &Snapshot) -> Result<()
         machine.memory_mut().store_bytes(addr_from, &page[..])?;
         machine.memory_mut().set_flag(page_index, page_flag)?;
     }
-
+    machine
+        .memory_mut()
+        .set_lr(&T::REG::from_u64(snapshot.load_reservation_address));
     Ok(())
 }
diff --git a/src/snapshot2.rs b/src/snapshot2.rs
@@ -96,6 +96,9 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
                 machine.memory_mut().set_flag(page, *flag)?;
             }
         }
+        machine
+            .memory_mut()
+            .set_lr(&M::REG::from_u64(snapshot.load_reservation_address));
         Ok(())
     }
 
@@ -204,6 +207,7 @@ impl<I: Clone + PartialEq, D: DataSource<I>> Snapshot2Context<I, D> {
             pc: machine.pc().to_u64(),
             cycles: machine.cycles(),
             max_cycles: machine.max_cycles(),
+            load_reservation_address: machine.memory().lr().to_u64(),
         })
     }
 
@@ -261,4 +265,5 @@ pub struct Snapshot2<I: Clone + PartialEq> {
     pub pc: u64,
     pub cycles: u64,
     pub max_cycles: u64,
+    pub load_reservation_address: u64,
 }
diff --git a/tests/programs/_build_all_native.sh b/tests/programs/_build_all_native.sh
@@ -53,6 +53,7 @@ riscv64-unknown-elf-as -o rorw_in_end_of_aot_block.o rorw_in_end_of_aot_block.S
 sh rvc_pageend.sh
 # TODO: sbinvi_aot_load_imm_bug
 riscv64-unknown-elf-as -o sc_after_sc.o sc_after_sc.S && riscv64-unknown-elf-ld -T sc_after_sc.lds -o sc_after_sc sc_after_sc.o && rm sc_after_sc.o
+riscv64-unknown-elf-as -o sc_after_snapshot.o sc_after_snapshot.S && riscv64-unknown-elf-ld -T sc_after_snapshot.lds -o sc_after_snapshot sc_after_snapshot.o && rm sc_after_snapshot.o
 riscv64-unknown-elf-as -o sc_only.o sc_only.S && riscv64-unknown-elf-ld -T sc_only.lds -o sc_only sc_only.o && rm sc_only.o
 # SKIP: simple
 riscv64-unknown-elf-gcc -o simple64 simple.c

diff --git a/tests/programs/sc_after_snapshot b/tests/programs/sc_after_snapshot
diff --git a/tests/programs/sc_after_snapshot.S b/tests/programs/sc_after_snapshot.S
@@ -0,0 +1,22 @@
+.global _start
+_start:
+_test_case_sc_after_snapshot:
+    la a0, n0         # a0 holds address of memory location n0
+    lr.d a2, (a0)
+    nop
+    nop
+    nop
+    nop
+    sc.d a3, a2, (a0)
+    bnez a3, fail     # sc.d must success
+done:
+    li a0, 0
+    li a7, 93
+    ecall
+fail:
+    li a0, 1
+    li a7, 93
+    ecall
+.section .data
+n0:
+    .dword 4      # Initialize to 4
diff --git a/tests/programs/sc_after_snapshot.lds b/tests/programs/sc_after_snapshot.lds
@@ -0,0 +1,7 @@
+SECTIONS
+{
+. = 0x100b0;
+.text : { *(.text) }
+. = 0x11000;
+.data : { *(.data) }
+}
diff --git a/tests/test_resume.rs b/tests/test_resume.rs
@@ -4,10 +4,12 @@ use bytes::Bytes;
 use ckb_vm::cost_model::constant_cycles;
 use ckb_vm::machine::asm::{AsmCoreMachine, AsmMachine};
 use ckb_vm::machine::trace::TraceMachine;
-use ckb_vm::machine::{DefaultCoreMachine, DefaultMachine, SupportMachine, VERSION0, VERSION1};
+use ckb_vm::machine::{
+    DefaultCoreMachine, DefaultMachine, SupportMachine, VERSION0, VERSION1, VERSION2,
+};
 use ckb_vm::memory::{sparse::SparseMemory, wxorx::WXorXMemory};
 use ckb_vm::snapshot::{make_snapshot, resume, Snapshot};
-use ckb_vm::{DefaultMachineBuilder, Error, ISA_IMC};
+use ckb_vm::{DefaultMachineBuilder, Error, ISA_A, ISA_IMC};
 use std::fs::File;
 use std::io::Read;
 
@@ -294,3 +296,30 @@ impl Machine {
         }
     }
 }
+
+#[test]
+pub fn test_sc_after_snapshot() {
+    let mut machine = machine_build::int_v2_imacb("tests/programs/sc_after_snapshot");
+    machine.machine.set_max_cycles(5);
+    let ret = machine.run();
+    assert!(ret.is_err());
+    assert_eq!(ret.unwrap_err(), Error::CyclesExceeded);
+
+    let snap = make_snapshot(&mut machine).unwrap();
+    let mut machine_new = TraceMachine::new(
+        DefaultMachineBuilder::new(
+            DefaultCoreMachine::<u64, WXorXMemory<SparseMemory<u64>>>::new(
+                ISA_IMC | ISA_A,
+                VERSION2,
+                u64::max_value(),
+            ),
+        )
+        .instruction_cycle_func(Box::new(constant_cycles))
+        .build(),
+    );
+    resume(&mut machine_new, &snap).unwrap();
+    machine_new.machine.set_max_cycles(20);
+    let ret = machine_new.run();
+    assert!(ret.is_ok());
+    assert_eq!(ret.unwrap(), 0);
+}
diff --git a/tests/test_resume2.rs b/tests/test_resume2.rs
@@ -6,12 +6,12 @@ use ckb_vm::elf::parse_elf;
 use ckb_vm::machine::asm::{AsmCoreMachine, AsmMachine};
 use ckb_vm::machine::trace::TraceMachine;
 use ckb_vm::machine::{
-    CoreMachine, DefaultCoreMachine, DefaultMachine, SupportMachine, VERSION0, VERSION1,
+    CoreMachine, DefaultCoreMachine, DefaultMachine, SupportMachine, VERSION0, VERSION1, VERSION2,
 };
 use ckb_vm::memory::{sparse::SparseMemory, wxorx::WXorXMemory};
 use ckb_vm::registers::{A0, A1, A7};
 use ckb_vm::snapshot2::{DataSource, Snapshot2, Snapshot2Context};
-use ckb_vm::{DefaultMachineBuilder, Error, Register, Syscalls, ISA_IMC};
+use ckb_vm::{DefaultMachineBuilder, Error, Register, Syscalls, ISA_A, ISA_IMC};
 use std::collections::HashMap;
 use std::fs::File;
 use std::io::Read;
@@ -379,7 +379,7 @@ impl MachineTy {
         match self {
             MachineTy::Asm => {
                 let context = Arc::new(Mutex::new(Snapshot2Context::new(data_source)));
-                let asm_core1 = AsmCoreMachine::new(ISA_IMC, version, 0);
+                let asm_core1 = AsmCoreMachine::new(ISA_IMC | ISA_A, version, 0);
                 let core1 = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core1)
                     .instruction_cycle_func(Box::new(constant_cycles))
                     .syscall(Box::new(InsertDataSyscall(context.clone())))
@@ -389,7 +389,9 @@ impl MachineTy {
             MachineTy::Interpreter => {
                 let context = Arc::new(Mutex::new(Snapshot2Context::new(data_source)));
                 let core_machine1 = DefaultCoreMachine::<u64, WXorXMemory<SparseMemory<u64>>>::new(
-                    ISA_IMC, version, 0,
+                    ISA_IMC | ISA_A,
+                    version,
+                    0,
                 );
                 Machine::Interpreter(
                     DefaultMachineBuilder::<DefaultCoreMachine<u64, WXorXMemory<SparseMemory<u64>>>>::new(
@@ -404,7 +406,9 @@ impl MachineTy {
             MachineTy::InterpreterWithTrace => {
                 let context = Arc::new(Mutex::new(Snapshot2Context::new(data_source)));
                 let core_machine1 = DefaultCoreMachine::<u64, WXorXMemory<SparseMemory<u64>>>::new(
-                    ISA_IMC, version, 0,
+                    ISA_IMC | ISA_A,
+                    version,
+                    0,
                 );
                 Machine::InterpreterWithTrace(
                     TraceMachine::new(
@@ -596,3 +600,23 @@ impl Machine {
         Ok(())
     }
 }
+
+#[test]
+pub fn test_sc_after_snapshot2() {
+    let data_source = load_program("tests/programs/sc_after_snapshot");
+
+    let mut machine1 = MachineTy::Interpreter.build(data_source.clone(), VERSION2);
+    machine1.set_max_cycles(5);
+    machine1.load_program(&vec!["main".into()]).unwrap();
+    let result1 = machine1.run();
+    assert!(result1.is_err());
+    assert_eq!(result1.unwrap_err(), Error::CyclesExceeded);
+    let snapshot = machine1.snapshot().unwrap();
+
+    let mut machine2 = MachineTy::Interpreter.build(data_source, VERSION2);
+    machine2.resume(snapshot).unwrap();
+    machine2.set_max_cycles(20);
+    let result2 = machine2.run();
+    assert!(result2.is_ok());
+    assert_eq!(result2.unwrap(), 0);
+}