crypto: poseidon2: Clean the stack in Poseidon utils (#9)

renegade-fi · Jan 27, 2025 · 1636986 · 1636986
1 parent 59a8817
commit 1636986
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/src/crypto/poseidon2/roundUtils.huff b/src/crypto/poseidon2/roundUtils.huff
@@ -35,6 +35,9 @@
     ADD_RC(<RC1>)           // [a + RC1, (b + RC2)^5, (c + RC3)^5, a, b, c]
     PUSH_PRIME() swap1      // [a + RC1, PRIME, (b + RC2)^5, (c + RC3)^5, a, b, c] 
     SBOX()                  // [(a + RC1)^5, (b + RC2)^5, (c + RC3)^5, a, b, c]
+
+    // Pop the old state off the stack
+    swap3 pop swap3 pop swap3 pop // [a', b', c']
 
     // Multiply the intermediate state by the external round MDS matrix
     EXTERNAL_MDS()
@@ -104,10 +107,15 @@
     PUSH_PRIME() dup2 PUSH_PRIME()      // [PRIME, sum, PRIME, sum, state[0], state[1], state[2]] 
     dup7 dup1 addmod                    // [state[2] * 2, sum, PRIME, sum, state[0], state[1], state[2]] 
     addmod                              // [state'[2], sum, state[0], state[1], state[2]]           
-    PUSH_PRIME() dup3                   // [sum, PRIME, state'[2], state[0], state[1], state[2]]
-    dup6 addmod                         // [state'[1], state'[2], sum, state[0], state[1], state[2]]
-    PUSH_PRIME() dup4                   // [sum, PRIME, state'[1], state'[2], sum, state[0], state[1], state[2]]
-    dup6 addmod                         // [state'[0], state'[1], state'[2], sum, state[0], state[1], state[2]]
+    swap4 pop                           // [sum, state[0], state[1], state'[2]]
+
+    PUSH_PRIME() dup2                   // [sum, PRIME, sum, state[0], state[1], state'[2]]
+    dup5 addmod                         // [state'[1], sum, state[0], state[1], state'[2]]
+    swap3 pop                           // [sum, state[0], state'[1], state'[2]]
+
+    PUSH_PRIME() dup2                   // [sum, PRIME, sum, state[0], state'[1], state'[2]]
+    dup4 addmod                         // [state'[0], sum, state[0], state'[1], state'[2]]
+    swap2 pop pop                       // [state'[0], state'[1], state'[2]]
 }
 
 /// @dev Apply the external MDS matrix to the sponge state