A series of performance improvements for parsing.

Author: GraylinKim

sc2reader/readers.py

@@ -232,21 +232,24 @@ def __call__(self, data, replay):
         data_length = data.length
         read_timestamp =  data.read_timestamp
         read_bits = data.read_bits
+        read_byte = data.read_byte
         read_bytes = data.read_bytes
         byte_align = data.byte_align
         append = game_events.append
         event_start = 0
 
         try:
             while event_start != data_length:
-                event = None
                 fstamp += read_timestamp()
                 pid = read_bits(5)
                 event_type = read_bits(7)
 
                 # Check for a lookup
                 if event_type in EVENT_DISPATCH:
                     event = EVENT_DISPATCH[event_type](data, fstamp, pid, event_type)
+                    if debug:
+                        event.bytes = data.read_range(event_start, data.tell())
+                    append(event)
 
                 # Otherwise maybe it is an unknown chunk
                 elif event_type == 0x26:
@@ -272,12 +275,6 @@ def __call__(self, data, replay):
                     raise ReadError("Event type {} unknown at position {}.".format(hex(event_type),hex(event_start)), event_type, event_start, replay, game_events, data)
 
                 byte_align()
-
-                if event:
-                    if debug:
-                        event.bytes = data.read_range(event_start, data.tell())
-                    append(event)
-
                 event_start = data.tell()
 
             return game_events
@@ -320,7 +317,7 @@ def player_selection_event(self, data, fstamp, pid, event_type):
         unit_types = [(data.read_short(BIG_ENDIAN) << 8 | data.read_byte(),data.read_bits(self.UNIT_INDEX_BITS)) for index in range(type_count)]
 
         unit_count = data.read_bits(self.UNIT_INDEX_BITS)
-        unit_ids = [data.read_bits(32) for index in range(unit_count)]
+        unit_ids = [data.read_int(BIG_ENDIAN) for index in range(unit_count)]
 
         unit_types = chain(*[[utype]*count for (utype, count) in unit_types])
         units = list(zip(unit_ids, unit_types))
@@ -364,21 +361,22 @@ def player_request_resource_event(self, data, fstamp, pid, event_type):
 
     def camera_event(self, data, fstamp, pid, event_type):
         # From https://github.com/Mischanix/sc2replay-csharp/wiki/replay.game.events
-        x = data.read_bits(16)/256.0
-        y = data.read_bits(16)/256.0
+        block = data.read_int(BIG_ENDIAN)
+        x = (block >> 16)/256.0
+        y = (block & 0xFFFF)/256.0
         distance = pitch = yaw = height = 0
         if data.read_bits(1):
-            distance = data.read_bits(16)/256.0
+            distance = data.read_short(BIG_ENDIAN)/256.0
         if data.read_bits(1):
             #Note: this angle is relative to the horizontal plane, but the editor shows the angle relative to the vertical plane. Subtract from 90 degrees to convert.
-            pitch = data.read_bits(16) #?
+            pitch = data.read_short(BIG_ENDIAN) #?
             pitch = 45 * (((((pitch * 0x10 - 0x2000) << 17) - 1) >> 17) + 1) / 4096.0
         if data.read_bits(1):
             #Note: this angle is the vector from the camera head to the camera target projected on to the x-y plane in positive coordinates. So, default is 90 degrees, while insert and delete produce 45 and 135 degrees by default.
-            yaw = data.read_bits(16) #?
+            yaw = data.read_short(BIG_ENDIAN) #?
             yaw = 45 * (((((yaw * 0x10 - 0x2000) << 17) - 1) >> 17) + 1) / 4096.0
         if data.read_bits(1):
-            height_offset = data.read_bits(16)/256.0
+            height_offset = data.read_short(BIG_ENDIAN)/256.0
         return CameraEvent(fstamp, pid, event_type, x, y, distance, pitch, yaw, height)
 
 
@@ -421,7 +419,7 @@ def player_ability_event(self, data, fstamp, pid, event_type):
 
         default_ability = not data.read_bits(1)
         if not default_ability:
-            ability = data.read_bits(16) << 5 | data.read_bits(5)
+            ability = data.read_short(BIG_ENDIAN) << 5 | data.read_bits(5)
             default_actor = not data.read_bits(1)
         else:
             ability = 0
@@ -430,17 +428,17 @@ def player_ability_event(self, data, fstamp, pid, event_type):
         if target_type == 1:
             x = data.read_bits(20)/4096.0
             y = data.read_bits(20)/4096.0
-            z = data.read_bits(32)
+            z = data.read_int(BIG_ENDIAN)
             z = (z>>1)/8192.0 * pow(-1, z & 0x1)
             unknown = data.read_bits(1)
             return LocationAbilityEvent(fstamp, pid, event_type, ability, (x, y, z))
 
         elif target_type == 2:
             player = team = None
 
-            data.read_bits(8)
-            data.read_bits(8)
-            unit = (data.read_bits(32), data.read_bits(16))
+            data.read_byte()
+            data.read_byte()
+            unit = (data.read_int(BIG_ENDIAN), data.read_short(BIG_ENDIAN))
 
             if self.ABILITY_TEAM_FLAG and data.read_bits(1):
                 team = data.read_bits(4)
@@ -450,13 +448,13 @@ def player_ability_event(self, data, fstamp, pid, event_type):
 
             x = data.read_bits(20)/4096.0
             y = data.read_bits(20)/4096.0
-            z = data.read_bits(32)
+            z = data.read_int(BIG_ENDIAN)
             z = (z>>1)/8192.0 * pow(-1, z & 0x1)
             unknown = data.read_bits(1)
             return TargetAbilityEvent(fstamp, pid, event_type, ability, unit, player, team, (x, y, z))
 
         elif target_type == 3:
