implement Blake2sLastBlock opcode in runner (#1932)

ohad-nir-starkware · web-flow · commit abe5f44f55e1 · 2025-02-17T18:51:47.000Z
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* feat: implement `Blake2sLastBlock` opcode in VM [#1932](https://github.com/lambdaclass/cairo-vm/pull/1932)
+
 * feat: implement `Blake2s` opcode in VM [#1927](https://github.com/lambdaclass/cairo-vm/pull/1927)
 
 * feat: remove `NonZeroReservedBits` from `VirtualMachineError` [#1948](https://github.com/lambdaclass/cairo-vm/pull/1948)
diff --git a/cairo_programs/stwo_exclusive_programs/blake2s_opcode_test.cairo b/cairo_programs/stwo_exclusive_programs/blake2s_opcode_test.cairo
@@ -1,18 +1,24 @@
 %builtins range_check bitwise
 
 from starkware.cairo.common.alloc import alloc
+from starkware.cairo.common.bool import FALSE, TRUE
 from starkware.cairo.common.cairo_blake2s.blake2s import STATE_SIZE_FELTS, INPUT_BLOCK_FELTS, _get_sigma
 from starkware.cairo.common.cairo_blake2s.packed_blake2s import N_PACKED_INSTANCES, blake2s_compress
 from starkware.cairo.common.cairo_builtins import BitwiseBuiltin
 
 const COUNTER = 64;
 const U32_MASK = 0xffffffff;
 
-// Tests the Blake2s opcode runner using a preexisting implementation within the repo as reference.
-// The initial state, a random message of 64 bytes and counter are used as input.
+// Tests the Blake2s and Blake2sLastBlock opcode runners using a preexisting implementation within the repo as reference.
+// The initial state, a random message of 64 bytes and a counter are used as input.
 // Both the opcode and the reference implementation are run on the same inputs and then their outputs are compared.
 // Before comparing the outputs, it is verified that the opcode runner has written the output to the correct location.
 func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
+    run_blake_test(is_last_block=FALSE);
+    run_blake_test(is_last_block=TRUE);
+    return ();
+}
+func run_blake_test{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}(is_last_block: felt) {
     alloc_locals;
 
     let (local random_message) = alloc();
@@ -52,7 +58,7 @@ func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
         h=input_state,
         message=random_message,
         t0=COUNTER,
-        f0=0,
+        f0=is_last_block * U32_MASK,
         sigma=sigma,
         output=cairo_output,
     );
@@ -76,7 +82,8 @@ func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
     assert bitwise_ptr[7].y = U32_MASK;
 
     // Run the blake2s opcode runner on the same inputs and store its output.
-    let vm_output = run_blake2s(
+    let vm_output = run_blake2s_opcode(
+        is_last_block = is_last_block,
         dst=COUNTER,
         op0=input_state,
         op1=random_message,
@@ -107,7 +114,7 @@ func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
     return ();
 }
 
-// Forces the runner to execute the Blake2s with the given operands.
+// Forces the runner to execute the Blake2s or Blake2sLastBlock opcode with the given operands.
 // op0 is a pointer to an array of 8 felts as u32 integers of the state.
 // op1 is a pointer to an array of 16 felts as u32 integers of the messsage.
 // dst is a felt representing a u32 of the counter.
@@ -116,7 +123,8 @@ func main{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
 // An instruction encoding is built from offsets -5, -4, -3 and flags which are all 0 except for
 // those denoting uses of fp as the base for operand addresses and flag_opcode_blake (16th flag).
 // The instruction is then written to [pc] and the runner is forced to execute Blake2s.
-func run_blake2s(
+func run_blake2s_opcode(
+    is_last_block: felt,
     dst: felt,
     op0: felt*,
     op1: felt*,
@@ -127,9 +135,9 @@ func run_blake2s(
     let offset0 = (2**15)-5;
     let offset1 = (2**15)-4;
     let offset2 = (2**15)-3;
-    static_assert dst == [fp -5];
-    static_assert op0 == [fp -4];
-    static_assert op1 == [fp -3];
+    static_assert dst == [fp - 5];
+    static_assert op0 == [fp - 4];
+    static_assert op1 == [fp - 3];
 
     // Set the flags for the instruction.
     let flag_dst_base_fp = 1;
@@ -147,17 +155,24 @@ func run_blake2s(
     let flag_opcode_call = 0;
     let flag_opcode_ret = 0;
     let flag_opcode_assert_eq = 0;
-    let flag_opcode_blake2s = 1;
 
-    // Build the instruction encoding.
-    let flag_num = flag_dst_base_fp+flag_op0_base_fp*(2**1)+flag_op1_imm*(2**2)+flag_op1_base_fp*(2**3)+flag_opcode_blake2s*(2**15);
-    let instruction_num = offset0 + offset1*(2**16) + offset2*(2**32) + flag_num*(2**48);
-    static_assert instruction_num==9226608988349300731;
+    let flag_num = flag_dst_base_fp+flag_op0_base_fp*(2**1)+flag_op1_imm*(2**2)+flag_op1_base_fp*(2**3);
+    let blake2s_opcode_extension_num = 1;
+    let blake2s_last_block_opcode_extension_num = 2;
+    let blake2s_instruction_num = offset0 + offset1*(2**16) + offset2*(2**32) + flag_num*(2**48) + blake2s_opcode_extension_num*(2**63);
+    let blake2s_last_block_instruction_num = offset0 + offset1*(2**16) + offset2*(2**32) + flag_num*(2**48) + blake2s_last_block_opcode_extension_num*(2**63);
+    static_assert blake2s_instruction_num==9226608988349300731;
+    static_assert blake2s_last_block_instruction_num==18449981025204076539;
 
     // Write the instruction to [pc] and point [ap] to the designated output.
     let (local vm_output) = alloc();
     assert [ap] = cast(vm_output, felt);
+
+    jmp last_block if is_last_block!=0;
     dw 9226608988349300731;
+    return cast([ap], felt*);
 
+    last_block:
+    dw 18449981025204076539;
     return cast([ap], felt*);
 }
diff --git a/vm/src/types/instruction.rs b/vm/src/types/instruction.rs
@@ -81,6 +81,7 @@ pub enum Opcode {
 pub enum OpcodeExtension {
     Stone,
     Blake,
+    BlakeFinalize,
 }
 
 impl Instruction {
diff --git a/vm/src/vm/decoding/decoder.rs b/vm/src/vm/decoding/decoder.rs
@@ -102,24 +102,28 @@ pub fn decode_instruction(encoded_instr: u128) -> Result<Instruction, VirtualMac
 
     let opcode_extension = match opcode_extension_num {
         0 => OpcodeExtension::Stone,
-        1 => {
-            if opcode != Opcode::NOp
-                || (op1_addr != Op1Addr::FP && op1_addr != Op1Addr::AP)
-                || res != Res::Op1
-                || pc_update != PcUpdate::Regular
-                || (ap_update_num != 0 && ap_update_num != 2)
-            {
-                return Err(VirtualMachineError::InvalidBlake2sFlags(flags & 0x7FFF));
-            };
-            OpcodeExtension::Blake
-        }
+        1 => OpcodeExtension::Blake,
+        2 => OpcodeExtension::BlakeFinalize,
         _ => {
             return Err(VirtualMachineError::InvalidOpcodeExtension(
                 opcode_extension_num,
             ))
         }
     };
 
+    let are_blake_flags_invalid = opcode != Opcode::NOp
+        || (op1_addr != Op1Addr::FP && op1_addr != Op1Addr::AP)
+        || res != Res::Op1
+        || pc_update != PcUpdate::Regular
+        || (ap_update_num != 0 && ap_update_num != 2);
+
+    if (opcode_extension == OpcodeExtension::Blake
+        || opcode_extension == OpcodeExtension::BlakeFinalize)
+        && are_blake_flags_invalid
+    {
+        return Err(VirtualMachineError::InvalidBlake2sFlags(flags & 0x7FFF));
+    }
+
     let ap_update = match (ap_update_num, opcode == Opcode::Call) {
         (0, true) => ApUpdate::Add2,
         (0, false) => ApUpdate::Regular,
diff --git a/vm/src/vm/vm_core.rs b/vm/src/vm/vm_core.rs
@@ -445,8 +445,13 @@ impl VirtualMachine {
             .memory
             .mark_as_accessed(operands_addresses.op1_addr);
 
-        if instruction.opcode_extension == OpcodeExtension::Blake {
-            self.handle_blake2s_instruction(&operands_addresses)?;
+        if instruction.opcode_extension == OpcodeExtension::Blake
+            || instruction.opcode_extension == OpcodeExtension::BlakeFinalize
+        {
+            self.handle_blake2s_instruction(
+                &operands_addresses,
+                instruction.opcode_extension == OpcodeExtension::BlakeFinalize,
+            )?;
         }
 
         self.update_registers(instruction, operands)?;
@@ -455,7 +460,7 @@ impl VirtualMachine {
         Ok(())
     }
 
-    /// Executes a Blake2s instruction.
+    /// Executes a Blake2s or Blake2sLastBlock instruction.
     /// Expects operands to be RelocatableValue and to point to segments of memory.
     /// op0 is expected to point to a sequence of 8 u32 values (state).
     /// op1 is expected to point to a sequence of 16 u32 values (message).
@@ -469,6 +474,7 @@ impl VirtualMachine {
     fn handle_blake2s_instruction(
         &mut self,
         operands_addresses: &OperandsAddresses,
+        is_last_block: bool,
     ) -> Result<(), VirtualMachineError> {
         let counter = self.segments.memory.get_u32(operands_addresses.dst_addr)?;
 
@@ -490,10 +496,12 @@ impl VirtualMachine {
         .try_into()
         .map_err(|_| VirtualMachineError::Blake2sInvalidOperand(1, 16))?;
 
+        let f0 = if is_last_block { 0xffffffff } else { 0 };
+
         let ap = self.run_context.get_ap();
         let output_address = self.segments.memory.get_relocatable(ap)?;
 
-        let new_state = blake2s_compress(&state, &message, counter, 0, 0, 0);
+        let new_state = blake2s_compress(&state, &message, counter, 0, f0, 0);
 
         for (i, &val) in new_state.iter().enumerate() {
             self.segments.memory.insert_as_accessed(
@@ -4496,7 +4504,7 @@ mod tests {
         };
 
         assert_matches!(
-            vm.handle_blake2s_instruction(&operands_addresses),
+            vm.handle_blake2s_instruction(&operands_addresses, false),
             Err(VirtualMachineError::Memory(MemoryError::UnknownMemoryCell(bx))) if *bx == (0, 7).into()
         );
     }
@@ -4528,7 +4536,7 @@ mod tests {
         };
 
         assert_matches!(
-            vm.handle_blake2s_instruction(&operands_addresses),
+            vm.handle_blake2s_instruction(&operands_addresses, false),
             Err(VirtualMachineError::Memory(MemoryError::UnknownMemoryCell(bx))) if *bx == (0, 8).into()
         );
     }
@@ -4568,7 +4576,7 @@ mod tests {
         };
 
         assert_matches!(
-            vm.handle_blake2s_instruction(&operands_addresses),
+            vm.handle_blake2s_instruction(&operands_addresses, false),
             Err(VirtualMachineError::Memory(MemoryError::InconsistentMemory(bx))) if *bx == ((0, 0).into(),0.into(),1848029226.into())
         );
     }
@@ -4621,7 +4629,10 @@ mod tests {
             ap: 0,
             fp: 0,
         };
-        assert_matches!(vm.handle_blake2s_instruction(&operands_addresses), Ok(()));
+        assert_matches!(
+            vm.handle_blake2s_instruction(&operands_addresses, false),
+            Ok(())
+        );
 
         let state: [u32; 8] = vm
             .get_u32_range((0, 0).into(), 8)
@@ -4647,6 +4658,84 @@ mod tests {
         assert_eq!(new_state, expected_new_state);
     }
 
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn handle_blake2s_last_block_instruction_ok() {
+        let mut vm = vm!();
+        vm.segments.memory = memory![
+            // State
+            ((0, 0), 0x6B08E647),
+            ((0, 1), 0xBB67AE85),
+            ((0, 2), 0x3C6EF372),
+            ((0, 3), 0xA54FF53A),
+            ((0, 4), 0x510E527F),
+            ((0, 5), 0x9B05688C),
+            ((0, 6), 0x1F83D9AB),
+            ((0, 7), 0x5BE0CD19),
+            // Message
+            ((0, 8), 930933030),
+            ((0, 9), 1766240503),
+            ((0, 10), 3660871006),
+            ((0, 11), 388409270),
+            ((0, 12), 1948594622),
+            ((0, 13), 3119396969),
+            ((0, 14), 3924579183),
+            ((0, 15), 2089920034),
+            ((0, 16), 3857888532),
+            ((0, 17), 929304360),
+            ((0, 18), 1810891574),
+            ((0, 19), 860971754),
+            ((0, 20), 1822893775),
+            ((0, 21), 2008495810),
+            ((0, 22), 2958962335),
+            ((0, 23), 2340515744),
+            // Counter
+            ((0, 24), 64),
+            // AP
+            ((1, 0), (0, 25)),
+            ((2, 0), (0, 0)),
+            ((2, 1), (0, 8))
+        ];
+        let operands_addresses = OperandsAddresses {
+            dst_addr: (0, 24).into(),
+            op0_addr: (2, 0).into(),
+            op1_addr: (2, 1).into(),
+        };
+        vm.run_context = RunContext {
+            pc: (0, 0).into(),
+            ap: 0,
+            fp: 0,
+        };
+        assert_matches!(
+            vm.handle_blake2s_instruction(&operands_addresses, true),
+            Ok(())
+        );
+
+        let state: [u32; 8] = vm
+            .get_u32_range((0, 0).into(), 8)
+            .unwrap()
+            .try_into()
+            .unwrap();
+        let message: [u32; 16] = vm
+            .get_u32_range((0, 8).into(), 16)
+            .unwrap()
+            .try_into()
+            .unwrap();
+        let counter = vm.segments.memory.get_u32((0, 24).into()).unwrap();
+
+        let expected_new_state: [u32; 8] =
+            blake2s_compress(&state, &message, counter, 0, 0xffffffff, 0)
+                .try_into()
+                .unwrap();
+
+        let new_state: [u32; 8] = vm
+            .get_u32_range((0, 25).into(), 8)
+            .unwrap()
+            .try_into()
+            .unwrap();
+        assert_eq!(new_state, expected_new_state);
+    }
+
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn get_traceback_entries_bad_usort() {

Original file line number	Diff line number	Diff line change
`@@ -81,6 +81,7 @@ pub enum Opcode {`
`81`	`81`	`pub enum OpcodeExtension {`
`82`	`82`	`Stone,`
`83`	`83`	`Blake,`
	`84`	`+ BlakeFinalize,`
`84`	`85`	`}`
`85`	`86`
`86`	`87`	`impl Instruction {`