Viewing files at commit 653060c4e6b1Commit detailsView latest

scripts/oakbench/purezstd.py 662 lines · 25.1 KB · python
1
"""Pure-Python Zstandard frame decoder, stdlib-only (ADR-0001).
2
3
Independently written from RFC 8878 ("Zstandard Compression and the
4
application/zstd Media Type"): frame/block format (section 3.1.1), literals
5
section with Huffman coding (3.1.1.3.1, 4.2), sequences section with FSE
6
(3.1.1.3.2, 4.1), sequence execution and repeat offsets (3.1.1.4-3.1.1.5),
7
and the optional XXH64 content checksum. Dictionaries are not supported (Oak
8
never uses them). It shares no code with libzstd, which both Oak and the
9
``zstandard`` Python package link, so a decode here is independent of the
10
subject's codec. Every malformed input raises ``ZstdError``; bounds raise
11
``ZstdBoundError`` so callers can separate harness limits from bad bytes.
12
13
``encode_raw`` builds a valid frame of uncompressed (Raw) blocks; tests use
14
it to author stored objects without a compressor dependency.
15
"""
16
from __future__ import annotations
17
18
import struct
19
20
MAGIC = 0xFD2FB528
21
MAX_BLOCK = 128 * 1024
22
23
24
class ZstdError(ValueError):
25
    """Positively malformed or inconsistent compressed bytes."""
26
27
28
class ZstdBoundError(ValueError):
29
    """Input is well-formed so far but exceeds a caller-imposed bound."""
30
31
32
# --- XXH64 (content checksum) -------------------------------------------------
33
_P1, _P2, _P3, _P4, _P5 = (11400714785074694791, 14029467366897019727, 1609587929392839161,
34
                           9650029242287828579, 2870177450012600261)
35
_U64 = 0xFFFFFFFFFFFFFFFF
36
37
38
def _rotl(x, r):
39
    return ((x << r) | (x >> (64 - r))) & _U64
40
41
42
def _round(acc, lane):
43
    return (_rotl((acc + lane * _P2) & _U64, 31) * _P1) & _U64
44
45
46
def xxh64(data: bytes, seed: int = 0) -> int:
47
    n = len(data)
48
    i = 0
49
    if n >= 32:
50
        v1, v2, v3, v4 = (seed + _P1 + _P2) & _U64, (seed + _P2) & _U64, seed, (seed - _P1) & _U64
51
        while i + 32 <= n:
52
            a, b, c, d = struct.unpack_from('<4Q', data, i)
53
            v1, v2, v3, v4 = _round(v1, a), _round(v2, b), _round(v3, c), _round(v4, d)
54
            i += 32
55
        h = (_rotl(v1, 1) + _rotl(v2, 7) + _rotl(v3, 12) + _rotl(v4, 18)) & _U64
56
        for v in (v1, v2, v3, v4):
57
            h = ((h ^ _round(0, v)) * _P1 + _P4) & _U64
58
    else:
59
        h = (seed + _P5) & _U64
60
    h = (h + n) & _U64
61
    while i + 8 <= n:
62
        (lane,) = struct.unpack_from('<Q', data, i)
63
        h = (_rotl(h ^ _round(0, lane), 27) * _P1 + _P4) & _U64
64
        i += 8
65
    if i + 4 <= n:
66
        (lane,) = struct.unpack_from('<I', data, i)
67
        h = (_rotl(h ^ (lane * _P1 & _U64), 23) * _P2 + _P3) & _U64
68
        i += 4
69
    while i < n:
70
        h = (_rotl(h ^ (data[i] * _P5 & _U64), 11) * _P1) & _U64
71
        i += 1
72
    h ^= h >> 33
73
    h = (h * _P2) & _U64
74
    h ^= h >> 29
75
    h = (h * _P3) & _U64
76
    h ^= h >> 32
77
    return h
78
79
80
# --- bit readers --------------------------------------------------------------
81
class _ForwardBits:
82
    """Little-endian forward bit reader (FSE table descriptions)."""
83
84
    def __init__(self, data: bytes, start: int, end: int):
85
        self.data, self.pos, self.end = data, start * 8, end * 8
