~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/index.py

  • Committer: Ian Clatworthy
  • Date: 2007-09-10 06:27:40 UTC
  • mto: (2818.1.1 ianc-integration)
  • mto: This revision was merged to the branch mainline in revision 2819.
  • Revision ID: ian.clatworthy@internode.on.net-20070910062740-kkj4776w6snhf0vl
more win32 path friendliness

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 cStringIO import StringIO
 
28
import re
 
29
 
 
30
from bzrlib.lazy_import import lazy_import
 
31
lazy_import(globals(), """
 
32
from bzrlib import trace
 
33
from bzrlib.trace import mutter
 
34
""")
 
35
from bzrlib import debug, errors
 
36
 
 
37
_OPTION_KEY_ELEMENTS = "key_elements="
 
38
_OPTION_LEN = "len="
 
39
_OPTION_NODE_REFS = "node_ref_lists="
 
40
_SIGNATURE = "Bazaar Graph Index 1\n"
 
41
 
 
42
 
 
43
_whitespace_re = re.compile('[\t\n\x0b\x0c\r\x00 ]')
 
44
_newline_null_re = re.compile('[\n\0]')
 
45
 
 
46
 
 
47
class GraphIndexBuilder(object):
 
48
    """A builder that can build a GraphIndex.
 
49
    
 
50
    The resulting graph has the structure:
 
51
    
 
52
    _SIGNATURE OPTIONS NODES NEWLINE
 
53
    _SIGNATURE     := 'Bazaar Graph Index 1' NEWLINE
 
54
    OPTIONS        := 'node_ref_lists=' DIGITS NEWLINE
 
55
    NODES          := NODE*
 
56
    NODE           := KEY NULL ABSENT? NULL REFERENCES NULL VALUE NEWLINE
 
57
    KEY            := Not-whitespace-utf8
 
58
    ABSENT         := 'a'
 
59
    REFERENCES     := REFERENCE_LIST (TAB REFERENCE_LIST){node_ref_lists - 1}
 
60
    REFERENCE_LIST := (REFERENCE (CR REFERENCE)*)?
 
61
    REFERENCE      := DIGITS  ; digits is the byte offset in the index of the
 
62
                              ; referenced key.
 
63
    VALUE          := no-newline-no-null-bytes
 
64
    """
 
65
 
 
66
    def __init__(self, reference_lists=0, key_elements=1):
 
67
        """Create a GraphIndex builder.
 
68
 
 
69
        :param reference_lists: The number of node references lists for each
 
70
            entry.
 
71
        :param key_elements: The number of bytestrings in each key.
 
72
        """
 
73
        self.reference_lists = reference_lists
 
74
        self._keys = set()
 
75
        self._nodes = {}
 
76
        self._nodes_by_key = {}
 
77
        self._key_length = key_elements
 
78
 
 
79
    def _check_key(self, key):
 
80
        """Raise BadIndexKey if key is not a valid key for this index."""
 
81
        if type(key) != tuple:
 
82
            raise errors.BadIndexKey(key)
 
83
        if self._key_length != len(key):
 
84
            raise errors.BadIndexKey(key)
 
85
        for element in key:
 
86
            if not element or _whitespace_re.search(element) is not None:
 
87
                raise errors.BadIndexKey(element)
 
88
 
 
89
    def add_node(self, key, value, references=()):
 
90
        """Add a node to the index.
 
91
 
 
92
        :param key: The key. keys are non-empty tuples containing
 
93
            as many whitespace-free utf8 bytestrings as the key length
 
94
            defined for this index.
 
95
        :param references: An iterable of iterables of keys. Each is a
 
96
            reference to another key.
 
97
        :param value: The value to associate with the key. It may be any
 
98
            bytes as long as it does not contain \0 or \n.
 
99
        """
 
100
        self._check_key(key)
 
101
        if _newline_null_re.search(value) is not None:
 
102
            raise errors.BadIndexValue(value)
 
103
        if len(references) != self.reference_lists:
 
104
            raise errors.BadIndexValue(references)
 
105
        node_refs = []
 
106
        for reference_list in references:
 
107
            for reference in reference_list:
 
108
                self._check_key(reference)
 
109
                if reference not in self._nodes:
 
110
                    self._nodes[reference] = ('a', (), '')
 
111
            node_refs.append(tuple(reference_list))
 
112
        if key in self._nodes and self._nodes[key][0] == '':
 
113
            raise errors.BadIndexDuplicateKey(key, self)
 
114
        self._nodes[key] = ('', tuple(node_refs), value)
 
115
        self._keys.add(key)
 
116
        if self._key_length > 1:
 
117
            key_dict = self._nodes_by_key
 
