/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar
3641.3.29 by John Arbash Meinel
Cleanup the copyright headers
1
# Copyright (C) 2008 Canonical Ltd
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
2
#
3
# This program is free software; you can redistribute it and/or modify
3641.3.29 by John Arbash Meinel
Cleanup the copyright headers
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.
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
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
3641.3.29 by John Arbash Meinel
Cleanup the copyright headers
15
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
16
#
17
18
"""B+Tree indices"""
19
20
import array
21
import bisect
22
from bisect import bisect_right
23
from copy import deepcopy
24
import math
25
import struct
26
import tempfile
27
import zlib
28
29
from bzrlib import (
30
    chunk_writer,
31
    debug,
32
    errors,
33
    index,
34
    lru_cache,
35
    osutils,
36
    trace,
37
    )
38
from bzrlib.index import _OPTION_NODE_REFS, _OPTION_KEY_ELEMENTS, _OPTION_LEN
39
from bzrlib.transport import get_transport
40
41
3641.3.3 by John Arbash Meinel
Change the header to indicate these indexes are
42
_BTSIGNATURE = "B+Tree Graph Index 2\n"
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
43
_OPTION_ROW_LENGTHS = "row_lengths="
44
_LEAF_FLAG = "type=leaf\n"
45
_INTERNAL_FLAG = "type=internal\n"
46
_INTERNAL_OFFSET = "offset="
47
48
_RESERVED_HEADER_BYTES = 120
49
_PAGE_SIZE = 4096
50
51
# 4K per page: 4MB - 1000 entries
52
_NODE_CACHE_SIZE = 1000
53
54
55
class _BuilderRow(object):
56
    """The stored state accumulated while writing out a row in the index.
57
58
    :ivar spool: A temporary file used to accumulate nodes for this row
59
        in the tree.
60
    :ivar nodes: The count of nodes emitted so far.
61
    """
62
63
    def __init__(self):
64
        """Create a _BuilderRow."""
65
        self.nodes = 0
66
        self.spool = tempfile.TemporaryFile()
67
        self.writer = None
68
69
    def finish_node(self, pad=True):
70
        byte_lines, _, padding = self.writer.finish()
71
        if self.nodes == 0:
72
            # padded note:
73
            self.spool.write("\x00" * _RESERVED_HEADER_BYTES)
74
        skipped_bytes = 0
75
        if not pad and padding:
76
            del byte_lines[-1]
77
            skipped_bytes = padding
78
        self.spool.writelines(byte_lines)
3644.2.3 by John Arbash Meinel
Do a bit more work to get all the tests to pass.
79
        remainder = (self.spool.tell() + skipped_bytes) % _PAGE_SIZE
80
        if remainder != 0:
81
            raise AssertionError("incorrect node length: %d, %d"
82
                                 % (self.spool.tell(), remainder))
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
83
        self.nodes += 1
84
        self.writer = None
85
86
87
class _InternalBuilderRow(_BuilderRow):
88
    """The stored state accumulated while writing out internal rows."""
89
90
    def finish_node(self, pad=True):
91
        if not pad:
92
            raise AssertionError("Must pad internal nodes only.")
93
        _BuilderRow.finish_node(self)
94
95
96
class _LeafBuilderRow(_BuilderRow):
97
    """The stored state accumulated while writing out a leaf rows."""
98
99
100
class BTreeBuilder(index.GraphIndexBuilder):
101
    """A Builder for B+Tree based Graph indices.
102
103
    The resulting graph has the structure:
104
105
    _SIGNATURE OPTIONS NODES
106
    _SIGNATURE     := 'B+Tree Graph Index 1' NEWLINE
107
    OPTIONS        := REF_LISTS KEY_ELEMENTS LENGTH
108
    REF_LISTS      := 'node_ref_lists=' DIGITS NEWLINE
109
    KEY_ELEMENTS   := 'key_elements=' DIGITS NEWLINE
110
    LENGTH         := 'len=' DIGITS NEWLINE
111
    ROW_LENGTHS    := 'row_lengths' DIGITS (COMMA DIGITS)*
112
    NODES          := NODE_COMPRESSED*
113
    NODE_COMPRESSED:= COMPRESSED_BYTES{4096}
114
    NODE_RAW       := INTERNAL | LEAF
115
    INTERNAL       := INTERNAL_FLAG POINTERS
116
    LEAF           := LEAF_FLAG ROWS
117
    KEY_ELEMENT    := Not-whitespace-utf8
118
    KEY            := KEY_ELEMENT (NULL KEY_ELEMENT)*
119
    ROWS           := ROW*
120
    ROW            := KEY NULL ABSENT? NULL REFERENCES NULL VALUE NEWLINE
121
    ABSENT         := 'a'
122
    REFERENCES     := REFERENCE_LIST (TAB REFERENCE_LIST){node_ref_lists - 1}
123
    REFERENCE_LIST := (REFERENCE (CR REFERENCE)*)?
124
    REFERENCE      := KEY
125
    VALUE          := no-newline-no-null-bytes
126
    """
127
128
    def __init__(self, reference_lists=0, key_elements=1, spill_at=100000):
129
        """See GraphIndexBuilder.__init__.
130
131
        :param spill_at: Optional parameter controlling the maximum number
132
            of nodes that BTreeBuilder will hold in memory.
133
        """
134
        index.GraphIndexBuilder.__init__(self, reference_lists=reference_lists,
135
            key_elements=key_elements)
136
        self._spill_at = spill_at
137
        self._backing_indices = []
3644.2.11 by John Arbash Meinel
Document the new form of _nodes and remove an unnecessary cast.
138
        # A map of {key: (node_refs, value)}
139
        self._nodes = {}
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
140
        # Indicate it hasn't been built yet
141
        self._nodes_by_key = None
3777.5.2 by John Arbash Meinel
Change the name to ChunkWriter.set_optimize()
142
        self._optimize_for_size = False
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
143
144
    def add_node(self, key, value, references=()):
