Extend SliceBox into Owner which can hold onto both Box<[T]> and Vec<T>

By providing type erasure for both Box<[T]> and Vec<T> we can avoid having to
transform Vec<T> and Array<A, D> into boxed slices which can potentially re-allocate.

Author: adamreichold

src/array.rs

@@ -17,7 +17,7 @@ use std::{iter::ExactSizeIterator, marker::PhantomData};
 use crate::convert::{ArrayExt, IntoPyArray, NpyIndex, ToNpyDims, ToPyArray};
 use crate::dtype::{DataType, Element};
 use crate::error::{FromVecError, NotContiguousError, ShapeError};
-use crate::slice_box::SliceBox;
+use crate::owner::Owner;
 
 /// A safe, static-typed interface for
 /// [NumPy ndarray](https://numpy.org/doc/stable/reference/arrays.ndarray.html).
@@ -442,34 +442,39 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
         Self::from_owned_ptr(py, ptr)
     }
 
-    pub(crate) unsafe fn from_boxed_slice<'py, ID>(
+    pub(crate) unsafe fn from_raw_parts<'py, ID, O>(
         py: Python<'py>,
         dims: ID,
         strides: *const npy_intp,
-        boxed_slice: Box<[T]>,
-        data_ptr: Option<*const T>,
+        data_ptr: *const T,
+        owner: O,
     ) -> &'py Self
     where
         ID: IntoDimension<Dim = D>,
+        Owner: From<O>,
     {
-        let dims = dims.into_dimension();
-        let data_ptr = data_ptr.unwrap_or_else(|| boxed_slice.as_ptr());
-        let container = SliceBox::new(boxed_slice);
-        let cell = pyo3::PyClassInitializer::from(container)
+        let owner = pyo3::PyClassInitializer::from(Owner::from(owner))
             .create_cell(py)
             .expect("Object creation failed.");
+
+        let dims = dims.into_dimension();
         let ptr = PY_ARRAY_API.PyArray_New(
             PY_ARRAY_API.get_type_object(npyffi::NpyTypes::PyArray_Type),
             dims.ndim_cint(),
             dims.as_dims_ptr(),
-            T::npy_type() as i32,
-            strides as *mut _,           // strides
-            data_ptr as _,               // data
-            mem::size_of::<T>() as i32,  // itemsize
-            npyffi::NPY_ARRAY_WRITEABLE, // flag
-            ptr::null_mut(),             // obj
+            T::npy_type() as c_int,
+            strides as *mut npy_intp,     // strides
+            data_ptr as *mut c_void,      // data
+            mem::size_of::<T>() as c_int, // itemsize
+            npyffi::NPY_ARRAY_WRITEABLE,  // flag
+            ptr::null_mut(),              // obj
         );
-        PY_ARRAY_API.PyArray_SetBaseObject(ptr as *mut npyffi::PyArrayObject, cell as _);
+
+        PY_ARRAY_API.PyArray_SetBaseObject(
+            ptr as *mut npyffi::PyArrayObject,
+            owner as *mut ffi::PyObject,
+        );
+
         Self::from_owned_ptr(py, ptr)
     }
 

src/convert.rs

@@ -37,17 +37,21 @@ impl<T: Element> IntoPyArray for Box<[T]> {
     type Item = T;
     type Dim = Ix1;
     fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
-        let len = self.len();
+        let dims = [self.len()];
         let strides = [mem::size_of::<T>() as npy_intp];
-        unsafe { PyArray::from_boxed_slice(py, [len], strides.as_ptr(), self, None) }
+        let data_ptr = self.as_ptr();
+        unsafe { PyArray::from_raw_parts(py, dims, strides.as_ptr(), data_ptr, self) }
     }
 }
 
 impl<T: Element> IntoPyArray for Vec<T> {
     type Item = T;
     type Dim = Ix1;
     fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
-        self.into_boxed_slice().into_pyarray(py)
+        let dims = [self.len()];
+        let strides = [mem::size_of::<T>() as npy_intp];
+        let data_ptr = self.as_ptr();
+        unsafe { PyArray::from_raw_parts(py, dims, strides.as_ptr(), data_ptr, self) }
     }
 }
 
@@ -59,21 +63,9 @@ where
     type Item = A;
     type Dim = D;
     fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