118
            if self.reference_lists:
 
119
                key_value = key, value, tuple(node_refs)
 
120
            else:
 
121
                key_value = key, value
 
122
            # possibly should do this on-demand, but it seems likely it is 
 
123
            # always wanted
 
124
            # For a key of (foo, bar, baz) create
 
125
            # _nodes_by_key[foo][bar][baz] = key_value
 
126
            for subkey in key[:-1]:
 
127
                key_dict = key_dict.setdefault(subkey, {})
 
128
            key_dict[key[-1]] = key_value
 
129
 
 
130
    def finish(self):
 
131
        lines = [_SIGNATURE]
 
132
        lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n')
 
133
        lines.append(_OPTION_KEY_ELEMENTS + str(self._key_length) + '\n')
 
134
        lines.append(_OPTION_LEN + str(len(self._keys)) + '\n')
 
135
        prefix_length = sum(len(x) for x in lines)
 
136
        # references are byte offsets. To avoid having to do nasty
 
137
        # polynomial work to resolve offsets (references to later in the 
 
138
        # file cannot be determined until all the inbetween references have
 
139
        # been calculated too) we pad the offsets with 0's to make them be
 
140
        # of consistent length. Using binary offsets would break the trivial
 
141
        # file parsing.
 
142
        # to calculate the width of zero's needed we do three passes:
 
143
        # one to gather all the non-reference data and the number of references.
 
144
        # one to pad all the data with reference-length and determine entry
 
145
        # addresses.
 
146
        # One to serialise.
 
147
        
 
148
        # forward sorted by key. In future we may consider topological sorting,
 
149
        # at the cost of table scans for direct lookup, or a second index for
 
150
        # direct lookup
 
151
        nodes = sorted(self._nodes.items())
 
152
        # if we do not prepass, we don't know how long it will be up front.
 
153
        expected_bytes = None
 
154
        # we only need to pre-pass if we have reference lists at all.
 
155
        if self.reference_lists:
 
156
            key_offset_info = []
 
157
            non_ref_bytes = prefix_length
 
158
            total_references = 0
 
159
            # TODO use simple multiplication for the constants in this loop.
 
160
            for key, (absent, references, value) in nodes:
 
161
                # record the offset known *so far* for this key:
 
162
                # the non reference bytes to date, and the total references to
 
163
                # date - saves reaccumulating on the second pass
 
164
                key_offset_info.append((key, non_ref_bytes, total_references))
 
165
                # key is literal, value is literal, there are 3 null's, 1 NL
 
166
                # key is variable length tuple, \x00 between elements
 
167
                non_ref_bytes += sum(len(element) for element in key)
 
168
                if self._key_length > 1:
 
169
                    non_ref_bytes += self._key_length - 1
 
170
                # value is literal bytes, there are 3 null's, 1 NL.
 
171
                non_ref_bytes += len(value) + 3 + 1
 
172
                # one byte for absent if set.
 
173
                if absent:
 
174
                    non_ref_bytes += 1
 
175
                elif self.reference_lists:
 
176
                    # (ref_lists -1) tabs
 
177
                    non_ref_bytes += self.reference_lists - 1
 
178
                    # (ref-1 cr's per ref_list)
 
179
                    for ref_list in references:
 
180
                        # how many references across the whole file?
 
181
                        total_references += len(ref_list)
 
182
                        # accrue reference separators
 
183
                        if ref_list:
 
184
                            non_ref_bytes += len(ref_list) - 1
 
185
            # how many digits are needed to represent the total byte count?
 
186
            digits = 1
 
187
            possible_total_bytes = non_ref_bytes + total_references*digits
 
188
            while 10 ** digits < possible_total_bytes:
 
189
                digits += 1
 
190
                possible_total_bytes = non_ref_bytes + total_references*digits
 
191
            expected_bytes = possible_total_bytes + 1 # terminating newline
 
192
            # resolve key addresses.
 
193
            key_addresses = {}
 
194
            for key, non_ref_bytes, total_references in key_offset_info:
 
195
                key_addresses[key] = non_ref_bytes + total_references*digits
 
196
            # serialise
 
197
            format_string = '%%0%sd' % digits
 
198
        for key, (absent, references, value) in nodes:
 
199
            flattened_references = []
 
200
            for ref_list in references:
 
201
                ref_addresses = []
 
202
                for reference in ref_list:
 
203
                    ref_addresses.append(format_string % key_addresses[reference])
 
204
                flattened_references.append('\r'.join(ref_addresses))
 
205
            string_key = '\x00'.join(key)
 
206
            lines.append("%s\x00%s\x00%s\x00%s\n" % (string_key, absent,
 
207
                '\t'.join(flattened_references), value))
 
