Add and Sub opcodes (#16)

* Add initial add opcode * Begin operand indirection * Add address for dest * Fix compilation error * Format code * Run clippy * Add target for asm (#17) * Add sub program * Add 'add' program * Update makefile with asm target * Update test setup, add failing sub test and 'add' test * Cargo fmt * Remove echo * Clippy * Update test setup, add failing sub test and 'add' test * Remove echo * Clippy * Add sub implementation * Make sub test work * Add should be zero and modify simple test * Modify tests * Format * Working tests * Lint --------- Co-authored-by: Gianbelinche <[email protected]> Co-authored-by: Francisco Krause Arnim <[email protected]> Co-authored-by: Francisco Krause Arnim <[email protected]>
lambdaclass · Jun 12, 2024 · f3a3034 · f3a3034
1 parent 0987aff
commit f3a3034
Show file tree

Hide file tree

Showing 12 changed files with 304 additions and 46 deletions.
diff --git a/Makefile b/Makefile
@@ -1,16 +1,30 @@
-.PHONY: clean
+.PHONY: clean lint compile-programs-asm compile-programs-yul
 
 ARTIFACTS_DIR=./program_artifacts
 PROGRAMS_DIR=./programs
-ZKSOLC_FLAGS=--asm --bin --yul --overwrite
+ZKSOLC_YUL_FLAGS=--asm --bin --yul --overwrite
+ZKSOLC_ASM_FLAGS=--zkasm --bin --overwrite
 
-PROGRAMS = $(wildcard $(PROGRAMS_DIR)/*.yul)
-ARTIFACTS = $(patsubst $(PROGRAMS_DIR)/%.yul, $(ARTIFACTS_DIR)/%.artifacts, $(PROGRAMS))
+YUL_PROGRAMS = $(wildcard $(PROGRAMS_DIR)/*.yul)
+ASM_PROGRAMS = $(wildcard $(PROGRAMS_DIR)/*.zasm)
+ARTIFACTS_YUL = $(patsubst $(PROGRAMS_DIR)/%.yul, $(ARTIFACTS_DIR)/%.artifacts.yul, $(YUL_PROGRAMS))
+ARTIFACTS_ASM = $(patsubst $(PROGRAMS_DIR)/%.zasm, $(ARTIFACTS_DIR)/%.artifacts.zasm, $(ASM_PROGRAMS))
 
-compile-programs: $(ARTIFACTS)
+compile-programs-asm: $(ARTIFACTS_ASM)
+compile-programs-yul: $(ARTIFACTS_YUL)
 
-$(ARTIFACTS_DIR)/%.artifacts: $(PROGRAMS_DIR)/%.yul
-	zksolc $(ZKSOLC_FLAGS) $< -o $@ --debug-output-dir $@
+compile-programs: clean compile-programs-asm compile-programs-yul
+
+$(ARTIFACTS_DIR)/%.artifacts.yul: $(PROGRAMS_DIR)/%.yul
+	zksolc $(ZKSOLC_YUL_FLAGS) $< -o $@ --debug-output-dir $@
+
+$(ARTIFACTS_DIR)/%.artifacts.zasm: $(PROGRAMS_DIR)/%.zasm
+	zksolc $(ZKSOLC_ASM_FLAGS)  $< -o $@ --debug-output-dir $@
 
 clean:
 	-rm -rf $(ARTIFACTS_DIR)
+
+lint:
+	cargo clippy --workspace --all-features --benches --examples --tests -- -D warnings
+test: clean compile-programs
+	cargo test
diff --git a/programs/add.zasm b/programs/add.zasm
@@ -0,0 +1,13 @@
+	.text
+	.file	"add.zasm"
+	.globl	__entry
+__entry:
+.func_begin0:
+	add	3, r0, r1
+	sstore	r0, r1
+	add	r0, r0, r1
+	ret
+.func_end0:
+
+	.note.GNU-stack
+	.rodata
diff --git a/programs/sub_and_add.zasm b/programs/sub_and_add.zasm
@@ -0,0 +1,14 @@
+	.text
+	.file	"sub_and_add.zasm"
+	.globl	__entry
+__entry:
+
+.func_begin0:
+    sub r1, r1, r1
+    add 1, r1, r1
+	sstore	r0, r1
+	ret
+
+.func_end0:
+	.note.GNU-stack
+	.rodata
diff --git a/programs/sub_should_be_zero.zasm b/programs/sub_should_be_zero.zasm
@@ -0,0 +1,12 @@
+	.text
+	.file	"sub.zasm"
+	.globl	__entry
+__entry:
+.func_begin0:
+    sub 3, r0, r1
+    sub 3, r1, r1
+	sstore	r0, r1
+	ret
+.func_end0:
+	.note.GNU-stack
+	.rodata
diff --git a/programs/sub_simple.zasm b/programs/sub_simple.zasm
@@ -0,0 +1,11 @@
+	.text
+	.file	"sub.zasm"
+	.globl	__entry
+__entry:
+.func_begin0:
+	sub	3, r0, r1
+	sstore	r0, r1
+	ret
+.func_end0:
+	.note.GNU-stack
+	.rodata
diff --git a/src/address_operands.rs b/src/address_operands.rs
@@ -0,0 +1,131 @@
+use u256::U256;
+use zkevm_opcode_defs::{ImmMemHandlerFlags, Operand, RegOrImmFlags};
+
+use crate::{state::VMState, value::TaggedValue, Opcode};
+
+fn only_reg_read(vm: &mut VMState, opcode: &Opcode) -> (U256, U256) {
+    let src0 = vm.get_register(opcode.src0_index);
+    let src1 = vm.get_register(opcode.src1_index);
+    (src0, src1)
+}
+
+fn only_imm16_read(vm: &mut VMState, opcode: &Opcode) -> (U256, U256) {
+    let src1 = vm.get_register(opcode.src1_index);
+    (U256::from(opcode.imm0), src1)
+}
+
+fn reg_and_imm_read(vm: &mut VMState, opcode: &Opcode) -> (U256, U256) {
+    let src0 = vm.get_register(opcode.src0_index);
+    let src1 = vm.get_register(opcode.src1_index);
+    let offset = opcode.imm0;
+
+    (src0 + U256::from(offset), src1)
+}
+
+pub fn address_operands_read(vm: &mut VMState, opcode: &Opcode) -> (U256, U256) {
+    match opcode.src0_operand_type {
+        Operand::RegOnly => only_reg_read(vm, opcode),
+        Operand::RegOrImm(variant) => match variant {
+            RegOrImmFlags::UseRegOnly => only_reg_read(vm, opcode),
+            RegOrImmFlags::UseImm16Only => only_imm16_read(vm, opcode),
+        },
+        Operand::Full(variant) => {
+            match variant {
+                ImmMemHandlerFlags::UseRegOnly => only_reg_read(vm, opcode),
+                ImmMemHandlerFlags::UseStackWithPushPop => {
+                    // stack-=[src0 + offset] + src1
+                    let (src0, src1) = reg_and_imm_read(vm, opcode);
+                    let res = vm.current_frame.stack[vm.sp() - src0.as_usize()].value;
+                    // TODO: Add push/pop stack
+                    (res, src1)
+                }
+                ImmMemHandlerFlags::UseStackWithOffset => {
+                    // stack[src0 + offset] + src1
+                    let (src0, src1) = reg_and_imm_read(vm, opcode);
+
+                    let res = vm.current_frame.stack[vm.sp() - src0.as_usize()].value;
+                    (res, src1)
+                }
+                ImmMemHandlerFlags::UseAbsoluteOnStack => {
+                    // stack=[src0 + offset] + src1
+                    let (src0, src1) = reg_and_imm_read(vm, opcode);
+                    let res = vm.current_frame.stack[src0.as_usize()].value;
+                    (res, src1)
+                }
+                ImmMemHandlerFlags::UseImm16Only => only_imm16_read(vm, opcode),
+                ImmMemHandlerFlags::UseCodePage => {
+                    let (src0, src1) = reg_and_imm_read(vm, opcode);
+
+                    let res = vm.current_frame.code_page[src0.as_usize()];
+                    (res, src1)
+                }
+            }
+        }
+    }
+}
+fn only_reg_write(vm: &mut VMState, opcode: &Opcode, res: U256) {
+    vm.set_register(opcode.dst0_index, res);
+}
+
+fn reg_and_imm_write(vm: &mut VMState, opcode: &Opcode) -> U256 {
+    let dst0 = vm.get_register(opcode.dst0_index);
+    let offset = opcode.imm0;
+
+    dst0 + U256::from(offset)
+}
+
+pub fn address_operands_store(vm: &mut VMState, opcode: &Opcode, res: U256) {
+    match opcode.dst0_operand_type {
+        Operand::RegOnly => {
+            only_reg_write(vm, opcode, res);
+        }
+        Operand::RegOrImm(variant) => match variant {
+            RegOrImmFlags::UseRegOnly => {
+                only_reg_write(vm, opcode, res);
+            }
+            RegOrImmFlags::UseImm16Only => {
+                panic!("dest cannot be imm16 only");
+            }
+        },
+        Operand::Full(variant) => {
+            match variant {
+                ImmMemHandlerFlags::UseRegOnly => {
+                    only_reg_write(vm, opcode, res);
+                }
+                ImmMemHandlerFlags::UseStackWithPushPop => {
+                    // stack-=[src0 + offset] + src1
+                    let src0 = reg_and_imm_write(vm, opcode);
+                    let sp = vm.sp();
+                    vm.current_frame.stack[sp - src0.as_usize()] = TaggedValue {
+                        value: res,
+                        is_pointer: false,
+                    };
+                    // TODO: Add push/pop stack
+                }
+                ImmMemHandlerFlags::UseStackWithOffset => {
+                    // stack[src0 + offset] + src1
+                    let src0 = reg_and_imm_write(vm, opcode);
+                    let sp = vm.sp();
+                    vm.current_frame.stack[sp - src0.as_usize()] = TaggedValue {
+                        value: res,
+                        is_pointer: false,
+                    };
+                }
+                ImmMemHandlerFlags::UseAbsoluteOnStack => {
+                    // stack=[src0 + offset] + src1
+                    let src0 = reg_and_imm_write(vm, opcode);
+                    vm.current_frame.stack[src0.as_usize()] = TaggedValue {
+                        value: res,
+                        is_pointer: false,
+                    };
+                }
+                ImmMemHandlerFlags::UseImm16Only => {
+                    panic!("dest cannot be imm16 only");
+                }
+                ImmMemHandlerFlags::UseCodePage => {
+                    panic!("dest cannot be code page");
+                }
+            }
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,9 +1,11 @@
+mod address_operands;
 mod op_handlers;
 mod opcode;
 pub mod state;
 mod value;
 
 use op_handlers::add::_add;
+use op_handlers::sub::_sub;
 pub use opcode::Opcode;
 use state::VMState;
 use u256::U256;
@@ -20,26 +22,25 @@ pub fn run_program(bin_path: &str) -> U256 {
     let bin = hex::decode(&encoded[2..]).unwrap();
 
     let mut program_code = vec![];
-    for raw_opcode_slice in bin.chunks(8) {
-        let mut raw_opcode_bytes: [u8; 8] = [0; 8];
-        raw_opcode_bytes.copy_from_slice(&raw_opcode_slice[..8]);
+    for raw_opcode_slice in bin.chunks(32) {
+        let mut raw_opcode_bytes: [u8; 32] = [0; 32];
+        raw_opcode_bytes.copy_from_slice(&raw_opcode_slice[..32]);
 
-        let raw_opcode_u64 = u64::from_be_bytes(raw_opcode_bytes);
-        let opcode = Opcode::from_raw_opcode(raw_opcode_u64, &opcode_table);
-        program_code.push(opcode);
+        let raw_opcode_u256 = U256::from_big_endian(&raw_opcode_bytes);
+        program_code.push(raw_opcode_u256);
     }
 
     let mut vm = VMState::new(program_code);
 
     loop {
-        let opcode = vm.current_frame.code_page[vm.current_frame.pc as usize].clone();
+        let opcode = vm.get_opcode(&opcode_table);
         match opcode.variant {
             Variant::Invalid(_) => todo!(),
             Variant::Nop(_) => todo!(),
             Variant::Add(_) => {
                 _add(&mut vm, opcode);
             }
-            Variant::Sub(_) => todo!(),
+            Variant::Sub(_) => _sub(&mut vm, opcode),
             Variant::Mul(_) => todo!(),
             Variant::Div(_) => todo!(),
             Variant::Jump(_) => todo!(),

diff --git a/src/op_handlers/add.rs b/src/op_handlers/add.rs
@@ -1,29 +1,21 @@
+use crate::address_operands::{address_operands_read, address_operands_store};
 use crate::{opcode::Opcode, state::VMState};
-use zkevm_opcode_defs::ImmMemHandlerFlags;
-use zkevm_opcode_defs::Operand;
+
+fn _add_reg_only(vm: &mut VMState, opcode: Opcode) {
+    // src0 + src1 -> dst0
+    let src0 = vm.get_register(opcode.src0_index);
+    let src1 = vm.get_register(opcode.src1_index);
+    vm.set_register(opcode.dst0_index, src0 + src1);
+}
+
+fn _add_imm16_only(vm: &mut VMState, opcode: Opcode) {
+    // imm0 + src0 -> dst0
+    let src1 = vm.get_register(opcode.src1_index);
+    vm.set_register(opcode.dst0_index, src1 + opcode.imm0);
+}
 
 pub fn _add(vm: &mut VMState, opcode: Opcode) {
-    match opcode.src0_operand_type {
-        Operand::RegOnly => todo!(),
-        Operand::RegOrImm(_) => todo!(),
-        Operand::Full(variant) => {
-            match variant {
-                ImmMemHandlerFlags::UseRegOnly => {
-                    // src0 + src1 -> dst0
-                    let src0 = vm.get_register(opcode.src0_index);
-                    let src1 = vm.get_register(opcode.src1_index);
-                    vm.set_register(opcode.dst0_index, src0 + src1);
-                }
-                ImmMemHandlerFlags::UseStackWithPushPop => todo!(),
-                ImmMemHandlerFlags::UseStackWithOffset => todo!(),
-                ImmMemHandlerFlags::UseAbsoluteOnStack => todo!(),
-                ImmMemHandlerFlags::UseImm16Only => {
-                    // imm0 + src0 -> dst0
-                    let src0 = vm.get_register(opcode.src0_index);
-                    vm.set_register(opcode.dst0_index, src0 + opcode.imm0);
-                }
-                ImmMemHandlerFlags::UseCodePage => todo!(),
-            }
-        }
-    }
+    let (src0, src1) = address_operands_read(vm, &opcode);
+    let res = src0 + src1;
+    address_operands_store(vm, &opcode, res);
 }
diff --git a/src/op_handlers/mod.rs b/src/op_handlers/mod.rs
@@ -1 +1,2 @@
 pub mod add;
+pub mod sub;
diff --git a/src/op_handlers/sub.rs b/src/op_handlers/sub.rs
@@ -0,0 +1,8 @@
+use crate::address_operands::{address_operands_read, address_operands_store};
+use crate::{opcode::Opcode, state::VMState};
+
+pub fn _sub(vm: &mut VMState, opcode: Opcode) {
+    let (src0, src1) = address_operands_read(vm, &opcode);
+    let res = src0 - src1;
+    address_operands_store(vm, &opcode, res);
+}
diff --git a/src/state.rs b/src/state.rs
@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use crate::{value::TaggedValue, Opcode};
 use u256::U256;
+use zkevm_opcode_defs::OpcodeVariant;
 
 #[derive(Debug)]
 pub struct CallFrame {
@@ -12,7 +13,7 @@ pub struct CallFrame {
     // TODO: this is a Vec of opcodes now but it's probably going to switch back to a
     // Vec<U256> later on, because I believe we have to record memory queries when
     // fetching code to execute. Check this
-    pub code_page: Vec<Opcode>,
+    pub code_page: Vec<U256>,
     pub pc: u64,
     // TODO: Storage is more complicated than this. We probably want to abstract it into a trait
     // to support in-memory vs on-disk storage, etc.
@@ -30,7 +31,7 @@ pub struct VMState {
 
 impl VMState {
     // TODO: The VM will probably not take the program to execute as a parameter later on.
-    pub fn new(program_code: Vec<Opcode>) -> Self {
+    pub fn new(program_code: Vec<U256>) -> Self {
         Self {
             registers: [U256::zero(); 15],
             flags: 0,
@@ -53,10 +54,27 @@ impl VMState {
 
         self.registers[(index - 1) as usize] = value;
     }
+
+    pub fn get_opcode(&self, opcode_table: &[OpcodeVariant]) -> Opcode {
+        let raw_opcode = self.current_frame.code_page[(self.current_frame.pc / 4) as usize];
+        let raw_opcode_64 = match self.current_frame.pc % 4 {
+            3 => (raw_opcode & u64::MAX.into()).as_u64(),
+            2 => ((raw_opcode >> 64) & u64::MAX.into()).as_u64(),
+            1 => ((raw_opcode >> 128) & u64::MAX.into()).as_u64(),
+            0 => ((raw_opcode >> 192) & u64::MAX.into()).as_u64(),
+            _ => panic!("This should never happen"),
+        };
+
+        Opcode::from_raw_opcode(raw_opcode_64, opcode_table)
+    }
+
+    pub fn sp(&self) -> usize {
+        self.current_frame.stack.len()
+    }
 }
 
 impl CallFrame {
-    pub fn new(program_code: Vec<Opcode>) -> Self {
+    pub fn new(program_code: Vec<U256>) -> Self {
         Self {
             stack: vec![],
             heap: vec![],