--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/rust/hg-pyo3/src/discovery.rs	Mon Jan 06 01:14:52 2025 +0100
@@ -0,0 +1,216 @@
+//! Discovery of common node sets
+use std::collections::HashSet;
+
+use hg::{discovery::PartialDiscovery as CorePartialDiscovery, Revision};
+use pyo3::{
+    intern, pyclass, pymethods,
+    types::{PyAnyMethods, PyDict, PyModule, PyModuleMethods, PyTuple},
+    Bound, Py, PyAny, PyObject, PyResult, Python,
+};
+use pyo3_sharedref::SharedByPyObject;
+
+use crate::{
+    exceptions::GraphError,
+    revision::{rev_pyiter_collect, PyRevision},
+    revlog::PySharedIndex,
+    utils::{new_submodule, py_rust_index_to_graph},
+};
+
+#[pyclass]
+struct PartialDiscovery {
+    inner: SharedByPyObject<CorePartialDiscovery<PySharedIndex>>,
+    idx: SharedByPyObject<PySharedIndex>,
+}
+
+#[pymethods]
+impl PartialDiscovery {
+    #[pyo3(signature = (repo, targetheads, respectsize, randomize=true))]
+    #[new]
+    fn new(
+        py: Python,
+        repo: &Bound<'_, PyAny>,
+        targetheads: &Bound<'_, PyAny>,
+        respectsize: bool,
+        randomize: bool,
+    ) -> PyResult<Self> {
+        let index = repo
+            .getattr(intern!(py, "changelog"))?
+            .getattr(intern!(py, "index"))?;
+        let cloned_index = py_rust_index_to_graph(&index.clone())?;
+        let index = py_rust_index_to_graph(&index)?;
+
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let target_heads = {
+            let borrowed_idx = unsafe { index.try_borrow(py)? };
+            rev_pyiter_collect(targetheads, &*borrowed_idx)?
+        };
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let lazy_disco = unsafe {
+            index.map(py, |idx| {
+                CorePartialDiscovery::new(
+                    idx,
+                    target_heads,
+                    respectsize,
+                    randomize,
+                )
+            })
+        };
+        Ok(Self {
+            inner: lazy_disco,
+            idx: cloned_index,
+        })
+    }
+
+    fn addcommons(
+        &mut self,
+        py: Python,
+        commons: &Bound<'_, PyAny>,
+    ) -> PyResult<PyObject> {
+        let commons = self.pyiter_to_vec(commons)?;
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };
+        inner
+            .add_common_revisions(commons)
+            .map_err(GraphError::from_hg)?;
+        Ok(py.None())
+    }
+
+    fn addmissings(
+        &mut self,
+        py: Python,
+        missings: &Bound<'_, PyAny>,
+    ) -> PyResult<PyObject> {
+        let missings = self.pyiter_to_vec(missings)?;
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };
+        inner
+            .add_missing_revisions(missings)
+            .map_err(GraphError::from_hg)?;
+        Ok(py.None())
+    }
+
+    fn addinfo(
+        &mut self,
+        py: Python,
+        sample: &Bound<'_, PyAny>,
+    ) -> PyResult<PyObject> {
+        let mut missing: Vec<Revision> = vec![];
+        let mut common: Vec<Revision> = vec![];
+        for info in sample.try_iter()? {
+            // info is a pair (Revision, bool)
+            let info = info?;
+            let info = info.downcast::<PyTuple>()?;
+            let rev: PyRevision = info.get_item(0)?.extract()?;
+            // This is fine since we're just using revisions as integers
+            // for the purposes of discovery
+            let rev = Revision(rev.0);
+            let known: bool = info.get_item(1)?.extract()?;
+            if known {
+                common.push(rev);
+            } else {
+                missing.push(rev);
+            }
+        }
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };
+        inner
+            .add_common_revisions(common)
+            .map_err(GraphError::from_hg)?;
+        inner
+            .add_missing_revisions(missing)
+            .map_err(GraphError::from_hg)?;
+        Ok(py.None())
+    }
+
+    fn hasinfo(&self, py: Python<'_>) -> PyResult<bool> {
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let inner = unsafe { self.inner.try_borrow(py)? };
+        Ok(inner.has_info())
+    }
+
+    fn iscomplete(&self, py: Python<'_>) -> PyResult<bool> {
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let inner = unsafe { self.inner.try_borrow(py)? };
+        Ok(inner.is_complete())
+    }
+
+    fn stats(&self, py: Python<'_>) -> PyResult<Py<PyDict>> {
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let inner = unsafe { self.inner.try_borrow(py)? };
+        let stats = inner.stats();
+        let as_dict = PyDict::new(py);
+        as_dict.set_item("undecided", stats.undecided)?;
+        Ok(as_dict.unbind())
+    }
+
+    fn commonheads(&self, py: Python<'_>) -> PyResult<HashSet<PyRevision>> {
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let inner = unsafe { self.inner.try_borrow(py)? };
+        let common_heads =
+            inner.common_heads().map_err(GraphError::from_hg)?;
+        Ok(common_heads.into_iter().map(Into::into).collect())
+    }
+
+    fn takefullsample(
+        &mut self,
+        py: Python,
+        _headrevs: &Bound<'_, PyAny>,
+        size: usize,
+    ) -> PyResult<Py<PyTuple>> {
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };
+        let sample =
+            inner.take_full_sample(size).map_err(GraphError::from_hg)?;
+        let as_pyrevision = sample.into_iter().map(|rev| PyRevision(rev.0));
+        Ok(PyTuple::new(py, as_pyrevision)?.unbind())
+    }
+
+    fn takequicksample(
+        &mut self,
+        py: Python,
+        headrevs: &Bound<'_, PyAny>,
+        size: usize,
+    ) -> PyResult<Py<PyTuple>> {
+        let revs = self.pyiter_to_vec(headrevs)?;
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let mut inner = unsafe { self.inner.try_borrow_mut(py)? };
+        let sample = inner
+            .take_quick_sample(revs, size)
+            .map_err(GraphError::from_hg)?;
+        let as_pyrevision = sample.into_iter().map(|rev| PyRevision(rev.0));
+        Ok(PyTuple::new(py, as_pyrevision)?.unbind())
+    }
+}
+
+impl PartialDiscovery {
+    /// Convert a Python iterator of revisions into a vector
+    fn pyiter_to_vec(
+        &self,
+        iter: &Bound<'_, PyAny>,
+    ) -> PyResult<Vec<Revision>> {
+        // Safety: we don't leak any reference derived form the "faked" one in
+        // `SharedByPyObject`
+        let index = unsafe { self.idx.try_borrow(iter.py())? };
+        rev_pyiter_collect(iter, &*index)
+    }
+}
+
+pub fn init_module<'py>(
+    py: Python<'py>,
+    package: &str,
+) -> PyResult<Bound<'py, PyModule>> {
+    let m = new_submodule(py, package, "discovery")?;
+    m.add_class::<PartialDiscovery>()?;
+    Ok(m)
+}

--- a/rust/hg-pyo3/src/lib.rs	Sun Jan 05 23:33:24 2025 +0100
+++ b/rust/hg-pyo3/src/lib.rs	Mon Jan 06 01:14:52 2025 +0100
@@ -3,6 +3,7 @@
 mod ancestors;
 mod convert_cpython;
 mod dagops;
+mod discovery;
 mod exceptions;
 mod node;
 mod revision;
@@ -23,6 +24,7 @@
 
     m.add_submodule(&ancestors::init_module(py, &dotted_name)?)?;
     m.add_submodule(&dagops::init_module(py, &dotted_name)?)?;
+    m.add_submodule(&discovery::init_module(py, &dotted_name)?)?;
     m.add_submodule(&revlog::init_module(py, &dotted_name)?)?;
     m.add("GraphError", py.get_type::<exceptions::GraphError>())?;
     Ok(())