small refactor in the vm (#40)

kkrt-labs · Nov 12, 2024 · 0a3ffaf · 0a3ffaf
1 parent 7ec6146
commit 0a3ffaf
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 22 deletions.
diff --git a/crates/brainfuck_vm/src/compiler.rs b/crates/brainfuck_vm/src/compiler.rs
@@ -2,15 +2,15 @@
 
 use stwo_prover::core::fields::m31::BaseField;
 
+#[derive(Debug, Clone, Default)]
 pub struct Compiler {
     code: Vec<char>,
     instructions: Vec<BaseField>,
 }
 
 impl Compiler {
     pub fn new(code: &str) -> Self {
-        let trimmed_code = code.chars().filter(|c| !c.is_whitespace()).collect();
-        Self { code: trimmed_code, instructions: vec![] }
+        Self { code: code.chars().filter(|c| !c.is_whitespace()).collect(), ..Default::default() }
     }
 
     pub fn compile(&mut self) -> Vec<BaseField> {
@@ -25,9 +25,8 @@ impl Compiler {
                 }
                 ']' => {
                     let start_pos = loop_stack.pop().unwrap();
-                    let loop_end_pos = self.instructions.len() + 1;
-                    self.instructions[start_pos] = BaseField::from((loop_end_pos - 1) as u32);
-                    self.instructions.push(BaseField::from((start_pos + 1) as u32));
+                    self.instructions[start_pos] = BaseField::from(self.instructions.len());
+                    self.instructions.push(BaseField::from(start_pos + 1));
                 }
                 _ => (),
             }

diff --git a/crates/brainfuck_vm/src/registers.rs b/crates/brainfuck_vm/src/registers.rs
@@ -1,9 +1,8 @@
 // Adapted from rkdud007 brainfuck-zkvm https://github.com/rkdud007/brainfuck-zkvm/blob/main/src/registers.rs
 
-use num_traits::identities::Zero;
 use stwo_prover::core::fields::m31::BaseField;
 
-#[derive(PartialEq, Eq, Clone)]
+#[derive(PartialEq, Eq, Clone, Default)]
 pub struct Registers {
     /// Clock Cycle Counter
     pub clk: BaseField,
@@ -21,23 +20,9 @@ pub struct Registers {
     pub mvi: BaseField,
 }
 
-impl Default for Registers {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 impl Registers {
     pub fn new() -> Self {
-        Self {
-            clk: BaseField::zero(),
-            ip: BaseField::zero(),
-            ci: BaseField::zero(),
-            ni: BaseField::zero(),
-            mp: BaseField::zero(),
-            mv: BaseField::zero(),
-            mvi: BaseField::zero(),
-        }
+        Self::default()
     }
 }