/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar

« back to all changes in this revision

Viewing changes to bzrlib/chk_map.py

  • Committer: Vincent Ladeuil
  • Date: 2009-07-15 07:32:26 UTC
  • mfrom: (4536 +trunk)
  • mto: (4536.1.1 integration)
  • mto: This revision was merged to the branch mainline in revision 4537.
  • Revision ID: v.ladeuil+lp@free.fr-20090715073226-a7ylxd6ctbzeu0o6
Merge trunk resolving conflicts

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 2008, 2009 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 
16
 
 
17
"""Persistent maps from tuple_of_strings->string using CHK stores.
 
18
 
 
19
Overview and current status:
 
20
 
 
21
The CHKMap class implements a dict from tuple_of_strings->string by using a trie
 
22
with internal nodes of 8-bit fan out; The key tuples are mapped to strings by
 
23
joining them by \x00, and \x00 padding shorter keys out to the length of the
 
24
longest key. Leaf nodes are packed as densely as possible, and internal nodes
 
25
are all an additional 8-bits wide leading to a sparse upper tree.
 
26
 
 
27
Updates to a CHKMap are done preferentially via the apply_delta method, to
 
28
allow optimisation of the update operation; but individual map/unmap calls are
 
29
possible and supported. All changes via map/unmap are buffered in memory until
 
30
the _save method is called to force serialisation of the tree. apply_delta
 
31
performs a _save implicitly.
 
32
 
 
33
TODO:
 
34
-----
 
35
 
 
36
Densely packed upper nodes.
 
37
 
 
38
"""
 
39
 
 
40
import heapq
 
41
 
 
42
from bzrlib import lazy_import
 
43
lazy_import.lazy_import(globals(), """
 
44
from bzrlib import versionedfile
 
45
""")
 
46
from bzrlib import (
 
47
    lru_cache,
 
48
    osutils,
 
49
    registry,
 
50
    trace,
 
51
    )
 
52
 
 
53
# approx 4MB
 
54
# If each line is 50 bytes, and you have 255 internal pages, with 255-way fan
 
55
# out, it takes 3.1MB to cache the layer.
 
56
_PAGE_CACHE_SIZE = 4*1024*1024
 
57
# We are caching bytes so len(value) is perfectly accurate
 
58
_page_cache = lru_cache.LRUSizeCache(_PAGE_CACHE_SIZE)
 
59
 
 
60
# If a ChildNode falls below this many bytes, we check for a remap
 
61
_INTERESTING_NEW_SIZE = 50
 
62
# If a ChildNode shrinks by more than this amount, we check for a remap
 
63
_INTERESTING_SHRINKAGE_LIMIT = 20
 
64
# If we delete more than this many nodes applying a delta, we check for a remap
 
65
_INTERESTING_DELETES_LIMIT = 5
 
66
 
 
67
 
 
68
def _search_key_plain(key):
 
69
    """Map the key tuple into a search string that just uses the key bytes."""
 
70
    return '\x00'.join(key)
 
71
 
 
72
 
 
73
search_key_registry = registry.Registry()
 
74
search_key_registry.register('plain', _search_key_plain)
 
75
 
 
76
 
 
77
class CHKMap(object):
 
78
    """A persistent map from string to string backed by a CHK store."""
 
79
 
 
80
    def __init__(self, store, root_key, search_key_func=None):
 
81
        """Create a CHKMap object.
 
82
 
 
83
        :param store: The store the CHKMap is stored in.
 
84
        :param root_key: The root key of the map. None to create an empty
 
85
            CHKMap.
 
86
        :param search_key_func: A function mapping a key => bytes. These bytes
 
87
            are then used by the internal nodes to split up leaf nodes into
 
88
            multiple pages.
 
89
        """
 
90
        self._store = store
 
91
        if search_key_func is None:
 
92
            search_key_func = _search_key_plain
 
93
        self._search_key_func = search_key_func
 
94
        if root_key is None:
 
95
            self._root_node = LeafNode(search_key_func=search_key_func)
 
96
        else:
 
97
            self._root_node = self._node_key(root_key)
 
98
 
 
99
    def apply_delta(self, delta):
 
100
        """Apply a delta to the map.
 
101
 
 
102
        :param delta: An iterable of old_key, new_key, new_value tuples.
 
103
            If new_key is not None, then new_key->new_value is inserted
 
104
            into the map; if old_key is not None, then the old mapping
 
105
            of old_key is removed.
 
106
        """
 
107
        delete_count = 0
 
108
        for old, new, value in delta:
 
109
            if old is not None and old != new:
 
110
                self.unmap(old, check_remap=False)
 
111
                delete_count += 1
 
112
        for old, new, value in delta:
 
113
            if new is not None:
 
114
                self.map(new, value)
 
115
        if delete_count > _INTERESTING_DELETES_LIMIT:
 
116
            trace.mutter("checking remap as %d deletions", delete_count)
 
117
            self._check_remap()
 
118
        return self._save()
 
119
 
 
120
    def _ensure_root(self):
 
121
        """Ensure that the root node is an object not a key."""
 
122
        if type(self._root_node) is tuple:
 
123
            # Demand-load the root
 
124
            self._root_node = self._get_node(self._root_node)
 
125
 
 
126
    def _get_node(self, node):
 
127
        """Get a node.
 
128
 
 
129
        Note that this does not update the _items dict in objects containing a
 
130
        reference to this node. As such it does not prevent subsequent IO being
 
131
        performed.
 
132
 
 
133
        :param node: A tuple key or node object.
 
134
        :return: A node object.
 
135
        """
 
136
        if type(node) is tuple:
 
137
            bytes = self._read_bytes(node)
 
138
            return _deserialise(bytes, node,
 
139
                search_key_func=self._search_key_func)
 
140
        else:
 
141
            return node
 
142
 
 
143
    def _read_bytes(self, key):
 
144
        try:
 
145
            return _page_cache[key]
 
146
        except KeyError:
 
147
            stream = self._store.get_record_stream([key], 'unordered', True)
 
148
            bytes = stream.next().get_bytes_as('fulltext')
 
149
            _page_cache[key] = bytes
 
150
            return bytes
 
151
 
 
152
    def _dump_tree(self, include_keys=False):
 
153
        """Return the tree in a string representation."""
 
154
        self._ensure_root()
 
155
        res = self._dump_tree_node(self._root_node, prefix='', indent='',
 
156
                                   include_keys=include_keys)
 
157
        res.append('') # Give a trailing '\n'
 
158
        return '\n'.join(res)
 
159
 
 
160
    def _dump_tree_node(self, node, prefix, indent, include_keys=True):
 
161
        """For this node and all children, generate a string representation."""
 
162
        result = []
 
163
        if not include_keys:
 
164
            key_str = ''
 
165
        else:
 
166
            node_key = node.key()
 
167
            if node_key is not None:
 
168
                key_str = ' %s' % (node_key[0],)
 
169
            else:
 
170
                key_str = ' None'
 
171
        result.append('%s%r %s%s' % (indent, prefix, node.__class__.__name__,
 
172
                                     key_str))
 
173
        if type(node) is InternalNode:
 
174
            # Trigger all child nodes to get loaded
 
175
            list(node._iter_nodes(self._store))
 
176
            for prefix, sub in sorted(node._items.iteritems()):
 
177
                result.extend(self._dump_tree_node(sub, prefix, indent + '  ',
 
178
                                                   include_keys=include_keys))
 
179
        else:
 
180
            for key, value in sorted(node._items.iteritems()):
 
181
                # Don't use prefix nor indent here to line up when used in
 
182
                # tests in conjunction with assertEqualDiff
 
183
                result.append('      %r %r' % (key, value))
 
184
        return result
 
185
 
 
186
    @classmethod
 
187
    def from_dict(klass, store, initial_value, maximum_size=0, key_width=1,
 
188
        search_key_func=None):
 
189
        """Create a CHKMap in store with initial_value as the content.
 
190
 
 
191
        :param store: The store to record initial_value in, a VersionedFiles
 
192
            object with 1-tuple keys supporting CHK key generation.
 
193
        :param initial_value: A dict to store in store. Its keys and values
 
194
            must be bytestrings.
 
195
        :param maximum_size: The maximum_size rule to apply to nodes. This
 
196
            determines the size at which no new data is added to a single node.
 
197
        :param key_width: The number of elements in each key_tuple being stored
 
198
            in this map.
 
199
        :param search_key_func: A function mapping a key => bytes. These bytes
 
200
            are then used by the internal nodes to split up leaf nodes into
 
201
            multiple pages.
 
202
        :return: The root chk of the resulting CHKMap.
 
203
        """
 
204
        root_key = klass._create_directly(store, initial_value,
 
205
            maximum_size=maximum_size, key_width=key_width,
 
206
            search_key_func=search_key_func)
 
207
        return root_key
 
208
 
 
209
    @classmethod
 
210
    def _create_via_map(klass, store, initial_value, maximum_size=0,
 
211
                        key_width=1, search_key_func=None):
 
212
        result = klass(store, None, search_key_func=search_key_func)
 
213
        result._root_node.set_maximum_size(maximum_size)
 
214
        result._root_node._key_width = key_width
 
215
        delta = []
 
216
        for key, value in initial_value.items():
 
217
            delta.append((None, key, value))
 
