rust/hg-cpython/src/ancestors.rs
changeset 51239 7eea2e4109ae
parent 50979 4c5f6e95df84
child 51240 59d81768ad6d
--- a/rust/hg-cpython/src/ancestors.rs	Fri Oct 27 23:29:29 2023 +0200
+++ b/rust/hg-cpython/src/ancestors.rs	Fri Oct 27 22:11:05 2023 +0200
@@ -34,15 +34,17 @@
 //! [`LazyAncestors`]: struct.LazyAncestors.html
 //! [`MissingAncestors`]: struct.MissingAncestors.html
 //! [`AncestorsIterator`]: struct.AncestorsIterator.html
-use crate::revlog::pyindex_to_graph;
+use crate::revlog::{py_rust_index_to_graph, pyindex_to_graph};
 use crate::PyRevision;
 use crate::{
     cindex::Index, conversion::rev_pyiter_collect, exceptions::GraphError,
+    revlog::PySharedIndex,
 };
 use cpython::{
-    ObjectProtocol, PyClone, PyDict, PyList, PyModule, PyObject, PyResult,
-    Python, PythonObject, ToPyObject,
+    ObjectProtocol, PyClone, PyDict, PyErr, PyList, PyModule, PyObject,
+    PyResult, Python, PythonObject, ToPyObject, UnsafePyLeaked,
 };
+
 use hg::MissingAncestors as CoreMissing;
 use hg::Revision;
 use std::cell::RefCell;
@@ -52,11 +54,43 @@
     LazyAncestors as VCGLazyAncestors,
 };
 