145
        """Add a node to the index.
146
147
        If adding the node causes the builder to reach its spill_at threshold,
148
        disk spilling will be triggered.
149
150
        :param key: The key. keys are non-empty tuples containing
151
            as many whitespace-free utf8 bytestrings as the key length
152
            defined for this index.
153
        :param references: An iterable of iterables of keys. Each is a
154
            reference to another key.
155
        :param value: The value to associate with the key. It may be any
156
            bytes as long as it does not contain \0 or \n.
157
        """
3644.2.9 by John Arbash Meinel
Refactor some code.
158
        # we don't care about absent_references
159
        node_refs, _ = self._check_key_ref_value(key, references, value)
3644.2.2 by John Arbash Meinel
the new btree index doesn't have 'absent' keys in its _nodes
160
        if key in self._nodes:
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
161
            raise errors.BadIndexDuplicateKey(key, self)
3644.2.11 by John Arbash Meinel
Document the new form of _nodes and remove an unnecessary cast.
162
        self._nodes[key] = (node_refs, value)
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
163
        self._keys.add(key)
3644.2.9 by John Arbash Meinel
Refactor some code.
164
        if self._nodes_by_key is not None and self._key_length > 1:
165
            self._update_nodes_by_key(key, value, node_refs)
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
166
        if len(self._keys) < self._spill_at:
167
            return
3644.2.9 by John Arbash Meinel
Refactor some code.
168
        self._spill_mem_keys_to_disk()
169
170
    def _spill_mem_keys_to_disk(self):
171
        """Write the in memory keys down to disk to cap memory consumption.
172
173
        If we already have some keys written to disk, we will combine them so
174
        as to preserve the sorted order.  The algorithm for combining uses
175
        powers of two.  So on the first spill, write all mem nodes into a
176
        single index. On the second spill, combine the mem nodes with the nodes
177
        on disk to create a 2x sized disk index and get rid of the first index.
178
        On the third spill, create a single new disk index, which will contain
179
        the mem nodes, and preserve the existing 2x sized index.  On the fourth,
180
        combine mem with the first and second indexes, creating a new one of
181
        size 4x. On the fifth create a single new one, etc.
182
        """
3644.2.8 by John Arbash Meinel
Two quick tweaks.
183
        iterators_to_combine = [self._iter_mem_nodes()]
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
184
        pos = -1
185
        for pos, backing in enumerate(self._backing_indices):
186
            if backing is None:
187
                pos -= 1
188
                break
189
            iterators_to_combine.append(backing.iter_all_entries())
190
        backing_pos = pos + 1
191
        new_backing_file, size = \
192
            self._write_nodes(self._iter_smallest(iterators_to_combine))
3644.2.9 by John Arbash Meinel
Refactor some code.
193
        dir_path, base_name = osutils.split(new_backing_file.name)
194
        # Note: The transport here isn't strictly needed, because we will use
195
        #       direct access to the new_backing._file object
196
        new_backing = BTreeGraphIndex(get_transport(dir_path),
197
                                      base_name, size)
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
198
        # GC will clean up the file
199
        new_backing._file = new_backing_file
200
        if len(self._backing_indices) == backing_pos:
201
            self._backing_indices.append(None)
202
        self._backing_indices[backing_pos] = new_backing
203
        for pos in range(backing_pos):
204
            self._backing_indices[pos] = None
205
        self._keys = set()
206
        self._nodes = {}
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
207
        self._nodes_by_key = None
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
208
209
    def add_nodes(self, nodes):
210
        """Add nodes to the index.
211
212
        :param nodes: An iterable of (key, node_refs, value) entries to add.
213
        """
214
        if self.reference_lists:
215
            for (key, value, node_refs) in nodes:
216
                self.add_node(key, value, node_refs)
217
        else:
218
            for (key, value) in nodes:
219
                self.add_node(key, value)
220
221
    def _iter_mem_nodes(self):
222
        """Iterate over the nodes held in memory."""
3644.2.8 by John Arbash Meinel
Two quick tweaks.
223
        nodes = self._nodes
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
224
        if self.reference_lists:
3644.2.8 by John Arbash Meinel
Two quick tweaks.
225
            for key in sorted(nodes):
226
                references, value = nodes[key]
227
                yield self, key, value, references
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
228
        else:
3644.2.8 by John Arbash Meinel
Two quick tweaks.
229
            for key in sorted(nodes):
230
                references, value = nodes[key]
231
                yield self, key, value
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
232
233
    def _iter_smallest(self, iterators_to_combine):
3641.3.9 by John Arbash Meinel
Special case around _iter_smallest when we have only
234
        if len(iterators_to_combine) == 1:
235
            for value in iterators_to_combine[0]:
236
                yield value
237
            return
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
238
        current_values = []
239
        for iterator in iterators_to_combine:
240
            try:
241
                current_values.append(iterator.next())
242
            except StopIteration:
243
                current_values.append(None)
244
        last = None
245
        while True:
246
            # Decorate candidates with the value to allow 2.4's min to be used.
247
            candidates = [(item[1][1], item) for item
248
                in enumerate(current_values) if item[1] is not None]
249
            if not len(candidates):
250
                return
251
            selected = min(candidates)
252
            # undecorate back to (pos, node)
253
            selected = selected[1]
254
            if last == selected[1][1]:
255
                raise errors.BadIndexDuplicateKey(last, self)
256
            last = selected[1][1]
257
            # Yield, with self as the index
258
            yield (self,) + selected[1][1:]
259
            pos = selected[0]
260
            try:
261
                current_values[pos] = iterators_to_combine[pos].next()
262
            except StopIteration:
263
                current_values[pos] = None
264
3641.3.11 by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state.
265
    def _add_key(self, string_key, line, rows):
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
266
        """Add a key to the current chunk.
267
268
        :param string_key: The key to add.
3641.3.11 by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state.
269
        :param line: The fully serialised key and value.
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
270
        """
271
        if rows[-1].writer is None:
272
            # opening a new leaf chunk;
273
            for pos, internal_row in enumerate(rows[:-1]):
274
                # flesh out any internal nodes that are needed to
3641.3.11 by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state.
275
                # preserve the height of the tree
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
276
                if internal_row.writer is None:
277
                    length = _PAGE_SIZE
278
                    if internal_row.nodes == 0:
