~bzr-pqm/bzr/bzr.dev

# Copyright (C) 2007, 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

"""Indexing facilities."""

__all__ = [

    'CombinedGraphIndex',

    'GraphIndex',

    'GraphIndexBuilder',

    'GraphIndexPrefixAdapter',

    'InMemoryGraphIndex',

    ]

from bisect import bisect_right

from cStringIO import StringIO

import re

import sys

from bzrlib.lazy_import import lazy_import

lazy_import(globals(), """

from bzrlib import trace

from bzrlib.bisect_multi import bisect_multi_bytes

from bzrlib.revision import NULL_REVISION

from bzrlib.trace import mutter

""")

from bzrlib import (

    debug,

    errors,

    )

from bzrlib.static_tuple import StaticTuple

_HEADER_READV = (0, 200)

_OPTION_KEY_ELEMENTS = "key_elements="

_OPTION_LEN = "len="

_OPTION_NODE_REFS = "node_ref_lists="

_SIGNATURE = "Bazaar Graph Index 1\n"

_whitespace_re = re.compile('[\t\n\x0b\x0c\r\x00 ]')

_newline_null_re = re.compile('[\n\0]')

def _has_key_from_parent_map(self, key):

    """Check if this index has one key.

    If it's possible to check for multiple keys at once through

    calling get_parent_map that should be faster.

    """

    return (key in self.get_parent_map([key]))

def _missing_keys_from_parent_map(self, keys):

    return set(keys) - set(self.get_parent_map(keys))

