/brz/remove-bazaar

# Copyright (C) 2008 Canonical Ltd

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

"""Persistent maps from string->string using CHK stores."""

import osutils

class CHKMap(object):

    """A persistent map from string to string backed by a CHK store."""

    def __init__(self, store, root_key):

        """Create a CHKMap object.

        :param store: The store the CHKMap is stored in.

        :param root_key: The root key of the map. None to create an empty

            CHKMap.

        """

        self._store = store

        if root_key is None:

            self._root_node = RootNode()

        else:

            self._root_node = root_key

    def apply_delta(self, delta):

        """Apply a delta to the map.

        :param delta: An iterable of old_key, new_key, new_value tuples.

            If new_key is not None, then new_key->new_value is inserted

            into the map; if old_key is not None, then the old mapping

            of old_key is removed.

        """

        for old, new, value in delta:

            if old is not None and old != new:

                # unmap

                self._unmap(old)

        for old, new, value in delta:

            if new is not None:

                # map

                self._map(new, value)

        return self._save()

    def _ensure_root(self):

        """Ensure that the root node is an object not a key."""

        if type(self._root_node) == tuple:

            # Demand-load the root

            bytes = self._read_bytes(self._root_node)

            root_key = self._root_node

            self._root_node = RootNode()

            self._root_node.deserialise(bytes, root_key)

    def _read_bytes(self, key):

        stream = self._store.get_record_stream([key], 'unordered', True)

        return stream.next().get_bytes_as('fulltext')

    @classmethod

    def from_dict(klass, store, initial_value, maximum_size=0):

        """Create a CHKMap in store with initial_value as the content.

        :param store: The store to record initial_value in, a VersionedFiles

            object with 1-tuple keys supporting CHK key generation.

        :param initial_value: A dict to store in store. Its keys and values

            must be bytestrings.

        :param maximum_size: The maximum_size rule to apply to nodes. This

            determines the size at which no new data is added to a single node.

        :return: The root chk of te resulting CHKMap.

        """

        result = CHKMap(store, None)

        result._root_node.set_maximum_size(maximum_size)

        delta = []

        for key, value in initial_value.items():

            delta.append((None, key, value))

        result.apply_delta(delta)

        return result._save()

    def iteritems(self, key_filter=None):

        """Iterate over the entire CHKMap's contents."""

        self._ensure_root()

        if key_filter is not None:

            for name, key, in self._root_node._nodes.iteritems():

                if name in key_filter:

                    bytes = self._read_bytes(key)

                    yield name, ValueNode.deserialise(bytes).value

        else:

            for name, key, in self._root_node._nodes.iteritems():

                bytes = self._read_bytes(key)

                yield name, ValueNode.deserialise(bytes).value

    def key(self):

        """Return the key for this map."""

        if isinstance(self._root_node, tuple):

            return self._root_node

        else:

            return self._root_node._key

    def __len__(self):

        self._ensure_root()

        return len(self._root_node)

    def _map(self, key, value):

        """Map key to value."""

        self._ensure_root()

        # Store the value

        bytes = ValueNode(value).serialise()

        # Experimental code to probe for keys rather than just adding; its not

        # clear if it is an improvement.

        #chk = ("sha1:%s" % osutils.sha_string(bytes),)

        #if not self._store.get_parent_map([key]):

        sha1, _, _ = self._store.add_lines((None,), (), osutils.split_lines(bytes))

        chk = ("sha1:" + sha1,)

        # And link into the root

        self._root_node.add_child(key, chk)

    def _unmap(self, key):

        """remove key from the map."""

        self._ensure_root()

        self._root_node.remove_child(key)

    def _save(self):

        """Save the map completely.

        :return: The key of the root node.

        """

        if type(self._root_node) == tuple:

            # Already saved.

            return self._root_node

        # TODO: flush root_nodes children?

        bytes = self._root_node.serialise()

        sha1, _, _ = self._store.add_lines((None,), (),

            osutils.split_lines(bytes))

        result = ("sha1:" + sha1,)

        self._root_node._key = result

        return result

