rust/hg-pyo3/src/dagops.rs
changeset 52413 20fe0bf9a9a5
parent 52407 c5128c541021
child 52414 233707101dae
--- 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)
 }