1	import Base
2	import ../../lib/common.bend as C
3	import ../../../src/crypto/blake/blake2b/types.bend as T
4	import ../../../src/crypto/argon2/types.bend as A
5	import ../blake/blake2b.bend as B
6	
7	# The Argon2 compression function G (RFC 9106 sections 3.5 and 3.6), as an
8	# executable specification. It shares with the implementation only the
9	# record types (a 64-bit lane T.W{lo, hi}, a 16-lane state T.V, a 128-lane
10	# block A.B of eight rows), and takes the 64-bit addition, XOR and rotation
11	# from the RFC 7693 BLAKE2b specification (spec/crypto/blake/blake2b.bend).
12	# Lanes of a block are numbered 0..127, lane n being row n / 16, position
13	# n mod 16; the 1024-byte block is the 128 lanes in order, little-endian.
14	
15	# ---------------------------------------------------------------- GB
16	
17	# The lane of n mod 2^64 (U32.from_nat keeps the low 32 bits).
18	def lane_of(+n: Nat) -> T.Lane:
19	  T.W{U32.from_nat(C.low(32n, n)), U32.from_nat(C.high(32n, n))}
20	
21	# fBlaMka(x, y) = x + y + 2 * trunc(x) * trunc(y), modulo 2^64, where
22	# trunc(x) * trunc(y) is the 64-bit product of the low 32 bits of x and y.
23	def fBlaMka(x: T.Lane, y: T.Lane) -> T.Lane:
24	  match x y:
25	    case T.W{+xl, +xh} T.W{+yl, +yh}:
26	      +p = lane_of(Nat.mul(U32.to_nat(xl), U32.to_nat(yl)))
27	      B.add(B.add(T.W{xl, xh}, T.W{yl, yh}), B.add(p, p))
28	
29	# GB(a, b, c, d) (RFC 9106 section 3.6).
30	def GB(+a: T.Lane, +b: T.Lane, +c: T.Lane, +d: T.Lane) -> T.Quad:
31	  +a1 = fBlaMka(a, b)
32	  +d1 = B.rotr(B.xor(d, a1), 32n)
33	  +c1 = fBlaMka(c, d1)
34	  +b1 = B.rotr(B.xor(b, c1), 24n)
35	  +a2 = fBlaMka(a1, b1)
36	  +d2 = B.rotr(B.xor(d1, a2), 16n)
37	  +c2 = fBlaMka(c1, d2)
38	  +b2 = B.rotr(B.xor(b1, c2), 63n)
39	  T.Q{a2, b2, c2, d2}
40	
41	# ---------------------------------------------------------------- P
42	
43	def apply(v: T.State, +a: Nat, +b: Nat, +c: Nat, +d: Nat, q: T.Quad) -> T.State:
44	  match q:
45	    case T.Q{qa, qb, qc, qd}:
46	      B.put(B.put(B.put(B.put(v, a, qa), b, qb), c, qc), d, qd)
47	
48	# (v_a, v_b, v_c, v_d) <- GB(v_a, v_b, v_c, v_d)
49	def GBw(+w: T.State, +a: Nat, +b: Nat, +c: Nat, +d: Nat) -> T.State:
50	  apply(w, a, b, c, d, GB(B.at(w, a), B.at(w, b), B.at(w, c), B.at(w, d)))
51	
52	def GBi(v: T.State, +a: Nat, +b: Nat, +c: Nat, +d: Nat) -> T.State:
53	  match v:
54	    case T.V{x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15}:
55	      GBw(T.V{x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15}, a, b, c, d)
56	
57	# The permutation P on sixteen lanes v0..v15: GB on the columns (0, 4, 8, 12),
58	# (1, 5, 9, 13), (2, 6, 10, 14), (3, 7, 11, 15), then on the diagonals
59	# (0, 5, 10, 15), (1, 6, 11, 12), (2, 7, 8, 13), (3, 4, 9, 14).
60	def P(v: T.State) -> T.State:
61	  v = GBi(v, 0n, 4n, 8n, 12n)
62	  v = GBi(v, 1n, 5n, 9n, 13n)
63	  v = GBi(v, 2n, 6n, 10n, 14n)
64	  v = GBi(v, 3n, 7n, 11n, 15n)
65	  v = GBi(v, 0n, 5n, 10n, 15n)
66	  v = GBi(v, 1n, 6n, 11n, 12n)
67	  v = GBi(v, 2n, 7n, 8n, 13n)
68	  GBi(v, 3n, 4n, 9n, 14n)
69	
70	# ---------------------------------------------------------------- blocks
71	
72	def row(b: A.Block, i: Nat) -> T.State:
73	  match b i:
74	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 0n:
75	      r0
76	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 1n:
77	      r1
78	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 2n:
79	      r2
80	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 3n:
81	      r3
82	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 4n:
83	      r4
84	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 5n:
85	      r5
86	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 6n:
87	      r6
88	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} _:
89	      r7
90	
91	def put_row(b: A.Block, i: Nat, s: T.State) -> A.Block:
92	  match b i:
93	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 0n:
94	      A.B{s, r1, r2, r3, r4, r5, r6, r7}
95	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 1n:
96	      A.B{r0, s, r2, r3, r4, r5, r6, r7}
97	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 2n:
98	      A.B{r0, r1, s, r3, r4, r5, r6, r7}
99	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 3n:
100	      A.B{r0, r1, r2, s, r4, r5, r6, r7}
101	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 4n:
102	      A.B{r0, r1, r2, r3, s, r5, r6, r7}
103	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 5n:
104	      A.B{r0, r1, r2, r3, r4, s, r6, r7}
105	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} 6n:
106	      A.B{r0, r1, r2, r3, r4, r5, s, r7}
107	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7} _:
108	      A.B{r0, r1, r2, r3, r4, r5, r6, s}
109	
110	# Lane n of a block, and the block with lane n replaced.
111	def lane(+b: A.Block, +n: Nat) -> T.Lane:
112	  B.at(row(b, Nat.div(n, 16n)), Nat.mod(n, 16n))
113	
114	def put_lane(+b: A.Block, +n: Nat, x: T.Lane) -> A.Block:
115	  +i = Nat.div(n, 16n)
116	  put_row(b, i, B.put(row(b, i), Nat.mod(n, 16n), x))
117	
118	def ix(xs: List<&2, Nat>, n: Nat) -> Nat:
119	  B.nth(Nat, xs, n, 0n)
120	
121	# The sixteen lanes at the given lane numbers, as a state.
122	def gather(+b: A.Block, +idx: List<&2, Nat>) -> T.State:
123	  T.V{lane(b, ix(idx, 0n)), lane(b, ix(idx, 1n)), lane(b, ix(idx, 2n)), lane(b, ix(idx, 3n)), lane(b, ix(idx, 4n)), lane(b, ix(idx, 5n)), lane(b, ix(idx, 6n)), lane(b, ix(idx, 7n)), lane(b, ix(idx, 8n)), lane(b, ix(idx, 9n)), lane(b, ix(idx, 10n)), lane(b, ix(idx, 11n)), lane(b, ix(idx, 12n)), lane(b, ix(idx, 13n)), lane(b, ix(idx, 14n)), lane(b, ix(idx, 15n))}
124	
125	def put_lanes(idx: List<&2, Nat>, xs: List<&2, T.Lane>, b: A.Block) -> A.Block:
126	  match idx xs:
127	    case n <> ns x <> rest:
128	      put_lanes(ns, rest, put_lane(b, n, x))
129	    case _ _:
130	      b
131	
132	# The block with the state's sixteen lanes stored at the given lane numbers.
133	def scatter(b: A.Block, idx: List<&2, Nat>, s: T.State) -> A.Block:
134	  match s:
135	    case T.V{s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15}:
136	      put_lanes(idx, [s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15], b)
137	
138	# first, first + step, .., n terms.
139	def seq(n: Nat, +first: Nat, +step: Nat) -> List<&2, Nat>:
140	  match n:
141	    case 0n:
142	      Nil{}
143	    case 1n+k:
144	      first <> seq(k, Nat.add(first, step), step)
145	
146	# Row i: lanes 16i, 16i + 1, .., 16i + 15.
147	def row_idx(+i: Nat) -> List<&2, Nat>:
148	  seq(16n, Nat.mul(16n, i), 1n)
149	
150	# Column i: the 16-byte registers i, i + 8, .., i + 56, that is lanes
151	# 2i, 2i + 1, 2i + 16, 2i + 17, .., 2i + 112, 2i + 113.
152	def col_pairs(n: Nat, +first: Nat) -> List<&2, Nat>:
153	  match n:
154	    case 0n:
155	      Nil{}
156	    case 1n+k:
157	      first <> Nat.add(first, 1n) <> col_pairs(k, Nat.add(first, 16n))
158	
159	def col_idx(+i: Nat) -> List<&2, Nat>:
160	  col_pairs(8n, Nat.mul(2n, i))
161	
162	# P applied, independently, to the lanes of each index list of b.
163	def perms(+b: A.Block, idxs: List<&2, List<&2, Nat>>) -> List<&2, T.State>:
164	  match idxs:
165	    case Nil{}:
166	      Nil{}
167	    case i <> rest:
168	      P(gather(b, i)) <> perms(b, rest)
169	
170	def scatter_all(idxs: List<&2, List<&2, Nat>>, ss: List<&2, T.State>, b: A.Block) -> A.Block:
171	  match idxs ss:
172	    case i <> ix s <> rest:
173	      scatter_all(ix, rest, scatter(b, i, s))
174	    case _ _:
175	      b
176	
177	def permute(+b: A.Block, +idxs: List<&2, List<&2, Nat>>) -> A.Block:
178	  scatter_all(idxs, perms(b, idxs), b)
179	
180	def rows_all() -> List<&2, List<&2, Nat>>:
181	  [row_idx(0n), row_idx(1n), row_idx(2n), row_idx(3n), row_idx(4n), row_idx(5n), row_idx(6n), row_idx(7n)]
182	
183	def cols_all() -> List<&2, List<&2, Nat>>:
184	  [col_idx(0n), col_idx(1n), col_idx(2n), col_idx(3n), col_idx(4n), col_idx(5n), col_idx(6n), col_idx(7n)]
185	
186	# (Q0, .., Q7) <- P(R0, .., R7), .., (Q56, .., Q63) <- P(R56, .., R63): P on
187	# each row of eight 16-byte registers.
188	def rows(b: A.Block) -> A.Block:
189	  match b:
190	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7}:
191	      permute(A.B{r0, r1, r2, r3, r4, r5, r6, r7}, rows_all())
192	
193	# (Z0, Z8, .., Z56) <- P(Q0, Q8, .., Q56), .., (Z7, .., Z63) <- P(Q7, .., Q63):
194	# P on each column.
195	def cols(b: A.Block) -> A.Block:
196	  match b:
197	    case A.B{r0, r1, r2, r3, r4, r5, r6, r7}:
198	      permute(A.B{r0, r1, r2, r3, r4, r5, r6, r7}, cols_all())
199	
200	def xor_state(s: T.State, t: T.State) -> T.State:
201	  match s t:
202	    case T.V{s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15} T.V{t0, t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12, t13, t14, t15}:
203	      T.V{B.xor(s0, t0), B.xor(s1, t1), B.xor(s2, t2), B.xor(s3, t3), B.xor(s4, t4), B.xor(s5, t5), B.xor(s6, t6), B.xor(s7, t7), B.xor(s8, t8), B.xor(s9, t9), B.xor(s10, t10), B.xor(s11, t11), B.xor(s12, t12), B.xor(s13, t13), B.xor(s14, t14), B.xor(s15, t15)}
204	
205	# X XOR Y, lane by lane.
206	def xor(x: A.Block, y: A.Block) -> A.Block:
207	  match x y:
208	    case A.B{x0, x1, x2, x3, x4, x5, x6, x7} A.B{y0, y1, y2, y3, y4, y5, y6, y7}:
209	      A.B{xor_state(x0, y0), xor_state(x1, y1), xor_state(x2, y2), xor_state(x3, y3), xor_state(x4, y4), xor_state(x5, y5), xor_state(x6, y6), xor_state(x7, y7)}
210	
211	# G(X, Y) = Z XOR R, where R = X XOR Y, Q = P on the rows of R and Z = P on
212	# the columns of Q (RFC 9106 section 3.5).
213	def G(x: A.Block, y: A.Block) -> A.Block:
214	  +r = xor(x, y)
215	  xor(cols(rows(r)), r)
216	
217	def zero_state() -> T.State:
218	  T.V{T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}, T.W{0, 0}}
219	
220	# The all-zero block.
221	def zero() -> A.Block:
222	  A.B{zero_state(), zero_state(), zero_state(), zero_state(), zero_state(), zero_state(), zero_state(), zero_state()}