86
87
    def read(self, n: int) -> int:
88
        if n == 0:
89
            return 0
90
        if self.pos + n > self.end:
91
            raise ZstdError('FSE table description overruns its section')
92
        byte = self.pos >> 3
93
        value = int.from_bytes(self.data[byte:byte + 8], 'little') >> (self.pos & 7)
94
        self.pos += n
95
        return value & ((1 << n) - 1)
96
97
    def peek(self, n: int) -> int:
98
        byte = self.pos >> 3
99
        value = int.from_bytes(self.data[byte:byte + 8], 'little') >> (self.pos & 7)
100
        available = self.end - self.pos
101
        if available < n:
102
            value &= (1 << max(available, 0)) - 1
103
        return value & ((1 << n) - 1)
104
105
    def consumed_bytes(self, start: int) -> int:
106
        return (self.pos + 7) // 8 - start
107
108
109
class _BackwardBits:
110
    """Reverse bitstream (RFC 8878 4.1/4.2.2): read from the end toward the start.
111
112
    ``offset`` may go negative: missing low bits read as zero, and callers use
113
    the sign to detect overflow exactly as the RFC's end-of-stream rules do.
114
    """
115
116
    def __init__(self, data: bytes):
117
        if not data:
118
            raise ZstdError('empty bitstream')
119
        last = data[-1]
120
        if last == 0:
121
            raise ZstdError('bitstream final byte lacks padding marker')
122
        self.data = data
123
        self.offset = len(data) * 8 - (8 - (last.bit_length() - 1))
124
125
    def _at(self, offset: int, n: int) -> int:
126
        if n == 0:
127
            return 0
128
        if offset >= 0:
129
            byte = offset >> 3
130
            value = int.from_bytes(self.data[byte:byte + 8], 'little') >> (offset & 7)
131
            return value & ((1 << n) - 1)
132
        if offset + n <= 0:
133
            return 0
134
        value = int.from_bytes(self.data[:8], 'little') << (-offset)
135
        return value & ((1 << n) - 1)
136
137
    def read(self, n: int) -> int:
138
        self.offset -= n
139
        return self._at(self.offset, n)
140
141
    def peek(self, n: int) -> int:
142
        return self._at(self.offset - n, n)
143
144
145
# --- FSE ----------------------------------------------------------------------
146
def _read_fse_description(data: bytes, start: int, end: int, max_log: int, max_symbol: int):
147
    """Normalized counts from an FSE table description; returns (probs, log, used_bytes)."""
148
    bits = _ForwardBits(data, start, end)
149
    log = bits.read(4) + 5
150
    if log > max_log:
151
        raise ZstdError('FSE accuracy log too large')
152
    size = 1 << log
153
    remaining = size + 1
154
    threshold = size
155
    nb = log + 1
156
    probs = []
157
    while remaining > 1:
158
        if len(probs) > max_symbol:
159
            raise ZstdError('FSE description has too many symbols')
160
        maximum = 2 * threshold - 1 - remaining
161
        low = bits.peek(nb - 1)
162
        if low < maximum:
163
            count = low
164
            bits.read(nb - 1)
165
        else:
166
            count = bits.peek(nb)
167
            if count >= threshold:
168
                count -= maximum
169
            bits.read(nb)
170
        count -= 1
171
        remaining -= abs(count)
172
        probs.append(count)
173
        if count == 0:
174
            while True:
175
                repeat = bits.read(2)
176
                probs.extend([0] * repeat)
177
                if repeat != 3:
178
                    break
179
        while remaining < threshold and threshold > 1:
180
            nb -= 1
181
            threshold >>= 1
182
    if remaining != 1 or len(probs) > max_symbol + 1:
183
        raise ZstdError('FSE description probabilities do not sum to table size')
184
    return probs, log, bits.consumed_bytes(start)
185
186
187
def _build_fse(probs, log):
188
    """Decoding table: list of (symbol, nb_bits, baseline) indexed by state."""
189
    size = 1 << log
190
    symbol = [0] * size
191
    high = size - 1
192
    next_count = [0] * len(probs)
193
    for s, p in enumerate(probs):
194
        if p == -1:
195
            symbol[high] = s