279
                        length -= _RESERVED_HEADER_BYTES # padded
3777.5.2 by John Arbash Meinel
Change the name to ChunkWriter.set_optimize()
280
                    internal_row.writer = chunk_writer.ChunkWriter(length, 0,
281
                        optimize_for_size=self._optimize_for_size)
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
282
                    internal_row.writer.write(_INTERNAL_FLAG)
283
                    internal_row.writer.write(_INTERNAL_OFFSET +
284
                        str(rows[pos + 1].nodes) + "\n")
285
            # add a new leaf
286
            length = _PAGE_SIZE
287
            if rows[-1].nodes == 0:
288
                length -= _RESERVED_HEADER_BYTES # padded
3777.5.2 by John Arbash Meinel
Change the name to ChunkWriter.set_optimize()
289
            rows[-1].writer = chunk_writer.ChunkWriter(length,
290
                optimize_for_size=self._optimize_for_size)
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
291
            rows[-1].writer.write(_LEAF_FLAG)
3641.3.11 by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state.
292
        if rows[-1].writer.write(line):
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
293
            # this key did not fit in the node:
294
            rows[-1].finish_node()
3641.3.11 by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state.
295
            key_line = string_key + "\n"
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
296
            new_row = True
297
            for row in reversed(rows[:-1]):
298
                # Mark the start of the next node in the node above. If it
299
                # doesn't fit then propogate upwards until we find one that
300
                # it does fit into.
3641.3.11 by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state.
301
                if row.writer.write(key_line):
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
302
                    row.finish_node()
303
                else:
304
                    # We've found a node that can handle the pointer.
305
                    new_row = False
306
                    break
307
            # If we reached the current root without being able to mark the
308
            # division point, then we need a new root:
309
            if new_row:
310
                # We need a new row
311
                if 'index' in debug.debug_flags:
312
                    trace.mutter('Inserting new global row.')
313
                new_row = _InternalBuilderRow()
314
                reserved_bytes = 0
315
                rows.insert(0, new_row)
316
                # This will be padded, hence the -100
317
                new_row.writer = chunk_writer.ChunkWriter(
318
                    _PAGE_SIZE - _RESERVED_HEADER_BYTES,
3777.5.2 by John Arbash Meinel
Change the name to ChunkWriter.set_optimize()
319
                    reserved_bytes,
320
                    optimize_for_size=self._optimize_for_size)
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
321
                new_row.writer.write(_INTERNAL_FLAG)
322
                new_row.writer.write(_INTERNAL_OFFSET +
323
                    str(rows[1].nodes - 1) + "\n")
3641.3.11 by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state.
324
                new_row.writer.write(key_line)
325
            self._add_key(string_key, line, rows)
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
326
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
327
    def _write_nodes(self, node_iterator):
328
        """Write node_iterator out as a B+Tree.
329
330
        :param node_iterator: An iterator of sorted nodes. Each node should
331
            match the output given by iter_all_entries.
332
        :return: A file handle for a temporary file containing a B+Tree for
333
            the nodes.
334
        """
335
        # The index rows - rows[0] is the root, rows[1] is the layer under it
336
        # etc.
337
        rows = []
338
        # forward sorted by key. In future we may consider topological sorting,
339
        # at the cost of table scans for direct lookup, or a second index for
340
        # direct lookup
341
        key_count = 0
342
        # A stack with the number of nodes of each size. 0 is the root node
343
        # and must always be 1 (if there are any nodes in the tree).
344
        self.row_lengths = []
345
        # Loop over all nodes adding them to the bottom row
346
        # (rows[-1]). When we finish a chunk in a row,
347
        # propogate the key that didn't fit (comes after the chunk) to the
348
        # row above, transitively.
349
        for node in node_iterator:
350
            if key_count == 0:
351
                # First key triggers the first row
352
                rows.append(_LeafBuilderRow())
353
            key_count += 1
3641.3.11 by John Arbash Meinel
Start working on an alternate way to track compressed_chunk state.
354
            # TODO: Flattening the node into a string key and a line should
355
            #       probably be put into a pyrex function. We can do a quick
356
            #       iter over all the entries to determine the final length,
357
            #       and then do a single malloc() rather than lots of
358
            #       intermediate mallocs as we build everything up.
359
            #       ATM 3 / 13s are spent flattening nodes (10s is compressing)
3641.3.30 by John Arbash Meinel
Rename _parse_btree to _btree_serializer
360
            string_key, line = _btree_serializer._flatten_node(node,
361
                                    self.reference_lists)
3641.3.8 by John Arbash Meinel
Move the add_key helper function into a separate func
362
            self._add_key(string_key, line, rows)
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
363
        for row in reversed(rows):
364
            pad = (type(row) != _LeafBuilderRow)
365
            row.finish_node(pad=pad)
366
        result = tempfile.NamedTemporaryFile()
367
        lines = [_BTSIGNATURE]
368
        lines.append(_OPTION_NODE_REFS + str(self.reference_lists) + '\n')
369
        lines.append(_OPTION_KEY_ELEMENTS + str(self._key_length) + '\n')
370
        lines.append(_OPTION_LEN + str(key_count) + '\n')
371
        row_lengths = [row.nodes for row in rows]
372
        lines.append(_OPTION_ROW_LENGTHS + ','.join(map(str, row_lengths)) + '\n')
373
        result.writelines(lines)
374
        position = sum(map(len, lines))
375
        root_row = True
376
        if position > _RESERVED_HEADER_BYTES:
377
            raise AssertionError("Could not fit the header in the"
378
                                 " reserved space: %d > %d"
379
                                 % (position, _RESERVED_HEADER_BYTES))
380
        # write the rows out:
381
        for row in rows:
382
            reserved = _RESERVED_HEADER_BYTES # reserved space for first node
383
            row.spool.flush()
384
            row.spool.seek(0)
385
            # copy nodes to the finalised file.
386
            # Special case the first node as it may be prefixed
387
            node = row.spool.read(_PAGE_SIZE)
388
            result.write(node[reserved:])
389
            result.write("\x00" * (reserved - position))
390
            position = 0 # Only the root row actually has an offset
