Patchwork [9,of,9] revset: add depth limit to descendants() (issue5374)

diff --git a/mercurial/dagop.py b/mercurial/dagop.py
--- a/mercurial/dagop.py
+++ b/mercurial/dagop.py
@@ -125,10 +125,38 @@  def _genrevdescendants(repo, revs, follo
                     yield i
                     break
 
-def revdescendants(repo, revs, followfirst):
+def _builddescendantsmap(repo, startrev, followfirst):
+    """Build map of 'rev -> child revs', offset from startrev"""
+    cl = repo.changelog
+    nullrev = node.nullrev
+    descmap = [[] for _rev in xrange(startrev, len(cl))]
+    for currev in cl.revs(startrev + 1):
+        p1rev, p2rev = cl.parentrevs(currev)
+        if p1rev >= startrev:
+            descmap[p1rev - startrev].append(currev)
+        if not followfirst and p2rev != nullrev and p2rev >= startrev:
+            descmap[p2rev - startrev].append(currev)
+    return descmap
+
+def _genrevdescendantsofdepth(repo, revs, followfirst, startdepth, stopdepth):
+    startrev = revs.min()
+    descmap = _builddescendantsmap(repo, startrev, followfirst)
+    def pfunc(rev):
+        return descmap[rev - startrev]
+    return _walkrevtree(pfunc, revs, startdepth, stopdepth, reverse=False)
+
+def revdescendants(repo, revs, followfirst, startdepth=None, stopdepth=None):
     """Like revlog.descendants() but supports additional options, includes
-    the given revs themselves, and returns a smartset"""
-    gen = _genrevdescendants(repo, revs, followfirst)
+    the given revs themselves, and returns a smartset
+
+    Scan ends at the stopdepth (exlusive) if specified. Revisions found
+    earlier than the startdepth are omitted.
+    """
+    if startdepth is None and stopdepth is None:
+        gen = _genrevdescendants(repo, revs, followfirst)
+    else:
+        gen = _genrevdescendantsofdepth(repo, revs, followfirst,
+                                        startdepth, stopdepth)
     return generatorset(gen, iterasc=True)
 
 def _reachablerootspure(repo, minroot, roots, heads, includepath):
diff --git a/mercurial/revset.py b/mercurial/revset.py
--- a/mercurial/revset.py
+++ b/mercurial/revset.py
@@ -595,23 +595,42 @@  def desc(repo, subset, x):
     return subset.filter(lambda r: matcher(repo[r].description()),
                          condrepr=('<desc %r>', ds))
 
-def _descendants(repo, subset, x, followfirst=False):
+def _descendants(repo, subset, x, followfirst=False, startdepth=None,
+                 stopdepth=None):
     roots = getset(repo, fullreposet(repo), x)
     if not roots:
         return baseset()
-    s = dagop.revdescendants(repo, roots, followfirst)
+    s = dagop.revdescendants(repo, roots, followfirst, startdepth, stopdepth)
     return subset & s
 
-@predicate('descendants(set)', safe=True)
+@predicate('descendants(set[, depth])', safe=True)
 def descendants(repo, subset, x):
     """Changesets which are descendants of changesets in set, including the
     given changesets themselves.
+
+    If depth is specified, the result only includes changesets up to
+    the specified generation.
     """
-    args = getargsdict(x, 'descendants', 'set')
+    # startdepth is for internal use only until we can decide the UI
+    args = getargsdict(x, 'descendants', 'set depth startdepth')
     if 'set' not in args:
         # i18n: "descendants" is a keyword
         raise error.ParseError(_('descendants takes at least 1 argument'))
-    return _descendants(repo, subset, args['set'])
+    startdepth = stopdepth = None
+    if 'startdepth' in args:
+        n = getinteger(args['startdepth'],
+                       "descendants expects an integer startdepth")
+        if n < 0:
+            raise error.ParseError("negative startdepth")
+        startdepth = n
+    if 'depth' in args:
+        # i18n: "descendants" is a keyword
+        n = getinteger(args['depth'], _("descendants expects an integer depth"))
+        if n < 0:
+            raise error.ParseError(_("negative depth"))
+        stopdepth = n + 1
+    return _descendants(repo, subset, args['set'],
+                        startdepth=startdepth, stopdepth=stopdepth)
 
 @predicate('_firstdescendants', safe=True)
 def _firstdescendants(repo, subset, x):
diff --git a/tests/test-revset.t b/tests/test-revset.t
--- a/tests/test-revset.t
+++ b/tests/test-revset.t
@@ -980,6 +980,60 @@  test descendants
   7
   8
 
+test descendants with depth limit
+
+ (depth=0 selects the node itself)
+
+  $ log 'descendants(0, depth=0)'
+  0
+  $ log 'null: & descendants(null, depth=0)'
+  -1
+
+ (p2 = null should be ignored)
+
+  $ log 'null: & descendants(null, depth=2)'
+  -1
+  0
+  1
+
+ (multiple paths: depth(6) = (2, 3))
+
+  $ log 'descendants(1+3, depth=2)'
+  1
+  2
+  3
+  4
+  5
+  6
+
+ (multiple paths: depth(5) = (1, 2), depth(6) = (2, 3))
+
+  $ log 'descendants(3+1, depth=2, startdepth=2)'
+  4
+  5
+  6
+
+ (multiple depths: depth(6) = (0, 2, 4), search for depth=2)
+
+  $ log 'descendants(0+3+6, depth=3, startdepth=1)'
+  1
+  2
+  3
+  4
+  5
+  6
+  7
+
+ (multiple depths: depth(6) = (0, 4), no match)
+
+  $ log 'descendants(0+6, depth=3, startdepth=1)'
+  1
+  2
+  3
+  4
+  5
+  7
+
 test author
 
   $ log 'author(bob)'