~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/index.py

  • Committer: Martin Pool
  • Date: 2005-05-09 04:38:31 UTC
  • Revision ID: mbp@sourcefrog.net-20050509043831-d45f7832b7d4d5b0
- better message when refusing to add symlinks (from mpe)

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
# Copyright (C) 2007 Canonical Ltd
2
 
#
3
 
# This program is free software; you can redistribute it and/or modify
4
 
# it under the terms of the GNU General Public License as published by
5
 
# the Free Software Foundation; either version 2 of the License, or
6
 
# (at your option) any later version.
7
 
#
8
 
# This program is distributed in the hope that it will be useful,
9
 
# but WITHOUT ANY WARRANTY; without even the implied warranty of
10
 
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11
 
# GNU General Public License for more details.
12
 
#
13
 
# You should have received a copy of the GNU General Public License
14
 
# along with this program; if not, write to the Free Software
15
 
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
16
 
 
17
 
"""Indexing facilities."""
18
 
 
19
 
__all__ = [
20
 
    'CombinedGraphIndex',
21
 
    'GraphIndex',
22
 
    'GraphIndexBuilder',
23
 
    'GraphIndexPrefixAdapter',
24
 
    'InMemoryGraphIndex',
25
 
    ]
26
 
 
27
 
from bisect import bisect_right
28
 
from cStringIO import StringIO
29
 
import re
30
 
 
31
 
from bzrlib.lazy_import import lazy_import
32
 
lazy_import(globals(), """
33
 
from bzrlib import trace
34
 
from bzrlib.bisect_multi import bisect_multi_bytes
35
 
from bzrlib.trace import mutter
36
 
""")
37
 
from bzrlib import debug, errors
38
 
 
39
 
_OPTION_KEY_ELEMENTS = "key_elements="
40
 
_OPTION_LEN = "len="
41
 
_OPTION_NODE_REFS = "node_ref_lists="
42
 
_SIGNATURE = "Bazaar Graph Index 1\n"
43
 
 
44
 
 
45
 
_whitespace_re = re.compile('[\t\n\x0b\x0c\r\x00 ]')
46
 
_newline_null_re = re.compile('[\n\0]')
47
 
 
48
 
 
49
 
class GraphIndexBuilder(object):
50
 
    """A builder that can build a GraphIndex.
51
 
    
52
 
    The resulting graph has the structure:
53
 
    
54
 
    _SIGNATURE OPTIONS NODES NEWLINE
55
 
    _SIGNATURE     := 'Bazaar Graph Index 1' NEWLINE
56
 
    OPTIONS        := 'node_ref_lists=' DIGITS NEWLINE
57
 
    NODES          := NODE*
58
 
    NODE           := KEY NULL ABSENT? NULL REFERENCES NULL VALUE NEWLINE
59
 
    KEY            := Not-whitespace-utf8
60
 
    ABSENT         := 'a'
61
 
    REFERENCES     := REFERENCE_LIST (TAB REFERENCE_LIST){node_ref_lists - 1}
62
 
    REFERENCE_LIST := (REFERENCE (CR REFERENCE)*)?
63
 
    REFERENCE      := DIGITS  ; digits is the byte offset in the index of the
64
 
                              ; referenced key.
65
 
    VALUE          := no-newline-no-null-bytes
66
 
    """
67
 
 
68
 
    def __init__(self, reference_lists=0, key_elements=1):
69
 
        """Create a GraphIndex builder.
70
 
 
71
 
        :param reference_lists: The number of node references lists for each
72
 
            entry.
73
 
        :param key_elements: The number of bytestrings in each key.
74
 
        """
75
 
        self.reference_lists = reference_lists
76
 
        self._keys = set()
77
 
        self._nodes = {}
78
 
        self._nodes_by_key = {}
79
 
        self._key_length = key_elements
80
 
 
81
 
    def _check_key(self, key):
82
 
        """Raise BadIndexKey if key is not a valid key for this index."""
83
 
        if type(key) != tuple:
84
 
            raise errors.BadIndexKey(key)
85
 
        if self._key_length != len(key):
86
 
            raise errors.BadIndexKey(key)
87
 
        for element in key:
88
 
            if not element or _whitespace_re.search(element) is not None:
89
 
                raise errors.BadIndexKey(element)
90
 
 
91
 
    def add_node(self, key, value, references=()):
92
 
        """Add a node to the index.
93
 
 
94
 
        :param key: The key. keys are non-empty tuples containing
95
 
            as many whitespace-free utf8 bytestrings as the key length
96
 
            defined for this index.
97
 
        :param references: An iterable of iterables of keys. Each is a
98
 
            reference to another key.
99
 
        :param value: The value to associate with the key. It may be any
100
 
            bytes as long as it does not contain \0 or \n.
101
 
        """
102
 
        self._check_key(key)
103
 
        if _newline_null_re.search(value) is not None:
104
 
            raise errors.BadIndexValue(value)
105
 
        if len(references) != self.reference_lists:
106
 
            raise errors.BadIndexValue(references)
107
 
        node_refs = []
108
 
        for reference_list in references:
109
 
            for reference in reference_list:
110
 
                self._check_key(reference)
111
 
                if reference not in self._nodes:
112
 
                    self._nodes[reference] = ('a', (), '')
113
 
            node_refs.append(tuple(reference_list))
114
 
        if key in self._nodes and self._nodes[key][0] == '':
115
 
            raise errors.BadIndexDuplicateKey(key, self)
116
 
        self._nodes[key] = ('', tuple(node_refs), value)
117
 
        self._keys.add(key)
118
 
        if self._key_length > 1:
119
 
            key_dict = self._nodes_by_key
120
 
            if self.reference_lists:
121
 
                key_value = key, value, tuple(node_refs)
122
 
            else:
123
 
                key_value = key, value
124
 
            # possibly should do this on-demand, but it seems likely it is 
125
 
            # always wanted
126
 
            # For a key of (foo, bar, baz) create
127
 
            # _nodes_by_key[foo][bar][baz] = key_value
128
 
            for subkey in key[:-1]:
129
 
                key_dict = key_dict.setdefault(subkey, {})
130
 
            key_dict[key[-1]] = key_value
131
 
 
132
 
    def finish(self):
133
 
        lines = [_SIGNATURE]
134
 
        lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n')
135
 
        lines.append(_OPTION_KEY_ELEMENTS + str(self._key_length) + '\n')
136
 
        lines.append(_OPTION_LEN + str(len(self._keys)) + '\n')
137
 
        prefix_length = sum(len(x) for x in lines)
138
 
        # references are byte offsets. To avoid having to do nasty
139
 
        # polynomial work to resolve offsets (references to later in the 
140
 
        # file cannot be determined until all the inbetween references have
141
 
        # been calculated too) we pad the offsets with 0's to make them be
142
 
        # of consistent length. Using binary offsets would break the trivial
143
 
        # file parsing.
144
 
        # to calculate the width of zero's needed we do three passes:
145
 
        # one to gather all the non-reference data and the number of references.
146
 
        # one to pad all the data with reference-length and determine entry
147
 
        # addresses.
148
 
        # One to serialise.
149
 
        
150
 
        # forward sorted by key. In future we may consider topological sorting,
151
 
        # at the cost of table scans for direct lookup, or a second index for
152
 
        # direct lookup
153
 
        nodes = sorted(self._nodes.items())
154
 
        # if we do not prepass, we don't know how long it will be up front.
155
 
        expected_bytes = None
156
 
        # we only need to pre-pass if we have reference lists at all.
157
 
        if self.reference_lists:
158
 
            key_offset_info = []