391
            copied_len = osutils.pumpfile(row.spool, result)
392
            if copied_len != (row.nodes - 1) * _PAGE_SIZE:
393
                if type(row) != _LeafBuilderRow:
3644.2.3 by John Arbash Meinel
Do a bit more work to get all the tests to pass.
394
                    raise AssertionError("Incorrect amount of data copied"
395
                        " expected: %d, got: %d"
396
                        % ((row.nodes - 1) * _PAGE_SIZE,
397
                           copied_len))
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
398
        result.flush()
399
        size = result.tell()
400
        result.seek(0)
401
        return result, size
402
403
    def finish(self):
404
        """Finalise the index.
405
406
        :return: A file handle for a temporary file containing the nodes added
407
            to the index.
408
        """
409
        return self._write_nodes(self.iter_all_entries())[0]
410
411
    def iter_all_entries(self):
412
        """Iterate over all keys within the index
413
414
        :return: An iterable of (index, key, reference_lists, value). There is no
415
            defined order for the result iteration - it will be in the most
416
            efficient order for the index (in this case dictionary hash order).
417
        """
418
        if 'evil' in debug.debug_flags:
419
            trace.mutter_callsite(3,
420
                "iter_all_entries scales with size of history.")
421
        # Doing serial rather than ordered would be faster; but this shouldn't
422
        # be getting called routinely anyway.
3644.2.8 by John Arbash Meinel
Two quick tweaks.
423
        iterators = [self._iter_mem_nodes()]
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
424
        for backing in self._backing_indices:
425
            if backing is not None:
426
                iterators.append(backing.iter_all_entries())
3641.3.9 by John Arbash Meinel
Special case around _iter_smallest when we have only
427
        if len(iterators) == 1:
428
            return iterators[0]
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
429
        return self._iter_smallest(iterators)
430
431
    def iter_entries(self, keys):
432
        """Iterate over keys within the index.
433
434
        :param keys: An iterable providing the keys to be retrieved.
435
        :return: An iterable of (index, key, value, reference_lists). There is no
436
            defined order for the result iteration - it will be in the most
437
            efficient order for the index (keys iteration order in this case).
438
        """
439
        keys = set(keys)
440
        if self.reference_lists:
441
            for key in keys.intersection(self._keys):
442
                node = self._nodes[key]
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
443
                yield self, key, node[1], node[0]
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
444
        else:
445
            for key in keys.intersection(self._keys):
446
                node = self._nodes[key]
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
447
                yield self, key, node[1]
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
448
        keys.difference_update(self._keys)
449
        for backing in self._backing_indices:
450
            if backing is None:
451
                continue
452
            if not keys:
453
                return
454
            for node in backing.iter_entries(keys):
455
                keys.remove(node[1])
456
                yield (self,) + node[1:]
457
458
    def iter_entries_prefix(self, keys):
459
        """Iterate over keys within the index using prefix matching.
460
461
        Prefix matching is applied within the tuple of a key, not to within
462
        the bytestring of each key element. e.g. if you have the keys ('foo',
463
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
464
        only the former key is returned.
465
466
        :param keys: An iterable providing the key prefixes to be retrieved.
467
            Each key prefix takes the form of a tuple the length of a key, but
468
            with the last N elements 'None' rather than a regular bytestring.
469
            The first element cannot be 'None'.
470
        :return: An iterable as per iter_all_entries, but restricted to the
471
            keys with a matching prefix to those supplied. No additional keys
472
            will be returned, and every match that is in the index will be
473
            returned.
474
        """
475
        # XXX: To much duplication with the GraphIndex class; consider finding
476
        # a good place to pull out the actual common logic.
477
        keys = set(keys)
478
        if not keys:
479
            return
480
        for backing in self._backing_indices:
481
            if backing is None:
482
                continue
483
            for node in backing.iter_entries_prefix(keys):
484
                yield (self,) + node[1:]
485
        if self._key_length == 1:
486
            for key in keys:
487
                # sanity check
488
                if key[0] is None:
489
                    raise errors.BadIndexKey(key)
490
                if len(key) != self._key_length:
491
                    raise errors.BadIndexKey(key)
492
                try:
493
                    node = self._nodes[key]
494
                except KeyError:
495
                    continue
496
                if self.reference_lists:
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
497
                    yield self, key, node[1], node[0]
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
498
                else:
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
499
                    yield self, key, node[1]
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
500
            return
501
        for key in keys:
502
            # sanity check
503
            if key[0] is None:
504
                raise errors.BadIndexKey(key)
505
            if len(key) != self._key_length:
506
                raise errors.BadIndexKey(key)
507
            # find what it refers to:
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
508
            key_dict = self._get_nodes_by_key()
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
509
            elements = list(key)
510
            # find the subdict to return
511
            try:
512
                while len(elements) and elements[0] is not None:
513
                    key_dict = key_dict[elements[0]]
514
                    elements.pop(0)
515
            except KeyError:
516
                # a non-existant lookup.
517
                continue
518
            if len(elements):
519
                dicts = [key_dict]
520
                while dicts:
521
                    key_dict = dicts.pop(-1)
522
                    # can't be empty or would not exist
523
                    item, value = key_dict.iteritems().next()
524
                    if type(value) == dict:
525
                        # push keys
526
                        dicts.extend(key_dict.itervalues())
527
                    else:
528
                        # yield keys
529
                        for value in key_dict.itervalues():
530
                            yield (self, ) + value
531
            else:
532
                yield (self, ) + key_dict
533
3644.2.1 by John Arbash Meinel
Change the IndexBuilders to not generate the nodes_by_key unless needed.
534
    def _get_nodes_by_key(self):
535
        if self._nodes_by_key is None:
536
            nodes_by_key = {}
537
            if self.reference_lists:
538
                for key, (references, value) in self._nodes.iteritems():
539
                    key_dict = nodes_by_key
540
                    for subkey in key[:-1]:
541
                        key_dict = key_dict.setdefault(subkey, {})
542
                    key_dict[key[-1]] = key, value, references
543
            else:
544
                for key, (references, value) in self._nodes.iteritems():
545
                    key_dict = nodes_by_key