196
            high -= 1
197
            next_count[s] = 1
198
        else:
199
            next_count[s] = p
200
    step = (size >> 1) + (size >> 3) + 3
201
    mask = size - 1
202
    position = 0
203
    for s, p in enumerate(probs):
204
        for _ in range(p if p > 0 else 0):
205
            symbol[position] = s
206
            position = (position + step) & mask
207
            while position > high:
208
                position = (position + step) & mask
209
    if position != 0:
210
        raise ZstdError('FSE table spread inconsistent')
211
    table = []
212
    for state in range(size):
213
        s = symbol[state]
214
        n = next_count[s]
215
        next_count[s] += 1
216
        nb = log - (n.bit_length() - 1)
217
        table.append((s, nb, (n << nb) - size))
218
    return table, log
219
220
221
def _rle_table(value):
222
    return [(value, 0, 0)], 0
223
224
225
# Predefined distributions (RFC 8878 3.1.1.3.2.2).
226
_LL_DEFAULT = _build_fse([4, 3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 2, 1, 1, 1,
227
                          1, 1, -1, -1, -1, -1], 6)
228
_ML_DEFAULT = _build_fse([1, 4, 3, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
229
                          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1, -1, -1, -1], 6)
230
_OF_DEFAULT = _build_fse([1, 1, 1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1,
231
                          -1], 5)
232
233
_LL_BASE = list(range(16)) + [16, 18, 20, 22, 24, 28, 32, 40, 48, 64, 128, 256, 512, 1024, 2048, 4096,
234
                              8192, 16384, 32768, 65536]