218
        root_key = result.apply_delta(delta)
 
219
        return root_key
 
220
 
 
221
    @classmethod
 
222
    def _create_directly(klass, store, initial_value, maximum_size=0,
 
223
                         key_width=1, search_key_func=None):
 
224
        node = LeafNode(search_key_func=search_key_func)
 
225
        node.set_maximum_size(maximum_size)
 
226
        node._key_width = key_width
 
227
        node._items = dict(initial_value)
 
228
        node._raw_size = sum([node._key_value_len(key, value)
 
229
                              for key,value in initial_value.iteritems()])
 
230
        node._len = len(node._items)
 
231
        node._compute_search_prefix()
 
232
        node._compute_serialised_prefix()
 
233
        if (node._len > 1
 
234
            and maximum_size
 
235
            and node._current_size() > maximum_size):
 
236
            prefix, node_details = node._split(store)
 
237
            if len(node_details) == 1:
 
238
                raise AssertionError('Failed to split using node._split')
 
239
            node = InternalNode(prefix, search_key_func=search_key_func)
 
240
            node.set_maximum_size(maximum_size)
 
241
            node._key_width = key_width
 
242
            for split, subnode in node_details:
 
243
                node.add_node(split, subnode)
 
244
        keys = list(node.serialise(store))
 
245
        return keys[-1]
 
246
 
 
247
    def iter_changes(self, basis):
 
248
        """Iterate over the changes between basis and self.
 
249
 
 
250
        :return: An iterator of tuples: (key, old_value, new_value). Old_value
 
251
            is None for keys only in self; new_value is None for keys only in
 
252
            basis.
 
253
        """
 
254
        # Overview:
 
255
        # Read both trees in lexographic, highest-first order.
 
256
        # Any identical nodes we skip
 
257
        # Any unique prefixes we output immediately.
 
258
        # values in a leaf node are treated as single-value nodes in the tree
 
259
        # which allows them to be not-special-cased. We know to output them
 
260
        # because their value is a string, not a key(tuple) or node.
 
261
        #
 
262
        # corner cases to beware of when considering this function:
 
263
        # *) common references are at different heights.
 
264
        #    consider two trees:
 
265
        #    {'a': LeafNode={'aaa':'foo', 'aab':'bar'}, 'b': LeafNode={'b'}}
 
266
        #    {'a': InternalNode={'aa':LeafNode={'aaa':'foo', 'aab':'bar'},
 
267
        #                        'ab':LeafNode={'ab':'bar'}}
 
268
        #     'b': LeafNode={'b'}}
 
269
        #    the node with aaa/aab will only be encountered in the second tree
 
270
        #    after reading the 'a' subtree, but it is encountered in the first
 
271
        #    tree immediately. Variations on this may have read internal nodes
 
272
        #    like this.  we want to cut the entire pending subtree when we
 
273
        #    realise we have a common node.  For this we use a list of keys -
 
274
        #    the path to a node - and check the entire path is clean as we
 
275
        #    process each item.
 
276
        if self._node_key(self._root_node) == self._node_key(basis._root_node):
 
277
            return
 
278
        self._ensure_root()
 
279
        basis._ensure_root()
 
280
        excluded_keys = set()
 
281
        self_node = self._root_node
 
282
        basis_node = basis._root_node
 
283
        # A heap, each element is prefix, node(tuple/NodeObject/string),
 
284
        # key_path (a list of tuples, tail-sharing down the tree.)
 
285
        self_pending = []
 
286
        basis_pending = []
 
287
        def process_node(node, path, a_map, pending):
 
288
            # take a node and expand it
 
289
            node = a_map._get_node(node)
 
290
            if type(node) == LeafNode:
 
291
                path = (node._key, path)
 
292
                for key, value in node._items.items():
 
293
                    # For a LeafNode, the key is a serialized_key, rather than
 
294
                    # a search_key, but the heap is using search_keys
 
295
                    search_key = node._search_key_func(key)
 
296
                    heapq.heappush(pending, (search_key, key, value, path))
 
297
            else:
 
298
                # type(node) == InternalNode
 
299
                path = (node._key, path)
 
300
                for prefix, child in node._items.items():
 
301
                    heapq.heappush(pending, (prefix, None, child, path))
 
302
        def process_common_internal_nodes(self_node, basis_node):
 
303
            self_items = set(self_node._items.items())
 
304
            basis_items = set(basis_node._items.items())
 
305
            path = (self_node._key, None)
 
306
            for prefix, child in self_items - basis_items:
 
307
                heapq.heappush(self_pending, (prefix, None, child, path))
 
308
            path = (basis_node._key, None)
 
309
            for prefix, child in basis_items - self_items:
 
310
                heapq.heappush(basis_pending, (prefix, None, child, path))
 
311
        def process_common_leaf_nodes(self_node, basis_node):
 
312
            self_items = set(self_node._items.items())
 
313
            basis_items = set(basis_node._items.items())
 
314
            path = (self_node._key, None)
 
315
            for key, value in self_items - basis_items:
 
316
                prefix = self._search_key_func(key)
 
317
                heapq.heappush(self_pending, (prefix, key, value, path))
 
318
            path = (basis_node._key, None)
 
319
            for key, value in basis_items - self_items:
 
320
                prefix = basis._search_key_func(key)
 
321
                heapq.heappush(basis_pending, (prefix, key, value, path))
 
322
        def process_common_prefix_nodes(self_node, self_path,
 
323
                                        basis_node, basis_path):
 
324
            # Would it be more efficient if we could request both at the same
 
325
            # time?
 
326
            self_node = self._get_node(self_node)
 
327
            basis_node = basis._get_node(basis_node)
 
328
            if (type(self_node) == InternalNode
 
329
                and type(basis_node) == InternalNode):
 
330
                # Matching internal nodes
 
331
                process_common_internal_nodes(self_node, basis_node)
 
332
            elif (type(self_node) == LeafNode
 
333
                  and type(basis_node) == LeafNode):
 
334
                process_common_leaf_nodes(self_node, basis_node)
 
335
            else:
 
336
                process_node(self_node, self_path, self, self_pending)
 
337
                process_node(basis_node, basis_path, basis, basis_pending)
 
338
        process_common_prefix_nodes(self_node, None, basis_node, None)
 
339
        self_seen = set()
 
340
        basis_seen = set()
 
341
        excluded_keys = set()
 
342
        def check_excluded(key_path):
 
343
            # Note that this is N^2, it depends on us trimming trees
 
344
            # aggressively to not become slow.
 
345
            # A better implementation would probably have a reverse map
 
346
            # back to the children of a node, and jump straight to it when
 
347
            # a common node is detected, the proceed to remove the already
 
348
            # pending children. bzrlib.graph has a searcher module with a
 
349
            # similar problem.
 
350
            while key_path is not None:
 
351
                key, key_path = key_path
 
352
                if key in excluded_keys:
 
353
                    return True
 
354
            return False
 
355
 
 
356
        loop_counter = 0
 
357
        while self_pending or basis_pending:
 
358
            loop_counter += 1
 
359
            if not self_pending:
 
360
                # self is exhausted: output remainder of basis
 
361
                for prefix, key, node, path in basis_pending:
 
362
                    if check_excluded(path):
 
363
                        continue
 
364
                    node = basis._get_node(node)
 
365
                    if key is not None:
 
366
                        # a value
 
367
                        yield (key, node, None)
 
368
                    else:
 
369
                        # subtree - fastpath the entire thing.
 
370
                        for key, value in node.iteritems(basis._store):
 
371
                            yield (key, value, None)
 
372
                return
 
373
            elif not basis_pending:
 
374
                # basis is exhausted: output remainder of self.
 
375
                for prefix, key, node, path in self_pending:
 
376
                    if check_excluded(path):
 
377
                        continue
 
378
                    node = self._get_node(node)
 
379
                    if key is not None:
 
380
                        # a value
 
381
                        yield (key, None, node)
 
382
                    else:
 
383
                        # subtree - fastpath the entire thing.
 
384
                        for key, value in node.iteritems(self._store):
 
385
                            yield (key, None, value)
 
386
                return
 
387
            else:
 
388
                # XXX: future optimisation - yield the smaller items
 
389
                # immediately rather than pushing everything on/off the
 
390
                # heaps. Applies to both internal nodes and leafnodes.
 
391
                if self_pending[0][0] < basis_pending[0][0]:
 
392
                    # expand self
 
393
                    prefix, key, node, path = heapq.heappop(self_pending)
 
394
                    if check_excluded(path):
 
395
                        continue
 
396
                    if key is not None:
 
397
                        # a value
 
398
                        yield (key, None, node)
 
399
                    else:
 
400
                        process_node(node, path, self, self_pending)
 
401
                        continue
 
402
                elif self_pending[0][0] > basis_pending[0][0]:
 
403
                    # expand basis
 
404
                    prefix, key, node, path = heapq.heappop(basis_pending)
 
405
                    if check_excluded(path):
 
406
                        continue
 
407
                    if key is not None:
 
408
                        # a value
 
409
                        yield (key, node, None)
 
410
                    else:
 
411
                        process_node(node, path, basis, basis_pending)
 
412
                        continue
 
413
                else:
 
414
                    # common prefix: possibly expand both
 
415
                    if self_pending[0][1] is None:
 
