rust/hg-cpython/src/discovery.rs
author Pierre-Yves David <pierre-yves.david@octobus.net>
Fri, 08 Nov 2024 17:08:11 +0100
branchstable
changeset 52216 fa58f4f97337
parent 51255 24d3298189d7
permissions -rw-r--r--
ci: shard the test run on mac os X This should comes with some benefit: - spread the load across more runner, - reduce the real-time CI run, - reduce the "retry" run when we need them. We start with the Mac jobs, but that would be tremendously useful for Windows too. For linux, we need to reduce the startup overhead for this to be worth it. Building smaller image and speeding up clone should help with that.

// discovery.rs
//
// Copyright 2018 Georges Racinet <gracinet@anybox.fr>
//
// This software may be used and distributed according to the terms of the
// GNU General Public License version 2 or any later version.

//! Bindings for the `hg::discovery` module provided by the
//! `hg-core` crate. From Python, this will be seen as `rustext.discovery`
//!
//! # Classes visible from Python:
//! - [`PartialDiscovery`] is the Rust implementation of
//!   `mercurial.setdiscovery.partialdiscovery`.

use crate::PyRevision;
use crate::{
    conversion::rev_pyiter_collect, exceptions::GraphError,
    revlog::PySharedIndex,
};
use cpython::{
    ObjectProtocol, PyClone, PyDict, PyModule, PyObject, PyResult, PyTuple,
    Python, PythonObject, ToPyObject, UnsafePyLeaked,
};
use hg::discovery::PartialDiscovery as CorePartialDiscovery;
use hg::Revision;
use std::collections::HashSet;

use std::cell::RefCell;

use crate::revlog::py_rust_index_to_graph;

py_class!(pub class PartialDiscovery |py| {
    data inner: RefCell<UnsafePyLeaked<CorePartialDiscovery<PySharedIndex>>>;
    data index: RefCell<UnsafePyLeaked<PySharedIndex>>;

    // `_respectsize` is currently only here to replicate the Python API and
    // will be used in future patches inside methods that are yet to be
    // implemented.
    def __new__(
        _cls,
        repo: PyObject,
        targetheads: PyObject,
        respectsize: bool,
        randomize: bool = true
    ) -> PyResult<PartialDiscovery> {
        Self::inner_new(py, repo, targetheads, respectsize, randomize)
    }

    def addcommons(&self, commons: PyObject) -> PyResult<PyObject> {
        self.inner_addcommons(py, commons)
    }

    def addmissings(&self, missings: PyObject) -> PyResult<PyObject> {
        self.inner_addmissings(py, missings)
    }

    def addinfo(&self, sample: PyObject) -> PyResult<PyObject> {
        self.inner_addinfo(py, sample)
    }

    def hasinfo(&self) -> PyResult<bool> {
        let leaked = self.inner(py).borrow();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let inner = unsafe { leaked.try_borrow(py)? };
        Ok(inner.has_info())
    }

    def iscomplete(&self) -> PyResult<bool> {
        let leaked = self.inner(py).borrow();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let inner = unsafe { leaked.try_borrow(py)? };
        Ok(inner.is_complete())
    }

    def stats(&self) -> PyResult<PyDict> {
        let leaked = self.inner(py).borrow();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let inner = unsafe { leaked.try_borrow(py)? };
        let stats = inner.stats();
        let as_dict: PyDict = PyDict::new(py);
        as_dict.set_item(py, "undecided",
                         stats.undecided.map(
                             |l| l.to_py_object(py).into_object())
                             .unwrap_or_else(|| py.None()))?;
        Ok(as_dict)
    }

    def commonheads(&self) -> PyResult<HashSet<PyRevision>> {
        let leaked = self.inner(py).borrow();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let inner = unsafe { leaked.try_borrow(py)? };
        let res = inner.common_heads()
                    .map_err(|e| GraphError::pynew(py, e))?;
        Ok(res.into_iter().map(Into::into).collect())
    }

    def takefullsample(&self, headrevs: PyObject,
                       size: usize) -> PyResult<PyObject> {
        self.inner_takefullsample(py, headrevs, size)
    }

    def takequicksample(&self, headrevs: PyObject,
                        size: usize) -> PyResult<PyObject> {
        self.inner_takequicksample(py, headrevs, size)
    }

});

impl PartialDiscovery {
    fn inner_new(
        py: Python,
        repo: PyObject,
        targetheads: PyObject,
        respectsize: bool,
        randomize: bool,
    ) -> PyResult<Self> {
        let index = repo.getattr(py, "changelog")?.getattr(py, "index")?;
        let cloned_index = py_rust_index_to_graph(py, index.clone_ref(py))?;
        let index = py_rust_index_to_graph(py, index)?;
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let target_heads = {
            let borrowed_idx = unsafe { index.try_borrow(py)? };
            rev_pyiter_collect(py, &targetheads, &*borrowed_idx)?
        };
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let lazy_disco = unsafe {
            index.map(py, |idx| {
                CorePartialDiscovery::new(
                    idx,
                    target_heads,
                    respectsize,
                    randomize,
                )
            })
        };
        Self::create_instance(
            py,
            RefCell::new(lazy_disco),
            RefCell::new(cloned_index),
        )
    }

