1	import Base
2	import ./types.bend as A
3	import ./sub.bend as SB
4	
5	# The Argon2 memory: a packed Array<U32> of 2^(8 + k) words, block b at words
6	# 256 b .. 256 b + 255 (lane n of the block at words 256 b + 2 n, low half,
7	# and 256 b + 2 n + 1, high half). Blocks are read and written word by word
8	# with Base's Array.get and Array.set, which the compiler turns into O(1)
9	# accesses of the packed buffer.
10	
11	# Words i - n .. i, prepended to acc: pair holds the memory and word i.
12	def rd(n: Nat, +i: Nat, acc: List<&2, U32>, pair: Array<U32> & U32) -> Array<U32> & List<&2, U32>:
13	  match n pair:
14	    case 0n Tuple{a, w}:
15	      (a, w <> acc)
16	    case 1n+p Tuple{a, w}:
17	      +j = Nat.sub(i, 1n)
18	      rd(p, j, w <> acc, Array.get(U32, a, U32.from_nat(j)))
19	
20	def read_fin(pair: Array<U32> & List<&2, U32>) -> Array<U32> & A.Block:
21	  (a, ws) = pair
22	  (a, SB.rows8(ws))
23	
24	# Block b; the memory is handed back.
25	def get(+b: Nat, a: Array<U32>) -> Array<U32> & A.Block:
26	  +last = Nat.add(Nat.mul(b, 256n), 255n)
27	  read_fin(rd(255n, last, Nil{}, Array.get(U32, a, U32.from_nat(last))))
28	
29	# The words stored from word i on.
30	def wr(ws: List<&2, U32>, +i: Nat, a: Array<U32>) -> Array<U32>:
31	  match ws:
32	    case Nil{}:
33	      a
34	    case w <> rest:
35	      wr(rest, Nat.add(i, 1n), Array.set(U32, a, U32.from_nat(i), w))
36	
37	# The memory with block b replaced by v.
38	def set(+b: Nat, v: A.Block, a: Array<U32>) -> Array<U32>:
39	  wr(SB.words(v), Nat.mul(b, 256n), a)
40	
41	def pw(k: Nat) -> Nat:
42	  match k:
43	    case 0n:
44	      1n
45	    case 1n+p:
46	      Nat.double(pw(p))
47	
48	# The smallest k >= k0 with 2^k >= n.
49	def levels_go(fuel: Nat, +n: Nat, +k: Nat, done: Bool) -> Nat:
50	  match fuel done:
51	    case _ True{}:
52	      k
53	    case 0n False{}:
54	      k
55	    case 1n+f False{}:
56	      levels_go(f, n, 1n+k, Nat.is_le(n, pw(1n+k)))
57	
58	def levels(+n: Nat) -> Nat:
59	  levels_go(32n, n, 0n, Nat.is_le(n, 1n))
60	
61	# A zeroed memory of 2^k blocks.
62	def new(+k: Nat) -> Array<U32>:
63	  Array.new(U32, Nat.add(8n, k), 0)