546
                    for subkey in key[:-1]:
547
                        key_dict = key_dict.setdefault(subkey, {})
548
                    key_dict[key[-1]] = key, value
549
            self._nodes_by_key = nodes_by_key
550
        return self._nodes_by_key
551
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
552
    def key_count(self):
553
        """Return an estimate of the number of keys in this index.
554
555
        For InMemoryGraphIndex the estimate is exact.
556
        """
557
        return len(self._keys) + sum(backing.key_count() for backing in
558
            self._backing_indices if backing is not None)
559
560
    def validate(self):
561
        """In memory index's have no known corruption at the moment."""
562
563
564
class _LeafNode(object):
565
    """A leaf node for a serialised B+Tree index."""
566
567
    def __init__(self, bytes, key_length, ref_list_length):
568
        """Parse bytes to create a leaf node object."""
569
        # splitlines mangles the \r delimiters.. don't use it.
3641.3.30 by John Arbash Meinel
Rename _parse_btree to _btree_serializer
570
        self.keys = dict(_btree_serializer._parse_leaf_lines(bytes,
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
571
            key_length, ref_list_length))
572
573
574
class _InternalNode(object):
575
    """An internal node for a serialised B+Tree index."""
576
577
    def __init__(self, bytes):
578
        """Parse bytes to create an internal node object."""
579
        # splitlines mangles the \r delimiters.. don't use it.
580
        self.keys = self._parse_lines(bytes.split('\n'))
581
582
    def _parse_lines(self, lines):
583
        nodes = []
584
        self.offset = int(lines[1][7:])
585
        for line in lines[2:]:
586
            if line == '':
587
                break
588
            nodes.append(tuple(line.split('\0')))
589
        return nodes
590
591
592
class BTreeGraphIndex(object):
593
    """Access to nodes via the standard GraphIndex interface for B+Tree's.
594
595
    Individual nodes are held in a LRU cache. This holds the root node in
596
    memory except when very large walks are done.
597
    """
598
599
    def __init__(self, transport, name, size):
600
        """Create a B+Tree index object on the index name.
601
602
        :param transport: The transport to read data for the index from.
603
        :param name: The file name of the index on transport.
604
        :param size: Optional size of the index in bytes. This allows
605
            compatibility with the GraphIndex API, as well as ensuring that
606
            the initial read (to read the root node header) can be done
607
            without over-reading even on empty indices, and on small indices
608
            allows single-IO to read the entire index.
609
        """
610
        self._transport = transport
611
        self._name = name
612
        self._size = size
613
        self._file = None
614
        self._page_size = transport.recommended_page_size()
615
        self._root_node = None
616
        # Default max size is 100,000 leave values
617
        self._leaf_value_cache = None # lru_cache.LRUCache(100*1000)
618
        self._leaf_node_cache = lru_cache.LRUCache(_NODE_CACHE_SIZE)
619
        self._internal_node_cache = lru_cache.LRUCache()
620
        self._key_count = None
621
        self._row_lengths = None
622
        self._row_offsets = None # Start of each row, [-1] is the end
623
624
    def __eq__(self, other):
625
        """Equal when self and other were created with the same parameters."""
626
        return (
627
            type(self) == type(other) and
628
            self._transport == other._transport and
629
            self._name == other._name and
630
            self._size == other._size)
631
632
    def __ne__(self, other):
633
        return not self.__eq__(other)
634
635
    def _get_root_node(self):
636
        if self._root_node is None:
637
            # We may not have a root node yet
638
            nodes = list(self._read_nodes([0]))
639
            if len(nodes):
640
                self._root_node = nodes[0][1]
641
        return self._root_node
642
643
    def _cache_nodes(self, nodes, cache):
644
        """Read nodes and cache them in the lru.
645
646
        The nodes list supplied is sorted and then read from disk, each node
647
        being inserted it into the _node_cache.
648
649
        Note: Asking for more nodes than the _node_cache can contain will
650
        result in some of the results being immediately discarded, to prevent
651
        this an assertion is raised if more nodes are asked for than are
652
        cachable.
653
654
        :return: A dict of {node_pos: node}
655
        """
656
        if len(nodes) > cache._max_cache:
657
            trace.mutter('Requesting %s > %s nodes, not all will be cached',
658
                         len(nodes), cache._max_cache)
659
        found = {}
660
        for node_pos, node in self._read_nodes(sorted(nodes)):
661
            if node_pos == 0: # Special case
662
                self._root_node = node
663
            else:
664
                cache.add(node_pos, node)
665
            found[node_pos] = node
666
        return found
667
3641.5.18 by John Arbash Meinel
Clean out the global state, good for prototyping and tuning, bad for production code.
668
    def _get_nodes(self, cache, node_indexes):
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
669
        found = {}
670
        needed = []
671
        for idx in node_indexes:
672
            if idx == 0 and self._root_node is not None:
673
                found[0] = self._root_node
674
                continue
675
            try:
676
                found[idx] = cache[idx]
677
            except KeyError:
678
                needed.append(idx)
679
        found.update(self._cache_nodes(needed, cache))
680
        return found
681
682
    def _get_internal_nodes(self, node_indexes):
683
        """Get a node, from cache or disk.
684
685
        After getting it, the node will be cached.
686
        """
3641.5.18 by John Arbash Meinel
Clean out the global state, good for prototyping and tuning, bad for production code.
687
        return self._get_nodes(self._internal_node_cache, node_indexes)
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
688
689
    def _get_leaf_nodes(self, node_indexes):
690
        """Get a bunch of nodes, from cache or disk."""
3641.5.18 by John Arbash Meinel
Clean out the global state, good for prototyping and tuning, bad for production code.
691
        found = self._get_nodes(self._leaf_node_cache, node_indexes)
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
692
        if self._leaf_value_cache is not None:
693
            for node in found.itervalues():
694
                for key, value in node.keys.iteritems():
695
                    if key in self._leaf_value_cache:
696
                        # Don't add the rest of the keys, we've seen this node
697
                        # before.
698
                        break
699
                    self._leaf_value_cache[key] = value
700
        return found
701
702
    def iter_all_entries(self):
