2
\$\begingroup\$

This is an implementation of rotation (cyclic shift) of many-byte groups to the left. Rotate right is supported as negative amount. I would like to request, if someone would be willing, to review this code for correctness and performance. Both Python 2 and 3 need to be supported.

This code is part of Kaitai Struct framework. Current implementation can be found here. Code below is in PR #28, unmerged.

There is a test case for this feature here.

Support code:

if PY2:
    range = xrange
    def integers2bytes(ints):
        return bytes(bytearray(ints))
    def bytes2integers(data):
        return bytearray(data)
else:
    def integers2bytes(ints):
        return bytes(ints)
    def bytes2integers(data):
        return data

Main code:

class KaitaiStream(object):

    # formula taken from: http://stackoverflow.com/a/812039
    precomputed_single_rotations = {amount: [(i << amount) & 0xff | (i >> (8-amount)) for i in range(256)] for amount in range(1,8)}

    @staticmethod
    def process_rotate_left(data, amount, group_size):
        amount = amount % (group_size * 8)
        if amount == 0:
            return data
        data_ints = bytes2integers(data)

        if group_size == 1:
            translate = KaitaiStream.precomputed_single_rotations[amount]
            return integers2bytes(translate[a] for a in data_ints)

        if len(data) % group_size != 0:
            raise Exception("data length must be a multiple of group size")

        if amount % 8 == 0:
            indices = [(i + amount//8) % group_size for i in range(group_size)]
            return integers2bytes(data_ints[i+k] for i in range(0,len(data),group_size) for k in indices)

        amount1 = amount % 8
        amount2 = 8 - amount1
        indices_pairs = [ ((i+amount//8) % group_size, (i+1+amount//8) % group_size) for i in range(group_size)]
        return integers2bytes((data_ints[i+k1] << amount1) & 0xff | (data_ints[i+k2] >> amount2) for i in range(0,len(data),group_size) for k1,k2 in indices_pairs)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.