208
        lines.append('\n')
 
209
        result = StringIO(''.join(lines))
 
210
        if expected_bytes and len(result.getvalue()) != expected_bytes:
 
211
            raise errors.BzrError('Failed index creation. Internal error:'
 
212
                ' mismatched output length and expected length: %d %d' %
 
213
                (len(result.getvalue()), expected_bytes))
 
214
        return StringIO(''.join(lines))
 
215
 
 
216
 
 
217
class GraphIndex(object):
 
218
    """An index for data with embedded graphs.
 
219
 
 
220
    The index maps keys to a list of key reference lists, and a value.
 
221
    Each node has the same number of key reference lists. Each key reference
 
222
    list can be empty or an arbitrary length. The value is an opaque NULL
 
223
    terminated string without any newlines. The storage of the index is 
 
224
    hidden in the interface: keys and key references are always tuples of
 
225
    bytestrings, never the internal representation (e.g. dictionary offsets).
 
226
 
 
227
    It is presumed that the index will not be mutated - it is static data.
 
228
 
 
229
    Successive iter_all_entries calls will read the entire index each time.
 
230
    Additionally, iter_entries calls will read the index linearly until the
 
231
    desired keys are found. XXX: This must be fixed before the index is
 
232
    suitable for production use. :XXX
 
233
    """
 
234
 
 
235
    def __init__(self, transport, name):
 
236
        """Open an index called name on transport.
 
237
 
 
238
        :param transport: A bzrlib.transport.Transport.
 
239
        :param name: A path to provide to transport API calls.
 
240
        """
 
241
        self._transport = transport
 
242
        self._name = name
 
243
        self._nodes = None
 
244
        self._key_count = None
 
245
        self._keys_by_offset = None
 
246
        self._nodes_by_key = None
 
247
 
 
248
    def _buffer_all(self):
 
249
        """Buffer all the index data.
 
250
 
 
251
        Mutates self._nodes and self.keys_by_offset.
 
252
        """
 
253
        if 'index' in debug.debug_flags:
 
254
            mutter('Reading entire index %s', self._transport.abspath(self._name))
 
255
        stream = self._transport.get(self._name)
 
256
        self._read_prefix(stream)
 
257
        expected_elements = 3 + self._key_length
 
258
        line_count = 0
 
259
        # raw data keyed by offset
 
260
        self._keys_by_offset = {}
 
261
        # ready-to-return key:value or key:value, node_ref_lists
 
262
        self._nodes = {}
 
263
        self._nodes_by_key = {}
 
264
        trailers = 0
 
265
        pos = stream.tell()
 
266
        for line in stream.readlines():
 
267
            if line == '\n':
 
268
                trailers += 1
 
269
                continue
 
270
            elements = line.split('\0')
 
271
            if len(elements) != expected_elements:
 
272
                raise errors.BadIndexData(self)
 
273
            # keys are tuples
 
274
            key = tuple(elements[:self._key_length])
 
275
            absent, references, value = elements[-3:]
 
276
            value = value[:-1] # remove the newline
 
277
            ref_lists = []
 
278
            for ref_string in references.split('\t'):
 
279
                ref_lists.append(tuple([
 
280
                    int(ref) for ref in ref_string.split('\r') if ref
 
281
                    ]))
 
282
            ref_lists = tuple(ref_lists)
 
283
            self._keys_by_offset[pos] = (key, absent, ref_lists, value)
 
284
            pos += len(line)
 
285
        for key, absent, references, value in self._keys_by_offset.itervalues():
 
286
            if absent:
 
287
                continue
 
288
            # resolve references:
 
289
            if self.node_ref_lists:
 
290
                node_refs = []
 
291
                for ref_list in references:
 
292
                    node_refs.append(tuple([self._keys_by_offset[ref][0] for ref in ref_list]))
 
293
                node_value = (value, tuple(node_refs))
 
294
            else:
 
295
                node_value = value
 
296
            self._nodes[key] = node_value
 
297
            if self._key_length > 1:
 
298
                subkey = list(reversed(key[:-1]))
 
299
                key_dict = self._nodes_by_key
 
300
                if self.node_ref_lists:
 
301
                    key_value = key, node_value[0], node_value[1]
 
302
                else:
 
303
                    key_value = key, node_value
 
304
                # possibly should do this on-demand, but it seems likely it is 
 
305
                # always wanted
 
306
                # For a key of (foo, bar, baz) create
 
307
                # _nodes_by_key[foo][bar][baz] = key_value
 
308
                for subkey in key[:-1]:
 
309
                    key_dict = key_dict.setdefault(subkey, {})
 
