Implement FlattenedArgsReader

nervosnetwork · Jan 16, 2025 · b1b4b00 · b1b4b00
1 parent 68a5331
commit b1b4b00
Show file tree

Hide file tree

Showing 20 changed files with 301 additions and 207 deletions.
diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml
@@ -69,7 +69,7 @@ jobs:
           ln -snf .. ckb-vm-test-suite/ckb-vm
           docker run --rm -v `pwd`:/code nervos/ckb-riscv-gnu-toolchain:bionic-20210804 cp -r /riscv /code/riscv
           cd ckb-vm-test-suite
-          git checkout 2ef749d0f580958043264a798e560fe25fec0c6e
+          git checkout 898edc351eeb4de974ca4f0ff8d1e4943a95aecb
           git submodule update --init --recursive
           RISCV=`pwd`/../riscv ./test.sh
 

diff --git a/benches/vm_benchmark.rs b/benches/vm_benchmark.rs
@@ -21,12 +21,13 @@ use std::fs;
 fn interpret_benchmark(c: &mut Criterion) {
     c.bench_function("interpret secp256k1_bench", |b| {
         let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
-        let args: Vec<Bytes> = vec!["secp256k1_bench",
-                                      "033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
-                                      "304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
-                                      "foo",
-                                      "bar"].into_iter().map(|a| a.into()).collect();
-
+        let args: Vec<Bytes> = vec![
+            "secp256k1_bench",
+            "033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
+            "304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
+            "foo",
+            "bar"
+        ].into_iter().map(|a| a.into()).collect();
         b.iter(|| run::<u64, SparseMemory<u64>>(&buffer, &args[..]).unwrap());
     });
 }
