From c11c52646281f0af0995199ddac723f1e49f15a9 Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Wed, 14 Sep 2022 11:35:28 -0700
Subject: [PATCH] cleanup and comments

---
 evm/src/cpu/kernel/asm/sha2/ops.asm  | 10 ++--------
 evm/src/cpu/kernel/asm/sha2/util.asm |  7 +++++++
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/evm/src/cpu/kernel/asm/sha2/ops.asm b/evm/src/cpu/kernel/asm/sha2/ops.asm
index 0953a9b5..4c57d491 100644
--- a/evm/src/cpu/kernel/asm/sha2/ops.asm
+++ b/evm/src/cpu/kernel/asm/sha2/ops.asm
@@ -3,7 +3,7 @@
     // stack: x, y
     ADD
     // stack: x + y
-    %and_const(0xFFFFFFFF)
+    %truncate_to_u32
     // stack: (x + y) & u32::MAX
 %endmacro
 
@@ -25,13 +25,7 @@
     // stack: 32 - rot, value, value >> rot
     SHL
     // stack: value << (32 - rot), value >> rot
-    PUSH 32
-    PUSH 1
-    SWAP1
-    SHL
-    // stack: 1 << 32, value << (32 - rot), value >> rot
-    SWAP1
-    MOD
+    %truncate_to_u32
     // stack: (value << (32 - rot)) % (1 << 32), value >> rot
     ADD
 %endmacro
diff --git a/evm/src/cpu/kernel/asm/sha2/util.asm b/evm/src/cpu/kernel/asm/sha2/util.asm
index f3f5b54c..28e71c12 100644
--- a/evm/src/cpu/kernel/asm/sha2/util.asm
+++ b/evm/src/cpu/kernel/asm/sha2/util.asm
@@ -1,11 +1,18 @@
+// We put the message schedule in memory starting at 64 * num_blocks + 2.
 %macro message_schedule_addr_from_num_blocks
     // stack: num_blocks
     %mul_const(64)
     %add_const(2)
 %endmacro
 
+// We use memory starting at 320 * num_blocks + 2 (after the message schedule
+// space) as scratch space to store stack values.
 %macro scratch_space_addr_from_num_blocks
     // stack: num_blocks
     %mul_const(320)
     %add_const(2)
+%endmacro
+
+%macro truncate_to_u32
+    %and_const(0xFFFFFFFF)
 %endmacro
\ No newline at end of file