310
                key_dict[key[-1]] = key_value
 
311
        # cache the keys for quick set intersections
 
312
        self._keys = set(self._nodes)
 
313
        if trailers != 1:
 
314
            # there must be one line - the empty trailer line.
 
315
            raise errors.BadIndexData(self)
 
316
 
 
317
    def iter_all_entries(self):
 
318
        """Iterate over all keys within the index.
 
319
 
 
320
        :return: An iterable of (key, value) or (key, value, reference_lists).
 
321
            The former tuple is used when there are no reference lists in the
 
322
            index, making the API compatible with simple key:value index types.
 
323
            There is no defined order for the result iteration - it will be in
 
324
            the most efficient order for the index.
 
325
        """
 
326
        if 'evil' in debug.debug_flags:
 
327
            trace.mutter_callsite(2,
 
328
                "iter_all_entries scales with size of history.")
 
329
        if self._nodes is None:
 
330
            self._buffer_all()
 
331
        if self.node_ref_lists:
 
332
            for key, (value, node_ref_lists) in self._nodes.iteritems():
 
333
                yield self, key, value, node_ref_lists
 
334
        else:
 
335
            for key, value in self._nodes.iteritems():
 
336
                yield self, key, value
 
337
 
 
338
    def _read_prefix(self, stream):
 
339
        signature = stream.read(len(self._signature()))
 
340
        if not signature == self._signature():
 
341
            raise errors.BadIndexFormatSignature(self._name, GraphIndex)
 
342
        options_line = stream.readline()
 
343
        if not options_line.startswith(_OPTION_NODE_REFS):
 
344
            raise errors.BadIndexOptions(self)
 
345
        try:
 
346
            self.node_ref_lists = int(options_line[len(_OPTION_NODE_REFS):-1])
 
347
        except ValueError:
 
348
            raise errors.BadIndexOptions(self)
 
349
        options_line = stream.readline()
 
350
        if not options_line.startswith(_OPTION_KEY_ELEMENTS):
 
351
            raise errors.BadIndexOptions(self)
 
352
        try:
 
353
            self._key_length = int(options_line[len(_OPTION_KEY_ELEMENTS):-1])
 
354
        except ValueError:
 
355
            raise errors.BadIndexOptions(self)
 
356
        options_line = stream.readline()
 
357
        if not options_line.startswith(_OPTION_LEN):
 
358
            raise errors.BadIndexOptions(self)
 
359
        try:
 
360
            self._key_count = int(options_line[len(_OPTION_LEN):-1])
 
361
        except ValueError:
 
362
            raise errors.BadIndexOptions(self)
 
363
 
 
364
    def iter_entries(self, keys):
 
365
        """Iterate over keys within the index.
 
366
 
 
367
        :param keys: An iterable providing the keys to be retrieved.
 
368
        :return: An iterable as per iter_all_entries, but restricted to the
 
369
            keys supplied. No additional keys will be returned, and every
 
370
            key supplied that is in the index will be returned.
 
371
        """
 
372
        keys = set(keys)
 
373
        if not keys:
 
374
            return
 
375
        if self._nodes is None:
 
376
            self._buffer_all()
 
377
        keys = keys.intersection(self._keys)
 
378
        if self.node_ref_lists:
 
379
            for key in keys:
 
380
                value, node_refs = self._nodes[key]
 
381
                yield self, key, value, node_refs
 
382
        else:
 
383
            for key in keys:
 
384
                yield self, key, self._nodes[key]
 
385
 
 
386
    def iter_entries_prefix(self, keys):
 
387
        """Iterate over keys within the index using prefix matching.
 
388
 
 
389
        Prefix matching is applied within the tuple of a key, not to within
 
390
        the bytestring of each key element. e.g. if you have the keys ('foo',
 
391
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
 
392
        only the former key is returned.
 
393
 
 
394
        :param keys: An iterable providing the key prefixes to be retrieved.
 
395
            Each key prefix takes the form of a tuple the length of a key, but
 
396
            with the last N elements 'None' rather than a regular bytestring.
 
397
            The first element cannot be 'None'.
 
398
        :return: An iterable as per iter_all_entries, but restricted to the
 
399
            keys with a matching prefix to those supplied. No additional keys
 
400
            will be returned, and every match that is in the index will be
 
401
            returned.
 
402
        """
 
403
        keys = set(keys)
 
404
        if not keys:
 
405
            return
 
406
        # load data - also finds key lengths
 
407
        if self._nodes is None:
 
408
            self._buffer_all()
 
409
        if self._key_length == 1:
 
410
            for key in keys:
 
411
                # sanity check
 
