nuvec_core/
view.rs

1//! View geometry: shapes, strides and offsets of arrays that share another array's data.
2//!
3//! Pure functions of `(shape, strides)`, used both when a view is recorded (to know its
4//! layout without evaluating) and when it is evaluated.
5
6use smallvec::SmallVec;
7
8use crate::shape::{self, Shape, Strides};
9use crate::{Error, Result};
10
11/// One slice specification per dimension: `start:stop:step` with Python semantics.
12///
13/// `None` bounds mean "from the beginning/end" in the step's direction. Negative indices
14/// count from the end, and out-of-range bounds are clamped.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct SliceArg {
17    /// First index, or `None` for the start in the step's direction.
18    pub start: Option<isize>,
19    /// One past the last index, or `None` for the end in the step's direction.
20    pub stop: Option<isize>,
21    /// Distance between selected indices; must not be zero.
22    pub step: isize,
23}
24
25impl SliceArg {
26    /// The whole dimension (`::`).
27    pub const fn full() -> Self {
28        Self {
29            start: None,
30            stop: None,
31            step: 1,
32        }
33    }
34
35    /// `start:stop:step`.
36    pub const fn new(start: Option<isize>, stop: Option<isize>, step: isize) -> Self {
37        Self { start, stop, step }
38    }
39}
40
41impl From<std::ops::Range<isize>> for SliceArg {
42    fn from(r: std::ops::Range<isize>) -> Self {
43        Self::new(Some(r.start), Some(r.end), 1)
44    }
45}
46
47impl From<std::ops::RangeFrom<isize>> for SliceArg {
48    fn from(r: std::ops::RangeFrom<isize>) -> Self {
49        Self::new(Some(r.start), None, 1)
50    }
51}
52
53impl From<std::ops::RangeTo<isize>> for SliceArg {
54    fn from(r: std::ops::RangeTo<isize>) -> Self {
55        Self::new(None, Some(r.end), 1)
56    }
57}
58
59impl From<std::ops::RangeFull> for SliceArg {
60    fn from(_: std::ops::RangeFull) -> Self {
61        Self::full()
62    }
63}
64
65/// A single index `i`, keeping the dimension with size 1 (`i:i+1`, or the last element for
66/// `-1`).
67impl From<isize> for SliceArg {
68    fn from(i: isize) -> Self {
69        Self::new(Some(i), if i == -1 { None } else { Some(i + 1) }, 1)
70    }
71}
72
73/// A normalised slice of one dimension.
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
75pub(crate) struct DimSlice {
76    pub(crate) start: usize,
77    pub(crate) step: isize,
78    pub(crate) len: usize,
79}
80
81/// Normalises one slice per dimension like Python's `slice.indices`.
82pub(crate) fn normalize_slice(shape: &[usize], args: &[SliceArg]) -> Result<Vec<DimSlice>> {
83    if args.len() != shape.len() {
84        return Err(Error::InvalidArgument(format!(
85            "slice needs one specification per dimension: got {} for shape {shape:?}",
86            args.len()
87        )));
88    }
89    shape
90        .iter()
91        .zip(args)
92        .map(|(&n, arg)| {
93            if arg.step == 0 {
94                return Err(Error::InvalidArgument("slice step must not be zero".into()));
95            }
96            let n = n as isize;
97            let (lower, upper) = if arg.step > 0 { (0, n) } else { (-1, n - 1) };
98            let clamp = |i: isize| {
99                let i = if i < 0 { i + n } else { i };
100                i.clamp(lower, upper)
101            };
102            let start = arg
103                .start
104                .map_or(if arg.step > 0 { lower } else { upper }, clamp);
105            let stop = arg
106                .stop
107                .map_or(if arg.step > 0 { upper } else { lower }, clamp);
108            let span = if arg.step > 0 {
109                stop - start
110            } else {
111                start - stop
112            };
113            let len = if span > 0 {
114                (span + arg.step.abs() - 1) / arg.step.abs()
115            } else {
116                0
117            };
118            Ok(if len == 0 {
119                DimSlice {
120                    start: 0,
121                    step: 1,
122                    len: 0,
123                }
124            } else {
125                DimSlice {
126                    start: start as usize,
127                    step: arg.step,
128                    len: len as usize,
129                }
130            })
131        })
132        .collect()
133}
134
135/// Strides and offset delta of a slice of a layout with `strides`.
136pub(crate) fn slice_geometry(strides: &[isize], slices: &[DimSlice]) -> (Strides, isize) {
137    let out = strides
138        .iter()
139        .zip(slices)
140        .map(|(&s, d)| s * d.step)
141        .collect();
142    let delta = strides
143        .iter()
144        .zip(slices)
145        .map(|(&s, d)| s * d.start as isize)
146        .sum();
147    (out, delta)
148}
149
150/// The broadcast shape of `a` and `b` (NumPy rules), or a shape mismatch for `op`.
151pub(crate) fn broadcast_shapes(op: &'static str, a: &[usize], b: &[usize]) -> Result<Shape> {
152    let ndim = a.len().max(b.len());
153    let dim = |s: &[usize], i: usize| {
154        if i + s.len() < ndim {
155            1
156        } else {
157            s[i + s.len() - ndim]
158        }
159    };
160    (0..ndim)
161        .map(|i| match (dim(a, i), dim(b, i)) {
162            (x, y) if x == y => Ok(x),
163            (1, y) => Ok(y),
164            (x, 1) => Ok(x),
165            _ => Err(Error::ShapeMismatch {
166                op,
167                lhs: a.to_vec(),
168                rhs: b.to_vec(),
169            }),
170        })
171        .collect()
172}
173
174/// Strides that present a layout of `shape`/`strides` as the larger shape `out`, or an
175/// error if `shape` cannot broadcast to `out`.
176pub(crate) fn broadcast_strides(
177    shape: &[usize],
178    strides: &[isize],
179    out: &[usize],
180) -> Result<Strides> {
181    let mismatch = || Error::ShapeMismatch {
182        op: "broadcast_to",
183        lhs: shape.to_vec(),
184        rhs: out.to_vec(),
185    };
186    if shape.len() > out.len() {
187        return Err(mismatch());
188    }
189    let lead = out.len() - shape.len();
190    (0..out.len())
191        .map(|i| {
192            if i < lead {
193                return Ok(0);
194            }
195            match (shape[i - lead], out[i]) {
196                (x, y) if x == y => Ok(strides[i - lead]),
197                (1, _) => Ok(0),
198                _ => Err(mismatch()),
199            }
200        })
201        .collect()
202}
203
204/// Strides for viewing a layout of `shape`/`strides` as `new_shape` without copying, or
205/// `None` if the data must be copied. NumPy's `_attempt_nocopy_reshape` (C order).
206///
207/// Requires both shapes to have the same element count.
208pub(crate) fn nocopy_reshape(
209    shape: &[usize],
210    strides: &[isize],
211    new_shape: &[usize],
212) -> Option<Strides> {
213    if shape.contains(&0) || new_shape.iter().all(|&d| d == 1) {
214        return Some(shape::row_contiguous_strides(new_shape));
215    }
216    let (old_dims, old_strides): (Vec<usize>, Vec<isize>) = shape
217        .iter()
218        .zip(strides)
219        .filter(|(d, _)| **d != 1)
220        .map(|(&d, &s)| (d, s))
221        .unzip();
222    let mut out: Strides = SmallVec::from_elem(1, new_shape.len());
223    let (mut oi, mut oj, mut ni, mut nj) = (0, 1, 0, 1);
224    while ni < new_shape.len() && oi < old_dims.len() {
225        let mut np = new_shape[ni];
226        let mut op = old_dims[oi];
227        while np != op {
228            if np < op {
229                np *= new_shape[nj];
230                nj += 1;
231            } else {
232                op *= old_dims[oj];
233                oj += 1;
234            }
235        }
236        // The merged old dimensions must be contiguous among themselves.
237        for ok in oi..oj - 1 {
238            if old_strides[ok] != old_dims[ok + 1] as isize * old_strides[ok + 1] {
239                return None;
240            }
241        }
242        out[nj - 1] = old_strides[oj - 1];
243        for nk in (ni + 1..nj).rev() {
244            out[nk - 1] = out[nk] * new_shape[nk] as isize;
245        }
246        ni = nj;
247        nj += 1;
248        oi = oj;
249        oj += 1;
250    }
251    Some(out)
252}
253
254/// Element count of `requested` with at most one `-1` inferred to give `size` elements.
255pub(crate) fn infer_shape(size: usize, requested: &[isize]) -> Result<Shape> {
256    let invalid = || {
257        Error::InvalidArgument(format!(
258            "cannot reshape {size} elements into shape {requested:?}"
259        ))
260    };
261    if requested.iter().filter(|&&d| d == -1).count() > 1 || requested.iter().any(|&d| d < -1) {
262        return Err(invalid());
263    }
264    let known = requested
265        .iter()
266        .filter(|&&d| d >= 0)
267        .try_fold(1usize, |acc, &d| acc.checked_mul(d as usize))
268        .ok_or_else(invalid)?;
269    let inferred = if requested.contains(&-1) {
270        if known == 0 || !size.is_multiple_of(known) {
271            return Err(invalid());
272        }
273        size / known
274    } else if known == size {
275        0
276    } else {
277        return Err(invalid());
278    };
279    Ok(requested
280        .iter()
281        .map(|&d| if d == -1 { inferred } else { d as usize })
282        .collect())
283}
284
285/// Shape and strides with axes reordered: output axis `i` is input axis `axes[i]`.
286pub(crate) fn permute(
287    shape: &[usize],
288    strides: &[isize],
289    axes: &[usize],
290) -> Result<(Shape, Strides)> {
291    let mut seen = vec![false; shape.len()];
292    let valid = axes.len() == shape.len()
293        && axes
294            .iter()
295            .all(|&a| a < shape.len() && !std::mem::replace(&mut seen[a], true));
296    if !valid {
297        return Err(Error::InvalidArgument(format!(
298            "transpose axes {axes:?} are not a permutation of {} dimensions",
299            shape.len()
300        )));
301    }
302    Ok((
303        axes.iter().map(|&a| shape[a]).collect(),
304        axes.iter().map(|&a| strides[a]).collect(),
305    ))
306}
307
308/// `axis` as an index into `ndim` dimensions; negative axes count from the end.
309pub(crate) fn normalize_axis(axis: isize, ndim: usize) -> Result<usize> {
310    let a = if axis < 0 { axis + ndim as isize } else { axis };
311    if (0..ndim as isize).contains(&a) {
312        Ok(a as usize)
313    } else {
314        Err(Error::InvalidArgument(format!(
315            "axis {axis} is out of range for {ndim} dimensions"
316        )))
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    fn slices(n: usize, start: Option<isize>, stop: Option<isize>, step: isize) -> DimSlice {
325        normalize_slice(&[n], &[SliceArg::new(start, stop, step)]).unwrap()[0]
326    }
327
328    #[test]
329    fn python_slice_semantics() {
330        let d = |start, step, len| DimSlice { start, step, len };
331        assert_eq!(slices(5, None, None, 1), d(0, 1, 5));
332        assert_eq!(slices(5, None, None, -1), d(4, -1, 5));
333        assert_eq!(slices(5, Some(-2), None, 1), d(3, 1, 2));
334        assert_eq!(slices(5, Some(1), Some(100), 2), d(1, 2, 2));
335        assert_eq!(slices(5, None, Some(-6), -2), d(4, -2, 3));
336        assert_eq!(slices(5, Some(3), Some(1), 1).len, 0);
337        assert_eq!(slices(5, Some(0), Some(5), 10), d(0, 10, 1));
338        assert!(normalize_slice(&[5], &[SliceArg::new(None, None, 0)]).is_err());
339        assert!(normalize_slice(&[5, 2], &[SliceArg::full()]).is_err());
340    }
341
342    #[test]
343    fn broadcasting() {
344        assert_eq!(
345            broadcast_shapes("t", &[3, 1], &[4]).unwrap().as_slice(),
346            &[3, 4]
347        );
348        assert_eq!(
349            broadcast_shapes("t", &[], &[2, 0]).unwrap().as_slice(),
350            &[2, 0]
351        );
352        assert!(broadcast_shapes("t", &[2, 3], &[4]).is_err());
353        assert_eq!(
354            broadcast_strides(&[3, 1], &[1, 1], &[2, 3, 4])
355                .unwrap()
356                .as_slice(),
357            &[0, 1, 0]
358        );
359        assert!(broadcast_strides(&[3], &[1], &[4]).is_err());
360    }
361
362    #[test]
363    fn reshape_views() {
364        // Contiguous: always a view.
365        assert_eq!(
366            nocopy_reshape(&[2, 3], &[3, 1], &[3, 2])
367                .unwrap()
368                .as_slice(),
369            &[2, 1]
370        );
371        assert_eq!(
372            nocopy_reshape(&[2, 3, 4], &[12, 4, 1], &[6, 4])
373                .unwrap()
374                .as_slice(),
375            &[4, 1]
376        );
377        // Transposed 2x3 -> 6: needs a copy.
378        assert!(nocopy_reshape(&[3, 2], &[1, 3], &[6]).is_none());
379        // Transposed, but splitting one dimension is fine.
380        assert_eq!(
381            nocopy_reshape(&[4, 2], &[1, 4], &[2, 2, 2])
382                .unwrap()
383                .as_slice(),
384            &[2, 1, 4]
385        );
386        // Unit dimensions are free (the stride of a size-1 dimension is arbitrary).
387        let s = nocopy_reshape(&[3, 2], &[1, 3], &[3, 1, 2]).unwrap();
388        assert_eq!((s[0], s[2]), (1, 3));
389    }
390
391    #[test]
392    fn shape_inference_and_axes() {
393        assert_eq!(infer_shape(6, &[-1, 2]).unwrap().as_slice(), &[3, 2]);
394        assert!(infer_shape(6, &[-1, -1]).is_err());
395        assert!(infer_shape(6, &[4, -1]).is_err());
396        assert!(infer_shape(0, &[0, -1]).is_err());
397        assert!(infer_shape(6, &[5]).is_err());
398        assert_eq!(normalize_axis(-1, 3).unwrap(), 2);
399        assert!(normalize_axis(3, 3).is_err());
400        assert!(permute(&[2, 3], &[3, 1], &[0, 0]).is_err());
401        let (s, st) = permute(&[2, 3, 4], &[12, 4, 1], &[2, 0, 1]).unwrap();
402        assert_eq!(
403            (s.as_slice(), st.as_slice()),
404            (&[4, 2, 3][..], &[1, 12, 4][..])
405        );
406    }
407}
nuvec_core/view.rs

nuvec_core/
view.rs