159
 
            non_ref_bytes = prefix_length
160
 
            total_references = 0
161
 
            # TODO use simple multiplication for the constants in this loop.
162
 
            for key, (absent, references, value) in nodes:
163
 
                # record the offset known *so far* for this key:
164
 
                # the non reference bytes to date, and the total references to
165
 
                # date - saves reaccumulating on the second pass
166
 
                key_offset_info.append((key, non_ref_bytes, total_references))
167
 
                # key is literal, value is literal, there are 3 null's, 1 NL
168
 
                # key is variable length tuple, \x00 between elements
169
 
                non_ref_bytes += sum(len(element) for element in key)
170
 
                if self._key_length > 1:
171
 
                    non_ref_bytes += self._key_length - 1
172
 
                # value is literal bytes, there are 3 null's, 1 NL.
173
 
                non_ref_bytes += len(value) + 3 + 1
174
 
                # one byte for absent if set.
175
 
                if absent:
176
 
                    non_ref_bytes += 1
177
 
                elif self.reference_lists:
178
 
                    # (ref_lists -1) tabs
179
 
                    non_ref_bytes += self.reference_lists - 1
180
 
                    # (ref-1 cr's per ref_list)
181
 
                    for ref_list in references:
182
 
                        # how many references across the whole file?
183
 
                        total_references += len(ref_list)
184
 
                        # accrue reference separators
185
 
                        if ref_list:
186
 
                            non_ref_bytes += len(ref_list) - 1
187
 
            # how many digits are needed to represent the total byte count?
188
 
            digits = 1
189
 
            possible_total_bytes = non_ref_bytes + total_references*digits
190
 
            while 10 ** digits < possible_total_bytes:
191
 
                digits += 1
192
 
                possible_total_bytes = non_ref_bytes + total_references*digits
193
 
            expected_bytes = possible_total_bytes + 1 # terminating newline
194
 
            # resolve key addresses.
195
 
            key_addresses = {}
196
 
            for key, non_ref_bytes, total_references in key_offset_info:
197
 
                key_addresses[key] = non_ref_bytes + total_references*digits
198
 
            # serialise
199
 
            format_string = '%%0%sd' % digits
200
 
        for key, (absent, references, value) in nodes:
201
 
            flattened_references = []
202
 
            for ref_list in references:
203
 
                ref_addresses = []
204
 
                for reference in ref_list:
205
 
                    ref_addresses.append(format_string % key_addresses[reference])
206
 
                flattened_references.append('\r'.join(ref_addresses))
207
 
            string_key = '\x00'.join(key)
208
 
            lines.append("%s\x00%s\x00%s\x00%s\n" % (string_key, absent,
209
 
                '\t'.join(flattened_references), value))
210
 
        lines.append('\n')
211
 
        result = StringIO(''.join(lines))
212
 
        if expected_bytes and len(result.getvalue()) != expected_bytes:
213
 
            raise errors.BzrError('Failed index creation. Internal error:'
214
 
                ' mismatched output length and expected length: %d %d' %
215
 
                (len(result.getvalue()), expected_bytes))
216
 
        return StringIO(''.join(lines))
217
 
 
218
 
 
219
 
class GraphIndex(object):
220
 
    """An index for data with embedded graphs.
221
 
 
222
 
    The index maps keys to a list of key reference lists, and a value.
223
 
    Each node has the same number of key reference lists. Each key reference
224
 
    list can be empty or an arbitrary length. The value is an opaque NULL
225
 
    terminated string without any newlines. The storage of the index is 
226
 
    hidden in the interface: keys and key references are always tuples of
227
 
    bytestrings, never the internal representation (e.g. dictionary offsets).
228
 
 
229
 
    It is presumed that the index will not be mutated - it is static data.
230
 
 
231
 
    Successive iter_all_entries calls will read the entire index each time.
232
 
    Additionally, iter_entries calls will read the index linearly until the
233
 
    desired keys are found. XXX: This must be fixed before the index is
234
 
    suitable for production use. :XXX
235
 
    """
236
 
 
237
 
    def __init__(self, transport, name, size):
238
 
        """Open an index called name on transport.
239
 
 
240
 
        :param transport: A bzrlib.transport.Transport.
241
 
        :param name: A path to provide to transport API calls.
242
 
        :param size: The size of the index in bytes. This is used for bisection
243
 
            logic to perform partial index reads. While the size could be
244
 
            obtained by statting the file this introduced an additional round
245
 
            trip as well as requiring stat'able transports, both of which are
246
 
            avoided by having it supplied. If size is None, then bisection
247
 
            support will be disabled and accessing the index will just stream
248
 
            all the data.
249
 
        """
250
 
        self._transport = transport
251
 
        self._name = name
252
 
        # Becomes a dict of key:(value, reference-list-byte-locations) used by
253
 
        # the bisection interface to store parsed but not resolved keys.
254
 
        self._bisect_nodes = None
255
 
        # Becomes a dict of key:(value, reference-list-keys) which are ready to
256
 
        # be returned directly to callers.
257
 
        self._nodes = None
258
 
        # a sorted list of slice-addresses for the parsed bytes of the file.
259
 
        # e.g. (0,1) would mean that byte 0 is parsed.
260
 
        self._parsed_byte_map = []
261
 
        # a sorted list of keys matching each slice address for parsed bytes
262
 
        # e.g. (None, 'foo@bar') would mean that the first byte contained no
263
 
        # key, and the end byte of the slice is the of the data for 'foo@bar'
264
 
        self._parsed_key_map = []
265
 
        self._key_count = None
266
 
        self._keys_by_offset = None
267
 
        self._nodes_by_key = None
268
 
        self._size = size
269
 
 
270
 
    def _buffer_all(self):
271
 
        """Buffer all the index data.
272
 
 
273
 
        Mutates self._nodes and self.keys_by_offset.
274
 
        """
275
 
        if 'index' in debug.debug_flags:
276
 
            mutter('Reading entire index %s', self._transport.abspath(self._name))
277
 
        stream = self._transport.get(self._name)
278
 
        self._read_prefix(stream)
279
 
        self._expected_elements = 3 + self._key_length
280
 
        line_count = 0
281
 
        # raw data keyed by offset
282
 
        self._keys_by_offset = {}
283
 
        # ready-to-return key:value or key:value, node_ref_lists
284
 
        self._nodes = {}
285
 
        self._nodes_by_key = {}
286
 
        trailers = 0
287
 
        pos = stream.tell()
288
 
        lines = stream.read().split('\n')
289
 
        del lines[-1]
290
 
        _, _, _, trailers = self._parse_lines(lines, pos)
291
 
        for key, absent, references, value in self._keys_by_offset.itervalues():
292
 
            if absent:
293
 
                continue
294
 
            # resolve references:
295
 
            if self.node_ref_lists:
296
 
                node_value = (value, self._resolve_references(references))
297
 
            else:
298
 
                node_value = value
299
 
            self._nodes[key] = node_value
300
 
            if self._key_length > 1:
301
 
                subkey = list(reversed(key[:-1]))
302
 
                key_dict = self._nodes_by_key
303
 
                if self.node_ref_lists:
304
 
                    key_value = key, node_value[0], node_value[1]
305
 
                else:
306
 
                    key_value = key, node_value
307
 
                # possibly should do this on-demand, but it seems likely it is 
308
 
                # always wanted
309
 
                # For a key of (foo, bar, baz) create
310
 
                # _nodes_by_key[foo][bar][baz] = key_value
311
 
                for subkey in key[:-1]:
312
 
                    key_dict = key_dict.setdefault(subkey, {})
