rust/hg-cpython/src/discovery.rs

// discovery.rs
//
// Copyright 2018 Georges Racinet <[e-mail]>
//
// 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)
}