From 950c57e6ffa9b84239e23d98d3274ab35314ebf6 Mon Sep 17 00:00:00 2001
From: mohanson <mohanson@outlook.com>
Date: Wed, 21 May 2025 11:17:55 +0800
Subject: [PATCH] Refactor examples/ckb_vm_runner

---
 Cargo.toml                |  1 +
 examples/ckb_vm_runner.rs | 73 ++++++++++++++++++++++++++++++---------
 2 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 63824653..da7c499b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -37,6 +37,7 @@ rand = "0.7.3"
 cc = "1.0"
 
 [dev-dependencies]
+argparse = "0.2"
 criterion = "0.5.1"
 proptest = "1.5.0"
 
diff --git a/examples/ckb_vm_runner.rs b/examples/ckb_vm_runner.rs
index 6e6a096c..b57527de 100644
--- a/examples/ckb_vm_runner.rs
+++ b/examples/ckb_vm_runner.rs
@@ -40,9 +40,9 @@ impl<Mac: SupportMachine> Syscalls<Mac> for DebugSyscall {
 }
 
 #[cfg(has_asm)]
-fn main_asm(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Error>> {
+fn main_asm64(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Error>> {
     let asm_core = <ckb_vm::machine::asm::AsmCoreMachine as SupportMachine>::new(
-        ckb_vm::ISA_IMC | ckb_vm::ISA_B | ckb_vm::ISA_MOP,
+        ckb_vm::ISA_IMC | ckb_vm::ISA_B | ckb_vm::ISA_A | ckb_vm::ISA_MOP,
         ckb_vm::machine::VERSION2,
         u64::MAX,
     );
@@ -55,7 +55,7 @@ fn main_asm(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Err
     let exit = machine.run();
     let cycles = machine.machine.cycles();
     println!(
-        "asm exit={:?} cycles={:?} r[a1]={:?}",
+        "asm64 exit={:?} cycles={:?} r[a1]={:?}",
         exit,
         cycles,
         machine.machine.registers()[ckb_vm::registers::A1]
@@ -64,9 +64,34 @@ fn main_asm(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Err
 }
 
 #[cfg(not(has_asm))]
-fn main_int(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Error>> {
+fn main_asm64(_: Bytes, _: Vec<Bytes>) -> Result<(), Box<dyn std::error::Error>> {
+    panic!("please use --features=asm to enable asm support.")
+}
+
+fn main_interpreter32(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Error>> {
+    let core_machine = ckb_vm::DefaultCoreMachine::<u32, ckb_vm::SparseMemory<u32>>::new(
+        ckb_vm::ISA_IMC | ckb_vm::ISA_B | ckb_vm::ISA_A | ckb_vm::ISA_MOP,
+        ckb_vm::machine::VERSION2,
+        u64::MAX,
+    );
+    let machine_builder = ckb_vm::RustDefaultMachineBuilder::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.into_iter().map(Ok))?;
+    let exit = machine.run();
+    let cycles = machine.cycles();
+    println!(
+        "interpreter32 exit={:?} cycles={:?} r[a1]={:?}",
+        exit,
+        cycles,
+        machine.registers()[ckb_vm::registers::A1]
+    );
+    std::process::exit(exit? as i32);
+}
+
+fn main_interpreter64(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Error>> {
     let core_machine = ckb_vm::DefaultCoreMachine::<u64, ckb_vm::SparseMemory<u64>>::new(
-        ckb_vm::ISA_IMC | ckb_vm::ISA_B | ckb_vm::ISA_MOP,
+        ckb_vm::ISA_IMC | ckb_vm::ISA_B | ckb_vm::ISA_A | ckb_vm::ISA_MOP,
         ckb_vm::machine::VERSION2,
         u64::MAX,
     );
@@ -77,7 +102,7 @@ fn main_int(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Err
     let exit = machine.run();
     let cycles = machine.cycles();
     println!(
-        "int exit={:?} cycles={:?} r[a1]={:?}",
+        "interpreter64 exit={:?} cycles={:?} r[a1]={:?}",
         exit,
         cycles,
         machine.registers()[ckb_vm::registers::A1]
@@ -86,16 +111,30 @@ fn main_int(code: Bytes, args: Vec<Bytes>) -> Result<(), Box<dyn std::error::Err
 }
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let args: Vec<String> = std::env::args().collect();
-    let code = std::fs::read(&args[1])?.into();
-    let riscv_args: Vec<Bytes> = if args.len() > 2 {
-        (&args[2..]).into_iter().map(|s| s.clone().into()).collect()
-    } else {
-        Vec::new()
-    };
-    #[cfg(has_asm)]
-    main_asm(code, riscv_args)?;
-    #[cfg(not(has_asm))]
-    main_int(code, riscv_args)?;
+    let mut mode = String::from("asm64");
+    let mut program_path = String::from("");
+    let mut program_args: Vec<String> = vec![];
+    {
+        let mut ap = argparse::ArgumentParser::new();
+        ap.set_description("CKB VM Runner");
+        ap.refer(&mut mode).add_option(
+            &["--mode"],
+            argparse::Store,
+            "[asm64, interpreter32, interpreter64]",
+        );
+        ap.refer(&mut program_path)
+            .add_argument("program_path", argparse::Store, "program path");
+        ap.refer(&mut program_args)
+            .add_argument("program_args", argparse::List, "program args");
+        ap.parse_args_or_exit();
+    }
+    let code = std::fs::read(&program_path)?.into();
+    let program_args = program_args.into_iter().map(|e| e.into()).collect();
+    match mode.as_str() {
+        "asm64" => main_asm64(code, program_args)?,
+        "interpreter32" => main_interpreter32(code, program_args)?,
+        "interpreter64" => main_interpreter64(code, program_args)?,
+        _ => panic!("mode should be one of [asm64, interpreter32, interpreter64]"),
+    }
     Ok(())
 }