313
 
                key_dict[key[-1]] = key_value
314
 
        # cache the keys for quick set intersections
315
 
        self._keys = set(self._nodes)
316
 
        if trailers != 1:
317
 
            # there must be one line - the empty trailer line.
318
 
            raise errors.BadIndexData(self)
319
 
 
320
 
    def iter_all_entries(self):
321
 
        """Iterate over all keys within the index.
322
 
 
323
 
        :return: An iterable of (index, key, value) or (index, key, value, reference_lists).
324
 
            The former tuple is used when there are no reference lists in the
325
 
            index, making the API compatible with simple key:value index types.
326
 
            There is no defined order for the result iteration - it will be in
327
 
            the most efficient order for the index.
328
 
        """
329
 
        if 'evil' in debug.debug_flags:
330
 
            trace.mutter_callsite(3,
331
 
                "iter_all_entries scales with size of history.")
332
 
        if self._nodes is None:
333
 
            self._buffer_all()
334
 
        if self.node_ref_lists:
335
 
            for key, (value, node_ref_lists) in self._nodes.iteritems():
336
 
                yield self, key, value, node_ref_lists
337
 
        else:
338
 
            for key, value in self._nodes.iteritems():
339
 
                yield self, key, value
340
 
 
341
 
    def _read_prefix(self, stream):
342
 
        signature = stream.read(len(self._signature()))
343
 
        if not signature == self._signature():
344
 
            raise errors.BadIndexFormatSignature(self._name, GraphIndex)
345
 
        options_line = stream.readline()
346
 
        if not options_line.startswith(_OPTION_NODE_REFS):
347
 
            raise errors.BadIndexOptions(self)
348
 
        try:
349
 
            self.node_ref_lists = int(options_line[len(_OPTION_NODE_REFS):-1])
350
 
        except ValueError:
351
 
            raise errors.BadIndexOptions(self)
352
 
        options_line = stream.readline()
353
 
        if not options_line.startswith(_OPTION_KEY_ELEMENTS):
354
 
            raise errors.BadIndexOptions(self)
355
 
        try:
356
 
            self._key_length = int(options_line[len(_OPTION_KEY_ELEMENTS):-1])
357
 
        except ValueError:
358
 
            raise errors.BadIndexOptions(self)
359
 
        options_line = stream.readline()
360
 
        if not options_line.startswith(_OPTION_LEN):
361
 
            raise errors.BadIndexOptions(self)
362
 
        try:
363
 
            self._key_count = int(options_line[len(_OPTION_LEN):-1])
364
 
        except ValueError:
365
 
            raise errors.BadIndexOptions(self)
366
 
 
367
 
    def _resolve_references(self, references):
368
 
        """Return the resolved key references for references.
369
 
        
370
 
        References are resolved by looking up the location of the key in the
371
 
        _keys_by_offset map and substituting the key name, preserving ordering.
372
 
 
373
 
        :param references: An iterable of iterables of key locations. e.g. 
374
 
            [[123, 456], [123]]
375
 
        :return: A tuple of tuples of keys.
376
 
        """
377
 
        node_refs = []
378
 
        for ref_list in references:
379
 
            node_refs.append(tuple([self._keys_by_offset[ref][0] for ref in ref_list]))
380
 
        return tuple(node_refs)
381
 
 
382
 
    def _find_index(self, range_map, key):
383
 
        """Helper for the _parsed_*_index calls.
384
 
 
385
 
        Given a range map - [(start, end), ...], finds the index of the range
386
 
        in the map for key if it is in the map, and if it is not there, the
387
 
        immediately preceeding range in the map.
388
 
        """
389
 
        result = bisect_right(range_map, key) - 1
390
 
        if result + 1 < len(range_map):
391
 
            # check the border condition, it may be in result + 1
392
 
            if range_map[result + 1][0] == key[0]:
393
 
                return result + 1
394
 
        return result
395
 
 
396
 
    def _parsed_byte_index(self, offset):
397
 
        """Return the index of the entry immediately before offset.
398
 
 
399
 
        e.g. if the parsed map has regions 0,10 and 11,12 parsed, meaning that
400
 
        there is one unparsed byte (the 11th, addressed as[10]). then:
401
 
        asking for 0 will return 0
402
 
        asking for 10 will return 0
403
 
        asking for 11 will return 1
404
 
        asking for 12 will return 1
405
 
        """
406
 
        key = (offset, 0)
407
 
        return self._find_index(self._parsed_byte_map, key)
408
 
 
409
 
    def _parsed_key_index(self, key):
410
 
        """Return the index of the entry immediately before key.
411
 
 
412
 
        e.g. if the parsed map has regions (None, 'a') and ('b','c') parsed,
413
 
        meaning that keys from None to 'a' inclusive, and 'b' to 'c' inclusive
414
 
        have been parsed, then:
415
 
        asking for '' will return 0
416
 
        asking for 'a' will return 0
417
 
        asking for 'b' will return 1
418
 
        asking for 'e' will return 1
419
 
        """
420
 
        search_key = (key, None)
421
 
        return self._find_index(self._parsed_key_map, search_key)
422
 
 
423
 
    def _is_parsed(self, offset):
424
 
        """Returns True if offset has been parsed."""
425
 
        index = self._parsed_byte_index(offset)
426
 
        if index == len(self._parsed_byte_map):
427
 
            return offset < self._parsed_byte_map[index - 1][1]
428
 
        start, end = self._parsed_byte_map[index]
429
 
        return offset >= start and offset < end
430
 
 
431
 
    def _iter_entries_from_total_buffer(self, keys):
432
 
        """Iterate over keys when the entire index is parsed."""
433
 
        keys = keys.intersection(self._keys)
434
 
        if self.node_ref_lists:
435
 
            for key in keys:
436
 
                value, node_refs = self._nodes[key]
437
 
                yield self, key, value, node_refs
438
 
        else:
439
 
            for key in keys:
440
 
                yield self, key, self._nodes[key]
441
 
 
442
 
    def iter_entries(self, keys):
443
 
        """Iterate over keys within the index.
444
 
 
445
 
        :param keys: An iterable providing the keys to be retrieved.
446
 
        :return: An iterable as per iter_all_entries, but restricted to the
447
 
            keys supplied. No additional keys will be returned, and every
448
 
            key supplied that is in the index will be returned.
449
 
        """
450
 
        # PERFORMANCE TODO: parse and bisect all remaining data at some
451
 
        # threshold of total-index processing/get calling layers that expect to
452
 
        # read the entire index to use the iter_all_entries  method instead.
453
 
        keys = set(keys)
454
 
        if not keys:
455
 
            return []
456
 
        if self._size is None and self._nodes is None:
457
 
            self._buffer_all()
458
 
        if self._nodes is not None:
459
 
            return self._iter_entries_from_total_buffer(keys)
460
 
        else:
461
 
            return (result[1] for result in bisect_multi_bytes(
462
 
                self._lookup_keys_via_location, self._size, keys))
463
 
 
464
 
    def iter_entries_prefix(self, keys):
465
 
        """Iterate over keys within the index using prefix matching.
466
 
 
467
 
        Prefix matching is applied within the tuple of a key, not to within
468
 
        the bytestring of each key element. e.g. if you have the keys ('foo',
469
 
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
470
 
        only the former key is returned.
471
 
 
472
 
        WARNING: Note that this method currently causes a full index parse
473
 
        unconditionally (which is reasonably appropriate as it is a means for
474
 
        thunking many small indices into one larger one and still supplies
475
 
        iter_all_entries at the thunk layer).
476
 
 
477
 
        :param keys: An iterable providing the key prefixes to be retrieved.
478
 
            Each key prefix takes the form of a tuple the length of a key, but
479
 
            with the last N elements 'None' rather than a regular bytestring.
480
 
            The first element cannot be 'None'.
481
 
        :return: An iterable as per iter_all_entries, but restricted to the
482
 
            keys with a matching prefix to those supplied. No additional keys
483
 
            will be returned, and every match that is in the index will be
484
 
            returned.
485
 
        """
