Bit-twiddling example of gem-drop logic

gemdrop.py

import numpy as np

def index_of_highest_bit(x):
    p = 1
    i = 0
    while p<x:
        p <<= 1
        i += 1
    return i

class GemBoard(object):
    def __init__(self, columns, rows, colors):
        if not 1<=rows<=32:
            raise ValueError('Rows must be in range 1 to 32')
        if not 1<=colors<=26:
            raise ValueError('Colors must be in range 1 to 26')
        self.holes = np.zeros((columns,1), dtype='u4')
        self.gems = np.zeros((columns,colors), dtype='u4')
        self.loopcount = index_of_highest_bit(rows+1)
        self.rows = rows
        self.columns = columns
        self.colors = colors

    def xor_smear(self, x):
        # Make series of 1s and 0s such that every
        # 1 in x causes a toggle in state.
        # 0b1000100100010111011 ->
        # 0b1111000111100101101
        # Assumes no more than 'rows' bits.
        p = 1
        for i in range(self.loopcount):
            x = x ^ (x>>p)
            p <<= 1
        return x

    def drop(self):
        # Drop all gems vertically to fill up holes.
        holes = self.holes
        gemdata = np.hstack([~holes, self.gems])
        # First time round the loop, select gems to drop by
        # 1 cell. Second time round, those to drop by 2 cells.
        # Third time round, 4 cells. It keeps doubling. The
        # selection masks carefully select cells so that by
        # the end every cell has been shifted by the correct
        # number of cells. For example, if a cell has exactly
        # six empty cells somewhere in the column below, it
        # should be selected for the size-2 drop, and then its
        # new position should be selected for the size-4 drop.
        for i in range(self.loopcount):
            # Find cells that need to be dropped by distance 1<<i.
            smear = self.xor_smear(holes)
            # Drop those cells.
            gemdata = gemdata & ~smear | (gemdata & smear) << (1<<i)
            # Select every second hole
            holes = holes & ~smear
        # We inverted holes at the start and uninvert it now.
        # That's because we want new holes to appear at the
        # top of the board as stuff falls.
        self.holes[:,0] = ~gemdata[:,0]
        self.gems[:,:] = gemdata[:,1:]

    def show(self):
        colors = np.array(list('-ABCDEFGHIJKLMNOPQRSTUVWXYZ'), dtype='S1')
        gemdata = np.hstack([self.holes,self.gems])
        row_dtype = 'S' + str(self.columns)
        result = []
        for row in range(self.rows):
            result.append(
                    colors[(gemdata & (1<<row)).argmax(1)]
                    .view(row_dtype)[0]
                    .decode('ascii'))
        return result

    def update(self, data):
        gemdata = np.zeros((self.columns,self.colors+1), dtype='u4')
        row_dtype = 'S' + str(self.columns)
        row_text = np.zeros(self.columns, dtype='S1')
        for i, line in enumerate(data):
            row_text.view(row_dtype)[:]=[line.replace('-','@').encode('ascii')]
            gemdata[np.arange(self.columns),row_text.view('u1')-64] |= 1 << i
        self.holes[:,0] = gemdata[:,0]
        self.gems[:,:] = gemdata[:,1:]

def example():
    board = GemBoard(10,7,3)
    board.update([
        'AAA----BBB',
        '-B--CC----',
        'AAA-CAA---',
        '-CCCC---C-',
        '--ACC---B-',
        'CABB------',
        '---B--CB--'])
    assert board.show() == [
        'AAA----BBB',
        '-B--CC----',
        'AAA-CAA---',
        '-CCCC---C-',
        '--ACC---B-',
        'CABB------',
        '---B--CB--']
    board.drop()
    for line in board.show():
        print(line)
    assert board.show() == [
        '----------',
        '----------',
        '-AA-------',
        '-BACC-----',
        'AACCC---B-',
        'ACABCCABC-',
        'CABBCACBBB']