rust-cpython: put leaked reference in PyLeakedRef

The next patch will make PyLeakedRef manage the lifetime of the underlying
object. leak_handle.data.take() will be removed soon.

rust/hg-cpython/src/dirstate/copymap.rs

@@ -13,6 +13,7 b' 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 b' impl CopyMap {'
 
 py_shared_iterator!(
     CopyMapKeysIterator,
-    PyLeakedRef,
+    PyLeakedRef<&'static RustDirstateMap>,
     CopyMapIter<'static>,
     CopyMap::translate_key,
     Option<PyBytes>
@@ -112,7 +113,7 b' py_shared_iterator!('
 
 py_shared_iterator!(
     CopyMapItemsIterator,
-    PyLeakedRef,
+    PyLeakedRef<&'static RustDirstateMap>,
     CopyMapIter<'static>,
     CopyMap::translate_key_value,
     Option<(PyBytes, PyBytes)>

rust/hg-cpython/src/dirstate/dirs_multiset.rs

@@ -92,8 +92,9 b' py_class!(pub class Dirs |py| {'
             })
     }
     def __iter__(&self) -> PyResult<DirsMultisetKeysIterator> {
-        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 b' impl Dirs {'
 
 py_shared_iterator!(
     DirsMultisetKeysIterator,
-    PyLeakedRef,
+    PyLeakedRef<&'static DirsMultiset>,
     DirsMultisetIter<'static>,
     Dirs::translate_key,
     Option<PyBytes>

rust/hg-cpython/src/dirstate/dirstate_map.rs

@@ -304,8 +304,9 b' py_class!(pub class DirstateMap |py| {'
     }
 
     def keys(&self) -> PyResult<DirstateMapKeysIterator> {
-        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 b' py_class!(pub class DirstateMap |py| {'
     }
 
     def items(&self) -> PyResult<DirstateMapItemsIterator> {
-        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 b' py_class!(pub class DirstateMap |py| {'
     }
 
     def __iter__(&self) -> PyResult<DirstateMapKeysIterator> {
-        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 b' py_class!(pub class DirstateMap |py| {'
     }
 
     def copymapiter(&self) -> PyResult<CopyMapKeysIterator> {
-        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 b' py_class!(pub class DirstateMap |py| {'
     }
 
     def copymapitemsiter(&self) -> PyResult<CopyMapItemsIterator> {
-        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 b' py_shared_ref!(DirstateMap, RustDirstate'
 
 py_shared_iterator!(
     DirstateMapKeysIterator,
-    PyLeakedRef,
+    PyLeakedRef<&'static RustDirstateMap>,
     StateMapIter<'static>,
     DirstateMap::translate_key,
     Option<PyBytes>
@@ -501,7 +506,7 b' py_shared_iterator!('
 
 py_shared_iterator!(
     DirstateMapItemsIterator,
-    PyLeakedRef,
+    PyLeakedRef<&'static RustDirstateMap>,
     StateMapIter<'static>,
     DirstateMap::translate_key_value,
     Option<(PyBytes, PyObject)>

rust/hg-cpython/src/ref_sharing.rs

@@ -187,23 +187,24 b" 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<PyLeakedRef<&'static T>> {
         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 b' 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<T> {
     _inner: PyObject,
+    pub data: Option<T>, // TODO: remove pub
     py_shared_state: &'static PySharedState,
 }
 
-impl PyLeakedRef {
+impl<T> PyLeakedRef<T> {
     /// # Safety
     ///
     /// The `py_shared_state` must be owned by the `inner` Python object.
@@ -332,16 +334,18 b' 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<T> Drop for PyLeakedRef<T> {
     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 b' impl Drop for PyLeakedRef {'
 ///     data inner: PySharedRefCell<MyStruct>;
 ///
 ///     def __iter__(&self) -> PyResult<MyTypeItemsIterator> {
-///         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 b' impl Drop for PyLeakedRef {'
 ///
 /// py_shared_iterator!(
 ///     MyTypeItemsIterator,
-///     PyLeakedRef,
+///     PyLeakedRef<&'static MyStruct>,
 ///     HashMap<'static, Vec<u8>, Vec<u8>>,
 ///     MyType::translate_key_value,
 ///     Option<(PyBytes, PyBytes)>