416
                        # process next self
 
417
                        read_self = True
 
418
                    else:
 
419
                        read_self = False
 
420
                    if basis_pending[0][1] is None:
 
421
                        # process next basis
 
422
                        read_basis = True
 
423
                    else:
 
424
                        read_basis = False
 
425
                    if not read_self and not read_basis:
 
426
                        # compare a common value
 
427
                        self_details = heapq.heappop(self_pending)
 
428
                        basis_details = heapq.heappop(basis_pending)
 
429
                        if self_details[2] != basis_details[2]:
 
430
                            yield (self_details[1],
 
431
                                basis_details[2], self_details[2])
 
432
                        continue
 
433
                    # At least one side wasn't a simple value
 
434
                    if (self._node_key(self_pending[0][2]) ==
 
435
                        self._node_key(basis_pending[0][2])):
 
436
                        # Identical pointers, skip (and don't bother adding to
 
437
                        # excluded, it won't turn up again.
 
438
                        heapq.heappop(self_pending)
 
439
                        heapq.heappop(basis_pending)
 
440
                        continue
 
441
                    # Now we need to expand this node before we can continue
 
442
                    if read_self and read_basis:
 
443
                        # Both sides start with the same prefix, so process
 
444
                        # them in parallel
 
445
                        self_prefix, _, self_node, self_path = heapq.heappop(
 
446
                            self_pending)
 
447
                        basis_prefix, _, basis_node, basis_path = heapq.heappop(
 
448
                            basis_pending)
 
449
                        if self_prefix != basis_prefix:
 
450
                            raise AssertionError(
 
451
                                '%r != %r' % (self_prefix, basis_prefix))
 
452
                        process_common_prefix_nodes(
 
453
                            self_node, self_path,
 
454
                            basis_node, basis_path)
 
455
                        continue
 
456
                    if read_self:
 
457
                        prefix, key, node, path = heapq.heappop(self_pending)
 
458
                        if check_excluded(path):
 
459
                            continue
 
460
                        process_node(node, path, self, self_pending)
 
461
                    if read_basis:
 
462
                        prefix, key, node, path = heapq.heappop(basis_pending)
 
463
                        if check_excluded(path):
 
464
                            continue
 
465
                        process_node(node, path, basis, basis_pending)
 
466
        # print loop_counter
 
467
 
 
468
    def iteritems(self, key_filter=None):
 
469
        """Iterate over the entire CHKMap's contents."""
 
470
        self._ensure_root()
 
471
        return self._root_node.iteritems(self._store, key_filter=key_filter)
 
472
 
 
473
    def key(self):
 
474
        """Return the key for this map."""
 
475
        if type(self._root_node) is tuple:
 
476
            return self._root_node
 
477
        else:
 
478
            return self._root_node._key
 
479
 
 
480
    def __len__(self):
 
481
        self._ensure_root()
 
482
        return len(self._root_node)
 
483
 
 
484
    def map(self, key, value):
 
485
        """Map a key tuple to value."""
 
486
        # Need a root object.
 
487
        self._ensure_root()
 
488
        prefix, node_details = self._root_node.map(self._store, key, value)
 
489
        if len(node_details) == 1:
 
490
            self._root_node = node_details[0][1]
 
491
        else:
 
492
            self._root_node = InternalNode(prefix,
 
493
                                search_key_func=self._search_key_func)
 
494
            self._root_node.set_maximum_size(node_details[0][1].maximum_size)
 
495
            self._root_node._key_width = node_details[0][1]._key_width
 
496
            for split, node in node_details:
 
497
                self._root_node.add_node(split, node)
 
498
 
 
499
    def _node_key(self, node):
 
500
        """Get the key for a node whether it's a tuple or node."""
 
501
        if type(node) is tuple:
 
502
            return node
 
503
        else:
 
504
            return node._key
 
505
 
 
506
    def unmap(self, key, check_remap=True):
 
507
        """remove key from the map."""
 
508
        self._ensure_root()
 
509
        if type(self._root_node) is InternalNode:
 
510
            unmapped = self._root_node.unmap(self._store, key,
 
511
                check_remap=check_remap)
 
512
        else:
 
513
            unmapped = self._root_node.unmap(self._store, key)
 
514
        self._root_node = unmapped
 
515
 
 
516
    def _check_remap(self):
 
517
        """Check if nodes can be collapsed."""
 
518
        self._ensure_root()
 
519
        if type(self._root_node) is InternalNode:
 
520
            self._root_node._check_remap(self._store)
 
521
 
 
522
    def _save(self):
 
523
        """Save the map completely.
 
524
 
 
525
        :return: The key of the root node.
 
526
        """
 
527
        if type(self._root_node) is tuple:
 
528
            # Already saved.
 
529
            return self._root_node
 
530
        keys = list(self._root_node.serialise(self._store))
 
531
        return keys[-1]
 
532
 
 
533
 
 
534
class Node(object):
 
535
    """Base class defining the protocol for CHK Map nodes.
 
536
 
 
537
    :ivar _raw_size: The total size of the serialized key:value data, before
 
538
        adding the header bytes, and without prefix compression.
 
539
    """
 
540
 
 
541
    def __init__(self, key_width=1):
 
542
        """Create a node.
 
543
 
 
544
        :param key_width: The width of keys for this node.
 
545
        """
 
546
        self._key = None
 
547
        # Current number of elements
 
548
        self._len = 0
 
549
        self._maximum_size = 0
 
550
        self._key_width = key_width
 
551
        # current size in bytes
 
552
        self._raw_size = 0
 
553
        # The pointers/values this node has - meaning defined by child classes.
 
554
        self._items = {}
 
555
        # The common search prefix
 
556
        self._search_prefix = None
 
557
 
 
558
    def __repr__(self):
 
559
        items_str = str(sorted(self._items))
 
560
        if len(items_str) > 20:
 
561
            items_str = items_str[:16] + '...]'
 
562
        return '%s(key:%s len:%s size:%s max:%s prefix:%s items:%s)' % (
 
563
            self.__class__.__name__, self._key, self._len, self._raw_size,
 
564
            self._maximum_size, self._search_prefix, items_str)
 
565
 
 
566
    def key(self):
 
567
        return self._key
 
568
 
 
569
    def __len__(self):
 
570
        return self._len
 
571
 
 
572
    @property
 
573
    def maximum_size(self):
 
574
        """What is the upper limit for adding references to a node."""
 
575
        return self._maximum_size
 
576
 
 
577
    def set_maximum_size(self, new_size):
 
578
        """Set the size threshold for nodes.
 
579
 
 
580
        :param new_size: The size at which no data is added to a node. 0 for
 
581
            unlimited.
 
582
        """
 
583
        self._maximum_size = new_size
 
584
 
 
585
    @classmethod
 
586
    def common_prefix(cls, prefix, key):
 
587
        """Given 2 strings, return the longest prefix common to both.
 
588
 
 
589
        :param prefix: This has been the common prefix for other keys, so it is
 
590
            more likely to be the common prefix in this case as well.
 
591
        :param key: Another string to compare to
 
592
        """
 
593
        if key.startswith(prefix):
 
594
            return prefix
 
595
        pos = -1
 
596
        # Is there a better way to do this?
 
597
        for pos, (left, right) in enumerate(zip(prefix, key)):
 
598
            if left != right:
 
599
                pos -= 1
 
600
                break
 
601
        common = prefix[:pos+1]
 
602
        return common
 
603
 
 
604
    @classmethod
 
605
    def common_prefix_for_keys(cls, keys):
 
606
        """Given a list of keys, find their common prefix.
 
607
 
 
608
        :param keys: An iterable of strings.
 
609
        :return: The longest common prefix of all keys.
 
610
        """
 
611
        common_prefix = None
 
612
        for key in keys:
 
613
            if common_prefix is None:
 
614
                common_prefix = key
 
615
                continue
 
616
            common_prefix = cls.common_prefix(common_prefix, key)
 
617
            if not common_prefix:
 
618
                # if common_prefix is the empty string, then we know it won't
 
619
                # change further
 
620
                return ''
 
621
        return common_prefix
 
622
 
 
623
 
 
624
# Singleton indicating we have not computed _search_prefix yet
 
625
_unknown = object()
 
626
 
 
627
class LeafNode(Node):
 
628
    """A node containing actual key:value pairs.
 
629
 
 
630
    :ivar _items: A dict of key->value items. The key is in tuple form.
 
631
    :ivar _size: The number of bytes that would be used by serializing all of
 
632
        the key/value pairs.
 
633
    """
 
634
 
 
635
    def __init__(self, search_key_func=None):
 
636
        Node.__init__(self)
 
637
        # All of the keys in this leaf node share this common prefix
 
638
        self._common_serialised_prefix = None
 
639
        self._serialise_key = '\x00'.join
 
640
        if search_key_func is None:
 
641
            self._search_key_func = _search_key_plain
 
642
        else:
 
643
            self._search_key_func = search_key_func
 
644
 
 
645
    def __repr__(self):
 
646
        items_str = str(sorted(self._items))
 
647
        if len(items_str) > 20:
 
648
            items_str = items_str[:16] + '...]'
 
649
        return \
 
650
            '%s(key:%s len:%s size:%s max:%s prefix:%s keywidth:%s items:%s)' \
 
