ndarray/dimension/
dynindeximpl.rs

1use crate::imp_prelude::*;
2use std::hash::{Hash, Hasher};
3use std::ops::{Deref, DerefMut, Index, IndexMut};
4use alloc::vec;
5use alloc::boxed::Box;
6use alloc::vec::Vec;
7const CAP: usize = 4;
8
9/// T is usize or isize
10#[derive(Debug)]
11enum IxDynRepr<T> {
12    Inline(u32, [T; CAP]),
13    Alloc(Box<[T]>),
14}
15
16impl<T> Deref for IxDynRepr<T> {
17    type Target = [T];
18    fn deref(&self) -> &[T] {
19        match *self {
20            IxDynRepr::Inline(len, ref ar) => {
21                debug_assert!(len as usize <= ar.len());
22                unsafe { ar.get_unchecked(..len as usize) }
23            }
24            IxDynRepr::Alloc(ref ar) => ar,
25        }
26    }
27}
28
29impl<T> DerefMut for IxDynRepr<T> {
30    fn deref_mut(&mut self) -> &mut [T] {
31        match *self {
32            IxDynRepr::Inline(len, ref mut ar) => {
33                debug_assert!(len as usize <= ar.len());
34                unsafe { ar.get_unchecked_mut(..len as usize) }
35            }
36            IxDynRepr::Alloc(ref mut ar) => ar,
37        }
38    }
39}
40
41/// The default is equivalent to `Self::from(&[0])`.
42impl Default for IxDynRepr<Ix> {
43    fn default() -> Self {
44        Self::copy_from(&[0])
45    }
46}
47
48use num_traits::Zero;
49
50impl<T: Copy + Zero> IxDynRepr<T> {
51    pub fn copy_from(x: &[T]) -> Self {
52        if x.len() <= CAP {
53            let mut arr = [T::zero(); CAP];
54            arr[..x.len()].copy_from_slice(x);
55            IxDynRepr::Inline(x.len() as _, arr)
56        } else {
57            Self::from(x)
58        }
59    }
60}
61
62impl<T: Copy + Zero> IxDynRepr<T> {
63    // make an Inline or Alloc version as appropriate
64    fn from_vec_auto(v: Vec<T>) -> Self {
65        if v.len() <= CAP {
66            Self::copy_from(&v)
67        } else {
68            Self::from_vec(v)
69        }
70    }
71}
72
73impl<T: Copy> IxDynRepr<T> {
74    fn from_vec(v: Vec<T>) -> Self {
75        IxDynRepr::Alloc(v.into_boxed_slice())
76    }
77
78    fn from(x: &[T]) -> Self {
79        Self::from_vec(x.to_vec())
80    }
81}
82
83impl<T: Copy> Clone for IxDynRepr<T> {
84    fn clone(&self) -> Self {
85        match *self {
86            IxDynRepr::Inline(len, arr) => IxDynRepr::Inline(len, arr),
87            _ => Self::from(&self[..]),
88        }
89    }
90}
91
92impl<T: Eq> Eq for IxDynRepr<T> {}
93
94impl<T: PartialEq> PartialEq for IxDynRepr<T> {
95    fn eq(&self, rhs: &Self) -> bool {
96        match (self, rhs) {
97            (&IxDynRepr::Inline(slen, ref sarr), &IxDynRepr::Inline(rlen, ref rarr)) => {
98                slen == rlen
99                    && (0..CAP as usize)
100                        .filter(|&i| i < slen as usize)
101                        .all(|i| sarr[i] == rarr[i])
102            }
103            _ => self[..] == rhs[..],
104        }
105    }
106}
107
108impl<T: Hash> Hash for IxDynRepr<T> {
109    fn hash<H: Hasher>(&self, state: &mut H) {
110        Hash::hash(&self[..], state)
111    }
112}
113
114/// Dynamic dimension or index type.
115///
116/// Use `IxDyn` directly. This type implements a dynamic number of
117/// dimensions or indices. Short dimensions are stored inline and don't need
118/// any dynamic memory allocation.
119#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
120pub struct IxDynImpl(IxDynRepr<Ix>);
121
122impl IxDynImpl {
123    pub(crate) fn insert(&self, i: usize) -> Self {
124        let len = self.len();
125        debug_assert!(i <= len);
126        IxDynImpl(if len < CAP {
127            let mut out = [1; CAP];
128            out[0..i].copy_from_slice(&self[0..i]);
129            out[i + 1..=len].copy_from_slice(&self[i..len]);
130            IxDynRepr::Inline((len + 1) as u32, out)
131        } else {
132            let mut out = Vec::with_capacity(len + 1);
133            out.extend_from_slice(&self[0..i]);
134            out.push(1);
135            out.extend_from_slice(&self[i..len]);
136            IxDynRepr::from_vec(out)
137        })
138    }
139
140    fn remove(&self, i: usize) -> Self {
141        IxDynImpl(match self.0 {
142            IxDynRepr::Inline(0, _) => IxDynRepr::Inline(0, [0; CAP]),
143            IxDynRepr::Inline(1, _) => IxDynRepr::Inline(0, [0; CAP]),
144            IxDynRepr::Inline(2, ref arr) => {
145                let mut out = [0; CAP];
146                out[0] = arr[1 - i];
147                IxDynRepr::Inline(1, out)
148            }
149            ref ixdyn => {
150                let len = ixdyn.len();
151                let mut result = IxDynRepr::copy_from(&ixdyn[..len - 1]);
152                for j in i..len - 1 {
153                    result[j] = ixdyn[j + 1]
154                }
155                result
156            }
157        })
158    }
159}
160
161impl<'a> From<&'a [Ix]> for IxDynImpl {
162    #[inline]
163    fn from(ix: &'a [Ix]) -> Self {
164        IxDynImpl(IxDynRepr::copy_from(ix))
165    }
166}
167
168impl From<Vec<Ix>> for IxDynImpl {
169    #[inline]
170    fn from(ix: Vec<Ix>) -> Self {
171        IxDynImpl(IxDynRepr::from_vec_auto(ix))
172    }
173}
174
175impl<J> Index<J> for IxDynImpl
176where
177    [Ix]: Index<J>,
178{
179    type Output = <[Ix] as Index<J>>::Output;
180    fn index(&self, index: J) -> &Self::Output {
181        &self.0[index]
182    }
183}
184
185impl<J> IndexMut<J> for IxDynImpl
186where
187    [Ix]: IndexMut<J>,
188{
189    fn index_mut(&mut self, index: J) -> &mut Self::Output {
190        &mut self.0[index]
191    }
192}
193
194impl Deref for IxDynImpl {
195    type Target = [Ix];
196    #[inline]
197    fn deref(&self) -> &[Ix] {
198        &self.0
199    }
200}
201
202impl DerefMut for IxDynImpl {
203    #[inline]
204    fn deref_mut(&mut self) -> &mut [Ix] {
205        &mut self.0
206    }
207}
208
209impl<'a> IntoIterator for &'a IxDynImpl {
210    type Item = &'a Ix;
211    type IntoIter = <&'a [Ix] as IntoIterator>::IntoIter;
212    #[inline]
213    fn into_iter(self) -> Self::IntoIter {
214        self[..].iter()
215    }
216}
217
218impl RemoveAxis for Dim<IxDynImpl> {
219    fn remove_axis(&self, axis: Axis) -> Self {
220        debug_assert!(axis.index() < self.ndim());
221        Dim::new(self.ix().remove(axis.index()))
222    }
223}
224
225impl IxDyn {
226    /// Create a new dimension value with `n` axes, all zeros
227    #[inline]
228    pub fn zeros(n: usize) -> IxDyn {
229        const ZEROS: &[usize] = &[0; 4];
230        if n <= ZEROS.len() {
231            Dim(&ZEROS[..n])
232        } else {
233            Dim(vec![0; n])
234        }
235    }
236}