[PATCH 1 of 2] revset: added cached generated list on generatorset

[Lucas Moscovicz]

# HG changeset patch
# User Lucas Moscovicz <lmoscovicz at fb.com>
# Date 1392259549 28800
#      Wed Feb 12 18:45:49 2014 -0800
# Node ID 4c90d85a7bc219eac7575beb7257619330f17c02
# Parent  86e9755cb0462f2c33b00b0cb362a63bd583f25d
revset: added cached generated list on generatorset

This allows to iterate the generatorset more than once.

diff --git a/mercurial/revset.py b/mercurial/revset.py
--- a/mercurial/revset.py
+++ b/mercurial/revset.py
@@ -2147,20 +2147,32 @@
         return set([r for r in self])
 
 class generatorset(object):
-    """Wrapper structure for generators that provides lazy membership."""
+    """Wrapper structure for generators that provides lazy membership and can
+    be iterated more than once.
+    When asked for membership it generates values until either it finds the
+    requested one or has gone through all the elements in the generator
+    """
     def __init__(self, gen):
         self._gen = gen
         self._iter = iter(gen)
         self._cache = {}
+        self._genlist = baseset([])
+        self._iterated = False
+
+    def _nextitem(self):
+        l = self._iter.next()
+        self._cache[l] = True
+        self._genlist.append(l)
+        return l
 
     def __contains__(self, x):
         if x in self._cache:
             return self._cache[x]
 
+        self._iterated = True
         while True:
             try:
-                l = self._iter.next()
-                self._cache[l] = True
+                l = self._nextitem()
                 if l == x:
                     return True
             except (StopIteration):
@@ -2170,9 +2182,21 @@
         return False
 
     def __iter__(self):
-        for item in self._gen:
-            self._cache[item] = True
-            yield item
+        if self._iterated:
+            for l in self._genlist:
+                yield l
+            while True:
+                try:
+                    item = self._nextitem()
+                    yield item
+                except (StopIteration):
+                    break
+        else:
+            self._iterated = True
+            for item in self._gen:
+                self._cache[item] = True
+                self._genlist.append(item)
+                yield item
 
     def set(self):
         return self