651
            % (self.__class__.__name__, self._key, self._len, self._raw_size,
 
652
            self._maximum_size, self._search_prefix, self._key_width, items_str)
 
653
 
 
654
    def _current_size(self):
 
655
        """Answer the current serialised size of this node.
 
656
 
 
657
        This differs from self._raw_size in that it includes the bytes used for
 
658
        the header.
 
659
        """
 
660
        if self._common_serialised_prefix is None:
 
661
            bytes_for_items = 0
 
662
            prefix_len = 0
 
663
        else:
 
664
            # We will store a single string with the common prefix
 
665
            # And then that common prefix will not be stored in any of the
 
666
            # entry lines
 
667
            prefix_len = len(self._common_serialised_prefix)
 
668
            bytes_for_items = (self._raw_size - (prefix_len * self._len))
 
669
        return (9 # 'chkleaf:\n'
 
670
            + len(str(self._maximum_size)) + 1
 
671
            + len(str(self._key_width)) + 1
 
672
            + len(str(self._len)) + 1
 
673
            + prefix_len + 1
 
674
            + bytes_for_items)
 
675
 
 
676
    @classmethod
 
677
    def deserialise(klass, bytes, key, search_key_func=None):
 
678
        """Deserialise bytes, with key key, into a LeafNode.
 
679
 
 
680
        :param bytes: The bytes of the node.
 
681
        :param key: The key that the serialised node has.
 
682
        """
 
683
        return _deserialise_leaf_node(bytes, key,
 
684
                                      search_key_func=search_key_func)
 
685
 
 
686
    def iteritems(self, store, key_filter=None):
 
687
        """Iterate over items in the node.
 
688
 
 
689
        :param key_filter: A filter to apply to the node. It should be a
 
690
            list/set/dict or similar repeatedly iterable container.
 
691
        """
 
692
        if key_filter is not None:
 
693
            # Adjust the filter - short elements go to a prefix filter. All
 
694
            # other items are looked up directly.
 
695
            # XXX: perhaps defaultdict? Profiling<rinse and repeat>
 
696
            filters = {}
 
697
            for key in key_filter:
 
698
                if len(key) == self._key_width:
 
699
                    # This filter is meant to match exactly one key, yield it
 
700
                    # if we have it.
 
701
                    try:
 
702
                        yield key, self._items[key]
 
703
                    except KeyError:
 
704
                        # This key is not present in this map, continue
 
705
                        pass
 
706
                else:
 
707
                    # Short items, we need to match based on a prefix
 
708
                    length_filter = filters.setdefault(len(key), set())
 
709
                    length_filter.add(key)
 
710
            if filters:
 
711
                filters = filters.items()
 
712
                for item in self._items.iteritems():
 
713
                    for length, length_filter in filters:
 
714
                        if item[0][:length] in length_filter:
 
715
                            yield item
 
716
                            break
 
717
        else:
 
718
            for item in self._items.iteritems():
 
719
                yield item
 
720
 
 
721
    def _key_value_len(self, key, value):
 
722
        # TODO: Should probably be done without actually joining the key, but
 
723
        #       then that can be done via the C extension
 
724
        return (len(self._serialise_key(key)) + 1
 
725
                + len(str(value.count('\n'))) + 1
 
726
                + len(value) + 1)
 
727
 
 
728
    def _search_key(self, key):
 
729
        return self._search_key_func(key)
 
730
 
 
731
    def _map_no_split(self, key, value):
 
732
        """Map a key to a value.
 
733
 
 
734
        This assumes either the key does not already exist, or you have already
 
735
        removed its size and length from self.
 
736
 
 
737
        :return: True if adding this node should cause us to split.
 
738
        """
 
739
        self._items[key] = value
 
740
        self._raw_size += self._key_value_len(key, value)
 
741
        self._len += 1
 
742
        serialised_key = self._serialise_key(key)
 
743
        if self._common_serialised_prefix is None:
 
744
            self._common_serialised_prefix = serialised_key
 
745
        else:
 
746
            self._common_serialised_prefix = self.common_prefix(
 
747
                self._common_serialised_prefix, serialised_key)
 
748
        search_key = self._search_key(key)
 
749
        if self._search_prefix is _unknown:
 
750
            self._compute_search_prefix()
 
751
        if self._search_prefix is None:
 
752
            self._search_prefix = search_key
 
753
        else:
 
754
            self._search_prefix = self.common_prefix(
 
755
                self._search_prefix, search_key)
 
756
        if (self._len > 1
 
757
            and self._maximum_size
 
758
            and self._current_size() > self._maximum_size):
 
759
            # Check to see if all of the search_keys for this node are
 
760
            # identical. We allow the node to grow under that circumstance
 
761
            # (we could track this as common state, but it is infrequent)
 
762
            if (search_key != self._search_prefix
 
763
                or not self._are_search_keys_identical()):
 
764
                return True
 
765
        return False
 
766
 
 
767
    def _split(self, store):
 
768
        """We have overflowed.
 
769
 
 
770
        Split this node into multiple LeafNodes, return it up the stack so that
 
771
        the next layer creates a new InternalNode and references the new nodes.
 
772
 
 
773
        :return: (common_serialised_prefix, [(node_serialised_prefix, node)])
 
774
        """
 
775
        if self._search_prefix is _unknown:
 
776
            raise AssertionError('Search prefix must be known')
 
777
        common_prefix = self._search_prefix
 
778
        split_at = len(common_prefix) + 1
 
779
        result = {}
 
780
        for key, value in self._items.iteritems():
 
781
            search_key = self._search_key(key)
 
782
            prefix = search_key[:split_at]
 
783
            # TODO: Generally only 1 key can be exactly the right length,
 
784
            #       which means we can only have 1 key in the node pointed
 
785
            #       at by the 'prefix\0' key. We might want to consider
 
786
            #       folding it into the containing InternalNode rather than
 
787
            #       having a fixed length-1 node.
 
788
            #       Note this is probably not true for hash keys, as they
 
789
            #       may get a '\00' node anywhere, but won't have keys of
 
790
            #       different lengths.
 
791
            if len(prefix) < split_at:
 
792
                prefix += '\x00'*(split_at - len(prefix))
 
793
            if prefix not in result:
 
794
                node = LeafNode(search_key_func=self._search_key_func)
 
795
                node.set_maximum_size(self._maximum_size)
 
796
                node._key_width = self._key_width
 
797
                result[prefix] = node
 
798
            else:
 
799
                node = result[prefix]
 
800
            sub_prefix, node_details = node.map(store, key, value)
 
801
            if len(node_details) > 1:
 
802
                if prefix != sub_prefix:
 
803
                    # This node has been split and is now found via a different
 
804
                    # path
 
805
                    result.pop(prefix)
 
806
                new_node = InternalNode(sub_prefix,
 
807
                    search_key_func=self._search_key_func)
 
808
                new_node.set_maximum_size(self._maximum_size)
 
809
                new_node._key_width = self._key_width
 
810
                for split, node in node_details:
 
811
                    new_node.add_node(split, node)
 
812
                result[prefix] = new_node
 
813
        return common_prefix, result.items()
 
814
 
 
815
    def map(self, store, key, value):
 
816
        """Map key to value."""
 
817
        if key in self._items:
 
818
            self._raw_size -= self._key_value_len(key, self._items[key])
 
819
            self._len -= 1
 
820
        self._key = None
 
821
        if self._map_no_split(key, value):
 
822
            return self._split(store)
 
823
        else:
 
824
            if self._search_prefix is _unknown:
 
825
                raise AssertionError('%r must be known' % self._search_prefix)
 
826
            return self._search_prefix, [("", self)]
 
827
 
 
828
    def serialise(self, store):
 
829
        """Serialise the LeafNode to store.
 
830
 
 
831
        :param store: A VersionedFiles honouring the CHK extensions.
 
832
        :return: An iterable of the keys inserted by this operation.
 
833
        """
 
834
        lines = ["chkleaf:\n"]
 
835
        lines.append("%d\n" % self._maximum_size)
 
836
        lines.append("%d\n" % self._key_width)
 
837
        lines.append("%d\n" % self._len)
 
838
        if self._common_serialised_prefix is None:
 
839
            lines.append('\n')
 
840
            if len(self._items) != 0:
 
841
                raise AssertionError('If _common_serialised_prefix is None'
 
842
                    ' we should have no items')
 
843
        else:
 
844
            lines.append('%s\n' % (self._common_serialised_prefix,))
 
845
            prefix_len = len(self._common_serialised_prefix)
 
846
        for key, value in sorted(self._items.items()):
 
847
            # Always add a final newline
 
848
            value_lines = osutils.chunks_to_lines([value + '\n'])
 
849
            serialized = "%s\x00%s\n" % (self._serialise_key(key),
 
850
                                         len(value_lines))
 
851
            if not serialized.startswith(self._common_serialised_prefix):
 
852
                raise AssertionError('We thought the common prefix was %r'
 
853
                    ' but entry %r does not have it in common'
 
854
                    % (self._common_serialised_prefix, serialized))
 
855
            lines.append(serialized[prefix_len:])
 
856
            lines.extend(value_lines)
 
857
        sha1, _, _ = store.add_lines((None,), (), lines)
 
858
        self._key = ("sha1:" + sha1,)
 
