1	import Base
2	import ../u64.bend as W
3	import ./chacha8/block.bend as B
4	
5	# ChaCha8: Go's math/rand/v2.ChaCha8, the C2SP chacha8rand generator
6	# (https://c2sp.org/chacha8rand; Go's internal/chacha8rand), bit for bit.
7	#
8	#   new(seed)          Some{g} for a 32-byte seed (a list of 32 bytes < 256),
9	#                      None otherwise (Go's NewChaCha8([32]byte))
10	#   of_key(k)          the generator of a key given as eight little-endian
11	#                      words (B.K{k0, ..., k7}); new(seed) is of_key of its words
12	#   next(g)            the next 64-bit output and the advanced generator
13	#                      (Go's (*ChaCha8).Uint64); the Source interface of
14	#                      rand.bend: R.uint64n(~ChaCha8, ~next, g, n), ...
15	#
16	# Each iteration keys ChaCha8 with the 32-byte input, runs the sixteen
17	# blocks 0..15 in four groups of four (chacha8/block.bend), outputs the first
18	# 124 64-bit words and takes the last four (32 bytes) as the next key: the
19	# key is erased every 992 bytes of output (C2SP "fast key erasure"; Go
20	# reseeds the same way in (*State).Refill). The generator keeps the key, the
21	# next group and the unread words of the current group, so next is one list
22	# pop, and one group computation every 32 (the last group of an iteration:
23	# 28) outputs.
24	#
25	# Proved to produce the C2SP stream of spec/math/random/chacha8rand.bend for
26	# every seed (proofs/math/random/chacha8/); tested against Go's vectors
27	# (tests/math/random/, tools/check_random.py).
28	
29	# the next group of the iteration: blocks 0-3, 4-7, 8-11 or 12-15
30	type Phase is Data:
31	  G0{}
32	  G1{}
33	  G2{}
34	  G3{}
35	
36	type ChaCha8 is Data:
37	  C{key: B.Key, phase: Phase, buf: List<&2, W.U64>}
38	
39	def zero() -> W.U64:
40	  W.U64{0, 0}
41	
42	def of_key(k: B.Key) -> ChaCha8:
43	  C{k, G0{}, []}
44	
45	# the key of the next iteration: the last four words of the last group
46	def rekey(l: List<&2, W.U64>) -> B.Key:
47	  match l:
48	    case [W.U64{k0, k1}, W.U64{k2, k3}, W.U64{k4, k5}, W.U64{k6, k7}]:
49	      B.K{k0, k1, k2, k3, k4, k5, k6, k7}
50	    case _:
51	      B.K{0, 0, 0, 0, 0, 0, 0, 0}
52	
53	# the last group: 28 words of output, then the next key
54	def last(+g: List<&2, W.U64>) -> ChaCha8:
55	  C{rekey(List.drop(&2, W.U64, g, 28n)), G0{}, List.take(&2, W.U64, g, 28n)}
56	
57	def refill(+k: B.Key, p: Phase) -> ChaCha8:
58	  match p:
59	    case G0{}:
60	      C{k, G1{}, B.group(k, 0)}
61	    case G1{}:
62	      C{k, G2{}, B.group(k, 4)}
63	    case G2{}:
64	      C{k, G3{}, B.group(k, 8)}
65	    case G3{}:
66	      last(B.group(k, 12))
67	
68	def pop(g: ChaCha8) -> W.U64 & ChaCha8:
69	  match g:
70	    case C{k, p, buf}:
71	      match buf:
72	        case Con{x, rest}:
73	          (x, C{k, p, rest})
74	        case Nil{}:
75	          (zero(), C{k, p, []})
76	
77	# the next 64-bit output (a refill always yields a nonempty buffer)
78	def next(g: ChaCha8) -> W.U64 & ChaCha8:
79	  match g:
80	    case C{+k, p, buf}:
81	      match buf:
82	        case Con{x, rest}:
83	          (x, C{k, p, rest})
84	        case Nil{}:
85	          pop(refill(k, p))
86	
87	# ---- seeds ----
88	
89	# the little-endian word of four bytes (Go's byteorder.LEUint32)
90	def le32(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> U32:
91	  U32.or(U32.or(U32.or(b0, U32.shln(b1, 8n)), U32.shln(b2, 16n)), U32.shln(b3, 24n))
92	
93	def all_bytes(bs: List<&2, U32>) -> Bool:
94	  match bs:
95	    case Nil{}:
96	      True{}
97	    case Con{+b, rest}:
98	      Bool.and(U32.is_lt(b, 256), all_bytes(rest))
99	
100	def key_of(bs: List<&2, U32>) -> Maybe<&2, B.Key>:
101	  match bs:
102	    case [a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3, d0, d1, d2, d3, e0, e1, e2, e3, f0, f1, f2, f3, g0, g1, g2, g3, h0, h1, h2, h3]:
103	      Some{B.K{le32(a0, a1, a2, a3), le32(b0, b1, b2, b3), le32(c0, c1, c2, c3), le32(d0, d1, d2, d3), le32(e0, e1, e2, e3), le32(f0, f1, f2, f3), le32(g0, g1, g2, g3), le32(h0, h1, h2, h3)}}
104	    case _:
105	      None{}
106	
107	def lift(m: Maybe<&2, B.Key>) -> Maybe<&2, ChaCha8>:
108	  match m:
109	    case None{}:
110	      None{}
111	    case Some{k}:
112	      Some{of_key(k)}
113	
114	def seed_ok(bs: List<&2, U32>, ok: Bool) -> Maybe<&2, ChaCha8>:
115	  match ok:
116	    case False{}:
117	      None{}
118	    case True{}:
119	      lift(key_of(bs))
120	
121	# Go's NewChaCha8(seed): the seed is 32 bytes, each below 256
122	def new(+seed: List<&2, U32>) -> Maybe<&2, ChaCha8>:
123	  seed_ok(seed, all_bytes(seed))