ndarray/
impl_raw_views.rs

1use num_complex::Complex;
2use std::mem;
3use std::ptr::NonNull;
4
5use crate::dimension::{self, stride_offset};
6use crate::extension::nonnull::nonnull_debug_checked_from_ptr;
7use crate::imp_prelude::*;
8use crate::is_aligned;
9use crate::shape_builder::{Strides, StrideShape};
10
11impl<A, D> RawArrayView<A, D>
12where
13    D: Dimension,
14{
15    /// Create a new `RawArrayView`.
16    ///
17    /// Unsafe because caller is responsible for ensuring that the array will
18    /// meet all of the invariants of the `ArrayBase` type.
19    #[inline]
20    pub(crate) unsafe fn new(ptr: NonNull<A>, dim: D, strides: D) -> Self {
21        RawArrayView::from_data_ptr(RawViewRepr::new(), ptr)
22            .with_strides_dim(strides, dim)
23    }
24
25    unsafe fn new_(ptr: *const A, dim: D, strides: D) -> Self {
26        Self::new(nonnull_debug_checked_from_ptr(ptr as *mut A), dim, strides)
27    }
28
29    /// Create an `RawArrayView<A, D>` from shape information and a raw pointer
30    /// to the elements.
31    ///
32    /// # Safety
33    ///
34    /// The caller is responsible for ensuring all of the following:
35    ///
36    /// * `ptr` must be non-null, and it must be safe to [`.offset()`] `ptr` by
37    ///   zero.
38    ///
39    /// * It must be safe to [`.offset()`] the pointer repeatedly along all
40    ///   axes and calculate the `count`s for the `.offset()` calls without
41    ///   overflow, even if the array is empty or the elements are zero-sized.
42    ///
43    ///   In other words,
44    ///
45    ///   * All possible pointers generated by moving along all axes must be in
46    ///     bounds or one byte past the end of a single allocation with element
47    ///     type `A`. The only exceptions are if the array is empty or the element
48    ///     type is zero-sized. In these cases, `ptr` may be dangling, but it must
49    ///     still be safe to [`.offset()`] the pointer along the axes.
50    ///
51    ///   * The offset in units of bytes between the least address and greatest
52    ///     address by moving along all axes must not exceed `isize::MAX`. This
53    ///     constraint prevents the computed offset, in bytes, from overflowing
54    ///     `isize` regardless of the starting point due to past offsets.
55    ///
56    ///   * The offset in units of `A` between the least address and greatest
57    ///     address by moving along all axes must not exceed `isize::MAX`. This
58    ///     constraint prevents overflow when calculating the `count` parameter to
59    ///     [`.offset()`] regardless of the starting point due to past offsets.
60    ///
61    /// * The product of non-zero axis lengths must not exceed `isize::MAX`.
62    /// 
63    /// * Strides must be non-negative.
64    ///
65    /// This function can use debug assertions to check some of these requirements,
66    /// but it's not a complete check.
67    ///
68    /// [`.offset()`]: https://doc.rust-lang.org/stable/std/primitive.pointer.html#method.offset
69    pub unsafe fn from_shape_ptr<Sh>(shape: Sh, ptr: *const A) -> Self
70    where
71        Sh: Into<StrideShape<D>>,
72    {
73        let shape = shape.into();
74        let dim = shape.dim;
75        if cfg!(debug_assertions) {
76            assert!(!ptr.is_null(), "The pointer must be non-null.");
77            if let Strides::Custom(strides) = &shape.strides {
78                dimension::strides_non_negative(strides).unwrap();
79                dimension::max_abs_offset_check_overflow::<A, _>(&dim, strides).unwrap();
80            } else {
81                dimension::size_of_shape_checked(&dim).unwrap();
82            }
83        }
84        let strides = shape.strides.strides_for_dim(&dim);
85        RawArrayView::new_(ptr, dim, strides)
86    }
87
88    /// Converts to a read-only view of the array.
89    ///
90    /// # Safety
91    ///
92    /// From a safety standpoint, this is equivalent to dereferencing a raw
93    /// pointer for every element in the array. You must ensure that all of the
94    /// data is valid, ensure that the pointer is aligned, and choose the
95    /// correct lifetime.
96    #[inline]
97    pub unsafe fn deref_into_view<'a>(self) -> ArrayView<'a, A, D> {
98        debug_assert!(
99            is_aligned(self.ptr.as_ptr()),
100            "The pointer must be aligned."
101        );
102        ArrayView::new(self.ptr, self.dim, self.strides)
103    }
104
105    /// Split the array view along `axis` and return one array pointer strictly
106    /// before the split and one array pointer after the split.
107    ///
108    /// **Panics** if `axis` or `index` is out of bounds.
109    pub fn split_at(self, axis: Axis, index: Ix) -> (Self, Self) {
110        assert!(index <= self.len_of(axis));
111        let left_ptr = self.ptr.as_ptr();
112        let right_ptr = if index == self.len_of(axis) {
113            self.ptr.as_ptr()
114        } else {
115            let offset = stride_offset(index, self.strides.axis(axis));
116            // The `.offset()` is safe due to the guarantees of `RawData`.
117            unsafe { self.ptr.as_ptr().offset(offset) }
118        };
119
120        let mut dim_left = self.dim.clone();
121        dim_left.set_axis(axis, index);
122        let left = unsafe { Self::new_(left_ptr, dim_left, self.strides.clone()) };
123
124        let mut dim_right = self.dim;
125        let right_len = dim_right.axis(axis) - index;
126        dim_right.set_axis(axis, right_len);
127        let right = unsafe { Self::new_(right_ptr, dim_right, self.strides) };
128
129        (left, right)
130    }
131
132    /// Cast the raw pointer of the raw array view to a different type
133    ///
134    /// **Panics** if element size is not compatible.
135    ///
136    /// Lack of panic does not imply it is a valid cast. The cast works the same
137    /// way as regular raw pointer casts.
138    ///
139    /// While this method is safe, for the same reason as regular raw pointer
140    /// casts are safe, access through the produced raw view is only possible
141    /// in an unsafe block or function.
142    pub fn cast<B>(self) -> RawArrayView<B, D> {
143        assert_eq!(
144            mem::size_of::<B>(),
145            mem::size_of::<A>(),
146            "size mismatch in raw view cast"
147        );
148        let ptr = self.ptr.cast::<B>();
149        unsafe { RawArrayView::new(ptr, self.dim, self.strides) }
150    }
151}
152
153impl<T, D> RawArrayView<Complex<T>, D>
154where
155    D: Dimension,
156{
157    /// Splits the view into views of the real and imaginary components of the
158    /// elements.
159    pub fn split_complex(self) -> Complex<RawArrayView<T, D>> {
160        // Check that the size and alignment of `Complex<T>` are as expected.
161        // These assertions should always pass, for arbitrary `T`.
162        assert_eq!(
163            mem::size_of::<Complex<T>>(),
164            mem::size_of::<T>().checked_mul(2).unwrap()
165        );
166        assert_eq!(mem::align_of::<Complex<T>>(), mem::align_of::<T>());
167
168        let dim = self.dim.clone();
169
170        // Double the strides. In the zero-sized element case and for axes of
171        // length <= 1, we leave the strides as-is to avoid possible overflow.
172        let mut strides = self.strides.clone();
173        if mem::size_of::<T>() != 0 {
174            for ax in 0..strides.ndim() {
175                if dim[ax] > 1 {
176                    strides[ax] = (strides[ax] as isize * 2) as usize;
177                }
178            }
179        }
180
181        let ptr_re: *mut T = self.ptr.as_ptr().cast();
182        let ptr_im: *mut T = if self.is_empty() {
183            // In the empty case, we can just reuse the existing pointer since
184            // it won't be dereferenced anyway. It is not safe to offset by
185            // one, since the allocation may be empty.
186            ptr_re
187        } else {
188            // In the nonempty case, we can safely offset into the first
189            // (complex) element.
190            unsafe { ptr_re.add(1) }
191        };
192
193        // `Complex` is `repr(C)` with only fields `re: T` and `im: T`. So, the
194        // real components of the elements start at the same pointer, and the
195        // imaginary components start at the pointer offset by one, with
196        // exactly double the strides. The new, doubled strides still meet the
197        // overflow constraints:
198        //
199        // - For the zero-sized element case, the strides are unchanged in
200        //   units of bytes and in units of the element type.
201        //
202        // - For the nonzero-sized element case:
203        //
204        //   - In units of bytes, the strides are unchanged. The only exception
205        //     is axes of length <= 1, but those strides are irrelevant anyway.
206        //
207        //   - Since `Complex<T>` for nonzero `T` is always at least 2 bytes,
208        //     and the original strides did not overflow in units of bytes, we
209        //     know that the new, doubled strides will not overflow in units of
210        //     `T`.
211        unsafe {
212            Complex {
213                re: RawArrayView::new_(ptr_re, dim.clone(), strides.clone()),
214                im: RawArrayView::new_(ptr_im, dim, strides),
215            }
216        }
217    }
218}
219
220impl<A, D> RawArrayViewMut<A, D>
221where
222    D: Dimension,
223{
224    /// Create a new `RawArrayViewMut`.
225    ///
226    /// Unsafe because caller is responsible for ensuring that the array will
227    /// meet all of the invariants of the `ArrayBase` type.
228    #[inline]
229    pub(crate) unsafe fn new(ptr: NonNull<A>, dim: D, strides: D) -> Self {
230        RawArrayViewMut::from_data_ptr(RawViewRepr::new(), ptr)
231            .with_strides_dim(strides, dim)
232    }
233
234    unsafe fn new_(ptr: *mut A, dim: D, strides: D) -> Self {
235        Self::new(nonnull_debug_checked_from_ptr(ptr), dim, strides)
236    }
237
238    /// Create an `RawArrayViewMut<A, D>` from shape information and a raw
239    /// pointer to the elements.
240    ///
241    /// # Safety
242    ///
243    /// The caller is responsible for ensuring all of the following:
244    ///
245    /// * `ptr` must be non-null, and it must be safe to [`.offset()`] `ptr` by
246    ///   zero.
247    ///
248    /// * It must be safe to [`.offset()`] the pointer repeatedly along all
249    ///   axes and calculate the `count`s for the `.offset()` calls without
250    ///   overflow, even if the array is empty or the elements are zero-sized.
251    ///
252    ///   In other words,
253    ///
254    ///   * All possible pointers generated by moving along all axes must be in
255    ///     bounds or one byte past the end of a single allocation with element
256    ///     type `A`. The only exceptions are if the array is empty or the element
257    ///     type is zero-sized. In these cases, `ptr` may be dangling, but it must
258    ///     still be safe to [`.offset()`] the pointer along the axes.
259    ///
260    ///   * The offset in units of bytes between the least address and greatest
261    ///     address by moving along all axes must not exceed `isize::MAX`. This
262    ///     constraint prevents the computed offset, in bytes, from overflowing
263    ///     `isize` regardless of the starting point due to past offsets.
264    ///
265    ///   * The offset in units of `A` between the least address and greatest
266    ///     address by moving along all axes must not exceed `isize::MAX`. This
267    ///     constraint prevents overflow when calculating the `count` parameter to
268    ///     [`.offset()`] regardless of the starting point due to past offsets.
269    ///
270    /// * The product of non-zero axis lengths must not exceed `isize::MAX`.
271    /// 
272    /// * Strides must be non-negative.
273    ///
274    /// This function can use debug assertions to check some of these requirements,
275    /// but it's not a complete check.
276    ///
277    /// [`.offset()`]: https://doc.rust-lang.org/stable/std/primitive.pointer.html#method.offset
278    pub unsafe fn from_shape_ptr<Sh>(shape: Sh, ptr: *mut A) -> Self
279    where
280        Sh: Into<StrideShape<D>>,
281    {
282        let shape = shape.into();
283        let dim = shape.dim;
284        if cfg!(debug_assertions) {
285            assert!(!ptr.is_null(), "The pointer must be non-null.");
286            if let Strides::Custom(strides) = &shape.strides {
287                dimension::strides_non_negative(strides).unwrap();
288                dimension::max_abs_offset_check_overflow::<A, _>(&dim, strides).unwrap();
289            } else {
290                dimension::size_of_shape_checked(&dim).unwrap();
291            }
292        }
293        let strides = shape.strides.strides_for_dim(&dim);
294        RawArrayViewMut::new_(ptr, dim, strides)
295    }
296
297    /// Converts to a non-mutable `RawArrayView`.
298    #[inline]
299    pub(crate) fn into_raw_view(self) -> RawArrayView<A, D> {
300        unsafe { RawArrayView::new(self.ptr, self.dim, self.strides) }
301    }
302
303    /// Converts to a read-only view of the array.
304    ///
305    /// # Safety
306    ///
307    /// From a safety standpoint, this is equivalent to dereferencing a raw
308    /// pointer for every element in the array. You must ensure that all of the
309    /// data is valid, ensure that the pointer is aligned, and choose the
310    /// correct lifetime.
311    #[inline]
312    pub unsafe fn deref_into_view<'a>(self) -> ArrayView<'a, A, D> {
313        debug_assert!(
314            is_aligned(self.ptr.as_ptr()),
315            "The pointer must be aligned."
316        );
317        ArrayView::new(self.ptr, self.dim, self.strides)
318    }
319
320    /// Converts to a mutable view of the array.
321    ///
322    /// # Safety
323    ///
324    /// From a safety standpoint, this is equivalent to dereferencing a raw
325    /// pointer for every element in the array. You must ensure that all of the
326    /// data is valid, ensure that the pointer is aligned, and choose the
327    /// correct lifetime.
328    #[inline]
329    pub unsafe fn deref_into_view_mut<'a>(self) -> ArrayViewMut<'a, A, D> {
330        debug_assert!(
331            is_aligned(self.ptr.as_ptr()),
332            "The pointer must be aligned."
333        );
334        ArrayViewMut::new(self.ptr, self.dim, self.strides)
335    }
336
337    /// Split the array view along `axis` and return one array pointer strictly
338    /// before the split and one array pointer after the split.
339    ///
340    /// **Panics** if `axis` or `index` is out of bounds.
341    pub fn split_at(self, axis: Axis, index: Ix) -> (Self, Self) {
342        let (left, right) = self.into_raw_view().split_at(axis, index);
343        unsafe {
344            (
345                Self::new(left.ptr, left.dim, left.strides),
346                Self::new(right.ptr, right.dim, right.strides),
347            )
348        }
349    }
350
351    /// Cast the raw pointer of the raw array view to a different type
352    ///
353    /// **Panics** if element size is not compatible.
354    ///
355    /// Lack of panic does not imply it is a valid cast. The cast works the same
356    /// way as regular raw pointer casts.
357    ///
358    /// While this method is safe, for the same reason as regular raw pointer
359    /// casts are safe, access through the produced raw view is only possible
360    /// in an unsafe block or function.
361    pub fn cast<B>(self) -> RawArrayViewMut<B, D> {
362        assert_eq!(
363            mem::size_of::<B>(),
364            mem::size_of::<A>(),
365            "size mismatch in raw view cast"
366        );
367        let ptr = self.ptr.cast::<B>();
368        unsafe { RawArrayViewMut::new(ptr, self.dim, self.strides) }
369    }
370}
371
372impl<T, D> RawArrayViewMut<Complex<T>, D>
373where
374    D: Dimension,
375{
376    /// Splits the view into views of the real and imaginary components of the
377    /// elements.
378    pub fn split_complex(self) -> Complex<RawArrayViewMut<T, D>> {
379        let Complex { re, im } = self.into_raw_view().split_complex();
380        unsafe {
381            Complex {
382                re: RawArrayViewMut::new(re.ptr, re.dim, re.strides),
383                im: RawArrayViewMut::new(im.ptr, im.dim, im.strides),
384            }
385        }
386    }
387}