859
        bytes = ''.join(lines)
 
860
        if len(bytes) != self._current_size():
 
861
            raise AssertionError('Invalid _current_size')
 
862
        _page_cache.add(self._key, bytes)
 
863
        return [self._key]
 
864
 
 
865
    def refs(self):
 
866
        """Return the references to other CHK's held by this node."""
 
867
        return []
 
868
 
 
869
    def _compute_search_prefix(self):
 
870
        """Determine the common search prefix for all keys in this node.
 
871
 
 
872
        :return: A bytestring of the longest search key prefix that is
 
873
            unique within this node.
 
874
        """
 
875
        search_keys = [self._search_key_func(key) for key in self._items]
 
876
        self._search_prefix = self.common_prefix_for_keys(search_keys)
 
877
        return self._search_prefix
 
878
 
 
879
    def _are_search_keys_identical(self):
 
880
        """Check to see if the search keys for all entries are the same.
 
881
 
 
882
        When using a hash as the search_key it is possible for non-identical
 
883
        keys to collide. If that happens enough, we may try overflow a
 
884
        LeafNode, but as all are collisions, we must not split.
 
885
        """
 
886
        common_search_key = None
 
887
        for key in self._items:
 
888
            search_key = self._search_key(key)
 
889
            if common_search_key is None:
 
890
                common_search_key = search_key
 
891
            elif search_key != common_search_key:
 
892
                return False
 
893
        return True
 
894
 
 
895
    def _compute_serialised_prefix(self):
 
896
        """Determine the common prefix for serialised keys in this node.
 
897
 
 
898
        :return: A bytestring of the longest serialised key prefix that is
 
899
            unique within this node.
 
900
        """
 
901
        serialised_keys = [self._serialise_key(key) for key in self._items]
 
902
        self._common_serialised_prefix = self.common_prefix_for_keys(
 
903
            serialised_keys)
 
904
        return self._common_serialised_prefix
 
905
 
 
906
    def unmap(self, store, key):
 
907
        """Unmap key from the node."""
 
908
        try:
 
909
            self._raw_size -= self._key_value_len(key, self._items[key])
 
910
        except KeyError:
 
911
            trace.mutter("key %s not found in %r", key, self._items)
 
912
            raise
 
913
        self._len -= 1
 
914
        del self._items[key]
 
915
        self._key = None
 
916
        # Recompute from scratch
 
917
        self._compute_search_prefix()
 
918
        self._compute_serialised_prefix()
 
919
        return self
 
920
 
 
921
 
 
922
class InternalNode(Node):
 
923
    """A node that contains references to other nodes.
 
924
 
 
925
    An InternalNode is responsible for mapping search key prefixes to child
 
926
    nodes.
 
927
 
 
928
    :ivar _items: serialised_key => node dictionary. node may be a tuple,
 
929
        LeafNode or InternalNode.
 
930
    """
 
931
 
 
932
    def __init__(self, prefix='', search_key_func=None):
 
933
        Node.__init__(self)
 
934
        # The size of an internalnode with default values and no children.
 
935
        # How many octets key prefixes within this node are.
 
936
        self._node_width = 0
 
937
        self._search_prefix = prefix
 
938
        if search_key_func is None:
 
939
            self._search_key_func = _search_key_plain
 
940
        else:
 
941
            self._search_key_func = search_key_func
 
942
 
 
943
    def add_node(self, prefix, node):
 
944
        """Add a child node with prefix prefix, and node node.
 
945
 
 
946
        :param prefix: The search key prefix for node.
 
947
        :param node: The node being added.
 
948
        """
 
949
        if self._search_prefix is None:
 
950
            raise AssertionError("_search_prefix should not be None")
 
951
        if not prefix.startswith(self._search_prefix):
 
952
            raise AssertionError("prefixes mismatch: %s must start with %s"
 
953
                % (prefix,self._search_prefix))
 
954
        if len(prefix) != len(self._search_prefix) + 1:
 
955
            raise AssertionError("prefix wrong length: len(%s) is not %d" %
 
956
                (prefix, len(self._search_prefix) + 1))
 
957
        self._len += len(node)
 
958
        if not len(self._items):
 
959
            self._node_width = len(prefix)
 
960
        if self._node_width != len(self._search_prefix) + 1:
 
961
            raise AssertionError("node width mismatch: %d is not %d" %
 
962
                (self._node_width, len(self._search_prefix) + 1))
 
963
        self._items[prefix] = node
 
964
        self._key = None
 
965
 
 
966
    def _current_size(self):
 
967
        """Answer the current serialised size of this node."""
 
968
        return (self._raw_size + len(str(self._len)) + len(str(self._key_width)) +
 
969
            len(str(self._maximum_size)))
 
970
 
 
971
    @classmethod
 
972
    def deserialise(klass, bytes, key, search_key_func=None):
 
973
        """Deserialise bytes to an InternalNode, with key key.
 
974
 
 
975
        :param bytes: The bytes of the node.
 
976
        :param key: The key that the serialised node has.
 
977
        :return: An InternalNode instance.
 
978
        """
 
979
        return _deserialise_internal_node(bytes, key,
 
980
                                          search_key_func=search_key_func)
 
981
 
 
982
    def iteritems(self, store, key_filter=None):
 
983
        for node, node_filter in self._iter_nodes(store, key_filter=key_filter):
 
984
            for item in node.iteritems(store, key_filter=node_filter):
 
985
                yield item
 
986
 
 
987
    def _iter_nodes(self, store, key_filter=None, batch_size=None):
 
988
        """Iterate over node objects which match key_filter.
 
989
 
 
990
        :param store: A store to use for accessing content.
 
991
        :param key_filter: A key filter to filter nodes. Only nodes that might
 
992
            contain a key in key_filter will be returned.
 
993
        :param batch_size: If not None, then we will return the nodes that had
 
994
            to be read using get_record_stream in batches, rather than reading
 
995
            them all at once.
 
996
        :return: An iterable of nodes. This function does not have to be fully
 
997
            consumed.  (There will be no pending I/O when items are being returned.)
 
998
        """
 
999
        # Map from chk key ('sha1:...',) to (prefix, key_filter)
 
1000
        # prefix is the key in self._items to use, key_filter is the key_filter
 
1001
        # entries that would match this node
 
1002
        keys = {}
 
1003
        shortcut = False
 
1004
        if key_filter is None:
 
1005
            # yielding all nodes, yield whatever we have, and queue up a read
 
1006
            # for whatever we are missing
 
1007
            shortcut = True
 
1008
            for prefix, node in self._items.iteritems():
 
1009
                if node.__class__ is tuple:
 
1010
                    keys[node] = (prefix, None)
 
1011
                else:
 
1012
                    yield node, None
 
1013
        elif len(key_filter) == 1:
 
1014
            # Technically, this path could also be handled by the first check
 
1015
            # in 'self._node_width' in length_filters. However, we can handle
 
1016
            # this case without spending any time building up the
 
1017
            # prefix_to_keys, etc state.
 
1018
 
 
1019
            # This is a bit ugly, but TIMEIT showed it to be by far the fastest
 
1020
            # 0.626us   list(key_filter)[0]
 
1021
            #       is a func() for list(), 2 mallocs, and a getitem
 
1022
            # 0.489us   [k for k in key_filter][0]
 
1023
            #       still has the mallocs, avoids the func() call
 
1024
            # 0.350us   iter(key_filter).next()
 
1025
            #       has a func() call, and mallocs an iterator
 
1026
            # 0.125us   for key in key_filter: pass
 
1027
            #       no func() overhead, might malloc an iterator
 
1028
            # 0.105us   for key in key_filter: break
 
1029
            #       no func() overhead, might malloc an iterator, probably
 
1030
            #       avoids checking an 'else' clause as part of the for
 
1031
            for key in key_filter:
 
1032
                break
 
1033
            search_prefix = self._search_prefix_filter(key)
 
1034
            if len(search_prefix) == self._node_width:
 
1035
                # This item will match exactly, so just do a dict lookup, and
 
1036
                # see what we can return
 
1037
                shortcut = True
 
1038
                try:
 
1039
                    node = self._items[search_prefix]
 
1040
                except KeyError:
 
1041
                    # A given key can only match 1 child node, if it isn't
 
1042
                    # there, then we can just return nothing
 
1043
                    return
 
1044
                if node.__class__ is tuple:
 
1045
                    keys[node] = (search_prefix, [key])
 
1046
                else:
 
1047
                    # This is loaded, and the only thing that can match,
 
1048
                    # return
 
1049
                    yield node, [key]
 
1050
                    return
 
1051
        if not shortcut:
 
1052
            # First, convert all keys into a list of search prefixes
 
1053
            # Aggregate common prefixes, and track the keys they come from
 
1054
            prefix_to_keys = {}
 
1055
            length_filters = {}
 
1056
            for key in key_filter:
 
1057
                search_prefix = self._search_prefix_filter(key)
 
1058
                length_filter = length_filters.setdefault(
 
1059
                                    len(search_prefix), set())
 
1060
                length_filter.add(search_prefix)
 
1061
                prefix_to_keys.setdefault(search_prefix, []).append(key)
 
1062
 
 
1063
            if (self._node_width in length_filters
 
1064
                and len(length_filters) == 1):
 
