Skip to content

Commit ebc0925

Browse files
committed
process rotate: precomputed translations, comprehension instead of a loop
1 parent a2f0c5a commit ebc0925

File tree

1 file changed

+12
-11
lines changed

1 file changed

+12
-11
lines changed

kaitaistruct.py

+12-11
Original file line numberDiff line numberDiff line change
@@ -356,18 +356,19 @@ def process_xor_many(data, key):
356356
else:
357357
return bytes(a ^ b for a, b in zip(data, itertools.cycle(key)))
358358

359+
# formula taken from: http://stackoverflow.com/a/812039
360+
precomputed_rotations = {amount:[(i << amount) & 0xff | (i >> (-amount & 7)) for i in range(256)] for amount in range(8)}
361+
359362
@staticmethod
360363
def process_rotate_left(data, amount, group_size):
361364
if group_size != 1:
362-
raise Exception(
363-
"unable to rotate group of %d bytes yet" %
364-
(group_size,)
365-
)
366-
367-
mask = group_size * 8 - 1
368-
anti_amount = -amount & mask
365+
raise Exception("unable to rotate groups other than 1 byte")
366+
amount = amount % 8
367+
if amount == 0:
368+
return data
369369

370-
r = bytearray(data)
371-
for i in range(len(r)):
372-
r[i] = (r[i] << amount) & 0xff | (r[i] >> anti_amount)
373-
return bytes(r)
370+
translate = KaitaiStream.precomputed_rotations[amount]
371+
if PY2:
372+
return bytes(bytearray(translate[a] for a in bytearray(data)))
373+
else:
374+
return bytes(translate[a] for a in data)

0 commit comments

Comments
 (0)