SHA3: clean up round functions

Primitive/Keyless/Hash/keccak.cry

@@ -129,39 +129,50 @@ private
     /**
      * One of the step mappings that's part of a round of Keccak-p.
      *
+     * The effect of this mapping is to XOR each bit in the state with the
+     * parity of two columns in the array. `C` computes the parities, `D`
+     * combines two of the parities for each column, and `A'` completes the
+     * transformation.
      * [FIPS-202] Section 3.2.1.
      */
     θ : State -> State
     θ A = A' where
-        C = [ xor a | a <- A ]
-        D = [ C @ x ^ (C @ y <<< 1)
-            // In the spec, this is denoted `(x-1) mod 5` for 0 <= x < 5.
-            | (x:[8]) <- [4,0,1,2,3]
-            // In the spec, this is denoted `(x+1) mod 5` for 0 <= x < 5.
-            | (y:[8]) <- [1,2,3,4,0]
-            ]
-        A' = [ [ a ^ (D @ x) | a <- A @ x ] | (x:[8]) <- [0 .. 4] ]
-
-        xor : {a, l} (fin a) => [a][l] -> [l]
-        xor xs = xors ! 0
-            where xors = [zero] # [ x ^ z | x <- xs | z <- xors ]
+        C = [ A@x@0 ^ A@x@1 ^ A@x@2 ^ A@x@3 ^ A@x@4
+            | x <- [0..4]]
+        D = [ C @ ((x - 1) % 5) ^ (C @ ((x + 1) % 5) <<< 1)
+            | x <- [0..4]]
+        A' = [[ A@x@y ^ D@x
+            | y <- [0..4]]
+            | x <- [0..4] ]
 
 
     /**
      * One of the step mappings that's part of a round of Keccak-p.
      *
+     * The effect of thuis mapping is to rotate the bits of each lane by a
+     * an _offset_ (computed in `set_lane`) that depends on the `x` and `y`
+     * coordinates of that lane.
      * [FIPS-202] Section 3.2.2.
      */
     ρ : State -> State
-    ρ A = groupBy`{5} [ a <<< r | a <- join A | (r:[8]) <- R ]
-        where
-            R = [
-                00, 36, 03, 41, 18,
-                01, 44, 10, 45, 02,
-                62, 06, 43, 15, 61,
-                28, 55, 25, 21, 56,
-                27, 20, 39, 08, 14
-            ]
+    ρ A = A' where
+        // Step 1.
+        A1 = [[ if (x == 0) && (y == 0) then A@0@0 else zero
+            | y <- [0..4]]
+            | x <- [0..4]]
+        // Step 2-3.
+        As = [((1,0), A1)] #
+            [ ((y, (2*x + 3*y) % 5), set_lane x y t Ai)
+            | ((x, y), Ai) <- As
+            | t <- [0..23]]
+        (_, A') = As ! 0
+
+        // Step 3a. Update the lane defined by x' and y'.
+        set_lane x' y' t Ai = [[
+            if (x' == x) && (y' == y) then A@x@y <<< ((t+1)*(t+2)/2)
+            else Ai@x@y
+            | y <- [0..4]]
+            | x <- [0..4]]
 
     /**
      * One of the step mappings that's part of a round of Keccak-p.