703
        """Iterate over all keys within the index.
704
705
        :return: An iterable of (index, key, value) or (index, key, value, reference_lists).
706
            The former tuple is used when there are no reference lists in the
707
            index, making the API compatible with simple key:value index types.
708
            There is no defined order for the result iteration - it will be in
709
            the most efficient order for the index.
710
        """
711
        if 'evil' in debug.debug_flags:
712
            trace.mutter_callsite(3,
713
                "iter_all_entries scales with size of history.")
714
        if not self.key_count():
715
            return
716
        start_of_leaves = self._row_offsets[-2]
717
        end_of_leaves = self._row_offsets[-1]
718
        needed_nodes = range(start_of_leaves, end_of_leaves)
719
        # We iterate strictly in-order so that we can use this function
720
        # for spilling index builds to disk.
721
        if self.node_ref_lists:
722
            for _, node in self._read_nodes(needed_nodes):
723
                for key, (value, refs) in sorted(node.keys.items()):
724
                    yield (self, key, value, refs)
725
        else:
726
            for _, node in self._read_nodes(needed_nodes):
727
                for key, (value, refs) in sorted(node.keys.items()):
728
                    yield (self, key, value)
729
730
    @staticmethod
731
    def _multi_bisect_right(in_keys, fixed_keys):
732
        """Find the positions where each 'in_key' would fit in fixed_keys.
733
734
        This is equivalent to doing "bisect_right" on each in_key into
735
        fixed_keys
736
737
        :param in_keys: A sorted list of keys to match with fixed_keys
738
        :param fixed_keys: A sorted list of keys to match against
739
        :return: A list of (integer position, [key list]) tuples.
740
        """
741
        if not in_keys:
742
            return []
743
        if not fixed_keys:
744
            # no pointers in the fixed_keys list, which means everything must
745
            # fall to the left.
746
            return [(0, in_keys)]
747
748
        # TODO: Iterating both lists will generally take M + N steps
749
        #       Bisecting each key will generally take M * log2 N steps.
750
        #       If we had an efficient way to compare, we could pick the method
751
        #       based on which has the fewer number of steps.
752
        #       There is also the argument that bisect_right is a compiled
753
        #       function, so there is even more to be gained.
754
        # iter_steps = len(in_keys) + len(fixed_keys)
755
        # bisect_steps = len(in_keys) * math.log(len(fixed_keys), 2)
756
        if len(in_keys) == 1: # Bisect will always be faster for M = 1
757
            return [(bisect_right(fixed_keys, in_keys[0]), in_keys)]
758
        # elif bisect_steps < iter_steps:
759
        #     offsets = {}
760
        #     for key in in_keys:
761
        #         offsets.setdefault(bisect_right(fixed_keys, key),
762
        #                            []).append(key)
763
        #     return [(o, offsets[o]) for o in sorted(offsets)]
764
        in_keys_iter = iter(in_keys)
765
        fixed_keys_iter = enumerate(fixed_keys)
766
        cur_in_key = in_keys_iter.next()
767
        cur_fixed_offset, cur_fixed_key = fixed_keys_iter.next()
768
769
        class InputDone(Exception): pass
770
        class FixedDone(Exception): pass
771
772
        output = []
773
        cur_out = []
774
775
        # TODO: Another possibility is that rather than iterating on each side,
776
        #       we could use a combination of bisecting and iterating. For
777
        #       example, while cur_in_key < fixed_key, bisect to find its
778
        #       point, then iterate all matching keys, then bisect (restricted
779
        #       to only the remainder) for the next one, etc.
780
        try:
781
            while True:
782
                if cur_in_key < cur_fixed_key:
783
                    cur_keys = []
784
                    cur_out = (cur_fixed_offset, cur_keys)
785
                    output.append(cur_out)
786
                    while cur_in_key < cur_fixed_key:
787
                        cur_keys.append(cur_in_key)
788
                        try:
789
                            cur_in_key = in_keys_iter.next()
790
                        except StopIteration:
791
                            raise InputDone
792
                    # At this point cur_in_key must be >= cur_fixed_key
793
                # step the cur_fixed_key until we pass the cur key, or walk off
794
                # the end
795
                while cur_in_key >= cur_fixed_key:
796
                    try:
797
                        cur_fixed_offset, cur_fixed_key = fixed_keys_iter.next()
798
                    except StopIteration:
799
                        raise FixedDone
800
        except InputDone:
801
            # We consumed all of the input, nothing more to do
802
            pass
803
        except FixedDone:
804
            # There was some input left, but we consumed all of fixed, so we
805
            # have to add one more for the tail
806
            cur_keys = [cur_in_key]
807
            cur_keys.extend(in_keys_iter)
808
            cur_out = (len(fixed_keys), cur_keys)
809
            output.append(cur_out)
810
        return output
811
812
    def iter_entries(self, keys):
813
        """Iterate over keys within the index.
814
815
        :param keys: An iterable providing the keys to be retrieved.
816
        :return: An iterable as per iter_all_entries, but restricted to the
817
            keys supplied. No additional keys will be returned, and every
818
            key supplied that is in the index will be returned.
819
        """
820
        # 6 seconds spent in miss_torture using the sorted() line.
821
        # Even with out of order disk IO it seems faster not to sort it when
822
        # large queries are being made.
823
        # However, now that we are doing multi-way bisecting, we need the keys
824
        # in sorted order anyway. We could change the multi-way code to not
825
        # require sorted order. (For example, it bisects for the first node,
826
        # does an in-order search until a key comes before the current point,
827
        # which it then bisects for, etc.)
828
        keys = frozenset(keys)
829
        if not keys:
830
            return
831
832
        if not self.key_count():
833
            return
834
835
        needed_keys = []
836
        if self._leaf_value_cache is None:
837
            needed_keys = keys
838
        else:
839
            for key in keys:
840
                value = self._leaf_value_cache.get(key, None)
841
                if value is not None:
842
                    # This key is known not to be here, skip it
843
                    value, refs = value
844
                    if self.node_ref_lists:
845
                        yield (self, key, value, refs)
846
                    else:
847
                        yield (self, key, value)
