mdarray/
array.rs

1use core::borrow::{Borrow, BorrowMut};
2use core::fmt::{Debug, Formatter, Result};
3use core::hash::{Hash, Hasher};
4use core::mem::{self, ManuallyDrop, MaybeUninit};
5use core::ops::{Deref, DerefMut, Index, IndexMut};
6use core::ptr;
7
8use crate::dim::Const;
9use crate::expr::{self, IntoExpr, Iter, Map, Zip};
10use crate::expr::{Apply, Expression, FromExpression, IntoExpression};
11use crate::index::SliceIndex;
12use crate::layout::{Dense, Layout};
13use crate::shape::{ConstShape, Shape};
14use crate::slice::Slice;
15use crate::tensor::Tensor;
16use crate::traits::Owned;
17use crate::view::{View, ViewMut};
18
19/// Multidimensional array with constant-sized dimensions and inline allocation.
20#[derive(Clone, Copy, Default)]
21#[repr(transparent)]
22pub struct Array<T, S: ConstShape>(pub S::Inner<T>);
23
24impl<T, S: ConstShape> Array<T, S> {
25    /// Creates an array from the given element.
26    #[inline]
27    pub fn from_elem(elem: T) -> Self
28    where
29        T: Clone,
30    {
31        Self::from_expr(expr::from_elem(S::default(), elem))
32    }
33
34    /// Creates an array with the results from the given function.
35    #[inline]
36    pub fn from_fn<F: FnMut(&[usize]) -> T>(f: F) -> Self {
37        Self::from_expr(expr::from_fn(S::default(), f))
38    }
39
40    /// Converts an array with a single element into the contained value.
41    ///
42    /// # Panics
43    ///
44    /// Panics if the array length is not equal to one.
45    #[inline]
46    pub fn into_scalar(self) -> T {
47        assert!(self.len() == 1, "invalid length");
48
49        self.into_shape::<()>().0
50    }
51
52    /// Converts the array into a reshaped array, which must have the same length.
53    ///
54    /// # Panics
55    ///
56    /// Panics if the array length is changed.
57    #[inline]
58    pub fn into_shape<I: ConstShape>(self) -> Array<T, I> {
59        assert!(I::default().len() == self.len(), "length must not change");
60
61        let me = ManuallyDrop::new(self);
62
63        unsafe { mem::transmute_copy(&me) }
64    }
65
66    /// Returns an array with the same shape, and the given closure applied to each element.
67    #[inline]
68    pub fn map<U, F: FnMut(T) -> U>(self, f: F) -> Array<U, S> {
69        self.apply(f)
70    }
71
72    /// Creates an array with uninitialized elements.
73    #[inline]
74    pub fn uninit() -> Array<MaybeUninit<T>, S> {
75        let array = <MaybeUninit<Self>>::uninit();
76
77        unsafe { mem::transmute_copy(&array) }
78    }
79
80    /// Creates an array with elements set to zero.
81    ///
82    /// Zero elements are created using `Default::default()`.
83    #[inline]
84    pub fn zeros() -> Self
85    where
86        T: Default,
87    {
88        let mut array = Self::uninit();
89
90        array.expr_mut().for_each(|x| {
91            _ = x.write(T::default());
92        });
93
94        unsafe { array.assume_init() }
95    }
96
97    #[inline]
98    fn from_expr<E: Expression<Item = T>>(expr: E) -> Self {
99        struct DropGuard<'a, T, S: ConstShape> {
100            array: &'a mut MaybeUninit<Array<T, S>>,
101            index: usize,
102        }
103
104        impl<T, S: ConstShape> Drop for DropGuard<'_, T, S> {
105            #[inline]
106            fn drop(&mut self) {
107                let ptr = self.array.as_mut_ptr() as *mut T;
108
109                unsafe {
110                    ptr::slice_from_raw_parts_mut(ptr, self.index).drop_in_place();
111                }
112            }
113        }
114
115        // Ensure that the shape is valid.
116        _ = expr.shape().with_dims(|dims| S::from_dims(dims));
117
118        let mut array = MaybeUninit::uninit();
119        let mut guard = DropGuard { array: &mut array, index: 0 };
120
121        let ptr = guard.array.as_mut_ptr() as *mut E::Item;
122
123        expr.for_each(|x| unsafe {
124            ptr.add(guard.index).write(x);
125            guard.index += 1;
126        });
127
128        mem::forget(guard);
129
130        unsafe { array.assume_init() }
131    }
132}
133
134impl<T, S: ConstShape> Array<MaybeUninit<T>, S> {
135    /// Converts the array element type from `MaybeUninit<T>` to `T`.
136    ///
137    /// # Safety
138    ///
139    /// All elements in the array must be initialized, or the behavior is undefined.
140    #[inline]
141    pub unsafe fn assume_init(self) -> Array<T, S> {
142        unsafe { mem::transmute_copy(&self) }
143    }
144}
145
146impl<'a, T, U, S: ConstShape> Apply<U> for &'a Array<T, S> {
147    type Output<F: FnMut(&'a T) -> U> = Map<Self::IntoExpr, F>;
148    type ZippedWith<I: IntoExpression, F: FnMut((&'a T, I::Item)) -> U> =
149        Map<Zip<Self::IntoExpr, I::IntoExpr>, F>;
150
151    #[inline]
152    fn apply<F: FnMut(&'a T) -> U>(self, f: F) -> Self::Output<F> {
153        self.expr().map(f)
154    }
155
156    #[inline]
157    fn zip_with<I: IntoExpression, F>(self, expr: I, f: F) -> Self::ZippedWith<I, F>
158    where
159        F: FnMut((&'a T, I::Item)) -> U,
160    {
161        self.expr().zip(expr).map(f)
162    }
163}
164
165impl<'a, T, U, S: ConstShape> Apply<U> for &'a mut Array<T, S> {
166    type Output<F: FnMut(&'a mut T) -> U> = Map<Self::IntoExpr, F>;
167    type ZippedWith<I: IntoExpression, F: FnMut((&'a mut T, I::Item)) -> U> =
168        Map<Zip<Self::IntoExpr, I::IntoExpr>, F>;
169
170    #[inline]
171    fn apply<F: FnMut(&'a mut T) -> U>(self, f: F) -> Self::Output<F> {
172        self.expr_mut().map(f)
173    }
174
175    #[inline]
176    fn zip_with<I: IntoExpression, F>(self, expr: I, f: F) -> Self::ZippedWith<I, F>
177    where
178        F: FnMut((&'a mut T, I::Item)) -> U,
179    {
180        self.expr_mut().zip(expr).map(f)
181    }
182}
183
184impl<T, U, S: ConstShape> Apply<U> for Array<T, S> {
185    type Output<F: FnMut(T) -> U> = Array<U, S>;
186    type ZippedWith<I: IntoExpression, F: FnMut((T, I::Item)) -> U> = Array<U, S>;
187
188    #[inline]
189    fn apply<F: FnMut(T) -> U>(self, f: F) -> Array<U, S> {
190        Array::from_expr(self.into_expr().map(f))
191    }
192
193    #[inline]
194    fn zip_with<I: IntoExpression, F>(self, expr: I, f: F) -> Array<U, S>
195    where
196        F: FnMut((T, I::Item)) -> U,
197    {
198        Array::from_expr(self.into_expr().zip(expr).map(f))
199    }
200}
201
202impl<T, U: ?Sized, S: ConstShape> AsMut<U> for Array<T, S>
203where
204    Slice<T, S>: AsMut<U>,
205{
206    #[inline]
207    fn as_mut(&mut self) -> &mut U {
208        (**self).as_mut()
209    }
210}
211
212impl<T, U: ?Sized, S: ConstShape> AsRef<U> for Array<T, S>
213where
214    Slice<T, S>: AsRef<U>,
215{
216    #[inline]
217    fn as_ref(&self) -> &U {
218        (**self).as_ref()
219    }
220}
221
222macro_rules! impl_as_mut_ref {
223    (($($xyz:tt),+), $array:tt) => {
224        impl<T, $(const $xyz: usize),+> AsMut<Array<T, ($(Const<$xyz>,)+)>> for $array {
225            #[inline]
226            fn as_mut(&mut self) -> &mut Array<T, ($(Const<$xyz>,)+)> {
227                unsafe { &mut *(self as *mut Self as *mut Array<T, ($(Const<$xyz>,)+)>) }
228            }
229        }
230
231        impl<T, $(const $xyz: usize),+> AsRef<Array<T, ($(Const<$xyz>,)+)>> for $array {
232            #[inline]
233            fn as_ref(&self) -> &Array<T, ($(Const<$xyz>,)+)> {
234                unsafe { &*(self as *const Self as *const Array<T, ($(Const<$xyz>,)+)>) }
235            }
236        }
237    };
238}
239
240impl_as_mut_ref!((X), [T; X]);
241impl_as_mut_ref!((X, Y), [[T; Y]; X]);
242impl_as_mut_ref!((X, Y, Z), [[[T; Z]; Y]; X]);
243impl_as_mut_ref!((X, Y, Z, W), [[[[T; W]; Z]; Y]; X]);
244impl_as_mut_ref!((X, Y, Z, W, U), [[[[[T; U]; W]; Z]; Y]; X]);
245impl_as_mut_ref!((X, Y, Z, W, U, V), [[[[[[T; V]; U]; W]; Z]; Y]; X]);
246
247impl<T, S: ConstShape> Borrow<Slice<T, S>> for Array<T, S> {
248    #[inline]
249    fn borrow(&self) -> &Slice<T, S> {
250        self
251    }
252}
253
254impl<T, S: ConstShape> BorrowMut<Slice<T, S>> for Array<T, S> {
255    #[inline]
256    fn borrow_mut(&mut self) -> &mut Slice<T, S> {
257        self
258    }
259}
260
261impl<T: Debug, S: ConstShape> Debug for Array<T, S> {
262    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
263        (**self).fmt(f)
264    }
265}
266
267impl<T, S: ConstShape> Deref for Array<T, S> {
268    type Target = Slice<T, S>;
269
270    #[inline]
271    fn deref(&self) -> &Self::Target {
272        _ = S::default().checked_len().expect("invalid length");
273
274        unsafe { &*(self as *const Self as *const Slice<T, S>) }
275    }
276}
277
278impl<T, S: ConstShape> DerefMut for Array<T, S> {
279    #[inline]
280    fn deref_mut(&mut self) -> &mut Self::Target {
281        _ = S::default().checked_len().expect("invalid length");
282
283        unsafe { &mut *(self as *mut Self as *mut Slice<T, S>) }
284    }
285}
286
287impl<T, S: ConstShape> From<Tensor<T, S>> for Array<T, S> {
288    #[inline]
289    fn from(value: Tensor<T, S>) -> Self {
290        Self::from_expr(value.into_expr())
291    }
292}
293
294impl<'a, T: 'a + Clone, S: ConstShape, L: Layout, I> From<I> for Array<T, S>
295where
296    I: IntoExpression<IntoExpr = View<'a, T, S, L>>,
297{
298    #[inline]
299    fn from(value: I) -> Self {
300        Self::from_expr(value.into_expr().cloned())
301    }
302}
303
304macro_rules! impl_from_array {
305    (($($xyz:tt),+), $array:tt) => {
306        impl<T: Clone $(,const $xyz: usize)+> From<&$array> for Array<T, ($(Const<$xyz>,)+)> {
307            #[inline]
308            fn from(array: &$array) -> Self {
309                Self(array.clone())
310            }
311        }
312
313        impl<T $(,const $xyz: usize)+> From<Array<T, ($(Const<$xyz>,)+)>> for $array {
314            #[inline]
315            fn from(array: Array<T, ($(Const<$xyz>,)+)>) -> Self {
316                array.0
317            }
318        }
319
320        impl<T $(,const $xyz: usize)+> From<$array> for Array<T, ($(Const<$xyz>,)+)> {
321            #[inline]
322            fn from(array: $array) -> Self {
323                Self(array)
324            }
325        }
326    };
327}
328
329impl_from_array!((X), [T; X]);
330impl_from_array!((X, Y), [[T; Y]; X]);
331impl_from_array!((X, Y, Z), [[[T; Z]; Y]; X]);
332impl_from_array!((X, Y, Z, W), [[[[T; W]; Z]; Y]; X]);
333impl_from_array!((X, Y, Z, W, U), [[[[[T; U]; W]; Z]; Y]; X]);
334impl_from_array!((X, Y, Z, W, U, V), [[[[[[T; V]; U]; W]; Z]; Y]; X]);
335
336impl<T, S: ConstShape> FromExpression<T, S> for Array<T, S> {
337    #[inline]
338    fn from_expr<I: IntoExpression<Item = T, Shape = S>>(expr: I) -> Self {
339        Self::from_expr(expr.into_expr())
340    }
341}
342
343impl<T: Hash, S: ConstShape> Hash for Array<T, S> {
344    #[inline]
345    fn hash<H: Hasher>(&self, state: &mut H) {
346        (**self).hash(state)
347    }
348}
349
350impl<T, S: ConstShape, I: SliceIndex<T, S, Dense>> Index<I> for Array<T, S> {
351    type Output = I::Output;
352
353    #[inline]
354    fn index(&self, index: I) -> &I::Output {
355        index.index(self)
356    }
357}
358
359impl<T, S: ConstShape, I: SliceIndex<T, S, Dense>> IndexMut<I> for Array<T, S> {
360    #[inline]
361    fn index_mut(&mut self, index: I) -> &mut I::Output {
362        index.index_mut(self)
363    }
364}
365
366impl<'a, T, S: ConstShape> IntoExpression for &'a Array<T, S> {
367    type Shape = S;
368    type IntoExpr = View<'a, T, S>;
369
370    #[inline]
371    fn into_expr(self) -> Self::IntoExpr {
372        self.expr()
373    }
374}
375
376impl<'a, T, S: ConstShape> IntoExpression for &'a mut Array<T, S> {
377    type Shape = S;
378    type IntoExpr = ViewMut<'a, T, S>;
379
380    #[inline]
381    fn into_expr(self) -> Self::IntoExpr {
382        self.expr_mut()
383    }
384}
385
386impl<T, S: ConstShape> IntoExpression for Array<T, S> {
387    type Shape = S;
388    type IntoExpr = IntoExpr<Array<ManuallyDrop<T>, S>>;
389
390    #[inline]
391    fn into_expr(self) -> Self::IntoExpr {
392        _ = S::default().checked_len().expect("invalid length");
393
394        let me = ManuallyDrop::new(self);
395
396        unsafe { IntoExpr::new(mem::transmute_copy(&me)) }
397    }
398}
399
400impl<'a, T, S: ConstShape> IntoIterator for &'a Array<T, S> {
401    type Item = &'a T;
402    type IntoIter = Iter<View<'a, T, S>>;
403
404    #[inline]
405    fn into_iter(self) -> Self::IntoIter {
406        self.iter()
407    }
408}
409
410impl<'a, T, S: ConstShape> IntoIterator for &'a mut Array<T, S> {
411    type Item = &'a mut T;
412    type IntoIter = Iter<ViewMut<'a, T, S>>;
413
414    #[inline]
415    fn into_iter(self) -> Self::IntoIter {
416        self.iter_mut()
417    }
418}
419
420impl<T, S: ConstShape> IntoIterator for Array<T, S> {
421    type Item = T;
422    type IntoIter = Iter<IntoExpr<Array<ManuallyDrop<T>, S>>>;
423
424    #[inline]
425    fn into_iter(self) -> Self::IntoIter {
426        self.into_expr().into_iter()
427    }
428}
429
430impl<T, S: ConstShape> Owned<T, S> for Array<T, S> {
431    type WithConst<const N: usize> = S::WithConst<T, N, Self>;
432
433    #[inline]
434    fn clone_from_slice(&mut self, slice: &Slice<T, S>)
435    where
436        T: Clone,
437    {
438        self.assign(slice);
439    }
440}
mdarray/array.rs

mdarray/
array.rs
