Rev 4821: Use a KnownGraph to implement the heads searches. in http://bazaar.launchpad.net/~jameinel/bzr/2.1.0b4-convert-kg-heads

[John Arbash Meinel]

At http://bazaar.launchpad.net/~jameinel/bzr/2.1.0b4-convert-kg-heads

------------------------------------------------------------
revno: 4821
revision-id: john at arbash-meinel.com-20091122053626-xgsgd01p2ld22hqj
parent: john at arbash-meinel.com-20091121194438-jkgufdz1il9ssbj0
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: 2.1.0b4-convert-kg-heads
timestamp: Sat 2009-11-21 23:36:26 -0600
message:
  Use a KnownGraph to implement the heads searches.
  
  This loads the full revision graph 1 time, and makes very fast heads() calls,
  versus lots of loads of part of the graph. Timing shows it to be
  much better.
-------------- next part --------------
=== modified file 'bzrlib/repository.py'
--- a/bzrlib/repository.py	2009-11-21 19:44:38 +0000
+++ b/bzrlib/repository.py	2009-11-22 05:36:26 +0000
@@ -38,6 +38,7 @@
     lru_cache,
     osutils,
     revision as _mod_revision,
+    static_tuple,
     symbol_versioning,
     tsort,
     ui,
@@ -3819,7 +3820,7 @@
                 basis_id, delta, current_revision_id, parents_parents)
             cache[current_revision_id] = parent_tree
 
-    def _fetch_batch(self, revision_ids, basis_id, cache):
+    def _fetch_batch(self, revision_ids, basis_id, cache, graph=None):
         """Fetch across a few revisions.
 
         :param revision_ids: The revisions to copy
@@ -3898,7 +3899,7 @@
             from bzrlib.fetch import _new_root_data_stream
             root_stream = _new_root_data_stream(
                 root_keys_to_create, self._revision_id_to_root_id, parent_map,
-                self.source)
+                self.source, graph=graph)
             to_texts.insert_record_stream(root_stream)
         to_texts.insert_record_stream(from_texts.get_record_stream(
             text_keys, self.target._format._fetch_order,
@@ -3961,13 +3962,30 @@
         cache[basis_id] = basis_tree
         del basis_tree # We don't want to hang on to it here
         hints = []
+        if self._converting_to_rich_root:
+            st = static_tuple.StaticTuple
+            revision_keys = [st(r_id).intern() for r_id in revision_ids]
+            known_graph = self.source.revisions.get_known_graph_ancestry(
+                            revision_keys)
+            class ThunkIdsToKeysHeads(object):
+                def __init__(self, kg):
+                    self.kg = kg
+                def heads(self, revision_ids):
+                    revision_keys = [(r,) for r in revision_ids]
+                    heads = self.kg.heads(revision_keys)
+                    return set([h[0] for h in heads])
+            graph = ThunkIdsToKeysHeads(known_graph)
+        else:
+            graph = None
+
         for offset in range(0, len(revision_ids), batch_size):
             self.target.start_write_group()
             try:
                 pb.update('Transferring revisions', offset,
                           len(revision_ids))
                 batch = revision_ids[offset:offset+batch_size]
-                basis_id = self._fetch_batch(batch, basis_id, cache)
+                basis_id = self._fetch_batch(batch, basis_id, cache,
+                                             graph=graph)
             except:
                 self.target.abort_write_group()
                 raise