-        let (strides, dim) = (self.npy_strides(), self.raw_dim());
-        let orig_ptr = self.as_ptr();
-        // Element of which size is 0 is not supported, but check it for future changes
-        let is_empty_or_size0 = self.is_empty() || std::mem::size_of::<A>() == 0;
-        let vec = self.into_raw_vec();
-        let offset = if is_empty_or_size0 {
-            0
-        } else {
-            (orig_ptr as usize - vec.as_ptr() as usize) / std::mem::size_of::<A>()
-        };
-        let mut boxed_slice = vec.into_boxed_slice();
-        // data_ptr is not always the pointer to the 1st element.
-        // See https://github.com/PyO3/rust-numpy/issues/182 for the detail.
-        let data_ptr = unsafe { boxed_slice.as_mut_ptr().add(offset) };
-        unsafe { PyArray::from_boxed_slice(py, dim, strides.as_ptr(), boxed_slice, Some(data_ptr)) }
+        let (strides, dims) = (self.npy_strides(), self.raw_dim());
+        let data_ptr = self.as_ptr();
+        unsafe { PyArray::from_raw_parts(py, dims, strides.as_ptr(), data_ptr, self) }
     }
 }
 

src/lib.rs

@@ -41,8 +41,8 @@ mod dtype;
 mod error;
 pub mod npyffi;
 pub mod npyiter;
+mod owner;
 mod readonly;
-mod slice_box;
 mod sum_products;
 
 pub use ndarray;

src/owner.rs

@@ -0,0 +1,108 @@
+use std::{mem, slice};
+
+use ndarray::{ArrayBase, Dimension, OwnedRepr};
+use pyo3::class::impl_::{PyClassImpl, ThreadCheckerStub};
+use pyo3::pyclass::PyClass;
+use pyo3::pyclass_slots::PyClassDummySlot;
+use pyo3::type_object::{LazyStaticType, PyTypeInfo};
+use pyo3::{ffi, types::PyAny, PyCell};
+
+use crate::dtype::Element;
+
+pub(crate) struct Owner {
+    ptr: *mut u8,
+    len: usize,
+    cap: usize,
+    drop: unsafe fn(*mut u8, usize, usize),
+}
+
+unsafe impl Send for Owner {}
+
+impl<T: Send> From<Box<[T]>> for Owner {
+    fn from(data: Box<[T]>) -> Self {
+        unsafe fn drop_boxed_slice<T>(ptr: *mut u8, len: usize, _cap: usize) {
+            let _ = Box::from_raw(slice::from_raw_parts_mut(ptr as *mut T, len) as *mut [T]);
+        }
+
+        let ptr = data.as_ptr() as *mut u8;
+        let len = data.len();
+        let cap = 0;
+        let drop = drop_boxed_slice::<T>;
+
+        mem::forget(data);
+
+        Self {
+            ptr,
+            len,
+            cap,
+            drop,
+        }
+    }
+}
+
+impl<T: Send> From<Vec<T>> for Owner {
+    fn from(data: Vec<T>) -> Self {
+        unsafe fn drop_vec<T>(ptr: *mut u8, len: usize, cap: usize) {
+            let _ = Vec::from_raw_parts(ptr as *mut T, len, cap);
+        }
+
+        let ptr = data.as_ptr() as *mut u8;
+        let len = data.len();
+        let cap = data.capacity();
+        let drop = drop_vec::<T>;
+
+        mem::forget(data);
+
+        Self {
+            ptr,
+            len,
+            cap,
+            drop,
+        }
+    }
+}
+
+impl<A, D> From<ArrayBase<OwnedRepr<A>, D>> for Owner
+where
+    A: Element,
+    D: Dimension,
+{
+    fn from(data: ArrayBase<OwnedRepr<A>, D>) -> Self {
+        Self::from(data.into_raw_vec())
+    }
+}
+
+impl Drop for Owner {
+    fn drop(&mut self) {
+        unsafe {
+            (self.drop)(self.ptr, self.len, self.cap);
+        }
+    }
+}
+
+impl PyClass for Owner {
+    type Dict = PyClassDummySlot;
+    type WeakRef = PyClassDummySlot;
+    type BaseNativeType = PyAny;
+}
+
+impl PyClassImpl for Owner {
+    const DOC: &'static str = "Memory store for a PyArray backed by a Rust type \0";
+
+    type BaseType = PyAny;
+    type Layout = PyCell<Self>;
+    type ThreadChecker = ThreadCheckerStub<Self>;
+}
+
+unsafe impl PyTypeInfo for Owner {
+    type AsRefTarget = PyCell<Self>;
+
+    const NAME: &'static str = "Owner";
+    const MODULE: Option<&'static str> = Some("_rust_numpy");
+
+    #[inline]
+    fn type_object_raw(py: pyo3::Python) -> *mut ffi::PyTypeObject {
+        static TYPE_OBJECT: LazyStaticType = LazyStaticType::new();
+        TYPE_OBJECT.get_or_init::<Self>(py)
+    }
+}