1065
                # all of the search prefixes match exactly _node_width. This
 
1066
                # means that everything is an exact match, and we can do a
 
1067
                # lookup into self._items, rather than iterating over the items
 
1068
                # dict.
 
1069
                search_prefixes = length_filters[self._node_width]
 
1070
                for search_prefix in search_prefixes:
 
1071
                    try:
 
1072
                        node = self._items[search_prefix]
 
1073
                    except KeyError:
 
1074
                        # We can ignore this one
 
1075
                        continue
 
1076
                    node_key_filter = prefix_to_keys[search_prefix]
 
1077
                    if node.__class__ is tuple:
 
1078
                        keys[node] = (search_prefix, node_key_filter)
 
1079
                    else:
 
1080
                        yield node, node_key_filter
 
1081
            else:
 
1082
                # The slow way. We walk every item in self._items, and check to
 
1083
                # see if there are any matches
 
1084
                length_filters = length_filters.items()
 
1085
                for prefix, node in self._items.iteritems():
 
1086
                    node_key_filter = []
 
1087
                    for length, length_filter in length_filters:
 
1088
                        sub_prefix = prefix[:length]
 
1089
                        if sub_prefix in length_filter:
 
1090
                            node_key_filter.extend(prefix_to_keys[sub_prefix])
 
1091
                    if node_key_filter: # this key matched something, yield it
 
1092
                        if node.__class__ is tuple:
 
1093
                            keys[node] = (prefix, node_key_filter)
 
1094
                        else:
 
1095
                            yield node, node_key_filter
 
1096
        if keys:
 
1097
            # Look in the page cache for some more bytes
 
1098
            found_keys = set()
 
1099
            for key in keys:
 
1100
                try:
 
1101
                    bytes = _page_cache[key]
 
1102
                except KeyError:
 
1103
                    continue
 
1104
                else:
 
1105
                    node = _deserialise(bytes, key,
 
1106
                        search_key_func=self._search_key_func)
 
1107
                    prefix, node_key_filter = keys[key]
 
1108
                    self._items[prefix] = node
 
1109
                    found_keys.add(key)
 
1110
                    yield node, node_key_filter
 
1111
            for key in found_keys:
 
1112
                del keys[key]
 
1113
        if keys:
 
1114
            # demand load some pages.
 
1115
            if batch_size is None:
 
1116
                # Read all the keys in
 
1117
                batch_size = len(keys)
 
1118
            key_order = list(keys)
 
1119
            for batch_start in range(0, len(key_order), batch_size):
 
1120
                batch = key_order[batch_start:batch_start + batch_size]
 
1121
                # We have to fully consume the stream so there is no pending
 
1122
                # I/O, so we buffer the nodes for now.
 
1123
                stream = store.get_record_stream(batch, 'unordered', True)
 
1124
                node_and_filters = []
 
1125
                for record in stream:
 
1126
                    bytes = record.get_bytes_as('fulltext')
 
1127
                    node = _deserialise(bytes, record.key,
 
1128
                        search_key_func=self._search_key_func)
 
1129
                    prefix, node_key_filter = keys[record.key]
 
1130
                    node_and_filters.append((node, node_key_filter))
 
1131
                    self._items[prefix] = node
 
1132
                    _page_cache.add(record.key, bytes)
 
1133
                for info in node_and_filters:
 
1134
                    yield info
 
1135
 
 
1136
    def map(self, store, key, value):
 
1137
        """Map key to value."""
 
1138
        if not len(self._items):
 
1139
            raise AssertionError("can't map in an empty InternalNode.")
 
1140
        search_key = self._search_key(key)
 
1141
        if self._node_width != len(self._search_prefix) + 1:
 
1142
            raise AssertionError("node width mismatch: %d is not %d" %
 
1143
                (self._node_width, len(self._search_prefix) + 1))
 
1144
        if not search_key.startswith(self._search_prefix):
 
1145
            # This key doesn't fit in this index, so we need to split at the
 
1146
            # point where it would fit, insert self into that internal node,
 
1147
            # and then map this key into that node.
 
1148
            new_prefix = self.common_prefix(self._search_prefix,
 
1149
                                            search_key)
 
1150
            new_parent = InternalNode(new_prefix,
 
1151
                search_key_func=self._search_key_func)
 
1152
            new_parent.set_maximum_size(self._maximum_size)
 
1153
            new_parent._key_width = self._key_width
 
1154
            new_parent.add_node(self._search_prefix[:len(new_prefix)+1],
 
1155
                                self)
 
1156
            return new_parent.map(store, key, value)
 
1157
        children = [node for node, _
 
1158
                          in self._iter_nodes(store, key_filter=[key])]
 
1159
        if children:
 
1160
            child = children[0]
 
1161
        else:
 
1162
            # new child needed:
 
1163
            child = self._new_child(search_key, LeafNode)
 
1164
        old_len = len(child)
 
1165
        if type(child) is LeafNode:
 
1166
            old_size = child._current_size()
 
1167
        else:
 
1168
            old_size = None
 
1169
        prefix, node_details = child.map(store, key, value)
 
1170
        if len(node_details) == 1:
 
1171
            # child may have shrunk, or might be a new node
 
1172
            child = node_details[0][1]
 
1173
            self._len = self._len - old_len + len(child)
 
1174
            self._items[search_key] = child
 
1175
            self._key = None
 
1176
            new_node = self
 
1177
            if type(child) is LeafNode:
 
1178
                if old_size is None:
 
1179
                    # The old node was an InternalNode which means it has now
 
1180
                    # collapsed, so we need to check if it will chain to a
 
1181
                    # collapse at this level.
 
1182
                    trace.mutter("checking remap as InternalNode -> LeafNode")
 
1183
                    new_node = self._check_remap(store)
 
1184
                else:
 
1185
                    # If the LeafNode has shrunk in size, we may want to run
 
1186
                    # a remap check. Checking for a remap is expensive though
 
1187
                    # and the frequency of a successful remap is very low.
 
1188
                    # Shrinkage by small amounts is common, so we only do the
 
1189
                    # remap check if the new_size is low or the shrinkage
 
1190
                    # amount is over a configurable limit.
 
1191
                    new_size = child._current_size()
 
1192
                    shrinkage = old_size - new_size
 
1193
                    if (shrinkage > 0 and new_size < _INTERESTING_NEW_SIZE
 
1194
                        or shrinkage > _INTERESTING_SHRINKAGE_LIMIT):
 
1195
                        trace.mutter(
 
1196
                            "checking remap as size shrunk by %d to be %d",
 
1197
                            shrinkage, new_size)
 
1198
                        new_node = self._check_remap(store)
 
1199
            if new_node._search_prefix is None:
 
1200
                raise AssertionError("_search_prefix should not be None")
 
1201
            return new_node._search_prefix, [('', new_node)]
 
1202
        # child has overflown - create a new intermediate node.
 
1203
        # XXX: This is where we might want to try and expand our depth
 
1204
        # to refer to more bytes of every child (which would give us
 
1205
        # multiple pointers to child nodes, but less intermediate nodes)
 
1206
        child = self._new_child(search_key, InternalNode)
 
1207
        child._search_prefix = prefix
 
1208
        for split, node in node_details:
 
1209
            child.add_node(split, node)
 
1210
        self._len = self._len - old_len + len(child)
 
1211
        self._key = None
 
1212
        return self._search_prefix, [("", self)]
 
1213
 
 
1214
    def _new_child(self, search_key, klass):
 
1215
        """Create a new child node of type klass."""
 
1216
        child = klass()
 
1217
        child.set_maximum_size(self._maximum_size)
 
1218
        child._key_width = self._key_width
 
1219
        child._search_key_func = self._search_key_func
 
1220
        self._items[search_key] = child
 
1221
        return child
 
1222
 
 
1223
    def serialise(self, store):
 
1224
        """Serialise the node to store.
 
1225
 
 
1226
        :param store: A VersionedFiles honouring the CHK extensions.
 
1227
        :return: An iterable of the keys inserted by this operation.
 
1228
        """
 
1229
        for node in self._items.itervalues():
 
1230
            if type(node) is tuple:
 
1231
                # Never deserialised.
 
1232
                continue
 
1233
            if node._key is not None:
 
1234
                # Never altered
 
1235
                continue
 
1236
            for key in node.serialise(store):
 
1237
                yield key
 
1238
        lines = ["chknode:\n"]
 
1239
        lines.append("%d\n" % self._maximum_size)
 
1240
        lines.append("%d\n" % self._key_width)
 
1241
        lines.append("%d\n" % self._len)
 
1242
        if self._search_prefix is None:
 
1243
            raise AssertionError("_search_prefix should not be None")
 
1244
        lines.append('%s\n' % (self._search_prefix,))
 
1245
        prefix_len = len(self._search_prefix)
 
1246
        for prefix, node in sorted(self._items.items()):
 
1247
            if type(node) is tuple:
 
1248
                key = node[0]
 
1249
            else:
 
1250
                key = node._key[0]
 
1251
            serialised = "%s\x00%s\n" % (prefix, key)
 
1252
            if not serialised.startswith(self._search_prefix):
 
1253
                raise AssertionError("prefixes mismatch: %s must start with %s"
 
1254
                    % (serialised, self._search_prefix))
 
1255
            lines.append(serialised[prefix_len:])
 