235
_LL_BITS = [0] * 16 + [1, 1, 1, 1, 2, 2, 3, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
236
_ML_BASE = [n + 3 for n in range(32)] + [35, 37, 39, 41, 43, 47, 51, 59, 67, 83, 99, 131, 259, 515, 1027,
237
                                         2051, 4099, 8195, 16387, 32771, 65539]
238
_ML_BITS = [0] * 32 + [1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
239
240
241
# --- Huffman ------------------------------------------------------------------
242
def _read_huffman(data: bytes, start: int, end: int):
243
    """Decode a Huffman tree description; returns ((symbols, bits, max_bits), used)."""
244
    if start >= end:
245
        raise ZstdError('missing Huffman tree description')
246
    header = data[start]
247
    if header >= 128:
248
        count = header - 127
249
        used = 1 + (count + 1) // 2
250
        if start + used > end:
251
            raise ZstdError('Huffman weights overrun')
252
        weights = []
253
        for i in range(count):
254
            byte = data[start + 1 + i // 2]
255
            weights.append(byte >> 4 if i % 2 == 0 else byte & 15)
256
    else:
257
        used = 1 + header
258
        if start + used > end:
259
            raise ZstdError('Huffman FSE weights overrun')
260
        probs, log, desc = _read_fse_description(data, start + 1, start + used, 6, 255)
261
        table, log = _build_fse(probs, log)
262
        stream = data[start + 1 + desc:start + used]
263
        bits = _BackwardBits(stream)
264
        state1 = bits.read(log)
265
        state2 = bits.read(log)
266
        weights = []
267
        while True:
268
            if len(weights) > 255:
269
                raise ZstdError('too many Huffman weights')
270
            s, nb, base = table[state1]
271
            weights.append(s)
272
            state1 = base + bits.read(nb)
273
            if bits.offset < 0:
274
                weights.append(table[state2][0])
275
                break
276
            s, nb, base = table[state2]
277
            weights.append(s)
278
            state2 = base + bits.read(nb)
279
            if bits.offset < 0:
280
                weights.append(table[state1][0])
281
                break
282
    if len(weights) > 255:
283
        raise ZstdError('too many Huffman weights')
284
    total = 0
285
    for w in weights:
286
        if w > 11:
287
            raise ZstdError('Huffman weight too large')
288
        if w:
289
            total += 1 << (w - 1)
290
    if total == 0:
291
        raise ZstdError('empty Huffman weights')
292
    max_bits = total.bit_length()
293
    left = (1 << max_bits) - total
294
    if left & (left - 1):
295
        raise ZstdError('Huffman weights do not complete a power of two')
296
    weights.append(left.bit_length())
297
    if max_bits > 11:
298
        raise ZstdError('Huffman code too long')
299
    bits_per_symbol = [max_bits + 1 - w if w else 0 for w in weights]
300
    rank_count = [0] * (max_bits + 2)
301
    for b in bits_per_symbol:
302
        if b:
303
            rank_count[b] += 1
304
    rank_index = [0] * (max_bits + 2)
305
    for b in range(max_bits, 0, -1):
306
        rank_index[b - 1] = rank_index[b] + rank_count[b] * (1 << (max_bits - b))
307
    size = 1 << max_bits
308
    symbols = [0] * size
309
    lengths = [0] * size
310
    for symbol, b in enumerate(bits_per_symbol):
311
        if not b:
312
            continue
313
        span = 1 << (max_bits - b)
314
        code = rank_index[b]
315
        if code + span > size:
316
            raise ZstdError('Huffman table overflow')
317
        for j in range(code, code + span):
318
            symbols[j] = symbol
319
            lengths[j] = b
320
        rank_index[b] += span
321
    return (symbols, lengths, max_bits), used
322
323
324
def _huffman_stream(stream: bytes, table, count: int, out: bytearray):
325
    symbols, lengths, max_bits = table
326
    bits = _BackwardBits(stream)
327
    for _ in range(count):
328
        index = bits.peek(max_bits)
329
        out.append(symbols[index])
330
        bits.offset -= lengths[index]
331
    if bits.offset != 0:
332
        raise ZstdError('Huffman stream not fully consumed')
333
334
335
# --- frames and blocks --------------------------------------------------------
336
class _FrameState:
337
    def __init__(self):
338
        self.huffman = None
339
        self.tables = [None, None, None]  # LL, OF, ML
340
        self.reps = [1, 4, 8]
341
342
343
def _literals(block: bytes, state: _FrameState):
344
    kind = block[0] & 3
345
    size_format = (block[0] >> 2) & 3
346
    if kind in (0, 1):
347
        if size_format in (0, 2):
348
            regen, header = block[0] >> 3, 1
349
        elif size_format == 1:
350
            if len(block) < 2:
351
                raise ZstdError('truncated literals header')
352
            regen, header = (block[0] | block[1] << 8) >> 4, 2
353
        else:
354
            if len(block) < 3:
355
                raise ZstdError('truncated literals header')
356
            regen, header = (block[0] | block[1] << 8 | block[2] << 16) >> 4, 3
357
        if kind == 0:
358
            if header + regen > len(block):
359
                raise ZstdError('raw literals overrun block')
360
            return bytes(block[header:header + regen]), header + regen
361
        if header + 1 > len(block):
362
            raise ZstdError('RLE literals overrun block')
363
        return bytes([block[header]]) * regen, header + 1
364
    header = {0: 3, 1: 3, 2: 4, 3: 5}[size_format]
365
    if len(block) < header:
366
        raise ZstdError('truncated literals header')
367
    h = int.from_bytes(block[:header], 'little')
368
    if size_format in (0, 1):
369
        regen, comp = (h >> 4) & 0x3FF, (h >> 14) & 0x3FF
370
    elif size_format == 2:
371
        regen, comp = (h >> 4) & 0x3FFF, (h >> 18) & 0x3FFF
372
    else:
373
        regen, comp = (h >> 4) & 0x3FFFF, (h >> 22) & 0x3FFFF
374
    if regen > MAX_BLOCK or header + comp > len(block):
375
        raise ZstdError('compressed literals overrun block')
376
    start, end = header, header + comp
377
    if kind == 2:
378
        state.huffman, used = _read_huffman(block, start, end)
379
        start += used
380
    elif state.huffman is None:
381
        raise ZstdError('treeless literals without a previous Huffman table')
382
    out = bytearray()
383
    if size_format == 0:
384
        _huffman_stream(block[start:end], state.huffman, regen, out)
385
    else:
386
        if start + 6 > end:
387
            raise ZstdError('truncated literal jump table')
388
        s1, s2, s3 = struct.unpack_from('<3H', block, start)
389
        start += 6
390
        s4 = end - start - s1 - s2 - s3
391
        if s4 < 1:
392
            raise ZstdError('invalid literal stream sizes')
393
        each = (regen + 3) // 4
394
        counts = (each, each, each, regen - 3 * each)
395
        if counts[3] < 0:
396
            raise ZstdError('invalid literal regenerated size')
397
        for size, count in zip((s1, s2, s3, s4), counts):
398
            _huffman_stream(block[start:start + size], state.huffman, count, out)
399
            start += size
400
    if len(out) != regen:
401
        raise ZstdError('literal size mismatch')
402
    return bytes(out), end
403
404
405
def _sequence_table(block, pos, mode, index, state, default, max_log, max_symbol):
406
    if mode == 0:
407
        table = default
408
    elif mode == 1:
409
        if pos >= len(block):
410
            raise ZstdError('truncated RLE sequence table')
411
        if block[pos] > max_symbol:
412
            raise ZstdError('RLE sequence symbol out of range')
413
        table = _rle_table(block[pos])
414
        pos += 1
415
    elif mode == 2:
416
        probs, log, used = _read_fse_description(block, pos, len(block), max_log, max_symbol)
417
        table = _build_fse(probs, log)
418
        pos += used
419
    else:
420
        table = state.tables[index]
421
        if table is None:
422
            raise ZstdError('repeat sequence table without a previous table')
423
    state.tables[index] = table
424
    return pos
425
426
427
def _compressed_block(block: bytes, state: _FrameState, out: bytearray, limit: int,
428
                      frame_start: int = 0, window: int = None):
429
    if not block:
430
        raise ZstdError('empty compressed block')
431
    literals, pos = _literals(block, state)
432
    if pos >= len(block):
433
        raise ZstdError('missing sequences section')
434
    b0 = block[pos]
435
    if b0 == 0:
436
        count, pos = 0, pos + 1
437
    elif b0 < 128:
438
        count, pos = b0, pos + 1
439
    elif b0 < 255:
440
        if pos + 2 > len(block):
441
            raise ZstdError('truncated sequence count')
442
        count, pos = ((b0 - 128) << 8) + block[pos + 1], pos + 2
443
    else:
444
        if pos + 3 > len(block):
445
            raise ZstdError('truncated sequence count')
446
        count, pos = block[pos + 1] + (block[pos + 2] << 8) + 0x7F00, pos + 3
447
    if count == 0:
448
        if pos != len(block):
449
            raise ZstdError('trailing bytes after empty sequences section')
450
        if len(out) + len(literals) > limit:
451
            raise ZstdBoundError('decompressed output exceeds bound')
452
        out += literals
453
        return
454
    if pos >= len(block):
455
        raise ZstdError('missing sequence modes')
456
    modes = block[pos]
457
    if modes & 3:
458
        raise ZstdError('reserved sequence mode bits set')
459
    pos += 1
460
    pos = _sequence_table(block, pos, modes >> 6, 0, state, _LL_DEFAULT, 9, 35)
461
    pos = _sequence_table(block, pos, (modes >> 4) & 3, 1, state, _OF_DEFAULT, 8, 31)
462
    pos = _sequence_table(block, pos, (modes >> 2) & 3, 2, state, _ML_DEFAULT, 9, 52)
463
    if pos >= len(block):
464
        raise ZstdError('missing sequence bitstream')
465
    (ll_table, ll_log), (of_table, of_log), (ml_table, ml_log) = state.tables
466
    bits = _BackwardBits(block[pos:])
467
    ll_state = bits.read(ll_log)
468
    of_state = bits.read(of_log)
469
    ml_state = bits.read(ml_log)
470
    reps = state.reps
471
    lit_pos = 0
472
    for i in range(count):
473
        of_code = of_table[of_state][0]
474
        ll_code = ll_table[ll_state][0]
475
        ml_code = ml_table[ml_state][0]
476
        if of_code > 31:
477
            raise ZstdError('offset code out of range')
478
        offset_value = (1 << of_code) + bits.read(of_code)
479
        ml = _ML_BASE[ml_code] + bits.read(_ML_BITS[ml_code])
480
        ll = _LL_BASE[ll_code] + bits.read(_LL_BITS[ll_code])
481
        if offset_value > 3:
482
            offset = offset_value - 3
483
            reps[2], reps[1], reps[0] = reps[1], reps[0], offset
484
        else:
485
            index = offset_value - 1 + (1 if ll == 0 else 0)
486
            if index == 0:
487
                offset = reps[0]
488
            else:
489
                offset = reps[index] if index < 3 else reps[0] - 1
490
                if index > 1:
491
                    reps[2] = reps[1]
492
                reps[1] = reps[0]
493
                reps[0] = offset
494
        if lit_pos + ll > len(literals):
495
            raise ZstdError('sequence literal length overruns literals')
496
        if len(out) + ll + ml > limit:
497
            raise ZstdBoundError('decompressed output exceeds bound')
498
        out += literals[lit_pos:lit_pos + ll]
499
        lit_pos += ll
500
        # Matches may reach only this frame's own output (no dictionary, no
501
        # previous frame) and never further back than the declared window.
502
        if offset == 0 or offset > len(out) - frame_start or (window is not None and offset > window):
503
            raise ZstdError('match offset outside frame window')
504
        start = len(out) - offset
505
        if offset >= ml:
506
            out += out[start:start + ml]
507
        else:
508
            for j in range(ml):
509
                out.append(out[start + j])
510
        if i != count - 1:
511
            s, nb, base = ll_table[ll_state]
512
            ll_state = base + bits.read(nb)
513
            s, nb, base = ml_table[ml_state]
514
            ml_state = base + bits.read(nb)
515
            s, nb, base = of_table[of_state]
516
            of_state = base + bits.read(nb)
517
    if bits.offset != 0:
518
        raise ZstdError('sequence bitstream not fully consumed')
519
    rest = literals[lit_pos:]
520
    if len(out) + len(rest) > limit:
521
        raise ZstdBoundError('decompressed output exceeds bound')
522
    out += rest
523
524
525
def frame_content_size(data: bytes):
526
    """Declared content size of the first frame, or None when absent."""
527
    return _frame_header(data, 0)[1]
528
529
530
def _frame_header(data: bytes, pos: int):
531
    if len(data) < pos + 5:
532
        raise ZstdError('truncated frame header')
533
    (magic,) = struct.unpack_from('<I', data, pos)
534
    if magic != MAGIC:
535
        raise ZstdError('not a zstd frame')
536
    descriptor = data[pos + 4]
537
    fcs_flag = descriptor >> 6
538
    single = (descriptor >> 5) & 1
539
    if descriptor & 0x08:
540
        raise ZstdError('reserved frame descriptor bit set')
541
    checksum = (descriptor >> 2) & 1
542
    dict_flag = descriptor & 3
543
    pos += 5
544
    window = None
545
    if not single:
546
        if pos >= len(data):
547
            raise ZstdError('truncated window descriptor')
548
        wd = data[pos]
549
        exponent, mantissa = wd >> 3, wd & 7
550
        base = 1 << (10 + exponent)
551
        window = base + (base // 8) * mantissa
552
        pos += 1
553
    dict_size = (0, 1, 2, 4)[dict_flag]
554
    if dict_size:
555
        if pos + dict_size > len(data):
556
            raise ZstdError('truncated dictionary id')
557
        if int.from_bytes(data[pos:pos + dict_size], 'little'):
558
            raise ZstdError('dictionary frames are not supported')
559
        pos += dict_size
560
    fcs_size = (1 if single else 0, 2, 4, 8)[fcs_flag]
561
    content_size = None
562
    if fcs_size:
563
        if pos + fcs_size > len(data):
564
            raise ZstdError('truncated frame content size')
565
        content_size = int.from_bytes(data[pos:pos + fcs_size], 'little')
566
        if fcs_size == 2:
567
            content_size += 256
568
        pos += fcs_size
569
    if single:
570
        window = content_size
571
    return pos, content_size, window, checksum
572
573
574
def decompress(data: bytes, max_output: int = 64 * 1024 * 1024, max_window: int = None) -> bytes:
575
    """Decode every frame in ``data`` (skippable frames ignored); trailing garbage is an error."""
576
    try:
577
        return _decompress(data, max_output, max_window)
578
    except (IndexError, KeyError, struct.error, OverflowError) as exc:
579
        # Malformed input must never escape as an instrument crash.
580
        raise ZstdError(f'malformed frame: {type(exc).__name__}') from exc
581
582
583
def _decompress(data, max_output, max_window):
584
    data = bytes(data)
585
    out = bytearray()
586
    pos = 0
587
    frames = 0
588
    while pos < len(data) or frames == 0:
589
        if len(data) - pos >= 4 and (struct.unpack_from('<I', data, pos)[0] & 0xFFFFFFF0) == 0x184D2A50:
590
            if pos + 8 > len(data):
591
                raise ZstdError('truncated skippable frame')
592
            (size,) = struct.unpack_from('<I', data, pos + 4)
593
            if pos + 8 + size > len(data):
594
                raise ZstdError('truncated skippable frame')
595
            pos += 8 + size
596
            frames += 1
597
            continue
598
        pos, content_size, window, checksum = _frame_header(data, pos)
599
        if max_window is not None and window is not None and window > max_window:
600
            raise ZstdBoundError('decompression window exceeds bound')
601
        if content_size is not None and content_size > max_output - len(out):
602
            raise ZstdBoundError('declared content size exceeds bound')
603
        frame_start = len(out)
604
        state = _FrameState()
605
        while True:
606
            if pos + 3 > len(data):
607
                raise ZstdError('truncated block header')
608
            header = data[pos] | data[pos + 1] << 8 | data[pos + 2] << 16
609
            pos += 3
610
            last, kind, size = header & 1, (header >> 1) & 3, header >> 3
611
            if kind == 3:
612
                raise ZstdError('reserved block type')
613
            if size > MAX_BLOCK:
614
                raise ZstdError('block exceeds maximum size')
615
            if kind == 1:
616
                if pos + 1 > len(data):
617
                    raise ZstdError('truncated RLE block')
618
                if len(out) + size > max_output:
619
                    raise ZstdBoundError('decompressed output exceeds bound')
620
                out += bytes([data[pos]]) * size
621
                pos += 1
622
            else:
623
                if pos + size > len(data):
624
                    raise ZstdError('truncated block')
625
                if kind == 0:
626
                    if len(out) + size > max_output:
627
                        raise ZstdBoundError('decompressed output exceeds bound')
628
                    out += data[pos:pos + size]
629
                else:
630
                    _compressed_block(data[pos:pos + size], state, out, max_output, frame_start, window)
631
                pos += size
632
            if last:
633
                break
634
        produced = len(out) - frame_start
635
        if content_size is not None and produced != content_size:
636
            raise ZstdError('frame content size mismatch')
637
        if checksum:
638
            if pos + 4 > len(data):
639
                raise ZstdError('truncated content checksum')
640
            (stored,) = struct.unpack_from('<I', data, pos)
641
            if xxh64(bytes(out[frame_start:])) & 0xFFFFFFFF != stored:
642
                raise ZstdError('content checksum mismatch')
643
            pos += 4
644
        frames += 1
645
    return bytes(out)
646
647
648
def encode_raw(payload: bytes, checksum: bool = True) -> bytes:
649
    """A valid zstd frame holding ``payload`` in Raw blocks (no compression)."""
650
    payload = bytes(payload)
651
    descriptor = (3 << 6) | (4 if checksum else 0)  # 8-byte content size, not single-segment
652
    out = bytearray(struct.pack('<IB', MAGIC, descriptor))
653
    out.append((17 - 10) << 3)  # window 128 KiB
654
    out += struct.pack('<Q', len(payload))
655
    blocks = [payload[i:i + MAX_BLOCK] for i in range(0, len(payload), MAX_BLOCK)] or [b'']
656
    for index, block in enumerate(blocks):
657
        last = 1 if index == len(blocks) - 1 else 0
658
        header = last | (0 << 1) | (len(block) << 3)
659
        out += header.to_bytes(3, 'little') + block
660
    if checksum:
661
        out += struct.pack('<I', xxh64(payload) & 0xFFFFFFFF)
662
    return bytes(out)