486
 
        keys = set(keys)
487
 
        if not keys:
488
 
            return
489
 
        # load data - also finds key lengths
490
 
        if self._nodes is None:
491
 
            self._buffer_all()
492
 
        if self._key_length == 1:
493
 
            for key in keys:
494
 
                # sanity check
495
 
                if key[0] is None:
496
 
                    raise errors.BadIndexKey(key)
497
 
                if len(key) != self._key_length:
498
 
                    raise errors.BadIndexKey(key)
499
 
                if self.node_ref_lists:
500
 
                    value, node_refs = self._nodes[key]
501
 
                    yield self, key, value, node_refs
502
 
                else:
503
 
                    yield self, key, self._nodes[key]
504
 
            return
505
 
        for key in keys:
506
 
            # sanity check
507
 
            if key[0] is None:
508
 
                raise errors.BadIndexKey(key)
509
 
            if len(key) != self._key_length:
510
 
                raise errors.BadIndexKey(key)
511
 
            # find what it refers to:
512
 
            key_dict = self._nodes_by_key
513
 
            elements = list(key)
514
 
            # find the subdict whose contents should be returned.
515
 
            try:
516
 
                while len(elements) and elements[0] is not None:
517
 
                    key_dict = key_dict[elements[0]]
518
 
                    elements.pop(0)
519
 
            except KeyError:
520
 
                # a non-existant lookup.
521
 
                continue
522
 
            if len(elements):
523
 
                dicts = [key_dict]
524
 
                while dicts:
525
 
                    key_dict = dicts.pop(-1)
526
 
                    # can't be empty or would not exist
527
 
                    item, value = key_dict.iteritems().next()
528
 
                    if type(value) == dict:
529
 
                        # push keys 
530
 
                        dicts.extend(key_dict.itervalues())
531
 
                    else:
532
 
                        # yield keys
533
 
                        for value in key_dict.itervalues():
534
 
                            # each value is the key:value:node refs tuple
535
 
                            # ready to yield.
536
 
                            yield (self, ) + value
537
 
            else:
538
 
                # the last thing looked up was a terminal element
539
 
                yield (self, ) + key_dict
540
 
 
541
 
    def key_count(self):
542
 
        """Return an estimate of the number of keys in this index.
543
 
        
544
 
        For GraphIndex the estimate is exact.
545
 
        """
546
 
        if self._key_count is None:
547
 
            # really this should just read the prefix
548
 
            self._buffer_all()
549
 
        return self._key_count
550
 
 
551
 
    def _lookup_keys_via_location(self, location_keys):
552
 
        """Public interface for implementing bisection.
553
 
 
554
 
        If _buffer_all has been called, then all the data for the index is in
555
 
        memory, and this method should not be called, as it uses a separate
556
 
        cache because it cannot pre-resolve all indices, which buffer_all does
557
 
        for performance.
558
 
 
559
 
        :param location_keys: A list of location(byte offset), key tuples.
560
 
        :return: A list of (location_key, result) tuples as expected by
561
 
            bzrlib.bisect_multi.bisect_multi_bytes.
562
 
        """
563
 
        # Possible improvements:
564
 
        #  - only bisect lookup each key once
565
 
        #  - sort the keys first, and use that to reduce the bisection window
566
 
        # ----- 
567
 
        # this progresses in three parts:
568
 
        # read data
569
 
        # parse it
570
 
        # attempt to answer the question from the now in memory data.
571
 
        # build the readv request
572
 
        # for each location, ask for 800 bytes - much more than rows we've seen
573
 
        # anywhere.
574
 
        readv_ranges = []
575
 
        for location, key in location_keys:
576
 
            # can we answer from cache?
577
 
            # - if we know the answer - yes
578
 
            index = self._parsed_key_index(key)
579
 
            if (len(self._parsed_key_map) and 
580
 
                self._parsed_key_map[index][0] <= key and
581
 
                (self._parsed_key_map[index][1] > key or
582
 
                 # end of the file has been parsed
583
 
                 self._parsed_byte_map[index][1] == self._size)):
584
 
                # the key has been parsed, so no lookup is needed
585
 
                continue
586
 
            # - if we have examined this part of the file already - yes
587
 
            index = self._parsed_byte_index(location)
588
 
            if (len(self._parsed_byte_map) and 
589
 
                self._parsed_byte_map[index][0] <= location and
590
 
                self._parsed_byte_map[index][1] > location):
591
 
                # the byte region has been parsed, so no read is needed.
592
 
                continue
593
 
            length = 800
594
 
            if location + length > self._size:
595
 
                length = self._size - location
596
 
            # todo, trim out parsed locations.
597
 
            if length > 0:
598
 
                readv_ranges.append((location, length))
599
 
        # read the header if needed
600
 
        if self._bisect_nodes is None:
601
 
            readv_ranges.append((0, 200))
602
 
        self._read_and_parse(readv_ranges)
603
 
        # generate results:
604
 
        #  - figure out <, >, missing, present
605
 
        #  - result present references so we can return them.
606
 
        result = []
607
 
        # keys that we cannot answer until we resolve references
608
 
        pending_references = []
609
 
        pending_locations = set()
610
 
        for location, key in location_keys:
611
 
            # can we answer from cache?
612
 
            index = self._parsed_key_index(key)
613
 
            if (self._parsed_key_map[index][0] <= key and
614
 
                (self._parsed_key_map[index][1] > key or
615
 
                 # end of the file has been parsed
616
 
                 self._parsed_byte_map[index][1] == self._size)):
617
 
                # the key has been parsed, so no lookup is needed
618
 
                if key in self._bisect_nodes:
619
 
                    if self.node_ref_lists:
620
 
                        # the references may not have been all parsed.
621
 
                        value, refs = self._bisect_nodes[key]
622
 
                        wanted_locations = []
623
 
                        for ref_list in refs:
624
 
                            for ref in ref_list:
625
 
                                if ref not in self._keys_by_offset:
626
 
                                    wanted_locations.append(ref)
627
 
                        if wanted_locations:
628
 
                            pending_locations.update(wanted_locations)
629
 
                            pending_references.append((location, key))
630
 
                            continue
631
 
                        result.append(((location, key), (self, key,
632
 
                            value, self._resolve_references(refs))))
633
 
                    else:
634
 
                        result.append(((location, key),
635
 
                            (self, key, self._bisect_nodes[key])))
636
 
                else:
637
 
                    result.append(((location, key), False))
638
 
                continue
639
 
            # no, is the key above or below the probed location:
640
 
            # get the range of the probed & parsed location
641
 
            index = self._parsed_byte_index(location)
642
 
            # if the key is below the start of the range, its below
643
 
            if key < self._parsed_key_map[index][0]:
644
 
                direction = -1
645
 
            else:
646
 
                direction = +1
647
 
            result.append(((location, key), direction))
648
 
        readv_ranges = []
649
 
        # lookup data to resolve references
650
 
        for location in pending_locations:
651
 
            length = 800
652
 
            if location + length > self._size:
653
 
                length = self._size - location
654
 
            # TODO: trim out parsed locations (e.g. if the 800 is into the