412
                if key[0] is None:
 
413
                    raise errors.BadIndexKey(key)
 
414
                if len(key) != self._key_length:
 
415
                    raise errors.BadIndexKey(key)
 
416
                if self.node_ref_lists:
 
417
                    value, node_refs = self._nodes[key]
 
418
                    yield self, key, value, node_refs
 
419
                else:
 
420
                    yield self, key, self._nodes[key]
 
421
            return
 
422
        for key in keys:
 
423
            # sanity check
 
424
            if key[0] is None:
 
425
                raise errors.BadIndexKey(key)
 
426
            if len(key) != self._key_length:
 
427
                raise errors.BadIndexKey(key)
 
428
            # find what it refers to:
 
429
            key_dict = self._nodes_by_key
 
430
            elements = list(key)
 
431
            # find the subdict whose contents should be returned.
 
432
            try:
 
433
                while len(elements) and elements[0] is not None:
 
434
                    key_dict = key_dict[elements[0]]
 
435
                    elements.pop(0)
 
436
            except KeyError:
 
437
                # a non-existant lookup.
 
438
                continue
 
439
            if len(elements):
 
440
                dicts = [key_dict]
 
441
                while dicts:
 
442
                    key_dict = dicts.pop(-1)
 
443
                    # can't be empty or would not exist
 
444
                    item, value = key_dict.iteritems().next()
 
445
                    if type(value) == dict:
 
446
                        # push keys 
 
447
                        dicts.extend(key_dict.itervalues())
 
448
                    else:
 
449
                        # yield keys
 
450
                        for value in key_dict.itervalues():
 
451
                            # each value is the key:value:node refs tuple
 
452
                            # ready to yield.
 
453
                            yield (self, ) + value
 
454
            else:
 
455
                # the last thing looked up was a terminal element
 
456
                yield (self, ) + key_dict
 
457
 
 
458
    def key_count(self):
 
459
        """Return an estimate of the number of keys in this index.
 
460
        
 
461
        For GraphIndex the estimate is exact.
 
462
        """
 
463
        if self._key_count is None:
 
464
            # really this should just read the prefix
 
465
            self._buffer_all()
 
466
        return self._key_count
 
467
 
 
468
    def _signature(self):
 
469
        """The file signature for this index type."""
 
470
        return _SIGNATURE
 
471
 
 
472
    def validate(self):
 
473
        """Validate that everything in the index can be accessed."""
 
474
        # iter_all validates completely at the moment, so just do that.
 
475
        for node in self.iter_all_entries():
 
476
            pass
 
477
 
 
478
 
 
479
class CombinedGraphIndex(object):
 
480
    """A GraphIndex made up from smaller GraphIndices.
 
481
    
 
482
    The backing indices must implement GraphIndex, and are presumed to be
 
483
    static data.
 
484
 
 
485
    Queries against the combined index will be made against the first index,
 
486
    and then the second and so on. The order of index's can thus influence
 
487
    performance significantly. For example, if one index is on local disk and a
 
488
    second on a remote server, the local disk index should be before the other
 
489
    in the index list.
 
490
    """
 
491
 
 
492
    def __init__(self, indices):
 
493
        """Create a CombinedGraphIndex backed by indices.
 
494
 
 
495
        :param indices: An ordered list of indices to query for data.
 
496
        """
 
497
        self._indices = indices
 
498
 
 
499
    def insert_index(self, pos, index):
 
500
        """Insert a new index in the list of indices to query.
 
501
 
 
502
        :param pos: The position to insert the index.
 
503
        :param index: The index to insert.
 
504
        """
 
505
        self._indices.insert(pos, index)
 
506
 
 
507
    def iter_all_entries(self):
 
508
        """Iterate over all keys within the index
 
509
 
 
510
        Duplicate keys across child indices are presumed to have the same
 
511
        value and are only reported once.
 
512
 
 
513
        :return: An iterable of (key, reference_lists, value). There is no
 
514
            defined order for the result iteration - it will be in the most
 
515
            efficient order for the index.
 
516
        """
 
517
        seen_keys = set()
 
518
        for index in self._indices:
 
519
            for node in index.iter_all_entries():
 
520
                if node[1] not in seen_keys:
 
521
                    yield node
 
522
                    seen_keys.add(node[1])
 
523
 
 
524
    def iter_entries(self, keys):
 
525
        """Iterate over keys within the index.
 
526
 
 
527
        Duplicate keys across child indices are presumed to have the same
 
528
        value and are only reported once.
 
529
 
 
530
        :param keys: An iterable providing the keys to be retrieved.
 
531
        :return: An iterable of (key, reference_lists, value). There is no
 
532
            defined order for the result iteration - it will be in the most
 
533
            efficient order for the index.
 
534
        """
 
