D5441: rust-cpython: binding for LazyAncestors

[gracinet (Georges Racinet)]

gracinet updated this revision to Diff 13007.
gracinet marked an inline comment as done.
gracinet edited the summary of this revision.

REPOSITORY
  rHG Mercurial

CHANGES SINCE LAST UPDATE
  https://phab.mercurial-scm.org/D5441?vs=12952&id=13007

REVISION DETAIL
  https://phab.mercurial-scm.org/D5441

AFFECTED FILES
  rust/hg-cpython/src/ancestors.rs
  tests/test-rust-ancestor.py

CHANGE DETAILS

diff --git a/tests/test-rust-ancestor.py b/tests/test-rust-ancestor.py
--- a/tests/test-rust-ancestor.py
+++ b/tests/test-rust-ancestor.py
@@ -9,7 +9,10 @@
     rustext = None
 else:
     # this would fail already without appropriate ancestor.__package__
-    from mercurial.rustext.ancestor import AncestorsIterator
+    from mercurial.rustext.ancestor import (
+        AncestorsIterator,
+        LazyAncestors
+    )
 
 try:
     from mercurial.cext import parsers as cparsers
@@ -71,6 +74,37 @@
         ait = AncestorsIterator(idx, [3], 0, False)
         self.assertEqual([r for r in ait], [2, 1, 0])
 
+    def testlazyancestors(self):
+        idx = self.parseindex()
+        start_count = sys.getrefcount(idx)  # should be 2 (see Python doc)
+        self.assertEqual({i: (r[5], r[6]) for i, r in enumerate(idx)},
+                         {0: (-1, -1),
+                          1: (0, -1),
+                          2: (1, -1),
+                          3: (2, -1)})
+        lazy = LazyAncestors(idx, [3], 0, True)
+        # we have two more references to the index:
+        # - in its inner iterator for __contains__ and __bool__
+        # - in the LazyAncestors instance itself (to spawn new iterators)
+        self.assertEqual(sys.getrefcount(idx), start_count + 2)
+
+        self.assertTrue(2 in lazy)
+        self.assertTrue(bool(lazy))
+        self.assertEqual(list(lazy), [3, 2, 1, 0])
+        # a second time to validate that we spawn new iterators
+        self.assertEqual(list(lazy), [3, 2, 1, 0])
+
+        # now let's watch the refcounts closer
+        ait = iter(lazy)
+        self.assertEqual(sys.getrefcount(idx), start_count + 3)
+        del ait
+        self.assertEqual(sys.getrefcount(idx), start_count + 2)
+        del lazy
+        self.assertEqual(sys.getrefcount(idx), start_count)
+
+        # let's check bool for an empty one
+        self.assertFalse(LazyAncestors(idx, [0], 0, False))
+
     def testrefcount(self):
         idx = self.parseindex()
         start_count = sys.getrefcount(idx)
@@ -87,7 +121,7 @@
         # and removing ref to the index after iterator init is no issue
         ait = AncestorsIterator(idx, [3], 0, True)
         del idx
-        self.assertEqual([r for r in ait], [3, 2, 1, 0])
+        self.assertEqual(list(ait), [3, 2, 1, 0])
 
     def testgrapherror(self):
         data = (data_non_inlined[:64 + 27] +
diff --git a/rust/hg-cpython/src/ancestors.rs b/rust/hg-cpython/src/ancestors.rs
--- a/rust/hg-cpython/src/ancestors.rs
+++ b/rust/hg-cpython/src/ancestors.rs
@@ -13,8 +13,8 @@
 };
 use exceptions::GraphError;
 use hg;
-use hg::AncestorsIterator as CoreIterator;
 use hg::Revision;
+use hg::{AncestorsIterator as CoreIterator, LazyAncestors as CoreLazy};
 use std::cell::RefCell;
 
 /// Utility function to convert a Python iterable into a Vec<Revision>
@@ -70,6 +70,37 @@
     }
 }
 
+py_class!(class LazyAncestors |py| {
+    data inner: RefCell<Box<CoreLazy<Index>>>;
+
+    def __contains__(&self, rev: Revision) -> PyResult<bool> {
+        self.inner(py)
+            .borrow_mut()
+            .contains(rev)
+            .map_err(|e| GraphError::pynew(py, e))
+    }
+
+    def __iter__(&self) -> PyResult<AncestorsIterator> {
+        AncestorsIterator::from_inner(py, self.inner(py).borrow().iter())
+    }
+
+    def __bool__(&self) -> PyResult<bool> {
+        Ok(!self.inner(py).borrow().is_empty())
+    }
+
+    def __new__(_cls, index: PyObject, initrevs: PyObject, stoprev: Revision,
+                inclusive: bool) -> PyResult<Self> {
+        let initvec = reviter_to_revvec(py, initrevs)?;
+
+        let lazy =
+            CoreLazy::new(Index::new(py, index)?, initvec, stoprev, inclusive)
+                .map_err(|e| GraphError::pynew(py, e))?;
+
+        Self::create_instance(py, RefCell::new(Box::new(lazy)))
+        }
+
+});
+
 /// Create the module, with __package__ given from parent
 pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
     let dotted_name = &format!("{}.ancestor", package);
@@ -81,6 +112,7 @@
         "Generic DAG ancestor algorithms - Rust implementation",
     )?;
     m.add_class::<AncestorsIterator>(py)?;
+    m.add_class::<LazyAncestors>(py)?;
 
     let sys = PyModule::import(py, "sys")?;
     let sys_modules: PyDict = sys.get(py, "modules")?.extract(py)?;



To: gracinet, #hg-reviewers, kevincox
Cc: yuja, durin42, kevincox, mercurial-devel