# HG changeset patch
# User Georges Racinet <georges.racinet@cloudcrane.io>
# Date 2024-11-30 19:58:09
# Node ID 20fe0bf9a9a599f012ec2912b74412c0e2ef7ff5
# Parent  6e8ba52857df2850d84ccc5a1af21577322c19bc

rust-pyo3: dagop submodule implementation

This is the first demonstration that the passing of objects through the
Python interpreter from `rustext` to `pyo3_rustext` actually works.

In the Python tests, we conflate the presence of the `pyo3_rustext` package
with that of `rustext`, as we do not plan to support building one and not
the other (we hope to convert fully to PyO3 soon).
The skipping is actually done by the base test class.

The implementation of `rank` is as trivial as it was in `hg-cpython`,
yet it required to be able to convert exceptions from `vcsgraph`.

diff --git a/rust/hg-pyo3/src/dagops.rs b/rust/hg-pyo3/src/dagops.rs
--- a/rust/hg-pyo3/src/dagops.rs
+++ b/rust/hg-pyo3/src/dagops.rs
@@ -11,12 +11,63 @@
 //! From Python, this will be seen as `mercurial.pyo3-rustext.dagop`
 use pyo3::prelude::*;
 
+use std::collections::HashSet;
+
+use hg::{dagops, Revision};
+
+use crate::convert_cpython::{from_cpython_pyerr, proxy_index_extract};
+use crate::exceptions::GraphError;
+use crate::revision::{rev_pyiter_collect, PyRevision};
 use crate::util::new_submodule;
 
+/// Using the the `index_proxy`, return heads out of any Python iterable of
+/// Revisions
+///
+/// This is the Rust counterpart for `mercurial.dagop.headrevs`
+#[pyfunction]
+pub fn headrevs(
+    index_proxy: &Bound<'_, PyAny>,
+    revs: &Bound<'_, PyAny>,
+) -> PyResult<HashSet<PyRevision>> {
+    let (py, py_leaked) = proxy_index_extract(index_proxy)?;
+    // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
+    let index = &*unsafe {
+        py_leaked
+            .try_borrow(py)
+            .map_err(|e| from_cpython_pyerr(py, e))?
+    };
+
+    let mut as_set: HashSet<Revision> = rev_pyiter_collect(revs, index)?;
+    dagops::retain_heads(index, &mut as_set).map_err(GraphError::from_hg)?;
+    Ok(as_set.into_iter().map(Into::into).collect())
+}
+
+/// Computes the rank, i.e. the number of ancestors including itself,
+/// of a node represented by its parents.
+///
+/// Currently, the pure Rust index supports only the REVLOGV1 format, hence
+/// the only possible return value is that the rank is unknown.
+///
+/// References:
+/// - C implementation, function `index_fast_rank()`.
+/// - `impl vcsgraph::graph::RankedGraph for Index` in `crate::cindex`.
+#[pyfunction]
+pub fn rank(
+    _index: &Bound<'_, PyAny>,
+    _p1r: PyRevision,
+    _p2r: PyRevision,
+) -> PyResult<()> {
+    Err(GraphError::from_vcsgraph(
+        vcsgraph::graph::GraphReadError::InconsistentGraphData,
+    ))
+}
+
 pub fn init_module<'py>(
     py: Python<'py>,
     package: &str,
 ) -> PyResult<Bound<'py, PyModule>> {
     let m = new_submodule(py, package, "dagop")?;
+    m.add_function(wrap_pyfunction!(headrevs, &m)?)?;
+    m.add_function(wrap_pyfunction!(rank, &m)?)?;
     Ok(m)
 }
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
@@ -5,11 +5,12 @@ from mercurial.node import wdirrev
 from mercurial.testing import revlog as revlogtesting
 
 try:
-    from mercurial import rustext
+    from mercurial import pyo3_rustext, rustext
 
     rustext.__name__  # trigger immediate actual import
+    pyo3_rustext.__name__
 except ImportError:
-    rustext = None
+    rustext = pyo3_rustext = None
 else:
     # this would fail already without appropriate ancestor.__package__
     from mercurial.rustext.ancestor import (
@@ -148,6 +149,17 @@ class rustancestorstest(revlogtesting.Ru
         idx = self.parserustindex()
         self.assertEqual(dagop.headrevs(idx, [1, 2, 3]), {3})
 
+    def testpyo3_headrevs(self):
+        idx = self.parserustindex()
+        self.assertEqual(pyo3_rustext.dagop.headrevs(idx, [1, 2, 3]), {3})
+
+    def testpyo3_rank(self):
+        idx = self.parserustindex()
+        try:
+            pyo3_rustext.dagop.rank(idx, 1, 2)
+        except pyo3_rustext.GraphError as exc:
+            self.assertEqual(exc.args, ("InconsistentGraphData",))
+
 
 if __name__ == '__main__':
     import silenttestrunner