1256
        sha1, _, _ = store.add_lines((None,), (), lines)
 
1257
        self._key = ("sha1:" + sha1,)
 
1258
        _page_cache.add(self._key, ''.join(lines))
 
1259
        yield self._key
 
1260
 
 
1261
    def _search_key(self, key):
 
1262
        """Return the serialised key for key in this node."""
 
1263
        # search keys are fixed width. All will be self._node_width wide, so we
 
1264
        # pad as necessary.
 
1265
        return (self._search_key_func(key) + '\x00'*self._node_width)[:self._node_width]
 
1266
 
 
1267
    def _search_prefix_filter(self, key):
 
1268
        """Serialise key for use as a prefix filter in iteritems."""
 
1269
        return self._search_key_func(key)[:self._node_width]
 
1270
 
 
1271
    def _split(self, offset):
 
1272
        """Split this node into smaller nodes starting at offset.
 
1273
 
 
1274
        :param offset: The offset to start the new child nodes at.
 
1275
        :return: An iterable of (prefix, node) tuples. prefix is a byte
 
1276
            prefix for reaching node.
 
1277
        """
 
1278
        if offset >= self._node_width:
 
1279
            for node in self._items.values():
 
1280
                for result in node._split(offset):
 
1281
                    yield result
 
1282
            return
 
1283
        for key, node in self._items.items():
 
1284
            pass
 
1285
 
 
1286
    def refs(self):
 
1287
        """Return the references to other CHK's held by this node."""
 
1288
        if self._key is None:
 
1289
            raise AssertionError("unserialised nodes have no refs.")
 
1290
        refs = []
 
1291
        for value in self._items.itervalues():
 
1292
            if type(value) is tuple:
 
1293
                refs.append(value)
 
1294
            else:
 
1295
                refs.append(value.key())
 
1296
        return refs
 
1297
 
 
1298
    def _compute_search_prefix(self, extra_key=None):
 
1299
        """Return the unique key prefix for this node.
 
1300
 
 
1301
        :return: A bytestring of the longest search key prefix that is
 
1302
            unique within this node.
 
1303
        """
 
1304
        self._search_prefix = self.common_prefix_for_keys(self._items)
 
1305
        return self._search_prefix
 
1306
 
 
1307
    def unmap(self, store, key, check_remap=True):
 
1308
        """Remove key from this node and it's children."""
 
1309
        if not len(self._items):
 
1310
            raise AssertionError("can't unmap in an empty InternalNode.")
 
1311
        children = [node for node, _
 
1312
                          in self._iter_nodes(store, key_filter=[key])]
 
1313
        if children:
 
1314
            child = children[0]
 
1315
        else:
 
1316
            raise KeyError(key)
 
1317
        self._len -= 1
 
1318
        unmapped = child.unmap(store, key)
 
1319
        self._key = None
 
1320
        search_key = self._search_key(key)
 
1321
        if len(unmapped) == 0:
 
1322
            # All child nodes are gone, remove the child:
 
1323
            del self._items[search_key]
 
1324
            unmapped = None
 
1325
        else:
 
1326
            # Stash the returned node
 
1327
            self._items[search_key] = unmapped
 
1328
        if len(self._items) == 1:
 
1329
            # this node is no longer needed:
 
1330
            return self._items.values()[0]
 
1331
        if type(unmapped) is InternalNode:
 
1332
            return self
 
1333
        if check_remap:
 
1334
            return self._check_remap(store)
 
1335
        else:
 
1336
            return self
 
1337
 
 
1338
    def _check_remap(self, store):
 
1339
        """Check if all keys contained by children fit in a single LeafNode.
 
1340
 
 
1341
        :param store: A store to use for reading more nodes
 
1342
        :return: Either self, or a new LeafNode which should replace self.
 
1343
        """
 
1344
        # Logic for how we determine when we need to rebuild
 
1345
        # 1) Implicitly unmap() is removing a key which means that the child
 
1346
        #    nodes are going to be shrinking by some extent.
 
1347
        # 2) If all children are LeafNodes, it is possible that they could be
 
1348
        #    combined into a single LeafNode, which can then completely replace
 
1349
        #    this internal node with a single LeafNode
 
1350
        # 3) If *one* child is an InternalNode, we assume it has already done
 
1351
        #    all the work to determine that its children cannot collapse, and
 
1352
        #    we can then assume that those nodes *plus* the current nodes don't
 
1353
        #    have a chance of collapsing either.
 
1354
        #    So a very cheap check is to just say if 'unmapped' is an
 
1355
        #    InternalNode, we don't have to check further.
 
1356
 
 
1357
        # TODO: Another alternative is to check the total size of all known
 
1358
        #       LeafNodes. If there is some formula we can use to determine the
 
1359
        #       final size without actually having to read in any more
 
1360
        #       children, it would be nice to have. However, we have to be
 
1361
        #       careful with stuff like nodes that pull out the common prefix
 
1362
        #       of each key, as adding a new key can change the common prefix
 
1363
        #       and cause size changes greater than the length of one key.
 
1364
        #       So for now, we just add everything to a new Leaf until it
 
1365
        #       splits, as we know that will give the right answer
 
1366
        new_leaf = LeafNode(search_key_func=self._search_key_func)
 
1367
        new_leaf.set_maximum_size(self._maximum_size)
 
1368
        new_leaf._key_width = self._key_width
 
1369
        # A batch_size of 16 was chosen because:
 
1370
        #   a) In testing, a 4k page held 14 times. So if we have more than 16
 
1371
        #      leaf nodes we are unlikely to hold them in a single new leaf
 
1372
        #      node. This still allows for 1 round trip
 
1373
        #   b) With 16-way fan out, we can still do a single round trip
 
1374
        #   c) With 255-way fan out, we don't want to read all 255 and destroy
 
1375
        #      the page cache, just to determine that we really don't need it.
 
1376
        for node, _ in self._iter_nodes(store, batch_size=16):
 
1377
            if type(node) is InternalNode:
 
1378
                # Without looking at any leaf nodes, we are sure
 
1379
                return self
 
1380
            for key, value in node._items.iteritems():
 
1381
                if new_leaf._map_no_split(key, value):
 
1382
                    return self
 
1383
        trace.mutter("remap generated a new LeafNode")
 
1384
        return new_leaf
 
1385
 
 
1386
 
 
1387
def _deserialise(bytes, key, search_key_func):
 
1388
    """Helper for repositorydetails - convert bytes to a node."""
 
1389
    if bytes.startswith("chkleaf:\n"):
 
1390
        node = LeafNode.deserialise(bytes, key, search_key_func=search_key_func)
 
1391
    elif bytes.startswith("chknode:\n"):
 
1392
        node = InternalNode.deserialise(bytes, key,
 
1393
            search_key_func=search_key_func)
 
1394
    else:
 
1395
        raise AssertionError("Unknown node type.")
 
1396
    return node
 
1397
 
 
1398
 
 
1399
class CHKMapDifference(object):
 
1400
    """Iterate the stored pages and key,value pairs for (new - old).
 
1401
 
 
1402
    This class provides a generator over the stored CHK pages and the
 
1403
    (key, value) pairs that are in any of the new maps and not in any of the
 
1404
    old maps.
 
1405
 
 
1406
    Note that it may yield chk pages that are common (especially root nodes),
 
1407
    but it won't yield (key,value) pairs that are common.
 
1408
    """
 
1409
 
 
1410
    def __init__(self, store, new_root_keys, old_root_keys,
 
1411
                 search_key_func, pb=None):
 
1412
        self._store = store
 
1413
        self._new_root_keys = new_root_keys
 
1414
        self._old_root_keys = old_root_keys
 
1415
        self._pb = pb
 
1416
        # All uninteresting chks that we have seen. By the time they are added
 
1417
        # here, they should be either fully ignored, or queued up for
 
1418
        # processing
 
1419
        self._all_old_chks = set(self._old_root_keys)
 
1420
        # All items that we have seen from the old_root_keys
 
1421
        self._all_old_items = set()
 
1422
        # These are interesting items which were either read, or already in the
 
1423
        # interesting queue (so we don't need to walk them again)
 
1424
        self._processed_new_refs = set()
 
1425
        self._search_key_func = search_key_func
 
1426
 
 
1427
        # The uninteresting and interesting nodes to be searched
 
1428
        self._old_queue = []
 
1429
        self._new_queue = []
 
1430
        # Holds the (key, value) items found when processing the root nodes,
 
1431
        # waiting for the uninteresting nodes to be walked
 
1432
        self._new_item_queue = []
 
1433
        self._state = None
 
1434
 
 
1435
    def _read_nodes_from_store(self, keys):
 
1436
        # We chose not to use _page_cache, because we think in terms of records
 
1437
        # to be yielded. Also, we expect to touch each page only 1 time during
 
1438
        # this code. (We may want to evaluate saving the raw bytes into the
 
1439
        # page cache, which would allow a working tree update after the fetch
 
1440
        # to not have to read the bytes again.)
 
1441
        stream = self._store.get_record_stream(keys, 'unordered', True)
 
1442
        for record in stream:
 
1443
            if self._pb is not None:
 
1444
                self._pb.tick()
 
1445
            if record.storage_kind == 'absent':
 
1446
                raise errors.NoSuchRevision(self._store, record.key)
 