    /// Convert a Python iterator of revisions into a vector
    fn pyiter_to_vec(
        &self,
        py: Python,
        iter: &PyObject,
    ) -> PyResult<Vec<Revision>> {
        let leaked = self.index(py).borrow();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let index = unsafe { leaked.try_borrow(py)? };
        rev_pyiter_collect(py, iter, &*index)
    }

    fn inner_addinfo(
        &self,
        py: Python,
        sample: PyObject,
    ) -> PyResult<PyObject> {
        let mut missing: Vec<Revision> = Vec::new();
        let mut common: Vec<Revision> = Vec::new();
        for info in sample.iter(py)? {
            // info is a pair (Revision, bool)
            let mut revknown = info?.iter(py)?;
            let rev: PyRevision = revknown.next().unwrap()?.extract(py)?;
            // This is fine since we're just using revisions as integers
            // for the purposes of discovery
            let rev = Revision(rev.0);
            let known: bool = revknown.next().unwrap()?.extract(py)?;
            if known {
                common.push(rev);
            } else {
                missing.push(rev);
            }
        }
        let mut leaked = self.inner(py).borrow_mut();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let mut inner = unsafe { leaked.try_borrow_mut(py)? };
        inner
            .add_common_revisions(common)
            .map_err(|e| GraphError::pynew(py, e))?;
        inner
            .add_missing_revisions(missing)
            .map_err(|e| GraphError::pynew(py, e))?;
        Ok(py.None())
    }

    fn inner_addcommons(
        &self,
        py: Python,
        commons: PyObject,
    ) -> PyResult<PyObject> {
        let commons_vec = self.pyiter_to_vec(py, &commons)?;
        let mut leaked = self.inner(py).borrow_mut();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let mut inner = unsafe { leaked.try_borrow_mut(py)? };
        inner
            .add_common_revisions(commons_vec)
            .map_err(|e| GraphError::pynew(py, e))?;
        Ok(py.None())
    }

    fn inner_addmissings(
        &self,
        py: Python,
        missings: PyObject,
    ) -> PyResult<PyObject> {
        let missings_vec = self.pyiter_to_vec(py, &missings)?;
        let mut leaked = self.inner(py).borrow_mut();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let mut inner = unsafe { leaked.try_borrow_mut(py)? };
        inner
            .add_missing_revisions(missings_vec)
            .map_err(|e| GraphError::pynew(py, e))?;
        Ok(py.None())
    }

    fn inner_takefullsample(
        &self,
        py: Python,
        _headrevs: PyObject,
        size: usize,
    ) -> PyResult<PyObject> {
        let mut leaked = self.inner(py).borrow_mut();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let mut inner = unsafe { leaked.try_borrow_mut(py)? };
        let sample = inner
            .take_full_sample(size)
            .map_err(|e| GraphError::pynew(py, e))?;
        let as_vec: Vec<PyObject> = sample
            .iter()
            .map(|rev| PyRevision(rev.0).to_py_object(py).into_object())
            .collect();
        Ok(PyTuple::new(py, as_vec.as_slice()).into_object())
    }

    fn inner_takequicksample(
        &self,
        py: Python,
        headrevs: PyObject,
        size: usize,
    ) -> PyResult<PyObject> {
        let revsvec = self.pyiter_to_vec(py, &headrevs)?;
        let mut leaked = self.inner(py).borrow_mut();
        // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked`
        let mut inner = unsafe { leaked.try_borrow_mut(py)? };
        let sample = inner
            .take_quick_sample(revsvec, size)
            .map_err(|e| GraphError::pynew(py, e))?;
        let as_vec: Vec<PyObject> = sample
            .iter()
            .map(|rev| PyRevision(rev.0).to_py_object(py).into_object())
            .collect();
        Ok(PyTuple::new(py, as_vec.as_slice()).into_object())
    }
}

/// Create the module, with __package__ given from parent
pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let dotted_name = &format!("{}.discovery", package);
    let m = PyModule::new(py, dotted_name)?;
    m.add(py, "__package__", package)?;
    m.add(
        py,
        "__doc__",
        "Discovery of common node sets - Rust implementation",
    )?;
    m.add_class::<PartialDiscovery>(py)?;

    let sys = PyModule::import(py, "sys")?;
    let sys_modules: PyDict = sys.get(py, "modules")?.extract(py)?;
    sys_modules.set_item(py, dotted_name, &m)?;
    // Example C code (see pyexpat.c and import.c) will "give away the
    // reference", but we won't because it will be consumed once the
    // Rust PyObject is dropped.
    Ok(m)
}