rust/hg-cpython/src/utils.rs
author Arseniy Alekseyev <aalekseyev@janestreet.com>
Mon, 12 Feb 2024 20:01:27 +0000
changeset 51394 b01e7d97e167
parent 49631 c7fb9b74e753
child 52042 28a0eb21ff04
permissions -rw-r--r--
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.

use cpython::exc::ValueError;
use cpython::{PyBytes, PyDict, PyErr, PyObject, PyResult, PyTuple, Python};
use hg::revlog::Node;

#[allow(unused)]
pub fn print_python_trace(py: Python) -> PyResult<PyObject> {
    eprintln!("===============================");
    eprintln!("Printing Python stack from Rust");
    eprintln!("===============================");
    let traceback = py.import("traceback")?;
    let sys = py.import("sys")?;
    let kwargs = PyDict::new(py);
    kwargs.set_item(py, "file", sys.get(py, "stderr")?)?;
    traceback.call(py, "print_stack", PyTuple::new(py, &[]), Some(&kwargs))
}

// Necessary evil for the time being, could maybe be moved to
// a TryFrom in Node itself
const NODE_BYTES_LENGTH: usize = 20;
type NodeData = [u8; NODE_BYTES_LENGTH];

/// Copy incoming Python bytes given as `PyObject` into `Node`,
/// doing the necessary checks
pub fn node_from_py_object<'a>(
    py: Python,
    bytes: &'a PyObject,
) -> PyResult<Node> {
    let as_py_bytes: &'a PyBytes = bytes.extract(py)?;
    node_from_py_bytes(py, as_py_bytes)
}

/// Clone incoming Python bytes given as `PyBytes` as a `Node`,
/// doing the necessary checks.
pub fn node_from_py_bytes(py: Python, bytes: &PyBytes) -> PyResult<Node> {
    <NodeData>::try_from(bytes.data(py))
        .map_err(|_| {
            PyErr::new::<ValueError, _>(
                py,
                format!("{}-byte hash required", NODE_BYTES_LENGTH),
            )
        })
        .map(Into::into)
}