[PATCH 3 of 6] ancestor: add function to generate ancestors lazily

[Siddharth Agarwal]

# HG changeset patch
# User Siddharth Agarwal <sid0 at fb.com>
# Date 1355513483 28800
# Node ID 74b0e1f732cf868681ace3e2bc44d5eb54e1729a
# Parent  35eb8d05863f3907b5bfd5f023a4746f91ef97d8
ancestor: add function to generate ancestors lazily

This will be used and performance tested in upcoming patches in this series.

diff -r 35eb8d05863f -r 74b0e1f732cf mercurial/ancestor.py
--- a/mercurial/ancestor.py	Fri Dec 14 10:23:18 2012 -0800
+++ b/mercurial/ancestor.py	Fri Dec 14 11:31:23 2012 -0800
@@ -316,3 +316,55 @@
 
     missing.reverse()
     return missing
+
+class lazyset(object):
+    def __init__(self, backingset, targetfn, args):
+        """Create a new lazyset. If the entity whose membership is being tested
+        isn't found in the backing set, targetfn is called. targetfn is
+        supposed to mutate backingset so that hopefully future calls to
+        targetfn won't be necessary."""
+        self._backingset = backingset
+        self._targetfn = targetfn
+        self._args = args
+
+    def __contains__(self, n):
+        return n in self._backingset or self._targetfn(n, *self._args)
+
+def lazyancestorset(revs, pfunc, increvs, stoprev=0):
+    """Return a lazyset generating ancestors for the given revs.
+
+    This is computed lazily starting from tip and working in descending rev
+    order. The lazyset only supports the membership operation.
+
+    revs should be an iterable. pfunc must return a list of parent revs for a
+    given rev. increvs is a boolean that indicates whether revs should be
+    included in the set. Revs lower than stoprev will not be generated."""
+    # Python's heap is a min-heap. Multiply all values by -1 to convert it
+    # into a max-heap.
+    visit = [-rev for rev in revs]
+    heapq.heapify(visit)
+    if increvs:
+        seen = set(revs)
+    else:
+        seen = set()
+
+    # This would ideally be a generator with send(), but that's 2.5+ only. The
+    # arguments are passed in rather than closed over because that makes a
+    # non-trivial speed difference.
+    def targetfn(target, visit, seen, stoprev, pfunc, heappush, heappop):
+        targetseen = False
+
+        while visit and -visit[0] >= target and not targetseen:
+            for parent in pfunc(-heappop(visit)):
+                if parent in seen or parent < stoprev:
+                    continue
+                heappush(visit, -parent)
+                seen.add(parent)
+                if parent == target:
+                    targetseen = True
+
+        return targetseen
+
+    # Reuse the seen set as the backing set
+    return lazyset(seen, targetfn, (visit, seen, stoprev, pfunc,
+                                    heapq.heappush, heapq.heappop))
diff -r 35eb8d05863f -r 74b0e1f732cf tests/test-ancestor.py
--- a/tests/test-ancestor.py	Fri Dec 14 10:23:18 2012 -0800
+++ b/tests/test-ancestor.py	Fri Dec 14 11:31:23 2012 -0800
@@ -70,5 +70,33 @@
     runmissingancestors([10, 11, 12], [13])
     runmissingancestors([13], [10, 11, 12])
 
+def genlazyancestorset(revs, increvs, stoprev=0):
+    print "%% lazy ancestor set for %s, including revs = %s" % (revs, increvs)
+    return ancestor.lazyancestorset(revs, pfunc, increvs, stoprev=stoprev)
+
+def isinlazyset(s, l):
+    for n in l:
+        print "%% is %d in set? %s" % (n, n in s)
+
+def test_lazyancestorset():
+    # Empty revs
+    s = genlazyancestorset([], False)
+    isinlazyset(s, [3, 0, -1])
+
+    # Standard example
+    s = genlazyancestorset([11, 13], False)
+    isinlazyset(s, [11, 13, 7, 9, 8, 3, 6, 4, 1, -1, 0])
+
+    # Including revs
+    s = genlazyancestorset([11, 13], True)
+    isinlazyset(s, [11, 13, 7, 9, 8, 3, 6, 4, 1, -1, 0])
+
+    # Test with stoprev
+    s = genlazyancestorset([11, 13], False, stoprev=6)
+    isinlazyset(s, [11, 13, 7, 9, 8, 3, 6, 4, 1, -1, 0])
+    s = genlazyancestorset([11, 13], True, stoprev=6)
+    isinlazyset(s, [11, 13, 7, 9, 8, 3, 6, 4, 1, -1, 0])
+
 if __name__ == '__main__':
     test_missingancestors()
+    test_lazyancestorset()
diff -r 35eb8d05863f -r 74b0e1f732cf tests/test-ancestor.py.out
--- a/tests/test-ancestor.py.out	Fri Dec 14 10:23:18 2012 -0800
+++ b/tests/test-ancestor.py.out	Fri Dec 14 11:31:23 2012 -0800
@@ -34,3 +34,55 @@
 [0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12]
 % ancestors of [13] and not of [10, 11, 12]
 [8, 13]
+% lazy ancestor set for [], including revs = False
+% is 3 in set? False
+% is 0 in set? False
+% is -1 in set? False
+% lazy ancestor set for [11, 13], including revs = False
+% is 11 in set? False
+% is 13 in set? False
+% is 7 in set? True
+% is 9 in set? False
+% is 8 in set? True
+% is 3 in set? True
+% is 6 in set? False
+% is 4 in set? True
+% is 1 in set? True
+% is -1 in set? False
+% is 0 in set? True
+% lazy ancestor set for [11, 13], including revs = True
+% is 11 in set? True
+% is 13 in set? True
+% is 7 in set? True
+% is 9 in set? False
+% is 8 in set? True
+% is 3 in set? True
+% is 6 in set? False
+% is 4 in set? True
+% is 1 in set? True
+% is -1 in set? False
+% is 0 in set? True
+% lazy ancestor set for [11, 13], including revs = False
+% is 11 in set? False
+% is 13 in set? False
+% is 7 in set? True
+% is 9 in set? False
+% is 8 in set? True
+% is 3 in set? False
+% is 6 in set? False
+% is 4 in set? False
+% is 1 in set? False
+% is -1 in set? False
+% is 0 in set? False
+% lazy ancestor set for [11, 13], including revs = True
+% is 11 in set? True
+% is 13 in set? True
+% is 7 in set? True
+% is 9 in set? False
+% is 8 in set? True
+% is 3 in set? False
+% is 6 in set? False
+% is 4 in set? False
+% is 1 in set? False
+% is -1 in set? False
+% is 0 in set? False