changeset 367:2674c789e0c3

Various optimizations
author Thibaut Girka <thib@sitedethib.com>
date Tue, 10 Jul 2012 00:25:08 +0200
parents 7dc012f631dc
children 71cd4461bb7f
files pytouhou/formats/exe.py pytouhou/utils/bitstream.pyx pytouhou/utils/lzss.pyx
diffstat 3 files changed, 24 insertions(+), 17 deletions(-) [+]
line wrap: on
line diff
--- a/pytouhou/formats/exe.py
+++ b/pytouhou/formats/exe.py
@@ -141,12 +141,11 @@ class SHT(object):
         data_va = pe_file.image_base + data_section.VirtualAddress
         data_size = data_section.SizeOfRawData
 
-        possible_character_records = list(cls.find_character_defs(pe_file))
-        if not possible_character_records:
+        try:
+            character_records_va = next(cls.find_character_defs(pe_file))
+        except StopIteration:
             raise InvalidExeException
 
-        character_records_va = possible_character_records[0]
-
         characters = []
         shots_offsets = {}
         for character in xrange(4):
--- a/pytouhou/utils/bitstream.pyx
+++ b/pytouhou/utils/bitstream.pyx
@@ -13,9 +13,10 @@
 ##
 
 cdef class BitStream:
-    cdef public io
-    cdef public int bits
-    cdef public unsigned char byte
+    cdef public object io
+    cdef unsigned int bits
+    cdef unsigned char byte
+    cdef bytes bytes
 
 
     def __init__(BitStream self, io):
@@ -48,19 +49,27 @@ cdef class BitStream:
 
     cpdef unsigned char read_bit(BitStream self):
         if not self.bits:
-            self.byte = ord(self.io.read(1))
+            self.bytes = self.io.read(1)
+            self.byte = (<unsigned char*> self.bytes)[0]
             self.bits = 8
         self.bits -= 1
         return (self.byte >> self.bits) & 0x01
 
 
-    cpdef unsigned int read(BitStream self, int nb_bits):
-        cdef unsigned int value = 0
-        cdef int i
+    cpdef unsigned int read(BitStream self, unsigned int nb_bits):
+        cdef unsigned int value = 0, read = 0
+        cdef unsigned int nb_bits2 = nb_bits
 
-        for i in range(nb_bits - 1, -1, -1):
-            value |= self.read_bit() << i
-        return value
+        while nb_bits2:
+            if not self.bits:
+                self.bytes = self.io.read(1)
+                self.byte = (<unsigned char*> self.bytes)[0]
+                self.bits = 8
+            read = self.bits if nb_bits2 > self.bits else nb_bits2
+            nb_bits2 -= read
+            self.bits -= read
+            value |= (self.byte >> self.bits) << nb_bits2
+        return value & ((1 << nb_bits) - 1)
 
 
     cpdef write_bit(BitStream self, bit):
--- a/pytouhou/utils/lzss.pyx
+++ b/pytouhou/utils/lzss.pyx
@@ -22,7 +22,7 @@ cpdef bytes decompress(object bitstream,
                        unsigned int length_size=4,
                        unsigned int minimum_match_length=3):
     cdef unsigned int i, ptr, dictionary_head, offset, length
-    cdef unsigned char flag, byte, *out_data, *dictionary
+    cdef unsigned char byte, *out_data, *dictionary
     cdef bytes _out_data
 
     out_data = <unsigned char*> malloc(size)
@@ -30,8 +30,7 @@ cpdef bytes decompress(object bitstream,
     dictionary_head, ptr = 1, 0
 
     while ptr < size:
-        flag = bitstream.read_bit()
-        if flag:
+        if bitstream.read_bit():
             # The `flag` bit is set, indicating the upcoming chunk of data is a literal
             # Add it to the uncompressed file, and store it in the dictionary
             byte = bitstream.read(8)