gabriel / muse public
compression.py python
275 lines 9.5 KB
Raw
sha256:6ed6b11aceb96bef850842a7aaac54c04843fa9fba5d1200e4bc3845e12d4142 docs: add muse-tag-guide mist (complete idiomatic guide wit… Sonnet 4.6 20 days ago
1 """Object compression and delta encoding for mpack wire transfer.
2
3 Tiers
4 -----
5 Tier 1a — zstd (preferred): ~3× faster than zlib at comparable ratios.
6 Available when the ``zstandard`` package is installed.
7 ``ZSTD_AVAILABLE`` is ``True`` when this tier is active.
8 Tier 1b — zlib (fallback): stdlib-only, always available.
9 Tier 2 — delta+zstd / delta+zlib: objects with a known base are
10 delta-encoded first, then compressed.
11
12 Use :func:`choose_compression` to select the best available encoding,
13 :func:`compress_and_encode` to compress and get the encoding name, and
14 :func:`decompress_frame` to decompress using the encoding stored in the
15 mpack object's ``enc`` field.
16
17 Receivers must support both ``zstd`` and ``zlib`` — the ``enc`` field on
18 each mpack object signals which was used.
19 """
20
21 import struct
22 import zlib
23
24 # ---------------------------------------------------------------------------
25 # zstd — optional fast compression tier
26 # ---------------------------------------------------------------------------
27
28 try:
29 import zstandard as _zstd_mod
30 ZSTD_AVAILABLE: bool = True
31 _ZSTD_COMPRESSOR = _zstd_mod.ZstdCompressor(level=3)
32 _ZSTD_DECOMPRESSOR = _zstd_mod.ZstdDecompressor()
33 except ImportError:
34 ZSTD_AVAILABLE = False
35 _zstd_mod = None # type: ignore[assignment]
36 _ZSTD_COMPRESSOR = None # type: ignore[assignment]
37 _ZSTD_DECOMPRESSOR = None # type: ignore[assignment]
38
39 def compress_zstd(data: bytes) -> bytes:
40 """Compress *data* with zstd at level 3 (speed/ratio balanced).
41
42 Level 3 delivers roughly the same compression ratio as zlib level 6
43 while compressing ~3× faster on source-code workloads.
44
45 Falls back to :func:`compress_zlib` when the ``zstandard`` package
46 is not installed (:data:`ZSTD_AVAILABLE` is ``False``).
47
48 Args:
49 data: Raw bytes to compress.
50
51 Returns:
52 Compressed bytes using zstd (or zlib as fallback).
53 """
54 if not ZSTD_AVAILABLE:
55 return compress_zlib(data)
56 return _ZSTD_COMPRESSOR.compress(data)
57
58 def decompress_zstd(data: bytes) -> bytes:
59 """Decompress zstd-compressed *data*.
60
61 Falls back to :func:`decompress_zlib` when the ``zstandard`` package
62 is not installed — consistent with :func:`compress_zstd`'s fallback.
63
64 Raises the underlying decompressor error on corrupt input so the
65 caller can return 400 to the remote rather than storing garbage.
66
67 Args:
68 data: zstd-compressed bytes.
69
70 Returns:
71 Original uncompressed bytes.
72 """
73 if not ZSTD_AVAILABLE:
74 return decompress_zlib(data)
75 return _ZSTD_DECOMPRESSOR.decompress(data)
76
77 def choose_compression() -> str:
78 """Return the best available compression encoding name.
79
80 Returns ``"zstd"`` when the ``zstandard`` package is installed,
81 ``"zlib"`` otherwise. Use the returned string as the ``enc`` field
82 value for mpack objects.
83 """
84 return "zstd" if ZSTD_AVAILABLE else "zlib"
85
86 def compress_and_encode(data: bytes) -> tuple[bytes, str]:
87 """Compress *data* with the best available algorithm.
88
89 Returns a ``(compressed_bytes, enc_name)`` tuple where ``enc_name``
90 is ``"zstd"`` or ``"zlib"`` — ready to write into a mpack object's
91 ``content`` and ``enc`` fields.
92
93 Args:
94 data: Raw bytes to compress.
95
96 Returns:
97 ``(compressed, enc)`` — compressed bytes and encoding name.
98 """
99 if ZSTD_AVAILABLE:
100 return _ZSTD_COMPRESSOR.compress(data), "zstd"
101 return zlib.compress(data, level=1), "zlib"
102
103 def decompress_frame(data: bytes, enc: str) -> bytes:
104 """Decompress *data* using the encoding specified by *enc*.
105
106 Dispatches on the ``enc`` field from a mpack object:
107
108 * ``"zstd"`` → zstd decompression
109 * ``"zlib"`` → zlib decompression
110 * ``"raw"`` → no decompression; *data* returned as-is
111
112 Args:
113 data: Compressed (or raw) bytes from a mpack object.
114 enc: Encoding name from the object's ``enc`` field.
115
116 Returns:
117 Decompressed bytes.
118
119 Raises:
120 ValueError: *enc* is not a recognised encoding name.
121 """
122 if enc == "zstd":
123 return decompress_zstd(data)
124 if enc == "zlib":
125 return zlib.decompress(data)
126 if enc == "raw":
127 return data
128 raise ValueError(
129 f"Unknown encoding {enc!r}. Supported: 'zstd', 'zlib', 'raw'. "
130 "Delta encodings ('delta+zstd', 'delta+zlib') require a base object "
131 "and must be handled by the caller."
132 )
133
134 # Minimum byte-run length to emit a COPY instruction rather than DATA.
135 # 32 bytes covers most source-code line edits without excessive overhead.
136 _COPY_THRESHOLD = 32
137
138 def compress_zlib(data: bytes) -> bytes:
139 """Compress *data* with zlib at level 1 (speed-optimised).
140
141 Level 1 gives ~60 % reduction on source text with negligible CPU overhead
142 at the throughput rates of a pack push (hundreds of objects per second).
143 """
144 return zlib.compress(data, level=1)
145
146 def decompress_zlib(data: bytes) -> bytes:
147 """Decompress zlib-compressed *data*.
148
149 Raises :class:`zlib.error` on corrupt input — let it propagate so the
150 server returns 400 to the client rather than storing garbage.
151 """
152 return zlib.decompress(data)
153
154 def compute_delta(base: bytes, target: bytes) -> bytes:
155 """Delta-encode *target* relative to *base*.
156
157 Returns a zlib-compressed instruction stream. Each instruction is one of:
158
159 * ``COPY``: ``b'\\x00' + struct.pack('>II', offset, length)``
160 — copy *length* bytes from *base* starting at *offset*.
161 * ``DATA``: ``b'\\x01' + struct.pack('>I', length) + data_bytes``
162 — emit *length* literal bytes verbatim.
163
164 Algorithm:
165 Build a hash table mapping every *_COPY_THRESHOLD*-byte chunk of
166 *base* to its last position. Walk *target* looking for matching
167 chunks; extend each match forward as far as possible, then emit a
168 single COPY instruction. Unmatched bytes accumulate in a DATA buffer
169 that is flushed whenever a COPY is emitted.
170
171 Raises ``ValueError`` when the compressed delta is no smaller than
172 ``zlib(target)`` — callers should fall back to raw encoding in that case.
173
174 Args:
175 base: Previous version of the object (already stored on the remote).
176 target: New version of the object to encode.
177
178 Returns:
179 Zlib-compressed binary instruction stream.
180 """
181 # Degenerate case: empty base → single DATA instruction.
182 if not base:
183 raw = b"\x01" + struct.pack(">I", len(target)) + target
184 return zlib.compress(raw, level=1)
185
186 # Build hash table: chunk → last position in base.
187 table: dict[bytes, int] = {}
188 for i in range(len(base) - _COPY_THRESHOLD + 1):
189 table[base[i : i + _COPY_THRESHOLD]] = i
190
191 parts: list[bytes] = []
192 data_buf = bytearray()
193 pos = 0
194
195 while pos < len(target):
196 # Try to find a matching chunk starting at *pos*.
197 if pos + _COPY_THRESHOLD <= len(target):
198 chunk = target[pos : pos + _COPY_THRESHOLD]
199 base_pos = table.get(chunk)
200 if base_pos is not None:
201 # Extend match forward as far as bytes agree.
202 length = _COPY_THRESHOLD
203 while (
204 pos + length < len(target)
205 and base_pos + length < len(base)
206 and target[pos + length] == base[base_pos + length]
207 ):
208 length += 1
209
210 # Flush accumulated DATA before emitting COPY.
211 if data_buf:
212 parts.append(
213 b"\x01" + struct.pack(">I", len(data_buf)) + bytes(data_buf)
214 )
215 data_buf = bytearray()
216
217 parts.append(b"\x00" + struct.pack(">II", base_pos, length))
218 pos += length
219 continue
220
221 # No match — accumulate literal byte.
222 data_buf.append(target[pos])
223 pos += 1
224
225 # Flush trailing DATA.
226 if data_buf:
227 parts.append(b"\x01" + struct.pack(">I", len(data_buf)) + bytes(data_buf))
228
229 delta = zlib.compress(b"".join(parts), level=1)
230 if len(delta) >= len(zlib.compress(target, level=1)):
231 raise ValueError("delta not profitable: delta >= zlib(target)")
232 return delta
233
234 def apply_delta(base: bytes, delta_compressed: bytes) -> bytes:
235 """Reconstruct the target by applying *delta_compressed* to *base*.
236
237 Inverse of :func:`compute_delta`. Parses the instruction stream produced
238 by that function and assembles the target byte-for-byte.
239
240 Args:
241 base: The base object bytes (fetched from storage).
242 delta_compressed: Zlib-compressed instruction stream.
243
244 Returns:
245 Reconstructed target bytes.
246
247 Raises:
248 zlib.error: *delta_compressed* is corrupt.
249 ValueError: Unknown instruction type in stream.
250 struct.error: Truncated instruction stream.
251 """
252 stream = zlib.decompress(delta_compressed)
253 result = bytearray()
254 pos = 0
255 n = len(stream)
256
257 while pos < n:
258 instr_type = stream[pos]
259 pos += 1
260
261 if instr_type == 0: # COPY
262 offset, length = struct.unpack_from(">II", stream, pos)
263 pos += 8
264 result.extend(base[offset : offset + length])
265
266 elif instr_type == 1: # DATA
267 (length,) = struct.unpack_from(">I", stream, pos)
268 pos += 4
269 result.extend(stream[pos : pos + length])
270 pos += length
271
272 else:
273 raise ValueError(f"unknown delta instruction type: {instr_type:#04x} at stream offset {pos - 1}")
274
275 return bytes(result)
File History 1 commit
sha256:6ed6b11aceb96bef850842a7aaac54c04843fa9fba5d1200e4bc3845e12d4142 docs: add muse-tag-guide mist (complete idiomatic guide wit… Sonnet 4.6 20 days ago