1447
            bytes = record.get_bytes_as('fulltext')
 
1448
            node = _deserialise(bytes, record.key,
 
1449
                                search_key_func=self._search_key_func)
 
1450
            if type(node) is InternalNode:
 
1451
                # Note we don't have to do node.refs() because we know that
 
1452
                # there are no children that have been pushed into this node
 
1453
                prefix_refs = node._items.items()
 
1454
                items = []
 
1455
            else:
 
1456
                prefix_refs = []
 
1457
                items = node._items.items()
 
1458
            yield record, node, prefix_refs, items
 
1459
 
 
1460
    def _read_old_roots(self):
 
1461
        old_chks_to_enqueue = []
 
1462
        all_old_chks = self._all_old_chks
 
1463
        for record, node, prefix_refs, items in \
 
1464
                self._read_nodes_from_store(self._old_root_keys):
 
1465
            # Uninteresting node
 
1466
            prefix_refs = [p_r for p_r in prefix_refs
 
1467
                                if p_r[1] not in all_old_chks]
 
1468
            new_refs = [p_r[1] for p_r in prefix_refs]
 
1469
            all_old_chks.update(new_refs)
 
1470
            self._all_old_items.update(items)
 
1471
            # Queue up the uninteresting references
 
1472
            # Don't actually put them in the 'to-read' queue until we have
 
1473
            # finished checking the interesting references
 
1474
            old_chks_to_enqueue.extend(prefix_refs)
 
1475
        return old_chks_to_enqueue
 
1476
 
 
1477
    def _enqueue_old(self, new_prefixes, old_chks_to_enqueue):
 
1478
        # At this point, we have read all the uninteresting and interesting
 
1479
        # items, so we can queue up the uninteresting stuff, knowing that we've
 
1480
        # handled the interesting ones
 
1481
        for prefix, ref in old_chks_to_enqueue:
 
1482
            not_interesting = True
 
1483
            for i in xrange(len(prefix), 0, -1):
 
1484
                if prefix[:i] in new_prefixes:
 
1485
                    not_interesting = False
 
1486
                    break
 
1487
            if not_interesting:
 
1488
                # This prefix is not part of the remaining 'interesting set'
 
1489
                continue
 
1490
            self._old_queue.append(ref)
 
1491
 
 
1492
    def _read_all_roots(self):
 
1493
        """Read the root pages.
 
1494
 
 
1495
        This is structured as a generator, so that the root records can be
 
1496
        yielded up to whoever needs them without any buffering.
 
1497
        """
 
1498
        # This is the bootstrap phase
 
1499
        if not self._old_root_keys:
 
1500
            # With no old_root_keys we can just shortcut and be ready
 
1501
            # for _flush_new_queue
 
1502
            self._new_queue = list(self._new_root_keys)
 
1503
            return
 
1504
        old_chks_to_enqueue = self._read_old_roots()
 
1505
        # filter out any root keys that are already known to be uninteresting
 
1506
        new_keys = set(self._new_root_keys).difference(self._all_old_chks)
 
1507
        # These are prefixes that are present in new_keys that we are
 
1508
        # thinking to yield
 
1509
        new_prefixes = set()
 
1510
        # We are about to yield all of these, so we don't want them getting
 
1511
        # added a second time
 
1512
        processed_new_refs = self._processed_new_refs
 
1513
        processed_new_refs.update(new_keys)
 
1514
        for record, node, prefix_refs, items in \
 
1515
                self._read_nodes_from_store(new_keys):
 
1516
            # At this level, we now know all the uninteresting references
 
1517
            # So we filter and queue up whatever is remaining
 
1518
            prefix_refs = [p_r for p_r in prefix_refs
 
1519
                           if p_r[1] not in self._all_old_chks
 
1520
                              and p_r[1] not in processed_new_refs]
 
1521
            refs = [p_r[1] for p_r in prefix_refs]
 
1522
            new_prefixes.update([p_r[0] for p_r in prefix_refs])
 
1523
            self._new_queue.extend(refs)
 
1524
            # TODO: We can potentially get multiple items here, however the
 
1525
            #       current design allows for this, as callers will do the work
 
1526
            #       to make the results unique. We might profile whether we
 
1527
            #       gain anything by ensuring unique return values for items
 
1528
            new_items = [item for item in items
 
1529
                               if item not in self._all_old_items]
 
1530
            self._new_item_queue.extend(new_items)
 
1531
            new_prefixes.update([self._search_key_func(item[0])
 
1532
                                 for item in new_items])
 
1533
            processed_new_refs.update(refs)
 
1534
            yield record
 
1535
        # For new_prefixes we have the full length prefixes queued up.
 
1536
        # However, we also need possible prefixes. (If we have a known ref to
 
1537
        # 'ab', then we also need to include 'a'.) So expand the
 
1538
        # new_prefixes to include all shorter prefixes
 
1539
        for prefix in list(new_prefixes):
 
1540
            new_prefixes.update([prefix[:i] for i in xrange(1, len(prefix))])
 
1541
        self._enqueue_old(new_prefixes, old_chks_to_enqueue)
 
1542
 
 
1543
    def _flush_new_queue(self):
 
1544
        # No need to maintain the heap invariant anymore, just pull things out
 
1545
        # and process them
 
1546
        refs = set(self._new_queue)
 
1547
        self._new_queue = []
 
1548
        # First pass, flush all interesting items and convert to using direct refs
 
1549
        all_old_chks = self._all_old_chks
 
1550
        processed_new_refs = self._processed_new_refs
 
1551
        all_old_items = self._all_old_items
 
1552
        new_items = [item for item in self._new_item_queue
 
1553
                           if item not in all_old_items]
 
1554
        self._new_item_queue = []
 
1555
        if new_items:
 
1556
            yield None, new_items
 
1557
        refs = refs.difference(all_old_chks)
 
1558
        while refs:
 
1559
            next_refs = set()
 
1560
            next_refs_update = next_refs.update
 
1561
            # Inlining _read_nodes_from_store improves 'bzr branch bzr.dev'
 
1562
            # from 1m54s to 1m51s. Consider it.
 
1563
            for record, _, p_refs, items in self._read_nodes_from_store(refs):
 
1564
                items = [item for item in items
 
1565
                         if item not in all_old_items]
 
1566
                yield record, items
 
1567
                next_refs_update([p_r[1] for p_r in p_refs])
 
1568
            next_refs = next_refs.difference(all_old_chks)
 
1569
            next_refs = next_refs.difference(processed_new_refs)
 
1570
            processed_new_refs.update(next_refs)
 
1571
            refs = next_refs
 
1572
 
 
1573
    def _process_next_old(self):
 
1574
        # Since we don't filter uninteresting any further than during
 
1575
        # _read_all_roots, process the whole queue in a single pass.
 
1576
        refs = self._old_queue
 
1577
        self._old_queue = []
 
1578
        all_old_chks = self._all_old_chks
 
1579
        for record, _, prefix_refs, items in self._read_nodes_from_store(refs):
 
1580
            self._all_old_items.update(items)
 
1581
            refs = [r for _,r in prefix_refs if r not in all_old_chks]
 
1582
            self._old_queue.extend(refs)
 
1583
            all_old_chks.update(refs)
 
1584
 
 
1585
    def _process_queues(self):
 
1586
        while self._old_queue:
 
1587
            self._process_next_old()
 
1588
        return self._flush_new_queue()
 
1589
 
 
1590
    def process(self):
 
1591
        for record in self._read_all_roots():
 
1592
            yield record, []
 
1593
        for record, items in self._process_queues():
 
1594
            yield record, items
 
1595
 
 
1596
 
 
1597
def iter_interesting_nodes(store, interesting_root_keys,
 
1598
                           uninteresting_root_keys, pb=None):
 
1599
    """Given root keys, find interesting nodes.
 
1600
 
 
1601
    Evaluate nodes referenced by interesting_root_keys. Ones that are also
 
1602
    referenced from uninteresting_root_keys are not considered interesting.
 
1603
 
 
1604
    :param interesting_root_keys: keys which should be part of the
 
1605
        "interesting" nodes (which will be yielded)
 
1606
    :param uninteresting_root_keys: keys which should be filtered out of the
 
1607
        result set.
 
1608
    :return: Yield
 
1609
        (interesting record, {interesting key:values})
 
1610
    """
 
1611
    iterator = CHKMapDifference(store, interesting_root_keys,
 
1612
                                uninteresting_root_keys,
 
1613
                                search_key_func=store._search_key_func,
 
1614
                                pb=pb)
 
1615
    return iterator.process()
 
1616
 
 
1617
 
 
1618
try:
 
1619
    from bzrlib._chk_map_pyx import (
 
1620
        _search_key_16,
 
1621
        _search_key_255,
 
1622
        _deserialise_leaf_node,
 
1623
        _deserialise_internal_node,
 
1624
        )
 
1625
except ImportError:
 
1626
    from bzrlib._chk_map_py import (
 
1627
        _search_key_16,
 
1628
        _search_key_255,
 
1629
        _deserialise_leaf_node,
 
1630
        _deserialise_internal_node,
 
1631
        )
 
1632
search_key_registry.register('hash-16-way', _search_key_16)
 
1633
search_key_registry.register('hash-255-way', _search_key_255)