import zlib.zchecksum as zchecksum import zlib.zhuff as zhuff import zlib.zdef as zdef import string import memhub class zdeflate_stream: pool: memhub.MemManager* writer: zdef.zbit_writer window: bytes window_pos: int window_size: int hash_head: int* hash_prev: int* level: int strategy: int wbits: int is_finished: int adler: uint32 def find_match(self, data: bytes, data_len: sizet, pos: sizet, best_dist: int*) -> int: if pos + 2 >= data_len: return 0 h: uint = zdeflate_hash(data + pos) chain: int = self.hash_head[h] best_len: int = zdef.ZDEFLATE_MIN_MATCH - 1 *best_dist = 0 max_chain: int = 128 if (self.level >= 5) else (32 if (self.level >= 2) else 4) limit: int = (1 << self.wbits) if (self.wbits > 0) else zdef.ZDEFLATE_WINDOW_SIZE if limit > zdef.ZDEFLATE_WINDOW_SIZE: limit = zdef.ZDEFLATE_WINDOW_SIZE attempts: int = 0 while chain >= 0 and attempts < max_chain: dist: int = int(pos) - chain if dist <= 0 or dist > limit: break match_len: int = 0 max_len: int = int(data_len - pos) if max_len > zdef.ZDEFLATE_MAX_MATCH: max_len = zdef.ZDEFLATE_MAX_MATCH while match_len < max_len and data[pos + match_len] == data[chain + match_len]: match_len += 1 if match_len > best_len: best_len = match_len *best_dist = dist if best_len >= zdef.ZDEFLATE_MAX_MATCH: break chain = self.hash_prev[chain & zdef.ZDEFLATE_WINDOW_MASK] if chain <= int(pos) - limit: break attempts += 1 return best_len if (best_len >= zdef.ZDEFLATE_MIN_MATCH) else 0 def update_hash(self, data: bytes, pos: sizet, data_len: sizet): # Only maintain the hash at positions that can still start a match: tail bytes with pos+2 >= data_len are never referenced by find_match, and zdeflate_hash reads pos..pos+2, so going past the end overflows (see ASan heap-buffer-overflow). if pos + 2 >= data_len: return h: uint = zdeflate_hash(data + pos) self.hash_prev[pos & zdef.ZDEFLATE_WINDOW_MASK] = self.hash_head[h] self.hash_head[h] = int(pos) def write_fixed_block(self, data: uint8*, length: sizet, final: int): lit_tree = zhuff.zhuff_tree() dist_tree = zhuff.zhuff_tree() lit_tree.build_fixed_lit_tree() dist_tree.build_fixed_dist_tree() self.writer.write_bits(1 if final else 0, 1) self.writer.write_bits(1, 2) pos: sizet = 0 while pos < length: best_dist: int = 0 match_len: int = 0 if self.level > 0 and pos + 2 < length: match_len = self.find_match(data, length, pos, &best_dist) if match_len >= zdef.ZDEFLATE_MIN_MATCH: len_sym: int = zdeflate_len_to_symbol(match_len) lit_tree.encode_symbol(len_sym, &self.writer) extra: int = zdef.zdeflate_len_extra_bits[len_sym - 257] if extra > 0: self.writer.write_bits(match_len - zdef.zdeflate_len_base[len_sym - 257], extra) dist_sym: int = zdeflate_dist_to_symbol(best_dist) dist_tree.encode_symbol(dist_sym, &self.writer) extra = zdef.zdeflate_dist_extra_bits[dist_sym] if extra > 0: self.writer.write_bits(best_dist - zdef.zdeflate_dist_base[dist_sym], extra) for i in range(int(match_len)): self.update_hash(data, pos + i, length) pos += match_len else: lit_tree.encode_symbol(data[pos], &self.writer) self.update_hash(data, pos, length) pos += 1 lit_tree.encode_symbol(256, &self.writer) def encode_block_data(self, data: uint8*, length: sizet, lit_tree: zhuff.zhuff_tree*, dist_tree: zhuff.zhuff_tree*): pos: sizet = 0 while pos < length: best_dist: int = 0 match_len: int = 0 if self.level > 0 and pos + 2 < length: match_len = self.find_match(data, length, pos, &best_dist) if match_len >= zdef.ZDEFLATE_MIN_MATCH: len_sym: int = zdeflate_len_to_symbol(match_len) lit_tree.encode_symbol(len_sym, &self.writer) extra: int = zdef.zdeflate_len_extra_bits[len_sym - 257] if extra > 0: self.writer.write_bits(match_len - zdef.zdeflate_len_base[len_sym - 257], extra) dist_sym: int = zdeflate_dist_to_symbol(best_dist) dist_tree.encode_symbol(dist_sym, &self.writer) extra = zdef.zdeflate_dist_extra_bits[dist_sym] if extra > 0: self.writer.write_bits(best_dist - zdef.zdeflate_dist_base[dist_sym], extra) for j in range(int(match_len)): self.update_hash(data, pos + j, length) pos += match_len else: lit_tree.encode_symbol(data[pos], &self.writer) self.update_hash(data, pos, length) pos += 1 lit_tree.encode_symbol(256, &self.writer) def write_dynamic_block(self, data: uint8*, length: sizet, final: int): lit_freqs: t.CArray[int, 288] dist_freqs: t.CArray[int, 32] zdeflate_count_freqs(lit_freqs, dist_freqs, self, data, length) hlit: int = 286 while hlit > 257 and lit_freqs[hlit - 1] == 0: hlit -= 1 hdist: int = 30 while hdist > 1 and dist_freqs[hdist - 1] == 0: hdist -= 1 lit_tree = zhuff.zhuff_tree() dist_tree = zhuff.zhuff_tree() lit_tree.build_codes(lit_freqs, hlit, zdef.ZDEFLATE_MAX_BITS) dist_tree.build_codes(dist_freqs, hdist, zdef.ZDEFLATE_MAX_BITS) all_lengths: [288 + 32]int for i in range(int(hlit)): all_lengths[i] = lit_tree.codes[i].bits for i in range(int(hdist)): all_lengths[hlit + i] = dist_tree.codes[i].bits total_lengths: int = hlit + hdist cl_freqs: [19]int zdeflate_count_cl_freqs(all_lengths, total_lengths, cl_freqs) cl_tree = zhuff.zhuff_tree() cl_tree.build_codes(cl_freqs, 19, 7) hclen: int = 19 while hclen > 4 and cl_tree.codes[zdef.zdeflate_codelen_order[hclen - 1]].bits == 0: hclen -= 1 self.writer.write_bits(1 if final else 0, 1) self.writer.write_bits(2, 2) self.writer.write_bits(hlit - 257, 5) self.writer.write_bits(hdist - 1, 5) self.writer.write_bits(hclen - 4, 4) for j in range(int(hclen)): self.writer.write_bits(cl_tree.codes[zdef.zdeflate_codelen_order[j]].bits, 3) zdeflate_write_cl_encoded(all_lengths, total_lengths, &self.writer, &cl_tree) self.encode_block_data(data, length, &lit_tree, &dist_tree) def compress_block(self, data: uint8*, length: sizet, final: int): if self.level == 0: self.write_stored_block(data, length, final) elif self.strategy == 4 or length < 128: self.write_fixed_block(data, length, final) else: self.write_dynamic_block(data, length, final) def write_stored_block(self, data: uint8*, length: sizet, final: int): self.writer.stored_block(data, length, final) # out is a "pointer to pointer" (i8**): see zinflate.vp decompress -- an embedded `t.CPtr[...]` collapses the declare side (VpiReader) to i8, mismatching the definition side's i8**. def compress(self, data: uint8*, length: sizet, out: bytes*, out_len: sizet*) -> int: if self.is_finished: return -2 if not data or length == 0: *out = None *out_len = 0 return 0 self.adler = zchecksum.zchecksum_adler32(data, length, self.adler) string.memset(self.hash_head, -1, zdef.ZDEFLATE_HASH_SIZE * int.__sizeof__()) string.memset(self.hash_prev, -1, zdef.ZDEFLATE_WINDOW_SIZE * int.__sizeof__()) self.compress_block(data, length, 0) *out_len = self.writer.total() *out = uint8*(zdef.zdef_alloc(self.pool, *out_len)) if not *out: return -4 string.memcpy(*out, self.writer.buf, *out_len) return 0 # out same as compress: i8** (declare side and definition side must match). def flush(self, mode: int, out: bytes*, out_len: sizet*) -> int: if mode == 4: if not self.is_finished: if self.level == 0: self.writer.write_bits(1, 1) self.writer.write_bits(0, 2) self.writer.align() else: lit_tree = zhuff.zhuff_tree() lit_tree.build_fixed_lit_tree() self.writer.write_bits(1, 1) self.writer.write_bits(1, 2) lit_tree.encode_symbol(256, &self.writer) self.is_finished = 1 if self.wbits >= 0: self.writer.align() adler: uint32 = self.adler self.writer.buf[self.writer.byte_pos + 0] = uint8((adler >> 24) & 0xFF) self.writer.buf[self.writer.byte_pos + 1] = uint8((adler >> 16) & 0xFF) self.writer.buf[self.writer.byte_pos + 2] = uint8((adler >> 8) & 0xFF) self.writer.buf[self.writer.byte_pos + 3] = uint8((adler >> 0) & 0xFF) self.writer.byte_pos += 4 elif mode == 2 or mode == 3: self.writer.write_bits(0, 1) self.writer.write_bits(0, 2) self.writer.align() *out_len = self.writer.total() *out = uint8*(zdef.zdef_alloc(self.pool, *out_len)) if not *out: return -4 string.memcpy(*out, self.writer.buf, *out_len) return 0 def set_dictionary(self, dict: uint8*, length: sizet) -> int: if not dict: return -2 use: sizet = length if use > zdef.ZDEFLATE_WINDOW_SIZE: dict += length - zdef.ZDEFLATE_WINDOW_SIZE use = zdef.ZDEFLATE_WINDOW_SIZE string.memcpy(self.window, dict, use) self.window_size = int(use) self.window_pos = int(use) self.adler = zchecksum.zchecksum_adler32(dict, length, self.adler) return 0 def copy(self) -> zdeflate_stream*: z: zdeflate_stream* = zdef.zdef_alloc(self.pool, zdeflate_stream.__sizeof__()) if not z: return None string.memcpy(z, self, zdeflate_stream.__sizeof__()) z.writer.buf = BYTEPTR(zdef.zdef_alloc(self.pool, self.writer.cap)) if not z.writer.buf: zdef.zdef_free(self.pool, z) return None string.memcpy(z.writer.buf, self.writer.buf, self.writer.cap) z.window = BYTEPTR(zdef.zdef_alloc(self.pool, zdef.ZDEFLATE_WINDOW_SIZE)) z.hash_head = INTPTR(zdef.zdef_alloc(self.pool, zdef.ZDEFLATE_HASH_SIZE * int.__sizeof__())) z.hash_prev = INTPTR(zdef.zdef_alloc(self.pool, zdef.ZDEFLATE_WINDOW_SIZE * int.__sizeof__())) if not z.window or not z.hash_head or not z.hash_prev: if z.writer.buf: zdef.zdef_free(self.pool, z.writer.buf) zdef.zdef_free(self.pool, z) return None string.memcpy(z.window, self.window, zdef.ZDEFLATE_WINDOW_SIZE) string.memcpy(z.hash_head, self.hash_head, zdef.ZDEFLATE_HASH_SIZE * int.__sizeof__()) string.memcpy(z.hash_prev, self.hash_prev, zdef.ZDEFLATE_WINDOW_SIZE * int.__sizeof__()) return z def destroy(self): if self.writer.buf: zdef.zdef_free(self.pool, self.writer.buf) if self.window: zdef.zdef_free(self.pool, self.window) if self.hash_head: zdef.zdef_free(self.pool, self.hash_head) if self.hash_prev: zdef.zdef_free(self.pool, self.hash_prev) zdef.zdef_free(self.pool, self) def zdeflate_count_freqs(lit_freqs: INTPTR, dist_freqs: INTPTR, s: zdeflate_stream*, data: uint8*, length: sizet): string.memset(lit_freqs, 0, 288 * int.__sizeof__()) string.memset(dist_freqs, 0, 32 * int.__sizeof__()) pos: sizet = 0 while pos < length: best_dist: int = 0 match_len: int = 0 if s.level > 0 and pos + 2 < length: match_len = zdeflate_stream.find_match(s, data, length, pos, &best_dist) if match_len >= zdef.ZDEFLATE_MIN_MATCH: len_sym: int = zdeflate_len_to_symbol(match_len) lit_freqs[len_sym] += 1 dist_freqs[zdeflate_dist_to_symbol(best_dist)] += 1 for i in range(int(match_len)): zdeflate_stream.update_hash(s, data, pos + i, length) pos += match_len else: lit_freqs[data[pos]] += 1 zdeflate_stream.update_hash(s, data, pos, length) pos += 1 lit_freqs[256] = 1 def zdeflate_hash(p: bytes) -> uint: return (uint(p[0]) ^ (uint(p[1]) << 5) ^ (uint(p[2]) << 10)) & zdef.ZDEFLATE_HASH_MASK def zdeflate_len_to_symbol(length: int) -> int: for i in range(29): if length <= zdef.zdeflate_len_base[i] + (1 << zdef.zdeflate_len_extra_bits[i]) - 1: return 257 + i return 285 def zdeflate_dist_to_symbol(dist: int) -> int: for i in range(30): if dist <= zdef.zdeflate_dist_base[i] + (1 << zdef.zdeflate_dist_extra_bits[i]) - 1: return i return 29 def zdeflate_count_cl_freqs(all_lengths: INTPTR, total: int, cl_freqs: INTPTR): string.memset(cl_freqs, 0, 19 * int.__sizeof__()) i: int = 0 while i < total: if all_lengths[i] == 0: run: int = 0 while i + run < total and all_lengths[i + run] == 0: run += 1 while run > 0: if run >= 11: count: int = run if count > 138: count = 138 cl_freqs[18] += 1 run -= count elif run >= 3: count: int = run if count > 10: count = 10 cl_freqs[17] += 1 run -= count else: cl_freqs[0] += 1 run -= 1 while i < total and all_lengths[i] == 0: i += 1 else: cl_freqs[all_lengths[i]] += 1 val: int = all_lengths[i] i += 1 run: int = 0 while i + run < total and all_lengths[i + run] == val: run += 1 while run >= 3: count: int = run if count > 6: count = 6 cl_freqs[16] += 1 run -= count i += count # Remaining 1~2 equal-length values are not silently skipped; the outer loop counts them literally (symbol val), so the covered total is exactly total def zdeflate_write_cl_encoded(all_lengths: INTPTR, total: int, w: zdef.zbit_writer*, cl_tree: zhuff.zhuff_tree*): i: int = 0 while i < total: if all_lengths[i] == 0: run: int = 0 while i + run < total and all_lengths[i + run] == 0: run += 1 while run > 0: if run >= 11: count: int = run if count > 138: count = 138 cl_tree.encode_symbol(18, w) w.write_bits(count - 11, 7) run -= count elif run >= 3: count: int = run if count > 10: count = 10 cl_tree.encode_symbol(17, w) w.write_bits(count - 3, 3) run -= count else: cl_tree.encode_symbol(0, w) run -= 1 while i < total and all_lengths[i] == 0: i += 1 else: cl_tree.encode_symbol(all_lengths[i], w) val: int = all_lengths[i] i += 1 run: int = 0 while i + run < total and all_lengths[i + run] == val: run += 1 while run >= 3: count: int = run if count > 6: count = 6 cl_tree.encode_symbol(16, w) w.write_bits(count - 3, 2) run -= count i += count # Remaining 1~2 equal-length values are not silently skipped; the outer loop encodes them literally (symbol val), so RLE coverage totals exactly total def zdeflate_create(pool: memhub.MemManager*, level: int, wbits: int, mem_level: int, strategy: int) -> zdeflate_stream*: n: sizet = zdeflate_stream.__sizeof__() s: zdeflate_stream* = zdef.zdef_alloc(pool, n) if not s: return None string.memset(s, 0, zdeflate_stream.__sizeof__()) s.pool = pool s.writer = zdef.zbit_writer(pool) s.window = BYTEPTR(zdef.zdef_alloc(pool, zdef.ZDEFLATE_WINDOW_SIZE)) s.hash_head = INTPTR(zdef.zdef_alloc(pool, zdef.ZDEFLATE_HASH_SIZE * int.__sizeof__())) s.hash_prev = INTPTR(zdef.zdef_alloc(pool, zdef.ZDEFLATE_WINDOW_SIZE * int.__sizeof__())) if not s.window or not s.hash_head or not s.hash_prev: s.destroy() return None string.memset(s.hash_head, -1, zdef.ZDEFLATE_HASH_SIZE * int.__sizeof__()) string.memset(s.hash_prev, -1, zdef.ZDEFLATE_WINDOW_SIZE * int.__sizeof__()) s.level = level s.wbits = wbits s.strategy = strategy s.is_finished = 0 s.adler = 1 if wbits > 0: s.writer.write_zlib_header(wbits, level) elif wbits == 0: s.writer.write_zlib_header(15, level) return s def zdeflate_one_shot(pool: memhub.MemManager*, data: uint8*, length: sizet, level: int, wbits: int, out_len: sizet*) -> uint8*: s: zdeflate_stream* = zdeflate_create(pool, level, wbits, 8, 0) if not s: return None string.memset(s.hash_head, -1, zdef.ZDEFLATE_HASH_SIZE * int.__sizeof__()) string.memset(s.hash_prev, -1, zdef.ZDEFLATE_WINDOW_SIZE * int.__sizeof__()) s.adler = zchecksum.zchecksum_adler32(data, length, 1) if length == 0: s.writer.write_bits(1, 1) s.writer.write_bits(1, 2) lit_tree = zhuff.zhuff_tree() lit_tree.build_fixed_lit_tree() lit_tree.encode_symbol(256, &s.writer) else: s.compress_block(data, length, 1) s.is_finished = 1 if wbits >= 0: s.writer.align() adler: uint32 = s.adler s.writer.buf[s.writer.byte_pos + 0] = uint8((adler >> 24) & 0xFF) s.writer.buf[s.writer.byte_pos + 1] = uint8((adler >> 16) & 0xFF) s.writer.buf[s.writer.byte_pos + 2] = uint8((adler >> 8) & 0xFF) s.writer.buf[s.writer.byte_pos + 3] = uint8((adler >> 0) & 0xFF) s.writer.byte_pos += 4 *out_len = s.writer.total() result: uint8* = uint8*(zdef.zdef_alloc(pool, *out_len)) if result: string.memcpy(result, s.writer.buf, *out_len) s.destroy() return result