Merge pull request #231 from adamreichold/simpler-slice-box

adamreichold · web-flow · commit 0efa2cde37da · 2022-01-06T17:05:33.000+01:00
RFC: Extend SliceBox and unify Array creation from owned data
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 - Unreleased
   - Support object arrays ([#216](https://github.com/PyO3/rust-numpy/pull/216))
+  - Support borrowing arrays that are part of other Python objects via `PyArray::borrow_from_array` ([#230](https://github.com/PyO3/rust-numpy/pull/216))
   - `PyArray::new` is now `unsafe`, as it produces uninitialized arrays ([#220](https://github.com/PyO3/rust-numpy/pull/220))
   - Fix thread-safety in internal API globals ([#222](https://github.com/PyO3/rust-numpy/pull/222))
   - Make arrays produced via `IntoPyArray`, i.e. those owning Rust data, writeable ([#235](https://github.com/PyO3/rust-numpy/pull/235))
diff --git a/src/array.rs b/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::slice_container::PySliceContainer;
 
 /// A safe, static-typed interface for
 /// [NumPy ndarray](https://numpy.org/doc/stable/reference/arrays.ndarray.html).
@@ -432,53 +432,71 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
             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
-            ptr::null_mut(),   // data
-            0,                 // itemsize
-            flag,              // flag
-            ptr::null_mut(),   // obj
+            T::npy_type() as c_int,
+            strides as *mut npy_intp, // strides
+            ptr::null_mut(),          // data
+            0,                        // itemsize
+            flag,                     // flag
+            ptr::null_mut(),          // obj
         );
         Self::from_owned_ptr(py, ptr)
     }
 
-    pub(crate) unsafe fn from_boxed_slice<'py, ID>(
+    unsafe fn new_with_data<'py, ID>(
         py: Python<'py>,
         dims: ID,
         strides: *const npy_intp,
-        boxed_slice: Box<[T]>,
-        data_ptr: Option<*const T>,
+        data_ptr: *const T,
+        container: *mut PyAny,
     ) -> &'py Self
     where
         ID: IntoDimension<Dim = D>,
     {
         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)
-            .create_cell(py)
-            .expect("Object creation failed.");
         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,
+            container as *mut ffi::PyObject,
         );
-        PY_ARRAY_API.PyArray_SetBaseObject(ptr as *mut npyffi::PyArrayObject, cell as _);
+
         Self::from_owned_ptr(py, ptr)
     }
 
-    /// Creates a NumPy array backed by `array` and ties its ownership to the Python object `owner`.
+    pub(crate) unsafe fn from_raw_parts<'py, ID, C>(
+        py: Python<'py>,
+        dims: ID,
+        strides: *const npy_intp,
+        data_ptr: *const T,
+        container: C,
+    ) -> &'py Self
+    where
+        ID: IntoDimension<Dim = D>,
+        PySliceContainer: From<C>,
+    {
+        let container = pyo3::PyClassInitializer::from(PySliceContainer::from(container))
+            .create_cell(py)
+            .expect("Object creation failed.");
+
+        Self::new_with_data(py, dims, strides, data_ptr, container as *mut PyAny)
+    }
+
+    /// Creates a NumPy array backed by `array` and ties its ownership to the Python object `container`.
     ///
     /// # Safety
     ///
-    /// `owner` is set as a base object of the returned array which must not be dropped until `owner` is dropped.
-    /// Furthermore, `array` must not be reallocated from the time this method is called and until `owner` is dropped.
+    /// `container` is set as a base object of the returned array which must not be dropped until `container` is dropped.
+    /// Furthermore, `array` must not be reallocated from the time this method is called and until `container` is dropped.
     ///
     /// # Example
     ///
@@ -503,33 +521,27 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     ///     }
     /// }
     /// ```
-    pub unsafe fn borrow_from_array<'py, S>(array: &ArrayBase<S, D>, owner: &'py PyAny) -> &'py Self
+    pub unsafe fn borrow_from_array<'py, S>(
+        array: &ArrayBase<S, D>,
+        container: &'py PyAny,
+    ) -> &'py Self
     where
         S: Data<Elem = T>,
     {
         let (strides, dims) = (array.npy_strides(), array.raw_dim());
         let data_ptr = array.as_ptr();
 
-        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 c_int,
-            strides.as_ptr() as *mut npy_intp, // strides
-            data_ptr as *mut c_void,           // data
-            mem::size_of::<T>() as c_int,      // itemsize
-            0,                                 // flag
-            ptr::null_mut(),                   // obj
-        );
-
-        mem::forget(owner.to_object(owner.py()));
+        let py = container.py();
 
-        PY_ARRAY_API.PyArray_SetBaseObject(
-            ptr as *mut npyffi::PyArrayObject,
-            owner as *const PyAny as *mut PyAny as *mut ffi::PyObject,
-        );
+        mem::forget(container.to_object(py));
 
-        Self::from_owned_ptr(owner.py(), ptr)
+        Self::new_with_data(
+            py,
+            dims,
+            strides.as_ptr(),
+            data_ptr,
+            container as *const PyAny as *mut PyAny,
+        )
     }
 
     /// Construct a new nd-dimensional array filled with 0.
diff --git a/src/convert.rs b/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) }
     }
 }
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -42,7 +42,7 @@ mod error;
 pub mod npyffi;
 pub mod npyiter;
 mod readonly;
-mod slice_box;
+mod slice_container;
 mod sum_products;
 
 pub use ndarray;
diff --git a/src/slice_box.rs b/src/slice_box.rs
diff --git a/src/slice_container.rs b/src/slice_container.rs
@@ -0,0 +1,113 @@
+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;
+
+/// Utility type to safely store Box<[_]> or Vec<_> on the Python heap
+pub(crate) struct PySliceContainer {
+    ptr: *mut u8,
+    len: usize,
+    cap: usize,
+    drop: unsafe fn(*mut u8, usize, usize),
+}
+
+unsafe impl Send for PySliceContainer {}
+
+impl<T: Send> From<Box<[T]>> for PySliceContainer {
+    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]);
+        }
+
+        // FIXME(adamreichold): Use `Box::into_raw` when
+        // `*mut [T]::{as_mut_ptr, len}` become stable and compatible with out MSRV.
+        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 PySliceContainer {
+    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);
+        }
+
+        // FIXME(adamreichold): Use `Vec::into_raw_parts`
+        // when it becomes stable and compatible with our MSRV.
+        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 PySliceContainer
+where
+    A: Element,
+    D: Dimension,
+{
+    fn from(data: ArrayBase<OwnedRepr<A>, D>) -> Self {
+        Self::from(data.into_raw_vec())
+    }
+}
+
+impl Drop for PySliceContainer {
+    fn drop(&mut self) {
+        unsafe {
+            (self.drop)(self.ptr, self.len, self.cap);
+        }
+    }
+}
+
+impl PyClass for PySliceContainer {
+    type Dict = PyClassDummySlot;
+    type WeakRef = PyClassDummySlot;
+    type BaseNativeType = PyAny;
+}
+
+impl PyClassImpl for PySliceContainer {
+    const DOC: &'static str = "Memory store for a PyArray backed by a Box<[_]> or a Vec<_> \0";
+
+    type BaseType = PyAny;
+    type Layout = PyCell<Self>;
+    type ThreadChecker = ThreadCheckerStub<Self>;
+}
+
+unsafe impl PyTypeInfo for PySliceContainer {
+    type AsRefTarget = PyCell<Self>;
+
+    const NAME: &'static str = "PySliceContainer";
+    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)
+    }
+}