535
        keys = set(keys)
 
536
        for index in self._indices:
 
537
            if not keys:
 
538
                return
 
539
            for node in index.iter_entries(keys):
 
540
                keys.remove(node[1])
 
541
                yield node
 
542
 
 
543
    def iter_entries_prefix(self, keys):
 
544
        """Iterate over keys within the index using prefix matching.
 
545
 
 
546
        Duplicate keys across child indices are presumed to have the same
 
547
        value and are only reported once.
 
548
 
 
549
        Prefix matching is applied within the tuple of a key, not to within
 
550
        the bytestring of each key element. e.g. if you have the keys ('foo',
 
551
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
 
552
        only the former key is returned.
 
553
 
 
554
        :param keys: An iterable providing the key prefixes to be retrieved.
 
555
            Each key prefix takes the form of a tuple the length of a key, but
 
556
            with the last N elements 'None' rather than a regular bytestring.
 
557
            The first element cannot be 'None'.
 
558
        :return: An iterable as per iter_all_entries, but restricted to the
 
559
            keys with a matching prefix to those supplied. No additional keys
 
560
            will be returned, and every match that is in the index will be
 
561
            returned.
 
562
        """
 
563
        keys = set(keys)
 
564
        if not keys:
 
565
            return
 
566
        seen_keys = set()
 
567
        for index in self._indices:
 
568
            for node in index.iter_entries_prefix(keys):
 
569
                if node[1] in seen_keys:
 
570
                    continue
 
571
                seen_keys.add(node[1])
 
572
                yield node
 
573
 
 
574
    def key_count(self):
 
575
        """Return an estimate of the number of keys in this index.
 
576
        
 
577
        For CombinedGraphIndex this is approximated by the sum of the keys of
 
578
        the child indices. As child indices may have duplicate keys this can
 
579
        have a maximum error of the number of child indices * largest number of
 
580
        keys in any index.
 
581
        """
 
582
        return sum((index.key_count() for index in self._indices), 0)
 
583
 
 
584
    def validate(self):
 
585
        """Validate that everything in the index can be accessed."""
 
586
        for index in self._indices:
 
587
            index.validate()
 
588
 
 
589
 
 
590
class InMemoryGraphIndex(GraphIndexBuilder):
 
591
    """A GraphIndex which operates entirely out of memory and is mutable.
 
592
 
 
593
    This is designed to allow the accumulation of GraphIndex entries during a
 
594
    single write operation, where the accumulated entries need to be immediately
 
595
    available - for example via a CombinedGraphIndex.
 
596
    """
 
597
 
 
598
    def add_nodes(self, nodes):
 
599
        """Add nodes to the index.
 
600
 
 
601
        :param nodes: An iterable of (key, node_refs, value) entries to add.
 
602
        """
 
603
        if self.reference_lists:
 
604
            for (key, value, node_refs) in nodes:
 
605
                self.add_node(key, value, node_refs)
 
606
        else:
 
607
            for (key, value) in nodes:
 
608
                self.add_node(key, value)
 
609
 
 
610
    def iter_all_entries(self):
 
611
        """Iterate over all keys within the index
 
612
 
 
613
        :return: An iterable of (key, reference_lists, value). There is no
 
614
            defined order for the result iteration - it will be in the most
 
615
            efficient order for the index (in this case dictionary hash order).
 
616
        """
 
617
        if 'evil' in debug.debug_flags:
 
618
            trace.mutter_callsite(2,
 
619
                "iter_all_entries scales with size of history.")
 
620
        if self.reference_lists:
 
621
            for key, (absent, references, value) in self._nodes.iteritems():
 
622
                if not absent:
 
623
                    yield self, key, value, references
 
624
        else:
 
625
            for key, (absent, references, value) in self._nodes.iteritems():
 
626
                if not absent:
 
627
                    yield self, key, value
 
628
 
 
629
    def iter_entries(self, keys):
 
630
        """Iterate over keys within the index.
 
631
 
 
632
        :param keys: An iterable providing the keys to be retrieved.
 
633
        :return: An iterable of (key, reference_lists, value). There is no
 
634
            defined order for the result iteration - it will be in the most
 
635
            efficient order for the index (keys iteration order in this case).
 
636
        """
 
637
        keys = set(keys)
 
638
        if self.reference_lists:
 
639
            for key in keys.intersection(self._keys):
 
640
                node = self._nodes[key]
 
641
                if not node[0]:
 
642
                    yield self, key, node[2], node[1]
 
643
        else:
 
644
            for key in keys.intersection(self._keys):
 
