1	# Bit readers and writers over Bytes, LSB-first (DEFLATE) or MSB-first (JPEG), and canonical Huffman decode tables. Source: https://github.com/paymog/bend-kit/tree/main/bits
2	import Base
3	import bend-kit-bytes@0.3.3.0/bytes.bend as Bytes
4	
5	# Bytes packs byte i at bits 8*(i%4) of word i/4, so LSB-first stream bit p is bit p%32 of
6	# word p/32, and MSB-first it is bit 31 - p%32 of that word byte-swapped.
7	# A U32 shift by 32 or more gives 0; the code below relies on it.
8	# ponytail: positions are U32 bit counts, so a stream holds under 512 MiB; widen to U64 if that bites.
9	
10	def mask(+n: U32) -> U32:
11	  (U32.shln(1, U32.to_nat(n)) - 1 : U32)
12	
13	def win.lsb.join(+s: U32, +lo: U32, r: Array<U32> & U32) -> Array<U32> & U32:
14	  (a, +hi) = r
15	  (a, (U32.shrn(lo, U32.to_nat(s)) .|. U32.shln(hi, U32.to_nat((32 - s : U32))) : U32))
16	
17	def win.msb.join(+s: U32, +lo: U32, r: Array<U32> & U32) -> Array<U32> & U32:
18	  (a, +hi) = r
19	  (a, (U32.shln(Bytes.bswap(lo), U32.to_nat(s)) .|. U32.shrn(Bytes.bswap(hi), U32.to_nat((32 - s : U32))) : U32))
20	
21	def win.lsb.at(+p: U32, r: Array<U32> & U32) -> Array<U32> & U32:
22	  (a, +lo) = r
23	  win.lsb.join((p .&. 31 : U32), lo, Array.get(U32, a, ((p >> 5n) + 1 : U32)))
24	
25	def win.msb.at(+p: U32, r: Array<U32> & U32) -> Array<U32> & U32:
26	  (a, +lo) = r
27	  win.msb.join((p .&. 31 : U32), lo, Array.get(U32, a, ((p >> 5n) + 1 : U32)))
28	
29	# The 32 stream bits from bit p, in stream order: first bit lowest LSB-first, highest MSB-first.
30	# Bits past the buffer's last word wrap to its first, so callers mask what lies past the end.
31	def win(msb: Bool, a: Array<U32>, +p: U32) -> Array<U32> & U32:
32	  match msb:
33	    case True{}:
34	      win.msb.at(p, Array.get(U32, a, (p >> 5n : U32)))
35	    case False{}:
36	      win.lsb.at(p, Array.get(U32, a, (p >> 5n : U32)))
37	
38	# The first n (at most 32) bits of a window as a number, first stream bit most significant MSB-first.
39	def field(+msb: Bool, +n: U32, +w: U32) -> U32:
40	  Bool.pick(U32, msb, U32.shrn(w, U32.to_nat((32 - n : U32))), (w .&. mask(n) : U32))
41	
42	# A reader owns the buffer of len bytes; pos is the next bit.
43	type Reader is Type:
44	  Reader{len: U32, buf: Array<U32>, pos: U32}
45	
46	def Reader.new(b: Bytes.Bytes) -> Reader:
47	  Bytes.Bytes{len, buf} = b
48	  Reader{len, buf, 0}
49	
50	# The buffer back, and the bit position.
51	def Reader.finish(r: Reader) -> Bytes.Bytes & U32:
52	  Reader{len, buf, pos} = r
53	  (Bytes.Bytes{len, buf}, pos)
54	
55	def Reader.left(r: Reader) -> Reader & U32:
56	  Reader{+len, buf, +pos} = r
57	  (Reader{len, buf, pos}, ((len << 3n) - pos : U32))
58	
59	def read.got(+msb: Bool, +len: U32, +pos: U32, +n: U32, r: Array<U32> & U32) -> Reader & Maybe<&2, U32>:
60	  (a, +w) = r
61	  (Reader{len, a, (pos + n : U32)}, Some{field(msb, n, w)})
62	
63	def read.if(ok: Bool, +msb: Bool, +len: U32, buf: Array<U32>, +pos: U32, +n: U32) -> Reader & Maybe<&2, U32>:
64	  match ok:
65	    case True{}:
66	      read.got(msb, len, pos, n, win(msb, buf, pos))
67	    case False{}:
68	      (Reader{len, buf, pos}, None{})
69	
70	def read(+msb: Bool, r: Reader, +n: U32) -> Reader & Maybe<&2, U32>:
71	  Reader{+len, buf, +pos} = r
72	  read.if(U32.is_le(n, U32.min(32, ((len << 3n) - pos : U32))), msb, len, buf, pos, n)
73	
74	# The next n bits (0 to 32) as a number, and the reader past them. None, with the reader
75	# unchanged, when n > 32 or fewer than n bits are left.
76	def Reader.lsb(r: Reader, +n: U32) -> Reader & Maybe<&2, U32>:
77	  read(False{}, r, n)
78	
79	def Reader.msb(r: Reader, +n: U32) -> Reader & Maybe<&2, U32>:
80	  read(True{}, r, n)
81	
82	# Bits past the end read as 0: the low avail bits stay LSB-first, the high ones MSB-first.
83	def peek.got(+msb: Bool, +len: U32, +pos: U32, +n: U32, +avail: U32, r: Array<U32> & U32) -> Reader & U32:
84	  (a, +w) = r
85	  +v = field(msb, n, w)
86	  (Reader{len, a, pos}, Bool.pick(U32, msb, (v .&. U32.shln(4294967295, U32.to_nat((n - avail : U32))) : U32), (v .&. mask(avail) : U32)))
87	
88	def peek(+msb: Bool, r: Reader, +n0: U32) -> Reader & U32:
89	  Reader{+len, buf, +pos} = r
90	  +n = U32.min(n0, 32)
91	  peek.got(msb, len, pos, n, U32.min(n, ((len << 3n) - pos : U32)), win(msb, buf, pos))
92	
93	# The next n bits (at most 32) without moving. Bits past the end read as 0.
94	def Reader.peek.lsb(r: Reader, +n: U32) -> Reader & U32:
95	  peek(False{}, r, n)
96	
97	def Reader.peek.msb(r: Reader, +n: U32) -> Reader & U32:
98	  peek(True{}, r, n)
99	
100	# Moves n bits on. False, with the reader unchanged, when fewer than n are left.
101	def Reader.skip(r: Reader, +n: U32) -> Reader & Bool:
102	  Reader{+len, buf, +pos} = r
103	  +ok = U32.is_le(n, ((len << 3n) - pos : U32))
104	  (Reader{len, buf, Bool.pick(U32, ok, (pos + n : U32), pos)}, ok)
105	
106	# Moves to the next byte boundary, as a DEFLATE stored block or a JPEG marker needs.
107	def Reader.align(r: Reader) -> Reader:
108	  Reader{len, buf, +pos} = r
109	  Reader{len, buf, ((pos + 7) .&. 4294967288 : U32)}
110	
111	# A writer's buf has 2^dep words, all 0 past bit pos.
112	type Writer is Type:
113	  Writer{buf: Array<U32>, dep: U32, pos: U32}
114	
115	def Writer.new() -> Writer:
116	  Writer{Array.new(U32, 0n, 0), 0, 0}
117	
118	# A write adds at most 32 bits, so one doubling always makes room.
119	def room(full: Bool, w: Writer) -> Writer:
120	  match full:
121	    case False{}:
122	      w
123	    case True{}:
124	      Writer{buf, +dep, pos} = w
125	      Writer{Bytes.dst(Bytes.copy(U32.shln(4, U32.to_nat(dep)), buf, Array.new(U32, U32.to_nat((dep + 1 : U32)), 0), 0, 0)), (dep + 1 : U32), pos}
126	
127	def or.of(+k: U32, +x: U32, r: Array<U32> & U32) -> Array<U32>:
128	  (a, +old) = r
129	  Array.set(U32, a, k, (old .|. x : U32))
130	
131	def or.if(some: Bool, a: Array<U32>, +k: U32, +x: U32) -> Array<U32>:
132	  match some:
133	    case True{}:
134	      or.of(k, x, Array.get(U32, a, k))
135	    case False{}:
136	      a
137	
138	# The parts of an n-bit v that land in word k and word k + 1, as stored words.
139	def parts(msb: Bool, +s: U32, +n: U32, +v: U32) -> U32 & U32:
140	  match msb:
141	    case True{}:
142	      +t = U32.shln(v, U32.to_nat((32 - n : U32)))
143	      (Bytes.bswap(U32.shrn(t, U32.to_nat(s))), Bytes.bswap(U32.shln(t, U32.to_nat((32 - s : U32)))))
144	    case False{}:
145	      +m = (v .&. mask(n) : U32)
146	      (U32.shln(m, U32.to_nat(s)), U32.shrn(m, U32.to_nat((32 - s : U32))))
147	
148	def write.put(+k: U32, x: U32 & U32, a: Array<U32>) -> Array<U32>:
149	  (+x0, +x1) = x
150	  or.if(U32.is_ne(x1, 0), or.if(U32.is_ne(x0, 0), a, k, x0), (k + 1 : U32), x1)
151	
152	def write.at(+msb: Bool, w: Writer, +n: U32, +v: U32) -> Writer:
153	  Writer{buf, dep, +pos} = w
154	  Writer{write.put((pos >> 5n : U32), parts(msb, (pos .&. 31 : U32), n, v), buf), dep, (pos + n : U32)}
155	
156	def write(+msb: Bool, w: Writer, +n0: U32, +v: U32) -> Writer:
157	  Writer{buf, +dep, +pos} = w
158	  +n = U32.min(n0, 32)
159	  write.at(msb, room(U32.is_lt(U32.shln(1, U32.to_nat(dep)), ((pos + n + 31) >> 5n : U32)), Writer{buf, dep, pos}), n, v)
160	
161	# Appends the low n bits of v (n is cut to 32). LSB-first sends v's lowest bit first;
162	# MSB-first sends bit n - 1 first.
163	def Writer.lsb(w: Writer, +n: U32, +v: U32) -> Writer:
164	  write(False{}, w, n, v)
165	
166	def Writer.msb(w: Writer, +n: U32, +v: U32) -> Writer:
167	  write(True{}, w, n, v)
168	
169	# Pads with 0 bits to the next byte boundary.
170	def Writer.align(w: Writer) -> Writer:
171	  Writer{buf, dep, +pos} = w
172	  Writer{buf, dep, ((pos + 7) .&. 4294967288 : U32)}
173	
174	# The bytes written; 0 bits fill the last byte.
175	def Writer.finish(w: Writer) -> Bytes.Bytes:
176	  Writer{buf, dep, +pos} = w
177	  Bytes.Bytes{((pos + 7) >> 3n : U32), buf}
178	
179	# A canonical Huffman table: entry i of 2^max is (sym << 5) | len for the code that the
180	# first max stream bits i start with, or 0 when no code does. msb is the stream's bit order.
181	# ponytail: one level of 2^max entries; a two-level table if building it dominates short blocks.
182	type Huff is Type:
183	  Huff{msb: Bool, max: U32, table: Array<U32>}
184	
185	# One symbol and the length of its code.
186	type Code is Data:
187	  Code{sym: U32, len: U32}
188	
189	def lens.max(xs: List<&2, Code>, +m: U32) -> U32:
190	  match xs:
191	    case Nil{}:
192	      m
193	    case Con{Code{sym, +len}, t}:
194	      lens.max(t, U32.max(m, len))
195	
196	def count.of(+len: U32, r: Array<U32> & U32) -> Array<U32>:
197	  (c, +v) = r
198	  Array.set(U32, c, len, (v + 1 : U32))
199	
200	def count.if(zero: Bool, c: Array<U32>, +len: U32) -> Array<U32>:
201	  match zero:
202	    case True{}:
203	      c
204	    case False{}:
205	      count.of(len, Array.get(U32, c, len))
206	
207	def count(xs: List<&2, Code>, c: Array<U32>) -> Array<U32>:
208	  match xs:
209	    case Nil{}:
210	      c
211	    case Con{Code{sym, +len}, t}:
212	      count(t, count.if(U32.is_eq(len, 0), c, len))
213	
214	# RFC 1951 §3.2.2: next[l] = (next[l-1] + count[l-1]) << 1. r holds count[l-1]; left counts
215	# the codes of length l - 1 still free, and a negative left means the lengths over-subscribe.
216	def next.go(k: Nat, +l: U32, +code: U32, +left: U32, +ok: Bool, n: Array<U32>, r: Array<U32> & U32) -> Array<U32> & Bool:
217	  match k:
218	    case 0n:
219	      (c, v) = r
220	      (n, ok)
221	    case 1n+p:
222	      (c, +v) = r
223	      +nx = ((code + v) << 1n : U32)
224	      next.go(p, (l + 1 : U32), nx, ((left - v) * 2 : U32), Bool.and(ok, U32.is_le(v, left)), Array.set(U32, n, l, nx), Array.get(U32, c, l))
225	
226	def fill(f: Nat, t: Array<U32>, +i: U32, +step: U32, +e: U32) -> Array<U32>:
227	  match f:
228	    case 0n:
229	      t
230	    case 1n+p:
231	      fill(p, Array.set(U32, t, i, e), (i + step : U32), step, e)
232	
233	def rev(k: Nat, +v: U32, +acc: U32) -> U32:
234	  match k:
235	    case 0n:
236	      acc
237	    case 1n+p:
238	      rev(p, U32.shrn(v, 1n), ((acc << 1n) .|. (v .&. 1) : U32))
239	
240	# MSB-first, a code fills a run of consecutive entries; LSB-first, its bits reversed fill every 2^len-th.
241	def assign.put(+msb: Bool, +max: U32, +sym: U32, +len: U32, r: Array<U32> & U32, t: Array<U32>) -> Array<U32> & Array<U32>:
242	  (n, +code) = r
243	  +lo = (max - len : U32)
244	  +start = Bool.pick(U32, msb, U32.shln(code, U32.to_nat(lo)), rev(U32.to_nat(len), code, 0))
245	  +step = Bool.pick(U32, msb, 1, U32.shln(1, U32.to_nat(len)))
246	  (Array.set(U32, n, len, (code + 1 : U32)), fill(U32.to_nat(U32.shln(1, U32.to_nat(lo))), t, start, step, ((sym << 5n) .|. len : U32)))
247	
248	def assign.one(zero: Bool, +msb: Bool, +max: U32, +sym: U32, +len: U32, st: Array<U32> & Array<U32>) -> Array<U32> & Array<U32>:
249	  match zero:
250	    case True{}:
251	      st
252	    case False{}:
253	      (n, t) = st
254	      assign.put(msb, max, sym, len, Array.get(U32, n, len), t)
255	
256	def assign(xs: List<&2, Code>, +msb: Bool, +max: U32, st: Array<U32> & Array<U32>) -> Array<U32>:
257	  match xs:
258	    case Nil{}:
259	      (n, t) = st
260	      t
261	    case Con{Code{+sym, +len}, t}:
262	      assign(t, msb, max, assign.one(U32.is_eq(len, 0), msb, max, sym, len, st))
263	
264	def build.next(+msb: Bool, +max: U32, xs: List<&2, Code>, r: Array<U32> & Bool) -> Maybe<&1, Huff>:
265	  (n, ok) = r
266	  match ok:
267	    case True{}:
268	      Some{Huff{msb, max, assign(xs, msb, max, (n, Array.new(U32, U32.to_nat(max), 0)))}}
269	    case False{}:
270	      None{}
271	
272	def build.short(ok: Bool, +msb: Bool, +max: U32, +xs: List<&2, Code>) -> Maybe<&1, Huff>:
273	  match ok:
274	    case True{}:
275	      build.next(msb, max, xs, next.go(17n, 1, 0, 1, True{}, Array.new(U32, 5n, 0), Array.get(U32, count(xs, Array.new(U32, 5n, 0)), 0)))
276	    case False{}:
277	      None{}
278	
279	# A table from codes. Codes go shortest first, ties in list order; length 0 means no code.
280	# None when a length exceeds 16 or the lengths over-subscribe. An incomplete set is allowed,
281	# and decode answers None for its unused codes. Symbols must be below 2^27.
282	def Huff.from_codes(+msb: Bool, +xs: List<&2, Code>) -> Maybe<&1, Huff>:
283	  +max = lens.max(xs, 0)
284	  build.short(U32.is_le(max, 16), msb, max, xs)
285	
286	def codes(xs: List<&2, U32>, +i: U32) -> List<&2, Code>:
287	  match xs:
288	    case Nil{}:
289	      Nil{}
290	    case Con{+len, t}:
291	      Con{Code{i, len}, codes(t, (i + 1 : U32))}
292	
293	# Symbol i has length lens[i], as DEFLATE sends them (RFC 1951 §3.2.2).
294	def Huff.from_lens(+msb: Bool, +lens: List<&2, U32>) -> Maybe<&1, Huff>:
295	  Huff.from_codes(msb, codes(lens, 0))
296	
297	# Bits past the end may be garbage in w: a code's entry repeats over every suffix, so only
298	# its own len bits pick it, and len <= left keeps them inside the stream.
299	def decode.hit(+msb: Bool, +max: U32, +len: U32, +pos: U32, +left: U32, buf: Array<U32>, r: Array<U32> & U32) -> Huff & Reader & Maybe<&2, U32>:
300	  (t, +e) = r
301	  +l = (e .&. 31 : U32)
302	  +ok = Bool.and(U32.is_ne(l, 0), U32.is_le(l, left))
303	  (Huff{msb, max, t}, Reader{len, buf, Bool.pick(U32, ok, (pos + l : U32), pos)}, Bool.pick(Maybe<&2, U32>, ok, Some{U32.shrn(e, 5n)}, None{}))
304	
305	def decode.look(+msb: Bool, +max: U32, +len: U32, +pos: U32, t: Array<U32>, r: Array<U32> & U32) -> Huff & Reader & Maybe<&2, U32>:
306	  (buf, +w) = r
307	  decode.hit(msb, max, len, pos, ((len << 3n) - pos : U32), buf, Array.get(U32, t, field(msb, max, w)))
308	
309	# One symbol, and the reader past its code. None, with the reader unchanged, when the next
310	# bits start no code or the input ends inside one.
311	def Huff.decode(h: Huff, r: Reader) -> Huff & Reader & Maybe<&2, U32>:
312	  Huff{+msb, +max, t} = h
313	  Reader{+len, buf, +pos} = r
314	  decode.look(msb, max, len, pos, t, win(msb, buf, pos))