Skip to content

Commit 01ce222

Browse files
mohansonmohanson
committed
Implement FlattenedArgsReader
1 parent 68a5331 commit 01ce222

19 files changed

+300
-206
lines changed

‎benches/vm_benchmark.rs‎

Lines changed: 41 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,13 @@ use std::fs;
2121
fn interpret_benchmark(c: &mut Criterion) {
2222
c.bench_function("interpret secp256k1_bench", |b| {
2323
let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
24-
let args: Vec<Bytes> = vec!["secp256k1_bench",
25-
"033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
26-
"304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
27-
"foo",
28-
"bar"].into_iter().map(|a| a.into()).collect();
29-
24+
let args: Vec<Bytes> = vec![
25+
"secp256k1_bench",
26+
"033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
27+
"304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
28+
"foo",
29+
"bar"
30+
].into_iter().map(|a| a.into()).collect();
3031
b.iter(|| run::<u64, SparseMemory<u64>>(&buffer, &args[..]).unwrap());
3132
});
3233
}
@@ -35,17 +36,18 @@ fn interpret_benchmark(c: &mut Criterion) {
3536
fn asm_benchmark(c: &mut Criterion) {
3637
c.bench_function("interpret secp256k1_bench via assembly", |b| {
3738
let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
38-
let args: Vec<Bytes> = vec!["secp256k1_bench",
39-
"033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
40-
"304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
41-
"foo",
42-
"bar"].into_iter().map(|a| a.into()).collect();
43-
39+
let args = [
40+
Ok("secp256k1_bench".into()),
41+
Ok("033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into()),
42+
Ok("304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into()),
43+
Ok("foo".into()),
44+
Ok("bar".into()),
45+
].into_iter();
4446
b.iter(|| {
4547
let asm_core = AsmCoreMachine::new(ISA_IMC, VERSION0, u64::MAX);
4648
let core = DefaultMachineBuilder::new(asm_core).build();
4749
let mut machine = AsmMachine::new(core);
48-
machine.load_program(&buffer, &args[..]).unwrap();
50+
machine.load_program(&buffer, args.clone()).unwrap();
4951
machine.run().unwrap()
5052
});
5153
});
@@ -55,17 +57,19 @@ fn asm_benchmark(c: &mut Criterion) {
5557
fn mop_benchmark(c: &mut Criterion) {
5658
c.bench_function("interpret secp256k1_bench via assembly mop", |b| {
5759
let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
58-
let args: Vec<Bytes> = vec!["secp256k1_bench",
59-
"033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
60-
"304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
61-
"foo",
62-
"bar"].into_iter().map(|a| a.into()).collect();
60+
let args = [
61+
Ok("secp256k1_bench".into()),
62+
Ok("033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into()),
63+
Ok("304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into()),
64+
Ok("foo".into()),
65+
Ok("bar".into()),
66+
].into_iter();
6367
b.iter(|| {
6468
let asm_core = AsmCoreMachine::new(ISA_IMC | ISA_B | ISA_MOP, VERSION2, u64::MAX);
6569
let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
6670
.build();
6771
let mut machine = AsmMachine::new(core);
68-
machine.load_program(&buffer, &args).unwrap();
72+
machine.load_program(&buffer, args.clone()).unwrap();
6973
machine.run().unwrap()
7074
});
7175
});
@@ -77,25 +81,27 @@ fn mop_memoized_benchmark(c: &mut Criterion) {
7781
let isa = ISA_IMC | ISA_B | ISA_MOP;
7882
let version = VERSION2;
7983
let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
80-
let args: Vec<Bytes> = vec!["secp256k1_bench",
81-
"033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
82-
"304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
83-
"foo",
84-
"bar"].into_iter().map(|a| a.into()).collect();
84+
let args = [
85+
Ok("secp256k1_bench".into()),
86+
Ok("033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into()),
87+
Ok("304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into()),
88+
Ok("foo".into()),
89+
Ok("bar".into()),
90+
].into_iter();
8591
let mut decoder = MemoizedFixedTraceDecoder::new(build_decoder::<u64>(isa, version));
8692
let asm_core = AsmCoreMachine::new(isa, version, u64::MAX);
8793
let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
8894
.build();
8995
let mut machine = AsmMachine::new(core);
90-
machine.load_program(&buffer, &args).unwrap();
96+
machine.load_program(&buffer, args.clone()).unwrap();
9197
machine.run_with_decoder(&mut decoder).unwrap();
9298

9399
b.iter(|| {
94100
let asm_core = AsmCoreMachine::new(isa, version, u64::MAX);
95101
let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
96102
.build();
97103
let mut machine = AsmMachine::new(core);
98-
machine.load_program(&buffer, &args).unwrap();
104+
machine.load_program(&buffer, args.clone()).unwrap();
99105
decoder.clear_traces();
100106
machine.run_with_decoder(&mut decoder).unwrap()
101107
});
@@ -108,25 +114,27 @@ fn mop_memoized_dynamic_benchmark(c: &mut Criterion) {
108114
let isa = ISA_IMC | ISA_B | ISA_MOP;
109115
let version = VERSION2;
110116
let buffer = fs::read("benches/data/secp256k1_bench").unwrap().into();
111-
let args: Vec<Bytes> = vec!["secp256k1_bench",
112-
"033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f",
113-
"304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3",
114-
"foo",
115-
"bar"].into_iter().map(|a| a.into()).collect();
117+
let args = [
118+
Ok("secp256k1_bench".into()),
119+
Ok("033f8cf9c4d51a33206a6c1c6b27d2cc5129daa19dbd1fc148d395284f6b26411f".into()),
120+
Ok("304402203679d909f43f073c7c1dcf8468a485090589079ee834e6eed92fea9b09b06a2402201e46f1075afa18f306715e7db87493e7b7e779569aa13c64ab3d09980b3560a3".into()),
121+
Ok("foo".into()),
122+
Ok("bar".into()),
123+
].into_iter();
116124
let mut decoder = MemoizedDynamicTraceDecoder::new(build_decoder::<u64>(isa, version));
117125
let asm_core = AsmCoreMachine::new(isa, version, u64::MAX);
118126
let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
119127
.build();
120128
let mut machine = AsmMachine::new(core);
121-
machine.load_program(&buffer, &args).unwrap();
129+
machine.load_program(&buffer, args.clone()).unwrap();
122130
machine.run_with_decoder(&mut decoder).unwrap();
123131

124132
b.iter(|| {
125133
let asm_core = AsmCoreMachine::new(isa, version, u64::MAX);
126134
let core = DefaultMachineBuilder::<Box<AsmCoreMachine>>::new(asm_core)
127135
.build();
128136
let mut machine = AsmMachine::new(core);
129-
machine.load_program(&buffer, &args).unwrap();
137+
machine.load_program(&buffer, args.clone()).unwrap();
130138
decoder.clear_traces();
131139
machine.run_with_decoder(&mut decoder).unwrap()
132140
});

‎examples/check_real_memory.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,10 @@ fn check_asm(memory_size: usize) -> Result<(), ()> {
168168
let core = DefaultMachineBuilder::new(asm_core).build();
169169
let mut machine = AsmMachine::new(core);
170170
machine
171-
.load_program(&Bytes::from(BIN_PATH_BUFFER), &vec![Bytes::from(BIN_NAME)])
171+
.load_program(
172+
&Bytes::from(BIN_PATH_BUFFER),
173+
[Ok(Bytes::from(BIN_NAME))].into_iter(),
174+
)
172175
.unwrap();
173176
let result = machine.run();
174177
assert!(result.is_ok());
@@ -192,7 +195,10 @@ fn check_asm_in_thread(memory_size: usize) -> Result<(), ()> {
192195
let core = DefaultMachineBuilder::new(asm_core).build();
193196
let mut machine = AsmMachine::new(core);
194197
machine
195-
.load_program(&Bytes::from(BIN_PATH_BUFFER), &vec![Bytes::from(BIN_NAME)])
198+
.load_program(
199+
&Bytes::from(BIN_PATH_BUFFER),
200+
[Ok(Bytes::from(BIN_NAME))].into_iter(),
201+
)
196202
.unwrap();
197203
let thread_join_handle = thread::spawn(move || {
198204
let result = machine.run();

‎examples/ckb_vm_runner.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ fn main_asm(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Err
4949
.syscall(Box::new(DebugSyscall {}))
5050
.build();
5151
let mut machine = ckb_vm::machine::asm::AsmMachine::new(core);
52-
machine.load_program(&code, &args)?;
52+
machine.load_program(&code, args.into_iter().map(Ok))?;
5353
let exit = machine.run();
5454
let cycles = machine.machine.cycles();
5555
println!(
@@ -71,7 +71,7 @@ fn main_int(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Err
7171
let machine_builder = ckb_vm::DefaultMachineBuilder::new(core_machine)
7272
.instruction_cycle_func(Box::new(estimate_cycles));
7373
let mut machine = machine_builder.syscall(Box::new(DebugSyscall {})).build();
74-
machine.load_program(&code, &args)?;
74+
machine.load_program(&code, args.into_iter().map(Ok))?;
7575
let exit = machine.run();
7676
let cycles = machine.cycles();
7777
println!(

‎src/lib.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ pub use crate::{
2222
instructions::{Instruction, Register},
2323
machine::{
2424
trace::TraceMachine, CoreMachine, DefaultCoreMachine, DefaultMachine,
25-
DefaultMachineBuilder, InstructionCycleFunc, Machine, SupportMachine,
25+
DefaultMachineBuilder, FlattenedArgsReader, InstructionCycleFunc, Machine, SupportMachine,
2626
},
2727
memory::{flat::FlatMemory, sparse::SparseMemory, wxorx::WXorXMemory, Memory},
2828
syscalls::Syscalls,
@@ -47,7 +47,7 @@ pub fn run<R: Register, M: Memory<REG = R> + Default>(
4747
WXorXMemory::new(M::default()),
4848
);
4949
let mut machine = TraceMachine::new(DefaultMachineBuilder::new(core_machine).build());
50-
machine.load_program(program, args)?;
50+
machine.load_program(program, args.into_iter().map(|e| Ok(e.clone())))?;
5151
machine.run()
5252
}
5353

@@ -63,7 +63,7 @@ pub fn run_with_memory<R: Register, M: Memory<REG = R>>(
6363
WXorXMemory::new(memory),
6464
);
6565
let mut machine = TraceMachine::new(DefaultMachineBuilder::new(core_machine).build());
66-
machine.load_program(program, args)?;
66+
machine.load_program(program, args.into_iter().map(|e| Ok(e.clone())))?;
6767
machine.run()
6868
}
6969

‎src/machine/asm/mod.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -666,15 +666,19 @@ impl AsmMachine {
666666
self.machine.inner.max_cycles = cycles;
667667
}
668668

669-
pub fn load_program(&mut self, program: &Bytes, args: &[Bytes]) -> Result<u64, Error> {
669+
pub fn load_program(
670+
&mut self,
671+
program: &Bytes,
672+
args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
673+
) -> Result<u64, Error> {
670674
self.machine.load_program(program, args)
671675
}
672676

673677
pub fn load_program_with_metadata(
674678
&mut self,
675679
program: &Bytes,
676680
metadata: &ProgramMetadata,
677-
args: &[Bytes],
681+
args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
678682
) -> Result<u64, Error> {
679683
self.machine
680684
.load_program_with_metadata(program, metadata, args)

‎src/machine/mod.rs‎

Lines changed: 75 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ use super::debugger::Debugger;
1212
use super::decoder::{build_decoder, InstDecoder};
1313
use super::elf::{parse_elf, LoadingAction, ProgramMetadata};
1414
use super::instructions::{execute, Instruction, Register};
15-
use super::memory::Memory;
15+
use super::memory::{load_c_string, Memory};
1616
use super::syscalls::Syscalls;
1717
use super::{
1818
registers::{A0, A7, REGISTER_ABI_NAMES, SP},
@@ -171,7 +171,7 @@ pub trait SupportMachine: CoreMachine {
171171

172172
fn initialize_stack(
173173
&mut self,
174-
args: &[Bytes],
174+
args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
175175
stack_start: u64,
176176
stack_size: u64,
177177
) -> Result<u64, Error> {
@@ -183,7 +183,7 @@ pub trait SupportMachine: CoreMachine {
183183
// reading "argc" will return an unexpected data. This situation is not very common.
184184
//
185185
// See https://github.com/nervosnetwork/ckb-vm/issues/106 for more details.
186-
if self.version() >= VERSION1 && args.is_empty() {
186+
if self.version() >= VERSION1 && args.len() == 0 {
187187
let argc_size = u64::from(Self::REG::BITS / 8);
188188
let origin_sp = stack_start + stack_size;
189189
let unaligned_sp_address = origin_sp - argc_size;
@@ -200,15 +200,21 @@ pub trait SupportMachine: CoreMachine {
200200
// of each argv object.
201201
let mut values = vec![Self::REG::from_u64(args.len() as u64)];
202202
for arg in args {
203+
let arg = arg?;
203204
let len = Self::REG::from_u64(arg.len() as u64 + 1);
204205
let address = self.registers()[SP].overflowing_sub(&len);
205206

206-
self.memory_mut().store_bytes(address.to_u64(), arg)?;
207+
self.memory_mut().store_bytes(address.to_u64(), &arg)?;
207208
self.memory_mut()
208209
.store_byte(address.to_u64() + arg.len() as u64, 1, 0)?;
209210

210211
values.push(address.clone());
211-
self.set_register(SP, address);
212+
self.set_register(SP, address.clone());
213+
214+
if self.version() >= VERSION2 && address.to_u64() < stack_start {
215+
// Provides an early exit to large argv array.
216+
return Err(Error::MemOutOfStack);
217+
}
212218
}
213219
if self.version() >= VERSION1 {
214220
// There are 2 standard requirements of the initialized stack:
@@ -246,7 +252,7 @@ pub trait SupportMachine: CoreMachine {
246252
self.set_register(SP, address);
247253
}
248254
if self.registers()[SP].to_u64() < stack_start {
249-
// args exceed stack size
255+
// Args exceed stack size.
250256
return Err(Error::MemOutOfStack);
251257
}
252258
Ok(stack_start + stack_size - self.registers()[SP].to_u64())
@@ -575,7 +581,11 @@ impl<Inner: CoreMachine> Display for DefaultMachine<Inner> {
575581
}
576582

577583
impl<Inner: SupportMachine> DefaultMachine<Inner> {
578-
pub fn load_program(&mut self, program: &Bytes, args: &[Bytes]) -> Result<u64, Error> {
584+
pub fn load_program(
585+
&mut self,
586+
program: &Bytes,
587+
args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
588+
) -> Result<u64, Error> {
579589
let elf_bytes = self.load_elf(program, true)?;
580590
let stack_bytes = self.initialize(args)?;
581591
let bytes = elf_bytes.checked_add(stack_bytes).ok_or_else(|| {
@@ -590,7 +600,7 @@ impl<Inner: SupportMachine> DefaultMachine<Inner> {
590600
&mut self,
591601
program: &Bytes,
592602
metadata: &ProgramMetadata,
593-
args: &[Bytes],
603+
args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
594604
) -> Result<u64, Error> {
595605
let elf_bytes = self.load_binary(program, metadata, true)?;
596606
let stack_bytes = self.initialize(args)?;
@@ -602,7 +612,10 @@ impl<Inner: SupportMachine> DefaultMachine<Inner> {
602612
Ok(bytes)
603613
}
604614

605-
fn initialize(&mut self, args: &[Bytes]) -> Result<u64, Error> {
615+
fn initialize(
616+
&mut self,
617+
args: impl ExactSizeIterator<Item = Result<Bytes, Error>>,
618+
) -> Result<u64, Error> {
606619
for syscall in &mut self.syscalls {
607620
syscall.initialize(&mut self.inner)?;
608621
}
@@ -766,6 +779,59 @@ impl Pause {
766779
}
767780
}
768781

782+
pub struct FlattenedArgsReader<'a, M: Memory> {
783+
memory: &'a mut M,
784+
argc: M::REG,
785+
argv: M::REG,
786+
cidx: M::REG,
787+
}
788+
789+
impl<'a, M: Memory> FlattenedArgsReader<'a, M> {
790+
pub fn new(memory: &'a mut M, argc: M::REG, argv: M::REG) -> Self {
791+
Self {
792+
memory,
793+
argc,
794+
argv,
795+
cidx: M::REG::zero(),
796+
}
797+
}
798+
}
799+
800+
impl<'a, M: Memory> Iterator for FlattenedArgsReader<'a, M> {
801+
type Item = Result<Bytes, Error>;
802+
803+
fn next(&mut self) -> Option<Self::Item> {
804+
if self.cidx.ge(&self.argc).to_u8() == 1 {
805+
return None;
806+
}
807+
let addr = match M::REG::BITS {
808+
32 => self.memory.load32(&self.argv),
809+
64 => self.memory.load64(&self.argv),
810+
_ => unreachable!(),
811+
};
812+
if let Err(err) = addr {
813+
return Some(Err(err));
814+
};
815+
let addr = addr.unwrap();
816+
let cstr = load_c_string(self.memory, &addr);
817+
if let Err(err) = cstr {
818+
return Some(Err(err));
819+
};
820+
let cstr = cstr.unwrap();
821+
self.cidx = self.cidx.overflowing_add(&M::REG::from_u8(1));
822+
self.argv = self
823+
.argv
824+
.overflowing_add(&M::REG::from_u8(M::REG::BITS / 8));
825+
Some(Ok(cstr))
826+
}
827+
}
828+
829+
impl<'a, M: Memory> ExactSizeIterator for FlattenedArgsReader<'a, M> {
830+
fn len(&self) -> usize {
831+
self.argc.to_u64() as usize
832+
}
833+
}
834+
769835
#[cfg(test)]
770836
mod tests {
771837
use std::sync::atomic::AtomicU8;

0 commit comments

Comments
 (0)