From 213ba8ff50835d99039b418dcaf31cd6f7698137 Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Tue, 7 Mar 2023 17:17:23 -0800
Subject: [PATCH] optimized initial hash value generation

---
 .../cpu/kernel/asm/hash/blake2b/addresses.asm | 16 -----
 .../kernel/asm/hash/blake2b/compression.asm   |  4 +-
 evm/src/cpu/kernel/asm/hash/blake2b/iv.asm    | 61 ++++++++++++++-----
 evm/src/cpu/kernel/tests/hash.rs              |  2 -
 4 files changed, 50 insertions(+), 33 deletions(-)

diff --git a/evm/src/cpu/kernel/asm/hash/blake2b/addresses.asm b/evm/src/cpu/kernel/asm/hash/blake2b/addresses.asm
index 9d65b9ed..f1d7c3e9 100644
--- a/evm/src/cpu/kernel/asm/hash/blake2b/addresses.asm
+++ b/evm/src/cpu/kernel/asm/hash/blake2b/addresses.asm
@@ -1,19 +1,3 @@
-// Load the initial hash value (the IV, but with params XOR'd into the first word).
-%macro blake2b_initial_hash_value
-    %blake2b_iv_i(7)
-    %blake2b_iv_i(6)
-    %blake2b_iv_i(5)
-    %blake2b_iv_i(4)
-    %blake2b_iv_i(3)
-    %blake2b_iv_i(2)
-    %blake2b_iv_i(1)
-    // stack: IV_1, IV_2, IV_3, IV_4, IV_5, IV_6, IV_7
-    PUSH 0x01010040 // params: key = 00, digest_size = 64 = 0x40
-    %blake2b_iv_i(0)
-    XOR
-    // stack: IV_0 ^ params, IV_1, IV_2, IV_3, IV_4, IV_5, IV_6, IV_7
-%endmacro
-
 // Address where the working version of the hash value is stored.
 %macro blake2b_hash_value_addr
     PUSH 0
diff --git a/evm/src/cpu/kernel/asm/hash/blake2b/compression.asm b/evm/src/cpu/kernel/asm/hash/blake2b/compression.asm
index 11e5389b..cd1f6a80 100644
--- a/evm/src/cpu/kernel/asm/hash/blake2b/compression.asm
+++ b/evm/src/cpu/kernel/asm/hash/blake2b/compression.asm
@@ -2,7 +2,9 @@ global blake2b_compression:
     // stack: retdest
     PUSH 0
     // stack: cur_block = 0, retdest
-    %blake2b_initial_hash_value
+    PUSH compression_loop
+    // stack: compression_loop, cur_block, retdest
+    %jump(blake2b_initial_hash_value)
 compression_loop:
     // stack: h_0, ..., h_7, cur_block, retdest
     
diff --git a/evm/src/cpu/kernel/asm/hash/blake2b/iv.asm b/evm/src/cpu/kernel/asm/hash/blake2b/iv.asm
index 94e9ba27..48df86a3 100644
--- a/evm/src/cpu/kernel/asm/hash/blake2b/iv.asm
+++ b/evm/src/cpu/kernel/asm/hash/blake2b/iv.asm
@@ -33,30 +33,63 @@ global blake2b_iv_const:
     BYTES 91, 224, 205, 25
     BYTES 19, 126, 33, 121
 
-%macro blake2b_iv
-    // stack: i, ...
+global blake2b_iv:
+    // stack: i, retdest
     PUSH blake2b_iv_const
-    // stack: blake2b_iv_const, i, ...
+    // stack: blake2b_iv_const, i, retdest
     SWAP1
-    // stack: i, blake2b_iv_const, ...
+    // stack: i, blake2b_iv_const, retdest
     %mul_const(8)
     ADD
-    // stack: blake2b_iv_const + 2 * i, ...
+    // stack: blake2b_iv_const + 2 * i, retdest
     DUP1
-    // stack: blake2b_iv_const + 2 * i, blake2b_iv_const + 2 * i, ...
+    // stack: blake2b_iv_const + 2 * i, blake2b_iv_const + 2 * i, retdest
     %add_const(4)
-    // stack: blake2b_iv_const + 2 * i + 1, blake2b_iv_const + 2 * i, ...
+    // stack: blake2b_iv_const + 2 * i + 1, blake2b_iv_const + 2 * i, retdest
     %mload_kernel_code_u32
     SWAP1
     %mload_kernel_code_u32
-    // stack: IV_i[32:], IV_i[:32], ...
+    // stack: IV_i[32:], IV_i[:32], retdest
     %shl_const(32)
-    // stack: IV_i[32:] << 32, IV_i[:32], ...
+    // stack: IV_i[32:] << 32, IV_i[:32], retdest
     ADD // OR
-    // stack: IV_i, ...
+    // stack: IV_i, retdest
+    SWAP1
+    JUMP
+
+%macro blake2b_iv
+    %stack (i) -> (i, %%after)
+    %jump(blake2b_iv)
+%%after:
 %endmacro
 
-%macro blake2b_iv_i(i)
-    PUSH $i
-    %blake2b_iv
-%endmacro
+// Load the initial hash value (the IV, but with params XOR'd into the first word).
+global blake2b_initial_hash_value:
+    // stack: retdest
+    PUSH 8
+    // stack: i=8, retdest
+blake2b_initial_hash_loop:
+    // stack: i, IV_i, ..., IV_7, retdest
+    %decrement
+    // stack: i-1, IV_i, ..., IV_7, retdest
+    PUSH blake2b_initial_hash_return
+    // stack: blake2b_initial_hash_return, i-1, IV_i, ..., IV_7, retdest
+    DUP2
+    // stack: i-1, blake2b_initial_hash_return, i-1, IV_i, ..., IV_7, retdest
+    %jump(blake2b_iv)
+blake2b_initial_hash_return:
+    // stack: IV_(i-1), i-1, IV_i, ..., IV_7, retdest
+    SWAP1
+    // stack: i-1, IV_(i-1), IV_i, ..., IV_7, retdest
+    DUP1
+    // stack: i-1, i-1, IV_(i-1), ..., IV_7, retdest
+    %jumpi(blake2b_initial_hash_loop)
+    // stack: i-1=0, IV_0, ..., IV_7, retdest
+    POP
+    // stack: IV_0, ..., IV_7, retdest
+    PUSH 0x01010040 // params: key = 00, digest_size = 64 = 0x40
+    XOR
+    // stack: IV_0 ^ params, IV_1, IV_2, IV_3, IV_4, IV_5, IV_6, IV_7, retdest
+    %stack(iv: 8, ret) -> (ret, iv)
+    JUMP
+
diff --git a/evm/src/cpu/kernel/tests/hash.rs b/evm/src/cpu/kernel/tests/hash.rs
index cc2c96eb..bc73ecd5 100644
--- a/evm/src/cpu/kernel/tests/hash.rs
+++ b/evm/src/cpu/kernel/tests/hash.rs
@@ -79,8 +79,6 @@ fn prepare_test<T>(
     // Run the interpeter
     let result = run_interpreter_with_memory(interpreter_setup).unwrap();
 
-    dbg!(result.stack().to_vec());
-
     Ok((expected, result.stack().to_vec()))
 }