class GraphIndexBuilder(object):

    """A builder that can build a GraphIndex.

    The resulting graph has the structure:

    _SIGNATURE OPTIONS NODES NEWLINE

    _SIGNATURE     := 'Bazaar Graph Index 1' NEWLINE

    OPTIONS        := 'node_ref_lists=' DIGITS NEWLINE

    NODES          := NODE*

    NODE           := KEY NULL ABSENT? NULL REFERENCES NULL VALUE NEWLINE

    KEY            := Not-whitespace-utf8

    ABSENT         := 'a'

    REFERENCES     := REFERENCE_LIST (TAB REFERENCE_LIST){node_ref_lists - 1}

    REFERENCE_LIST := (REFERENCE (CR REFERENCE)*)?

    REFERENCE      := DIGITS  ; digits is the byte offset in the index of the

                              ; referenced key.

    VALUE          := no-newline-no-null-bytes

    """

    def __init__(self, reference_lists=0, key_elements=1):

        """Create a GraphIndex builder.

        :param reference_lists: The number of node references lists for each

            entry.

        :param key_elements: The number of bytestrings in each key.

        """

        self.reference_lists = reference_lists

        # A dict of {key: (absent, ref_lists, value)}

        self._nodes = {}

        # Keys that are referenced but not actually present in this index

        self._absent_keys = set()

        self._nodes_by_key = None

        self._key_length = key_elements

        self._optimize_for_size = False

        self._combine_backing_indices = True

    def _check_key(self, key):

        """Raise BadIndexKey if key is not a valid key for this index."""

        if type(key) not in (tuple, StaticTuple):

            raise errors.BadIndexKey(key)

        if self._key_length != len(key):

            raise errors.BadIndexKey(key)

        for element in key:

            if not element or _whitespace_re.search(element) is not None:

                raise errors.BadIndexKey(element)

    def _external_references(self):

        """Return references that are not present in this index.

        """

        keys = set()

        refs = set()

        # TODO: JAM 2008-11-21 This makes an assumption about how the reference

        #       lists are used. It is currently correct for pack-0.92 through

        #       1.9, which use the node references (3rd column) second

        #       reference list as the compression parent. Perhaps this should

        #       be moved into something higher up the stack, since it

        #       makes assumptions about how the index is used.

        if self.reference_lists > 1:

            for node in self.iter_all_entries():

                keys.add(node[1])

                refs.update(node[3][1])

            return refs - keys

        else:

            # If reference_lists == 0 there can be no external references, and

            # if reference_lists == 1, then there isn't a place to store the

            # compression parent

            return set()

    def _get_nodes_by_key(self):

        if self._nodes_by_key is None:

            nodes_by_key = {}

            if self.reference_lists:

                for key, (absent, references, value) in self._nodes.iteritems():

                    if absent:

                        continue

                    key_dict = nodes_by_key

                    for subkey in key[:-1]:

                        key_dict = key_dict.setdefault(subkey, {})

                    key_dict[key[-1]] = key, value, references

            else:

                for key, (absent, references, value) in self._nodes.iteritems():

                    if absent:

                        continue

                    key_dict = nodes_by_key

                    for subkey in key[:-1]:

                        key_dict = key_dict.setdefault(subkey, {})

                    key_dict[key[-1]] = key, value

            self._nodes_by_key = nodes_by_key

        return self._nodes_by_key

    def _update_nodes_by_key(self, key, value, node_refs):

        """Update the _nodes_by_key dict with a new key.

        For a key of (foo, bar, baz) create

        _nodes_by_key[foo][bar][baz] = key_value

        """

        if self._nodes_by_key is None:

            return

        key_dict = self._nodes_by_key

        if self.reference_lists:

            key_value = StaticTuple(key, value, node_refs)

        else:

            key_value = StaticTuple(key, value)

        for subkey in key[:-1]:

            key_dict = key_dict.setdefault(subkey, {})

        key_dict[key[-1]] = key_value

    def _check_key_ref_value(self, key, references, value):

        """Check that 'key' and 'references' are all valid.

        :param key: A key tuple. Must conform to the key interface (be a tuple,

            be of the right length, not have any whitespace or nulls in any key

            element.)

        :param references: An iterable of reference lists. Something like

            [[(ref, key)], [(ref, key), (other, key)]]

        :param value: The value associate with this key. Must not contain

            newlines or null characters.

        :return: (node_refs, absent_references)

            node_refs   basically a packed form of 'references' where all

                        iterables are tuples

            absent_references   reference keys that are not in self._nodes.

                                This may contain duplicates if the same key is

                                referenced in multiple lists.

        """

        as_st = StaticTuple.from_sequence

        self._check_key(key)

        if _newline_null_re.search(value) is not None:

            raise errors.BadIndexValue(value)

        if len(references) != self.reference_lists:

            raise errors.BadIndexValue(references)

        node_refs = []

        absent_references = []

        for reference_list in references:

            for reference in reference_list:

                # If reference *is* in self._nodes, then we know it has already

                # been checked.

                if reference not in self._nodes:

                    self._check_key(reference)

                    absent_references.append(reference)

            reference_list = as_st([as_st(ref).intern()

                                    for ref in reference_list])

            node_refs.append(reference_list)

        return as_st(node_refs), absent_references

    def add_node(self, key, value, references=()):

        """Add a node to the index.

        :param key: The key. keys are non-empty tuples containing

            as many whitespace-free utf8 bytestrings as the key length

            defined for this index.

        :param references: An iterable of iterables of keys. Each is a

            reference to another key.

        :param value: The value to associate with the key. It may be any

            bytes as long as it does not contain \0 or \n.

        """

        (node_refs,

         absent_references) = self._check_key_ref_value(key, references, value)

        if key in self._nodes and self._nodes[key][0] != 'a':

            raise errors.BadIndexDuplicateKey(key, self)

        for reference in absent_references:

            # There may be duplicates, but I don't think it is worth worrying

            # about

            self._nodes[reference] = ('a', (), '')

        self._absent_keys.update(absent_references)

        self._absent_keys.discard(key)

        self._nodes[key] = ('', node_refs, value)

        if self._nodes_by_key is not None and self._key_length > 1:

            self._update_nodes_by_key(key, value, node_refs)

    def clear_cache(self):

        """See GraphIndex.clear_cache()

        This is a no-op, but we need the api to conform to a generic 'Index'

        abstraction.

        """

    def finish(self):

        lines = [_SIGNATURE]

        lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n')

        lines.append(_OPTION_KEY_ELEMENTS + str(self._key_length) + '\n')

        key_count = len(self._nodes) - len(self._absent_keys)

        lines.append(_OPTION_LEN + str(key_count) + '\n')

        prefix_length = sum(len(x) for x in lines)

        # references are byte offsets. To avoid having to do nasty

        # polynomial work to resolve offsets (references to later in the

        # file cannot be determined until all the inbetween references have

        # been calculated too) we pad the offsets with 0's to make them be

        # of consistent length. Using binary offsets would break the trivial

        # file parsing.

        # to calculate the width of zero's needed we do three passes:

        # one to gather all the non-reference data and the number of references.

        # one to pad all the data with reference-length and determine entry

        # addresses.

        # One to serialise.

        # forward sorted by key. In future we may consider topological sorting,

        # at the cost of table scans for direct lookup, or a second index for

        # direct lookup

        nodes = sorted(self._nodes.items())

        # if we do not prepass, we don't know how long it will be up front.

        expected_bytes = None

        # we only need to pre-pass if we have reference lists at all.

        if self.reference_lists:

            key_offset_info = []

            non_ref_bytes = prefix_length

            total_references = 0

            # TODO use simple multiplication for the constants in this loop.

            for key, (absent, references, value) in nodes:

                # record the offset known *so far* for this key:

                # the non reference bytes to date, and the total references to

                # date - saves reaccumulating on the second pass

                key_offset_info.append((key, non_ref_bytes, total_references))

                # key is literal, value is literal, there are 3 null's, 1 NL

                # key is variable length tuple, \x00 between elements

                non_ref_bytes += sum(len(element) for element in key)

                if self._key_length > 1:

                    non_ref_bytes += self._key_length - 1

                # value is literal bytes, there are 3 null's, 1 NL.

                non_ref_bytes += len(value) + 3 + 1

                # one byte for absent if set.

                if absent:

                    non_ref_bytes += 1

                elif self.reference_lists:

                    # (ref_lists -1) tabs

                    non_ref_bytes += self.reference_lists - 1

                    # (ref-1 cr's per ref_list)

                    for ref_list in references:

                        # how many references across the whole file?

                        total_references += len(ref_list)

                        # accrue reference separators

                        if ref_list:

                            non_ref_bytes += len(ref_list) - 1

            # how many digits are needed to represent the total byte count?

            digits = 1

            possible_total_bytes = non_ref_bytes + total_references*digits

            while 10 ** digits < possible_total_bytes:

                digits += 1

                possible_total_bytes = non_ref_bytes + total_references*digits

            expected_bytes = possible_total_bytes + 1 # terminating newline

            # resolve key addresses.

            key_addresses = {}

            for key, non_ref_bytes, total_references in key_offset_info:

                key_addresses[key] = non_ref_bytes + total_references*digits

            # serialise

            format_string = '%%0%sd' % digits

        for key, (absent, references, value) in nodes:

            flattened_references = []

            for ref_list in references:

                ref_addresses = []

                for reference in ref_list:

                    ref_addresses.append(format_string % key_addresses[reference])

                flattened_references.append('\r'.join(ref_addresses))

            string_key = '\x00'.join(key)

            lines.append("%s\x00%s\x00%s\x00%s\n" % (string_key, absent,

                '\t'.join(flattened_references), value))

        lines.append('\n')

        result = StringIO(''.join(lines))

        if expected_bytes and len(result.getvalue()) != expected_bytes:

            raise errors.BzrError('Failed index creation. Internal error:'

                ' mismatched output length and expected length: %d %d' %

                (len(result.getvalue()), expected_bytes))

        return result

    def set_optimize(self, for_size=None, combine_backing_indices=None):

        """Change how the builder tries to optimize the result.

        :param for_size: Tell the builder to try and make the index as small as

            possible.

        :param combine_backing_indices: If the builder spills to disk to save

            memory, should the on-disk indices be combined. Set to True if you

            are going to be probing the index, but to False if you are not. (If

            you are not querying, then the time spent combining is wasted.)

        :return: None

        """

        # GraphIndexBuilder itself doesn't pay attention to the flag yet, but

        # other builders do.

        if for_size is not None:

            self._optimize_for_size = for_size

        if combine_backing_indices is not None:

            self._combine_backing_indices = combine_backing_indices

    def find_ancestry(self, keys, ref_list_num):

        """See CombinedGraphIndex.find_ancestry()"""

        pending = set(keys)

        parent_map = {}

        missing_keys = set()

        while pending:

            next_pending = set()

            for _, key, value, ref_lists in self.iter_entries(pending):

                parent_keys = ref_lists[ref_list_num]

                parent_map[key] = parent_keys

                next_pending.update([p for p in parent_keys if p not in

                                     parent_map])

                missing_keys.update(pending.difference(parent_map))

            pending = next_pending

        return parent_map, missing_keys