+// Error propagation for an [`UnsafePyLeaked`] wrapping a [`Result`]
+//
+// It would be nice for UnsharedPyLeaked to provide this directly as a variant
+// of the `map` method with a signature such as:
+//
+// ```
+//   unafe fn map_or_err(py: Python,
+//                       f: impl FnOnce(T) -> Result(U, E),
+//                       convert_err: impl FnOnce(Python, E) -> PyErr)
+// ```
+//
+// This would spare users of the `cpython` crate the additional `unsafe` deref
+// to inspect the error and return it outside `UnsafePyLeaked`, and the
+// subsequent unwrapping that this function performs.
+fn pyleaked_or_map_err<T, E: std::fmt::Debug + Copy>(
+    py: Python,
+    leaked: UnsafePyLeaked<Result<T, E>>,
+    convert_err: impl FnOnce(Python, E) -> PyErr,
+) -> PyResult<UnsafePyLeaked<T>> {
+    // Result.inspect_err is unstable in Rust 1.61
+    if let Err(e) = *unsafe { leaked.try_borrow(py)? } {
+        return Err(convert_err(py, e));
+    }
+    Ok(unsafe {
+        leaked.map(py, |res| {
+            res.expect("Error case should have already be treated")
+        })
+    })
+}
+
 py_class!(pub class AncestorsIterator |py| {
-    data inner: RefCell<Box<VCGAncestorsIterator<Index>>>;
+    data inner: RefCell<UnsafePyLeaked<VCGAncestorsIterator<PySharedIndex>>>;
 
     def __next__(&self) -> PyResult<Option<PyRevision>> {
-        match self.inner(py).borrow_mut().next() {
+        let mut leaked = self.inner(py).borrow_mut();
+        let mut inner = unsafe { leaked.try_borrow_mut(py)? };
+        match inner.next() {
             Some(Err(e)) => Err(GraphError::pynew_from_vcsgraph(py, e)),
             None => Ok(None),
             Some(Ok(r)) => Ok(Some(PyRevision(r))),
@@ -64,7 +98,9 @@
     }
 
     def __contains__(&self, rev: PyRevision) -> PyResult<bool> {
-        self.inner(py).borrow_mut().contains(rev.0)
+        let mut leaked = self.inner(py).borrow_mut();
+        let mut inner = unsafe { leaked.try_borrow_mut(py)? };
+        inner.contains(rev.0)
             .map_err(|e| GraphError::pynew_from_vcsgraph(py, e))
     }
 
@@ -79,16 +115,7 @@
         stoprev: PyRevision,
         inclusive: bool
     ) -> PyResult<AncestorsIterator> {
-        let index = pyindex_to_graph(py, index)?;
-        let initvec: Vec<_> = rev_pyiter_collect(py, &initrevs, &index)?;
-        let ait = VCGAncestorsIterator::new(
-            index,
-            initvec.into_iter().map(|r| r.0),
-            stoprev.0,
-            inclusive,
-        )
-        .map_err(|e| GraphError::pynew_from_vcsgraph(py, e))?;
-        AncestorsIterator::from_inner(py, ait)
+        Self::inner_new(py, index, initrevs, stoprev, inclusive)
     }
 
 });
@@ -96,28 +123,71 @@
 impl AncestorsIterator {
     pub fn from_inner(
         py: Python,
-        ait: VCGAncestorsIterator<Index>,
+        ait: UnsafePyLeaked<VCGAncestorsIterator<PySharedIndex>>,
     ) -> PyResult<Self> {
-        Self::create_instance(py, RefCell::new(Box::new(ait)))
+        Self::create_instance(py, RefCell::new(ait))
+    }
+
+    pub fn inner_new(
+        py: Python,
+        index: PyObject,
+        initrevs: PyObject,
+        stoprev: PyRevision,
+        inclusive: bool,
+    ) -> PyResult<AncestorsIterator> {
+        let index = py_rust_index_to_graph(py, index)?;
+        let initvec: Vec<_> = {
+            let borrowed_idx = unsafe { index.try_borrow(py)? };
+            rev_pyiter_collect(py, &initrevs, &*borrowed_idx)?
+        };
+        let res_ait = unsafe {
+            index.map(py, |idx| {
+                VCGAncestorsIterator::new(
+                    idx,
+                    initvec.into_iter().map(|r| r.0),
+                    stoprev.0,
+                    inclusive,
+                )
+            })
+        };
+        let ait =
+            pyleaked_or_map_err(py, res_ait, GraphError::pynew_from_vcsgraph)?;
+        AncestorsIterator::from_inner(py, ait)
     }
 }
 
 py_class!(pub class LazyAncestors |py| {
-    data inner: RefCell<Box<VCGLazyAncestors<Index>>>;
+    data inner: RefCell<UnsafePyLeaked<
+        RefCell<VCGLazyAncestors<PySharedIndex>>
+        >>;
+    data index: PyObject;
+    data initrevs: PyObject;
+    data stoprev: PyRevision;
+    data inclusive: bool;
 
     def __contains__(&self, rev: PyRevision) -> PyResult<bool> {
-        self.inner(py)
-            .borrow_mut()
-            .contains(rev.0)
+        let leaked = self.inner(py).borrow();
+        let inner: &RefCell<VCGLazyAncestors<PySharedIndex>> =
+            &*unsafe { leaked.try_borrow(py)? };
+        let inner_mut: &mut VCGLazyAncestors<PySharedIndex> =
+            &mut *inner.borrow_mut();
+        inner_mut.contains(rev.0)
             .map_err(|e| GraphError::pynew_from_vcsgraph(py, e))
     }
 
     def __iter__(&self) -> PyResult<AncestorsIterator> {
-        AncestorsIterator::from_inner(py, self.inner(py).borrow().iter())
+        let index = self.index(py).clone_ref(py);
+        let initrevs = self.initrevs(py).clone_ref(py);
+        AncestorsIterator::inner_new(py, index, initrevs,
+                                     *self.stoprev(py),
+                                     *self.inclusive(py))
     }
 
     def __bool__(&self) -> PyResult<bool> {
-        Ok(!self.inner(py).borrow().is_empty())
+        let leaked = self.inner(py).borrow();
+        let inner = unsafe { leaked.try_borrow(py)? };
+        let empty = inner.borrow().is_empty();
+        Ok(!empty)
     }
 
     def __new__(
@@ -127,19 +197,27 @@
         stoprev: PyRevision,
         inclusive: bool
     ) -> PyResult<Self> {
-        let index = pyindex_to_graph(py, index)?;
-        let initvec: Vec<_> = rev_pyiter_collect(py, &initrevs, &index)?;
+        let cloned_index = index.clone_ref(py);
+        let index = py_rust_index_to_graph(py, index)?;
+        let initvec: Vec<_> = {
+            let borrowed_idx =  unsafe {index.try_borrow(py)?};
+            rev_pyiter_collect(py, &initrevs, &*borrowed_idx)?
+        };
 
-        let lazy =
-            VCGLazyAncestors::new(
-                index,
+        let res_lazy =
+            unsafe { index.map(py, |idx| VCGLazyAncestors::new(
+                idx,
                 initvec.into_iter().map(|r| r.0),
                 stoprev.0,
                 inclusive
-            )
-            .map_err(|e| GraphError::pynew_from_vcsgraph(py, e))?;
-
-        Self::create_instance(py, RefCell::new(Box::new(lazy)))
+            ))};
+        let lazy = pyleaked_or_map_err(py, res_lazy,
+                                       GraphError::pynew_from_vcsgraph)?;
+        let lazy_cell = unsafe { lazy.map(py, RefCell::new)};
+        let res = Self::create_instance(
+            py, RefCell::new(lazy_cell),
+            cloned_index, initrevs, stoprev, inclusive)?;
+        Ok(res)
         }
 
 });