Prevent mutable borrow after immutable one(s)

src/lib.rs

@@ -11,9 +11,10 @@ extern crate cpython;
 use cpython::exc::RuntimeError;
 use cpython::*;
 
-use std::cell::{RefCell, RefMut};
+use std::cell::{Cell, RefCell, RefMut};
 use std::collections::hash_set::Iter;
 use std::collections::HashSet;
+use std::ops::{Deref, DerefMut};
 
 type Inner = HashSet<u32>;
 
@@ -21,10 +22,15 @@ py_exception!(shared_ref, AlreadyBorrowed, RuntimeError);
 
 py_class!(class RustSet |py| {
     data hs: RefCell<Inner>;
-    data leak_count: RefCell<usize>;
+    data leak_count: Cell<usize>;
+    data mutably_borrowed: Cell<bool>;
 
     def __new__(_cls) -> PyResult<RustSet> {
-        Self::create_instance(py, RefCell::new(Inner::new()), RefCell::new(0))
+        Self::create_instance(py, RefCell::new(
+            Inner::new()),
+            Cell::new(0),
+            Cell::new(false),
+        )
     }
 
     def __contains__(&self, v: u32) -> PyResult<bool> {
@@ -45,14 +51,14 @@ py_class!(class RustSet |py| {
     }
 
     def get_leak_count(&self) -> PyResult<usize> {
-        Ok(*self.leak_count(py).borrow())
+        Ok(self.leak_count(py).get())
     }
 
     def __iter__(&self) -> PyResult<RustSetIterator> {
         RustSetIterator::create_instance(
             py,
             RefCell::new(Some(RustSetLeakedRef::new(py, &self))),
-            RefCell::new(self.leak_immutable(py).iter()),
+            RefCell::new(self.leak_immutable(py)?.iter()),
         )
     }
 
@@ -71,28 +77,62 @@ impl RustSet {
     /// Return a reference to the inner set, with an artificial static lifetime
     ///
     /// We need to be protected by the GIL for thread-safety
-    fn leak_immutable(&self, py: Python) -> &'static Inner {
+    fn leak_immutable(&self, py: Python) -> PyResult<&'static Inner> {
+        if self.mutably_borrowed(py).get() {
+            return Err(AlreadyBorrowed::new(
+                py,
+                "Cannot borrow immutably while there is a \
+                 mutable reference in Python objects",
+            ));
+        }
         let ptr = self.hs(py).as_ptr();
-        *self.leak_count(py).borrow_mut() += 1;
-        unsafe { &*ptr }
+        let leak_count = self.leak_count(py).get();
+        self.leak_count(py).replace(leak_count + 1);
+        unsafe { Ok(&*ptr) }
     }
 
     /// Borrow a mutable reference to the inner set.
     ///
-    /// Raise the explicit `AlreadyBorrowed` exception if there are immutable
-    // references currently in use.
-    fn borrow_mut<'a>(&'a self, py: Python<'a>) -> PyResult<RefMut<Inner>> {
-        match *self.leak_count(py).borrow() {
-            0 => Ok(self.hs(py).borrow_mut()),
+    /// Raise the explicit `AlreadyBorrowed` exception if there are other
+    /// references currently in use.
+    fn borrow_mut<'a>(&'a self, py: Python<'a>) -> PyResult<PyRefMut<Inner>> {
+        if self.mutably_borrowed(py).get() {
+            return Err(AlreadyBorrowed::new(
+                py,
+                "Cannot borrow mutably while there exists another \
+                 mutable reference in a Python object",
+            ));
+        }
+        match self.leak_count(py).get() {
+            0 => {
+                self.mutably_borrowed(py).replace(true);
+                Ok(PyRefMut::new(py, self.hs(py).borrow_mut(), &self))
+            }
+            // TODO
+            // For now, this works differently than Python references
+            // in the case of iterators.
+            // Python does not complain when the data an iterator
+            // points to is modified if the iterator is never used
+            // afterwards.
+            // Here, we are stricter than this by refusing to give a
+            // mutable reference if it is already borrowed.
+            // While the additional safety might be argued for, it
+            // breaks valid programming patterns in Python and we need
+            // to fix this issue down the line.
             _ => Err(AlreadyBorrowed::new(
                 py,
-                "Can't mutate while there are immutable references in Python objects",
+                "Cannot borrow mutably while there are \
+                 immutable references in Python objects",
             )),
         }
     }
 
-    fn decrease_leak_count(&self, py: Python) {
-        *self.leak_count(py).borrow_mut() -= 1;
+    fn decrease_leak_count(&self, py: Python, mutable: bool) {
+        let leak_count = self.leak_count(py).get();
+        self.leak_count(py).replace(leak_count.saturating_sub(1));
+        if mutable {
+            self.mutably_borrowed(py).replace(false);
+        }
     }
 }
 
@@ -116,7 +156,42 @@ impl Drop for RustSetLeakedRef {
     fn drop(&mut self) {
         let gil = Python::acquire_gil();
         let py = gil.python();
-        self.rs.decrease_leak_count(py);
+        self.rs.decrease_leak_count(py, false);
+    }
+}
+
+struct PyRefMut<'a, T> {
+    inner: RefMut<'a, T>,
+    rs: RustSet,
+}
+
+impl<'a, T> PyRefMut<'a, T> {
+    fn new(py: Python, rm: RefMut<'a, T>, rs: &RustSet) -> Self {
+        Self {
+            inner: rm,
+            rs: rs.clone_ref(py),
+        }
+    }
+}
+
+impl<'a, T> Deref for PyRefMut<'a, T> {
+    type Target = RefMut<'a, T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+impl<'a, T> DerefMut for PyRefMut<'a, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+impl<'a, T> Drop for PyRefMut<'a, T> {
+    fn drop(&mut self) {
+        let gil = Python::acquire_gil();
+        let py = gil.python();
+        self.rs.decrease_leak_count(py, true);
     }
 }
 
@@ -145,13 +220,18 @@ py_class!(class RustSetIterator |py| {
 
 });
 
-py_module_initializer!(shared_ref, initshared_ref, PyInit_shared_ref, |py, m| {
-    m.add(
-        py,
-        "__doc__",
-        "Demonstrator of shared Rust references within Python",
-    )?;
-    m.add_class::<RustSet>(py)?;
-    m.add(py, "AlreadyBorrowed", py.get_type::<AlreadyBorrowed>())?;
-    Ok(())
-});
+py_module_initializer!(
+    shared_ref,
+    initshared_ref,
+    PyInit_shared_ref,
+    |py, m| {
+        m.add(
+            py,
+            "__doc__",
+            "Demonstrator of shared Rust references within Python",
+        )?;
+        m.add_class::<RustSet>(py)?;
+        m.add(py, "AlreadyBorrowed", py.get_type::<AlreadyBorrowed>())?;
+        Ok(())
+    }
+);