353 lines
13 KiB
Plaintext
353 lines
13 KiB
Plaintext
import zlib.zdef as zdef
|
|
import string
|
|
import memhub
|
|
|
|
|
|
define ZHUFF_MAX_CODES = 288
|
|
define ZHUFF_MAX_BITS = 15
|
|
|
|
class zhuff_code:
|
|
code: uint
|
|
bits: int
|
|
|
|
class zhuff_tree:
|
|
pool: memhub.MemManager*
|
|
codes: [ZHUFF_MAX_CODES]zhuff_code
|
|
count: int
|
|
max_bits: int
|
|
|
|
def __init__(self):
|
|
self.pool = None
|
|
|
|
def build_codes(self, freqs: int*, count: int, max_bits: int):
|
|
lengths: [ZHUFF_MAX_CODES]int
|
|
bl_count: [ZHUFF_MAX_BITS + 1]int
|
|
next_code: [ZHUFF_MAX_BITS + 1]uint
|
|
self.count = count
|
|
self.max_bits = max_bits
|
|
self.build_code_lengths(lengths, freqs, count, max_bits)
|
|
string.memset(bl_count, 0, bl_count.__sizeof__())
|
|
for i in range(int(count)):
|
|
if lengths[i] > 0:
|
|
bl_count[lengths[i]] += 1
|
|
code: uint = 0
|
|
next_code[0] = 0
|
|
for bits in range(1, int(max_bits + 1)):
|
|
code = (code + bl_count[bits - 1]) << 1
|
|
next_code[bits] = code
|
|
for i in range(int(count)):
|
|
self.codes[i].bits = lengths[i]
|
|
if lengths[i] > 0:
|
|
self.codes[i].code = next_code[lengths[i]]
|
|
next_code[lengths[i]] += 1
|
|
else:
|
|
self.codes[i].code = 0
|
|
|
|
def build_fixed_lit_tree(self):
|
|
lengths: [288]int
|
|
bl_count: [16]int
|
|
next_code: [16]uint
|
|
self.get_fixed_lit_lengths(lengths)
|
|
string.memset(bl_count, 0, bl_count.__sizeof__())
|
|
for i in range(288):
|
|
if lengths[i] > 0:
|
|
bl_count[lengths[i]] += 1
|
|
code: uint = 0
|
|
next_code[0] = 0
|
|
for bits in range(1, int(9 + 1)):
|
|
code = (code + bl_count[bits - 1]) << 1
|
|
next_code[bits] = code
|
|
self.count = 288
|
|
self.max_bits = 9
|
|
for i in range(288):
|
|
self.codes[i].bits = lengths[i]
|
|
if lengths[i] > 0:
|
|
self.codes[i].code = next_code[lengths[i]]
|
|
next_code[lengths[i]] += 1
|
|
else:
|
|
self.codes[i].code = 0
|
|
|
|
def build_fixed_dist_tree(self):
|
|
lengths: [32]int
|
|
bl_count: [16]int
|
|
next_code: [16]uint
|
|
self.get_fixed_dist_lengths(lengths)
|
|
string.memset(bl_count, 0, bl_count.__sizeof__())
|
|
for i in range(32):
|
|
if lengths[i] > 0:
|
|
bl_count[lengths[i]] += 1
|
|
code: uint = 0
|
|
next_code[0] = 0
|
|
for bits in range(1, int(5 + 1)):
|
|
code = (code + bl_count[bits - 1]) << 1
|
|
next_code[bits] = code
|
|
self.count = 32
|
|
self.max_bits = 5
|
|
for i in range(32):
|
|
self.codes[i].bits = lengths[i]
|
|
if lengths[i] > 0:
|
|
self.codes[i].code = next_code[lengths[i]]
|
|
next_code[lengths[i]] += 1
|
|
else:
|
|
self.codes[i].code = 0
|
|
|
|
def encode_symbol(self, symbol: int, writer: zdef.zbit_writer*):
|
|
if symbol < 0 or symbol >= self.count:
|
|
return
|
|
if self.codes[symbol].bits == 0:
|
|
return
|
|
writer.write_bits_rev(self.codes[symbol].code, self.codes[symbol].bits)
|
|
|
|
|
|
def build_code_lengths(self, lengths: int*, freqs: int*, count: int, max_bits: int):
|
|
bl_count: [ZHUFF_MAX_BITS + 2]int
|
|
sort_count: int = 0
|
|
string.memset(bl_count, 0, bl_count.__sizeof__())
|
|
string.memset(lengths, 0, int.__sizeof__() * count)
|
|
for i in range(int(count)):
|
|
if freqs[i] > 0:
|
|
sort_count += 1
|
|
if sort_count == 0: return
|
|
if sort_count == 1:
|
|
for i in range(int(count)):
|
|
if freqs[i] > 0:
|
|
lengths[i] = 1
|
|
return
|
|
items: int* = int*(zdef.zdef_alloc(self.pool, sort_count * int.__sizeof__()))
|
|
item_freqs: int* = int*(zdef.zdef_alloc(self.pool, sort_count * int.__sizeof__()))
|
|
idx: int = 0
|
|
for i in range(int(count)):
|
|
if freqs[i] > 0:
|
|
items[idx] = i
|
|
item_freqs[idx] = freqs[i]
|
|
idx += 1
|
|
for i in range(int(sort_count)):
|
|
key_item: int = items[i]
|
|
key_freq: int = item_freqs[i]
|
|
j: int = i - 1
|
|
while j >= 0 and item_freqs[j] > key_freq:
|
|
items[j + 1] = items[j]
|
|
item_freqs[j + 1] = item_freqs[j]
|
|
j -= 1
|
|
items[j + 1] = key_item
|
|
item_freqs[j + 1] = key_freq
|
|
parent: int* = int*(zdef.zdef_alloc(self.pool, (sort_count * 2) * int.__sizeof__()))
|
|
for i in range(int(sort_count * 2)):
|
|
parent[i] = -1
|
|
heap: int* = int*(zdef.zdef_alloc(self.pool, (sort_count + 1) * int.__sizeof__()))
|
|
heap_size: int = 0
|
|
for i in range(int(sort_count)):
|
|
heap[heap_size] = i
|
|
heap_size += 1
|
|
pos: int = heap_size - 1
|
|
while pos > 0:
|
|
par: int = (pos - 1) / 2
|
|
if item_freqs[heap[par]] > item_freqs[heap[pos]]:
|
|
tmp: int = heap[par]
|
|
heap[par] = heap[pos]
|
|
heap[pos] = tmp
|
|
pos = par
|
|
else: break
|
|
combined_freq: int* = int*(zdef.zdef_alloc(self.pool, (sort_count * 2) * int.__sizeof__()))
|
|
for i in range(int(sort_count)):
|
|
combined_freq[i] = item_freqs[i]
|
|
internal: int = sort_count
|
|
while heap_size > 1:
|
|
a: int = heap[0]
|
|
heap[0] = heap[heap_size - 1]
|
|
heap_size -= 1
|
|
pos: int = 0
|
|
while True:
|
|
left: int = 2 * pos + 1
|
|
right: int = 2 * pos + 2
|
|
smallest: int = pos
|
|
if left < heap_size and combined_freq[heap[left]] < combined_freq[heap[smallest]]:
|
|
smallest = left
|
|
if right < heap_size and combined_freq[heap[right]] < combined_freq[heap[smallest]]:
|
|
smallest = right
|
|
if smallest != pos:
|
|
tmp: int = heap[pos]
|
|
heap[pos] = heap[smallest]
|
|
heap[smallest] = tmp
|
|
pos = smallest
|
|
else: break
|
|
b: int = heap[0]
|
|
heap[0] = heap[heap_size - 1]
|
|
heap_size -= 1
|
|
pos = 0
|
|
while True:
|
|
left: int = 2 * pos + 1
|
|
right: int = 2 * pos + 2
|
|
smallest: int = pos
|
|
if left < heap_size and combined_freq[heap[left]] < combined_freq[heap[smallest]]:
|
|
smallest = left
|
|
if right < heap_size and combined_freq[heap[right]] < combined_freq[heap[smallest]]:
|
|
smallest = right
|
|
if smallest != pos:
|
|
tmp: int = heap[pos]
|
|
heap[pos] = heap[smallest]
|
|
heap[smallest] = tmp
|
|
pos = smallest
|
|
else: break
|
|
if internal >= sort_count * 2 - 1: break
|
|
combined_freq[internal] = combined_freq[a] + combined_freq[b]
|
|
parent[a] = internal
|
|
parent[b] = internal
|
|
heap[heap_size] = internal
|
|
heap_size += 1
|
|
pos = heap_size - 1
|
|
while pos > 0:
|
|
par: int = (pos - 1) / 2
|
|
if combined_freq[heap[par]] > combined_freq[heap[pos]]:
|
|
tmp: int = heap[par]
|
|
heap[par] = heap[pos]
|
|
heap[pos] = tmp
|
|
pos = par
|
|
else: break
|
|
internal += 1
|
|
for i in range(int(sort_count)):
|
|
node: int = i
|
|
length: int = 0
|
|
while parent[node] >= 0:
|
|
length += 1
|
|
node = parent[node]
|
|
lengths[items[i]] = length
|
|
string.memset(bl_count, 0, bl_count.__sizeof__())
|
|
for i in range(int(count)):
|
|
if lengths[i] > 0:
|
|
bl_count[lengths[i]] += 1
|
|
overflow: int = 0
|
|
for i in range(int(count)):
|
|
if lengths[i] > max_bits:
|
|
overflow += 1
|
|
bl_count[lengths[i]] -= 1
|
|
lengths[i] = max_bits
|
|
bl_count[max_bits] += 1
|
|
while overflow > 0:
|
|
bits: int = max_bits - 1
|
|
while bits > 0 and bl_count[bits] == 0:
|
|
bits -= 1
|
|
if bits == 0: break
|
|
bl_count[bits] -= 1
|
|
bl_count[bits + 1] += 2
|
|
bl_count[max_bits] -= 1
|
|
overflow -= 2
|
|
sorted_items: int* = int*(zdef.zdef_alloc(self.pool, sort_count * int.__sizeof__()))
|
|
si: int = 0
|
|
for i in range(int(count)):
|
|
if freqs[i] > 0:
|
|
sorted_items[si] = i
|
|
si += 1
|
|
for i in range(int(si - 1)):
|
|
for j in range(int(i + 1), int(si)):
|
|
if lengths[sorted_items[i]] < lengths[sorted_items[j]]:
|
|
tmp: int = sorted_items[i]
|
|
sorted_items[i] = sorted_items[j]
|
|
sorted_items[j] = tmp
|
|
sidx: int = 0
|
|
for bits in range(int(max_bits), 1, -1):
|
|
n: int = bl_count[bits]
|
|
while n > 0 and sidx < si:
|
|
lengths[sorted_items[sidx]] = bits
|
|
sidx += 1
|
|
n -= 1
|
|
zdef.zdef_free(self.pool, sorted_items)
|
|
zdef.zdef_free(self.pool, heap)
|
|
zdef.zdef_free(self.pool, parent)
|
|
zdef.zdef_free(self.pool, combined_freq)
|
|
zdef.zdef_free(self.pool, item_freqs)
|
|
zdef.zdef_free(self.pool, items)
|
|
|
|
def get_fixed_lit_lengths(self, lengths: int*):
|
|
for i in range( 0, int(143 + 1)): lengths[i] = 8
|
|
for i in range(int(143 + 1), int(255 + 1)): lengths[i] = 9
|
|
for i in range(int(255 + 1), int(279 + 1)): lengths[i] = 7
|
|
for i in range(int(279 + 1), int(287 + 1)): lengths[i] = 8
|
|
|
|
def get_fixed_dist_lengths(self, lengths: int*):
|
|
for i in range(32):
|
|
lengths[i] = 5
|
|
|
|
def build_tree_from_lengths(self, lengths: int*, count: int, max_bits: int):
|
|
bl_count: [16]int
|
|
next_code: [16]uint
|
|
|
|
string.memset(bl_count, 0, bl_count.__sizeof__())
|
|
for i in range(int(count)):
|
|
if lengths[i] > 0:
|
|
bl_count[lengths[i]] += 1
|
|
|
|
code: uint = 0
|
|
next_code[0] = 0
|
|
for bits in range(1, int(max_bits + 1)):
|
|
code = (code + bl_count[bits - 1]) << 1
|
|
next_code[bits] = code
|
|
|
|
self.count = count
|
|
self.max_bits = max_bits
|
|
for i in range(int(count)):
|
|
self.codes[i].bits = lengths[i]
|
|
if lengths[i] > 0:
|
|
self.codes[i].code = next_code[lengths[i]]
|
|
next_code[lengths[i]] += 1
|
|
else:
|
|
self.codes[i].code = 0
|
|
|
|
|
|
class zhuff_decode_node:
|
|
children: [2]int
|
|
symbol: int
|
|
|
|
|
|
class zhuff_decode_tree:
|
|
pool: memhub.MemManager*
|
|
nodes: [2 * ZHUFF_MAX_CODES]zhuff_decode_node
|
|
node_count: int
|
|
root: int
|
|
|
|
def __init__(self):
|
|
self.pool = None
|
|
|
|
def build_decode_tree(self, ht: zhuff_tree*):
|
|
self.node_count = 1
|
|
self.root = 0
|
|
self.nodes[0].children[0] = -1
|
|
self.nodes[0].children[1] = -1
|
|
self.nodes[0].symbol = -1
|
|
for i in range(int(ht.count)):
|
|
if ht.codes[i].bits == 0: continue
|
|
node: int = 0
|
|
for bit in range(int(ht.codes[i].bits - 1), -1, -1):
|
|
dir: int = (ht.codes[i].code >> bit) & 1
|
|
if bit > 0:
|
|
# Internal node: traverse or create
|
|
if self.nodes[node].children[dir] == -1:
|
|
new_node: int = self.node_count
|
|
self.node_count += 1
|
|
self.nodes[new_node].children[0] = -1
|
|
self.nodes[new_node].children[1] = -1
|
|
self.nodes[new_node].symbol = -1
|
|
self.nodes[node].children[dir] = new_node
|
|
node = self.nodes[node].children[dir]
|
|
else:
|
|
# Leaf: bit 0 - set symbol on the child
|
|
if self.nodes[node].children[dir] == -1:
|
|
leaf: int = self.node_count
|
|
self.node_count += 1
|
|
self.nodes[leaf].children[0] = -1
|
|
self.nodes[leaf].children[1] = -1
|
|
self.nodes[leaf].symbol = i
|
|
self.nodes[node].children[dir] = leaf
|
|
else:
|
|
self.nodes[self.nodes[node].children[dir]].symbol = i
|
|
|
|
def decode_symbol(self, reader: zdef.zbit_reader*) -> int:
|
|
node: int = self.root
|
|
while self.nodes[node].symbol == -1:
|
|
bit: uint
|
|
if reader.read_bits(1, &bit) != 0: return -1
|
|
dir: int = int(bit)
|
|
if self.nodes[node].children[dir] == -1: return -1
|
|
node = self.nodes[node].children[dir]
|
|
return self.nodes[node].symbol
|