persist.py

import itertools


class PersistentList(object):

    BIT_WIDTH = 5  # number of bits to access each level

    def __init__(self, contents=()):
        root = self._contents = []
        tail = self._tail = []

        index, height, width = -1, 1, self.WIDTH
        for index, element in enumerate(contents):
            tail.append(element)

            if len(tail) == width:
                if len(root) == width:
                    height += 1
                    root = self._contents = [root]
                root.append(tail)
                tail = self._tail = []

        self._height = height + 1 if root else height
        self._size = index + 1

    def __len__(self):
        return self._size

    def __getitem__(self, i):
        contents = self._contents_with(i)
        return contents[i & self.WIDTH - 1]

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return NotImplemented
        return self._contents == other._contents and self._tail == other._tail

    def __ne__(self, other):
        return not self == other

    def __add__(self, other):
        result = self
        for element in other:
            result = result.cons(element)
        return result

    def __mul__(self, times):
        result = self
        for i in range(times - 1):
            for j in self:
                result = result.cons(j)
        return result

    def __reversed__(self):
        for contents in reversed(self._contents):
            for element in reversed(contents):
                yield element

    def __repr__(self):
        return "{0.__class__.__name__}({1})".format(self, list(self))

    def _contents_with(self, i):
        if i >= self._tail_offset:
            return self._tail
        contents = self._contents
        mask = self._height * self.BIT_WIDTH
        for j in reversed(range(self._height)):
            contents = contents[(i >> mask) & self.BIT_WIDTH]
        return contents

    @property
    def WIDTH(self):
        return 2 ** self.BIT_WIDTH

    @property
    def _tail_offset(self):
        return len(self) >> self.BIT_WIDTH << self.BIT_WIDTH

    def cons(self, i):
        new = self.copy()

        if len(self._tail) == self.WIDTH:
            if len(self._contents) == self.WIDTH:
                new._height += 1
                new._contents = [new._contents, [new._tail]]
            else:
                new._contents.append(new._tail)

            new._tail = []

        new._tail.append(i)
        new._size += 1
        return new

    def copy(self):
        new = self.__class__()
        new._contents = self._contents
        new._tail = list(self._tail)
        new._height = self._height
        new._size = self._size
        return new

    def insert(self, i, j):
        new = self.copy()
        new._contents.insert(i, j)
        return new

    def remove(self, i):
        new = self.copy()
        new._contents.remove(i)
        return new

    def replace(self, i, j):
        contents = list(self)
        contents[i] = j
        return self.__class__(contents)

    def without(self, i):
        if isinstance(i, slice):
            raise TypeError("Can't remove slices. Use + instead.")

        contents = self[:i] + self[i + 1:]
        return self.__class__(contents)

    def index(self, i):
        return self._contents.index(i)

    def count(self, i):
        return self._contents.count(i)