Rev 268: Include the known_graph code. in lp:~jameinel/bzr-fastimport/trunk

[John Arbash Meinel]

At lp:~jameinel/bzr-fastimport/trunk

------------------------------------------------------------
revno: 268 [merge]
revision-id: john at arbash-meinel.com-20100105072414-gy7y1iw47jajujna
parent: john at arbash-meinel.com-20091211145504-ez9cj300ty9g4ly4
parent: john at arbash-meinel.com-20091201151519-key0pustue6p1q1c
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: trunk
timestamp: Tue 2010-01-05 01:24:14 -0600
message:
  Include the known_graph code.
modified:
  revision_store.py              revisionloader.py-20080215075452-zn0u2pky2jnf7j07-1
-------------- next part --------------
=== modified file 'revision_store.py'
--- a/revision_store.py	2009-12-09 20:14:46 +0000
+++ b/revision_store.py	2010-01-05 07:24:14 +0000
@@ -18,8 +18,16 @@
 
 import cStringIO
 
-from bzrlib import errors, inventory, knit, lru_cache, osutils, trace
-from bzrlib import revision as _mod_revision
+from bzrlib import (
+    errors,
+    graph as _mod_graph,
+    inventory,
+    knit,
+    lru_cache,
+    osutils,
+    revision as _mod_revision,
+    trace,
+    )
 
 
 class _TreeShim(object):
@@ -159,6 +167,8 @@
         :param repository: the target repository
         """
         self.repo = repo
+        self._graph = None
+        self._use_known_graph = True
         self._supports_chks = getattr(repo._format, 'supports_chks', False)
 
     def expects_rich_root(self):
@@ -355,6 +365,28 @@
             parents=rev.parent_ids, config=None, timestamp=rev.timestamp,
             timezone=rev.timezone, committer=rev.committer,
             revprops=rev.properties, revision_id=rev.revision_id)
+        if self._graph is None and self._use_known_graph:
+            if (getattr(_mod_graph, 'GraphThunkIdsToKeys', None) is None
+                or getattr(_mod_graph.KnownGraph, 'add_node', None) is None):
+                self._use_known_graph = False
+            else:
+                self._graph = self.repo.revisions.get_known_graph_ancestry(
+                    [(r,) for r in rev.parent_ids])
+        if self._graph is not None:
+            orig_heads = builder._heads
+            def thunked_heads(file_id, revision_ids):
+                # self._graph thinks in terms of keys, not ids, so translate
+                # them
+                # old_res = orig_heads(file_id, revision_ids)
+                if len(revision_ids) < 2:
+                    res = set(revision_ids)
+                else:
+                    res = set([h[0] for h in
+                              self._graph.heads([(r,) for r in revision_ids])])
+                # if old_res != res:
+                #     import pdb; pdb.set_trace()
+                return res
+            builder._heads = thunked_heads
 
         if rev.parent_ids:
             basis_rev_id = rev.parent_ids[0]
@@ -379,6 +411,10 @@
         rev.inv_sha1 = builder.inv_sha1
         builder.repository.add_revision(builder._new_revision_id, rev,
             builder.new_inventory, builder._config)
+        if self._graph is not None:
+            # TODO: Use StaticTuple and .intern() for these things
+            self._graph.add_node((builder._new_revision_id,),
+                                 [(p,) for p in rev.parent_ids])
 
         if signature is not None:
             raise AssertionError('signatures not guaranteed yet')