ndarray_layout/
lib.rs

1#![doc = include_str!("../README.md")]
2#![deny(warnings, missing_docs)]
3
4/// An array layout allow N dimensions inlined.
5pub struct ArrayLayout<const N: usize> {
6    ndim: usize,
7    content: Union<N>,
8}
9
10unsafe impl<const N: usize> Send for ArrayLayout<N> {}
11unsafe impl<const N: usize> Sync for ArrayLayout<N> {}
12
13union Union<const N: usize> {
14    ptr: NonNull<usize>,
15    _inlined: (isize, [usize; N], [isize; N]),
16}
17
18impl<const N: usize> Clone for ArrayLayout<N> {
19    #[inline]
20    fn clone(&self) -> Self {
21        Self::new(self.shape(), self.strides(), self.offset())
22    }
23}
24
25impl<const N: usize> PartialEq for ArrayLayout<N> {
26    #[inline]
27    fn eq(&self, other: &Self) -> bool {
28        self.ndim == other.ndim && self.content().as_slice() == other.content().as_slice()
29    }
30}
31
32impl<const N: usize> Eq for ArrayLayout<N> {}
33
34impl<const N: usize> Drop for ArrayLayout<N> {
35    fn drop(&mut self) {
36        if let Some(ptr) = self.ptr_allocated() {
37            unsafe { dealloc(ptr.cast().as_ptr(), layout(self.ndim)) }
38        }
39    }
40}
41
42/// 元信息存储顺序。
43#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
44pub enum Endian {
45    /// 大端序，范围更大的维度在元信息中更靠前的位置。
46    BigEndian,
47    /// 小端序，范围更小的维度在元信息中更靠前的位置。
48    LittleEndian,
49}
50
51impl<const N: usize> ArrayLayout<N> {
52    /// Creates a new Layout with the given shape, strides, and offset.
53    ///
54    /// ```rust
55    /// # use ndarray_layout::ArrayLayout;
56    /// let layout = ArrayLayout::<4>::new(&[2, 3, 4], &[12, -4, 1], 20);
57    /// assert_eq!(layout.offset(), 20);
58    /// assert_eq!(layout.shape(), &[2, 3, 4]);
59    /// assert_eq!(layout.strides(), &[12, -4, 1]);
60    /// ```
61    pub fn new(shape: &[usize], strides: &[isize], offset: isize) -> Self {
62        // check
63        assert_eq!(
64            shape.len(),
65            strides.len(),
66            "shape and strides must have the same length"
67        );
68
69        let mut ans = Self::with_ndim(shape.len());
70        let mut content = ans.content_mut();
71        content.set_offset(offset);
72        content.copy_shape(shape);
73        content.copy_strides(strides);
74        ans
75    }
76
77    /// Creates a new contiguous Layout with the given shape.
78    ///
79    /// ```rust
80    /// # use ndarray_layout::{Endian, ArrayLayout};
81    /// let layout = ArrayLayout::<4>::new_contiguous(&[2, 3, 4], Endian::LittleEndian, 4);
82    /// assert_eq!(layout.offset(), 0);
83    /// assert_eq!(layout.shape(), &[2, 3, 4]);
84    /// assert_eq!(layout.strides(), &[4, 8, 24]);
85    /// ```
86    pub fn new_contiguous(shape: &[usize], endian: Endian, element_size: usize) -> Self {
87        let mut ans = Self::with_ndim(shape.len());
88        let mut content = ans.content_mut();
89        content.set_offset(0);
90        content.copy_shape(shape);
91        let mut mul = element_size as isize;
92        let push = |i| {
93            content.set_stride(i, mul);
94            mul *= shape[i] as isize;
95        };
96        match endian {
97            Endian::BigEndian => (0..shape.len()).rev().for_each(push),
98            Endian::LittleEndian => (0..shape.len()).for_each(push),
99        }
100        ans
101    }
102
103    /// Gets offset.
104    #[inline]
105    pub const fn ndim(&self) -> usize {
106        self.ndim
107    }
108
109    /// Gets offset.
110    #[inline]
111    pub fn offset(&self) -> isize {
112        self.content().offset()
113    }
114
115    /// Gets shape.
116    #[inline]
117    pub fn shape(&self) -> &[usize] {
118        self.content().shape()
119    }
120
121    /// Gets strides.
122    #[inline]
123    pub fn strides(&self) -> &[isize] {
124        self.content().strides()
125    }
126
127    /// Calculates the number of elements in the array.
128    ///
129    /// ```rust
130    /// # use ndarray_layout::{Endian::BigEndian, ArrayLayout};
131    /// let layout = ArrayLayout::<4>::new_contiguous(&[2, 3, 4], BigEndian, 20);
132    /// assert_eq!(layout.num_elements(), 24);
133    /// ```
134    #[inline]
135    pub fn num_elements(&self) -> usize {
136        self.shape().iter().product()
137    }
138
139    /// Calculates the offset of element at the given `index`.
140    ///
141    /// ```rust
142    /// # use ndarray_layout::{Endian::BigEndian, ArrayLayout};
143    /// let layout = ArrayLayout::<4>::new_contiguous(&[2, 3, 4], BigEndian, 4);
144    /// assert_eq!(layout.element_offset(22, BigEndian), 88); // 88 <- (22 % 4 * 4) + (22 / 4 % 3 * 16) + (22 / 4 / 3 % 2 * 48)
145    /// ```
146    pub fn element_offset(&self, index: usize, endian: Endian) -> isize {
147        fn offset_forwards(
148            mut rem: usize,
149            shape: impl IntoIterator<Item = usize>,
150            strides: impl IntoIterator<Item = isize>,
151        ) -> isize {
152            let mut ans = 0;
153            for (d, s) in zip(shape, strides) {
154                ans += s * (rem % d) as isize;
155                rem /= d
156            }
157            ans
158        }
159
160        let shape = self.shape().iter().cloned();
161        let strides = self.strides().iter().cloned();
162        self.offset()
163            + match endian {
164                Endian::BigEndian => offset_forwards(index, shape.rev(), strides.rev()),
165                Endian::LittleEndian => offset_forwards(index, shape, strides),
166            }
167    }
168
169    /// Calculates the range of data in bytes to determine the location of the memory area that the array needs to access.
170    pub fn data_range(&self) -> RangeInclusive<isize> {
171        let content = self.content();
172        let mut start = content.offset();
173        let mut end = content.offset();
174        for (&d, s) in zip(content.shape(), content.strides()) {
175            use std::cmp::Ordering::{Equal, Greater, Less};
176            let i = d as isize - 1;
177            match s.cmp(&0) {
178                Equal => {}
179                Less => start += s * i,
180                Greater => end += s * i,
181            }
182        }
183        start..=end
184    }
185}
186
187mod fmt;
188mod transform;
189pub use transform::{BroadcastArg, IndexArg, MergeArg, SliceArg, Split, TileArg};
190
191use std::{
192    alloc::{Layout, alloc, dealloc},
193    iter::zip,
194    ops::RangeInclusive,
195    ptr::{NonNull, copy_nonoverlapping},
196    slice::from_raw_parts,
197};
198
199impl<const N: usize> ArrayLayout<N> {
200    #[inline]
201    fn ptr_allocated(&self) -> Option<NonNull<usize>> {
202        const { assert!(N > 0) }
203        if self.ndim > N {
204            Some(unsafe { self.content.ptr })
205        } else {
206            None
207        }
208    }
209
210    #[inline]
211    fn content(&self) -> Content<false> {
212        Content {
213            ptr: self
214                .ptr_allocated()
215                .unwrap_or(unsafe { NonNull::new_unchecked(&self.content as *const _ as _) }),
216            ndim: self.ndim,
217        }
218    }
219
220    #[inline]
221    fn content_mut(&mut self) -> Content<true> {
222        Content {
223            ptr: self
224                .ptr_allocated()
225                .unwrap_or(unsafe { NonNull::new_unchecked(&self.content as *const _ as _) }),
226            ndim: self.ndim,
227        }
228    }
229
230    /// Create a new ArrayLayout with the given dimensions.
231    #[inline]
232    fn with_ndim(ndim: usize) -> Self {
233        Self {
234            ndim,
235            content: if ndim <= N {
236                Union {
237                    _inlined: (0, [0; N], [0; N]),
238                }
239            } else {
240                Union {
241                    ptr: unsafe { NonNull::new_unchecked(alloc(layout(ndim)).cast()) },
242                }
243            },
244        }
245    }
246}
247
248struct Content<const MUT: bool> {
249    ptr: NonNull<usize>,
250    ndim: usize,
251}
252
253impl<const MUT: bool> Content<MUT> {
254    #[inline]
255    fn as_slice(&self) -> &[usize] {
256        unsafe { from_raw_parts(self.ptr.as_ptr(), 1 + self.ndim * 2) }
257    }
258
259    #[inline]
260    fn offset(&self) -> isize {
261        unsafe { self.ptr.cast().read() }
262    }
263
264    #[inline]
265    fn shape<'a>(&self) -> &'a [usize] {
266        unsafe { from_raw_parts(self.ptr.add(1).as_ptr(), self.ndim) }
267    }
268
269    #[inline]
270    fn strides<'a>(&self) -> &'a [isize] {
271        unsafe { from_raw_parts(self.ptr.add(1 + self.ndim).cast().as_ptr(), self.ndim) }
272    }
273}
274
275impl Content<true> {
276    #[inline]
277    fn set_offset(&mut self, val: isize) {
278        unsafe { self.ptr.cast().write(val) }
279    }
280
281    #[inline]
282    fn set_shape(&mut self, idx: usize, val: usize) {
283        assert!(idx < self.ndim);
284        unsafe { self.ptr.add(1 + idx).write(val) }
285    }
286
287    #[inline]
288    fn set_stride(&mut self, idx: usize, val: isize) {
289        assert!(idx < self.ndim);
290        unsafe { self.ptr.add(1 + idx + self.ndim).cast().write(val) }
291    }
292
293    #[inline]
294    fn copy_shape(&mut self, val: &[usize]) {
295        assert!(val.len() == self.ndim);
296        unsafe { copy_nonoverlapping(val.as_ptr(), self.ptr.add(1).as_ptr(), self.ndim) }
297    }
298
299    #[inline]
300    fn copy_strides(&mut self, val: &[isize]) {
301        assert!(val.len() == self.ndim);
302        unsafe {
303            copy_nonoverlapping(
304                val.as_ptr(),
305                self.ptr.add(1 + self.ndim).cast().as_ptr(),
306                self.ndim,
307            )
308        }
309    }
310}
311
312#[inline]
313fn layout(ndim: usize) -> Layout {
314    Layout::array::<usize>(1 + ndim * 2).unwrap()
315}
ndarray_layout/lib.rs

ndarray_layout/
lib.rs
