jpeg.bend

gralha-jpeg-decode@0.1.0.1 · 749 lines · 127 definitions

1# JPEG decoder: baseline and progressive Huffman, subsampling, color transforms.
2import Base
3import gralha-text-scan@0.1.0.1/scan.bend as U
4import gralha-bytes@0.1.0.1/bytes.bend as Bin
5import gralha-image@0.1.0.1/image.bend as Img
6
7# JPEG entropy bytes are MSB-first and escape FF as FF 00.
8type JPEGSpectral is Data:
9  JPEGSpectral{start: U32, end: U32, high: U32, low: U32, eob: U32}
10
11def jpeg.spectral.default() -> JPEGSpectral:
12  JPEGSpectral{0, 63, 0, 0, 0}
13
14def jpeg.spectral.eob(scan: JPEGSpectral, eob: U32) -> JPEGSpectral:
15  JPEGSpectral{start, end, high, low, previous} = scan
16  JPEGSpectral{start, end, high, low, eob}
17
18type JPEGStream is Data:
19  JPEGStream{data: Bin.Bytes, position: U32, word: U32, count: U32, scan: JPEGSpectral, valid: Bool}
20
21type JPEGRead is Data:
22  JPEGRead{value: U32, stream: JPEGStream}
23
24def jpeg.stream.valid(stream: JPEGStream) -> Bool:
25  JPEGStream{data, position, word, count, scan, valid} = stream
26  valid
27
28def jpeg.stream.invalid(stream: JPEGStream) -> JPEGStream:
29  JPEGStream{data, position, word, count, scan, valid} = stream
30  JPEGStream{data, position, word, count, scan, False{}}
31
32def jpeg.stream.position(stream: JPEGStream) -> U32:
33  JPEGStream{data, position, word, count, scan, valid} = stream
34  position
35
36def jpeg.stream.fill(+stream: JPEGStream) -> JPEGStream:
37  JPEGStream{data, position, word, count, scan, valid} = stream
38  +v = Bin.bytes.get(data, position)
39  +escaped = U32.is_eq(v, 255)
40  +next = (position + U.choose(U32, escaped, 2, 1) : U32)
41  U.branch(JPEGStream, U32.is_gt(count, 0), _ => stream,
42    _ => JPEGStream{data, next, v, 8, scan, valid && U32.is_le(next, Bin.bytes.size(data)) && (Bool.not(escaped) || U32.is_eq(Bin.bytes.get(data, U32.inc(position)), 0))})
43
44def jpeg.stream.take(+stream: JPEGStream) -> JPEGRead:
45  JPEGStream{data, position, word, count, scan, valid} = stream
46  U.branch(JPEGRead, valid && U32.is_gt(count, 0),
47    _ => JPEGRead{U32.and(U32.shrn(word, U32.to_nat((count - 1 : U32))), 1), JPEGStream{data, position, word, (count - 1 : U32), scan, valid}},
48    _ => JPEGRead{0, jpeg.stream.invalid(stream)})
49
50def jpeg.bits.next(read: JPEGRead, acc: U32, next: JPEGStream -> U32 -> JPEGRead) -> JPEGRead:
51  JPEGRead{bit, stream} = read
52  next(stream)((acc * 2 + bit : U32))
53
54@unsafe
55def jpeg.bits.loop(n: Nat, stream: JPEGStream, acc: U32) -> JPEGRead:
56  match n:
57    case 0n: JPEGRead{acc, stream}
58    case 1n+d: jpeg.bits.next(jpeg.stream.take(jpeg.stream.fill(stream)), acc, input => decoded => jpeg.bits.loop(d, input, decoded))
59
60def jpeg.bits(stream: JPEGStream, count: U32) -> JPEGRead:
61  jpeg.bits.loop(U32.to_nat(count), stream, 0)
62
63def jpeg.huff.next(read: JPEGRead, zero: Bin.Huffman, one: Bin.Huffman, next: Bin.Huffman -> JPEGStream -> JPEGRead) -> JPEGRead:
64  JPEGRead{bit, stream} = read
65  next(U.choose(Bin.Huffman, U32.is_eq(bit, 0), zero, one))(stream)
66
67@unsafe
68def jpeg.huff.read(node: Bin.Huffman, stream: JPEGStream) -> JPEGRead:
69  match node:
70    case Bin.HuffEmpty{}: JPEGRead{0, jpeg.stream.invalid(stream)}
71    case Bin.HuffBad{}: JPEGRead{0, jpeg.stream.invalid(stream)}
72    case Bin.HuffLeaf{symbol}: JPEGRead{symbol, stream}
73    case Bin.HuffNode{zero, one}: jpeg.huff.next(jpeg.bits(stream, 1), zero, one, node => input => jpeg.huff.read(node, input))
74
75# Quantized coefficients use two's-complement U32 storage, with explicit
76# signed conversion before floating-point reconstruction.
77def jpeg.signed(+value: U32) -> F32:
78  U.branch(F32, U32.is_ge(value, 2147483648), _ => F32.neg(U32.to_f32((0 - value : U32))), _ => U32.to_f32(value))
79
80def jpeg.extend(+value: U32, +count: U32) -> U32:
81  +base = U32.shln(1, U32.to_nat(count))
82  U.choose(U32, U32.is_gt(count, 0) && U32.is_lt(value, (base / 2 : U32)), (value + 1 - base : U32), value)
83
84type JPEGMarker is Data:
85  JPEGMarker{code: U32, position: U32, valid: Bool}
86
87@unsafe
88def jpeg.marker.byte(+src: Bin.Bytes, +at: U32) -> JPEGMarker:
89  +v = Bin.bytes.get(src, at)
90  U.branch(JPEGMarker, U32.is_ge(at, Bin.bytes.size(src)), _ => JPEGMarker{0, at, False{}},
91    _ => U.branch(JPEGMarker, U32.is_eq(v, 255), _ => jpeg.marker.byte(src, U32.inc(at)), _ => JPEGMarker{v, U32.inc(at), Bool.not(U32.is_eq(v, 0))}))
92
93def jpeg.marker(+src: Bin.Bytes, +at: U32) -> JPEGMarker:
94  U.branch(JPEGMarker, U32.is_lt(at, Bin.bytes.size(src)) && U32.is_eq(Bin.bytes.get(src, at), 255), _ => jpeg.marker.byte(src, U32.inc(at)), _ => JPEGMarker{0, at, False{}})
95
96def jpeg.zigzag(index: U32) -> U32:
97  U.uints.nth(Con{0, Con{1, Con{8, Con{16, Con{9, Con{2, Con{3, Con{10, Con{17, Con{24, Con{32, Con{25, Con{18, Con{11, Con{4, Con{5, Con{12, Con{19, Con{26, Con{33, Con{40, Con{48, Con{41, Con{34, Con{27, Con{20, Con{13, Con{6, Con{7, Con{14, Con{21, Con{28, Con{35, Con{42, Con{49, Con{56, Con{57, Con{50, Con{43, Con{36, Con{29, Con{22, Con{15, Con{23, Con{30, Con{37, Con{44, Con{51, Con{58, Con{59, Con{52, Con{45, Con{38, Con{31, Con{39, Con{46, Con{53, Con{60, Con{61, Con{54, Con{47, Con{55, Con{62, Con{63, Nil{}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}, U32.to_nat(index), 0)
98
99def jpeg.stream.eob(stream: JPEGStream, eob: U32) -> JPEGStream:
100  JPEGStream{data, position, word, count, scan, valid} = stream
101  JPEGStream{data, position, word, count, jpeg.spectral.eob(scan, eob), valid}
102
103type JPEGQuant is Data:
104  JPEGQuant{id: U32, values: Bin.Bytes}
105
106type JPEGHuff is Data:
107  JPEGHuff{id: U32, ac: Bool, node: Bin.Huffman}
108
109type JPEGComponent is Data:
110  JPEGComponent{id: U32, h: U32, v: U32, quant: U32, dc: U32, ac: U32, predictor: U32, bw: U32, bh: U32, coefficients: Bin.Bytes, samples: Bin.Bytes}
111
112type JPEGFrame is Data:
113  JPEGFrame{w: U32, h: U32, hm: U32, vm: U32, progressive: Bool, components: List<&2, JPEGComponent>, quants: List<&2, JPEGQuant>, huffs: List<&2, JPEGHuff>, restart: U32, adobe: U32, valid: Bool}
114
115def jpeg.frame.empty() -> JPEGFrame:
116  JPEGFrame{0, 0, 1, 1, False{}, Nil{}, Nil{}, Nil{}, 0, 255, True{}}
117
118def jpeg.frame.invalid(frame: JPEGFrame) -> JPEGFrame:
119  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
120  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, False{}}
121
122def jpeg.frame.valid(frame: JPEGFrame) -> Bool:
123  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
124  valid
125
126def jpeg.quant.find(xs: List<&2, JPEGQuant>, +id: U32) -> Bin.Bytes:
127  match xs:
128    case Nil{}: Bin.bytes.empty()
129    case Con{JPEGQuant{key, values}, rest}: U.branch(Bin.Bytes, U32.is_eq(key, id), _ => values, _ => jpeg.quant.find(rest, id))
130
131def jpeg.huff.find(xs: List<&2, JPEGHuff>, +id: U32, +ac: Bool) -> Bin.Huffman:
132  match xs:
133    case Nil{}: Bin.HuffEmpty{}
134    case Con{JPEGHuff{key, kind, node}, rest}: U.branch(Bin.Huffman, U32.is_eq(key, id) && U.choose(Bool, kind, ac, Bool.not(ac)), _ => node, _ => jpeg.huff.find(rest, id, ac))
135
136def jpeg.quant.read(n: Nat, +src: Bin.Bytes, +at: U32, +stride: U32, +index: U32, out: Bin.Bytes) -> Bin.Bytes:
137  match n:
138    case 0n: out
139    case 1n+d: jpeg.quant.read(d, src, (at + stride : U32), stride, U32.inc(index), Bin.bytes.set(out, jpeg.zigzag(index), Bin.bytes.be(src, at, U32.to_nat(stride), 0)))
140
141@unsafe
142def jpeg.quant.segment(+src: Bin.Bytes, +at: U32, +end: U32, +frame: JPEGFrame) -> JPEGFrame:
143  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
144  +pq = Bin.bytes.get(src, at)
145  +stride = U32.inc((pq / 16 : U32))
146  +next = (at + 1 + 64 * stride : U32)
147  U.branch(JPEGFrame, U32.is_eq(at, end), _ => frame,
148    _ => U.branch(JPEGFrame, U32.is_gt(at, end) || U32.is_gt(pq, 19) || U32.is_gt((pq % 16 : U32), 3) || U32.is_gt(next, end), _ => jpeg.frame.invalid(frame),
149      _ => jpeg.quant.segment(src, next, end, JPEGFrame{w, h, hm, vm, progressive, components, Con{JPEGQuant{(pq % 16 : U32), jpeg.quant.read(64n, src, U32.inc(at), stride, 0, Bin.bytes.zero(64))}, quants}, huffs, restart, adobe, valid})))
150
151type JPEGCoding is Data:
152  JPEGCoding{node: Bin.Huffman, code: U32, at: U32, valid: Bool}
153
154def jpeg.huff.symbols(n: Nat, +src: Bin.Bytes, +bits: U32, +coding: JPEGCoding) -> JPEGCoding:
155  match n:
156    case 0n: coding
157    case 1n+d:
158      JPEGCoding{node, cw0, at, valid} = coding
159      +cw = cw0
160      jpeg.huff.symbols(d, src, bits, JPEGCoding{Bin.huff.insert(node, U32.to_nat(bits), cw, Bin.bytes.get(src, at)), U32.inc(cw), U32.inc(at), valid && U32.is_lt(cw, (U32.shln(1, U32.to_nat(bits)) - 1 : U32))})
161
162def jpeg.huff.shift(coding: JPEGCoding) -> JPEGCoding:
163  JPEGCoding{node, cw0, at, valid} = coding
164  JPEGCoding{node, (cw0 * 2 : U32), at, valid}
165
166def jpeg.huff.lengths(n: Nat, +src: Bin.Bytes, +counts: U32, +bits: U32, +coding: JPEGCoding) -> JPEGCoding:
167  match n:
168    case 0n: coding
169    case 1n+d: jpeg.huff.lengths(d, src, U32.inc(counts), U32.inc(bits), jpeg.huff.shift(jpeg.huff.symbols(U32.to_nat(Bin.bytes.get(src, counts)), src, bits, coding)))
170
171def jpeg.huff.install(coding: JPEGCoding, +frame: JPEGFrame, spec: U32, end: U32, next: U32 -> JPEGFrame -> JPEGFrame) -> JPEGFrame:
172  JPEGCoding{node, cw0, at, ok} = coding
173  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
174  +position = at
175  +tag = spec
176  next(position)(JPEGFrame{w, h, hm, vm, progressive, components, quants, Con{JPEGHuff{(tag % 16 : U32), U32.is_ge(tag, 16), node}, huffs}, restart, adobe, valid && ok && U32.is_le(position, end)})
177
178@unsafe
179def jpeg.huff.segment(+src: Bin.Bytes, +at: U32, +end: U32, +frame: JPEGFrame) -> JPEGFrame:
180  +spec = Bin.bytes.get(src, at)
181  U.branch(JPEGFrame, U32.is_eq(at, end), _ => frame,
182    _ => U.branch(JPEGFrame, Bool.not(jpeg.frame.valid(frame)) || U32.is_gt((at + 17 : U32), end) || U32.is_gt(spec, 19) || U32.is_gt((spec % 16 : U32), 3), _ => jpeg.frame.invalid(frame),
183      _ => jpeg.huff.install(jpeg.huff.lengths(16n, src, U32.inc(at), 1, JPEGCoding{Bin.HuffEmpty{}, 0, (at + 17 : U32), True{}}), frame, spec, end, pos => updated => jpeg.huff.segment(src, pos, end, updated))))
184
185def jpeg.ceil(+n: U32, +d: U32) -> U32:
186  ((n + d - 1) / d : U32)
187
188def jpeg.components.read(n: Nat, +src: Bin.Bytes, +at: U32) -> List<&2, JPEGComponent>:
189  match n:
190    case 0n: Nil{}
191    case 1n+d:
192      +hv = Bin.bytes.get(src, U32.inc(at))
193      Con{JPEGComponent{Bin.bytes.get(src, at), (hv / 16 : U32), (hv % 16 : U32), Bin.bytes.get(src, (at + 2 : U32)), 0, 0, 0, 0, 0, Bin.bytes.empty(), Bin.bytes.empty()}, jpeg.components.read(d, src, (at + 3 : U32))}
194
195def jpeg.components.max(xs: List<&2, JPEGComponent>, +vertical: Bool, +acc: U32) -> U32:
196  match xs:
197    case Nil{}: acc
198    case Con{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}, rest}:
199      +val = U.choose(U32, vertical, v, h)
200      jpeg.components.max(rest, vertical, U.choose(U32, U32.is_gt(val, acc), val, acc))
201
202def jpeg.components.has(xs: List<&2, JPEGComponent>, +key: U32) -> Bool:
203  match xs:
204    case Nil{}: False{}
205    case Con{JPEGComponent{id, _, _, _, _, _, _, _, _, _, _}, rest}: U32.is_eq(id, key) || jpeg.components.has(rest, key)
206
207def jpeg.components.units(xs: List<&2, JPEGComponent>) -> U32:
208  match xs:
209    case Nil{}: 0
210    case Con{JPEGComponent{_, h, v, _, _, _, _, _, _, _, _}, rest}: (h * v + jpeg.components.units(rest) : U32)
211
212def jpeg.components.valid(+xs: List<&2, JPEGComponent>) -> Bool:
213  match xs:
214    case Nil{}: True{}
215    case Con{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}, rest}:
216      +hh = h
217      +vv = v
218      U32.is_ge(hh, 1) && U32.is_le(hh, 4) && U32.is_ge(vv, 1) && U32.is_le(vv, 4) && U32.is_le(quant, 3) && Bool.not(jpeg.components.has(rest, id)) && jpeg.components.valid(rest)
219
220def jpeg.components.allocate(xs: List<&2, JPEGComponent>, +cols: U32, +rows: U32) -> List<&2, JPEGComponent>:
221  match xs:
222    case Nil{}: Nil{}
223    case Con{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}, rest}:
224      +hh = h
225      +vv = v
226      Con{JPEGComponent{id, hh, vv, quant, dc, ac, predictor, (cols * hh : U32), (rows * vv : U32), Bin.bytes.zero((cols * hh * rows * vv * 64 : U32)), Bin.bytes.zero(64)}, jpeg.components.allocate(rest, cols, rows)}
227
228def jpeg.frame.start(+src: Bin.Bytes, +at: U32, +end: U32, +frame: JPEGFrame, +kind: U32) -> JPEGFrame:
229  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
230  +ww = Bin.bytes.be(src, (at + 3 : U32), 2n, 0)
231  +hh = Bin.bytes.be(src, U32.inc(at), 2n, 0)
232  +count = Bin.bytes.get(src, (at + 5 : U32))
233  +cs = jpeg.components.read(U32.to_nat(U.choose(U32, U32.is_le(count, 4), count, 0)), src, (at + 6 : U32))
234  +mh = jpeg.components.max(cs, False{}, 1)
235  +mv = jpeg.components.max(cs, True{}, 1)
236  +ok = valid && U32.is_eq(w, 0) && U32.is_eq(Bin.bytes.get(src, at), 8) && U32.is_gt(ww, 0) && U32.is_le(ww, 2048) && U32.is_gt(hh, 0) && U32.is_le(hh, 2048) && (U32.is_eq(count, 1) || U32.is_eq(count, 3) || U32.is_eq(count, 4)) && U32.is_eq((at + 6 + count * 3 : U32), end) && jpeg.components.valid(cs) && U32.is_le(jpeg.components.units(cs), 10)
237  U.branch(JPEGFrame, ok, _ => JPEGFrame{ww, hh, mh, mv, U32.is_eq(kind, 194), jpeg.components.allocate(cs, jpeg.ceil(ww, (8 * mh : U32)), jpeg.ceil(hh, (8 * mv : U32))), quants, huffs, restart, adobe, ok}, _ => jpeg.frame.invalid(frame))
238
239def jpeg.quant.valid(n: Nat, +index: U32, +quant: Bin.Bytes) -> Bool:
240  match n:
241    case 0n: U32.is_eq(Bin.bytes.size(quant), 64)
242    case 1n+d: U32.is_gt(Bin.bytes.get(quant, index), 0) && jpeg.quant.valid(d, U32.inc(index), quant)
243
244type JPEGBlock is Data:
245  JPEGBlock{stream: JPEGStream, coefficients: Bin.Bytes, predictor: U32}
246
247def jpeg.block.invalid(block: JPEGBlock) -> JPEGBlock:
248  JPEGBlock{stream, coefficients, predictor} = block
249  JPEGBlock{jpeg.stream.invalid(stream), coefficients, predictor}
250
251def jpeg.progressive.dc.bits(read: JPEGRead, count: U32, coefficients: Bin.Bytes, predictor: U32, offset: U32, low: U32) -> JPEGBlock:
252  JPEGRead{word, stream} = read
253  +dc = (predictor + jpeg.extend(word, count) : U32)
254  JPEGBlock{stream, Bin.bytes.set(coefficients, offset, U32.shln(dc, U32.to_nat(low))), dc}
255
256def jpeg.progressive.dc.first(+read: JPEGRead, coefficients: Bin.Bytes, predictor: U32, offset: U32, low: U32) -> JPEGBlock:
257  JPEGRead{count, stream} = read
258  +n = count
259  +ok = U32.is_le(n, 11)
260  jpeg.progressive.dc.bits(jpeg.bits(U.choose(JPEGStream, ok, stream, jpeg.stream.invalid(stream)), U.choose(U32, ok, n, 0)), n, coefficients, predictor, offset, low)
261
262def jpeg.progressive.dc.refine(read: JPEGRead, +coefficients: Bin.Bytes, predictor: U32, +offset: U32, low: U32) -> JPEGBlock:
263  JPEGRead{bit, stream} = read
264  JPEGBlock{stream, Bin.bytes.set(coefficients, offset, U32.or(Bin.bytes.get(coefficients, offset), U32.shln(bit, U32.to_nat(low)))), predictor}
265
266def jpeg.progressive.ac.bits(read: JPEGRead, count: U32, coefficients: Bin.Bytes, predictor: U32, offset: U32, index: U32, low: U32) -> JPEGBlock:
267  JPEGRead{word, stream} = read
268  JPEGBlock{stream, Bin.bytes.set(coefficients, (offset + jpeg.zigzag(index) : U32), U32.shln(jpeg.extend(word, count), U32.to_nat(low))), predictor}
269
270def jpeg.progressive.eob(read: JPEGRead, run: U32, coefficients: Bin.Bytes, predictor: U32, decrement: U32) -> JPEGBlock:
271  JPEGRead{word, stream} = read
272  JPEGBlock{jpeg.stream.eob(stream, (U32.shln(1, U32.to_nat(run)) + word - decrement : U32)), coefficients, predictor}
273
274type JPEGRefined is Data:
275  JPEGRefined{index: U32, block: JPEGBlock}
276
277def jpeg.progressive.first.symbol(read: JPEGRead, +index: U32, +end: U32, +offset: U32, +low: U32, +coefficients: Bin.Bytes, +predictor: U32) -> JPEGRefined:
278  JPEGRead{symbol, stream} = read
279  +rs = symbol
280  +input = stream
281  +run = (rs / 16 : U32)
282  +count = (rs % 16 : U32)
283  +eob = U32.is_eq(count, 0) && U32.is_lt(run, 15)
284  +zero = U32.is_eq(rs, 240)
285  +pos = (index + run : U32)
286  +ok = jpeg.stream.valid(input) && (eob || (zero && U32.is_le(pos, end)) || (U32.is_gt(count, 0) && U32.is_le(count, 10) && U32.is_le(pos, end)))
287  JPEGRefined{U.choose(U32, eob || Bool.not(ok), U32.inc(end), U32.inc(pos)),
288    U.branch(JPEGBlock, Bool.not(ok), _ => JPEGBlock{jpeg.stream.invalid(input), coefficients, predictor},
289      _ => U.branch(JPEGBlock, eob, _ => jpeg.progressive.eob(jpeg.bits(input, run), run, coefficients, predictor, 1),
290        _ => U.branch(JPEGBlock, zero, _ => JPEGBlock{input, coefficients, predictor}, _ => jpeg.progressive.ac.bits(jpeg.bits(input, count), count, coefficients, predictor, offset, pos, low))))}
291
292def jpeg.refined.next(refined: JPEGRefined, next: U32 -> JPEGBlock -> JPEGBlock) -> JPEGBlock:
293  JPEGRefined{index, block} = refined
294  next(index)(block)
295
296@unsafe
297def jpeg.progressive.first(+index: U32, +end: U32, +offset: U32, +low: U32, +table: Bin.Huffman, +block: JPEGBlock) -> JPEGBlock:
298  JPEGBlock{stream, coefficients, predictor} = block
299  U.branch(JPEGBlock, U32.is_gt(index, end) || Bool.not(jpeg.stream.valid(stream)), _ => block,
300    _ => jpeg.refined.next(jpeg.progressive.first.symbol(jpeg.huff.read(table, stream), index, end, offset, low, coefficients, predictor), pos => decoded => jpeg.progressive.first(pos, end, offset, low, table, decoded)))
301
302def jpeg.progressive.refine.bit(read: JPEGRead, +coefficients: Bin.Bytes, predictor: U32, +at: U32, +p1: U32) -> JPEGBlock:
303  JPEGRead{bit, stream} = read
304  +old = Bin.bytes.get(coefficients, at)
305  +negative = U32.is_ge(old, 2147483648)
306  +magnitude = U.choose(U32, negative, (0 - old : U32), old)
307  +adjust = U32.is_eq(bit, 1) && U32.is_eq(U32.and(magnitude, p1), 0)
308  JPEGBlock{stream, Bin.bytes.set(coefficients, at, U.choose(U32, adjust, (old + U.choose(U32, negative, (0 - p1 : U32), p1) : U32), old)), predictor}
309
310# Walk over zero positions while refining the previously nonzero coefficients.
311# A new coefficient is inserted only after the requested zero run is consumed.
312@unsafe
313def jpeg.progressive.walk(+index: U32, +end: U32, +run: U32, +fresh: U32, +offset: U32, +p1: U32, +block: JPEGBlock) -> JPEGRefined:
314  JPEGBlock{stream, coefficients, predictor} = block
315  +at = (offset + jpeg.zigzag(index) : U32)
316  +nonzero = Bool.not(U32.is_eq(Bin.bytes.get(coefficients, at), 0))
317  +done = U32.is_gt(index, end) || (U32.is_eq(run, 0) && U32.is_eq(fresh, 0)) || Bool.not(jpeg.stream.valid(stream))
318  U.branch(JPEGRefined, done, _ => JPEGRefined{index, U.choose(JPEGBlock, U32.is_gt(index, end) && Bool.not(U32.is_eq(fresh, 0)), jpeg.block.invalid(block), block)},
319    _ => U.branch(JPEGRefined, nonzero,
320      _ => jpeg.progressive.walk(U32.inc(index), end, run, fresh, offset, p1, jpeg.progressive.refine.bit(jpeg.bits(stream, 1), coefficients, predictor, at, p1)),
321      _ => U.branch(JPEGRefined, U32.is_gt(run, 0), _ => jpeg.progressive.walk(U32.inc(index), end, (run - 1 : U32), fresh, offset, p1, block),
322        _ => JPEGRefined{U32.inc(index), JPEGBlock{stream, Bin.bytes.set(coefficients, at, fresh), predictor}})))
323
324def jpeg.progressive.refine.fresh(read: JPEGRead, index: U32, end: U32, run: U32, offset: U32, +p1: U32, coefficients: Bin.Bytes, predictor: U32) -> JPEGRefined:
325  JPEGRead{bit, stream} = read
326  jpeg.progressive.walk(index, end, run, U.choose(U32, U32.is_eq(bit, 1), p1, (0 - p1 : U32)), offset, p1, JPEGBlock{stream, coefficients, predictor})
327
328def jpeg.progressive.refine.eob(block: JPEGBlock, index: U32, end: U32, offset: U32, p1: U32) -> JPEGRefined:
329  jpeg.progressive.walk(index, end, 64, 0, offset, p1, block)
330
331def jpeg.progressive.refine.symbol(read: JPEGRead, +index: U32, +end: U32, +offset: U32, +p1: U32, +coefficients: Bin.Bytes, +predictor: U32) -> JPEGRefined:
332  JPEGRead{symbol, stream} = read
333  +rs = symbol
334  +input = stream
335  +run = (rs / 16 : U32)
336  +count = (rs % 16 : U32)
337  +ok = jpeg.stream.valid(input) && U32.is_le(count, 1)
338  U.branch(JPEGRefined, Bool.not(ok), _ => JPEGRefined{U32.inc(end), JPEGBlock{jpeg.stream.invalid(input), coefficients, predictor}},
339    _ => U.branch(JPEGRefined, U32.is_eq(count, 1), _ => jpeg.progressive.refine.fresh(jpeg.bits(input, 1), index, end, run, offset, p1, coefficients, predictor),
340      _ => U.branch(JPEGRefined, U32.is_eq(run, 15), _ => jpeg.progressive.walk(index, end, 16, 0, offset, p1, JPEGBlock{input, coefficients, predictor}),
341        _ => jpeg.progressive.refine.eob(jpeg.progressive.eob(jpeg.bits(input, run), run, coefficients, predictor, 1), index, end, offset, p1))))
342
343@unsafe
344def jpeg.progressive.refine(+index: U32, +end: U32, +offset: U32, +p1: U32, +table: Bin.Huffman, +block: JPEGBlock) -> JPEGBlock:
345  JPEGBlock{stream, coefficients, predictor} = block
346  U.branch(JPEGBlock, U32.is_gt(index, end) || Bool.not(jpeg.stream.valid(stream)), _ => block,
347    _ => jpeg.refined.next(jpeg.progressive.refine.symbol(jpeg.huff.read(table, stream), index, end, offset, p1, coefficients, predictor), pos => decoded => jpeg.progressive.refine(pos, end, offset, p1, table, decoded)))
348
349def jpeg.refined.block(refined: JPEGRefined) -> JPEGBlock:
350  JPEGRefined{index, block} = refined
351  block
352
353def jpeg.progressive.decode(+component: JPEGComponent, +tables: List<&2, JPEGHuff>, +offset: U32, +stream: JPEGStream) -> JPEGBlock:
354  JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples} = component
355  JPEGStream{data, position, word, count, JPEGSpectral{start, end, high, low, eob}, valid} = stream
356  +p1 = U32.shln(1, U32.to_nat(low))
357  U.branch(JPEGBlock, U32.is_eq(start, 0),
358    _ => U.branch(JPEGBlock, U32.is_eq(high, 0), _ => jpeg.progressive.dc.first(jpeg.huff.read(jpeg.huff.find(tables, dc, False{}), stream), coefficients, predictor, offset, low), _ => jpeg.progressive.dc.refine(jpeg.bits(stream, 1), coefficients, predictor, offset, low)),
359    _ => U.branch(JPEGBlock, U32.is_gt(eob, 0),
360      _ => U.branch(JPEGBlock, U32.is_eq(high, 0), _ => JPEGBlock{jpeg.stream.eob(stream, (eob - 1 : U32)), coefficients, predictor},
361        _ => jpeg.refined.block(jpeg.progressive.walk(start, end, 64, 0, offset, p1, JPEGBlock{jpeg.stream.eob(stream, (eob - 1 : U32)), coefficients, predictor}))),
362      _ => U.branch(JPEGBlock, U32.is_eq(high, 0), _ => jpeg.progressive.first(start, end, offset, low, jpeg.huff.find(tables, ac, True{}), JPEGBlock{stream, coefficients, predictor}),
363        _ => jpeg.progressive.refine(start, end, offset, p1, jpeg.huff.find(tables, ac, True{}), JPEGBlock{stream, coefficients, predictor}))))
364
365def jpeg.block.dc.bits(read: JPEGRead, +count: U32, coefficients: Bin.Bytes, predictor: U32, +offset: U32) -> JPEGBlock:
366  JPEGRead{word, stream} = read
367  +dc = (predictor + jpeg.extend(word, count) : U32)
368  JPEGBlock{stream, Bin.bytes.set(coefficients, offset, dc), dc}
369
370def jpeg.block.dc(+read: JPEGRead, coefficients: Bin.Bytes, predictor: U32, offset: U32) -> JPEGBlock:
371  JPEGRead{count, stream} = read
372  +n = count
373  +ok = U32.is_le(n, 11)
374  jpeg.block.dc.bits(jpeg.bits(U.choose(JPEGStream, ok, stream, jpeg.stream.invalid(stream)), U.choose(U32, ok, n, 0)), n, coefficients, predictor, offset)
375
376def jpeg.block.ac.bits(read: JPEGRead, count: U32, coefficients: Bin.Bytes, predictor: U32, offset: U32, +index: U32, next: U32 -> JPEGBlock -> JPEGBlock) -> JPEGBlock:
377  JPEGRead{word, stream} = read
378  next(U32.inc(index))(JPEGBlock{stream, Bin.bytes.set(coefficients, (offset + jpeg.zigzag(index) : U32), jpeg.extend(word, count)), predictor})
379
380def jpeg.block.ac.symbol(read: JPEGRead, +coefficients: Bin.Bytes, +predictor: U32, +offset: U32, +index: U32, next: U32 -> JPEGBlock -> JPEGBlock) -> JPEGBlock:
381  JPEGRead{symbol, stream} = read
382  +rs = symbol
383  +input = stream
384  +count = (rs % 16 : U32)
385  +pos = (index + rs / 16 : U32)
386  +end = U32.is_eq(rs, 0)
387  +zero = U32.is_eq(rs, 240)
388  +ok = jpeg.stream.valid(input) && (end || (zero && U32.is_le((index + 16 : U32), 64)) || (U32.is_gt(count, 0) && U32.is_le(count, 10) && U32.is_lt(pos, 64)))
389  +skip = end || Bool.not(ok)
390  +result = U.branch(JPEGBlock, skip || zero, _ => JPEGBlock{U.choose(JPEGStream, ok, input, jpeg.stream.invalid(input)), coefficients, predictor},
391    _ => jpeg.block.ac.bits(jpeg.bits(input, count), count, coefficients, predictor, offset, pos, _ => block => block))
392  next(U.choose(U32, skip, 64, U.choose(U32, zero, (index + 16 : U32), U32.inc(pos))))(result)
393
394@unsafe
395def jpeg.block.ac(+index: U32, +table: Bin.Huffman, +offset: U32, +block: JPEGBlock) -> JPEGBlock:
396  JPEGBlock{stream, coefficients, predictor} = block
397  U.branch(JPEGBlock, U32.is_ge(index, 64) || Bool.not(jpeg.stream.valid(stream)), _ => block,
398    _ => jpeg.block.ac.symbol(jpeg.huff.read(table, stream), coefficients, predictor, offset, index, pos => decoded => jpeg.block.ac(pos, table, offset, decoded)))
399
400def jpeg.block.decode(+component: JPEGComponent, +tables: List<&2, JPEGHuff>, offset: U32, +stream: JPEGStream) -> JPEGBlock:
401  JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples} = component
402  JPEGStream{data, position, word, count, JPEGSpectral{start, end, high, low, eob}, valid} = stream
403  +base = offset
404  U.branch(JPEGBlock, U32.is_eq(start, 0) && U32.is_eq(end, 63),
405    _ => jpeg.block.ac(1, jpeg.huff.find(tables, ac, True{}), base, jpeg.block.dc(jpeg.huff.read(jpeg.huff.find(tables, dc, False{}), stream), coefficients, predictor, base)),
406    _ => jpeg.progressive.decode(component, tables, base, stream))
407
408type JPEGPart is Data:
409  JPEGPart{component: JPEGComponent, stream: JPEGStream}
410
411def jpeg.component.block(block: JPEGBlock, component: JPEGComponent) -> JPEGPart:
412  JPEGBlock{stream, coefficients, predictor} = block
413  JPEGComponent{id, h, v, quant, dc, ac, old, bw, bh, previous, samples} = component
414  JPEGPart{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}, stream}
415
416def jpeg.part.next(part: JPEGPart, next: JPEGComponent -> JPEGStream -> JPEGPart) -> JPEGPart:
417  JPEGPart{component, stream} = part
418  next(component)(stream)
419
420@unsafe
421def jpeg.component.blocks(n: Nat, +index: U32, +cols: U32, +x: U32, +y: U32, +tables: List<&2, JPEGHuff>, +component: JPEGComponent, stream: JPEGStream) -> JPEGPart:
422  match n:
423    case 0n: JPEGPart{component, stream}
424    case 1n+d:
425      JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples} = component
426      +offset = (((y + index / cols) * bw + x + index % cols) * 64 : U32)
427      jpeg.part.next(jpeg.component.block(jpeg.block.decode(component, tables, offset, stream), component), item => input => jpeg.component.blocks(d, U32.inc(index), cols, x, y, tables, item, input))
428
429type JPEGScan is Data:
430  JPEGScan{components: List<&2, JPEGComponent>, stream: JPEGStream}
431
432def jpeg.scan.prepend(scan: JPEGScan, component: JPEGComponent) -> JPEGScan:
433  JPEGScan{components, stream} = scan
434  JPEGScan{Con{component, components}, stream}
435
436def jpeg.scan.part(part: JPEGPart, next: JPEGStream -> JPEGScan) -> JPEGScan:
437  JPEGPart{component, stream} = part
438  jpeg.scan.prepend(next(stream), component)
439
440@unsafe
441def jpeg.mcu(+xs: List<&2, JPEGComponent>, +tables: List<&2, JPEGHuff>, +single: Bool, +x: U32, +y: U32, stream: JPEGStream) -> JPEGScan:
442  match xs:
443    case Nil{}: JPEGScan{Nil{}, stream}
444    case Con{component, rest}:
445      JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples} = component
446      +c = {component : JPEGComponent}
447      +hh = U.choose(U32, single, 1, h)
448      +vv = U.choose(U32, single, 1, v)
449      jpeg.scan.part(jpeg.component.blocks(U32.to_nat((hh * vv : U32)), 0, hh, (x * hh : U32), (y * vv : U32), tables, c, stream), input => jpeg.mcu(rest, tables, single, x, y, input))
450
451def jpeg.components.reset(xs: List<&2, JPEGComponent>) -> List<&2, JPEGComponent>:
452  match xs:
453    case Nil{}: Nil{}
454    case Con{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}, rest}:
455      Con{JPEGComponent{id, h, v, quant, dc, ac, 0, bw, bh, coefficients, samples}, jpeg.components.reset(rest)}
456
457def jpeg.restart.marker(marker: JPEGMarker, scan: JPEGScan, expected: U32) -> JPEGScan:
458  JPEGMarker{tag, at, ok} = marker
459  JPEGScan{components, JPEGStream{data, position, word, count, spectral, valid}} = scan
460  JPEGScan{jpeg.components.reset(components), JPEGStream{data, at, 0, 0, jpeg.spectral.eob(spectral, 0), valid && ok && U32.is_eq(tag, expected)}}
461
462def jpeg.restart(+scan: JPEGScan, expected: U32) -> JPEGScan:
463  JPEGScan{components, JPEGStream{data, position, word, count, spectral, valid}} = scan
464  jpeg.restart.marker(jpeg.marker(data, position), scan, expected)
465
466def jpeg.scan.next(scan: JPEGScan, next: List<&2, JPEGComponent> -> JPEGStream -> JPEGScan) -> JPEGScan:
467  JPEGScan{components, stream} = scan
468  next(components)(stream)
469
470def jpeg.scan.active(+scan: JPEGScan, next: List<&2, JPEGComponent> -> JPEGStream -> JPEGScan) -> JPEGScan:
471  JPEGScan{components, stream} = scan
472  U.branch(JPEGScan, jpeg.stream.valid(stream), _ => next(components)(stream), _ => scan)
473
474@unsafe
475def jpeg.mcus(n: Nat, +index: U32, +cols: U32, +single: Bool, +tables: List<&2, JPEGHuff>, +interval: U32, +scan: JPEGScan) -> JPEGScan:
476  match n:
477    case 0n: scan
478    case 1n+d:
479      +reset = U32.is_gt(interval, 0) && U32.is_gt(index, 0) && U32.is_eq((index % interval : U32), 0)
480      +ready = U.branch(JPEGScan, reset, _ => jpeg.restart(scan, (208 + (index / interval - 1) % 8 : U32)), _ => scan)
481      jpeg.scan.active(ready, cs => input => jpeg.scan.next(jpeg.mcu(cs, tables, single, (index % cols : U32), (index / cols : U32), input), parts => decoded => jpeg.mcus(d, U32.inc(index), cols, single, tables, interval, JPEGScan{parts, decoded})))
482
483def jpeg.component.find(+xs: List<&2, JPEGComponent>, +id: U32) -> Maybe<&2, JPEGComponent>:
484  match xs:
485    case Nil{}: None{}
486    case Con{component, rest}:
487      JPEGComponent{key, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples} = component
488      +c = {component : JPEGComponent}
489      U.branch(Maybe<&2, JPEGComponent>, U32.is_eq(key, id), _ => Some{c}, _ => jpeg.component.find(rest, id))
490
491def jpeg.component.tables(found: Maybe<&2, JPEGComponent>, spec: U32) -> JPEGComponent:
492  match found:
493    case None{}: JPEGComponent{255, 1, 1, 0, 4, 4, 0, 0, 0, Bin.bytes.empty(), Bin.bytes.empty()}
494    case Some{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}}:
495      +tag = spec
496      JPEGComponent{id, h, v, quant, (tag / 16 : U32), (tag % 16 : U32), 0, bw, bh, coefficients, samples}
497
498def jpeg.scan.select(n: Nat, +src: Bin.Bytes, +at: U32, +xs: List<&2, JPEGComponent>) -> List<&2, JPEGComponent>:
499  match n:
500    case 0n: Nil{}
501    case 1n+d: Con{jpeg.component.tables(jpeg.component.find(xs, Bin.bytes.get(src, at)), Bin.bytes.get(src, U32.inc(at))), jpeg.scan.select(d, src, (at + 2 : U32), xs)}
502
503def jpeg.component.or(found: Maybe<&2, JPEGComponent>, fallback: JPEGComponent) -> JPEGComponent:
504  match found:
505    case None{}: fallback
506    case Some{component}: component
507
508def jpeg.components.merge(+xs: List<&2, JPEGComponent>, +updated: List<&2, JPEGComponent>) -> List<&2, JPEGComponent>:
509  match xs:
510    case Nil{}: Nil{}
511    case Con{component, rest}:
512      JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples} = component
513      +c = {component : JPEGComponent}
514      Con{jpeg.component.or(jpeg.component.find(updated, id), c), jpeg.components.merge(rest, updated)}
515
516def jpeg.scan.extent(xs: List<&2, JPEGComponent>, w: U32, h: U32, hm: U32, vm: U32, +vertical: Bool) -> U32:
517  match xs:
518    case Nil{}: 0
519    case Con{JPEGComponent{id, hh, vv, quant, dc, ac, predictor, bw, bh, coefficients, samples}, rest}:
520      jpeg.ceil((U.choose(U32, vertical, h, w) * U.choose(U32, vertical, vv, hh) : U32), (U.choose(U32, vertical, vm, hm) * 8 : U32))
521
522type JPEGDecoded is Data:
523  JPEGDecoded{frame: JPEGFrame, at: U32}
524
525def jpeg.scan.complete(stream: JPEGStream) -> Bool:
526  JPEGStream{data, position, word, count, JPEGSpectral{start, end, high, low, eob}, valid} = stream
527  valid && U32.is_eq(eob, 0)
528
529def jpeg.scan.finish(scan: JPEGScan, frame: JPEGFrame) -> JPEGDecoded:
530  JPEGScan{updated, stream} = scan
531  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
532  +input = stream
533  JPEGDecoded{JPEGFrame{w, h, hm, vm, progressive, jpeg.components.merge(components, updated), quants, huffs, restart, adobe, valid && jpeg.scan.complete(input)}, jpeg.stream.position(input)}
534
535# The sample buffer holds approximation history until IDCT replaces it.
536# Zero means unseen; low+1 records the last transmitted coefficient bit.
537def jpeg.history.valid(n: Nat, +index: U32, +history: Bin.Bytes, +high: U32) -> Bool:
538  match n:
539    case 0n: True{}
540    case 1n+d: U32.is_eq(Bin.bytes.get(history, index), U.choose(U32, U32.is_eq(high, 0), 0, U32.inc(high))) && jpeg.history.valid(d, U32.inc(index), history, high)
541
542def jpeg.history.set(n: Nat, +index: U32, +low: U32, history: Bin.Bytes) -> Bin.Bytes:
543  match n:
544    case 0n: history
545    case 1n+d: jpeg.history.set(d, U32.inc(index), low, Bin.bytes.set(history, index, U32.inc(low)))
546
547def jpeg.scan.ready(+xs: List<&2, JPEGComponent>, +start: U32, +end: U32, +high: U32, +quants: List<&2, JPEGQuant>) -> Bool:
548  match xs:
549    case Nil{}: True{}
550    case Con{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}, rest}:
551      U32.is_gt(bw, 0) && (U32.is_eq(start, 0) || U32.is_gt(Bin.bytes.get(samples, 0), 0)) && U32.is_le(dc, 3) && U32.is_le(ac, 3) && Bool.not(jpeg.components.has(rest, id)) && jpeg.quant.valid(64n, 0, jpeg.quant.find(quants, quant)) && jpeg.history.valid(U32.to_nat((end - start + 1 : U32)), start, samples, high) && jpeg.scan.ready(rest, start, end, high, quants)
552
553def jpeg.scan.histories(xs: List<&2, JPEGComponent>, +start: U32, +end: U32, +low: U32) -> List<&2, JPEGComponent>:
554  match xs:
555    case Nil{}: Nil{}
556    case Con{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}, rest}:
557      Con{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, jpeg.history.set(U32.to_nat((end - start + 1 : U32)), start, low, samples)}, jpeg.scan.histories(rest, start, end, low)}
558
559def jpeg.scan.decode(+src: Bin.Bytes, +at: U32, +end: U32, +frame: JPEGFrame) -> JPEGDecoded:
560  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
561  +count = Bin.bytes.get(src, at)
562  +cs = jpeg.scan.select(U32.to_nat(U.choose(U32, U32.is_le(count, 4), count, 0)), src, U32.inc(at), components)
563  +single = U32.is_eq(count, 1)
564  +cols = U.choose(U32, single, jpeg.scan.extent(cs, w, h, hm, vm, False{}), jpeg.ceil(w, (hm * 8 : U32)))
565  +rows = U.choose(U32, single, jpeg.scan.extent(cs, w, h, hm, vm, True{}), jpeg.ceil(h, (vm * 8 : U32)))
566  +ss = Bin.bytes.get(src, (end - 3 : U32))
567  +se = Bin.bytes.get(src, (end - 2 : U32))
568  +approx = Bin.bytes.get(src, (end - 1 : U32))
569  +ah = (approx / 16 : U32)
570  +al = (approx % 16 : U32)
571  +sequential = U32.is_eq(ss, 0) && U32.is_eq(se, 63) && U32.is_eq(approx, 0)
572  +spectral = U32.is_le(ss, se) && U32.is_le(se, 63) && U.choose(Bool, U32.is_eq(ss, 0), U32.is_eq(se, 0), single) && U32.is_le(al, 13) && (U32.is_eq(ah, 0) || U32.is_eq(ah, U32.inc(al)))
573  +ok = valid && U.choose(Bool, progressive, spectral, sequential) && U32.is_gt(w, 0) && U32.is_gt(count, 0) && U32.is_le(count, 4) && U32.is_eq((at + 1 + count * 2 + 3 : U32), end)
574  U.branch(JPEGDecoded, ok, _ => U.branch(JPEGDecoded, jpeg.scan.ready(cs, ss, se, ah, quants), _ => jpeg.scan.finish(jpeg.mcus(U32.to_nat((cols * rows : U32)), 0, cols, single, huffs, restart, JPEGScan{jpeg.scan.histories(cs, ss, se, al), JPEGStream{src, end, 0, 0, JPEGSpectral{ss, se, ah, al, 0}, True{}}}), frame), _ => JPEGDecoded{jpeg.frame.invalid(frame), end}), _ => JPEGDecoded{jpeg.frame.invalid(frame), end})
575
576def jpeg.frame.misc(+frame: JPEGFrame, +src: Bin.Bytes, +at: U32, +end: U32, +tag: U32) -> JPEGFrame:
577  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
578  +dri = U32.is_eq(tag, 221)
579  +app14 = U32.is_eq(tag, 238) && U32.is_ge((end - at : U32), 12) && U32.is_eq(Bin.bytes.be(src, at, 4n, 0), 1097101154) && U32.is_eq(Bin.bytes.get(src, (at + 4 : U32)), 101)
580  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, U.choose(U32, dri, Bin.bytes.be(src, at, 2n, 0), restart), U.choose(U32, app14, Bin.bytes.get(src, (at + 11 : U32)), adobe), valid && (Bool.not(dri) || U32.is_eq(end, (at + 2 : U32)))}
581
582def jpeg.file.segment(+src: Bin.Bytes, +at: U32, +end: U32, +tag: U32, +frame: JPEGFrame) -> JPEGDecoded:
583  U.branch(JPEGDecoded, U32.is_eq(tag, 218), _ => jpeg.scan.decode(src, at, end, frame),
584    _ => JPEGDecoded{U.branch(JPEGFrame, U32.is_eq(tag, 219), _ => jpeg.quant.segment(src, at, end, frame),
585      _ => U.branch(JPEGFrame, U32.is_eq(tag, 196), _ => jpeg.huff.segment(src, at, end, frame),
586        _ => U.branch(JPEGFrame, U32.is_eq(tag, 192) || U32.is_eq(tag, 193) || U32.is_eq(tag, 194), _ => jpeg.frame.start(src, at, end, frame, tag),
587          _ => U.branch(JPEGFrame, (U32.is_ge(tag, 192) && U32.is_le(tag, 207)) || U32.is_eq(tag, 220), _ => jpeg.frame.invalid(frame), _ => jpeg.frame.misc(frame, src, at, end, tag))))), end})
588
589def jpeg.file.next(decoded: JPEGDecoded, next: U32 -> JPEGFrame -> JPEGFrame) -> JPEGFrame:
590  JPEGDecoded{frame, at} = decoded
591  next(at)(frame)
592
593def jpeg.file.continue(+decoded: JPEGDecoded, stop: Bool, seen: Bool, next: U32 -> JPEGFrame -> Bool -> JPEGFrame) -> JPEGFrame:
594  JPEGDecoded{frame, at} = decoded
595  U.branch(JPEGFrame, stop, _ => frame, _ => next(at)(frame)(seen))
596
597def jpeg.file.marker(marker: JPEGMarker, +src: Bin.Bytes, +frame: JPEGFrame, +seen: Bool, next: U32 -> JPEGFrame -> Bool -> JPEGFrame) -> JPEGFrame:
598  JPEGMarker{tag, at, valid} = marker
599  +kind = tag
600  +pos = at
601  +ok = valid
602  +size = Bin.bytes.be(src, pos, 2n, 0)
603  +end = (pos + size : U32)
604  +done = U32.is_eq(kind, 217)
605  +invalid = Bool.not(ok) || Bool.not(jpeg.frame.valid(frame)) || U32.is_lt(kind, 192) || (U32.is_ge(kind, 208) && U32.is_le(kind, 216)) || (Bool.not(done) && (U32.is_lt(size, 2) || U32.is_gt(end, Bin.bytes.size(src))))
606  +stop = done || invalid
607  +result = U.branch(JPEGDecoded, stop, _ => JPEGDecoded{U.choose(JPEGFrame, done && seen && Bool.not(invalid), frame, jpeg.frame.invalid(frame)), pos}, _ => jpeg.file.segment(src, (pos + 2 : U32), end, kind, frame))
608  jpeg.file.continue(result, stop, seen || U32.is_eq(kind, 218), next)
609
610@unsafe
611def jpeg.file(+src: Bin.Bytes, +at: U32, frame: JPEGFrame, seen: Bool) -> JPEGFrame:
612  jpeg.file.marker(jpeg.marker(src, at), src, frame, seen, pos => parsed => scanned => jpeg.file(src, pos, parsed, scanned))
613
614def jpeg.info(+src: Bin.Bytes) -> JPEGFrame:
615  U.branch(JPEGFrame, U32.is_eq(Bin.bytes.be(src, 0, 2n, 0), 65496), _ => jpeg.file(src, 2, jpeg.frame.empty(), False{}), _ => jpeg.frame.invalid(jpeg.frame.empty()))
616
617# Separable, normalized 8x8 inverse DCT. Coefficients remain quantized until
618# every sequential scan has been decoded; each component has a padded plane.
619def jpeg.idct.basis(+x: U32, +u: U32) -> F32:
620  +k = ((2 * x + 1) * u % 32 : U32)
621  +fold = U.choose(U32, U32.is_gt(k, 16), (32 - k : U32), k)
622  +negative = U32.is_gt(fold, 8)
623  +i = U.choose(U32, negative, (16 - fold : U32), fold)
624  +cosine = U.nth(Con{1.0, Con{0.9807852804032304, Con{0.9238795325112867, Con{0.8314696123025452, Con{0.7071067811865476, Con{0.5555702330196023, Con{0.3826834323650898, Con{0.1950903220161283, Con{0.0, Nil{}}}}}}}}}}, U32.to_nat(i), 0.0)
625  (0.5 * U.choose(F32, U32.is_eq(u, 0), 0.7071067811865476, 1.0) * U.choose(F32, negative, F32.neg(cosine), cosine) : F32)
626
627def jpeg.idct.horizontal(n: Nat, +u: U32, +x: U32, +row: U32, +offset: U32, +coefficients: Bin.Bytes, +quant: Bin.Bytes, acc: F32) -> F32:
628  match n:
629    case 0n: acc
630    case 1n+d:
631      +i = (row * 8 + u : U32)
632      jpeg.idct.horizontal(d, U32.inc(u), x, row, offset, coefficients, quant, (acc + jpeg.signed(Bin.bytes.get(coefficients, (offset + i : U32))) * U32.to_f32(Bin.bytes.get(quant, i)) * jpeg.idct.basis(x, u) : F32))
633
634def jpeg.idct.rows(n: Nat, +index: U32, +offset: U32, +coefficients: Bin.Bytes, +quant: Bin.Bytes) -> List<&2, F32>:
635  match n:
636    case 0n: Nil{}
637    case 1n+d: Con{jpeg.idct.horizontal(8n, 0, (index % 8 : U32), (index / 8 : U32), offset, coefficients, quant, 0.0), jpeg.idct.rows(d, U32.inc(index), offset, coefficients, quant)}
638
639def jpeg.idct.vertical(n: Nat, +v: U32, +x: U32, +y: U32, +rows: List<&2, F32>, acc: F32) -> F32:
640  match n:
641    case 0n: acc
642    case 1n+d: jpeg.idct.vertical(d, U32.inc(v), x, y, rows, (acc + U.nth(rows, U32.to_nat((v * 8 + x : U32)), 0.0) * jpeg.idct.basis(y, v) : F32))
643
644def jpeg.idct.pixels(n: Nat, +index: U32, +x: U32, +y: U32, +stride: U32, +rows: List<&2, F32>, out: Bin.Bytes) -> Bin.Bytes:
645  match n:
646    case 0n: out
647    case 1n+d:
648      +xx = (index % 8 : U32)
649      +yy = (index / 8 : U32)
650      jpeg.idct.pixels(d, U32.inc(index), x, y, stride, rows, Bin.bytes.set(out, ((y + yy) * stride + x + xx : U32), F32.to_u32(U.clamp(F32.floor((jpeg.idct.vertical(8n, 0, xx, yy, rows, 0.0) + 128.5 : F32)), 0.0, 255.0))))
651
652def jpeg.idct.blocks(n: Nat, +index: U32, +cols: U32, +coefficients: Bin.Bytes, +quant: Bin.Bytes, out: Bin.Bytes) -> Bin.Bytes:
653  match n:
654    case 0n: out
655    case 1n+d: jpeg.idct.blocks(d, U32.inc(index), cols, coefficients, quant, jpeg.idct.pixels(64n, 0, (index % cols * 8 : U32), (index / cols * 8 : U32), (cols * 8 : U32), jpeg.idct.rows(64n, 0, (index * 64 : U32), coefficients, quant), out))
656
657def jpeg.idct.components(xs: List<&2, JPEGComponent>, +quants: List<&2, JPEGQuant>) -> List<&2, JPEGComponent>:
658  match xs:
659    case Nil{}: Nil{}
660    case Con{JPEGComponent{id, h, v, quant, dc, ac, predictor, bw, bh, coefficients, samples}, rest}:
661      +cs = coefficients
662      +qq = quant
663      +ww = bw
664      +hh = bh
665      Con{JPEGComponent{id, h, v, qq, dc, ac, predictor, ww, hh, cs, jpeg.idct.blocks(U32.to_nat((ww * hh : U32)), 0, ww, cs, jpeg.quant.find(quants, qq), Bin.bytes.zero((ww * hh * 64 : U32)))}, jpeg.idct.components(rest, quants)}
666
667def jpeg.sample.byte(+samples: Bin.Bytes, +stride: U32, +x: U32, +y: U32) -> F32:
668  U32.to_f32(Bin.bytes.get(samples, (y * stride + x : U32)))
669
670def jpeg.sample(component: JPEGComponent, +w: U32, +h: U32, +hm: U32, +vm: U32, +x: U32, +y: U32) -> F32:
671  JPEGComponent{id, ch, cv, quant, dc, ac, predictor, bw, bh, coefficients, samples} = component
672  +hh = ch
673  +vv = cv
674  +data = samples
675  +stride = (bw * 8 : U32)
676  +width = jpeg.ceil((w * hh : U32), hm)
677  +height = jpeg.ceil((h * vv : U32), vm)
678  # libjpeg's horizontally doubled narrow planes use box reconstruction.
679  +box = U32.is_eq((hh * 2 : U32), hm) && U32.is_le(width, 2) && (U32.is_eq(vv, vm) || U32.is_eq((vv * 2 : U32), vm))
680  +px = U.choose(F32, box, U32.to_f32((x * hh / hm : U32)), U.clamp(((U32.to_f32(x) + 0.5) * U32.to_f32(hh) / U32.to_f32(hm) - 0.5 : F32), 0.0, U32.to_f32((width - 1 : U32))))
681  +py = U.choose(F32, box, U32.to_f32((y * vv / vm : U32)), U.clamp(((U32.to_f32(y) + 0.5) * U32.to_f32(vv) / U32.to_f32(vm) - 0.5 : F32), 0.0, U32.to_f32((height - 1 : U32))))
682  +x0 = F32.to_u32(F32.floor(px))
683  +y0 = F32.to_u32(F32.floor(py))
684  +x1 = U.choose(U32, U32.is_lt(U32.inc(x0), width), U32.inc(x0), x0)
685  +y1 = U.choose(U32, U32.is_lt(U32.inc(y0), height), U32.inc(y0), y0)
686  +fx = (px - U32.to_f32(x0) : F32)
687  +fy = (py - U32.to_f32(y0) : F32)
688  (jpeg.sample.byte(data, stride, x0, y0) * (1.0 - fx) * (1.0 - fy) + jpeg.sample.byte(data, stride, x1, y0) * fx * (1.0 - fy) + jpeg.sample.byte(data, stride, x0, y1) * (1.0 - fx) * fy + jpeg.sample.byte(data, stride, x1, y1) * fx * fy : F32)
689
690def jpeg.sample.channels(xs: List<&2, JPEGComponent>, +w: U32, +h: U32, +hm: U32, +vm: U32, +x: U32, +y: U32) -> List<&2, F32>:
691  match xs:
692    case Nil{}: Nil{}
693    case Con{c, rest}: Con{jpeg.sample(c, w, h, hm, vm, x, y), jpeg.sample.channels(rest, w, h, hm, vm, x, y)}
694
695def jpeg.channel(x: F32) -> F32:
696  (U.clamp(F32.floor((x + 0.5 : F32)), 0.0, 255.0) / 255.0 : F32)
697
698def jpeg.ycbcr(+y: F32, +cb: F32, +cr: F32) -> Img.Color:
699  Img.Color{jpeg.channel((y + 1.402 * (cr - 128.0) : F32)), jpeg.channel((y - 0.344136 * (cb - 128.0) - 0.714136 * (cr - 128.0) : F32)), jpeg.channel((y + 1.772 * (cb - 128.0) : F32)), 1.0}
700
701def jpeg.ycck(c: Img.Color, +k: F32) -> Img.Color:
702  Img.Color{r, g, b, a} = c
703  Img.Color{jpeg.channel(((1.0 - r) * k : F32)), jpeg.channel(((1.0 - g) * k : F32)), jpeg.channel(((1.0 - b) * k : F32)), 1.0}
704
705def jpeg.rgb.ids(xs: List<&2, JPEGComponent>) -> Bool:
706  match xs:
707    case Con{JPEGComponent{a, _, _, _, _, _, _, _, _, _, _}, Con{JPEGComponent{b, _, _, _, _, _, _, _, _, _, _}, Con{JPEGComponent{c, _, _, _, _, _, _, _, _, _, _}, Nil{}}}}: U32.is_eq(a, 82) && U32.is_eq(b, 71) && U32.is_eq(c, 66)
708    case _: False{}
709
710def jpeg.color(xs: List<&2, F32>, +adobe: U32) -> Img.Color:
711  match xs:
712    case Con{gray, Nil{}}:
713      +v = jpeg.channel(gray)
714      Img.Color{v, v, v, 1.0}
715    case Con{a, Con{b, Con{c, Nil{}}}}:
716      +aa = a
717      +bb = b
718      +cc = c
719      U.choose(Img.Color, U32.is_eq(adobe, 0), Img.Color{jpeg.channel(aa), jpeg.channel(bb), jpeg.channel(cc), 1.0}, jpeg.ycbcr(aa, bb, cc))
720    case Con{a, Con{b, Con{c, Con{k, Nil{}}}}}:
721      +aa = a
722      +bb = b
723      +cc = c
724      +kk = k
725      U.branch(Img.Color, U32.is_eq(adobe, 2), _ => jpeg.ycck(jpeg.ycbcr(aa, bb, cc), kk),
726        _ => Img.Color{jpeg.channel((aa * kk / 255.0 : F32)), jpeg.channel((bb * kk / 255.0 : F32)), jpeg.channel((cc * kk / 255.0 : F32)), 1.0})
727    case _: Img.clear()
728
729def jpeg.paint(n: Nat, +index: U32, +w: U32, +h: U32, +hm: U32, +vm: U32, +components: List<&2, JPEGComponent>, +adobe: U32, +depth: Nat, out: Img.Pixels) -> Img.Pixels:
730  match n:
731    case 0n: out
732    case 1n+d:
733      +x = (index % w : U32)
734      +y = (index / w : U32)
735      jpeg.paint(d, U32.inc(index), w, h, hm, vm, components, adobe, depth, Img.pixels.set(out, depth, x, y, jpeg.color(jpeg.sample.channels(components, w, h, hm, vm, x, y), adobe)))
736
737def jpeg.components.complete(xs: List<&2, JPEGComponent>) -> Bool:
738  match xs:
739    case Nil{}: True{}
740    case Con{JPEGComponent{_, _, _, _, _, _, _, _, _, _, samples}, rest}: U32.is_gt(Bin.bytes.get(samples, 0), 0) && jpeg.components.complete(rest)
741
742def jpeg.decode.frame(+frame: JPEGFrame) -> Img.Picture:
743  JPEGFrame{w, h, hm, vm, progressive, components, quants, huffs, restart, adobe, valid} = frame
744  +depth = Img.picture.depth(U.choose(U32, U32.is_gt(w, h), w, h), 1, 0n)
745  U.branch(Img.Picture, valid && U32.is_gt(w, 0) && U32.is_gt(h, 0) && jpeg.components.complete(components),
746    _ => Img.Picture{w, h, depth, jpeg.paint(U32.to_nat((w * h : U32)), 0, w, h, hm, vm, jpeg.idct.components(components, quants), U.choose(U32, U32.is_eq(adobe, 255) && jpeg.rgb.ids(components), 0, adobe), depth, Img.Flat{Img.clear()})}, _ => Img.picture.blank())
747
748def jpeg.decode(src: Bin.Bytes) -> Img.Picture:
749  jpeg.decode.frame(jpeg.info(src))