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}