Merge pull request #233 from adamreichold/fix-slice-box-type-confusion

Fix SliceBox-induced type confusion

Author: davidhewitt

src/array.rs

@@ -448,8 +448,8 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
         ID: IntoDimension<Dim = D>,
     {
         let dims = dims.into_dimension();
+        let data_ptr = data_ptr.unwrap_or(boxed_slice.as_ptr());
         let container = SliceBox::new(boxed_slice);
-        let data_ptr = data_ptr.unwrap_or_else(|| container.data.as_ptr());
         let cell = pyo3::PyClassInitializer::from(container)
             .create_cell(py)
             .expect("Object creation failed.");

src/npyiter.rs

@@ -186,6 +186,7 @@ impl<'py, T: Element> NpySingleIterBuilder<'py, T, ReadWrite> {
 
 impl<'py, T: Element, I: IterMode> NpySingleIterBuilder<'py, T, I> {
     /// Sets a flag to this builder, returning `self`.
+    #[must_use]
     pub fn set(mut self, flag: NpyIterFlag) -> Self {
         self.flags |= flag.to_c_enum();
         self
@@ -388,6 +389,7 @@ impl<'py, T: Element> NpyMultiIterBuilder<'py, T, ()> {
     }
 
     /// Set a flag to this builder, returning `self`.
+    #[must_use]
     pub fn set(mut self, flag: NpyIterFlag) -> Self {
         self.flags |= flag.to_c_enum();
         self

src/slice_box.rs

@@ -1,49 +1,58 @@
 use pyo3::class::impl_::{PyClassImpl, ThreadCheckerStub};
 use pyo3::pyclass::PyClass;
 use pyo3::pyclass_slots::PyClassDummySlot;
-use pyo3::{ffi, type_object, types::PyAny, PyCell};
+use pyo3::type_object::{LazyStaticType, PyTypeInfo};
+use pyo3::{ffi, types::PyAny, PyCell};
 
-pub(crate) struct SliceBox<T> {
-    pub(crate) data: Box<[T]>,
+pub(crate) struct SliceBox {
+    ptr: *mut [u8],
+    drop: unsafe fn(*mut [u8]),
 }
 
-impl<T> SliceBox<T> {
-    pub(crate) fn new(data: Box<[T]>) -> Self {
-        Self { data }
+unsafe impl Send for SliceBox {}
+
+impl SliceBox {
+    pub(crate) fn new<T: Send>(data: Box<[T]>) -> Self {
+        unsafe fn drop_boxed_slice<T>(ptr: *mut [u8]) {
+            let _ = Box::from_raw(ptr as *mut [T]);
+        }
+
+        let ptr = Box::into_raw(data) as *mut [u8];
+        let drop = drop_boxed_slice::<T>;
+
+        Self { ptr, drop }
+    }
+}
+
+impl Drop for SliceBox {
+    fn drop(&mut self) {
+        unsafe {
+            (self.drop)(self.ptr);
+        }
     }
 }
 
-impl<T> PyClass for SliceBox<T>
-where
-    T: Send,
-{
+impl PyClass for SliceBox {
     type Dict = PyClassDummySlot;
     type WeakRef = PyClassDummySlot;
     type BaseNativeType = PyAny;
 }
 
-impl<T> PyClassImpl for SliceBox<T>
-where
-    T: Send,
-{
+impl PyClassImpl for SliceBox {
     const DOC: &'static str = "Memory store for PyArray using rust's Box<[T]> \0";
 
     type BaseType = PyAny;
     type Layout = PyCell<Self>;
     type ThreadChecker = ThreadCheckerStub<Self>;
 }
 
-unsafe impl<T> type_object::PyTypeInfo for SliceBox<T>
-where
-    T: Send,
-{
+unsafe impl PyTypeInfo for SliceBox {
     type AsRefTarget = PyCell<Self>;
     const NAME: &'static str = "SliceBox";
     const MODULE: Option<&'static str> = Some("_rust_numpy");
 
     #[inline]
     fn type_object_raw(py: pyo3::Python) -> *mut ffi::PyTypeObject {
-        use pyo3::type_object::LazyStaticType;
         static TYPE_OBJECT: LazyStaticType = LazyStaticType::new();
         TYPE_OBJECT.get_or_init::<Self>(py)
     }

tests/to_py.rs

@@ -232,3 +232,27 @@ fn to_pyarray_object_array() {
         }
     })
 }
+
+#[test]
+fn slice_box_type_confusion() {
+    use ndarray::Array2;
+    use pyo3::{
+        types::{PyDict, PyString},
+        ToPyObject,
+    };
+
+    pyo3::Python::with_gil(|py| {
+        let mut nd_arr = Array2::from_shape_fn((2, 3), |(_, _)| py.None());
+        nd_arr[(0, 2)] = PyDict::new(py).to_object(py);
+        nd_arr[(1, 0)] = PyString::new(py, "Hello:)").to_object(py);
+
+        let _py_arr = nd_arr.into_pyarray(py);
+
+        // Dropping `_arr` used to trigger a segmentation fault due to calling `Py_DECREF`
+        // on 1, 2 and 3 interpreted as pointers into the Python heap
+        // after having created a `SliceBox<PyObject>` backing `_py_arr`,
+        // c.f. https://github.com/PyO3/rust-numpy/issues/232.
+        let vec = vec![1, 2, 3];
+        let _arr = vec.into_pyarray(py);
+    });
+}