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