655
 
            # parsed region trim it, and dont use the adjust_for_latency
656
 
            # facility)
657
 
            if length > 0:
658
 
                readv_ranges.append((location, length))
659
 
        self._read_and_parse(readv_ranges)
660
 
        for location, key in pending_references:
661
 
            # answer key references we had to look-up-late.
662
 
            index = self._parsed_key_index(key)
663
 
            value, refs = self._bisect_nodes[key]
664
 
            result.append(((location, key), (self, key,
665
 
                value, self._resolve_references(refs))))
666
 
        return result
667
 
 
668
 
    def _parse_header_from_bytes(self, bytes):
669
 
        """Parse the header from a region of bytes.
670
 
 
671
 
        :param bytes: The data to parse.
672
 
        :return: An offset, data tuple such as readv yields, for the unparsed
673
 
            data. (which may length 0).
674
 
        """
675
 
        signature = bytes[0:len(self._signature())]
676
 
        if not signature == self._signature():
677
 
            raise errors.BadIndexFormatSignature(self._name, GraphIndex)
678
 
        lines = bytes[len(self._signature()):].splitlines()
679
 
        options_line = lines[0]
680
 
        if not options_line.startswith(_OPTION_NODE_REFS):
681
 
            raise errors.BadIndexOptions(self)
682
 
        try:
683
 
            self.node_ref_lists = int(options_line[len(_OPTION_NODE_REFS):])
684
 
        except ValueError:
685
 
            raise errors.BadIndexOptions(self)
686
 
        options_line = lines[1]
687
 
        if not options_line.startswith(_OPTION_KEY_ELEMENTS):
688
 
            raise errors.BadIndexOptions(self)
689
 
        try:
690
 
            self._key_length = int(options_line[len(_OPTION_KEY_ELEMENTS):])
691
 
        except ValueError:
692
 
            raise errors.BadIndexOptions(self)
693
 
        options_line = lines[2]
694
 
        if not options_line.startswith(_OPTION_LEN):
695
 
            raise errors.BadIndexOptions(self)
696
 
        try:
697
 
            self._key_count = int(options_line[len(_OPTION_LEN):])
698
 
        except ValueError:
699
 
            raise errors.BadIndexOptions(self)
700
 
        # calculate the bytes we have processed
701
 
        header_end = (len(signature) + len(lines[0]) + len(lines[1]) +
702
 
            len(lines[2]) + 3)
703
 
        self._parsed_bytes(0, None, header_end, None)
704
 
        # setup parsing state
705
 
        self._expected_elements = 3 + self._key_length
706
 
        # raw data keyed by offset
707
 
        self._keys_by_offset = {}
708
 
        # keys with the value and node references
709
 
        self._bisect_nodes = {}
710
 
        return header_end, bytes[header_end:]
711
 
 
712
 
    def _parse_region(self, offset, data):
713
 
        """Parse node data returned from a readv operation.
714
 
 
715
 
        :param offset: The byte offset the data starts at.
716
 
        :param data: The data to parse.
717
 
        """
718
 
        # trim the data.
719
 
        # end first:
720
 
        end = offset + len(data)
721
 
        high_parsed = offset
722
 
        while True:
723
 
            # Trivial test - if the current index's end is within the
724
 
            # low-matching parsed range, we're done.
725
 
            index = self._parsed_byte_index(high_parsed)
726
 
            if end < self._parsed_byte_map[index][1]:
727
 
                return
728
 
            # print "[%d:%d]" % (offset, end), \
729
 
            #     self._parsed_byte_map[index:index + 2]
730
 
            high_parsed, last_segment = self._parse_segment(
731
 
                offset, data, end, index)
732
 
            if last_segment:
733
 
                return
734
 
 
735
 
    def _parse_segment(self, offset, data, end, index):
736
 
        """Parse one segment of data.
737
 
 
738
 
        :param offset: Where 'data' begins in the file.
739
 
        :param data: Some data to parse a segment of.
740
 
        :param end: Where data ends
741
 
        :param index: The current index into the parsed bytes map.
742
 
        :return: True if the parsed segment is the last possible one in the
743
 
            range of data.
744
 
        :return: high_parsed_byte, last_segment.
745
 
            high_parsed_byte is the location of the highest parsed byte in this
746
 
            segment, last_segment is True if the parsed segment is the last
747
 
            possible one in the data block.
748
 
        """
749
 
        # default is to use all data
750
 
        trim_end = None
751
 
        # accomodate overlap with data before this.
752
 
        if offset < self._parsed_byte_map[index][1]:
753
 
            # overlaps the lower parsed region
754
 
            # skip the parsed data
755
 
            trim_start = self._parsed_byte_map[index][1] - offset
756
 
            # don't trim the start for \n
757
 
            start_adjacent = True
758
 
        elif offset == self._parsed_byte_map[index][1]:
759
 
            # abuts the lower parsed region
760
 
            # use all data
761
 
            trim_start = None
762
 
            # do not trim anything
763
 
            start_adjacent = True
764
 
        else:
765
 
            # does not overlap the lower parsed region
766
 
            # use all data
767
 
            trim_start = None
768
 
            # but trim the leading \n
769
 
            start_adjacent = False
770
 
        if end == self._size:
771
 
            # lines up to the end of all data:
772
 
            # use it all
773
 
            trim_end = None
774
 
            # do not strip to the last \n
775
 
            end_adjacent = True
776
 
            last_segment = True
777
 
        elif index + 1 == len(self._parsed_byte_map):
778
 
            # at the end of the parsed data
779
 
            # use it all
780
 
            trim_end = None
781
 
            # but strip to the last \n
782
 
            end_adjacent = False
783
 
            last_segment = True
784
 
        elif end == self._parsed_byte_map[index + 1][0]:
785
 
            # buts up against the next parsed region
786
 
            # use it all
787
 
            trim_end = None
788
 
            # do not strip to the last \n
789
 
            end_adjacent = True
790
 
            last_segment = True
791
 
        elif end > self._parsed_byte_map[index + 1][0]:
792
 
            # overlaps into the next parsed region
793
 
            # only consider the unparsed data
794
 
            trim_end = self._parsed_byte_map[index + 1][0] - offset
795
 
            # do not strip to the last \n as we know its an entire record
796
 
            end_adjacent = True
797
 
            last_segment = end < self._parsed_byte_map[index + 1][1]
798
 
        else:
799
 
            # does not overlap into the next region
800
 
            # use it all
801
 
            trim_end = None
802
 
            # but strip to the last \n
803
 
            end_adjacent = False
804
 
            last_segment = True
805
 
        # now find bytes to discard if needed
806
 
        if not start_adjacent:
807
 
            # work around python bug in rfind
808
 
            if trim_start is None:
809
 
                trim_start = data.find('\n') + 1
810
 
            else:
811
 
                trim_start = data.find('\n', trim_start) + 1
812
 
            assert trim_start != 0, 'no \n was present'
813
 
            # print 'removing start', offset, trim_start, repr(data[:trim_start])
814
 
        if not end_adjacent:
815
 
            # work around python bug in rfind
816
 
            if trim_end is None:
817
 
                trim_end = data.rfind('\n') + 1
818
 
            else:
819
 
                trim_end = data.rfind('\n', None, trim_end) + 1
820
 
            assert trim_end != 0, 'no \n was present'
821
 
            # print 'removing end', offset, trim_end, repr(data[trim_end:])
822
 
        # adjust offset and data to the parseable data.
823
 
        trimmed_data = data[trim_start:trim_end]
824
 
        assert trimmed_data, 'read unneeded data [%d:%d] from [%d:%d]' % (
825
 
            trim_start, trim_end, offset, offset + len(data))
