/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: Ian Clatworthy
  • Date: 2009-07-13 06:58:49 UTC
  • mto: (4527.1.1 integration)
  • mto: This revision was merged to the branch mainline in revision 4529.
  • Revision ID: ian.clatworthy@canonical.com-20090713065849-n7g2rsjyl6dt1mgv
Apply review feedback

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
def _find_children_info(store, interesting_keys, uninteresting_keys, pb):
 
1400
    """Read the associated records, and determine what is interesting."""
 
1401
    uninteresting_keys = set(uninteresting_keys)
 
1402
    chks_to_read = uninteresting_keys.union(interesting_keys)
 
1403
    next_uninteresting = set()
 
1404
    next_interesting = set()
 
1405
    next_interesting_intersection = None
 
1406
    uninteresting_items = set()
 
1407
    interesting_items = set()
 
1408
    interesting_to_yield = []
 
1409
    for record in store.get_record_stream(chks_to_read, 'unordered', True):
 
1410
        # records_read.add(record.key())
 
1411
        if pb is not None:
 
1412
            pb.tick()
 
1413
        bytes = record.get_bytes_as('fulltext')
 
1414
        # We don't care about search_key_func for this code, because we only
 
1415
        # care about external references.
 
1416
        node = _deserialise(bytes, record.key, search_key_func=None)
 
1417
        if record.key in uninteresting_keys:
 
1418
            if type(node) is InternalNode:
 
1419
                next_uninteresting.update(node.refs())
 
1420
            else:
 
1421
                # We know we are at a LeafNode, so we can pass None for the
 
1422
                # store
 
1423
                uninteresting_items.update(node.iteritems(None))
 
1424
        else:
 
1425
            interesting_to_yield.append(record.key)
 
1426
            if type(node) is InternalNode:
 
1427
                if next_interesting_intersection is None:
 
1428
                    next_interesting_intersection = set(node.refs())
 
1429
                else:
 
1430
                    next_interesting_intersection = \
 
1431
                        next_interesting_intersection.intersection(node.refs())
 
1432
                next_interesting.update(node.refs())
 
1433
            else:
 
1434
                interesting_items.update(node.iteritems(None))
 
1435
    return (next_uninteresting, uninteresting_items,
 
1436
            next_interesting, interesting_to_yield, interesting_items,
 
1437
            next_interesting_intersection)
 
1438
 
 
1439
 
 
1440
def _find_all_uninteresting(store, interesting_root_keys,
 
1441
                            uninteresting_root_keys, pb):
 
1442
    """Determine the full set of uninteresting keys."""
 
1443
    # What about duplicates between interesting_root_keys and
 
1444
    # uninteresting_root_keys?
 
1445
    if not uninteresting_root_keys:
 
1446
        # Shortcut case. We know there is nothing uninteresting to filter out
 
1447
        # So we just let the rest of the algorithm do the work
 
1448
        # We know there is nothing uninteresting, and we didn't have to read
 
1449
        # any interesting records yet.
 
1450
        return (set(), set(), set(interesting_root_keys), [], set())
 
1451
    all_uninteresting_chks = set(uninteresting_root_keys)
 
1452
    all_uninteresting_items = set()
 
1453
 
 
1454
    # First step, find the direct children of both the interesting and
 
1455
    # uninteresting set
 
1456
    (uninteresting_keys, uninteresting_items,
 
1457
     interesting_keys, interesting_to_yield,
 
1458
     interesting_items, interesting_intersection,
 
1459
     ) = _find_children_info(store, interesting_root_keys,
 
1460
                                              uninteresting_root_keys,
 
1461
                                              pb=pb)
 
1462
    all_uninteresting_chks.update(uninteresting_keys)
 
1463
    all_uninteresting_items.update(uninteresting_items)
 
1464
    del uninteresting_items
 
1465
    # Do not examine in detail pages common to all interesting trees.
 
1466
    # Pages that are common to all interesting trees will have their
 
1467
    # older versions found via the uninteresting tree traversal. Some pages
 
1468
    # found via the interesting trees traversal will be uninteresting for
 
1469
    # other of the interesting trees, which is why we require the pages to be
 
1470
    # common for us to trim them.
 
1471
    if interesting_intersection is not None:
 
1472
        uninteresting_keys.difference_update(interesting_intersection)
 
1473
 
 
1474
    # Second, find the full set of uninteresting bits reachable by the
 
1475
    # uninteresting roots
 
1476
    chks_to_read = uninteresting_keys
 
1477
    while chks_to_read:
 
1478
        next_chks = set()
 
1479
        for record in store.get_record_stream(chks_to_read, 'unordered', False):
 
1480
            # TODO: Handle 'absent'
 
1481
            if pb is not None:
 
1482
                pb.tick()
 
1483
            bytes = record.get_bytes_as('fulltext')
 
1484
            # We don't care about search_key_func for this code, because we
 
1485
            # only care about external references.
 
1486
            node = _deserialise(bytes, record.key, search_key_func=None)
 
1487
            if type(node) is InternalNode:
 
1488
                # uninteresting_prefix_chks.update(node._items.iteritems())
 
1489
                chks = node._items.values()
 
1490
                # TODO: We remove the entries that are already in
 
1491
                #       uninteresting_chks ?
 
1492
                next_chks.update(chks)
 
1493
                all_uninteresting_chks.update(chks)
 
