~bzr-pqm/bzr/bzr.dev

# Copyright (C) 2008, 2009 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

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

Overview and current status:

The CHKMap class implements a dict from tuple_of_strings->string by using a trie

with internal nodes of 8-bit fan out; The key tuples are mapped to strings by

joining them by \x00, and \x00 padding shorter keys out to the length of the

longest key. Leaf nodes are packed as densely as possible, and internal nodes

are all an additional 8-bits wide leading to a sparse upper tree.

Updates to a CHKMap are done preferentially via the apply_delta method, to

allow optimisation of the update operation; but individual map/unmap calls are

possible and supported. All changes via map/unmap are buffered in memory until

the _save method is called to force serialisation of the tree. apply_delta

performs a _save implicitly.

TODO:

-----

Densely packed upper nodes.

"""

import heapq

from bzrlib import lazy_import

lazy_import.lazy_import(globals(), """

from bzrlib import versionedfile

""")

from bzrlib import (

    lru_cache,

    osutils,

    registry,

    trace,

    )

# approx 4MB

# If each line is 50 bytes, and you have 255 internal pages, with 255-way fan

# out, it takes 3.1MB to cache the layer.

_PAGE_CACHE_SIZE = 4*1024*1024

# We are caching bytes so len(value) is perfectly accurate

_page_cache = lru_cache.LRUSizeCache(_PAGE_CACHE_SIZE)

# If a ChildNode falls below this many bytes, we check for a remap

_INTERESTING_NEW_SIZE = 50

# If a ChildNode shrinks by more than this amount, we check for a remap

_INTERESTING_SHRINKAGE_LIMIT = 20

# If we delete more than this many nodes applying a delta, we check for a remap

_INTERESTING_DELETES_LIMIT = 5

def _search_key_plain(key):

    """Map the key tuple into a search string that just uses the key bytes."""

    return '\x00'.join(key)

search_key_registry = registry.Registry()

search_key_registry.register('plain', _search_key_plain)

class CHKMap(object):

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

    def __init__(self, store, root_key, search_key_func=None):

        """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.

        :param search_key_func: A function mapping a key => bytes. These bytes

            are then used by the internal nodes to split up leaf nodes into

            multiple pages.

        """

        self._store = store

        if search_key_func is None:

            search_key_func = _search_key_plain

        self._search_key_func = search_key_func

        if root_key is None:

            self._root_node = LeafNode(search_key_func=search_key_func)

        else:

            self._root_node = self._node_key(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.

        """

        delete_count = 0

        for old, new, value in delta:

            if old is not None and old != new:

                self.unmap(old, check_remap=False)

                delete_count += 1

        for old, new, value in delta:

            if new is not None:

                self.map(new, value)

        if delete_count > _INTERESTING_DELETES_LIMIT:

            trace.mutter("checking remap as %d deletions", delete_count)

            self._check_remap()

        return self._save()

    def _ensure_root(self):

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

        if type(self._root_node) is tuple:

            # Demand-load the root

            self._root_node = self._get_node(self._root_node)

    def _get_node(self, node):

        """Get a node.

        Note that this does not update the _items dict in objects containing a

        reference to this node. As such it does not prevent subsequent IO being

        performed.

        :param node: A tuple key or node object.

        :return: A node object.

        """

        if type(node) is tuple:

            bytes = self._read_bytes(node)

            return _deserialise(bytes, node,

                search_key_func=self._search_key_func)

        else:

            return node

    def _read_bytes(self, key):

        try:

            return _page_cache[key]

        except KeyError:

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

            bytes = stream.next().get_bytes_as('fulltext')

            _page_cache[key] = bytes

            return bytes

    def _dump_tree(self, include_keys=False):

        """Return the tree in a string representation."""

        self._ensure_root()

        res = self._dump_tree_node(self._root_node, prefix='', indent='',

                                   include_keys=include_keys)

        res.append('') # Give a trailing '\n'

        return '\n'.join(res)

    def _dump_tree_node(self, node, prefix, indent, include_keys=True):

        """For this node and all children, generate a string representation."""

        result = []

        if not include_keys:

            key_str = ''

        else:

            node_key = node.key()

            if node_key is not None:

                key_str = ' %s' % (node_key[0],)

            else:

                key_str = ' None'

        result.append('%s%r %s%s' % (indent, prefix, node.__class__.__name__,

                                     key_str))

        if type(node) is InternalNode:

            # Trigger all child nodes to get loaded

            list(node._iter_nodes(self._store))

            for prefix, sub in sorted(node._items.iteritems()):

                result.extend(self._dump_tree_node(sub, prefix, indent + '  ',

                                                   include_keys=include_keys))

        else:

            for key, value in sorted(node._items.iteritems()):

                # Don't use prefix nor indent here to line up when used in

                # tests in conjunction with assertEqualDiff

                result.append('      %r %r' % (key, value))

        return result

    @classmethod

    def from_dict(klass, store, initial_value, maximum_size=0, key_width=1,

        search_key_func=None):

        """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.

        :param key_width: The number of elements in each key_tuple being stored

            in this map.

        :param search_key_func: A function mapping a key => bytes. These bytes

            are then used by the internal nodes to split up leaf nodes into

            multiple pages.

        :return: The root chk of the resulting CHKMap.

        """

        root_key = klass._create_directly(store, initial_value,

            maximum_size=maximum_size, key_width=key_width,

            search_key_func=search_key_func)

        return root_key

    @classmethod

    def _create_via_map(klass, store, initial_value, maximum_size=0,

                        key_width=1, search_key_func=None):

        result = klass(store, None, search_key_func=search_key_func)

        result._root_node.set_maximum_size(maximum_size)

        result._root_node._key_width = key_width

        delta = []

        for key, value in initial_value.items():

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

        root_key = result.apply_delta(delta)

        return root_key

    @classmethod

    def _create_directly(klass, store, initial_value, maximum_size=0,

                         key_width=1, search_key_func=None):

        node = LeafNode(search_key_func=search_key_func)

        node.set_maximum_size(maximum_size)

        node._key_width = key_width

        node._items = dict(initial_value)

        node._raw_size = sum([node._key_value_len(key, value)

                              for key,value in initial_value.iteritems()])

        node._len = len(node._items)

        node._compute_search_prefix()

        node._compute_serialised_prefix()

        if (node._len > 1

            and maximum_size

            and node._current_size() > maximum_size):

            prefix, node_details = node._split(store)

            if len(node_details) == 1:

                raise AssertionError('Failed to split using node._split')

            node = InternalNode(prefix, search_key_func=search_key_func)

            node.set_maximum_size(maximum_size)

            node._key_width = key_width

            for split, subnode in node_details:

                node.add_node(split, subnode)

        keys = list(node.serialise(store))

        return keys[-1]

    def iter_changes(self, basis):

        """Iterate over the changes between basis and self.

        :return: An iterator of tuples: (key, old_value, new_value). Old_value

            is None for keys only in self; new_value is None for keys only in

            basis.

        """

        # Overview:

        # Read both trees in lexographic, highest-first order.

        # Any identical nodes we skip

        # Any unique prefixes we output immediately.

        # values in a leaf node are treated as single-value nodes in the tree

        # which allows them to be not-special-cased. We know to output them

        # because their value is a string, not a key(tuple) or node.

        #

        # corner cases to beware of when considering this function:

        # *) common references are at different heights.

        #    consider two trees:

        #    {'a': LeafNode={'aaa':'foo', 'aab':'bar'}, 'b': LeafNode={'b'}}

        #    {'a': InternalNode={'aa':LeafNode={'aaa':'foo', 'aab':'bar'},

        #                        'ab':LeafNode={'ab':'bar'}}

        #     'b': LeafNode={'b'}}

        #    the node with aaa/aab will only be encountered in the second tree

        #    after reading the 'a' subtree, but it is encountered in the first

        #    tree immediately. Variations on this may have read internal nodes

        #    like this.  we want to cut the entire pending subtree when we

        #    realise we have a common node.  For this we use a list of keys -

        #    the path to a node - and check the entire path is clean as we

        #    process each item.

        if self._node_key(self._root_node) == self._node_key(basis._root_node):

            return

        self._ensure_root()

        basis._ensure_root()

        excluded_keys = set()

        self_node = self._root_node

        basis_node = basis._root_node

        # A heap, each element is prefix, node(tuple/NodeObject/string),

        # key_path (a list of tuples, tail-sharing down the tree.)

        self_pending = []

        basis_pending = []

        def process_node(node, path, a_map, pending):

            # take a node and expand it

            node = a_map._get_node(node)

            if type(node) == LeafNode:

                path = (node._key, path)

                for key, value in node._items.items():

                    # For a LeafNode, the key is a serialized_key, rather than

                    # a search_key, but the heap is using search_keys

                    search_key = node._search_key_func(key)

                    heapq.heappush(pending, (search_key, key, value, path))

            else:

                # type(node) == InternalNode

                path = (node._key, path)

                for prefix, child in node._items.items():

                    heapq.heappush(pending, (prefix, None, child, path))

        def process_common_internal_nodes(self_node, basis_node):

            self_items = set(self_node._items.items())

            basis_items = set(basis_node._items.items())

            path = (self_node._key, None)

            for prefix, child in self_items - basis_items:

                heapq.heappush(self_pending, (prefix, None, child, path))

            path = (basis_node._key, None)

            for prefix, child in basis_items - self_items:

                heapq.heappush(basis_pending, (prefix, None, child, path))

        def process_common_leaf_nodes(self_node, basis_node):

            self_items = set(self_node._items.items())

            basis_items = set(basis_node._items.items())

            path = (self_node._key, None)

            for key, value in self_items - basis_items:

                prefix = self._search_key_func(key)

                heapq.heappush(self_pending, (prefix, key, value, path))

            path = (basis_node._key, None)

            for key, value in basis_items - self_items:

                prefix = basis._search_key_func(key)

                heapq.heappush(basis_pending, (prefix, key, value, path))

        def process_common_prefix_nodes(self_node, self_path,

                                        basis_node, basis_path):

            # Would it be more efficient if we could request both at the same

            # time?

            self_node = self._get_node(self_node)

            basis_node = basis._get_node(basis_node)

            if (type(self_node) == InternalNode

                and type(basis_node) == InternalNode):

                # Matching internal nodes

                process_common_internal_nodes(self_node, basis_node)

            elif (type(self_node) == LeafNode

                  and type(basis_node) == LeafNode):

                process_common_leaf_nodes(self_node, basis_node)

            else:

                process_node(self_node, self_path, self, self_pending)

                process_node(basis_node, basis_path, basis, basis_pending)

        process_common_prefix_nodes(self_node, None, basis_node, None)

        self_seen = set()

        basis_seen = set()

        excluded_keys = set()

        def check_excluded(key_path):

            # Note that this is N^2, it depends on us trimming trees

            # aggressively to not become slow.

            # A better implementation would probably have a reverse map

            # back to the children of a node, and jump straight to it when

            # a common node is detected, the proceed to remove the already

            # pending children. bzrlib.graph has a searcher module with a

            # similar problem.

            while key_path is not None:

                key, key_path = key_path

                if key in excluded_keys:

                    return True

            return False

        loop_counter = 0

        while self_pending or basis_pending:

            loop_counter += 1

            if not self_pending:

                # self is exhausted: output remainder of basis

                for prefix, key, node, path in basis_pending:

                    if check_excluded(path):

                        continue

                    node = basis._get_node(node)

                    if key is not None:

                        # a value

                        yield (key, node, None)

                    else:

                        # subtree - fastpath the entire thing.

                        for key, value in node.iteritems(basis._store):

                            yield (key, value, None)

                return

            elif not basis_pending:

                # basis is exhausted: output remainder of self.

                for prefix, key, node, path in self_pending:

                    if check_excluded(path):

                        continue

                    node = self._get_node(node)

                    if key is not None:

                        # a value

                        yield (key, None, node)

                    else:

                        # subtree - fastpath the entire thing.

                        for key, value in node.iteritems(self._store):

                            yield (key, None, value)

                return

            else:

                # XXX: future optimisation - yield the smaller items

                # immediately rather than pushing everything on/off the

                # heaps. Applies to both internal nodes and leafnodes.

                if self_pending[0][0] < basis_pending[0][0]:

                    # expand self

                    prefix, key, node, path = heapq.heappop(self_pending)

                    if check_excluded(path):

                        continue

                    if key is not None:

                        # a value

                        yield (key, None, node)

                    else:

                        process_node(node, path, self, self_pending)

                        continue

                elif self_pending[0][0] > basis_pending[0][0]:

                    # expand basis

                    prefix, key, node, path = heapq.heappop(basis_pending)

                    if check_excluded(path):

                        continue

                    if key is not None:

                        # a value

                        yield (key, node, None)

                    else:

                        process_node(node, path, basis, basis_pending)

                        continue

                else:

                    # common prefix: possibly expand both

                    if self_pending[0][1] is None:

                        # process next self

                        read_self = True

                    else:

                        read_self = False

                    if basis_pending[0][1] is None:

                        # process next basis

                        read_basis = True

                    else:

                        read_basis = False

                    if not read_self and not read_basis:

                        # compare a common value

                        self_details = heapq.heappop(self_pending)

                        basis_details = heapq.heappop(basis_pending)

                        if self_details[2] != basis_details[2]:

                            yield (self_details[1],

                                basis_details[2], self_details[2])

                        continue

                    # At least one side wasn't a simple value

                    if (self._node_key(self_pending[0][2]) ==

                        self._node_key(basis_pending[0][2])):

                        # Identical pointers, skip (and don't bother adding to

                        # excluded, it won't turn up again.

                        heapq.heappop(self_pending)

                        heapq.heappop(basis_pending)

                        continue

                    # Now we need to expand this node before we can continue

                    if read_self and read_basis:

                        # Both sides start with the same prefix, so process

                        # them in parallel

                        self_prefix, _, self_node, self_path = heapq.heappop(

                            self_pending)

                        basis_prefix, _, basis_node, basis_path = heapq.heappop(

                            basis_pending)

                        if self_prefix != basis_prefix:

                            raise AssertionError(

                                '%r != %r' % (self_prefix, basis_prefix))

                        process_common_prefix_nodes(

                            self_node, self_path,

                            basis_node, basis_path)

                        continue

                    if read_self:

                        prefix, key, node, path = heapq.heappop(self_pending)

                        if check_excluded(path):

                            continue

                        process_node(node, path, self, self_pending)

                    if read_basis:

                        prefix, key, node, path = heapq.heappop(basis_pending)

                        if check_excluded(path):

                            continue

                        process_node(node, path, basis, basis_pending)

        # print loop_counter

    def iteritems(self, key_filter=None):

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

        self._ensure_root()

        return self._root_node.iteritems(self._store, key_filter=key_filter)

    def key(self):

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

        if type(self._root_node) is 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 a key tuple to value."""

        # Need a root object.

        self._ensure_root()

        prefix, node_details = self._root_node.map(self._store, key, value)

        if len(node_details) == 1:

            self._root_node = node_details[0][1]

        else:

            self._root_node = InternalNode(prefix,

                                search_key_func=self._search_key_func)

            self._root_node.set_maximum_size(node_details[0][1].maximum_size)

            self._root_node._key_width = node_details[0][1]._key_width

            for split, node in node_details:

                self._root_node.add_node(split, node)

    def _node_key(self, node):

        """Get the key for a node whether it's a tuple or node."""

        if type(node) is tuple:

            return node

        else:

            return node._key

    def unmap(self, key, check_remap=True):

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

        self._ensure_root()

        if type(self._root_node) is InternalNode:

            unmapped = self._root_node.unmap(self._store, key,

                check_remap=check_remap)

        else:

            unmapped = self._root_node.unmap(self._store, key)

        self._root_node = unmapped

    def _check_remap(self):

        """Check if nodes can be collapsed."""

        self._ensure_root()

        if type(self._root_node) is InternalNode:

            self._root_node._check_remap(self._store)

    def _save(self):

        """Save the map completely.

        :return: The key of the root node.

        """

        if type(self._root_node) is tuple:

            # Already saved.

            return self._root_node

        keys = list(self._root_node.serialise(self._store))

        return keys[-1]

