Fix a drop bug in PyReadwriteArray::resize and optimize borrowing by caching address and key.

Author: adamreichold

src/borrow.rs

@@ -180,7 +180,7 @@ use crate::dtype::Element;
 use crate::error::{BorrowError, NotContiguousError};
 use crate::npyffi::{self, PyArrayObject, NPY_ARRAY_WRITEABLE};
 
-#[derive(PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
 struct BorrowKey {
     /// exclusive range of lowest and highest address covered by array
     range: (usize, usize),
@@ -375,10 +375,16 @@ static BORROW_FLAGS: BorrowFlags = BorrowFlags::new();
 /// i.e. that only shared references into the interior of the array can be created safely.
 ///
 /// See the [module-level documentation](self) for more.
-pub struct PyReadonlyArray<'py, T, D>(&'py PyArray<T, D>)
+#[repr(C)]
+pub struct PyReadonlyArray<'py, T, D>
 where
     T: Element,
-    D: Dimension;
+    D: Dimension,
+{
+    array: &'py PyArray<T, D>,
+    address: usize,
+    key: BorrowKey,
+}
 
 /// Read-only borrow of a one-dimensional array.
 pub type PyReadonlyArray1<'py, T> = PyReadonlyArray<'py, T, Ix1>;
@@ -409,7 +415,7 @@ where
     type Target = PyArray<T, D>;
 
     fn deref(&self) -> &Self::Target {
-        self.0
+        self.array
     }
 }
 
@@ -426,27 +432,30 @@ where
     D: Dimension,
 {
     pub(crate) fn try_new(array: &'py PyArray<T, D>) -> Result<Self, BorrowError> {
-        let py = array.py();
         let address = base_address(array);
         let key = BorrowKey::from_array(array);
 
-        BORROW_FLAGS.acquire(py, address, key)?;
+        BORROW_FLAGS.acquire(array.py(), address, key)?;
 
-        Ok(Self(array))
+        Ok(Self {
+            array,
+            address,
+            key,
+        })
     }
 
     /// Provides an immutable array view of the interior of the NumPy array.
     #[inline(always)]
     pub fn as_array(&self) -> ArrayView<T, D> {
         // SAFETY: Global borrow flags ensure aliasing discipline.
-        unsafe { self.0.as_array() }
+        unsafe { self.array.as_array() }
     }
 
     /// Provide an immutable slice view of the interior of the NumPy array if it is contiguous.
     #[inline(always)]
     pub fn as_slice(&self) -> Result<&[T], NotContiguousError> {
         // SAFETY: Global borrow flags ensure aliasing discipline.
-        unsafe { self.0.as_slice() }
+        unsafe { self.array.as_slice() }
     }
 
     /// Provide an immutable reference to an element of the NumPy array if the index is within bounds.
@@ -455,7 +464,7 @@ where
     where
         I: NpyIndex<Dim = D>,
     {
-        unsafe { self.0.get(index) }
+        unsafe { self.array.get(index) }
     }
 }
 
@@ -465,7 +474,15 @@ where
     D: Dimension,
 {
     fn clone(&self) -> Self {
-        Self::try_new(self.0).unwrap()
+        BORROW_FLAGS
+            .acquire(self.array.py(), self.address, self.key)
+            .unwrap();
+
+        Self {
+            array: self.array,
+            address: self.address,
+            key: self.key,
+        }
     }
 }
 
@@ -475,11 +492,7 @@ where
     D: Dimension,
 {
     fn drop(&mut self) {
-        let py = self.0.py();
-        let address = base_address(self.0);
-        let key = BorrowKey::from_array(self.0);
-
-        BORROW_FLAGS.release(py, address, key);
+        BORROW_FLAGS.release(self.array.py(), self.address, self.key);
     }
 }
 
@@ -505,10 +518,16 @@ where
 /// i.e. that only a single exclusive reference into the interior of the array can be created safely.
 ///
 /// See the [module-level documentation](self) for more.
-pub struct PyReadwriteArray<'py, T, D>(&'py PyArray<T, D>)
+#[repr(C)]
+pub struct PyReadwriteArray<'py, T, D>
 where
     T: Element,
-    D: Dimension;
+    D: Dimension,
+{
+    array: &'py PyArray<T, D>,
+    address: usize,
+    key: BorrowKey,
+}
 
 /// Read-write borrow of a one-dimensional array.
 pub type PyReadwriteArray1<'py, T> = PyReadwriteArray<'py, T, Ix1>;