826
 
        if trim_start:
827
 
            offset += trim_start
828
 
        # print "parsing", repr(trimmed_data)
829
 
        # splitlines mangles the \r delimiters.. don't use it.
830
 
        lines = trimmed_data.split('\n')
831
 
        del lines[-1]
832
 
        pos = offset
833
 
        first_key, last_key, nodes, _ = self._parse_lines(lines, pos)
834
 
        for key, value in nodes:
835
 
            self._bisect_nodes[key] = value
836
 
        self._parsed_bytes(offset, first_key,
837
 
            offset + len(trimmed_data), last_key)
838
 
        return offset + len(trimmed_data), last_segment
839
 
 
840
 
    def _parse_lines(self, lines, pos):
841
 
        key = None
842
 
        first_key = None
843
 
        trailers = 0
844
 
        nodes = []
845
 
        for line in lines:
846
 
            if line == '':
847
 
                # must be at the end
848
 
                if self._size:
849
 
                    assert self._size == pos + 1, "%s %s" % (self._size, pos)
850
 
                trailers += 1
851
 
                continue
852
 
            elements = line.split('\0')
853
 
            if len(elements) != self._expected_elements:
854
 
                raise errors.BadIndexData(self)
855
 
            # keys are tuples
856
 
            key = tuple(elements[:self._key_length])
857
 
            if first_key is None:
858
 
                first_key = key
859
 
            absent, references, value = elements[-3:]
860
 
            ref_lists = []
861
 
            for ref_string in references.split('\t'):
862
 
                ref_lists.append(tuple([
863
 
                    int(ref) for ref in ref_string.split('\r') if ref
864
 
                    ]))
865
 
            ref_lists = tuple(ref_lists)
866
 
            self._keys_by_offset[pos] = (key, absent, ref_lists, value)
867
 
            pos += len(line) + 1 # +1 for the \n
868
 
            if absent:
869
 
                continue
870
 
            if self.node_ref_lists:
871
 
                node_value = (value, ref_lists)
872
 
            else:
873
 
                node_value = value
874
 
            nodes.append((key, node_value))
875
 
            # print "parsed ", key
876
 
        return first_key, key, nodes, trailers
877
 
 
878
 
    def _parsed_bytes(self, start, start_key, end, end_key):
879
 
        """Mark the bytes from start to end as parsed.
880
 
 
881
 
        Calling self._parsed_bytes(1,2) will mark one byte (the one at offset
882
 
        1) as parsed.
883
 
 
884
 
        :param start: The start of the parsed region.
885
 
        :param end: The end of the parsed region.
886
 
        """
887
 
        index = self._parsed_byte_index(start)
888
 
        new_value = (start, end)
889
 
        new_key = (start_key, end_key)
890
 
        if index == -1:
891
 
            # first range parsed is always the beginning.
892
 
            self._parsed_byte_map.insert(index, new_value)
893
 
            self._parsed_key_map.insert(index, new_key)
894
 
            return
895
 
        # four cases:
896
 
        # new region
897
 
        # extend lower region
898
 
        # extend higher region
899
 
        # combine two regions
900
 
        if (index + 1 < len(self._parsed_byte_map) and
901
 
            self._parsed_byte_map[index][1] == start and
902
 
            self._parsed_byte_map[index + 1][0] == end):
903
 
            # combine two regions
904
 
            self._parsed_byte_map[index] = (self._parsed_byte_map[index][0],
905
 
                self._parsed_byte_map[index + 1][1])
906
 
            self._parsed_key_map[index] = (self._parsed_key_map[index][0],
907
 
                self._parsed_key_map[index + 1][1])
908
 
            del self._parsed_byte_map[index + 1]
909
 
            del self._parsed_key_map[index + 1]
910
 
        elif self._parsed_byte_map[index][1] == start:
911
 
            # extend the lower entry
912
 
            self._parsed_byte_map[index] = (
913
 
                self._parsed_byte_map[index][0], end)
914
 
            self._parsed_key_map[index] = (
915
 
                self._parsed_key_map[index][0], end_key)
916
 
        elif (index + 1 < len(self._parsed_byte_map) and
917
 
            self._parsed_byte_map[index + 1][0] == end):
918
 
            # extend the higher entry
919
 
            self._parsed_byte_map[index + 1] = (
920
 
                start, self._parsed_byte_map[index + 1][1])
921
 
            self._parsed_key_map[index + 1] = (
922
 
                start_key, self._parsed_key_map[index + 1][1])
923
 
        else:
924
 
            # new entry
925
 
            self._parsed_byte_map.insert(index + 1, new_value)
926
 
            self._parsed_key_map.insert(index + 1, new_key)
927
 
 
928
 
    def _read_and_parse(self, readv_ranges):
929
 
        """Read the the ranges and parse the resulting data.
930
 
 
931
 
        :param readv_ranges: A prepared readv range list.
932
 
        """
933
 
        if readv_ranges:
934
 
            readv_data = self._transport.readv(self._name, readv_ranges, True,
935
 
                self._size)
936
 
            # parse
937
 
            for offset, data in readv_data:
938
 
                if self._bisect_nodes is None:
939
 
                    # this must be the start
940
 
                    assert offset == 0
941
 
                    offset, data = self._parse_header_from_bytes(data)
942
 
                # print readv_ranges, "[%d:%d]" % (offset, offset + len(data))
943
 
                self._parse_region(offset, data)
944
 
 
945
 
    def _signature(self):
946
 
        """The file signature for this index type."""
947
 
        return _SIGNATURE
948
 
 
949
 
    def validate(self):
950
 
        """Validate that everything in the index can be accessed."""
951
 
        # iter_all validates completely at the moment, so just do that.
952
 
        for node in self.iter_all_entries():
953
 
            pass
954
 
 
955
 
 
956
 
class CombinedGraphIndex(object):
957
 
    """A GraphIndex made up from smaller GraphIndices.
958
 
    
959
 
    The backing indices must implement GraphIndex, and are presumed to be
960
 
    static data.
961
 
 
962
 
    Queries against the combined index will be made against the first index,
963
 
    and then the second and so on. The order of index's can thus influence
964
 
    performance significantly. For example, if one index is on local disk and a
965
 
    second on a remote server, the local disk index should be before the other
966
 
    in the index list.
967
 
    """
968
 
 
969
 
    def __init__(self, indices):
970
 
        """Create a CombinedGraphIndex backed by indices.
971
 
 
972
 
        :param indices: An ordered list of indices to query for data.
973
 
        """
974
 
        self._indices = indices
975
 
 
976
 
    def __repr__(self):
977
 
        return "%s(%s)" % (
978
 
                self.__class__.__name__,
979
 
                ', '.join(map(repr, self._indices)))
980
 
 
981
 
    def insert_index(self, pos, index):
982
 
        """Insert a new index in the list of indices to query.
983
 
 
984
 
        :param pos: The position to insert the index.
985
 
        :param index: The index to insert.
986
 
        """
987
 
        self._indices.insert(pos, index)
988
 
 
989
 
    def iter_all_entries(self):
990
 
        """Iterate over all keys within the index
991
 
 
992
 
        Duplicate keys across child indices are presumed to have the same
993
 
        value and are only reported once.
994
 
 
995
 
        :return: An iterable of (index, key, reference_lists, value).
996
 
            There is no defined order for the result iteration - it will be in
997
 
            the most efficient order for the index.
998
 
        """
999
 
        seen_keys = set()
1000
 
        for index in self._indices:
1001
 
            for node in index.iter_all_entries():