class Node(object):

    """Base class defining the protocol for CHK Map nodes."""

    def __init__(self, key_width=1):

        """Create a node.

        :param key_width: The width of keys for this node.

        """

        self._key = None

        # Current number of elements

        self._len = 0

        self._maximum_size = 0

        self._key_width = 1

        # current size in bytes

        self._size = 0

        # The pointers/values this node has - meaning defined by child classes.

        self._items = {}

    def __len__(self):

        return self._len

    @property

    def maximum_size(self):

        """What is the upper limit for adding references to a node."""

        return self._maximum_size

    def set_maximum_size(self, new_size):

        """Set the size threshold for nodes.

        :param new_size: The size at which no data is added to a node. 0 for

            unlimited.

        """

        self._maximum_size = new_size

class LeafNode(Node):

    """A node containing actual key:value pairs."""

    def __init__(self):

        Node.__init__(self)

        # The size of a leaf node with default values and no children.

        self._size = 12

    def _current_size(self):

        """Answer the current serialised size of this node."""

        return (self._size + len(str(self._len)) + len(str(self._key_width)) +

            len(str(self._maximum_size)))

    @classmethod

    def deserialise(klass, bytes, key):

        """Deserialise bytes, with key key, into a LeafNode.

        :param bytes: The bytes of the node.

        :param key: The key that the serialised node has.

        """

        result = LeafNode()

        lines = bytes.splitlines()

        items = {}

        if lines[0] != 'chkleaf:':

            raise ValueError("not a serialised leaf node: %r" % bytes)

        maximum_size = int(lines[1])

        width = int(lines[2])

        length = int(lines[3])

        for line in lines[4:]:

            elements = line.split('\x00')

            items[tuple(elements[:-1])] = elements[-1]

        if len(items) != length:

            raise AssertionError("item count mismatch")

        result._items = items

        result._len = length

        result._maximum_size = maximum_size

        result._key = key

        result._key_width = width

        result._size = len(bytes)

        return result

    def iteritems(self, key_filter=None):

        if key_filter is not None:

            for item in self._items.iteritems():

                if item[0] in key_filter:

                    yield item

        else:

            for item in self._items.iteritems():

                yield item

    def key(self):

        return self._key

    def map(self, key, value):

        """Map key to value."""

        if key in self._items:

            self._size -= 2 + len('\x00'.join(key)) + len(self._items[key])

            self._len -= 1

        self._items[key] = value

        self._size += 2 + len('\x00'.join(key)) + len(value)

        self._len += 1

        self._key = None

        if (self._maximum_size and self._current_size() > self._maximum_size and

            self._len > 1):

            result = InternalNode()

            result.set_maximum_size(self._maximum_size)

            result._key_width = self._key_width

            result._add_node(self)

            return result

        else:

            return self

    def serialise(self, store):

        """Serialise the tree to store.

        :param store: A VersionedFiles honouring the CHK extensions.

        :return: An iterable of the keys inserted by this operation.

        """

        lines = ["chkleaf:\n"]

        lines.append("%d\n" % self._maximum_size)

        lines.append("%d\n" % self._key_width)

        lines.append("%d\n" % self._len)

        for key, value in sorted(self._items.items()):

            lines.append("%s\x00%s\n" % ('\x00'.join(key), value))

        sha1, _, _ = store.add_lines((None,), (), lines)

        self._key = ("sha1:" + sha1,)

        return [self._key]

    def unmap(self, key):

        """Unmap key from the node."""

        self._size -= 2 + len('\x00'.join(key)) + len(self._items[key])

        self._len -= 1

        del self._items[key]

        self._key = None

        return self