645
                node = self._nodes[key]
 
646
                if not node[0]:
 
647
                    yield self, key, node[2]
 
648
 
 
649
    def iter_entries_prefix(self, keys):
 
650
        """Iterate over keys within the index using prefix matching.
 
651
 
 
652
        Prefix matching is applied within the tuple of a key, not to within
 
653
        the bytestring of each key element. e.g. if you have the keys ('foo',
 
654
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
 
655
        only the former key is returned.
 
656
 
 
657
        :param keys: An iterable providing the key prefixes to be retrieved.
 
658
            Each key prefix takes the form of a tuple the length of a key, but
 
659
            with the last N elements 'None' rather than a regular bytestring.
 
660
            The first element cannot be 'None'.
 
661
        :return: An iterable as per iter_all_entries, but restricted to the
 
662
            keys with a matching prefix to those supplied. No additional keys
 
663
            will be returned, and every match that is in the index will be
 
664
            returned.
 
665
        """
 
666
        # XXX: To much duplication with the GraphIndex class; consider finding
 
667
        # a good place to pull out the actual common logic.
 
668
        keys = set(keys)
 
669
        if not keys:
 
670
            return
 
671
        if self._key_length == 1:
 
672
            for key in keys:
 
673
                # sanity check
 
674
                if key[0] is None:
 
675
                    raise errors.BadIndexKey(key)
 
676
                if len(key) != self._key_length:
 
677
                    raise errors.BadIndexKey(key)
 
678
                node = self._nodes[key]
 
679
                if node[0]:
 
680
                    continue 
 
681
                if self.reference_lists:
 
682
                    yield self, key, node[2], node[1]
 
683
                else:
 
684
                    yield self, key, node[2]
 
685
            return
 
686
        for key in keys:
 
687
            # sanity check
 
688
            if key[0] is None:
 
689
                raise errors.BadIndexKey(key)
 
690
            if len(key) != self._key_length:
 
691
                raise errors.BadIndexKey(key)
 
692
            # find what it refers to:
 
693
            key_dict = self._nodes_by_key
 
694
            elements = list(key)
 
695
            # find the subdict to return
 
696
            try:
 
697
                while len(elements) and elements[0] is not None:
 
698
                    key_dict = key_dict[elements[0]]
 
699
                    elements.pop(0)
 
700
            except KeyError:
 
701
                # a non-existant lookup.
 
702
                continue
 
703
            if len(elements):
 
704
                dicts = [key_dict]
 
705
                while dicts:
 
706
                    key_dict = dicts.pop(-1)
 
707
                    # can't be empty or would not exist
 
708
                    item, value = key_dict.iteritems().next()
 
709
                    if type(value) == dict:
 
710
                        # push keys 
 
711
                        dicts.extend(key_dict.itervalues())
 
712
                    else:
 
713
                        # yield keys
 
714
                        for value in key_dict.itervalues():
 
715
                            yield (self, ) + value
 
716
            else:
 
717
                yield (self, ) + key_dict
 
718
 
 
719
    def key_count(self):
 
720
        """Return an estimate of the number of keys in this index.
 
721
        
 
722
        For InMemoryGraphIndex the estimate is exact.
 
723
        """
 
724
        return len(self._keys)
 
725
 
 
726
    def validate(self):
 
727
        """In memory index's have no known corruption at the moment."""
 
728
 
 
729
 
 
730
class GraphIndexPrefixAdapter(object):
 
731
    """An adapter between GraphIndex with different key lengths.
 
732
 
 
733
    Queries against this will emit queries against the adapted Graph with the
 
734
    prefix added, queries for all items use iter_entries_prefix. The returned
 
735
    nodes will have their keys and node references adjusted to remove the 
 
736
    prefix. Finally, an add_nodes_callback can be supplied - when called the
 
737
    nodes and references being added will have prefix prepended.
 
738
    """
 
739
 
 
740
    def __init__(self, adapted, prefix, missing_key_length,
 
741
        add_nodes_callback=None):
 
742
        """Construct an adapter against adapted with prefix."""
 
743
        self.adapted = adapted
 
744
        self.prefix_key = prefix + (None,)*missing_key_length
 
745
        self.prefix = prefix
 
746
        self.prefix_len = len(prefix)
 
747
        self.add_nodes_callback = add_nodes_callback
 
748
 
 
749
    def add_nodes(self, nodes):
 
750
        """Add nodes to the index.
 
751
 
 
752
        :param nodes: An iterable of (key, node_refs, value) entries to add.
 
753
        """
 
754
        # save nodes in case its an iterator
 
755
        nodes = tuple(nodes)
 
756
        translated_nodes = []
 