@@ -561,27 +580,30 @@ where
             return Err(BorrowError::NotWriteable);
         }
 
-        let py = array.py();
         let address = base_address(array);
         let key = BorrowKey::from_array(array);
 
-        BORROW_FLAGS.acquire_mut(py, address, key)?;
+        BORROW_FLAGS.acquire_mut(array.py(), address, key)?;
 
-        Ok(Self(array))
+        Ok(Self {
+            array,
+            address,
+            key,
+        })
     }
 
     /// Provides a mutable array view of the interior of the NumPy array.
     #[inline(always)]
     pub fn as_array_mut(&mut self) -> ArrayViewMut<T, D> {
         // SAFETY: Global borrow flags ensure aliasing discipline.
-        unsafe { self.0.as_array_mut() }
+        unsafe { self.array.as_array_mut() }
     }
 
     /// Provide a mutable slice view of the interior of the NumPy array if it is contiguous.
     #[inline(always)]
     pub fn as_slice_mut(&mut self) -> Result<&mut [T], NotContiguousError> {
         // SAFETY: Global borrow flags ensure aliasing discipline.
-        unsafe { self.0.as_slice_mut() }
+        unsafe { self.array.as_slice_mut() }
     }
 
     /// Provide a mutable reference to an element of the NumPy array if the index is within bounds.
@@ -590,7 +612,7 @@ where
     where
         I: NpyIndex<Dim = D>,
     {
-        unsafe { self.0.get_mut(index) }
+        unsafe { self.array.get_mut(index) }
     }
 }
 
@@ -616,23 +638,16 @@ where
     /// });
     /// ```
     pub fn resize(self, new_elems: usize) -> PyResult<Self> {
-        let py = self.0.py();
-        let address = base_address(self.0);
-        let key = BorrowKey::from_array(self.0);
-
-        BORROW_FLAGS.release_mut(py, address, key);
+        let array = self.array;
 
         // SAFETY: Ownership of `self` proves exclusive access to the interior of the array.
         unsafe {
-            self.0.resize(new_elems)?;
+            array.resize(new_elems)?;
         }
 
-        let address = base_address(self.0);
-        let key = BorrowKey::from_array(self.0);
-
-        BORROW_FLAGS.acquire_mut(py, address, key)?;
+        drop(self);
 
-        Ok(self)
+        Ok(Self::try_new(array).unwrap())
     }
 }
 
@@ -642,11 +657,7 @@ where
     D: Dimension,
 {
     fn drop(&mut self) {
-        let py = self.0.py();
-        let address = base_address(self.0);
-        let key = BorrowKey::from_array(self.0);
-
-        BORROW_FLAGS.release_mut(py, address, key);
+        BORROW_FLAGS.release_mut(self.array.py(), self.address, self.key);
     }
 }
 
@@ -1275,4 +1286,43 @@ mod tests {
             }
         });
     }
+
+    #[test]
+    #[should_panic(expected = "AlreadyBorrowed")]
+    fn cannot_clone_exclusive_borrow_via_deref() {
+        Python::with_gil(|py| {
+            let array = PyArray::<f64, _>::zeros(py, (3, 2, 1), false);
+
+            let exclusive = array.readwrite();
+            let _shared = exclusive.clone();
+        });
+    }
+
+    #[test]
+    fn failed_resize_does_not_double_release() {
+        Python::with_gil(|py| {
+            let array = PyArray::<f64, _>::zeros(py, 10, false);
+
+            // The view will make the internal reference check of `PyArray_Resize` fail.
+            let locals = [("array", array)].into_py_dict(py);
+            let _view = py
+                .eval("array[:]", None, Some(locals))
+                .unwrap()
+                .downcast::<PyArray1<f64>>()
+                .unwrap();
+
+            let exclusive = array.readwrite();
+            assert!(exclusive.resize(100).is_err());
+        });
+    }
+
+    #[test]
+    fn ineffective_resize_does_not_conflict() {
+        Python::with_gil(|py| {
+            let array = PyArray::<f64, _>::zeros(py, 10, false);
+
+            let exclusive = array.readwrite();
+            assert!(exclusive.resize(10).is_ok());
+        });
+    }
 }