1002
 
                if node[1] not in seen_keys:
1003
 
                    yield node
1004
 
                    seen_keys.add(node[1])
1005
 
 
1006
 
    def iter_entries(self, keys):
1007
 
        """Iterate over keys within the index.
1008
 
 
1009
 
        Duplicate keys across child indices are presumed to have the same
1010
 
        value and are only reported once.
1011
 
 
1012
 
        :param keys: An iterable providing the keys to be retrieved.
1013
 
        :return: An iterable of (index, key, reference_lists, value). There is no
1014
 
            defined order for the result iteration - it will be in the most
1015
 
            efficient order for the index.
1016
 
        """
1017
 
        keys = set(keys)
1018
 
        for index in self._indices:
1019
 
            if not keys:
1020
 
                return
1021
 
            for node in index.iter_entries(keys):
1022
 
                keys.remove(node[1])
1023
 
                yield node
1024
 
 
1025
 
    def iter_entries_prefix(self, keys):
1026
 
        """Iterate over keys within the index using prefix matching.
1027
 
 
1028
 
        Duplicate keys across child indices are presumed to have the same
1029
 
        value and are only reported once.
1030
 
 
1031
 
        Prefix matching is applied within the tuple of a key, not to within
1032
 
        the bytestring of each key element. e.g. if you have the keys ('foo',
1033
 
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
1034
 
        only the former key is returned.
1035
 
 
1036
 
        :param keys: An iterable providing the key prefixes to be retrieved.
1037
 
            Each key prefix takes the form of a tuple the length of a key, but
1038
 
            with the last N elements 'None' rather than a regular bytestring.
1039
 
            The first element cannot be 'None'.
1040
 
        :return: An iterable as per iter_all_entries, but restricted to the
1041
 
            keys with a matching prefix to those supplied. No additional keys
1042
 
            will be returned, and every match that is in the index will be
1043
 
            returned.
1044
 
        """
1045
 
        keys = set(keys)
1046
 
        if not keys:
1047
 
            return
1048
 
        seen_keys = set()
1049
 
        for index in self._indices:
1050
 
            for node in index.iter_entries_prefix(keys):
1051
 
                if node[1] in seen_keys:
1052
 
                    continue
1053
 
                seen_keys.add(node[1])
1054
 
                yield node
1055
 
 
1056
 
    def key_count(self):
1057
 
        """Return an estimate of the number of keys in this index.
1058
 
        
1059
 
        For CombinedGraphIndex this is approximated by the sum of the keys of
1060
 
        the child indices. As child indices may have duplicate keys this can
1061
 
        have a maximum error of the number of child indices * largest number of
1062
 
        keys in any index.
1063
 
        """
1064
 
        return sum((index.key_count() for index in self._indices), 0)
1065
 
 
1066
 
    def validate(self):
1067
 
        """Validate that everything in the index can be accessed."""
1068
 
        for index in self._indices:
1069
 
            index.validate()
1070
 
 
1071
 
 
1072
 
class InMemoryGraphIndex(GraphIndexBuilder):
1073
 
    """A GraphIndex which operates entirely out of memory and is mutable.
1074
 
 
1075
 
    This is designed to allow the accumulation of GraphIndex entries during a
1076
 
    single write operation, where the accumulated entries need to be immediately
1077
 
    available - for example via a CombinedGraphIndex.
1078
 
    """
1079
 
 
1080
 
    def add_nodes(self, nodes):
1081
 
        """Add nodes to the index.
1082
 
 
1083
 
        :param nodes: An iterable of (key, node_refs, value) entries to add.
1084
 
        """
1085
 
        if self.reference_lists:
1086
 
            for (key, value, node_refs) in nodes:
1087
 
                self.add_node(key, value, node_refs)
1088
 
        else:
1089
 
            for (key, value) in nodes:
1090
 
                self.add_node(key, value)
1091
 
 
1092
 
    def iter_all_entries(self):
1093
 
        """Iterate over all keys within the index
1094
 
 
1095
 
        :return: An iterable of (index, key, reference_lists, value). There is no
1096
 
            defined order for the result iteration - it will be in the most
1097
 
            efficient order for the index (in this case dictionary hash order).
1098
 
        """
1099
 
        if 'evil' in debug.debug_flags:
1100
 
            trace.mutter_callsite(3,
1101
 
                "iter_all_entries scales with size of history.")
1102
 
        if self.reference_lists:
1103
 
            for key, (absent, references, value) in self._nodes.iteritems():
1104
 
                if not absent:
1105
 
                    yield self, key, value, references
1106
 
        else:
1107
 
            for key, (absent, references, value) in self._nodes.iteritems():
1108
 
                if not absent:
1109
 
                    yield self, key, value
1110
 
 
1111
 
    def iter_entries(self, keys):
1112
 
        """Iterate over keys within the index.
1113
 
 
1114
 
        :param keys: An iterable providing the keys to be retrieved.
1115
 
        :return: An iterable of (index, key, reference_lists, value). There is no
1116
 
            defined order for the result iteration - it will be in the most
1117
 
            efficient order for the index (keys iteration order in this case).
1118
 
        """
1119
 
        keys = set(keys)
1120
 
        if self.reference_lists:
1121
 
            for key in keys.intersection(self._keys):
1122
 
                node = self._nodes[key]
1123
 
                if not node[0]:
1124
 
                    yield self, key, node[2], node[1]
1125
 
        else:
1126
 
            for key in keys.intersection(self._keys):
1127
 
                node = self._nodes[key]
1128
 
                if not node[0]:
1129
 
                    yield self, key, node[2]
1130
 
 
1131
 
    def iter_entries_prefix(self, keys):
1132
 
        """Iterate over keys within the index using prefix matching.
1133
 
 
1134
 
        Prefix matching is applied within the tuple of a key, not to within
1135
 
        the bytestring of each key element. e.g. if you have the keys ('foo',
1136
 
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
1137
 
        only the former key is returned.
1138
 
 
1139
 
        :param keys: An iterable providing the key prefixes to be retrieved.
1140
 
            Each key prefix takes the form of a tuple the length of a key, but
1141
 
            with the last N elements 'None' rather than a regular bytestring.
1142
 
            The first element cannot be 'None'.
1143
 
        :return: An iterable as per iter_all_entries, but restricted to the
1144
 
            keys with a matching prefix to those supplied. No additional keys
1145
 
            will be returned, and every match that is in the index will be
1146
 
            returned.
1147
 
        """
1148
 
        # XXX: To much duplication with the GraphIndex class; consider finding
1149
 
        # a good place to pull out the actual common logic.
1150
 
        keys = set(keys)
1151
 
        if not keys:
1152
 
            return
1153
 
        if self._key_length == 1:
1154
 
            for key in keys:
1155
 
                # sanity check
1156
 
                if key[0] is None:
1157
 
                    raise errors.BadIndexKey(key)
1158
 
                if len(key) != self._key_length:
1159
 
                    raise errors.BadIndexKey(key)
1160
 
                node = self._nodes[key]
1161
 
                if node[0]:
1162
 
                    continue 
1163
 
                if self.reference_lists:
1164
 
                    yield self, key, node[2], node[1]
1165
 
                else:
1166
 
                    yield self, key, node[2]
1167
 
            return
1168
 
        for key in keys:
1169
 
            # sanity check
1170
 
            if key[0] is None:
1171
 
                raise errors.BadIndexKey(key)
1172
 
            if len(key) != self._key_length:
1173
 
                raise errors.BadIndexKey(key)
1174
 
            # find what it refers to:
1175
 
            key_dict = self._nodes_by_key