class Node(object):

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

    :ivar _raw_size: The total size of the serialized key:value data, before

        adding the header bytes, and without prefix compression.

    """

    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 = key_width

        # current size in bytes

        self._raw_size = 0

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

        self._items = {}

        # The common search prefix

        self._search_prefix = None

    def __repr__(self):

        items_str = str(sorted(self._items))

        if len(items_str) > 20:

            items_str = items_str[:16] + '...]'

        return '%s(key:%s len:%s size:%s max:%s prefix:%s items:%s)' % (

            self.__class__.__name__, self._key, self._len, self._raw_size,

            self._maximum_size, self._search_prefix, items_str)

    def key(self):

        return self._key

    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

    @classmethod

    def common_prefix(cls, prefix, key):

        """Given 2 strings, return the longest prefix common to both.

        :param prefix: This has been the common prefix for other keys, so it is

            more likely to be the common prefix in this case as well.

        :param key: Another string to compare to

        """

        if key.startswith(prefix):

            return prefix

        pos = -1

        # Is there a better way to do this?

        for pos, (left, right) in enumerate(zip(prefix, key)):

            if left != right:

                pos -= 1

                break

        common = prefix[:pos+1]

        return common

    @classmethod

    def common_prefix_for_keys(cls, keys):

        """Given a list of keys, find their common prefix.

        :param keys: An iterable of strings.

        :return: The longest common prefix of all keys.

        """

        common_prefix = None

        for key in keys:

            if common_prefix is None:

                common_prefix = key

                continue

            common_prefix = cls.common_prefix(common_prefix, key)

            if not common_prefix:

                # if common_prefix is the empty string, then we know it won't

                # change further

                return ''

        return common_prefix

# Singleton indicating we have not computed _search_prefix yet

_unknown = object()

class LeafNode(Node):

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

    :ivar _items: A dict of key->value items. The key is in tuple form.

    :ivar _size: The number of bytes that would be used by serializing all of

        the key/value pairs.

    """

    def __init__(self, search_key_func=None):

        Node.__init__(self)

        # All of the keys in this leaf node share this common prefix

        self._common_serialised_prefix = None

        self._serialise_key = '\x00'.join

        if search_key_func is None:

            self._search_key_func = _search_key_plain

        else:

            self._search_key_func = search_key_func

    def __repr__(self):

        items_str = str(sorted(self._items))

        if len(items_str) > 20:

            items_str = items_str[:16] + '...]'

        return \

            '%s(key:%s len:%s size:%s max:%s prefix:%s keywidth:%s items:%s)' \

            % (self.__class__.__name__, self._key, self._len, self._raw_size,

            self._maximum_size, self._search_prefix, self._key_width, items_str)

    def _current_size(self):

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

        This differs from self._raw_size in that it includes the bytes used for

        the header.

        """

        if self._common_serialised_prefix is None:

            bytes_for_items = 0

            prefix_len = 0

        else:

            # We will store a single string with the common prefix

            # And then that common prefix will not be stored in any of the

            # entry lines

            prefix_len = len(self._common_serialised_prefix)

            bytes_for_items = (self._raw_size - (prefix_len * self._len))

        return (9 # 'chkleaf:\n'

            + len(str(self._maximum_size)) + 1

            + len(str(self._key_width)) + 1

            + len(str(self._len)) + 1

            + prefix_len + 1

            + bytes_for_items)

    @classmethod

    def deserialise(klass, bytes, key, search_key_func=None):

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

        :param bytes: The bytes of the node.

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

        """

        return _deserialise_leaf_node(bytes, key,

                                      search_key_func=search_key_func)

    def iteritems(self, store, key_filter=None):

        """Iterate over items in the node.

        :param key_filter: A filter to apply to the node. It should be a

            list/set/dict or similar repeatedly iterable container.

        """

        if key_filter is not None:

            # Adjust the filter - short elements go to a prefix filter. All

            # other items are looked up directly.

            # XXX: perhaps defaultdict? Profiling<rinse and repeat>

            filters = {}

            for key in key_filter:

                if len(key) == self._key_width:

                    # This filter is meant to match exactly one key, yield it

                    # if we have it.

                    try:

                        yield key, self._items[key]

                    except KeyError:

                        # This key is not present in this map, continue

                        pass

                else:

                    # Short items, we need to match based on a prefix

                    length_filter = filters.setdefault(len(key), set())

                    length_filter.add(key)

            if filters:

                filters = filters.items()

                for item in self._items.iteritems():

                    for length, length_filter in filters:

                        if item[0][:length] in length_filter:

                            yield item

                            break

        else:

            for item in self._items.iteritems():

                yield item

    def _key_value_len(self, key, value):

        # TODO: Should probably be done without actually joining the key, but

        #       then that can be done via the C extension

        return (len(self._serialise_key(key)) + 1

                + len(str(value.count('\n'))) + 1

                + len(value) + 1)

    def _search_key(self, key):

        return self._search_key_func(key)

    def _map_no_split(self, key, value):

        """Map a key to a value.

        This assumes either the key does not already exist, or you have already

        removed its size and length from self.

        :return: True if adding this node should cause us to split.

        """

        self._items[key] = value

        self._raw_size += self._key_value_len(key, value)

        self._len += 1

        serialised_key = self._serialise_key(key)

        if self._common_serialised_prefix is None:

            self._common_serialised_prefix = serialised_key