848
                else:
849
                    needed_keys.append(key)
850
851
        last_key = None
852
        needed_keys = keys
853
        if not needed_keys:
854
            return
855
        # 6 seconds spent in miss_torture using the sorted() line.
856
        # Even with out of order disk IO it seems faster not to sort it when
857
        # large queries are being made.
858
        needed_keys = sorted(needed_keys)
859
860
        nodes_and_keys = [(0, needed_keys)]
861
862
        for row_pos, next_row_start in enumerate(self._row_offsets[1:-1]):
863
            node_indexes = [idx for idx, s_keys in nodes_and_keys]
864
            nodes = self._get_internal_nodes(node_indexes)
865
866
            next_nodes_and_keys = []
867
            for node_index, sub_keys in nodes_and_keys:
868
                node = nodes[node_index]
869
                positions = self._multi_bisect_right(sub_keys, node.keys)
870
                node_offset = next_row_start + node.offset
871
                next_nodes_and_keys.extend([(node_offset + pos, s_keys)
872
                                           for pos, s_keys in positions])
873
            nodes_and_keys = next_nodes_and_keys
874
        # We should now be at the _LeafNodes
875
        node_indexes = [idx for idx, s_keys in nodes_and_keys]
876
877
        # TODO: We may *not* want to always read all the nodes in one
878
        #       big go. Consider setting a max size on this.
879
880
        nodes = self._get_leaf_nodes(node_indexes)
881
        for node_index, sub_keys in nodes_and_keys:
882
            if not sub_keys:
883
                continue
884
            node = nodes[node_index]
885
            for next_sub_key in sub_keys:
886
                if next_sub_key in node.keys:
887
                    value, refs = node.keys[next_sub_key]
888
                    if self.node_ref_lists:
889
                        yield (self, next_sub_key, value, refs)
890
                    else:
891
                        yield (self, next_sub_key, value)
892
893
    def iter_entries_prefix(self, keys):
894
        """Iterate over keys within the index using prefix matching.
895
896
        Prefix matching is applied within the tuple of a key, not to within
897
        the bytestring of each key element. e.g. if you have the keys ('foo',
898
        'bar'), ('foobar', 'gam') and do a prefix search for ('foo', None) then
899
        only the former key is returned.
900
901
        WARNING: Note that this method currently causes a full index parse
902
        unconditionally (which is reasonably appropriate as it is a means for
903
        thunking many small indices into one larger one and still supplies
904
        iter_all_entries at the thunk layer).
905
906
        :param keys: An iterable providing the key prefixes to be retrieved.
907
            Each key prefix takes the form of a tuple the length of a key, but
908
            with the last N elements 'None' rather than a regular bytestring.
909
            The first element cannot be 'None'.
910
        :return: An iterable as per iter_all_entries, but restricted to the
911
            keys with a matching prefix to those supplied. No additional keys
912
            will be returned, and every match that is in the index will be
913
            returned.
914
        """
915
        keys = sorted(set(keys))
916
        if not keys:
917
            return
918
        # Load if needed to check key lengths
919
        if self._key_count is None:
920
            self._get_root_node()
921
        # TODO: only access nodes that can satisfy the prefixes we are looking
922
        # for. For now, to meet API usage (as this function is not used by
923
        # current bzrlib) just suck the entire index and iterate in memory.
924
        nodes = {}
925
        if self.node_ref_lists:
926
            if self._key_length == 1:
927
                for _1, key, value, refs in self.iter_all_entries():
928
                    nodes[key] = value, refs
929
            else:
930
                nodes_by_key = {}
931
                for _1, key, value, refs in self.iter_all_entries():
932
                    key_value = key, value, refs
933
                    # For a key of (foo, bar, baz) create
934
                    # _nodes_by_key[foo][bar][baz] = key_value
935
                    key_dict = nodes_by_key
936
                    for subkey in key[:-1]:
937
                        key_dict = key_dict.setdefault(subkey, {})
938
                    key_dict[key[-1]] = key_value
939
        else:
940
            if self._key_length == 1:
941
                for _1, key, value in self.iter_all_entries():
942
                    nodes[key] = value
943
            else:
944
                nodes_by_key = {}
945
                for _1, key, value in self.iter_all_entries():
946
                    key_value = key, value
947
                    # For a key of (foo, bar, baz) create
948
                    # _nodes_by_key[foo][bar][baz] = key_value
949
                    key_dict = nodes_by_key
950
                    for subkey in key[:-1]:
951
                        key_dict = key_dict.setdefault(subkey, {})
952
                    key_dict[key[-1]] = key_value
953
        if self._key_length == 1:
954
            for key in keys:
955
                # sanity check
956
                if key[0] is None:
957
                    raise errors.BadIndexKey(key)
958
                if len(key) != self._key_length:
959
                    raise errors.BadIndexKey(key)
960
                try:
961
                    if self.node_ref_lists:
962
                        value, node_refs = nodes[key]
963
                        yield self, key, value, node_refs
964
                    else:
965
                        yield self, key, nodes[key]
966
                except KeyError:
967
                    pass
968
            return
969
        for key in keys:
970
            # sanity check
971
            if key[0] is None:
972
                raise errors.BadIndexKey(key)
973
            if len(key) != self._key_length:
974
                raise errors.BadIndexKey(key)
975
            # find what it refers to:
976
            key_dict = nodes_by_key
977
            elements = list(key)
978
            # find the subdict whose contents should be returned.
979
            try:
980
                while len(elements) and elements[0] is not None:
981
                    key_dict = key_dict[elements[0]]
982
                    elements.pop(0)
983
            except KeyError:
984
                # a non-existant lookup.
985
                continue
986
            if len(elements):
987
                dicts = [key_dict]
988
                while dicts:
989
                    key_dict = dicts.pop(-1)
990
                    # can't be empty or would not exist
991
                    item, value = key_dict.iteritems().next()
992
                    if type(value) == dict:
993
                        # push keys
994
                        dicts.extend(key_dict.itervalues())
995
                    else:
996
                        # yield keys
997
                        for value in key_dict.itervalues():
998
                            # each value is the key:value:node refs tuple