1494
            else:
 
1495
                all_uninteresting_items.update(node._items.iteritems())
 
1496
        chks_to_read = next_chks
 
1497
    return (all_uninteresting_chks, all_uninteresting_items,
 
1498
            interesting_keys, interesting_to_yield, interesting_items)
 
1499
 
 
1500
 
 
1501
def iter_interesting_nodes(store, interesting_root_keys,
 
1502
                           uninteresting_root_keys, pb=None):
 
1503
    """Given root keys, find interesting nodes.
 
1504
 
 
1505
    Evaluate nodes referenced by interesting_root_keys. Ones that are also
 
1506
    referenced from uninteresting_root_keys are not considered interesting.
 
1507
 
 
1508
    :param interesting_root_keys: keys which should be part of the
 
1509
        "interesting" nodes (which will be yielded)
 
1510
    :param uninteresting_root_keys: keys which should be filtered out of the
 
1511
        result set.
 
1512
    :return: Yield
 
1513
        (interesting record, {interesting key:values})
 
1514
    """
 
1515
    # TODO: consider that it may be more memory efficient to use the 20-byte
 
1516
    #       sha1 string, rather than tuples of hexidecimal sha1 strings.
 
1517
    # TODO: Try to factor out a lot of the get_record_stream() calls into a
 
1518
    #       helper function similar to _read_bytes. This function should be
 
1519
    #       able to use nodes from the _page_cache as well as actually
 
1520
    #       requesting bytes from the store.
 
1521
 
 
1522
    (all_uninteresting_chks, all_uninteresting_items, interesting_keys,
 
1523
     interesting_to_yield, interesting_items) = _find_all_uninteresting(store,
 
1524
        interesting_root_keys, uninteresting_root_keys, pb)
 
1525
 
 
1526
    # Now that we know everything uninteresting, we can yield information from
 
1527
    # our first request
 
1528
    interesting_items.difference_update(all_uninteresting_items)
 
1529
    interesting_to_yield = set(interesting_to_yield) - all_uninteresting_chks
 
1530
    if interesting_items:
 
1531
        yield None, interesting_items
 
1532
    if interesting_to_yield:
 
1533
        # We request these records again, rather than buffering the root
 
1534
        # records, most likely they are still in the _group_cache anyway.
 
1535
        for record in store.get_record_stream(interesting_to_yield,
 
1536
                                              'unordered', False):
 
1537
            yield record, []
 
1538
    all_uninteresting_chks.update(interesting_to_yield)
 
1539
    interesting_keys.difference_update(all_uninteresting_chks)
 
1540
 
 
1541
    chks_to_read = interesting_keys
 
1542
    counter = 0
 
1543
    while chks_to_read:
 
1544
        next_chks = set()
 
1545
        for record in store.get_record_stream(chks_to_read, 'unordered', False):
 
1546
            counter += 1
 
1547
            if pb is not None:
 
1548
                pb.update('find chk pages', counter)
 
1549
            # TODO: Handle 'absent'?
 
1550
            bytes = record.get_bytes_as('fulltext')
 
1551
            # We don't care about search_key_func for this code, because we
 
1552
            # only care about external references.
 
1553
            node = _deserialise(bytes, record.key, search_key_func=None)
 
1554
            if type(node) is InternalNode:
 
1555
                # all_uninteresting_chks grows large, as it lists all nodes we
 
1556
                # don't want to process (including already seen interesting
 
1557
                # nodes).
 
1558
                # small.difference_update(large) scales O(large), but
 
1559
                # small.difference(large) scales O(small).
 
1560
                # Also, we know we just _deserialised this node, so we can
 
1561
                # access the dict directly.
 
1562
                chks = set(node._items.itervalues()).difference(
 
1563
                            all_uninteresting_chks)
 
1564
                # Is set() and .difference_update better than:
 
1565
                # chks = [chk for chk in node.refs()
 
1566
                #              if chk not in all_uninteresting_chks]
 
1567
                next_chks.update(chks)
 
1568
                # These are now uninteresting everywhere else
 
1569
                all_uninteresting_chks.update(chks)
 
1570
                interesting_items = []
 
1571
            else:
 
1572
                interesting_items = [item for item in node._items.iteritems()
 
1573
                                     if item not in all_uninteresting_items]
 
1574
                # TODO: Do we need to filter out items that we have already
 
1575
                #       seen on other pages? We don't really want to buffer the
 
1576
                #       whole thing, but it does mean that callers need to
 
1577
                #       understand they may get duplicate values.
 
1578
                # all_uninteresting_items.update(interesting_items)
 
1579
            yield record, interesting_items
 
1580
        chks_to_read = next_chks
 
1581
 
 
1582
 
 
1583
try:
 
1584
    from bzrlib._chk_map_pyx import (
 
1585
        _search_key_16,
 
1586
        _search_key_255,
 
1587
        _deserialise_leaf_node,
 
1588
        _deserialise_internal_node,
 
1589
        )
 
1590
except ImportError:
 
1591
    from bzrlib._chk_map_py import (
 
1592
        _search_key_16,
 
1593
        _search_key_255,
 
1594
        _deserialise_leaf_node,
 
1595
        _deserialise_internal_node,
 
1596
        )
 
1597
search_key_registry.register('hash-16-way', _search_key_16)
 
1598
search_key_registry.register('hash-255-way', _search_key_255)