757
        try:
 
758
            # Add prefix_key to each reference node_refs is a tuple of tuples,
 
759
            # so split it apart, and add prefix_key to the internal reference
 
760
            for (key, value, node_refs) in nodes:
 
761
                adjusted_references = (
 
762
                    tuple(tuple(self.prefix + ref_node for ref_node in ref_list)
 
763
                        for ref_list in node_refs))
 
764
                translated_nodes.append((self.prefix + key, value,
 
765
                    adjusted_references))
 
766
        except ValueError:
 
767
            # XXX: TODO add an explicit interface for getting the reference list
 
768
            # status, to handle this bit of user-friendliness in the API more 
 
769
            # explicitly.
 
770
            for (key, value) in nodes:
 
771
                translated_nodes.append((self.prefix + key, value))
 
772
        self.add_nodes_callback(translated_nodes)
 
773
 
 
774
    def add_node(self, key, value, references=()):
 
775
        """Add a node to the index.
 
776
 
 
777
        :param key: The key. keys are non-empty tuples containing
 
778
            as many whitespace-free utf8 bytestrings as the key length
 
779
            defined for this index.
 
780
        :param references: An iterable of iterables of keys. Each is a
 
781
            reference to another key.
 
782
        :param value: The value to associate with the key. It may be any
 
783
            bytes as long as it does not contain \0 or \n.
 
784
        """
 
785
        self.add_nodes(((key, value, references), ))
 
786
 
 
787
    def _strip_prefix(self, an_iter):
 
788
        """Strip prefix data from nodes and return it."""
 
789
        for node in an_iter:
 
790
            # cross checks
 
791
            if node[1][:self.prefix_len] != self.prefix:
 
792
                raise errors.BadIndexData(self)
 
793
            for ref_list in node[3]:
 
794
                for ref_node in ref_list:
 
795
                    if ref_node[:self.prefix_len] != self.prefix:
 
796
                        raise errors.BadIndexData(self)
 
797
            yield node[0], node[1][self.prefix_len:], node[2], (
 
798
                tuple(tuple(ref_node[self.prefix_len:] for ref_node in ref_list)
 
799
                for ref_list in node[3]))
 
800
 
 
801
    def iter_all_entries(self):
 
802
        """Iterate over all keys within the index
 
803
 
 
804
        iter_all_entries is implemented against the adapted index using
 
805
        iter_entries_prefix.
 
806
 
 
807
        :return: An iterable of (key, reference_lists, value). There is no
 
808
            defined order for the result iteration - it will be in the most
 
809
            efficient order for the index (in this case dictionary hash order).
 
810
        """
 
811
        return self._strip_prefix(self.adapted.iter_entries_prefix([self.prefix_key]))
 
812
 
 
813
    def iter_entries(self, keys):
 
814
        """Iterate over keys within the index.
 
815
 
 
816
        :param keys: An iterable providing the keys to be retrieved.
 
817
        :return: An iterable of (key, reference_lists, value). There is no
 
818
            defined order for the result iteration - it will be in the most
 
819
            efficient order for the index (keys iteration order in this case).
 
820
        """
 
821
        return self._strip_prefix(self.adapted.iter_entries(
 
822
            self.prefix + key for key in keys))
 
823
 
 
824
    def iter_entries_prefix(self, keys):
 
825
        """Iterate over keys within the index using prefix matching.
 
826
 
 
827
        Prefix matching is applied within the tuple of a key, not to within
 
828
        the bytestring of each key element. e.g. if you have the keys ('foo',
 
829
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
 
830
        only the former key is returned.
 
831
 
 
832
        :param keys: An iterable providing the key prefixes to be retrieved.
 
833
            Each key prefix takes the form of a tuple the length of a key, but
 
834
            with the last N elements 'None' rather than a regular bytestring.
 
835
            The first element cannot be 'None'.
 
836
        :return: An iterable as per iter_all_entries, but restricted to the
 
837
            keys with a matching prefix to those supplied. No additional keys
 
838
            will be returned, and every match that is in the index will be
 
839
            returned.
 
840
        """
 
841
        return self._strip_prefix(self.adapted.iter_entries_prefix(
 
842
            self.prefix + key for key in keys))
 
843
 
 
844
    def key_count(self):
 
845
        """Return an estimate of the number of keys in this index.
 
846
        
 
847
        For GraphIndexPrefixAdapter this is relatively expensive - key
 
848
        iteration with the prefix is done.
 
849
        """
 
850
        return len(list(self.iter_all_entries()))
 
851
 
 
852
    def validate(self):
 
853
        """Call the adapted's validate."""
 
854
        self.adapted.validate()