class InternalNode(Node):

    """A node that contains references to other nodes."""

    def __init__(self):

        Node.__init__(self)

        # The size of an internalnode with default values and no children.

        # self._size = 12

        # How many octets key prefixes within this node are.

        self._node_width = 0

    def _add_node(self, node):

        """Add a node to the InternalNode.

        :param node: An existing node to add. The node will be examined to see

            if it is over or undersized and rebalanced if needed across this

            nodes children.

        """

        if self._len == 0:

            # new tree level, we're being populated by upspill from a overfull

            # tree.

            # Cheap-to-code-but-slow?

            elements = {}

            max_width = 0

            # suck in all the values

            for key, value in node.iteritems():

                # We work on the serialised keys

                serialised_key = '\x00'.join(key)

                elements[serialised_key] = (key, value)

                max_width = max(len(serialised_key), max_width)

            # Determine the maximum common key width we will internally handle.

            # Start with the full key width; if that exceeds our node size

            # shrink it until we are within the node limit.

            self._node_width = max_width

            width = self._node_width

            # Populate all the resulting keys:

            items = self._items

            for serialised_key, key_value in elements.iteritems():

                actual_key = self._serialised_key(key_value[0])

                child = items.get(actual_key, None)

                if not child:

                    child = LeafNode()

                    child.set_maximum_size(self._maximum_size)

                    child._key_width = self._key_width

                    items[actual_key] = child

                child.map(key_value[0], key_value[1])

                self._len += 1

        else:

            raise NotImplementedError(self._add_node)

    def _current_size(self):

        """Answer the current serialised size of this node."""

        return (self._size + len(str(self._len)) + len(str(self._key_width)) +

            len(str(self._maximum_size)))

    def iteritems(self, key_filter=None):

        if key_filter is None:

            for child in self._items.itervalues():

                for item in child.iteritems():

                    yield item

        else:

            serialised_filter = set([self._serialised_key(key) for key in

                key_filter])

            for key, child in self._items.iteritems():

                if key in serialised_filter:

                    for item in child.iteritems(key_filter):

                        yield item

    def map(self, key, value):

        """Map key to value."""

        serialised_key = self._serialised_key(key)

        try:

            child = self._items[serialised_key]

        except KeyError:

            child = LeafNode()

            child.set_maximum_size(self._maximum_size)

            child._key_width = self._key_width

            self._items[serialised_key] = child

        old_len = len(child)

        new_child = child.map(key, value)

        # TODO: rebalance/enforce balance

        if new_child is not child:

            # The child has exceeded its size; if we take more bytes off the

            # key prefix for the child, that may fit into our node;

            # How many more bytes can we fit?

            remaining_size = max(0, self.maximum_size - self._current_size())

            size_per_row = (self._node_width + 45 + 2)

            # without increasing the key width:

            extra_rows = remaining_size / size_per_row

            if extra_rows:

                # What is the minimum node width increase to split new_child:

                offset_bytes = [1]

                offset = self._node_width - 1

                while len(offset_bytes) == 1 and offset < new_child._node_width:

                    offset += 1

                    offset_bytes = set(child_key[offset] for child_key in

                        new_child._items.keys())

                if len(offset_bytes) > 1:

                    # We've found the fan out point

                    increase = self._node_width - offset

                    # calculate how many more pointers we need to carry

                    new_keys = len(offset_bytes)

                    for subnode in self._items.values():

                        new_keys += subnode._key_count(self._node_width, offset)

                    if (new_keys * (offset + 45 + 2) +

                        self._prelude_size() > self._maximum_size):

                        # can't fit it all, accept the new child

                        self._items[serialised_key] = new_child

                    else:

                        # increasing the 

                        pass

                else:

                    # it didn't fan out! wtf!

                    raise AssertionError("no fan out")

            else:

                # leave it split

                self._items[serialised_key] = new_child

        self._len += 1

        return self

    def _serialised_key(self, key):

        """Return the serialised key for key in this node."""

        return ('\x00'.join(key) + '\x00'*self._node_width)[:self._node_width]

    def _key_count(self, parent_width, cut_width):

        """Return the number of keys in/under this node between two widths.

        :param parent_width: The start offset in keys to consider.

        :param cut_width: The width to stop considering at.

        """

        if cut_width > self._node_width:

            raise NotImplementedError(self._key_count)

        # Generate a list of unique substrings

    def unmap(self, key):

        """Remove key from this node and it's children."""

        serialised_key = self._serialised_key(key)

        child = self._items[serialised_key]

        new_child = child.unmap(key)

        # TODO shrink/rebalance

        if not len(new_child):

            del self._items[serialised_key]

            if len(self._items) == 1:

                return self._items.values()[0]

        elif new_child is not child:

            self._items[serialised_key] = new_child

        return self

