starknet_os: fix blake2s finalize opcode command in naive blake

Author: meship-starkware

crates/apollo_starknet_os_program/src/cairo/starkware/starknet/core/os/naive_blake.cairo

@@ -1,18 +1,43 @@
 from starkware.cairo.common.alloc import alloc
 from starkware.cairo.common.cairo_blake2s.blake2s import blake_with_opcode
-from starkware.cairo.common.cairo_blake2s.blake2s import BLAKE2S_FINALIZE_INSTRUCTION
 
 // Computes blake2s of `input` of size 16 felts, representing 32 bits each.
 // The initial state is the standard BLAKE2s IV XORed with the parameter block P[0] = 0x01010020.
 func blake_with_opcode_for_single_16_length_word(data: felt*, out: felt*, initial_state: felt*) {
+    const BLAKE2S_FINALIZE_OPCODE_EXT = 2;
+    const OP0_REG = 1;
+    const OP1_FP = 3;
+    const AP_ADD1 = 11;
+    // Flags: 000100000001010
+    // - State and data are referenced from fp
+    // - Counter is referenced from ap
+    // - Increment ap by 1 after instruction
+    const BLAKE2S_FLAGS = 2 ** OP0_REG + 2 ** OP1_FP + 2 ** AP_ADD1;
+
+    const STATE_OFFSET = -3;
+    const MESSAGE_OFFSET = -5;
+    const COUNTER_OFFSET = -1;
+
+    const COUNTER_SHIFT = 1;
+    const STATE_SHIFT = 2 ** 16;
+    const MESSAGE_SHIFT = 2 ** 32;
+    const FLAGS_SHIFT = 2 ** 48;
+    const OPCODE_EXT_SHIFT = 2 ** 63;
+
+    const POS_STATE_OFFSET = 2 ** 15 + STATE_OFFSET;
+    const POS_MESSAGE_OFFSET = 2 ** 15 + MESSAGE_OFFSET;
+    const POS_COUNTER_OFFSET = 2 ** 15 + COUNTER_OFFSET;
+
+    const BLAKE2S_FINALIZE_INSTRUCTION = POS_COUNTER_OFFSET * COUNTER_SHIFT + POS_STATE_OFFSET *
+        STATE_SHIFT + POS_MESSAGE_OFFSET * MESSAGE_SHIFT + BLAKE2S_FLAGS * FLAGS_SHIFT +
+        BLAKE2S_FINALIZE_OPCODE_EXT * OPCODE_EXT_SHIFT;
+
     tempvar counter = 64;
-    [ap] = initial_state, ap++;
-    [ap] = data, ap++;
-    [ap] = counter, ap++;
     [ap] = out;
+    static_assert [ap + COUNTER_OFFSET] == counter;
+    static_assert [fp + STATE_OFFSET] == initial_state;
+    static_assert [fp + MESSAGE_OFFSET] == data;
     dw BLAKE2S_FINALIZE_INSTRUCTION;
-    // Increment AP after blake opcode.
-    ap += 1;
     return ();
 }