view rust/hg-cpython/src/discovery.rs @ 51397:b01e7d97e167

revlog: add a Rust implementation of `headrevsdiff`

Python implementation of `headrevsdiff` can be very slow in the worst
case compared with the `heads` computation it replaces, since the
latter is done in Rust.

Even the average case of this Python implementation is still
noticeable in the profiles.

This patch makes the computation much much faster by doing it in Rust.

// 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)
}