@@ -35,17 +36,18 @@ fn interpret_benchmark(c: &mut Criterion) {
 fn asm_benchmark(c: &mut Criterion) {
     c.bench_function("interpret secp256k1_bench via assembly", |b| {
         let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
-        let args: Vec<Bytes> = vec!["secp256k1_bench",
-                                      "033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
-                                      "304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
-                                      "foo",
-                                      "bar"].into_iter().map(|a| a.into()).collect();
-
+        let args = [
+            Ok("secp256k1_bench".into()),
+            Ok("033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into()),
+            Ok("304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into()),
+            Ok("foo".into()),
+            Ok("bar".into()),
+        ].into_iter();
         b.iter(|| {
             let asm_core = AsmCoreMachine::new(ISA_IMC, VERSION0, u64::MAX);
             let core = DefaultMachineBuilder::new(asm_core).build();
             let mut machine = AsmMachine::new(core);
-            machine.load_program(&buffer, &args[..]).unwrap();
+            machine.load_program(&buffer, args.clone()).unwrap();
             machine.run().unwrap()
         });
     });
@@ -55,17 +57,19 @@ fn asm_benchmark(c: &mut Criterion) {
 fn mop_benchmark(c: &mut Criterion) {
     c.bench_function("interpret secp256k1_bench via assembly mop", |b| {
         let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
-        let args: Vec<Bytes> = vec!["secp256k1_bench",
-                                      "033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
-                                      "304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
-                                      "foo",
-                                      "bar"].into_iter().map(|a| a.into()).collect();
+        let args = [
+            Ok("secp256k1_bench".into()),
+            Ok("033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into()),
+            Ok("304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into()),
+            Ok("foo".into()),
+            Ok("bar".into()),
+        ].into_iter();
         b.iter(|| {
             let asm_core = AsmCoreMachine::new(ISA_IMC | ISA_B | ISA_MOP, VERSION2, u64::MAX);
             let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
                 .build();
             let mut machine = AsmMachine::new(core);
-            machine.load_program(&buffer, &args).unwrap();
+            machine.load_program(&buffer, args.clone()).unwrap();
             machine.run().unwrap()
         });
     });
@@ -77,25 +81,27 @@ fn mop_memoized_benchmark(c: &mut Criterion) {
         let isa = ISA_IMC | ISA_B | ISA_MOP;
         let version = VERSION2;
         let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
-        let args: Vec<Bytes> = vec!["secp256k1_bench",
-                                      "033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
-                                      "304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
-                                      "foo",
-                                      "bar"].into_iter().map(|a| a.into()).collect();
+        let args = [
+            Ok("secp256k1_bench".into()),
+            Ok("033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into()),
+            Ok("304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into()),
+            Ok("foo".into()),
+            Ok("bar".into()),
+        ].into_iter();
         let mut decoder = MemoizedFixedTraceDecoder::new(build_decoder::<u64>(isa, version));
         let asm_core = AsmCoreMachine::new(isa, version, u64::MAX);
         let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
             .build();
         let mut machine = AsmMachine::new(core);
-        machine.load_program(&buffer, &args).unwrap();
+        machine.load_program(&buffer, args.clone()).unwrap();
         machine.run_with_decoder(&mut decoder).unwrap();
 
         b.iter(|| {
             let asm_core = AsmCoreMachine::new(isa, version, u64::MAX);
             let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
                 .build();
             let mut machine = AsmMachine::new(core);
-            machine.load_program(&buffer, &args).unwrap();
+            machine.load_program(&buffer, args.clone()).unwrap();
             decoder.clear_traces();
             machine.run_with_decoder(&mut decoder).unwrap()
         });
@@ -108,25 +114,27 @@ fn mop_memoized_dynamic_benchmark(c: &mut Criterion) {
         let isa = ISA_IMC | ISA_B | ISA_MOP;
         let version = VERSION2;
         let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
-        let args: Vec<Bytes> = vec!["secp256k1_bench",
-                                      "033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
-                                      "304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
-                                      "foo",
-                                      "bar"].into_iter().map(|a| a.into()).collect();
+        let args = [
+            Ok("secp256k1_bench".into()),
+            Ok("033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into()),
+            Ok("304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into()),
+            Ok("foo".into()),
+            Ok("bar".into()),
+        ].into_iter();
         let mut decoder = MemoizedDynamicTraceDecoder::new(build_decoder::<u64>(isa, version));
         let asm_core = AsmCoreMachine::new(isa, version, u64::MAX);
         let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
             .build();
         let mut machine = AsmMachine::new(core);
-        machine.load_program(&buffer, &args).unwrap();
+        machine.load_program(&buffer, args.clone()).unwrap();
         machine.run_with_decoder(&mut decoder).unwrap();
 
         b.iter(|| {
             let asm_core = AsmCoreMachine::new(isa, version, u64::MAX);
             let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
                 .build();
             let mut machine = AsmMachine::new(core);
-            machine.load_program(&buffer, &args).unwrap();
+            machine.load_program(&buffer, args.clone()).unwrap();
             decoder.clear_traces();
             machine.run_with_decoder(&mut decoder).unwrap()
         });

diff --git a/examples/check_real_memory.rs b/examples/check_real_memory.rs
@@ -168,7 +168,10 @@ fn check_asm(memory_size: usize) -> Result<(), ()> {
         let core = DefaultMachineBuilder::new(asm_core).build();
         let mut machine = AsmMachine::new(core);
         machine
-            .load_program(&Bytes::from(BIN_PATH_BUFFER), &vec![Bytes::from(BIN_NAME)])
+            .load_program(
+                &Bytes::from(BIN_PATH_BUFFER),
+                [Ok(Bytes::from(BIN_NAME))].into_iter(),
+            )
             .unwrap();
         let result = machine.run();
         assert!(result.is_ok());
@@ -192,7 +195,10 @@ fn check_asm_in_thread(memory_size: usize) -> Result<(), ()> {
         let core = DefaultMachineBuilder::new(asm_core).build();
         let mut machine = AsmMachine::new(core);
         machine
-            .load_program(&Bytes::from(BIN_PATH_BUFFER), &vec![Bytes::from(BIN_NAME)])
+            .load_program(
+                &Bytes::from(BIN_PATH_BUFFER),
+                [Ok(Bytes::from(BIN_NAME))].into_iter(),
+            )
             .unwrap();
         let thread_join_handle = thread::spawn(move || {
             let result = machine.run();

diff --git a/examples/ckb_vm_runner.rs b/examples/ckb_vm_runner.rs
@@ -49,7 +49,7 @@ fn main_asm(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Err
         .syscall(Box::new(DebugSyscall {}))
         .build();
     let mut machine = ckb_vm::machine::asm::AsmMachine::new(core);
-    machine.load_program(&code, &args)?;
+    machine.load_program(&code, args.into_iter().map(Ok))?;
     let exit = machine.run();
     let cycles = machine.machine.cycles();
     println!(
@@ -71,7 +71,7 @@ fn main_int(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Err
     let machine_builder = ckb_vm::DefaultMachineBuilder::new(core_machine)
         .instruction_cycle_func(Box::new(estimate_cycles));
     let mut machine = machine_builder.syscall(Box::new(DebugSyscall {})).build();
-    machine.load_program(&code, &args)?;
+    machine.load_program(&code, args.into_iter().map(Ok))?;
     let exit = machine.run();
     let cycles = machine.cycles();
     println!(

diff --git a/src/lib.rs b/src/lib.rs
@@ -22,7 +22,7 @@ pub use crate::{
     instructions::{Instruction, Register},
     machine::{
         trace::TraceMachine, CoreMachine, DefaultCoreMachine, DefaultMachine,
-        DefaultMachineBuilder, InstructionCycleFunc, Machine, SupportMachine,
+        DefaultMachineBuilder, FlattenedArgsReader, InstructionCycleFunc, Machine, SupportMachine,
     },
     memory::{flat::FlatMemory, sparse::SparseMemory, wxorx::WXorXMemory, Memory},
     syscalls::Syscalls,
@@ -47,7 +47,7 @@ pub fn run<R: Register, M: Memory<REG = R> + Default>(
         WXorXMemory::new(M::default()),
     );
     let mut machine = TraceMachine::new(DefaultMachineBuilder::new(core_machine).build());
-    machine.load_program(program, args)?;
+    machine.load_program(program, args.iter().map(|e| Ok(e.clone())))?;
     machine.run()
 }
 
@@ -63,7 +63,7 @@ pub fn run_with_memory<R: Register, M: Memory<REG = R>>(
         WXorXMemory::new(memory),
     );
     let mut machine = TraceMachine::new(DefaultMachineBuilder::new(core_machine).build());
-    machine.load_program(program, args)?;
+    machine.load_program(program, args.iter().map(|e| Ok(e.clone())))?;
     machine.run()
 }
 

diff --git a/src/machine/asm/mod.rs b/src/machine/asm/mod.rs
@@ -666,15 +666,19 @@ impl AsmMachine {
         self.machine.inner.max_cycles = cycles;
     }
 
-    pub fn load_program(&mut self, program: &Bytes, args: &[Bytes]) -> Result<u64, Error> {
+    pub fn load_program(
+        &mut self,
+        program: &Bytes,
+        args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
+    ) -> Result<u64, Error> {
         self.machine.load_program(program, args)
     }
 
     pub fn load_program_with_metadata(
         &mut self,
         program: &Bytes,
         metadata: &ProgramMetadata,
-        args: &[Bytes],
+        args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
     ) -> Result<u64, Error> {
         self.machine
             .load_program_with_metadata(program, metadata, args)

diff --git a/src/machine/mod.rs b/src/machine/mod.rs
@@ -12,7 +12,7 @@ use super::debugger::Debugger;
 use super::decoder::{build_decoder, InstDecoder};
 use super::elf::{parse_elf, LoadingAction, ProgramMetadata};
 use super::instructions::{execute, Instruction, Register};
-use super::memory::Memory;
+use super::memory::{load_c_string, Memory};
 use super::syscalls::Syscalls;
 use super::{
     registers::{A0, A7, REGISTER_ABI_NAMES, SP},
@@ -171,7 +171,7 @@ pub trait SupportMachine: CoreMachine {
 
     fn initialize_stack(
         &mut self,
-        args: &[Bytes],
+        args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
         stack_start: u64,
         stack_size: u64,
     ) -> Result<u64, Error> {
@@ -183,7 +183,7 @@ pub trait SupportMachine: CoreMachine {
         // reading "argc" will return an unexpected data. This situation is not very common.
         //
         // See https://github.com/nervosnetwork/ckb-vm/issues/106 for more details.
-        if self.version() >= VERSION1 && args.is_empty() {
+        if self.version() >= VERSION1 && args.len() == 0 {
             let argc_size = u64::from(Self::REG::BITS / 8);
             let origin_sp = stack_start + stack_size;
             let unaligned_sp_address = origin_sp - argc_size;
@@ -200,15 +200,21 @@ pub trait SupportMachine: CoreMachine {
         // of each argv object.
         let mut values = vec![Self::REG::from_u64(args.len() as u64)];
         for arg in args {
+            let arg = arg?;
             let len = Self::REG::from_u64(arg.len() as u64 + 1);
             let address = self.registers()[SP].overflowing_sub(&len);
 
-            self.memory_mut().store_bytes(address.to_u64(), arg)?;
+            self.memory_mut().store_bytes(address.to_u64(), &arg)?;
             self.memory_mut()
                 .store_byte(address.to_u64() + arg.len() as u64, 1, 0)?;
 
             values.push(address.clone());
-            self.set_register(SP, address);
+            self.set_register(SP, address.clone());
+
+            if self.version() >= VERSION2 && address.to_u64() < stack_start {
+                // Provides an early exit to large argv array.
+                return Err(Error::MemOutOfStack);
+            }
         }
         if self.version() >= VERSION1 {
             // There are 2 standard requirements of the initialized stack:
@@ -246,7 +252,7 @@ pub trait SupportMachine: CoreMachine {
             self.set_register(SP, address);
         }
         if self.registers()[SP].to_u64() < stack_start {
-            // args exceed stack size
+            // Args exceed stack size.
             return Err(Error::MemOutOfStack);
         }
         Ok(stack_start + stack_size - self.registers()[SP].to_u64())
@@ -575,7 +581,11 @@ impl<Inner: CoreMachine> Display for DefaultMachine<Inner> {
 }
 
 impl<Inner: SupportMachine> DefaultMachine<Inner> {
-    pub fn load_program(&mut self, program: &Bytes, args: &[Bytes]) -> Result<u64, Error> {
+    pub fn load_program(
+        &mut self,
+        program: &Bytes,
+        args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
+    ) -> Result<u64, Error> {
         let elf_bytes = self.load_elf(program, true)?;
         let stack_bytes = self.initialize(args)?;
         let bytes = elf_bytes.checked_add(stack_bytes).ok_or_else(|| {
@@ -590,7 +600,7 @@ impl<Inner: SupportMachine> DefaultMachine<Inner> {
         &mut self,
         program: &Bytes,
         metadata: &ProgramMetadata,
-        args: &[Bytes],
+        args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
     ) -> Result<u64, Error> {
         let elf_bytes = self.load_binary(program, metadata, true)?;
         let stack_bytes = self.initialize(args)?;
@@ -602,7 +612,10 @@ impl<Inner: SupportMachine> DefaultMachine<Inner> {
         Ok(bytes)
     }
 
-    fn initialize(&mut self, args: &[Bytes]) -> Result<u64, Error> {
+    fn initialize(
+        &mut self,
+        args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
+    ) -> Result<u64, Error> {
         for syscall in &mut self.syscalls {
             syscall.initialize(&mut self.inner)?;
         }
@@ -766,6 +779,59 @@ impl Pause {
     }
 }
 
+pub struct FlattenedArgsReader<'a, M: Memory> {
+    memory: &'a mut M,
+    argc: M::REG,
+    argv: M::REG,
+    cidx: M::REG,
+}
+
+impl<'a, M: Memory> FlattenedArgsReader<'a, M> {
+    pub fn new(memory: &'a mut M, argc: M::REG, argv: M::REG) -> Self {
+        Self {
+            memory,
+            argc,
+            argv,
+            cidx: M::REG::zero(),
+        }
+    }
+}
+
+impl<'a, M: Memory> Iterator for FlattenedArgsReader<'a, M> {
+    type Item = Result<Bytes, Error>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.cidx.ge(&self.argc).to_u8() == 1 {
+            return None;
+        }
+        let addr = match M::REG::BITS {
+            32 => self.memory.load32(&self.argv),
+            64 => self.memory.load64(&self.argv),
+            _ => unreachable!(),
+        };
+        if let Err(err) = addr {
+            return Some(Err(err));
+        };
+        let addr = addr.unwrap();
+        let cstr = load_c_string(self.memory, &addr);
+        if let Err(err) = cstr {
+            return Some(Err(err));
+        };
+        let cstr = cstr.unwrap();
+        self.cidx = self.cidx.overflowing_add(&M::REG::from_u8(1));
+        self.argv = self
+            .argv
+            .overflowing_add(&M::REG::from_u8(M::REG::BITS / 8));
+        Some(Ok(cstr))
+    }
+}
+
+impl<'a, M: Memory> ExactSizeIterator for FlattenedArgsReader<'a, M> {
+    fn len(&self) -> usize {
+        self.argc.to_u64() as usize
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::atomic::AtomicU8;