1176
 
            elements = list(key)
1177
 
            # find the subdict to return
1178
 
            try:
1179
 
                while len(elements) and elements[0] is not None:
1180
 
                    key_dict = key_dict[elements[0]]
1181
 
                    elements.pop(0)
1182
 
            except KeyError:
1183
 
                # a non-existant lookup.
1184
 
                continue
1185
 
            if len(elements):
1186
 
                dicts = [key_dict]
1187
 
                while dicts:
1188
 
                    key_dict = dicts.pop(-1)
1189
 
                    # can't be empty or would not exist
1190
 
                    item, value = key_dict.iteritems().next()
1191
 
                    if type(value) == dict:
1192
 
                        # push keys 
1193
 
                        dicts.extend(key_dict.itervalues())
1194
 
                    else:
1195
 
                        # yield keys
1196
 
                        for value in key_dict.itervalues():
1197
 
                            yield (self, ) + value
1198
 
            else:
1199
 
                yield (self, ) + key_dict
1200
 
 
1201
 
    def key_count(self):
1202
 
        """Return an estimate of the number of keys in this index.
1203
 
        
1204
 
        For InMemoryGraphIndex the estimate is exact.
1205
 
        """
1206
 
        return len(self._keys)
1207
 
 
1208
 
    def validate(self):
1209
 
        """In memory index's have no known corruption at the moment."""
1210
 
 
1211
 
 
1212
 
class GraphIndexPrefixAdapter(object):
1213
 
    """An adapter between GraphIndex with different key lengths.
1214
 
 
1215
 
    Queries against this will emit queries against the adapted Graph with the
1216
 
    prefix added, queries for all items use iter_entries_prefix. The returned
1217
 
    nodes will have their keys and node references adjusted to remove the 
1218
 
    prefix. Finally, an add_nodes_callback can be supplied - when called the
1219
 
    nodes and references being added will have prefix prepended.
1220
 
    """
1221
 
 
1222
 
    def __init__(self, adapted, prefix, missing_key_length,
1223
 
        add_nodes_callback=None):
1224
 
        """Construct an adapter against adapted with prefix."""
1225
 
        self.adapted = adapted
1226
 
        self.prefix_key = prefix + (None,)*missing_key_length
1227
 
        self.prefix = prefix
1228
 
        self.prefix_len = len(prefix)
1229
 
        self.add_nodes_callback = add_nodes_callback
1230
 
 
1231
 
    def add_nodes(self, nodes):
1232
 
        """Add nodes to the index.
1233
 
 
1234
 
        :param nodes: An iterable of (key, node_refs, value) entries to add.
1235
 
        """
1236
 
        # save nodes in case its an iterator
1237
 
        nodes = tuple(nodes)
1238
 
        translated_nodes = []
1239
 
        try:
1240
 
            # Add prefix_key to each reference node_refs is a tuple of tuples,
1241
 
            # so split it apart, and add prefix_key to the internal reference
1242
 
            for (key, value, node_refs) in nodes:
1243
 
                adjusted_references = (
1244
 
                    tuple(tuple(self.prefix + ref_node for ref_node in ref_list)
1245
 
                        for ref_list in node_refs))
1246
 
                translated_nodes.append((self.prefix + key, value,
1247
 
                    adjusted_references))
1248
 
        except ValueError:
1249
 
            # XXX: TODO add an explicit interface for getting the reference list
1250
 
            # status, to handle this bit of user-friendliness in the API more 
1251
 
            # explicitly.
1252
 
            for (key, value) in nodes:
1253
 
                translated_nodes.append((self.prefix + key, value))
1254
 
        self.add_nodes_callback(translated_nodes)
1255
 
 
1256
 
    def add_node(self, key, value, references=()):
1257
 
        """Add a node to the index.
1258
 
 
1259
 
        :param key: The key. keys are non-empty tuples containing
1260
 
            as many whitespace-free utf8 bytestrings as the key length
1261
 
            defined for this index.
1262
 
        :param references: An iterable of iterables of keys. Each is a
1263
 
            reference to another key.
1264
 
        :param value: The value to associate with the key. It may be any
1265
 
            bytes as long as it does not contain \0 or \n.
1266
 
        """
1267
 
        self.add_nodes(((key, value, references), ))
1268
 
 
1269
 
    def _strip_prefix(self, an_iter):
1270
 
        """Strip prefix data from nodes and return it."""
1271
 
        for node in an_iter:
1272
 
            # cross checks
1273
 
            if node[1][:self.prefix_len] != self.prefix:
1274
 
                raise errors.BadIndexData(self)
1275
 
            for ref_list in node[3]:
1276
 
                for ref_node in ref_list:
1277
 
                    if ref_node[:self.prefix_len] != self.prefix:
1278
 
                        raise errors.BadIndexData(self)
1279
 
            yield node[0], node[1][self.prefix_len:], node[2], (
1280
 
                tuple(tuple(ref_node[self.prefix_len:] for ref_node in ref_list)
1281
 
                for ref_list in node[3]))
1282
 
 
1283
 
    def iter_all_entries(self):
1284
 
        """Iterate over all keys within the index
1285
 
 
1286
 
        iter_all_entries is implemented against the adapted index using
1287
 
        iter_entries_prefix.
1288
 
 
1289
 
        :return: An iterable of (index, key, reference_lists, value). There is no
1290
 
            defined order for the result iteration - it will be in the most
1291
 
            efficient order for the index (in this case dictionary hash order).
1292
 
        """
1293
 
        return self._strip_prefix(self.adapted.iter_entries_prefix([self.prefix_key]))
1294
 
 
1295
 
    def iter_entries(self, keys):
1296
 
        """Iterate over keys within the index.
1297
 
 
1298
 
        :param keys: An iterable providing the keys to be retrieved.
1299
 
        :return: An iterable of (key, reference_lists, value). There is no
1300
 
            defined order for the result iteration - it will be in the most
1301
 
            efficient order for the index (keys iteration order in this case).
1302
 
        """
1303
 
        return self._strip_prefix(self.adapted.iter_entries(
1304
 
            self.prefix + key for key in keys))
1305
 
 
1306
 
    def iter_entries_prefix(self, keys):
1307
 
        """Iterate over keys within the index using prefix matching.
1308
 
 
1309
 
        Prefix matching is applied within the tuple of a key, not to within
1310
 
        the bytestring of each key element. e.g. if you have the keys ('foo',
1311
 
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
1312
 
        only the former key is returned.
1313
 
 
1314
 
        :param keys: An iterable providing the key prefixes to be retrieved.
1315
 
            Each key prefix takes the form of a tuple the length of a key, but
1316
 
            with the last N elements 'None' rather than a regular bytestring.
1317
 
            The first element cannot be 'None'.
1318
 
        :return: An iterable as per iter_all_entries, but restricted to the
1319
 
            keys with a matching prefix to those supplied. No additional keys
1320
 
            will be returned, and every match that is in the index will be
1321
 
            returned.
1322
 
        """
1323
 
        return self._strip_prefix(self.adapted.iter_entries_prefix(
1324
 
            self.prefix + key for key in keys))
1325
 
 
1326
 
    def key_count(self):
1327
 
        """Return an estimate of the number of keys in this index.
1328
 
        
1329
 
        For GraphIndexPrefixAdapter this is relatively expensive - key
1330
 
        iteration with the prefix is done.
1331
 
        """
1332
 
        return len(list(self.iter_all_entries()))
1333
 
 
1334
 
    def validate(self):
1335
 
        """Call the adapted's validate."""
1336
 
        self.adapted.validate()