--- a/rust/hg-pyo3/src/dagops.rs Thu Dec 05 14:36:40 2024 +0100
+++ b/rust/hg-pyo3/src/dagops.rs Sat Nov 30 20:58:09 2024 +0100
@@ -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)
}