D5584: rust: MissingAncestors.basesheads()

[gracinet (Georges Racinet)]

gracinet created this revision.
Herald added subscribers: mercurial-devel, kevincox, durin42.
Herald added a reviewer: hg-reviewers.

REVISION SUMMARY
  This new API method on `MissingAncestors` leverages directly the
  Rust implementation for relative heads of a set, and also
  lowers the cost of returning the results to Python in the context of
  discovery.
  
  These interchange costs can probably be further reduced by implementing
  the `partialdiscovery` class in Rust, but that will be investigated in the
  5.0 development cycle.

REPOSITORY
  rHG Mercurial

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

AFFECTED FILES
  rust/hg-core/src/ancestors.rs
  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
@@ -114,6 +114,7 @@
         missanc.addbases({2})
         self.assertEqual(missanc.bases(), {1, 2})
         self.assertEqual(missanc.missingancestors([3]), [3])
+        self.assertEqual(missanc.basesheads(), {2})
 
     def testmissingancestorsremove(self):
         idx = self.parseindex()
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
@@ -142,7 +142,10 @@
 py_class!(pub class MissingAncestors |py| {
     data inner: RefCell<Box<CoreMissing<Index>>>;
 
-    def __new__(_cls, index: PyObject, bases: PyObject) -> PyResult<MissingAncestors> {
+    def __new__(_cls,
+                index: PyObject,
+                bases: PyObject)
+    -> PyResult<MissingAncestors> {
         let bases_vec: Vec<Revision> = rev_pyiter_collect(py, &bases)?;
         let inner = CoreMissing::new(Index::new(py, index)?, bases_vec);
         MissingAncestors::create_instance(py, RefCell::new(Box::new(inner)))
@@ -166,6 +169,11 @@
         py_set(py, self.inner(py).borrow().get_bases())
     }
 
+    def basesheads(&self) -> PyResult<PyObject> {
+        let inner = self.inner(py).borrow();
+        py_set(py, &inner.bases_heads().map_err(|e| GraphError::pynew(py, e))?)
+    }
+
     def removeancestorsfrom(&self, revs: PyObject) -> PyResult<PyObject> {
         let mut inner = self.inner(py).borrow_mut();
         // this is very lame: we convert to a Rust set, update it in place
diff --git a/rust/hg-core/src/ancestors.rs b/rust/hg-core/src/ancestors.rs
--- a/rust/hg-core/src/ancestors.rs
+++ b/rust/hg-core/src/ancestors.rs
@@ -10,6 +10,7 @@
 use super::{Graph, GraphError, Revision, NULL_REVISION};
 use std::cmp::max;
 use std::collections::{BinaryHeap, HashSet};
+use crate::dagops;
 
 /// Iterator over the ancestors of a given list of revisions
 /// This is a generic type, defined and implemented for any Graph, so that
@@ -229,6 +230,19 @@
         &self.bases
     }
 
+    /// Computes the relative heads of current bases.
+    ///
+    /// The object is still usable after this.
+    pub fn bases_heads(&self) -> Result<HashSet<Revision>, GraphError> {
+        dagops::heads(&self.graph, self.bases.iter())
+    }
+
+    /// Consumes the object and returns the relative heads of its bases.
+    pub fn into_bases_heads(mut self) -> Result<HashSet<Revision>, GraphError> {
+        dagops::retain_heads(&self.graph, &mut self.bases)?;
+        Ok(self.bases)
+    }
+
     pub fn add_bases(
         &mut self,
         new_bases: impl IntoIterator<Item = Revision>,
@@ -556,8 +570,8 @@
     }
 
     #[test]
-    /// Test constructor, add/get bases
-    fn test_missing_bases() {
+    /// Test constructor, add/get bases and heads
+    fn test_missing_bases() -> Result<(), GraphError> {
         let mut missing_ancestors =
             MissingAncestors::new(SampleGraph, [5, 3, 1, 3].iter().cloned());
         let mut as_vec: Vec<Revision> =
@@ -569,6 +583,11 @@
         as_vec = missing_ancestors.get_bases().iter().cloned().collect();
         as_vec.sort();
         assert_eq!(as_vec, [1, 3, 5, 7, 8]);
+
+        as_vec = missing_ancestors.bases_heads()?.iter().cloned().collect();
+        as_vec.sort();
+        assert_eq!(as_vec, [3, 5, 7, 8]);
+        Ok(())
     }
 
     fn assert_missing_remove(



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