research/poseidon: Clean up and add explanations.

script/research/poseidon/poseidon.py

@@ -3,11 +3,16 @@ import numpy
 from finite_fields.finitefield import IntegersModP
 from constants import MDS_matrix, round_constants
 
+# Width
 T = 3
+# Full rounds
 R_F = 8
+# Partial rounds
 R_P = 56
+# Sponge rate
 RATE = 2
 
+# pallas
 p = 0x40000000000000000000000000000000224698fc094cf91b992d30ed00000001
 Fp = IntegersModP(p)
 
@@ -35,7 +40,7 @@ def perm(inp):
             state_words[i] = state_words[i] + rcf[0]
             rcf.pop(0)
         for i in range(0, T):
-            state_words[i] = (state_words[i])**5
+            state_words[i] = (state_words[i])**5  # sbox
         state_words = numpy.array(numpy.dot(MDS_MATRIX, state_words))
 
     # Middle partial rounds
@@ -44,7 +49,7 @@ def perm(inp):
         for i in range(0, T):
             state_words[i] = state_words[i] + rcf[0]
             rcf.pop(0)
-        state_words[0] = (state_words[0])**5
+        state_words[0] = (state_words[0])**5  # sbox
         state_words = numpy.array(numpy.dot(MDS_MATRIX, state_words))
 
     # Last full rounds
@@ -54,17 +59,17 @@ def perm(inp):
             state_words[i] = state_words[i] + rcf[0]
             rcf.pop(0)
         for i in range(0, T):
-            state_words[i] = (state_words[i])**5
+            state_words[i] = (state_words[i])**5  # sbox
         state_words = numpy.array(numpy.dot(MDS_MATRIX, state_words))
 
     return state_words
 
 
-def debug(w, n, s, m):
+def debug(n, s, m):
     if enable_debug:
-        print(f"State {w} {n} absorb:")
+        print(f"State {n} absorb:")
         pprint([hex(int(i)) for i in s])
-        print(f"Mode {w} {n} absorb:")
+        print(f"Mode {n} absorb:")
         pprint([hex(int(i)) if i is not None else None for i in m])
 
 
@@ -79,12 +84,13 @@ def poseidon_hash(messages):
     output = [None] * RATE
     state = [Fp(0)] * T
 
+    # Capacity value is L ⋅ 2^64 + (o-1) where o is the output length
     initial_capacity_element = Fp(L << 64)
     state[RATE] = initial_capacity_element
 
-    # absorb sponge
+    # This outermost loop absorbs the messages in the sponge.
     for n, value in enumerate(messages):
-        debug("before", n + 1, state, mode)
+        debug(f"before {n+1}", state, mode)
         loop = False  # Use this to mark we should reiterate
         for i in range(0, len(mode)):
             if mode[i] is None:
@@ -93,7 +99,7 @@ def poseidon_hash(messages):
                 break
 
         if loop:
-            debug("after", n + 1, state, mode)
+            debug(f"after {n+1}", state, mode)
             continue
 
         # zip short-circuits when one iterator completes, so this will
@@ -101,29 +107,31 @@ def poseidon_hash(messages):
         for i, _ in enumerate(zip(state, mode)):
             state[i] += mode[i]
 
+        # Permutation of the current state
         state = perm(state)
 
-        output = [None] * RATE
         for i, _ in enumerate(zip(output, state)):
             output[i] = state[i]
 
+        # Reinit sponge with the current message as the first value.
         mode = [None] * RATE
         mode[0] = value
 
-        debug("after", n + 1, state, mode)
+        debug(f"after {n+1}", state, mode)
 
-    debug("before", "final", state, mode)
+    debug("before final", state, mode)
     for i, _ in enumerate(zip(state, mode)):
         state[i] += mode[i]
 
+    # Permutation of the final state
     state = perm(state)
 
-    output = [None] * RATE
     for i, _ in enumerate(zip(output, state)):
         output[i] = state[i]
 
+    # Sponge now has the output, so the first element is our hash.
     mode = output
-    debug("after", "final", state, mode)
+    debug("after final", state, mode)
     return output[0]
 
 
@@ -138,7 +146,8 @@ if __name__ == "__main__":
     #output_words = perm(input_words)
     #print([hex(int(i)) for i in output_words])
 
-    words = [Fp(1), Fp(2)]
-    #words = [Fp(1), Fp(2), Fp(3), Fp(4), Fp(5)]
-    h = poseidon_hash(words)
-    print(hex(int(h)))
+    words = []
+    for i in range(0, 10):
+        words.append(Fp(i))
+        h = poseidon_hash(words.copy())
+        print(hex(int(h)))