[PATCH] Add ancestors and descendants to revlog

[Stefano Tortarolo]

# HG changeset patch
# User Stefano Tortarolo <stefano.tortarolo at gmail.com>
# Date 1216484390 -7200
# Node ID 6ccea5af0524a4047a75b04dddc3c7718f86cc36
# Parent  1a9577da9d02a386526c6a6791d6d57f56b28126
Add ancestors and descendants to revlog
This patch adds two methods to revlog:
- ancestors: given a list of revisions returns their ancestors
- descendants: given a list of revisions return their descendants

diff --git a/mercurial/revlog.py b/mercurial/revlog.py
--- a/mercurial/revlog.py
+++ b/mercurial/revlog.py
@@ -595,6 +595,27 @@
                     reachable[p] = 1
                     visit.append(p)
         return reachable
+
+    def ancestors(self, *revs):
+        'Generate the ancestors of revs using a breadth-first visit'
+        visit = list(revs)
+        seen = util.set([nullrev])
+        while visit:
+            for parent in self.parentrevs(visit.pop(0)):
+                if parent not in seen:
+                    visit.append(parent)
+                    seen.add(parent)
+                    yield parent
+
+    def descendants(self, *revs):
+        'Generate the descendants of revs in topological order'
+        seen = util.set(revs)
+        for i in xrange(min(revs) + 1, len(self)):
+            for x in self.parentrevs(i):
+                if x != nullrev and x in seen:
+                    seen.add(i)
+                    yield i
+                    break
 
     def nodesbetween(self, roots=None, heads=None):
         """Return a tuple containing three elements. Elements 1 and 2 contain
diff --git a/tests/test-revlog-ancestry.py b/tests/test-revlog-ancestry.py
new file mode 100644
--- /dev/null
+++ b/tests/test-revlog-ancestry.py
@@ -0,0 +1,74 @@
+import os
+from mercurial import hg, ui, merge
+from hgext import graphlog
+
+u = ui.ui()
+
+repo = hg.repository(u, 'test1', create=1)
+os.chdir('test1')
+
+def commit(text, time):
+    repo.commit(text=text, date="%d 0" % time)
+
+def addcommit(name, time):
+    f = file(name, 'w')
+    f.write('%s\n' % name)
+    f.close()
+    repo.add([name])
+    commit(name, time)
+
+def update(rev):
+    merge.update(repo, rev, False, True, False)
+
+def merge_(rev):
+    merge.update(repo, rev, True, False, False)
+
+if __name__ == '__main__':
+    addcommit("A", 0)
+    addcommit("B", 1)
+
+    update(0)
+    addcommit("C", 2)
+
+    merge_(1)
+    commit("D", 3)
+
+    update(2)
+    addcommit("E", 4)
+    addcommit("F", 5)
+
+    update(3)
+    addcommit("G", 6)
+
+    merge_(5)
+    commit("H", 7)
+
+    update(5)
+    addcommit("I", 8)
+
+    # Ancestors
+    print 'Ancestors of 5'
+    for r in repo.changelog.ancestors(5):
+        print r, 
+
+    print '\nAncestors of 6 and 5'
+    for r in repo.changelog.ancestors(6, 5):
+        print r, 
+
+    print '\nAncestors of 5 and 4'
+    for r in repo.changelog.ancestors(5, 4):
+        print r, 
+
+    # Descendants
+    print '\n\nDescendants of 5'
+    for r in repo.changelog.descendants(5):
+        print r, 
+
+    print '\nDescendants of 5 and 3'
+    for r in repo.changelog.descendants(5, 3):
+        print r, 
+
+    print '\nDescendants of 5 and 4'
+    for r in repo.changelog.descendants(5, 4):
+        print r, 
+
diff --git a/tests/test-revlog-ancestry.py.out b/tests/test-revlog-ancestry.py.out
new file mode 100644
--- /dev/null
+++ b/tests/test-revlog-ancestry.py.out
@@ -0,0 +1,13 @@
+Ancestors of 5
+4 2 0 
+Ancestors of 6 and 5
+3 4 2 1 0 
+Ancestors of 5 and 4
+4 2 0 
+
+Descendants of 5
+7 8 
+Descendants of 5 and 3
+6 7 8 
+Descendants of 5 and 4
+5 7 8