class RootNode(Node):

    """A root node in a CHKMap."""

    def __init__(self):

        Node.__init__(self)

        self._nodes = {}

        self._size = 12

        self.prefix_width = 0

    def add_child(self, name, child):

        """Add a child to the node.

        If the node changes, it's key is reset.

        :param name: The name of the child. A bytestring.

        :param child: The child, a key tuple for the childs value.

        """

        if self._maximum_size and self._current_size() >= self._maximum_size:

            return False

        if name in self._nodes:

            self.remove_child(name)

        self._nodes[name] = child

        self._len += 1

        self._key = None

        self._size += len(name) + len(child[0]) + 2

        return True

    def _current_size(self):

        """Answer the current serialised size of this node."""

        return (self._size + len(str(self._maximum_size)) + len(str(self._len))

            + len(str(self.prefix_width)))

    def deserialise(self, bytes, key):

        """Set the nodes value to that contained in bytes.

        :param bytes: The bytes of the node.

        :param key: The key that the serialised node has.

        """

        lines = bytes.splitlines()

        nodes = {}

        if lines[0] != 'chkroot:':

            raise ValueError("not a serialised root node: %r" % bytes)

        maximum_size = int(lines[1])

        prefix_width = int(lines[2])

        length = int(lines[3])

        for line in lines[4:]:

            name, value = line.split('\x00')

            nodes[name] = (value,)

        self._nodes = nodes

        self._len = length

        self._maximum_size = maximum_size

        self._key = key

        self.prefix_width = prefix_width

    def refs(self):

        """Get the CHK key references this node holds."""

        return self._nodes.values()

    def remove_child(self, name):

        """Remove name from the node.

        If the node changes, it's key is reset.

        :param name: The name to remove from the node.

        """

        node = self._nodes.pop(name)

        self._size -= 2 + len(name) + len(node[0])

        self._len -= 1

        self._key = None

    def serialise(self):

        """Flatten the node to a bytestring.

        :return: A bytestring.

        """

        lines = ["chkroot:\n"]

        lines.append("%d\n" % self._maximum_size)

        lines.append("%d\n" % self.prefix_width)

        lines.append("%d\n" % self._len)

        for name, child in sorted(self._nodes.items()):

            lines.append("%s\x00%s\n" % (name, child[0]))

        return "".join(lines)

class ValueNode(object):

    """A value in a CHKMap."""

    def __init__(self, value):

        """Create a ValueNode.

        :param value: The value of this node, must be a bytestring.

        """

        self.value = value

    @classmethod

    def deserialise(klass, bytes):

        """Get a ValueNode from a serialised bytestring.

        :param bytes: The bytes returned by an earlier serialisation.

        """

        if not bytes.startswith("chkvalue:\n"):

            raise ValueError("not a chkvalue %r" % bytes)

        return ValueNode(bytes[10:])

    def serialise(self):

        """Flatten the value to a bytestring.

        :return: A bytestring.

        """

        return "chkvalue:\n" + self.value

    def refs(self):

        """ValueNodes have no refs within the dict."""

        return []

def _deserialise(bytes, key):

    """Helper for repositorydetails - convert bytes to a node."""

    if bytes.startswith("chkvalue:\n"):

        return ValueNode.deserialise(bytes)

    elif bytes.startswith("chkroot:\n"):

        result = RootNode()

        result.deserialise(bytes, key)

        return result

    else:

        raise AssertionError("Unknown node type.")