fixed problem with compatibility on Arch Linux
[unME11.git] / HuffDec11.py
1 import os, sys, struct, zlib
2
3 class Error(Exception): pass
4
5 def cwDec(w): # Convert 16-bit value to string codeword
6   return bin(0x10000 | w).rstrip('0')[3:-1]
7
8 def cwEnc(cw): # Convert string codeword to 16-bit value
9   return int((cw+'1').ljust(16, '0'), 2)
10
11 #***************************************************************************
12 #***************************************************************************
13 #***************************************************************************
14
15 def HuffTabReader_bin(ab):
16   fmtRec = struct.Struct("<HB")
17   o = 0
18   while o < len(ab):
19     w, cb = fmtRec.unpack_from(ab, o)
20     o += fmtRec.size
21     v = ab[o:o+cb]
22     assert len(v) == cb
23     o += cb
24     yield(cwDec(w), cb, v)
25
26 #***************************************************************************
27 #***************************************************************************
28 #***************************************************************************
29
30 class HuffNode(object):
31   def __init__(self, cw, hd):
32     self.cw = cw # String codeword value
33     self.w = cwEnc(cw) # Encoded codeword value
34     if hd:
35       self.nBits = len(cw) # Length of codeword in bits
36       self.cb = hd.dLen.get(cw, None)
37       self.av = [d.get(cw, None) for d in hd.adTab]
38     else:
39       self.nBits = None # Actual length of codeword is unknown
40
41 #***************************************************************************
42 #***************************************************************************
43 #***************************************************************************
44
45 class HuffDecoder(object):
46   NAMES = ("Code", "Data")
47   DUMP_KNOWN = 0
48   DUMP_LEN = 1
49   DUMP_ALL = 2
50   fmtInt = struct.Struct("<L")
51   baseDir = os.path.split(__file__)[0]
52   BLOCK_SIZE = 0x1000 # 4K bytes
53
54   def __init__(self):
55     with open(os.path.join(self.baseDir, "huff11.bin"), "rb") as fi: self.unpackTables(zlib.decompress(fi.read(), -15)) # Load from compressed version
56     self.prepareMap()
57
58   def loadTable(self, items):
59     sv = set() # Set for values
60     d = {}
61     for cw, cb, v in items:
62       if cw in d: raise Error("Codeword %s already defined" % cw)
63
64       if cb is None: continue
65       cbKnown = self.dLen.get(cw, None)
66       if cbKnown is None: self.dLen[cw] = cb
67       elif cb != cbKnown: raise Error("Codeword %s sequence length %d != know %d" % (cw, cb, cbKnown))
68
69       if v is None: continue
70       assert len(v) == cb
71       d[cw] = v # Remember value
72       sv.add(v)
73
74     self.adTab.append(d)
75
76   def unpackTables(self, ab):
77     n, = self.fmtInt.unpack_from(ab)
78     o = self.fmtInt.size
79     self.dLen, self.adTab = {}, []
80     for i in xrange(n):
81       cb, = self.fmtInt.unpack_from(ab, o)
82       o += self.fmtInt.size
83       data = ab[o:o+cb]
84       assert len(data) == cb
85       o += cb
86       self.loadTable(HuffTabReader_bin(data))
87
88   def propagateMap(self, node):
89     cw = node.cw
90     for idx in xrange(int(cw[::-1], 2), len(self.aMap), 1<<len(cw)):
91       assert self.aMap[idx] is None
92       self.aMap[idx] = node
93
94   def prepareMap(self):
95     aCW = sorted(self.dLen.keys())[::-1]
96     minBits, maxBits = len(aCW[0]), len(aCW[-1])
97     self.aMap = [None]*(1<<maxBits) # 2**maxBits map
98     aCW.append('0'*(maxBits+1)) # Longer than max
99     nBits = minBits # Current length
100     e = int(aCW[0], 2)|1 # End value for current length
101     for o in xrange(1, len(aCW)):
102       nextBits = len(aCW[o])
103       if nextBits == nBits: continue # Run until length change
104       assert nextBits > nBits # Length must increase
105       s = int(aCW[o-1], 2) # Start value for current length
106       for i in xrange(s, e+1):
107         cw = bin(i)[2:].zfill(nBits)
108         self.propagateMap(HuffNode(cw, self))
109       e = int(aCW[o], 2)|1 # End value for next length
110       for i in xrange(e/2 + 1, s): # Handle values with unknown codeword length
111         cw = bin(i)[2:].zfill(nBits)
112         self.propagateMap(HuffNode(cw, None))
113       nBits = nextBits
114     for v in self.aMap: assert v is not None
115
116   def enumCW(self, ab):
117     v = int(bin(int("01"+ab.encode("hex"), 16))[3:][::-1], 2) # Reversed bits
118     cb = 0
119     while cb < self.BLOCK_SIZE: # Block length
120       node = self.aMap[v & 0x7FFF]
121       if node.nBits is None: raise Error("Unknown codeword %s* length" % node.cw)
122       yield node
123       v >>= node.nBits
124       if node.cb is not None: cb += node.cb
125
126   def decompressChunk(self, ab, iTab):
127     r = []
128     cb = 0
129     for node in self.enumCW(ab):
130       v = node.av[iTab]
131       if v is None: raise Error("Unknown sequence for codeword %s in table #%d" % (node.cw, iTab))
132       r.append(v)
133       cb += len(v)
134       if cb >= self.BLOCK_SIZE: break
135     return "".join(r)
136
137   def decompress(self, ab, length):
138     nChunks, left = divmod(length, self.BLOCK_SIZE)
139     assert 0 == left
140     aOfs = list(struct.unpack_from("<%dL" % nChunks, ab))
141     aOpt = [0]*nChunks
142     for i in xrange(nChunks):
143       aOpt[i], aOfs[i] = divmod(aOfs[i], 0x40000000)
144
145     base = nChunks*4
146     aOfs.append(len(ab) - base)
147     r = []
148     for i, opt in enumerate(aOpt):
149       iTab, bCompr = divmod(opt, 2)
150       assert 1 == bCompr
151       unpacked = self.decompressChunk(ab[base + aOfs[i]: base + aOfs[i+1]], iTab)
152       assert len(unpacked) == self.BLOCK_SIZE
153       r.append(unpacked)
154     return "".join(r)
155