bits.bend

bend-kit-bits@0.1.0.1 · 314 lines · 55 definitions

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
2import Base
3import 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
10def mask(+n: U32) -> U32:
11  (U32.shln(1, U32.to_nat(n)) - 1 : U32)
12
13def 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
17def 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
21def 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
25def 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.
31def 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.
39def 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.
43type Reader is Type:
44  Reader{len: U32, buf: Array<U32>, pos: U32}
45
46def 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.
51def Reader.finish(r: Reader) -> Bytes.Bytes & U32:
52  Reader{len, buf, pos} = r
53  (Bytes.Bytes{len, buf}, pos)
54
55def Reader.left(r: Reader) -> Reader & U32:
56  Reader{+len, buf, +pos} = r
57  (Reader{len, buf, pos}, ((len << 3n) - pos : U32))
58
59def 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
63def 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
70def 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.
76def Reader.lsb(r: Reader, +n: U32) -> Reader & Maybe<&2, U32>:
77  read(False{}, r, n)
78
79def 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.
83def 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
88def 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.
94def Reader.peek.lsb(r: Reader, +n: U32) -> Reader & U32:
95  peek(False{}, r, n)
96
97def 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.
101def 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.
107def 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.
112type Writer is Type:
113  Writer{buf: Array<U32>, dep: U32, pos: U32}
114
115def 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.
119def 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
127def 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
131def 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.
139def 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
148def 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
152def 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
156def 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.
163def Writer.lsb(w: Writer, +n: U32, +v: U32) -> Writer:
164  write(False{}, w, n, v)
165
166def 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.
170def 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.
175def 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.
182type Huff is Type:
183  Huff{msb: Bool, max: U32, table: Array<U32>}
184
185# One symbol and the length of its code.
186type Code is Data:
187  Code{sym: U32, len: U32}
188
189def 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
196def count.of(+len: U32, r: Array<U32> & U32) -> Array<U32>:
197  (c, +v) = r
198  Array.set(U32, c, len, (v + 1 : U32))
199
200def 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
207def 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.
216def 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
226def 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
233def 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.
241def 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
248def 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
256def 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
264def 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
272def 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.
282def 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
286def 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).
294def 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.
299def 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
305def 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.
311def 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))