999
                            # ready to yield.
1000
                            yield (self, ) + value
1001
            else:
1002
                # the last thing looked up was a terminal element
1003
                yield (self, ) + key_dict
1004
1005
    def key_count(self):
1006
        """Return an estimate of the number of keys in this index.
1007
1008
        For BTreeGraphIndex the estimate is exact as it is contained in the
1009
        header.
1010
        """
1011
        if self._key_count is None:
1012
            self._get_root_node()
1013
        return self._key_count
1014
1015
    def _parse_header_from_bytes(self, bytes):
1016
        """Parse the header from a region of bytes.
1017
1018
        :param bytes: The data to parse.
1019
        :return: An offset, data tuple such as readv yields, for the unparsed
1020
            data. (which may be of length 0).
1021
        """
1022
        signature = bytes[0:len(self._signature())]
1023
        if not signature == self._signature():
1024
            raise errors.BadIndexFormatSignature(self._name, BTreeGraphIndex)
1025
        lines = bytes[len(self._signature()):].splitlines()
1026
        options_line = lines[0]
1027
        if not options_line.startswith(_OPTION_NODE_REFS):
1028
            raise errors.BadIndexOptions(self)
1029
        try:
1030
            self.node_ref_lists = int(options_line[len(_OPTION_NODE_REFS):])
1031
        except ValueError:
1032
            raise errors.BadIndexOptions(self)
1033
        options_line = lines[1]
1034
        if not options_line.startswith(_OPTION_KEY_ELEMENTS):
1035
            raise errors.BadIndexOptions(self)
1036
        try:
1037
            self._key_length = int(options_line[len(_OPTION_KEY_ELEMENTS):])
1038
        except ValueError:
1039
            raise errors.BadIndexOptions(self)
1040
        options_line = lines[2]
1041
        if not options_line.startswith(_OPTION_LEN):
1042
            raise errors.BadIndexOptions(self)
1043
        try:
1044
            self._key_count = int(options_line[len(_OPTION_LEN):])
1045
        except ValueError:
1046
            raise errors.BadIndexOptions(self)
1047
        options_line = lines[3]
1048
        if not options_line.startswith(_OPTION_ROW_LENGTHS):
1049
            raise errors.BadIndexOptions(self)
1050
        try:
1051
            self._row_lengths = map(int, [length for length in
1052
                options_line[len(_OPTION_ROW_LENGTHS):].split(',')
1053
                if len(length)])
1054
        except ValueError:
1055
            raise errors.BadIndexOptions(self)
1056
        offsets = []
1057
        row_offset = 0
1058
        for row in self._row_lengths:
1059
            offsets.append(row_offset)
1060
            row_offset += row
1061
        offsets.append(row_offset)
1062
        self._row_offsets = offsets
1063
1064
        # calculate the bytes we have processed
1065
        header_end = (len(signature) + sum(map(len, lines[0:4])) + 4)
1066
        return header_end, bytes[header_end:]
1067
1068
    def _read_nodes(self, nodes):
1069
        """Read some nodes from disk into the LRU cache.
1070
1071
        This performs a readv to get the node data into memory, and parses each
1072
        node, the yields it to the caller. The nodes are requested in the
1073
        supplied order. If possible doing sort() on the list before requesting
1074
        a read may improve performance.
1075
1076
        :param nodes: The nodes to read. 0 - first node, 1 - second node etc.
1077
        :return: None
1078
        """
1079
        ranges = []
1080
        for index in nodes:
1081
            offset = index * _PAGE_SIZE
1082
            size = _PAGE_SIZE
1083
            if index == 0:
1084
                # Root node - special case
1085
                if self._size:
1086
                    size = min(_PAGE_SIZE, self._size)
1087
                else:
1088
                    stream = self._transport.get(self._name)
1089
                    start = stream.read(_PAGE_SIZE)
1090
                    # Avoid doing this again
1091
                    self._size = len(start)
1092
                    size = min(_PAGE_SIZE, self._size)
1093
            else:
1094
                size = min(size, self._size - offset)
1095
            ranges.append((offset, size))
1096
        if not ranges:
1097
            return
1098
        if self._file is None:
1099
            data_ranges = self._transport.readv(self._name, ranges)
1100
        else:
1101
            data_ranges = []
1102
            for offset, size in ranges:
1103
                self._file.seek(offset)
1104
                data_ranges.append((offset, self._file.read(size)))
1105
        for offset, data in data_ranges:
1106
            if offset == 0:
1107
                # extract the header
1108
                offset, data = self._parse_header_from_bytes(data)
1109
                if len(data) == 0:
1110
                    continue
1111
            bytes = zlib.decompress(data)
1112
            if bytes.startswith(_LEAF_FLAG):
1113
                node = _LeafNode(bytes, self._key_length, self.node_ref_lists)
1114
            elif bytes.startswith(_INTERNAL_FLAG):
1115
                node = _InternalNode(bytes)
1116
            else:
1117
                raise AssertionError("Unknown node type for %r" % bytes)
1118
            yield offset / _PAGE_SIZE, node
1119
1120
    def _signature(self):
1121
        """The file signature for this index type."""
1122
        return _BTSIGNATURE
1123
1124
    def validate(self):
1125
        """Validate that everything in the index can be accessed."""
1126
        # just read and parse every node.
1127
        self._get_root_node()
1128
        if len(self._row_lengths) > 1:
1129
            start_node = self._row_offsets[1]
1130
        else:
1131
            # We shouldn't be reading anything anyway
1132
            start_node = 1
1133
        node_end = self._row_offsets[-1]
1134
        for node in self._read_nodes(range(start_node, node_end)):
1135
            pass
1136
1137
1138
try:
3641.3.30 by John Arbash Meinel
Rename _parse_btree to _btree_serializer
1139
    from bzrlib import _btree_serializer_c as _btree_serializer
3641.3.1 by John Arbash Meinel
Bring in the btree_index and chunk_writer code and their tests.
1140
except ImportError:
3641.3.30 by John Arbash Meinel
Rename _parse_btree to _btree_serializer
1141
    from bzrlib import _btree_serializer_py as _btree_serializer