diff --git a/rust/hg-pyo3/src/convert_cpython.rs b/rust/hg-pyo3/src/convert_cpython.rs --- a/rust/hg-pyo3/src/convert_cpython.rs +++ b/rust/hg-pyo3/src/convert_cpython.rs @@ -16,6 +16,7 @@ use cpython::ObjectProtocol; use cpython::PythonObject; use lazy_static::lazy_static; +use hg::revlog::index::Index as CoreIndex; use rusthg::revlog::{InnerRevlog, PySharedIndex}; /// Force cpython's GIL handle with the appropriate lifetime @@ -165,9 +166,21 @@ pub fn py_rust_index_to_graph( Ok(unsafe { leaked.map(py, |idx| PySharedIndex { inner: &idx.index }) }) } -pub(crate) fn proxy_index_extract<'py>( +/// Full extraction of the proxy index object as received in PyO3 to a +/// [`CoreIndex`] reference. +/// +/// The safety invariants to maintain are those of the underlying +/// [`UnsafePyLeaked::try_borrow`]: the caller must not leak the inner +/// reference. +pub(crate) unsafe fn proxy_index_extract<'py>( index_proxy: &Bound<'py, PyAny>, -) -> PyResult<(cpython::Python<'py>, cpython::UnsafePyLeaked)> { +) -> PyResult<&'py CoreIndex> { let (py, idx_proxy) = to_cpython_py_object(index_proxy); - Ok((py, py_rust_index_to_graph(py, idx_proxy)?)) + let py_leaked = py_rust_index_to_graph(py, idx_proxy)?; + let py_shared = &*unsafe { + py_leaked + .try_borrow(py) + .map_err(|e| from_cpython_pyerr(py, e))? + }; + Ok(py_shared.inner) } diff --git a/rust/hg-pyo3/src/dagops.rs b/rust/hg-pyo3/src/dagops.rs --- a/rust/hg-pyo3/src/dagops.rs +++ b/rust/hg-pyo3/src/dagops.rs @@ -15,7 +15,7 @@ use std::collections::HashSet; use hg::{dagops, Revision}; -use crate::convert_cpython::{from_cpython_pyerr, proxy_index_extract}; +use crate::convert_cpython::proxy_index_extract; use crate::exceptions::GraphError; use crate::revision::{rev_pyiter_collect, PyRevision}; use crate::util::new_submodule; @@ -29,14 +29,8 @@ pub fn headrevs( index_proxy: &Bound<'_, PyAny>, revs: &Bound<'_, PyAny>, ) -> PyResult> { - let (py, py_leaked) = proxy_index_extract(index_proxy)?; // Safety: we don't leak the "faked" reference out of `UnsafePyLeaked` - let index = &*unsafe { - py_leaked - .try_borrow(py) - .map_err(|e| from_cpython_pyerr(py, e))? - }; - + let index = unsafe { proxy_index_extract(index_proxy)? }; let mut as_set: HashSet = rev_pyiter_collect(revs, index)?; dagops::retain_heads(index, &mut as_set).map_err(GraphError::from_hg)?; Ok(as_set.into_iter().map(Into::into).collect())