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> { + 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 = 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> { 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