-            unit_id = data.read_bits(32)
+            unit_id = data.read_int(BIG_ENDIAN)
             unknown = data.read_bits(1)
             return SelfAbilityEvent(fstamp, pid, event_type, ability, unit_id)
 

sc2reader/utils.py

@@ -126,8 +126,8 @@ def peek(self, length):
         Optimized bit-shift aware read methods
     '''
     def read_byte(self):
-        if self.bytes_left() == 0:
-            raise EOFError("Cannot read byte; no bytes remaining")
+        #if self.bytes_left() == 0:
+        #    raise EOFError("Cannot read byte; no bytes remaining")
 
         if self.bit_shift==0:
             return ord(self.read(1))
@@ -139,8 +139,8 @@ def read_byte(self):
             return hi_bits | lo_bits
 
     def read_short(self, endian=LITTLE_ENDIAN):
-        if self.bytes_left() < 2:
-            raise EOFError("Cannot read short; only {} bytes left in buffer".format(self.left))
+        #if self.bytes_left() < 2:
+        #    raise EOFError("Cannot read short; only {} bytes left in buffer".format(self.left))
 
         if self.bit_shift == 0:
             return struct.unpack(endian+'H', self.read(2))[0]
@@ -156,8 +156,8 @@ def read_short(self, endian=LITTLE_ENDIAN):
             return number
 
     def read_int(self, endian=LITTLE_ENDIAN):
-        if self.bytes_left() < 4:
-            raise EOFError("Cannot read int; only {} bytes left in buffer".format(self.left))
+        #if self.bytes_left() < 4:
+        #    raise EOFError("Cannot read int; only {} bytes left in buffer".format(self.left))
 
         if self.bit_shift == 0:
             return struct.unpack(endian+'I', self.read(4))[0]
@@ -173,34 +173,60 @@ def read_int(self, endian=LITTLE_ENDIAN):
             return number
 
     def read_bits(self, bits):
-        if self.bytes_left()*8 < bits-(8-self.bit_shift):
-            raise EOFError("Cannot read {} bits. only {} bits left in buffer.".format(bits, (self.length-self.tell()+1)*8-self.bit_shift))
-
-        if self.bit_shift!=0:
-            first_bits = min(8-self.bit_shift, bits)
-            bits = bits-first_bits
-            mask = self.lo_masks[first_bits] << self.bit_shift
-            result = ((self.bit_buffer & mask) >> self.bit_shift) << bits
-            self.bit_shift = (self.bit_shift + first_bits) % 8
+        #if self.bytes_left()*8 < bits-(8-self.bit_shift):
+        #    raise EOFError("Cannot read {} bits. only {} bits left in buffer.".format(bits, (self.length-self.tell()+1)*8-self.bit_shift))
+        bit_shift = self.bit_shift
+        if bit_shift!=0:
+            bits_left = 8-bit_shift
+
+            # Read it all and continue
+            if bits_left < bits:
+                bits -= bits_left
+                result = (self.bit_buffer >> bit_shift) << bits
+
+            # Read part and return
+            elif bits_left > bits:
+                self.bit_shift+=bits
+                return (self.bit_buffer >> bit_shift) & self.lo_masks[bits]
+
+            # Read all and return
+            else:
+                self.bit_shift = 0
+                return self.bit_buffer >> bit_shift
+
         else:
             result = 0
 
-        if bits > 8:
+        if bits >= 8:
             bytes = bits/8
-            for byte in struct.unpack("B"*bytes, self.read(bytes)):
+
+            if bytes == 1:
                 bits -= 8
-                result |= byte << bits
+                result |= ord(self.read(1)) << bits
+
+            elif bytes == 2:
+                bits -= 16
+                result |= struct.unpack(">H",self.read(2))[0] << bits
+
+            elif bytes == 4:
+                bits -= 32
+                result |= struct.unpack(">I",self.read(4))[0] << bits
+
+            else:
+                for byte in struct.unpack("B"*bytes, self.read(bytes)):
+                    bits -= 8
+                    result |= byte << bits
 
         if bits != 0:
-            self.bit_shift = bits
             self.bit_buffer = ord(self.read(1))
             result |= self.bit_buffer & self.lo_masks[bits]
 
+        self.bit_shift = bits
         return result
 
     def read_bytes(self, bytes):
-        if self.bytes_left() < bytes:
-            raise EOFError("Cannot read {} bytes. only {} bytes left in buffer.".format(bytes, self.length-self.tell()))
+        #if self.bytes_left() < bytes:
+        #    raise EOFError("Cannot read {} bytes. only {} bytes left in buffer.".format(bytes, self.length-self.tell()))
 
         if self.bit_shift==0:
             return self.read(bytes)
@@ -251,19 +277,16 @@ def read_timestamp(self):
         The least significant 2 bits of the first byte specify how many extra
         bytes the timestamp has.
         """
-        return self.read_bits(self.read_bits(2)*8+6)
-        """faster???
         first = self.read_byte()
         time,count = first >> 2, first & 0x03
         if count == 0:
             return time
         elif count == 1:
             return time << 8 | self.read_byte()
         elif count == 2:
-            return time << 16 | self.read_short()
+            return time << 16 | self.read_short(BIG_ENDIAN)
         elif count == 3:
-            return time << 24 | self.read_short() << 8 | self.read_byte()
-        """
+            return time << 24 | self.read_short(BIG_ENDIAN) << 8 | self.read_byte()
 
     def read_data_struct(self, indent=0, key=None):
         """