Patchwork [1,of,6,V3] rust-cpython: put leaked reference in PyLeakedRef

login
register
mail settings
Submitter Yuya Nishihara
Date Oct. 17, 2019, 1:06 p.m.
Message ID <30ad653deb7fbe2196f3.1571317592@mimosa>
Download mbox | patch
Permalink /patch/42451/
State New
Headers show

Comments

Yuya Nishihara - Oct. 17, 2019, 1:06 p.m.
# HG changeset patch
# User Yuya Nishihara <yuya@tcha.org>
# Date 1568552779 -32400
#      Sun Sep 15 22:06:19 2019 +0900
# Node ID 30ad653deb7fbe2196f3d2d9d97f9f650c8b7952
# Parent  2c70dd03b74395fbec276a2e2d1d3b8347268a5d
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.
via Mercurial-devel - Oct. 18, 2019, 6:23 p.m.
Queuing this, thanks.

On Thu, Oct 17, 2019 at 6:14 AM Yuya Nishihara <yuya@tcha.org> wrote:

> # HG changeset patch
> # User Yuya Nishihara <yuya@tcha.org>
> # Date 1568552779 -32400
> #      Sun Sep 15 22:06:19 2019 +0900
> # Node ID 30ad653deb7fbe2196f3d2d9d97f9f650c8b7952
> # Parent  2c70dd03b74395fbec276a2e2d1d3b8347268a5d
> 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.
>
> 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<PyBytes>
> @@ -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<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 @@ impl Dirs {
>
>  py_shared_iterator!(
>      DirsMultisetKeysIterator,
> -    PyLeakedRef,
> +    PyLeakedRef<&'static DirsMultiset>,
>      DirsMultisetIter<'static>,
>      Dirs::translate_key,
>      Option<PyBytes>
> 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<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 @@ 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 @@ 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 @@ 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 @@ 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 @@ py_shared_ref!(DirstateMap, RustDirstate
>
>  py_shared_iterator!(
>      DirstateMapKeysIterator,
> -    PyLeakedRef,
> +    PyLeakedRef<&'static RustDirstateMap>,
>      StateMapIter<'static>,
>      DirstateMap::translate_key,
>      Option<PyBytes>
> @@ -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<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 @@ 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 @@ 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 @@ 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 @@ 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)>
> _______________________________________________
> Mercurial-devel mailing list
> Mercurial-devel@mercurial-scm.org
> https://www.mercurial-scm.org/mailman/listinfo/mercurial-devel
>
Raphaël Gomès - Oct. 19, 2019, 7:41 p.m.
Great work on this series, thanks.

I'll admit it was a little hard to follow all of the intermediary steps 
needed, but it's better that way for debugging I think.

I'll jump on the follow-up.

On 10/17/19 3:06 PM, Yuya Nishihara wrote:
> # HG changeset patch
> # User Yuya Nishihara <yuya@tcha.org>
> # Date 1568552779 -32400
> #      Sun Sep 15 22:06:19 2019 +0900
> # Node ID 30ad653deb7fbe2196f3d2d9d97f9f650c8b7952
> # Parent  2c70dd03b74395fbec276a2e2d1d3b8347268a5d
> 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.
>
> 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<PyBytes>
> @@ -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<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 @@ impl Dirs {
>   
>   py_shared_iterator!(
>       DirsMultisetKeysIterator,
> -    PyLeakedRef,
> +    PyLeakedRef<&'static DirsMultiset>,
>       DirsMultisetIter<'static>,
>       Dirs::translate_key,
>       Option<PyBytes>
> 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<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 @@ 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 @@ 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 @@ 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 @@ 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 @@ py_shared_ref!(DirstateMap, RustDirstate
>   
>   py_shared_iterator!(
>       DirstateMapKeysIterator,
> -    PyLeakedRef,
> +    PyLeakedRef<&'static RustDirstateMap>,
>       StateMapIter<'static>,
>       DirstateMap::translate_key,
>       Option<PyBytes>
> @@ -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<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 @@ 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 @@ 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 @@ 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 @@ 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)>

Patch

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<PyBytes>
@@ -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<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 @@  impl Dirs {
 
 py_shared_iterator!(
     DirsMultisetKeysIterator,
-    PyLeakedRef,
+    PyLeakedRef<&'static DirsMultiset>,
     DirsMultisetIter<'static>,
     Dirs::translate_key,
     Option<PyBytes>
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<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 @@  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 @@  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 @@  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 @@  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 @@  py_shared_ref!(DirstateMap, RustDirstate
 
 py_shared_iterator!(
     DirstateMapKeysIterator,
-    PyLeakedRef,
+    PyLeakedRef<&'static RustDirstateMap>,
     StateMapIter<'static>,
     DirstateMap::translate_key,
     Option<PyBytes>
@@ -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<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 @@  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 @@  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 @@  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 @@  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)>