1	# Generated by tools/generators/argon2_gen.py; do not edit by hand.
2	import Base
3	import ../../math/w64.bend as X
4	import ../blake/blake2b/types.bend as T
5	import ../blake/blake2b/lane.bend as L
6	import ./types.bend as A
7	
8	# The low half of a lane (trunc in RFC 9106).
9	def lo(x: T.Lane) -> U32:
10	  match x:
11	    case T.W{l,h}: l
12	
13	# The full 64-bit product of two U32 (src/math/w64.bend, proved).
14	def prod(+a: U32, +b: U32) -> T.Lane:
15	  +m = X.mul32(a,b)
16	  T.W{X.lo(m),X.hi(m)}
17	
18	# fBlaMka(x, y) = x + y + 2 * trunc(x) * trunc(y) mod 2^64 (RFC 9106 section 3.6).
19	def bla(x: T.Lane, y: T.Lane) -> T.Lane:
20	  match x y:
21	    case T.W{+xl,+xh} T.W{+yl,+yh}:
22	      +p = prod(xl,yl)
23	      L.add(L.add(T.W{xl,xh},T.W{yl,yh}),L.add(p,p))
24	
25	# GB (RFC 9106 section 3.6), one assignment pair per function.
26	def gb4(+a2: T.Lane, +b1: T.Lane, +c2: T.Lane, +d2: T.Lane) -> T.Quad:
27	  T.Q{a2,L.rot63(L.xor(b1,c2)),c2,d2}
28	
29	def gb3(+a2: T.Lane, +b1: T.Lane, +c1: T.Lane, +d1: T.Lane) -> T.Quad:
30	  +d2 = L.rot16(L.xor(d1,a2))
31	  gb4(a2,b1,bla(c1,d2),d2)
32	
33	def gb2(+a1: T.Lane, +b: T.Lane, +c1: T.Lane, +d1: T.Lane) -> T.Quad:
34	  +b1 = L.rot24(L.xor(b,c1))
35	  gb3(bla(a1,b1),b1,c1,d1)
36	
37	def gb1(+a1: T.Lane, +b: T.Lane, +c: T.Lane, +d: T.Lane) -> T.Quad:
38	  +d1 = L.rot32(L.xor(d,a1))
39	  gb2(a1,b,bla(c,d1),d1)
40	
41	def gb(+a: T.Lane, +b: T.Lane, +c: T.Lane, +d: T.Lane) -> T.Quad:
42	  gb1(bla(a,b),b,c,d)
43	
44	def put0(q: T.Quad, x1: T.Lane, x2: T.Lane, x3: T.Lane, x5: T.Lane, x6: T.Lane, x7: T.Lane, x9: T.Lane, x10: T.Lane, x11: T.Lane, x13: T.Lane, x14: T.Lane, x15: T.Lane) -> T.State:
45	  match q:
46	    case T.Q{ya,yb,yc,yd}:
47	      T.V{ya,x1,x2,x3,yb,x5,x6,x7,yc,x9,x10,x11,yd,x13,x14,x15}
48	
49	# GB on positions 0, 4, 8, 12.
50	def step0(v: T.State) -> T.State:
51	  match v:
52	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15}:
53	      put0(gb(x0,x4,x8,x12),x1,x2,x3,x5,x6,x7,x9,x10,x11,x13,x14,x15)
54	
55	def put1(q: T.Quad, x0: T.Lane, x2: T.Lane, x3: T.Lane, x4: T.Lane, x6: T.Lane, x7: T.Lane, x8: T.Lane, x10: T.Lane, x11: T.Lane, x12: T.Lane, x14: T.Lane, x15: T.Lane) -> T.State:
56	  match q:
57	    case T.Q{ya,yb,yc,yd}:
58	      T.V{x0,ya,x2,x3,x4,yb,x6,x7,x8,yc,x10,x11,x12,yd,x14,x15}
59	
60	# GB on positions 1, 5, 9, 13.
61	def step1(v: T.State) -> T.State:
62	  match v:
63	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15}:
64	      put1(gb(x1,x5,x9,x13),x0,x2,x3,x4,x6,x7,x8,x10,x11,x12,x14,x15)
65	
66	def put2(q: T.Quad, x0: T.Lane, x1: T.Lane, x3: T.Lane, x4: T.Lane, x5: T.Lane, x7: T.Lane, x8: T.Lane, x9: T.Lane, x11: T.Lane, x12: T.Lane, x13: T.Lane, x15: T.Lane) -> T.State:
67	  match q:
68	    case T.Q{ya,yb,yc,yd}:
69	      T.V{x0,x1,ya,x3,x4,x5,yb,x7,x8,x9,yc,x11,x12,x13,yd,x15}
70	
71	# GB on positions 2, 6, 10, 14.
72	def step2(v: T.State) -> T.State:
73	  match v:
74	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15}:
75	      put2(gb(x2,x6,x10,x14),x0,x1,x3,x4,x5,x7,x8,x9,x11,x12,x13,x15)
76	
77	def put3(q: T.Quad, x0: T.Lane, x1: T.Lane, x2: T.Lane, x4: T.Lane, x5: T.Lane, x6: T.Lane, x8: T.Lane, x9: T.Lane, x10: T.Lane, x12: T.Lane, x13: T.Lane, x14: T.Lane) -> T.State:
78	  match q:
79	    case T.Q{ya,yb,yc,yd}:
80	      T.V{x0,x1,x2,ya,x4,x5,x6,yb,x8,x9,x10,yc,x12,x13,x14,yd}
81	
82	# GB on positions 3, 7, 11, 15.
83	def step3(v: T.State) -> T.State:
84	  match v:
85	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15}:
86	      put3(gb(x3,x7,x11,x15),x0,x1,x2,x4,x5,x6,x8,x9,x10,x12,x13,x14)
87	
88	def put4(q: T.Quad, x1: T.Lane, x2: T.Lane, x3: T.Lane, x4: T.Lane, x6: T.Lane, x7: T.Lane, x8: T.Lane, x9: T.Lane, x11: T.Lane, x12: T.Lane, x13: T.Lane, x14: T.Lane) -> T.State:
89	  match q:
90	    case T.Q{ya,yb,yc,yd}:
91	      T.V{ya,x1,x2,x3,x4,yb,x6,x7,x8,x9,yc,x11,x12,x13,x14,yd}
92	
93	# GB on positions 0, 5, 10, 15.
94	def step4(v: T.State) -> T.State:
95	  match v:
96	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15}:
97	      put4(gb(x0,x5,x10,x15),x1,x2,x3,x4,x6,x7,x8,x9,x11,x12,x13,x14)
98	
99	def put5(q: T.Quad, x0: T.Lane, x2: T.Lane, x3: T.Lane, x4: T.Lane, x5: T.Lane, x7: T.Lane, x8: T.Lane, x9: T.Lane, x10: T.Lane, x13: T.Lane, x14: T.Lane, x15: T.Lane) -> T.State:
100	  match q:
101	    case T.Q{ya,yb,yc,yd}:
102	      T.V{x0,ya,x2,x3,x4,x5,yb,x7,x8,x9,x10,yc,yd,x13,x14,x15}
103	
104	# GB on positions 1, 6, 11, 12.
105	def step5(v: T.State) -> T.State:
106	  match v:
107	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15}:
108	      put5(gb(x1,x6,x11,x12),x0,x2,x3,x4,x5,x7,x8,x9,x10,x13,x14,x15)
109	
110	def put6(q: T.Quad, x0: T.Lane, x1: T.Lane, x3: T.Lane, x4: T.Lane, x5: T.Lane, x6: T.Lane, x9: T.Lane, x10: T.Lane, x11: T.Lane, x12: T.Lane, x14: T.Lane, x15: T.Lane) -> T.State:
111	  match q:
112	    case T.Q{ya,yb,yc,yd}:
113	      T.V{x0,x1,ya,x3,x4,x5,x6,yb,yc,x9,x10,x11,x12,yd,x14,x15}
114	
115	# GB on positions 2, 7, 8, 13.
116	def step6(v: T.State) -> T.State:
117	  match v:
118	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15}:
119	      put6(gb(x2,x7,x8,x13),x0,x1,x3,x4,x5,x6,x9,x10,x11,x12,x14,x15)
120	
121	def put7(q: T.Quad, x0: T.Lane, x1: T.Lane, x2: T.Lane, x5: T.Lane, x6: T.Lane, x7: T.Lane, x8: T.Lane, x10: T.Lane, x11: T.Lane, x12: T.Lane, x13: T.Lane, x15: T.Lane) -> T.State:
122	  match q:
123	    case T.Q{ya,yb,yc,yd}:
124	      T.V{x0,x1,x2,ya,yb,x5,x6,x7,x8,yc,x10,x11,x12,x13,yd,x15}
125	
126	# GB on positions 3, 4, 9, 14.
127	def step7(v: T.State) -> T.State:
128	  match v:
129	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15}:
130	      put7(gb(x3,x4,x9,x14),x0,x1,x2,x5,x6,x7,x8,x10,x11,x12,x13,x15)
131	
132	# The permutation P on sixteen lanes.
133	def p(v: T.State) -> T.State:
134	  step7(step6(step5(step4(step3(step2(step1(step0(v))))))))
135	
136	# P on each row: lanes 16i .. 16i+15.
137	def rows(b: A.Block) -> A.Block:
138	  match b:
139	    case A.B{r0,r1,r2,r3,r4,r5,r6,r7}:
140	      A.B{p(r0),p(r1),p(r2),p(r3),p(r4),p(r5),p(r6),p(r7)}
141	
142	# The eight permuted columns back into rows: column i holds lanes 2i, 2i+1 of every row.
143	def untr(c0: T.State, c1: T.State, c2: T.State, c3: T.State, c4: T.State, c5: T.State, c6: T.State, c7: T.State) -> A.Block:
144	  match c0 c1 c2 c3 c4 c5 c6 c7:
145	    case T.V{y0_0,y0_1,y0_2,y0_3,y0_4,y0_5,y0_6,y0_7,y0_8,y0_9,y0_10,y0_11,y0_12,y0_13,y0_14,y0_15} T.V{y1_0,y1_1,y1_2,y1_3,y1_4,y1_5,y1_6,y1_7,y1_8,y1_9,y1_10,y1_11,y1_12,y1_13,y1_14,y1_15} T.V{y2_0,y2_1,y2_2,y2_3,y2_4,y2_5,y2_6,y2_7,y2_8,y2_9,y2_10,y2_11,y2_12,y2_13,y2_14,y2_15} T.V{y3_0,y3_1,y3_2,y3_3,y3_4,y3_5,y3_6,y3_7,y3_8,y3_9,y3_10,y3_11,y3_12,y3_13,y3_14,y3_15} T.V{y4_0,y4_1,y4_2,y4_3,y4_4,y4_5,y4_6,y4_7,y4_8,y4_9,y4_10,y4_11,y4_12,y4_13,y4_14,y4_15} T.V{y5_0,y5_1,y5_2,y5_3,y5_4,y5_5,y5_6,y5_7,y5_8,y5_9,y5_10,y5_11,y5_12,y5_13,y5_14,y5_15} T.V{y6_0,y6_1,y6_2,y6_3,y6_4,y6_5,y6_6,y6_7,y6_8,y6_9,y6_10,y6_11,y6_12,y6_13,y6_14,y6_15} T.V{y7_0,y7_1,y7_2,y7_3,y7_4,y7_5,y7_6,y7_7,y7_8,y7_9,y7_10,y7_11,y7_12,y7_13,y7_14,y7_15}:
146	      A.B{T.V{y0_0,y0_1,y1_0,y1_1,y2_0,y2_1,y3_0,y3_1,y4_0,y4_1,y5_0,y5_1,y6_0,y6_1,y7_0,y7_1},T.V{y0_2,y0_3,y1_2,y1_3,y2_2,y2_3,y3_2,y3_3,y4_2,y4_3,y5_2,y5_3,y6_2,y6_3,y7_2,y7_3},T.V{y0_4,y0_5,y1_4,y1_5,y2_4,y2_5,y3_4,y3_5,y4_4,y4_5,y5_4,y5_5,y6_4,y6_5,y7_4,y7_5},T.V{y0_6,y0_7,y1_6,y1_7,y2_6,y2_7,y3_6,y3_7,y4_6,y4_7,y5_6,y5_7,y6_6,y6_7,y7_6,y7_7},T.V{y0_8,y0_9,y1_8,y1_9,y2_8,y2_9,y3_8,y3_9,y4_8,y4_9,y5_8,y5_9,y6_8,y6_9,y7_8,y7_9},T.V{y0_10,y0_11,y1_10,y1_11,y2_10,y2_11,y3_10,y3_11,y4_10,y4_11,y5_10,y5_11,y6_10,y6_11,y7_10,y7_11},T.V{y0_12,y0_13,y1_12,y1_13,y2_12,y2_13,y3_12,y3_13,y4_12,y4_13,y5_12,y5_13,y6_12,y6_13,y7_12,y7_13},T.V{y0_14,y0_15,y1_14,y1_15,y2_14,y2_15,y3_14,y3_15,y4_14,y4_15,y5_14,y5_15,y6_14,y6_15,y7_14,y7_15}}
147	
148	# P on each column: lanes 2i, 2i+1, 2i+16, 2i+17, .., 2i+112, 2i+113.
149	def cols(b: A.Block) -> A.Block:
150	  match b:
151	    case A.B{T.V{l0,l1,l2,l3,l4,l5,l6,l7,l8,l9,l10,l11,l12,l13,l14,l15},T.V{l16,l17,l18,l19,l20,l21,l22,l23,l24,l25,l26,l27,l28,l29,l30,l31},T.V{l32,l33,l34,l35,l36,l37,l38,l39,l40,l41,l42,l43,l44,l45,l46,l47},T.V{l48,l49,l50,l51,l52,l53,l54,l55,l56,l57,l58,l59,l60,l61,l62,l63},T.V{l64,l65,l66,l67,l68,l69,l70,l71,l72,l73,l74,l75,l76,l77,l78,l79},T.V{l80,l81,l82,l83,l84,l85,l86,l87,l88,l89,l90,l91,l92,l93,l94,l95},T.V{l96,l97,l98,l99,l100,l101,l102,l103,l104,l105,l106,l107,l108,l109,l110,l111},T.V{l112,l113,l114,l115,l116,l117,l118,l119,l120,l121,l122,l123,l124,l125,l126,l127}}:
152	      untr(p(T.V{l0,l1,l16,l17,l32,l33,l48,l49,l64,l65,l80,l81,l96,l97,l112,l113}),p(T.V{l2,l3,l18,l19,l34,l35,l50,l51,l66,l67,l82,l83,l98,l99,l114,l115}),p(T.V{l4,l5,l20,l21,l36,l37,l52,l53,l68,l69,l84,l85,l100,l101,l116,l117}),p(T.V{l6,l7,l22,l23,l38,l39,l54,l55,l70,l71,l86,l87,l102,l103,l118,l119}),p(T.V{l8,l9,l24,l25,l40,l41,l56,l57,l72,l73,l88,l89,l104,l105,l120,l121}),p(T.V{l10,l11,l26,l27,l42,l43,l58,l59,l74,l75,l90,l91,l106,l107,l122,l123}),p(T.V{l12,l13,l28,l29,l44,l45,l60,l61,l76,l77,l92,l93,l108,l109,l124,l125}),p(T.V{l14,l15,l30,l31,l46,l47,l62,l63,l78,l79,l94,l95,l110,l111,l126,l127}))
153	
154	# Lane-wise XOR of two rows, and of two blocks (row by row).
155	def xs(s: T.State, t: T.State) -> T.State:
156	  match s t:
157	    case T.V{x0,x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x11,x12,x13,x14,x15} T.V{y0,y1,y2,y3,y4,y5,y6,y7,y8,y9,y10,y11,y12,y13,y14,y15}:
158	      T.V{L.xor(x0,y0),L.xor(x1,y1),L.xor(x2,y2),L.xor(x3,y3),L.xor(x4,y4),L.xor(x5,y5),L.xor(x6,y6),L.xor(x7,y7),L.xor(x8,y8),L.xor(x9,y9),L.xor(x10,y10),L.xor(x11,y11),L.xor(x12,y12),L.xor(x13,y13),L.xor(x14,y14),L.xor(x15,y15)}
159	
160	def xor(x: A.Block, y: A.Block) -> A.Block:
161	  match x y:
162	    case A.B{x0,x1,x2,x3,x4,x5,x6,x7} A.B{y0,y1,y2,y3,y4,y5,y6,y7}:
163	      A.B{xs(x0,y0),xs(x1,y1),xs(x2,y2),xs(x3,y3),xs(x4,y4),xs(x5,y5),xs(x6,y6),xs(x7,y7)}
164	
165	# G(X, Y) = P(R) XOR R with R = X XOR Y, P on the rows then the columns
166	# (RFC 9106 section 3.5).
167	def fin(+r: A.Block) -> A.Block:
168	  xor(cols(rows(r)),r)
169	
170	def compress(x: A.Block, y: A.Block) -> A.Block:
171	  fin(xor(x,y))
172	
173	# Version 0x13 passes after the first: G(X, Y) XOR the block being overwritten.
174	def compress_xor(x: A.Block, y: A.Block, old: A.Block) -> A.Block:
175	  xor(compress(x,y),old)
176	
177	# The all-zero block.
178	def zero() -> A.Block:
179	  A.B{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}},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}},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}},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}},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}},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}},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}},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}}}