Fix tests when running under miri

There were a couple of violations of stacked borrow that
where fixed by using `Box<i64>` instead of inline i64 for the
shape/stride values.

Co-Authored-By: Guillaume Fraux <guillaume.fraux@epfl.ch>

Author: HaoZeke

.github/workflows/tests.yml

@@ -34,3 +34,16 @@ jobs:
         run: |
           pip install numpy
           cargo test --all-features
+
+  miri:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+
+      - name: Setup toolchain
+        run: |
+          rustup update nightly && rustup default nightly
+          rustup component add miri
+
+      - name: Run tests under miri
+        run: cargo miri test --all-features

src/data_types.rs

@@ -247,6 +247,7 @@ mod tests {
 
     #[test]
     #[cfg(feature = "half")]
+    #[cfg_attr(miri, ignore)]
     fn test_half_precision() {
         // Test that half::f16 correctly maps to kDLFloat with 16 bits
         let dtype = half::f16::get_dlpack_data_type();

src/pyo3.rs

@@ -20,6 +20,9 @@
 //! # Examples
 //!
 //! ```
+//! # #[cfg(miri)] fn main() {}
+//! # #[cfg(not(miri))]
+//! # fn main() {
 //! use pyo3::prelude::*;
 //! use pyo3::types::IntoPyDict;
 //! use pyo3::ffi::c_str;
@@ -55,6 +58,7 @@
 //!
 //!     assert_eq!(array, ndarray::arr2(&[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]));
 //! });
+//! # }
 //! ```
 
 use crate::sys::{self, DLManagedTensorVersioned};
@@ -334,6 +338,7 @@ mod tests {
     macro_rules! test_numpy_to_ndarray_via_dlpack_dtype {
         ($test_name:ident, $rust_type:ty, $np_dtype:expr) => {
             #[test]
+            #[cfg_attr(miri, ignore)]
             fn $test_name() -> PyResult<()> {
                 Python::initialize();
                 Python::attach(|py| {
@@ -370,6 +375,7 @@ result_capsule = array.__dlpack__()
     macro_rules! test_ndarray_to_numpy_via_dlpack_dtype {
         ($test_name:ident, $rust_type:ty, $np_dtype:expr) => {
             #[test]
+            #[cfg_attr(miri, ignore)]
             fn $test_name() -> PyResult<()> {
                 Python::initialize();
                 Python::attach(|py| {
@@ -414,6 +420,7 @@ assert np.allclose(array, expected)
     test_ndarray_to_numpy_via_dlpack_dtype!(test_to_numpy_i64, i64, "int64");
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_null_strides_fails_conversion() -> PyResult<()> {
         Python::initialize();
         Python::attach(|py| {
@@ -456,6 +463,7 @@ assert np.allclose(array, expected)
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_v1_0_null_strides_allowed() -> PyResult<()> {
         Python::initialize();
         Python::attach(|py| {

src/sync.rs

@@ -24,8 +24,12 @@ struct ManagerContextMutex<T> where T: 'static {
     #[borrows(array)]
     #[covariant]
     lock: MutexGuard<'this, Vec<T>>,
-    shape: i64,
-    stride: i64,
+    // Use Box<i64> so that pointers derived via with_*_mut target heap
+    // memory rather than inline struct fields. This avoids Stacked Borrows
+    // violations when multiple with_*_mut calls each create exclusive
+    // reborrows of the ouroboros struct.
+    shape: Box<i64>,
+    stride: Box<i64>,
 }
 
 unsafe extern "C" fn mutex_deleter_fn<T>(manager: *mut sys::DLManagedTensorVersioned) where T: 'static {
@@ -41,24 +45,24 @@ impl<T> TryFrom<Arc<Mutex<Vec<T>>>> for DLPackTensor where T: GetDLPackDataType
         let ctx = ManagerContextMutexBuilder {
             array: array,
             lock_builder: |array| { array.lock().expect("could not lock the mutex") },
-            shape: 0,
-            stride: 1,
+            shape: Box::new(0),
+            stride: Box::new(1),
         };
         let mut ctx = Box::new(ctx.build());
 
         // set the shape after acquiring the lock to avoid deadlocks
         let shape = ctx.with_lock(|lock| lock.len() as i64);
-        ctx.with_shape_mut(|v| *v = shape);
+        ctx.with_shape_mut(|v| **v = shape);
 
         // extract pointers out of the boxed context to use in the DLPack tensor
         let mut shape_ptr = std::ptr::null_mut();
         ctx.with_shape_mut(|shape| {
-            shape_ptr = shape as *mut i64;
+            shape_ptr = shape.as_mut();
         });
 
         let mut stride_ptr = std::ptr::null_mut();
         ctx.with_stride_mut(|stride| {
-            stride_ptr = stride as *mut i64;
+            stride_ptr = stride.as_mut();
         });
 
         let mut data = std::ptr::null_mut();
@@ -102,8 +106,8 @@ struct ManagerContextRwLock<T> where T: 'static {
     #[borrows(array)]
     #[covariant]
     lock: RwLockWriteGuard<'this, Vec<T>>,
-    shape: i64,
-    stride: i64,
+    shape: Box<i64>,
+    stride: Box<i64>,
 }
 
 unsafe extern "C" fn rwlock_deleter_fn<T>(manager: *mut sys::DLManagedTensorVersioned) where T: 'static {
@@ -119,24 +123,24 @@ impl<T> TryFrom<Arc<RwLock<Vec<T>>>> for DLPackTensor where T: GetDLPackDataType
         let ctx = ManagerContextRwLockBuilder {
             array: array,
             lock_builder: move |array| { array.write().expect("could not lock the rwlock") },
-            shape: 0,
-            stride: 1,
+            shape: Box::new(0),
+            stride: Box::new(1),
         };
         let mut ctx = Box::new(ctx.build());
 
         // set the shape after acquiring the lock to avoid deadlocks
         let shape = ctx.with_lock(|lock| lock.len() as i64);
-        ctx.with_shape_mut(|v| *v = shape);
+        ctx.with_shape_mut(|v| **v = shape);
 
         // extract pointers out of the boxed context to use in the DLPack tensor
         let mut shape_ptr = std::ptr::null_mut();
         ctx.with_shape_mut(|shape| {
-            shape_ptr = shape as *mut i64;
+            shape_ptr = shape.as_mut();
         });
 
         let mut stride_ptr = std::ptr::null_mut();
         ctx.with_stride_mut(|stride| {
-            stride_ptr = stride as *mut i64;
+            stride_ptr = stride.as_mut();
         });
 
         let mut data = std::ptr::null_mut();
@@ -215,4 +219,27 @@ mod tests {
         let lock = data.read().unwrap();
         assert_eq!(&*lock, &[1, 42, 3]);
     }
+
+    // Last-ref tests: the tensor holds the only Arc reference, so dropping
+    // it actually deallocates the ManagerContext via the deleter function.
+
+    #[test]
+    fn test_mutex_last_arc_ref() {
+        let data = Arc::new(Mutex::new(vec![1i32, 2, 3]));
+
+        let mut tensor: DLPackTensor = data.try_into().unwrap();
+        let tensor_mut_ref = tensor.as_mut();
+        let slice: &mut [i32] = tensor_mut_ref.try_into().unwrap();
+        assert_eq!(slice, &[1, 2, 3]);
+    }
+
+    #[test]
+    fn test_rwlock_last_arc_ref() {
+        let data = Arc::new(RwLock::new(vec![1i32, 2, 3]));
+
+        let mut tensor: DLPackTensor = data.try_into().unwrap();
+        let tensor_mut_ref = tensor.as_mut();
+        let slice: &mut [i32] = tensor_mut_ref.try_into().unwrap();
+        assert_eq!(slice, &[1, 2, 3]);
+    }
 }

src/vec.rs

@@ -125,8 +125,8 @@ impl<'a, T> TryFrom<DLPackTensorRefMut<'a>> for &'a mut [T] where T: DLPackPoint
 
 struct ManagerContext<T> {
     array: T,
-    shape: i64,
-    stride: i64,
+    shape: Box<i64>,
+    stride: Box<i64>,
 }
 
 unsafe extern "C" fn deleter_fn<T>(manager: *mut sys::DLManagedTensorVersioned) {
@@ -144,12 +144,12 @@ macro_rules! impl_try_from {
                 let len = value.len();
                 let mut ctx = Box::new(ManagerContext {
                     array: value,
-                    shape: len as i64,
-                    stride: 1,
+                    shape: Box::new(len as i64),
+                    stride: Box::new(1),
                 });
 
-                let shape_ptr = &mut ctx.shape;
-                let stride_ptr = &mut ctx.stride;
+                let shape_ptr = ctx.shape.as_mut();
+                let stride_ptr = ctx.stride.as_mut();
 
                 let dl_tensor = sys::DLTensor {
                     data: ctx.array.as_ptr().cast_mut().cast(),