Rev 12: Write a helper function that can turn 2 lists of keys into a position list. in http://bzr.arbash-meinel.com/plugins/index2

[John Arbash Meinel]

At http://bzr.arbash-meinel.com/plugins/index2

------------------------------------------------------------
revno: 12
revision-id: john at arbash-meinel.com-20080701231134-73vr4qkr6dogf3h0
parent: john at arbash-meinel.com-20080701192734-9mbeqa4zszat3z1f
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: index2
timestamp: Tue 2008-07-01 18:11:34 -0500
message:
  Write a helper function that can turn 2 lists of keys into a position list.
-------------- next part --------------
=== modified file 'btree_index.py'
--- a/btree_index.py	2008-07-01 19:27:34 +0000
+++ b/btree_index.py	2008-07-01 23:11:34 +0000
@@ -354,6 +354,57 @@
                 for key, (value, refs) in node.keys.items():
                     yield (self, key, value)
 
+    @staticmethod
+    def _positions_for_keys(in_keys, fixed_keys):
+        """Find the positions where each 'in_key' would fit in fixed_keys.
+
+        This is equivalent to doing "bisect_right" on each in_key into
+        fixed_keys
+
+        :param in_keys: A sorted list of keys to match with fixed_keys
+        :param fixed_keys: A sorted list of keys to match against
+        :return: A list of integer positions.
+        """
+        if not in_keys:
+            return []
+        if not fixed_keys:
+            raise AssertionError('we must have entries in fixed_keys')
+        in_keys_iter = iter(in_keys)
+        fixed_keys_iter = enumerate(fixed_keys)
+        cur_in_key = in_keys_iter.next()
+        cur_fixed_offset, cur_fixed_key = fixed_keys_iter.next()
+
+        class InputDone(Exception): pass
+        class FixedDone(Exception): pass
+
+        output = []
+
+        try:
+            while True:
+                if cur_in_key < cur_fixed_key:
+                    output.append(cur_fixed_offset)
+                    while cur_in_key < cur_fixed_key:
+                        try:
+                            cur_in_key = in_keys_iter.next()
+                        except StopIteration:
+                            raise InputDone
+                    # At this point cur_in_key must be >= cur_fixed_key
+                # step the cur_fixed_key until we pass the cur key, or walk off
+                # the end
+                while cur_in_key >= cur_fixed_key:
+                    try:
+                        cur_fixed_offset, cur_fixed_key = fixed_keys_iter.next()
+                    except StopIteration:
+                        raise FixedDone
+        except InputDone:
+            # We consumed all of the input, nothing more to do
+            pass
+        except FixedDone:
+            # There was some input left, but we consumed all of fixed, so we
+            # have to add one more for the tail
+            output.append(len(fixed_keys))
+        return output
+
     def iter_entries(self, keys):
         """Iterate over keys within the index.
 
@@ -369,6 +420,13 @@
             self._cache_nodes([0])
         if not self._key_count:
             return
+        node = self._get_node(0)
+        for row_length in self._row_lengths:
+            next_nodes = []
+
+            # Split up the keys based on which nodes will work best next
+            # Both the node keys and are internal list are both sorted, so we
+            # just step both of them to find the right locations.
         for key in keys:
             # repeated single-bisection : 'make it work'
             # FUTURE: look within the current node for the next key,

=== modified file 'tests/test_btree_index.py'
--- a/tests/test_btree_index.py	2008-07-01 09:51:16 +0000
+++ b/tests/test_btree_index.py	2008-07-01 23:11:34 +0000
@@ -19,6 +19,7 @@
 
 import zlib
 
+from bzrlib import tests
 from bzrlib.plugins import index2
 from bzrlib.plugins.index2 import errors, btree_index
 from bzrlib.tests import TestCaseWithTransport
@@ -1116,3 +1117,32 @@
             ('11', '44'): ('value:4', ((), (('11', 'ref00'),)))
             }, node.keys)
 
+
+class TestPositionsForKeys(tests.TestCase):
+
+    def assertPositionsForKeys(self, offsets, search_keys, fixed_keys):
+        self.assertEqual(offsets,
+                         btree_index.BTreeGraphIndex._positions_for_keys(
+                            search_keys, fixed_keys))
+
+    def test_after(self):
+        self.assertPositionsForKeys([1], ['b'], ['a'])
+        self.assertPositionsForKeys([3], ['e', 'f', 'g'], ['a', 'b', 'c'])
+
+    def test_before(self):
+        self.assertPositionsForKeys([0], ['a'], ['b'])
+        self.assertPositionsForKeys([0], ['a', 'b', 'c', 'd'], ['e', 'f', 'g'])
+
+    def test_exact(self):
+        self.assertPositionsForKeys([1], ['a'], ['a'])
+        self.assertPositionsForKeys([1, 2], ['a', 'b'], ['a', 'b'])
+        self.assertPositionsForKeys([1, 3], ['a', 'c'], ['a', 'b', 'c'])
+
+    def test_inbetween(self):
+        self.assertPositionsForKeys([1], ['b'], ['a', 'c'])
+        self.assertPositionsForKeys([1, 2], ['b', 'c', 'd', 'f', 'g'],
+                                    ['a', 'e', 'h'])
+
+    def test_mixed(self):
+        self.assertPositionsForKeys([0, 2, 4], ['a', 'b', 'd', 'e', 'g', 'h'],
+                                    ['c', 'd', 'f', 'g'])