diff --git a/rust/hg-cpython/src/dirstate/copymap.rs b/rust/hg-cpython/src/dirstate/copymap.rs --- a/rust/hg-cpython/src/dirstate/copymap.rs +++ b/rust/hg-cpython/src/dirstate/copymap.rs @@ -13,6 +13,7 @@ use std::cell::RefCell; use crate::dirstate::dirstate_map::DirstateMap; use crate::ref_sharing::PyLeakedRef; +use hg::DirstateMap as RustDirstateMap; use hg::{utils::hg_path::HgPathBuf, CopyMapIter}; py_class!(pub class CopyMap |py| { @@ -104,7 +105,7 @@ impl CopyMap { py_shared_iterator!( CopyMapKeysIterator, - PyLeakedRef, + PyLeakedRef<&'static RustDirstateMap>, CopyMapIter<'static>, CopyMap::translate_key, Option @@ -112,7 +113,7 @@ py_shared_iterator!( py_shared_iterator!( CopyMapItemsIterator, - PyLeakedRef, + PyLeakedRef<&'static RustDirstateMap>, CopyMapIter<'static>, CopyMap::translate_key_value, Option<(PyBytes, PyBytes)> diff --git a/rust/hg-cpython/src/dirstate/dirs_multiset.rs b/rust/hg-cpython/src/dirstate/dirs_multiset.rs --- a/rust/hg-cpython/src/dirstate/dirs_multiset.rs +++ b/rust/hg-cpython/src/dirstate/dirs_multiset.rs @@ -92,8 +92,9 @@ py_class!(pub class Dirs |py| { }) } def __iter__(&self) -> PyResult { - let (leak_handle, leaked_ref) = + let mut leak_handle = unsafe { self.inner_shared(py).leak_immutable()? }; + let leaked_ref = leak_handle.data.take().unwrap(); DirsMultisetKeysIterator::from_inner( py, leak_handle, @@ -125,7 +126,7 @@ impl Dirs { py_shared_iterator!( DirsMultisetKeysIterator, - PyLeakedRef, + PyLeakedRef<&'static DirsMultiset>, DirsMultisetIter<'static>, Dirs::translate_key, Option diff --git a/rust/hg-cpython/src/dirstate/dirstate_map.rs b/rust/hg-cpython/src/dirstate/dirstate_map.rs --- a/rust/hg-cpython/src/dirstate/dirstate_map.rs +++ b/rust/hg-cpython/src/dirstate/dirstate_map.rs @@ -304,8 +304,9 @@ py_class!(pub class DirstateMap |py| { } def keys(&self) -> PyResult { - let (leak_handle, leaked_ref) = + let mut leak_handle = unsafe { self.inner_shared(py).leak_immutable()? }; + let leaked_ref = leak_handle.data.take().unwrap(); DirstateMapKeysIterator::from_inner( py, leak_handle, @@ -314,8 +315,9 @@ py_class!(pub class DirstateMap |py| { } def items(&self) -> PyResult { - let (leak_handle, leaked_ref) = + let mut leak_handle = unsafe { self.inner_shared(py).leak_immutable()? }; + let leaked_ref = leak_handle.data.take().unwrap(); DirstateMapItemsIterator::from_inner( py, leak_handle, @@ -324,8 +326,9 @@ py_class!(pub class DirstateMap |py| { } def __iter__(&self) -> PyResult { - let (leak_handle, leaked_ref) = + let mut leak_handle = unsafe { self.inner_shared(py).leak_immutable()? }; + let leaked_ref = leak_handle.data.take().unwrap(); DirstateMapKeysIterator::from_inner( py, leak_handle, @@ -443,8 +446,9 @@ py_class!(pub class DirstateMap |py| { } def copymapiter(&self) -> PyResult { - let (leak_handle, leaked_ref) = + let mut leak_handle = unsafe { self.inner_shared(py).leak_immutable()? }; + let leaked_ref = leak_handle.data.take().unwrap(); CopyMapKeysIterator::from_inner( py, leak_handle, @@ -453,8 +457,9 @@ py_class!(pub class DirstateMap |py| { } def copymapitemsiter(&self) -> PyResult { - let (leak_handle, leaked_ref) = + let mut leak_handle = unsafe { self.inner_shared(py).leak_immutable()? }; + let leaked_ref = leak_handle.data.take().unwrap(); CopyMapItemsIterator::from_inner( py, leak_handle, @@ -493,7 +498,7 @@ py_shared_ref!(DirstateMap, RustDirstate py_shared_iterator!( DirstateMapKeysIterator, - PyLeakedRef, + PyLeakedRef<&'static RustDirstateMap>, StateMapIter<'static>, DirstateMap::translate_key, Option @@ -501,7 +506,7 @@ py_shared_iterator!( py_shared_iterator!( DirstateMapItemsIterator, - PyLeakedRef, + PyLeakedRef<&'static RustDirstateMap>, StateMapIter<'static>, DirstateMap::translate_key_value, Option<(PyBytes, PyObject)> diff --git a/rust/hg-cpython/src/ref_sharing.rs b/rust/hg-cpython/src/ref_sharing.rs --- a/rust/hg-cpython/src/ref_sharing.rs +++ b/rust/hg-cpython/src/ref_sharing.rs @@ -187,23 +187,24 @@ impl<'a, T> PySharedRef<'a, T> { self.data.borrow_mut(self.py) } - /// Returns a leaked reference and its management object. + /// Returns a leaked reference temporarily held by its management object. /// /// # Safety /// /// It's up to you to make sure that the management object lives /// longer than the leaked reference. Otherwise, you'll get a /// dangling reference. - pub unsafe fn leak_immutable( - &self, - ) -> PyResult<(PyLeakedRef, &'static T)> { + pub unsafe fn leak_immutable(&self) -> PyResult> { let (static_ref, static_state_ref) = self .data .py_shared_state .leak_immutable(self.py, self.data)?; - let leak_handle = - PyLeakedRef::new(self.py, self.owner, static_state_ref); - Ok((leak_handle, static_ref)) + Ok(PyLeakedRef::new( + self.py, + self.owner, + static_ref, + static_state_ref, + )) } } @@ -318,12 +319,13 @@ macro_rules! py_shared_ref { /// /// In truth, this does not represent leaked references themselves; /// it is instead useful alongside them to manage them. -pub struct PyLeakedRef { +pub struct PyLeakedRef { _inner: PyObject, + pub data: Option, // TODO: remove pub py_shared_state: &'static PySharedState, } -impl PyLeakedRef { +impl PyLeakedRef { /// # Safety /// /// The `py_shared_state` must be owned by the `inner` Python object. @@ -332,16 +334,18 @@ impl PyLeakedRef { pub unsafe fn new( py: Python, inner: &PyObject, + data: T, py_shared_state: &'static PySharedState, ) -> Self { Self { _inner: inner.clone_ref(py), + data: Some(data), py_shared_state, } } } -impl Drop for PyLeakedRef { +impl Drop for PyLeakedRef { fn drop(&mut self) { // py_shared_state should be alive since we do have // a Python reference to the owner object. Taking GIL makes @@ -379,8 +383,9 @@ impl Drop for PyLeakedRef { /// data inner: PySharedRefCell; /// /// def __iter__(&self) -> PyResult { -/// let (leak_handle, leaked_ref) = +/// let mut leak_handle = /// unsafe { self.inner_shared(py).leak_immutable()? }; +/// let leaked_ref = leak_handle.data.take().unwrap(); /// MyTypeItemsIterator::from_inner( /// py, /// leak_handle, @@ -406,7 +411,7 @@ impl Drop for PyLeakedRef { /// /// py_shared_iterator!( /// MyTypeItemsIterator, -/// PyLeakedRef, +/// PyLeakedRef<&'static MyStruct>, /// HashMap<'static, Vec, Vec>, /// MyType::translate_key_value, /// Option<(PyBytes, PyBytes)>