class GraphIndex(object):

    """An index for data with embedded graphs.

    The index maps keys to a list of key reference lists, and a value.

    Each node has the same number of key reference lists. Each key reference

    list can be empty or an arbitrary length. The value is an opaque NULL

    terminated string without any newlines. The storage of the index is

    hidden in the interface: keys and key references are always tuples of

    bytestrings, never the internal representation (e.g. dictionary offsets).

    It is presumed that the index will not be mutated - it is static data.

    Successive iter_all_entries calls will read the entire index each time.

    Additionally, iter_entries calls will read the index linearly until the

    desired keys are found. XXX: This must be fixed before the index is

    suitable for production use. :XXX

    """

    def __init__(self, transport, name, size, unlimited_cache=False):

        """Open an index called name on transport.

        :param transport: A bzrlib.transport.Transport.

        :param name: A path to provide to transport API calls.

        :param size: The size of the index in bytes. This is used for bisection

            logic to perform partial index reads. While the size could be

            obtained by statting the file this introduced an additional round

            trip as well as requiring stat'able transports, both of which are

            avoided by having it supplied. If size is None, then bisection

            support will be disabled and accessing the index will just stream

            all the data.

        """

        self._transport = transport

        self._name = name

        # Becomes a dict of key:(value, reference-list-byte-locations) used by

        # the bisection interface to store parsed but not resolved keys.

        self._bisect_nodes = None

        # Becomes a dict of key:(value, reference-list-keys) which are ready to

        # be returned directly to callers.

        self._nodes = None

        # a sorted list of slice-addresses for the parsed bytes of the file.

        # e.g. (0,1) would mean that byte 0 is parsed.

        self._parsed_byte_map = []

        # a sorted list of keys matching each slice address for parsed bytes

        # e.g. (None, 'foo@bar') would mean that the first byte contained no

        # key, and the end byte of the slice is the of the data for 'foo@bar'

        self._parsed_key_map = []

        self._key_count = None

        self._keys_by_offset = None

        self._nodes_by_key = None

        self._size = size

        # The number of bytes we've read so far in trying to process this file

        self._bytes_read = 0

    def __eq__(self, other):

        """Equal when self and other were created with the same parameters."""

        return (

            type(self) == type(other) and

            self._transport == other._transport and

            self._name == other._name and

            self._size == other._size)

    def __ne__(self, other):

        return not self.__eq__(other)

    def __repr__(self):

        return "%s(%r)" % (self.__class__.__name__,

            self._transport.abspath(self._name))

    def _buffer_all(self, stream=None):

        """Buffer all the index data.

        Mutates self._nodes and self.keys_by_offset.

        """

        if self._nodes is not None:

            # We already did this

            return

        if 'index' in debug.debug_flags:

            mutter('Reading entire index %s', self._transport.abspath(self._name))

        if stream is None:

            stream = self._transport.get(self._name)

        self._read_prefix(stream)

        self._expected_elements = 3 + self._key_length

        line_count = 0

        # raw data keyed by offset

        self._keys_by_offset = {}

        # ready-to-return key:value or key:value, node_ref_lists

        self._nodes = {}

        self._nodes_by_key = None

        trailers = 0

        pos = stream.tell()

        lines = stream.read().split('\n')

        stream.close()

        del lines[-1]

        _, _, _, trailers = self._parse_lines(lines, pos)

        for key, absent, references, value in self._keys_by_offset.itervalues():

            if absent:

                continue

            # resolve references:

            if self.node_ref_lists:

                node_value = (value, self._resolve_references(references))

            else:

                node_value = value

            self._nodes[key] = node_value

        # cache the keys for quick set intersections

        if trailers != 1:

            # there must be one line - the empty trailer line.

            raise errors.BadIndexData(self)

    def clear_cache(self):

        """Clear out any cached/memoized values.

        This can be called at any time, but generally it is used when we have

        extracted some information, but don't expect to be requesting any more

        from this index.

        """

    def external_references(self, ref_list_num):

        """Return references that are not present in this index.

        """

        self._buffer_all()

        if ref_list_num + 1 > self.node_ref_lists:

            raise ValueError('No ref list %d, index has %d ref lists'

                % (ref_list_num, self.node_ref_lists))

        refs = set()

        nodes = self._nodes

        for key, (value, ref_lists) in nodes.iteritems():

            ref_list = ref_lists[ref_list_num]

            refs.update([ref for ref in ref_list if ref not in nodes])

        return refs

    def _get_nodes_by_key(self):

        if self._nodes_by_key is None:

            nodes_by_key = {}

            if self.node_ref_lists:

                for key, (value, references) in self._nodes.iteritems():

                    key_dict = nodes_by_key

                    for subkey in key[:-1]:

                        key_dict = key_dict.setdefault(subkey, {})

                    key_dict[key[-1]] = key, value, references

            else:

                for key, value in self._nodes.iteritems():

                    key_dict = nodes_by_key

                    for subkey in key[:-1]:

                        key_dict = key_dict.setdefault(subkey, {})

                    key_dict[key[-1]] = key, value

            self._nodes_by_key = nodes_by_key

        return self._nodes_by_key

    def iter_all_entries(self):

        """Iterate over all keys within the index.

        :return: An iterable of (index, key, value) or (index, key, value, reference_lists).

            The former tuple is used when there are no reference lists in the

            index, making the API compatible with simple key:value index types.

            There is no defined order for the result iteration - it will be in

            the most efficient order for the index.

        """

        if 'evil' in debug.debug_flags:

            trace.mutter_callsite(3,

                "iter_all_entries scales with size of history.")

        if self._nodes is None:

            self._buffer_all()

        if self.node_ref_lists:

            for key, (value, node_ref_lists) in self._nodes.iteritems():

                yield self, key, value, node_ref_lists

        else:

            for key, value in self._nodes.iteritems():

                yield self, key, value

    def _read_prefix(self, stream):

        signature = stream.read(len(self._signature()))

        if not signature == self._signature():

            raise errors.BadIndexFormatSignature(self._name, GraphIndex)

        options_line = stream.readline()

        if not options_line.startswith(_OPTION_NODE_REFS):

            raise errors.BadIndexOptions(self)

        try:

            self.node_ref_lists = int(options_line[len(_OPTION_NODE_REFS):-1])

        except ValueError:

            raise errors.BadIndexOptions(self)

        options_line = stream.readline()

        if not options_line.startswith(_OPTION_KEY_ELEMENTS):

            raise errors.BadIndexOptions(self)

        try:

            self._key_length = int(options_line[len(_OPTION_KEY_ELEMENTS):-1])

        except ValueError:

            raise errors.BadIndexOptions(self)

        options_line = stream.readline()

        if not options_line.startswith(_OPTION_LEN):

            raise errors.BadIndexOptions(self)

        try:

            self._key_count = int(options_line[len(_OPTION_LEN):-1])

        except ValueError:

            raise errors.BadIndexOptions(self)

    def _resolve_references(self, references):

        """Return the resolved key references for references.

        References are resolved by looking up the location of the key in the

        _keys_by_offset map and substituting the key name, preserving ordering.

        :param references: An iterable of iterables of key locations. e.g.

            [[123, 456], [123]]

        :return: A tuple of tuples of keys.

        """

        node_refs = []

        for ref_list in references:

            node_refs.append(tuple([self._keys_by_offset[ref][0] for ref in ref_list]))

        return tuple(node_refs)

    def _find_index(self, range_map, key):

        """Helper for the _parsed_*_index calls.

        Given a range map - [(start, end), ...], finds the index of the range

        in the map for key if it is in the map, and if it is not there, the

        immediately preceeding range in the map.

        """

        result = bisect_right(range_map, key) - 1

        if result + 1 < len(range_map):

            # check the border condition, it may be in result + 1

            if range_map[result + 1][0] == key[0]:

                return result + 1

        return result

    def _parsed_byte_index(self, offset):

        """Return the index of the entry immediately before offset.

        e.g. if the parsed map has regions 0,10 and 11,12 parsed, meaning that

        there is one unparsed byte (the 11th, addressed as[10]). then:

        asking for 0 will return 0

        asking for 10 will return 0

        asking for 11 will return 1

        asking for 12 will return 1

        """

        key = (offset, 0)

        return self._find_index(self._parsed_byte_map, key)

    def _parsed_key_index(self, key):

        """Return the index of the entry immediately before key.

        e.g. if the parsed map has regions (None, 'a') and ('b','c') parsed,

        meaning that keys from None to 'a' inclusive, and 'b' to 'c' inclusive

        have been parsed, then:

        asking for '' will return 0

        asking for 'a' will return 0

        asking for 'b' will return 1

        asking for 'e' will return 1

        """

        search_key = (key, None)

        return self._find_index(self._parsed_key_map, search_key)

    def _is_parsed(self, offset):

        """Returns True if offset has been parsed."""

        index = self._parsed_byte_index(offset)

        if index == len(self._parsed_byte_map):

            return offset < self._parsed_byte_map[index - 1][1]

        start, end = self._parsed_byte_map[index]

        return offset >= start and offset < end

    def _iter_entries_from_total_buffer(self, keys):

        """Iterate over keys when the entire index is parsed."""

        # Note: See the note in BTreeBuilder.iter_entries for why we don't use

        #       .intersection() here

        nodes = self._nodes

        keys = [key for key in keys if key in nodes]

        if self.node_ref_lists:

            for key in keys:

                value, node_refs = nodes[key]

                yield self, key, value, node_refs

        else:

            for key in keys:

                yield self, key, nodes[key]

    def iter_entries(self, keys):

        """Iterate over keys within the index.

        :param keys: An iterable providing the keys to be retrieved.

        :return: An iterable as per iter_all_entries, but restricted to the

            keys supplied. No additional keys will be returned, and every

            key supplied that is in the index will be returned.

        """

        keys = set(keys)

        if not keys:

            return []

        if self._size is None and self._nodes is None:

            self._buffer_all()

        # We fit about 20 keys per minimum-read (4K), so if we are looking for

        # more than 1/20th of the index its likely (assuming homogenous key

        # spread) that we'll read the entire index. If we're going to do that,

        # buffer the whole thing. A better analysis might take key spread into

        # account - but B+Tree indices are better anyway.

        # We could look at all data read, and use a threshold there, which will

        # trigger on ancestry walks, but that is not yet fully mapped out.

        if self._nodes is None and len(keys) * 20 > self.key_count():

            self._buffer_all()

        if self._nodes is not None:

            return self._iter_entries_from_total_buffer(keys)

        else:

            return (result[1] for result in bisect_multi_bytes(

                self._lookup_keys_via_location, self._size, keys))