Skip to main content

rustymath_core/
structure.rs

1//! Plans for structural operations: take and drop, rotate and reverse, windows, transpose and catenate. A plan gives the result's layout, and `Steps` that read the result's items from the argument's items in order. `gather::walk` copies a compact array by a plan, and basedpl copies its other arrays by the same plan. Each plan checks its arguments.
2use crate::agreement::{Agreement, Mapping};
3use crate::error::{Error, Kind};
4use crate::gather::{IndexSource, IndexValue, Steps, Table, Writes};
5use crate::layout::{Keys, Lanes, Layout, agreed, generated_len, names_conflict, shown, strides};
6use crate::{buffer::View, float::Float};
7use std::{borrow::Cow, sync::Arc};
8
9/// A structural result's layout, and the steps that read its items from the argument. The steps' lengths multiply to the result's item count, outermost first. One entry can cover several result axes.
10#[derive(Clone, Debug)]
11pub struct Plan { pub layout: Layout, pub steps: Vec<Steps> }
12
13/// A count that an argument gives, such as a take's or a diagonal's offset: its sign and magnitude.
14#[derive(Clone, Copy, Debug)]
15pub struct Count { pub negative: bool, pub magnitude: usize }
16
17impl Count {
18    /// The count `n`. A magnitude past `usize` saturates.
19    pub fn of(n: i64) -> Self { Count { negative: n < 0, magnitude: usize::try_from(n.unsigned_abs()).unwrap_or(usize::MAX) } }
20}
21
22/// Written index parts, before the selected axes and their lengths are known.
23#[derive(Clone, Debug, PartialEq)]
24pub enum IndexPart<T> { Axis(T), All, NewAxis, Rest }
25impl<T> IndexPart<T> {
26    /// The part with its axis payload borrowed.
27    pub fn as_ref(&self) -> IndexPart<&T> {
28        match self {
29            IndexPart::Axis(t) => IndexPart::Axis(t), IndexPart::All => IndexPart::All, IndexPart::NewAxis => IndexPart::NewAxis,
30            IndexPart::Rest => IndexPart::Rest
31        }
32    }
33    /// The part with `f` applied to its axis payload. The structural markers stay as they are.
34    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> IndexPart<U> {
35        match self {
36            IndexPart::Axis(t) => IndexPart::Axis(f(t)), IndexPart::All => IndexPart::All, IndexPart::NewAxis => IndexPart::NewAxis,
37            IndexPart::Rest => IndexPart::Rest
38        }
39    }
40    /// The part with fallible `f` applied to its axis payload, or `f`'s error.
41    pub fn try_map<U, E>(self, f: impl FnOnce(T) -> Result<U, E>) -> Result<IndexPart<U>, E> {
42        Ok(match self {
43            IndexPart::Axis(t) => IndexPart::Axis(f(t)?), IndexPart::All => IndexPart::All, IndexPart::NewAxis => IndexPart::NewAxis,
44            IndexPart::Rest => IndexPart::Rest
45        })
46    }
47    /// The axis payload, or `None` for a structural marker.
48    pub fn axis(&self) -> Option<&T> { match self { IndexPart::Axis(t) => Some(t), _ => None } }
49}
50
51/// Expands Rest once geometry is known, with one source claim per consuming part. Axis payloads identify the original written part.
52pub fn index_claims<'a>(rank: usize, parts: &[IndexPart<(Option<&'a str>, usize)>]) -> Result<(Vec<IndexPart<usize>>, Vec<(Option<&'a str>, usize)>), Error> {
53    let used: usize = parts.iter().map(|p| match p { IndexPart::Axis((_, n)) => *n, IndexPart::All => 1, _ => 0 }).sum();
54    if used > rank { return Err(Error::new(Kind::Index, "too many axis selectors")); }
55    if parts.iter().filter(|p| matches!(p, IndexPart::Rest)).count() > 1 { return Err(Error::new(Kind::Index, "only one rest selector is allowed")); }
56    let mut written = Vec::new();
57    let mut claims = Vec::new();
58    for (i, part) in parts.iter().enumerate() {
59        match part {
60            IndexPart::Axis(claim) => {
61                written.push(IndexPart::Axis(i));
62                claims.push(*claim);
63            }
64            IndexPart::All => {
65                written.push(IndexPart::All);
66                claims.push((None, 1));
67            }
68            IndexPart::NewAxis => written.push(IndexPart::NewAxis),
69            IndexPart::Rest => {
70                for _ in used..rank {
71                    written.push(IndexPart::All);
72                    claims.push((None, 1));
73                }
74            }
75        }
76    }
77    Ok((written, claims))
78}
79
80/// Routes written parts to source axes without padding them with implicit selectors. Axis payloads give their name and consumed width; returned payloads identify the original part.
81pub fn addressed_parts(
82    layout: &Layout, parts: &[IndexPart<(Option<&str>, usize)>], explicit: Option<&[usize]>) -> Result<Vec<(usize, IndexPart<usize>)>, Error> {
83    let rank = layout.shape().len();
84    let (written, claims) = index_claims(rank, parts)?;
85    let axes = match explicit {
86        Some(axes) if axes.len() == claims.len() => axes.to_vec(),
87        Some(_) => return Err(Error::new(Kind::Length, "one axis is needed for each consuming selector")), None => addressed(layout, &claims)?
88    };
89    if axes.iter().zip(&claims).any(|(&axis, &(_, width))| axis.checked_add(width).is_none_or(|end| end > rank)) {
90        return Err(Error::new(Kind::Index, "a selector spans beyond the array's axes"));
91    }
92    let last = axes.iter().zip(&claims).map(|(&a, &(_, n))| a + n).max().unwrap_or(0);
93    let mut next = 0;
94    let mut routed = written.into_iter().map(|part| {
95        let axis = axes.get(next).copied().unwrap_or(last);
96        if !matches!(part, IndexPart::NewAxis) { next += 1; }
97        (axis, part)
98    }).collect::<Vec<_>>();
99    routed.sort_by_key(|(axis, _)| *axis);
100    Ok(routed)
101}
102
103/// Plans a selection by written parts on `layout`, as `⌷` and the Python frontends read brackets. `claims` give each part's axis name and width, as for `addressed_parts`, and `explicit` its axes. `decode(i, axis)` gives the selector of written part `i` at source `axis`, with whether it reads one position. Axes that no part reaches stay whole, and a part wider than one axis reads its axes as one. Also gives whether the selection reads one item: one position on every source axis, in order. `broadcast` pairs index arrays as `select` does. `err` converts this module's errors to the caller's.
104pub fn plan_parts<E>(
105    layout: &Layout, claims: &[IndexPart<(Option<&str>, usize)>], explicit: Option<&[usize]>, broadcast: bool,
106    decode: impl FnMut(usize, usize) -> Result<(Selector, bool), E>, err: impl Fn(Error) -> E) -> Result<(Plan, bool), E> {
107    let r = routes(layout, claims, explicit, decode, &err)?;
108    Ok((plan_routes(layout, &r, broadcast).map_err(&err)?, r.one))
109}
110
111/// The selectors of a selection by written parts (`routes`): one for each axis of the layout that `sources` regroups the source into, with the source axes that each reads, and whether the selection reads one item, one position on every source axis in order.
112#[derive(Clone, Debug)]
113pub struct Routes { pub selectors: Vec<Selector>, pub sources: Vec<Vec<usize>>, pub one: bool }
114
115/// The routes of a selection by written parts on `layout`, as `plan_parts` takes them. `claims` give each part's axis name and width, as for `addressed_parts`, and `explicit` its axes. `decode(i, axis)` gives the selector of written part `i` at source `axis`, with whether it reads one position. Axes that no part reaches stay whole (`Selector::All`), and a part wider than one axis reads its axes as one.
116pub fn routes<E>(
117    layout: &Layout, claims: &[IndexPart<(Option<&str>, usize)>], explicit: Option<&[usize]>,
118    mut decode: impl FnMut(usize, usize) -> Result<(Selector, bool), E>, err: impl Fn(Error) -> E) -> Result<Routes, E> {
119    let shape = layout.shape();
120    let (mut one, mut next, mut sources, mut selectors) = (claims.iter().any(|p| matches!(p, IndexPart::Axis(_))), 0, Vec::<Vec<usize>>::new(), Vec::new());
121    for (axis, part) in addressed_parts(layout, claims, explicit).map_err(&err)? {
122        one &= next == axis;
123        for a in next..axis {
124            sources.push(vec![a]);
125            selectors.push(Selector::All);
126        }
127        next = axis;
128        let (selector, width, single) = match part {
129            IndexPart::NewAxis => {
130                selectors.push(Selector::NewAxis);
131                one = false;
132                continue;
133            }
134            IndexPart::Axis(i) => {
135                let IndexPart::Axis((_, width)) = claims[i] else { unreachable!("an axis refers to its written claim") };
136                let (selector, single) = decode(i, axis)?;
137                (selector, width, single)
138            }
139            _ => (Selector::All, 1, false),
140        };
141        one &= single;
142        sources.push((axis..axis + width).collect());
143        selectors.push(selector);
144        next += width;
145    }
146    one &= next == shape.len();
147    sources.extend((next..shape.len()).map(|a| vec![a]));
148    Ok(Routes { selectors, sources, one })
149}
150
151/// The plan of the selection that `routes` gives on `layout`, as `plan_parts` plans it.
152pub fn plan_routes(layout: &Layout, routes: &Routes, broadcast: bool) -> Result<Plan, Error> {
153    planned(route_layout(layout, routes)?.as_ref(), routes.selectors.clone(), broadcast)
154}
155
156fn route_layout<'a>(layout: &'a Layout, routes: &Routes) -> Result<Cow<'a, Layout>, Error> {
157    let shape = layout.shape();
158    if routes.sources.len() == shape.len() { return Ok(Cow::Borrowed(layout)); }
159    regrouped(layout, routes.sources.iter().map(|s| s.iter().map(|&a| shape[a]).product()).collect(), &routes.sources, &[]).map(Cow::Owned)
160}
161
162/// One independent axis selector. Positions are zero-based and already normalized; a range can walk backwards. `All` keeps the whole axis, whatever its length. An index array replaces its axis by the axes of its layout. Its positions, a [`Table`], lie within the axis or are `Steps::FILL`, which reads the caller's fill. The code that builds an index array checks its positions, as `index_positions` and `index_table` do. `NewAxis` inserts a length-one axis without consuming an input axis.
163#[derive(Clone, Debug)]
164pub enum Selector { At(usize), Range { start: usize, len: usize, step: isize }, All, Indices { layout: Layout, positions: Table }, NewAxis }
165
166/// Selects with one selector for each leading axis of `layout`, in order, keeping the axes that no selector mentions. A single position removes its axis, and a range keeps it in place. An index array's axes replace its axis. An unnamed vector keeps the axis's name and the keys at the positions it selects if none repeats, as a named vector without keys of its own keeps those keys, and an unnamed index array of another rank has plain axes. Index arrays pair, as arguments pair in arithmetic: all of them with `broadcast`, as `.[]` reads them, and otherwise those that share an axis name, as `⌷` reads them. Paired index arrays put their axes where the first of them was. A range pairs with an index array that has its axis's name. Positions that step evenly along each axis read as strides (`strided`), whether one index array gives them or paired arrays sum to them. One index array whose positions read a stretch of its axis past the ends, as edge modes and positions that count back from the end give them, reads as one padded or rotated entry.
167pub fn select(layout: &Layout, selectors: Vec<Selector>, broadcast: bool) -> Result<Plan, Error> { planned(layout, selectors, broadcast) }
168
169fn planned(layout: &Layout, selectors: Vec<Selector>, broadcast: bool) -> Result<Plan, Error> {
170    let (shape, stride) = (layout.shape(), strides(layout.shape()));
171    let mut parts = Vec::with_capacity(shape.len());
172    let mut axis = 0;
173    for selector in selectors {
174        if matches!(selector, Selector::NewAxis) { parts.push(Part::new(Steps::along(1, 0), Layout::from([1]), 0, 1)); continue; }
175        if axis == shape.len() { return Err(Error::new(Kind::Index, "too many axis selectors")); }
176        parts.push(Part::of(selector, layout, axis, stride[axis])?);
177        axis += 1;
178    }
179    parts.extend((axis..shape.len()).map(|a| Part::new(Steps::along(shape[a], 1), layout.axes(a..a + 1), stride[a], shape[a])));
180    let names: Vec<Arc<str>> = parts.iter().filter(|p| p.array).flat_map(|p| p.axes.names().iter().flatten().cloned()).collect();
181    for part in parts.iter_mut().filter(|p| !p.array && p.axes.name(0).is_some_and(|n| names.contains(n))) { part.array = true; }
182    // The index arrays that pair, each group in order.
183    let mut groups: Vec<Vec<usize>> = Vec::new();
184    let shares = |p: usize, q: usize| parts[p].axes.names().iter().flatten().any(|n| parts[q].axes.names().contains(&Some(n.clone())));
185    for p in (0..parts.len()).filter(|&p| parts[p].array) {
186        let hits: Vec<usize> = (0..groups.len()).filter(|&g| broadcast || groups[g].iter().any(|&q| shares(p, q))).collect();
187        let Some((&first, rest)) = hits.split_first() else { groups.push(vec![p]); continue; };
188        for &g in rest.iter().rev() {
189            let moved = groups.remove(g);
190            groups[first].extend(moved);
191        }
192        groups[first].push(p);
193        groups[first].sort_unstable();
194    }
195    let mut group_of = vec![usize::MAX; parts.len()];
196    for (g, members) in groups.iter().enumerate() { for &p in members { group_of[p] = g; } }
197    // At most one entry for each axis that a part reads along, or one for a part with none.
198    let (mut result, mut steps) = (Layout::EMPTY, Vec::with_capacity(parts.iter().map(|p| p.axes.shape().len().max(1)).sum()));
199    for p in 0..parts.len() {
200        // Paired index arrays read where the first of them was.
201        if let Some(group) = groups.get(group_of[p]).filter(|g| g.len() > 1) {
202            if group[0] == p { result = result.concat(&joined(&parts, group, broadcast, &mut steps)?); }
203            continue;
204        }
205        let part = &mut parts[p];
206        result = result.concat(part.alone.as_ref().unwrap_or(&part.axes));
207        let entry = std::mem::replace(&mut part.steps, Steps::along(0, 0)).scaled(part.stride);
208        if part.array { strided(entry, part.axes.shape(), Some((part.size, part.stride)), &mut steps); } else { steps.push(entry); }
209    }
210    generated_len(result.shape())?;
211    Ok(Plan { layout: result, steps })
212}
213
214/// One selector's share of a plan: positions along its axis, before scaling by the axis's `stride`, the axis's `size`, and the result axes they run along. An index array that pairs with nothing takes `alone` as its axes, where it has them.
215struct Part { steps: Steps, axes: Layout, stride: usize, size: usize, array: bool, alone: Option<Layout> }
216
217impl Part {
218    fn new(steps: Steps, axes: Layout, stride: usize, size: usize) -> Self { Part { steps, axes, stride, size, array: false, alone: None } }
219
220    fn of(selector: Selector, layout: &Layout, axis: usize, stride: usize) -> Result<Self, Error> {
221        let size = layout.shape()[axis];
222        let outside = || Error::new(Kind::Index, "axis selector is outside its axis");
223        match selector {
224            Selector::At(i) if i < size => Ok(Part::new(Steps::Stride { start: i, len: 1, step: 0 }, Layout::EMPTY, stride, size)),
225            Selector::At(_) => Err(outside()),
226            Selector::All => Self::of(Selector::Range { start: 0, len: size, step: 1 }, layout, axis, stride),
227            Selector::Range { start, len, step } => {
228                if len > 0 && !(start < size && start.checked_add_signed(step.saturating_mul((len - 1) as isize)).is_some_and(|i| i < size)) {
229                    return Err(outside());
230                }
231                let steps = Steps::Stride { start, len, step };
232                let axes = layout.select(axis, steps.offsets())?.axes(axis..axis + 1);
233                Ok(Part::new(steps, axes, stride, size))
234            }
235            Selector::Indices { layout: within, positions } => {
236                if generated_len(within.shape())? != positions.len() {
237                    return Err(Error::new(
238                        Kind::Length, format!("index array shape {:?} needs {} positions, got {}", within.shape(), within.count(), positions.len())));
239                }
240                let steps = Steps::Table(positions);
241                let selected = layout.select(axis, steps.offsets())?;
242                Ok(Part { steps, alone: alone(&within, &selected, axis)?, axes: within, stride, size, array: true })
243            }
244            Selector::NewAxis => unreachable!("a new axis reads no positions"),
245        }
246    }
247}
248
249/// The axes that an index array laid out as `within` takes where it pairs with nothing, from `selected`, the layout whose `axis` holds the positions that it reads: a vector without names takes that axis, a named vector without keys takes its keys, and an index array of another rank without names has plain axes. `None` keeps `within`.
250fn alone(within: &Layout, selected: &Layout, axis: usize) -> Result<Option<Layout>, Error> {
251    Ok(match within.shape().len() {
252        1 if !within.has_names() => Some(selected.axes(axis..axis + 1)),
253        1 if within.keys(0).is_none() => Some(within.clone().with_keys(vec![selected.keys(axis).cloned()])?),
254        _ if !within.has_names() => Some(Layout::from(within.shape().to_vec())), _ => None
255    })
256}
257
258/// The axes that paired index arrays agree on, with the entries that read them pushed onto `steps`, one for each run of those axes: the fewest consecutive axes that hold all the axes of each array in the run. An array alone in its run reads through its own positions, as an array that pairs with nothing does. Where every array in a run steps evenly through its positions (`Steps::even`), and the run's axes have no keys, the run reads as one `Stride` for each of its axes, whose step is the sum of each array's step along that axis times the array's stride, and no table is built. Otherwise the arrays add their scaled positions into one table, where a fill in any of them makes a fill. An array that lies along the whole run adds its positions in chunks, and one that broadcasts along it adds them in the runs that a walk through its mapping reaches. The table reads as strides where its offsets are affine (`strided`).
259fn joined(parts: &[Part], group: &[usize], broadcast: bool, steps: &mut Vec<Steps>) -> Result<Layout, Error> {
260    let mut frame = parts[group[0]].axes.clone();
261    let mut pairs = vec![(0..frame.shape().len()).collect::<Vec<_>>()];
262    for &p in &group[1..] {
263        let axes = Agreement::paired_axes(&frame, &parts[p].axes, broadcast);
264        frame = Agreement::paired(&frame, &parts[p].axes, &axes)?.layout;
265        pairs.push(axes);
266    }
267    let rank = frame.shape().len();
268    // A run ends after axis `a` unless some array has axes on both sides of it. An array with no axes reads in the first run.
269    let crosses = |a: usize| pairs.iter().any(|axes| axes.iter().any(|&b| b <= a) && axes.iter().any(|&b| b > a));
270    let mut start = 0;
271    for end in (1..=rank).filter(|&end| end == rank || !crosses(end - 1)).chain((rank == 0).then_some(0)) {
272        let run = start..end;
273        start = end;
274        let lens = &frame.shape()[run.clone()];
275        let members = || (0..group.len()).filter(|&m| pairs[m].first().map_or(run.start == 0, |a| run.contains(a)));
276        let mut lone = members();
277        if let (Some(m), None) = (lone.next(), lone.next()) && pairs[m].iter().copied().eq(run.clone()) {
278            let Part { steps: own, stride, size, .. } = &parts[group[m]];
279            strided(own.clone().scaled(*stride), lens, Some((*size, *stride)), steps);
280            continue;
281        }
282        let at = steps.len();
283        steps.extend(lens.iter().map(|&len| Steps::Stride { start: 0, len, step: 0 }));
284        let mut first = 0;
285        let summed = !lens.is_empty() && run.clone().all(|a| frame.keys(a).is_none())
286            && members().all(|m| {
287                let Part { steps: own, axes, stride, .. } = &parts[group[m]];
288                let Some(Steps::Stride { start, step, .. }) = own.even() else { return false };
289                first += start * stride;
290                // The positions step through the array's own axes in row-major order, so each axis steps by the array's step times the axis's width.
291                let mut width = step * *stride as isize;
292                for (&a, &len) in pairs[m].iter().zip(axes.shape()).rev() {
293                    if len > 1 && let Steps::Stride { step, .. } = &mut steps[at + a - run.start] { *step += width; }
294                    width *= len as isize;
295                }
296                true
297            });
298        if summed { Steps::settle(&mut steps[at..], first); continue; }
299        steps.truncate(at);
300        let within = if run.len() == rank { Cow::Borrowed(&frame) } else { Cow::Owned(frame.axes(run.clone())) };
301        let len = generated_len(lens)?;
302        let mut table = vec![0; len];
303        for m in members() {
304            let Part { steps: own, axes, stride, .. } = &parts[group[m]];
305            let add = |t: &mut usize, o: Option<usize>| {
306                *t = match o { Some(o) if *t != Steps::FILL => *t + o * stride, _ => Steps::FILL }
307            };
308            let shifted: Vec<usize>;
309            let along = if run.start == 0 { &pairs[m] } else { shifted = pairs[m].iter().map(|&a| a - run.start).collect(); &shifted };
310            match (Mapping::new(axes, along, &within, len), own) {
311                (Mapping::Linear(1), Steps::Table(positions)) => positions.chunks(|j0, offsets| {
312                    for (t, &o) in table[j0..].iter_mut().zip(offsets) { add(t, Some(o).filter(|&o| o != Steps::FILL)) }
313                }),
314                (map, _) => {
315                    let (base, tables) = map.to_walk(lens);
316                    Writes::Walk { base, steps: &tables }.runs_while(&mut |k, reads| {
317                        for (t, &j) in table[k..].iter_mut().zip(reads) { add(t, (j != Steps::FILL).then(|| own.get(j)).flatten()) }
318                        true
319                    });
320                }
321            }
322        }
323        strided(Steps::table(table), lens, None, steps);
324    }
325    Ok(frame)
326}
327
328/// Pushes `entry`, which reads the axes of `lens`, onto `into`: as one `Stride` for each of those axes where its offsets are affine (`Table::affine`), as one `Padded` or `Rotated` entry where `lens` is one axis read past the ends of the source axis whose size and stride `axis` gives (`Table::edged`), or else as it is. A copy reads strides in place, and the transpose of a copy through strides reads back where a table would scatter.
329fn strided(entry: Steps, lens: &[usize], axis: Option<(usize, usize)>, into: &mut Vec<Steps>) {
330    if let Steps::Table(table) = &entry {
331        if table.affine(lens, into) { return; }
332        if let ([_], Some((size, stride))) = (lens, axis) && let Some(edged) = table.edged(size, stride) { return into.push(edged); }
333    }
334    into.push(entry);
335}
336
337/// The name of the one named axis of `layout`, where it has exactly one.
338pub fn sole_name(layout: &Layout) -> Option<&str> {
339    let mut names = layout.names().iter().flatten();
340    names.next().filter(|_| names.next().is_none()).map(|n| &**n)
341}
342
343/// The first axis of `layout` that each selector addresses, given the name of each selector's one named axis, where it has exactly one, and the number of axes it covers, where more axes than remain cover the rest, as a coordinate past the last axis does. A selector whose name `layout` has addresses the axis with that name, unless another selector has the same name. Selectors with the same name take free axes in order, as unnamed selectors and NumPy's index arrays do. Each other selector takes the first free axes whose names don't conflict with its name. Either way a selector claims every axis it covers. Two selectors that claim one axis, or a selector with no free axis, give an Index error.
344pub fn addressed(layout: &Layout, selectors: &[(Option<&str>, usize)]) -> Result<Vec<usize>, Error> {
345    let rank = layout.shape().len();
346    let (mut free, mut axes) = (vec![true; rank], vec![rank; selectors.len()]);
347    let span = |a: usize, width: usize| a..(a + width).min(rank);
348    let names: Vec<_> = selectors.iter().map(|&(n, _)| n.filter(|_| selectors.iter().filter(|s| s.0 == n).count() == 1)).collect();
349    for (k, &(_, width)) in selectors.iter().enumerate() {
350        let Some((n, a)) = names[k].and_then(|n| Some((n, layout.axis_named(n)?))) else { continue };
351        if span(a, width).any(|b| !free[b]) { return Err(Error::new(Kind::Index, format!("two selectors address axis {}", shown(n)))); }
352        free[span(a, width)].fill(false);
353        axes[k] = a;
354    }
355    for (k, &(_, width)) in selectors.iter().enumerate() {
356        if axes[k] != rank { continue; }
357        let name = names[k];
358        let fits = |a: usize| a < rank && span(a, width).all(|b| free[b] && !names_conflict(name, layout.name(b).map(|n| &**n)));
359        let missing =
360            || Error::new(Kind::Index, name.map_or_else(|| "too many axis selectors".into(), |n| format!("no axis for a selector named {}", shown(n))));
361        let a = (0..rank).find(|&a| fits(a)).ok_or_else(missing)?;
362        free[span(a, width)].fill(false);
363        axes[k] = a;
364    }
365    Ok(axes)
366}
367
368/// [`select`] with `selector` on `axis` and every other axis whole. A `NewAxis` goes in before `axis`.
369pub fn select_axis(layout: &Layout, axis: usize, selector: Selector) -> Result<Plan, Error> {
370    let selectors: Vec<_> = layout.shape().iter().take(axis).map(|_| Selector::All).chain([selector]).collect();
371    select(layout, selectors, false)
372}
373
374/// The position of key `name` on `axis` of `layout`, or `Steps::FILL` where the axis lacks it: a missing key reads the fill and writes nothing, as a position outside the axis does. BPL and both Python bindings read keys by this rule.
375pub fn key_position(layout: &Layout, axis: usize, name: &str) -> Result<usize, Error> {
376    let keys = layout.keys(axis).ok_or_else(|| Error::new(Kind::Index, "axis has no keys"))?;
377    Ok(keys.position(name).unwrap_or(Steps::FILL))
378}
379
380/// The selector that keys give on `axis` of `layout`, by [`key_position`]: one key's position, which removes the axis, where `within` is `None`, and otherwise an index array of the keys' positions, laid out as `within`.
381pub fn key_selector<'a>(layout: &Layout, axis: usize, names: impl IntoIterator<Item = &'a str>, within: Option<Layout>) -> Result<Selector, Error> {
382    let positions = names.into_iter().map(|n| key_position(layout, axis, n)).collect::<Result<Vec<_>, _>>()?;
383    if within.is_none() && let [p] = positions[..] && p != Steps::FILL { return Ok(Selector::At(p)); }
384    Ok(Selector::Indices { layout: within.unwrap_or(Layout::EMPTY), positions: Table::new(positions) })
385}
386
387/// For each axis of `layout`, the keys that `wanted` names on it and it lacks, each once, in order. `wanted` pairs an axis with the keys that one selector names on it. A write adds these keys before it plans ([`extended`]).
388pub fn missing_keys<'a, K: IntoIterator<Item = &'a str>>(layout: &Layout, wanted: impl IntoIterator<Item = (usize, K)>) -> Result<Vec<Vec<Arc<str>>>, Error> {
389    let mut added: Vec<Vec<Arc<str>>> = vec![vec![]; layout.shape().len()];
390    for (axis, names) in wanted {
391        for name in names {
392            let lacked = added.get_mut(axis).ok_or_else(|| Error::new(Kind::Rank, "too many axis selectors"))?;
393            if layout.keys(axis).is_none_or(|k| k.position(name).is_none()) && !lacked.iter().any(|n| &**n == name) { lacked.push(name.into()); }
394        }
395    }
396    Ok(added)
397}
398
399/// The plan that appends the keys in `added`, a list for each axis, to the axes of `layout`, as [`missing_keys`] gives them. An axis without keys takes them after a position without a key for each of its items. The new positions read `Steps::FILL`, and the others read the items in place.
400pub fn extended(layout: &Layout, added: &[Vec<Arc<str>>]) -> Result<Plan, Error> {
401    let (mut shape, mut keys) = (layout.shape().to_vec(), layout.all_keys());
402    for (axis, names) in added.iter().enumerate().filter(|(_, names)| !names.is_empty()) {
403        let mut all = keys[axis].as_ref().map_or_else(|| vec![None; shape[axis]], |k| k.names().to_vec());
404        all.extend(names.iter().cloned().map(Some));
405        shape[axis] = all.len();
406        keys[axis] = Some(Keys::partial(all)?);
407    }
408    generated_len(&shape)?;
409    let steps = shape.iter().zip(layout.shape()).zip(strides(layout.shape())).map(|((&len, &size), stride)| Steps::clipped(0, len, size, stride)).collect();
410    Ok(Plan { layout: Layout::from(shape).with_keys(keys)?.inherit_names(layout.names().to_vec()), steps })
411}
412
413/// A table that reads the index values in `values` as positions along an axis of length `size`, as [`positions`] reads them. Booleans and integers that all lie inside the axis read in place, each converted as it's read: the one pass that checks them writes nothing, and also finds whether they give evenly spaced positions, so that a walk can read them as one stride. Floats, read within tolerance `t`, and values with any outside the axis become the list of positions that [`positions`] gives, with the fill position for each value that names none.
414pub fn index_table(values: Arc<dyn IndexSource>, size: usize, t: crate::tolerant::Tolerance) -> Result<Table, Error> {
415    /// `None` where some value is outside the axis. Otherwise the first position and the step where the positions are evenly spaced.
416    fn check<I: IndexValue>(values: &[I], size: usize) -> Option<Option<(usize, isize)>> {
417        let (lo, hi) = I::bounds(values);
418        if !values.is_empty() && !(lo >= -(size as i64) && hi < size as i64) { return None; }
419        let at = |k: usize| values[k].position();
420        if values.len() < 2 { return Some(None); }
421        // Values of one sign give evenly spaced positions where they're evenly spaced themselves. Checking the second step first skips the full pass for values in no order.
422        let step = at(1) - at(0);
423        let even = (lo >= 0 || hi < 0) && (values.len() < 3 || at(2) - at(1) == step)
424            && values.windows(2).fold(true, |even, w| even & (w[1].position() - w[0].position() == step));
425        Some(even.then(|| (counted_back(at(0), size), step as isize)))
426    }
427    let checked = match values.view() { View::Floats(_) => None, view => crate::gather::with_index_values!(view, |x| check(x, size)) };
428    match checked { Some(even) => Ok(Table::indices(values, size, even)), None => positions(values.view(), size, t).map(Table::new) }
429}
430
431/// Position `i` on an axis of length `size`, where a negative `i` counts back from the end. Adding `size` to a negative position without a branch keeps a loop over positions vectorised. A position outside the axis is then at least `size`.
432#[inline(always)]
433pub fn position_recipe<F: Float, O: crate::recipe::Ops<F> + ?Sized>(o: &O, value: O::I, size: i64) -> (O::I, O::B) {
434    use crate::recipe::{Cmp, Integer};
435    let (zero, limit) = (o.integer(0), o.integer(size));
436    let negative = o.integer_compare(Cmp::Less, value, zero);
437    let position = o.integer_op(Integer::Add, value, o.integer_select(negative, limit, zero));
438    let valid = o.and(o.not(o.integer_compare(Cmp::Less, position, zero)), o.integer_compare(Cmp::Less, position, limit));
439    (position, valid)
440}
441pub(crate) fn counted_back(i: i64, size: usize) -> usize {
442    let ops = crate::recipe::IntegerOps::<i64>::new(crate::domain::Variant::checked(crate::buffer::DType::Integer(crate::buffer::Width::I64)));
443    position_recipe::<f64, _>(&ops, i as u64, size as i64).0 as usize
444}
445
446/// The positions that index values give along an axis of length `size`. A negative value counts back from the end. A value outside the axis, a float that isn't a whole number within tolerance `t` (`float_position`), or a NaN gives the fill position (`Steps::FILL`), which reads the fill and writes nothing. Items that aren't numbers are a Domain error.
447pub fn positions(view: View<'_>, size: usize, t: crate::tolerant::Tolerance) -> Result<Vec<usize>, Error> {
448    let at = |p: Option<i64>| p.map(|p| counted_back(p, size)).filter(|&p| p < size).unwrap_or(Steps::FILL);
449    Ok(match view {
450        View::Booleans(b) => b.iter().map(|&i| at(Some(i.into()))).collect(),
451        View::Integers(ix) => with_ints!(ix, |i| i.iter().map(|&i| at(Some(i.into()))).collect()),
452        View::Floats(x) => with_floats!(x, |x| x.iter().map(|&y| at(float_position(y, t))).collect()),
453        _ => return Err(Error::new(Kind::Domain, "indices must be real numbers")),
454    })
455}
456
457/// The message of the Domain error for an index value that isn't a whole number.
458pub const NOT_WHOLE: &str = "index must be a whole number";
459
460/// The Domain error for an index value that isn't a whole number.
461fn not_whole() -> Error { Error::new(Kind::Domain, NOT_WHOLE) }
462
463/// `$int` on the items of a Boolean or integer view, `$float` on a float view's items, or the Domain error for any other view.
464macro_rules! index_view {
465    ($view:expr, |$i:ident| $int:expr, |$f:ident| $float:expr) => {
466        match $view {
467            View::Booleans($i) => $int,
468            View::Integers(ix) => with_ints!(ix, |$i| $int),
469            View::Floats(x) => with_floats!(x, |$f| $float),
470            _ => Err(not_whole()),
471        }
472    };
473}
474
475/// The whole number that a float index value names: the one it equals within tolerance `t`, the comparison tolerance, which only Mathish makes nonzero. `None` for a NaN or a value that isn't whole, which reads the fill. A whole number past `i64` saturates, and so lies outside every axis. The one rule for float index values: [`positions`], [`signed_positions`], [`index_table`] and [`Edge::positions`] read floats through it.
476pub fn float_position<F: Float>(y: F, t: crate::tolerant::Tolerance) -> Option<i64> { crate::tolerant::Tolerant::whole(t, y.wide()).map(|n| n as i64) }
477
478/// The signed position that each index value gives, before any axis applies: integers as they are, and floats as `float_position` reads them within tolerance `t`, with `None` for a NaN or a value that isn't whole. Items that aren't real numbers have no position. [`positions`] reads the values and checks them against an axis in one pass.
479pub fn signed_positions(view: View<'_>, t: crate::tolerant::Tolerance) -> Vec<Option<i64>> {
480    match view {
481        View::Booleans(b) => b.iter().map(|&i| Some(i.into())).collect(), View::Integers(ix) => with_ints!(ix, |i| i.iter().map(|&i| Some(i.into())).collect()),
482        View::Floats(x) => with_floats!(x, |x| x.iter().map(|&y| float_position(y, t)).collect()), _ => vec![None; view.len()]
483    }
484}
485
486/// The positions along an axis of length `size` that signed positions give, as [`positions`] gives them, with `Steps::FILL` for `None` and for a position outside the axis.
487pub fn counted_positions(values: &[Option<i64>], size: usize) -> Vec<usize> {
488    values.iter().map(|&q| q.map(|i| counted_back(i, size)).filter(|&p| p < size).unwrap_or(Steps::FILL)).collect()
489}
490
491/// How a position outside an axis maps to one inside it, as edge modes of padding do. `Fill` gives no position, which reads the fill.
492#[derive(Clone, Copy, Debug, PartialEq, Eq)]
493pub enum Edge { Fill, Clamp, Wrap, Reflect, Symmetric }
494
495/// The positions that [`Edge::positions`] gives: each item's position on the axis, or floats where some item reads the fill, with NaN there.
496#[derive(Debug, PartialEq)]
497pub enum Edged { Positions(Vec<usize>), Floats(Vec<f64>) }
498
499impl Edge {
500    /// Each mode with its name.
501    pub const MODES: [(&'static str, Edge); 5] =
502        [("fill", Edge::Fill), ("clamp", Edge::Clamp), ("wrap", Edge::Wrap), ("reflect", Edge::Reflect), ("symmetric", Edge::Symmetric)];
503    /// The mode named `name`.
504    pub fn named(name: &str) -> Option<Edge> { Self::MODES.iter().find(|(n, _)| *n == name).map(|&(_, e)| e) }
505    /// This mode's name.
506    pub fn name(self) -> &'static str { Self::MODES.iter().find(|&&(_, e)| e == self).expect("every mode has a name").0 }
507    /// The position on an axis of `n` items that position `p` reads, or `None` for the fill. A position inside the axis stays as it is in every mode. A negative `p` is outside the axis here. `Reflect` leaves out the edge item, as NumPy's `reflect` does, so `-1` reads 1. `Symmetric` repeats it, as NumPy's `symmetric` does, so `-1` reads 0. An empty axis gives `None` in every mode.
508    #[inline]
509    pub fn position(self, p: i64, n: usize) -> Option<usize> {
510        if (p as u64) < n as u64 { return Some(p as usize); }
511        let n = i64::try_from(n).ok().filter(|&n| n > 0)?;
512        let q = match self {
513            Edge::Fill => p, Edge::Clamp => p.clamp(0, n - 1), Edge::Wrap => p.rem_euclid(n), Edge::Reflect if n == 1 => 0,
514            Edge::Reflect => (n - 1) - ((n - 1) - p.rem_euclid(2 * (n - 1))).abs(), Edge::Symmetric => { let r = p.rem_euclid(2 * n); r.min(2 * n - 1 - r) }
515        };
516        usize::try_from(q).ok().filter(|&q| q < n as usize)
517    }
518    /// The position on an axis of `n` items that each index value in `view` reads (`position`): integers as they are, and floats as `float_position` reads them within tolerance `t`. A NaN, or a float that isn't a whole number, reads the fill.
519    pub fn positions(self, view: View<'_>, n: usize, t: crate::tolerant::Tolerance) -> Result<Edged, Error> {
520        index_view!(view, |ix| Ok(self.mapped(n, false, ix.iter().map(|&i| Some(i.into())))), |x| {
521            let signed: Vec<_> = x.iter().map(|&y| float_position(y, t)).collect();
522            let unread = signed.iter().any(Option::is_none);
523            Ok(self.mapped(n, unread, signed.into_iter()))
524        })
525    }
526    /// What [`Edge::positions`] gives for `signed` positions, each `None` for a NaN, where `nan` says whether any is. Where `Fill` or an empty axis reads the fill for some position, or a position is NaN, the result is floats.
527    pub fn mapped(self, n: usize, nan: bool, signed: impl ExactSizeIterator<Item = Option<i64>>) -> Edged {
528        if self == Edge::Fill || nan || (n == 0 && signed.len() > 0) {
529            return Edged::Floats(signed.map(|q| q.and_then(|p| self.position(p, n)).map_or(f64::NAN, |p| p as f64)).collect());
530        }
531        Edged::Positions(signed.map(|q| q.and_then(|p| self.position(p, n)).expect("a mode other than fill reads a position on an axis with items")).collect())
532    }
533}
534
535/// Repeats each cell along `axis` by its count: one count for every cell, or one for each cell. A negative count gives that many fills in place of its cell. Where the axis has one cell, every count repeats it. A fill has no key, and repeating a keyed position is an error, as it is for selection.
536pub fn repeat(layout: &Layout, axis: usize, counts: &[i64]) -> Result<Plan, Error> { replicated(layout, axis, counts, false) }
537
538/// Undoes `repeat` with the same counts: one cell for each count. A positive count gives the first cell it made, and any other count a fill. One count must divide the number of cells. More counts must use every cell, unless the axis has one cell.
539pub fn expand(layout: &Layout, axis: usize, counts: &[i64]) -> Result<Plan, Error> { replicated(layout, axis, counts, true) }
540
541/// `repeat`'s items with no plan to walk: `view`, the items of `layout`, with each cell along `axis` written as many times as its count, a run at a time (`gather::repeat`, or `gather::positions` for counts of 0 and 1). `None` where a count is negative or `axis` has keys, which `repeat`'s plan serves.
542pub fn repeated<C: crate::search::Key>(view: View<'_>, layout: &Layout, axis: usize, counts: &[C]) -> Result<Option<(Layout, crate::buffer::Buffer)>, Error> {
543    use crate::buffer::Element;
544    let (lo, hi) = if counts.is_empty() { (0, 0) } else { C::range(counts) };
545    if layout.keys(axis).is_some() || lo < 0 { return Ok(None); }
546    let (lanes, len, shape) = replication(layout, axis, counts, false)?;
547    let (total, inner, cell) = (shape[axis], lanes.inner, |j: usize| if lanes.len == 1 { 0 } else { j });
548    let mask = counts.len() > 1 && hi <= 1;
549    let buffer = with_dtype!(view.dtype(), T => {
550        let items = T::of(view).expect("a view's items have its dtype");
551        let mut out = Vec::with_capacity(lanes.outer * total * inner);
552        for block in items.chunks((lanes.len * inner).max(1)) {
553            if inner > 1 { for j in 0..len { for _ in 0..counts[if counts.len() == 1 { 0 } else { j }].key() { out.extend_from_slice(&block[cell(j) * inner..][..inner]); } } }
554            else if mask { crate::gather::positions(&mut out, counts, total, |j| block[cell(j)]) }
555            else { crate::gather::repeat(&mut out, len, counts, total, |j| block[cell(j)]) }
556        }
557        T::buffer(out)
558    });
559    Ok(Some((unkeyed(layout, shape)?, buffer)))
560}
561
562/// The lanes along `axis` of `layout`, the number of cells that `counts` repeat or undo, and the result's shape, for `repeat`, or `expand` where `inverse`.
563fn replication<C: crate::search::Key>(layout: &Layout, axis: usize, counts: &[C], inverse: bool) -> Result<(Lanes, usize, Vec<usize>), Error> {
564    let lanes = Lanes::new(layout.shape(), axis)?;
565    let single = counts.len() == 1;
566    let size = |n: i64| Count::of(n).magnitude;
567    let len = match (inverse, single) {
568        (false, true) => lanes.len,
569        (false, false) if lanes.len == 1 || counts.len() == lanes.len => counts.len(),
570        (false, false) => return Err(Error::new(Kind::Length, "repeat needs one count, or one count for each cell")),
571        (true, true) => match size(counts[0].key()) {
572            0 if lanes.len == 0 => 1, n if n > 0 && lanes.len % n == 0 => lanes.len / n,
573            _ => return Err(Error::new(Kind::Domain, "one count must divide the number of cells to invert"))
574        },
575        (true, false) => counts.len(),
576    };
577    // The greatest count times their number bounds the sum, which then runs unchecked.
578    let most = if counts.is_empty() { 0 } else { let (lo, hi) = C::range(counts); lo.unsigned_abs().max(hi.unsigned_abs()) };
579    let total = if inverse { Some(len) } else {
580        most.checked_mul(len as u64)
581            .and_then(|_| usize::try_from(if single { most * len as u64 } else { counts.iter().map(|n| n.key().unsigned_abs()).sum() }).ok())
582    };
583    let mut shape = layout.shape().to_vec();
584    shape[axis] = total.ok_or_else(|| Error::new(Kind::Limit, "repeat length overflow"))?;
585    generated_len(&shape)?;
586    Ok((lanes, len, shape))
587}
588
589/// `shape` with `layout`'s keys on its other axes and its names, as a repeat with no keys along its axis lays out.
590fn unkeyed(layout: &Layout, shape: Vec<usize>) -> Result<Layout, Error> {
591    Ok(Layout::from(shape).with_keys(layout.all_keys())?.inherit_names(layout.names().to_vec()))
592}
593
594/// The plan of `repeat`, or of `expand` where `inverse`. Repeating, each count uses up one cell. Undoing, each count uses up the cells it made and gives back one.
595fn replicated(layout: &Layout, axis: usize, counts: &[i64], inverse: bool) -> Result<Plan, Error> {
596    let (lanes, len, shape) = replication(layout, axis, counts, inverse)?;
597    let single = counts.len() == 1;
598    let count = |j: usize| counts[if single { 0 } else { j }];
599    let size = |n: i64| Count::of(n).magnitude;
600    let keyed = layout.keys(axis).is_some();
601    let rows = Steps::along(lanes.outer, lanes.len * lanes.inner);
602    if !inverse && single && counts[0] >= 0 && !keyed {
603        return Ok(Plan { layout: unkeyed(layout, shape)?,
604            steps: vec![rows, Steps::along(lanes.len, lanes.inner), Steps::along(size(counts[0]), 0), Steps::along(lanes.inner, 1)] });
605    }
606    let (table, positions) = if inverse {
607        let (mut table, mut positions, mut used) = (Vec::with_capacity(len), Vec::new(), 0usize);
608        for j in 0..len {
609            let n = count(j);
610            let cell = (n > 0).then_some(if lanes.len == 1 { 0 } else { used });
611            used = used.saturating_add(size(n));
612            table.push(cell.map_or(Steps::FILL, |c| c * lanes.inner));
613            if keyed { positions.push(cell); }
614        }
615        if lanes.len != 1 && used != lanes.len { return Err(Error::new(Kind::Length, "the counts must use every cell to invert")); }
616        (table, positions)
617    } else {
618        // The table repeats each count's cell by the count's magnitude, with a fill in place of the cell where the count is negative. A mask of counts 0 and 1 takes `gather::positions`, which skips blocks of zeros.
619        let cell = |j: usize| if lanes.len == 1 { 0 } else { j };
620        let (mut table, mut positions) = (Vec::with_capacity(shape[axis]), Vec::new());
621        if !single && !keyed && counts.iter().all(|&n| n == 0 || n == 1) {
622            crate::gather::positions(&mut table, counts, shape[axis], |j| cell(j) * lanes.inner)
623        } else {
624            let sizes: std::borrow::Cow<'_, [i64]> =
625                if counts.iter().all(|&n| n >= 0) { counts.into() } else { counts.iter().map(|&n| size(n) as i64).collect() };
626            let at = |j: usize| (count(j) > 0).then_some(cell(j));
627            crate::gather::repeat(&mut table, len, &sizes, shape[axis], |j| at(j).map_or(Steps::FILL, |c| c * lanes.inner));
628            if keyed { crate::gather::repeat(&mut positions, len, &sizes, shape[axis], at); }
629        }
630        (table, positions)
631    };
632    let result = if keyed { layout.select(axis, positions.into_iter())? } else { unkeyed(layout, shape)? };
633    Ok(Plan { layout: result, steps: vec![rows, Steps::table(table), Steps::along(lanes.inner, 1)] })
634}
635
636/// The lengths of `shape` with `widths` added before and after each axis.
637pub fn padded(shape: &[usize], widths: &[(usize, usize)]) -> Result<Vec<usize>, Error> {
638    if widths.len() != shape.len() { return Err(Error::new(Kind::Rank, "padding needs one pair for each axis")); }
639    Ok(shape.iter().zip(widths).map(|(&len, &(before, after))| len.saturating_add(before).saturating_add(after)).collect())
640}
641
642/// The steps that pad each source axis under its edge mode.
643pub fn pad_steps(shape: &[usize], widths: &[(usize, usize)], edge: Edge) -> Result<Vec<Steps>, Error> {
644    let lengths = padded(shape, widths)?;
645    let stride = strides(shape);
646    Ok((0..shape.len()).map(|a| Steps::padded(-(widths[a].0 as i128), lengths[a], shape[a], stride[a], edge)).collect())
647}
648
649/// Pads every axis by `(before, after)` positions, to the lengths that `padded` gives. Each position past an end reads by `edge`: the item type's fill, or a position on the axis. Existing axis keys remain on their positions. Padding a keyed axis by another mode than `Edge::Fill` repeats keyed positions, which is an error. An axis that pads by nothing keeps a plain walk.
650pub fn pad(layout: &Layout, widths: &[(usize, usize)], edge: Edge) -> Result<Plan, Error> {
651    let shape = padded(layout.shape(), widths)?;
652    generated_len(&shape)?;
653    let mut result = layout.clone();
654    for (axis, ((&size, &len), &(before, _))) in layout.shape().iter().zip(&shape).zip(widths).enumerate() {
655        if len != size { result = result.select(axis, Steps::padded(-(before as i128), len, size, 1, edge).offsets())?; }
656    }
657    Ok(Plan { layout: result, steps: pad_steps(layout.shape(), widths, edge)? })
658}
659
660/// The lengths of `shape` after taking or dropping `counts`, each along its axis: a take keeps `|k|` positions, and a drop keeps `max(0, n-|k|)`. A later count for the same axis replaces an earlier one.
661pub fn taken(shape: &[usize], counts: &[(usize, Count)], take: bool) -> Result<Vec<usize>, Error> {
662    if counts.len() > shape.len() { return Err(Error::new(Kind::Rank, "counts exceed argument rank")); }
663    let mut lengths = shape.to_vec();
664    for &(axis, count) in counts {
665        let Some(&len) = shape.get(axis) else { return Err(Error::new(Kind::Domain, "axis is outside array rank")) };
666        lengths[axis] = if take { count.magnitude } else { len.saturating_sub(count.magnitude) };
667    }
668    Ok(lengths)
669}
670
671/// Where each axis of `shape` starts after taking or dropping `counts`: negative where a take pads before the axis.
672fn take_starts(shape: &[usize], counts: &[(usize, Count)], take: bool) -> Vec<i128> {
673    let mut starts = vec![0; shape.len()];
674    for &(axis, Count { negative, magnitude: n }) in counts {
675        let size = shape[axis];
676        starts[axis] = match (take, negative) { (true, true) => size as i128 - n as i128, (false, false) => n.min(size) as i128, _ => 0 };
677    }
678    starts
679}
680
681/// The steps that read `counts` taken from `shape`, or dropped where `take` is false, one entry for each axis (`take`).
682pub fn take_steps(shape: &[usize], counts: &[(usize, Count)], take: bool) -> Result<Vec<Steps>, Error> {
683    let lengths = taken(shape, counts, take)?;
684    let (starts, stride) = (take_starts(shape, counts, take), strides(shape));
685    Ok((0..shape.len()).map(|a| Steps::padded(starts[a], lengths[a], shape[a], stride[a], Edge::Fill)).collect())
686}
687
688/// Takes `count` items along `axis` for each `(axis, count)` in `counts`, at the lengths that `taken` gives, or drops them when `take` is false. A negative count works from the end. A take past an end reads fill there. A later count for the same axis replaces an earlier one.
689pub fn take(layout: &Layout, counts: &[(usize, Count)], take: bool) -> Result<Plan, Error> {
690    let old = layout.shape();
691    let shape = taken(old, counts, take)?;
692    generated_len(&shape)?;
693    let starts = take_starts(old, counts, take);
694    let mut result = layout.clone();
695    for a in (0..old.len()).filter(|&a| counts.iter().any(|&(axis, _)| axis == a)) {
696        result = result.select(a, Steps::padded(starts[a], shape[a], old[a], 1, Edge::Fill).offsets())?;
697    }
698    Ok(Plan { layout: result, steps: take_steps(old, counts, take)? })
699}
700
701/// The lengths of `shape` rotated along `axis`, which are its own. `frame`, the shape of counts that give one count for each lane, must match `shape` without `axis`, which is one of its axes.
702pub fn rotated<'a>(shape: &'a [usize], axis: usize, frame: Option<&[usize]>) -> Result<&'a [usize], Error> {
703    if axis >= shape.len() { return Err(Error::new(Kind::Rank, "rotation axis is outside array rank")); }
704    let others = || shape[..axis].iter().chain(&shape[axis + 1..]);
705    if frame.is_some_and(|f| !f.iter().eq(others())) { return Err(Error::new(Kind::Length, "rotation counts must match the axis frame")); }
706    Ok(shape)
707}
708
709/// The steps that rotate an axis by its count, or reverse it when no count is given.
710pub fn rotate_steps(shape: &[usize], axis: usize, count: Option<Count>) -> Vec<Steps> {
711    let stride = strides(shape);
712    (0..shape.len()).map(|a| match count {
713        _ if a != axis => Steps::along(shape[a], stride[a]),
714        None => Steps::Stride { start: shape[a].saturating_sub(1) * stride[a], len: shape[a], step: -(stride[a] as isize) },
715        Some(c) => {
716            let n = shape[a];
717            let r = if n == 0 { 0 } else { c.magnitude % n };
718            Steps::rotated(if c.negative && r != 0 { n - r } else { r }, n, stride[a])
719        }
720    }).collect()
721}
722
723/// Rotates the items along `axis`, or reverses them when `counts` is `None`. `counts` gives the counts' shape and items: one count for every lane, or one for each lane in the shape of `layout` without `axis`. A lane starts at position `count` of the axis, counted modulo its length. Keys on the axis move with their items when every lane rotates alike, and are dropped otherwise.
724pub fn rotate(layout: &Layout, axis: usize, counts: Option<(&[usize], &[i64])>) -> Result<Plan, Error> {
725    let shape = layout.shape();
726    let lanes = Lanes::new(shape, axis)?;
727    let Lanes { outer, len, inner } = lanes;
728    let shift = |n: i64| if len == 0 { 0 } else { n.rem_euclid(len as i64) as usize };
729    let count = match counts {
730        None => None,
731        Some((frame, ns)) => {
732            rotated(shape, axis, (ns.len() != 1).then_some(frame))?;
733            if len != 0 && (ns.is_empty() || ns.iter().any(|&n| shift(n) != shift(ns[0]))) {
734                // Each lane has its own count, so the offsets don't separate by axis, and the axis loses its keys.
735                let offsets =
736                    (0..outer).flat_map(|i| (0..len).flat_map(move |j| (0..inner).map(move |k| lanes.offset(i, (j + shift(ns[i * inner + k])) % len, k))));
737                return Ok(Plan { layout: layout.clone().without_keys(axis), steps: vec![Steps::table(offsets.collect())] });
738            }
739            Some(shift(ns.first().copied().unwrap_or(0)))
740        }
741    };
742    let source = |j: usize| match count { None => len - 1 - j, Some(s) => (j + s) % len };
743    Ok(Plan { layout: layout.select(axis, (0..len).map(|j| Some(source(j))))?,
744        steps: rotate_steps(shape, axis, count.map(|s| Count { negative: false, magnitude: s })) })
745}
746
747/// One axis of windows: windows of `size` items, each `step` items after the last. A padded axis has a window centred on every `step`th position, and reads past its ends by its edge mode.
748#[derive(Clone, Copy, Debug, PartialEq, Eq)]
749pub struct WindowAxis { pub size: usize, pub step: usize, pub edge: Option<Edge> }
750
751impl WindowAxis {
752    /// Whether the axis pads, with a window centred on its positions.
753    pub fn padded(&self) -> bool { self.edge.is_some() }
754}
755
756impl WindowAxis {
757    /// Where window `i` starts along the axis. It's negative when padding comes before the first item.
758    pub fn start(&self, i: usize) -> isize { (i * self.step) as isize - if self.padded() { ((self.size - 1) / 2) as isize } else { 0 } }
759}
760
761/// The lengths of the windows of `shape` that `axes` give along its leading axes: the number of windows along each of those axes, then a window's lengths. Along an axis of length `n`, `⌊(n-size)÷step⌋+1` windows fit, or none where `size` is longer. A padded window is centred on each position, which gives `⌈(n-[size even])÷step⌉` windows. Its size must be less than twice the length of an axis that isn't empty.
762pub fn windowed(shape: &[usize], axes: &[WindowAxis]) -> Result<Vec<usize>, Error> {
763    if axes.len() > shape.len() { return Err(Error::new(Kind::Rank, "window specification has more axes than its argument")); }
764    let mut lengths = Vec::with_capacity(axes.len() + shape.len());
765    for (w, &n) in axes.iter().zip(shape) {
766        lengths.push(match w.padded() {
767            false if n < w.size => 0, false => (n - w.size) / w.step + 1,
768            true if n == 0 || w.size / 2 < n => n.saturating_sub(usize::from(w.size.is_multiple_of(2))).div_ceil(w.step),
769            true => return Err(Error::new(Kind::Domain, "padded window is too large for the argument"))
770        });
771    }
772    lengths.extend(axes.iter().map(|w| w.size));
773    lengths.extend_from_slice(&shape[axes.len()..]);
774    Ok(lengths)
775}
776
777/// The windows' frame steps followed by the steps within each window.
778pub fn window_steps(lengths: &[usize], axes: &[WindowAxis], source: &[usize]) -> Vec<Steps> {
779    let stride = strides(source);
780    let frame = axes.iter().enumerate().map(|(a, w)| Steps::along(lengths[a], w.step * stride[a]));
781    frame.chain(lengths[axes.len()..].iter().zip(&stride).map(|(&len, &stride)| Steps::along(len, stride))).collect()
782}
783
784/// Pads the positions that windows read, or returns None when no axis pads or there are no windows.
785pub fn window_padding_steps(shape: &[usize], axes: &[WindowAxis], lengths: &[usize]) -> Option<Vec<Steps>> {
786    if !axes.iter().any(WindowAxis::padded) || lengths.contains(&0) { return None; }
787    let stride = strides(shape);
788    Some((0..shape.len()).map(|a| {
789        let Some(w) = axes.get(a) else { return Steps::along(shape[a], stride[a]) };
790        let before = if w.padded() { (w.size - 1) / 2 } else { 0 };
791        let len = (lengths[a] - 1) * w.step + w.size;
792        Steps::padded(-(before as i128), len, shape[a], stride[a], w.edge.unwrap_or(Edge::Fill))
793    }).collect())
794}
795
796/// Windows over the leading axes of an array of one shape: a `WindowAxis` for each, and the lengths that `windowed` gives, the frame that counts the windows along each axis followed by the shape of one window.
797#[derive(Clone, Debug)]
798pub struct Windows { shape: Vec<usize>, axes: Vec<WindowAxis>, lengths: Vec<usize> }
799
800impl Windows {
801    /// The windows of `sizes` and `moves` over the leading axes of `shape`. Each size is a `Count`: its magnitude is the window's size, and `edge` gives the edge mode of its axis from it, or `None` for an axis that doesn't pad. An axis with no move moves by 1.
802    pub fn new(sizes: &[i64], moves: &[i64], edge: impl Fn(Count) -> Option<Edge>, shape: &[usize]) -> Result<Self, Error> {
803        if moves.len() > sizes.len() { return Err(Error::new(Kind::Length, "window specification has more movements than sizes")); }
804        let axes = sizes.iter().enumerate().map(|(a, &size)| {
805            let step = match moves.get(a) {
806                None => 1, Some(&m) if m > 0 => Count::of(m).magnitude, Some(_) => return Err(Error::new(Kind::Domain, "window movements must be positive"))
807            };
808            let size = Count::of(size);
809            let axis = WindowAxis { size: size.magnitude, step, edge: edge(size) };
810            if axis.padded() && axis.size == 0 { return Err(Error::new(Kind::Domain, "padded window sizes must be positive")); }
811            Ok(axis)
812        }).collect::<Result<Vec<_>, _>>()?;
813        let lengths = windowed(shape, &axes)?;
814        Ok(Self { shape: shape.to_vec(), axes, lengths })
815    }
816    /// The number of windows along each windowed axis.
817    pub fn frame(&self) -> &[usize] { &self.lengths[..self.axes.len()] }
818    /// The shape of one window: the windows' sizes, then the argument's remaining axes.
819    pub fn cell(&self) -> &[usize] { &self.lengths[self.axes.len()..] }
820    /// The keys of the frame's axes. A padded axis keeps the key of the position each window centres on, and an unpadded axis has none.
821    fn frame_keys(&self, layout: &Layout) -> Result<Vec<Option<Arc<Keys>>>, Error> {
822        self.axes.iter().enumerate().map(|(a, w)| match layout.keys(a) {
823            Some(k) if w.padded() => k.select((0..self.lengths[a]).map(|i| Some(i * w.step))).map(Some), _ => Ok(None)
824        }).collect()
825    }
826    /// The keys of one window's axes: none on a windowed axis, so that every window has one layout, and `layout`'s own on the others.
827    fn cell_keys(&self, layout: &Layout) -> Vec<Option<Arc<Keys>>> {
828        (0..self.cell().len()).map(|a| if a < self.axes.len() { None } else { layout.keys(a).cloned() }).collect()
829    }
830    /// The plan that pads each padded axis by its edge mode, or `None` when no axis pads or there are no windows (`window_padding_steps`). Each padded axis of the result holds only the positions its windows read, and the result's unpadded windows are the padded windows of the argument.
831    pub fn padding(&self) -> Option<Plan> {
832        let steps = window_padding_steps(&self.shape, &self.axes, &self.lengths)?;
833        let padded: Vec<_> = steps.iter().map(Steps::len).collect();
834        Some(Plan { layout: padded.into(), steps })
835    }
836    /// The plan of every window over `layout`, which has the windows' shape: the frame's axes, then each window's axes. The windows read an array of shape `source`: `layout`'s shape, or the shape of `padding`'s result. The frame's axes have the keys that `frame_keys` gives, and each window's the keys that `cell_keys` gives.
837    pub fn plan(&self, layout: &Layout, source: &[usize]) -> Result<Plan, Error> {
838        let n = self.axes.len();
839        generated_len(&self.lengths)?;
840        generated_len(self.cell())?;
841        let steps = window_steps(&self.lengths, &self.axes, source);
842        let keys = self.frame_keys(layout)?.into_iter().chain(self.cell_keys(layout)).collect();
843        let names = (0..n).chain(0..self.cell().len()).map(|a| layout.name(a).cloned()).collect();
844        Ok(Plan { layout: Layout::from(self.lengths.clone()).with_keys(keys)?.inherit_names(names).unique_names(), steps })
845    }
846}
847
848/// How a copy walks its source: one entry for each axis of the copy, the source axis that each entry walks, and the offset of the positions on source axes that no entry walks. Axes that walk one source axis together read windows, as [`Windows::plan`] lays them out. Every reader of windows from `Steps` classifies them here: [`Walk::read`] gives what the axes along one source axis read.
849#[derive(Clone, Debug)]
850pub struct Walk {
851    /// The offset of the positions on source axes that no entry walks.
852    pub base: usize,
853    /// One entry for each axis of the copy. The first entry along each walked axis holds that axis's position in its start.
854    pub steps: Vec<Steps>,
855    /// The source axis that each entry walks, or `None` for an entry whose offsets are all 0, as for an axis of one item or a spread.
856    pub along: Vec<Option<usize>>,
857    /// The source's strides.
858    pub strides: Vec<usize>,
859}
860
861impl Walk {
862    /// The walk that `tables` give for a copy of `shape` from a source of shape `source`. `None` where an entry doesn't walk one source axis, where a walked axis's position can't move into the start of a `Stride`, as for a `Padded` or `Rotated` entry over part of an axis, and where the source is empty.
863    pub fn new(tables: &[Steps], shape: &[usize], source: &[usize]) -> Option<Self> {
864        let (base, mut steps) = Steps::per_axis(tables, shape)?;
865        let strides = strides(source);
866        if base >= crate::layout::count(source) { return None; }
867        let along: Vec<Option<usize>> = steps.iter().map(|e| e.axis_of(source, &strides)).collect::<Option<_>>()?;
868        let mut rest = 0;
869        for (a, p) in crate::layout::coordinates(base, source).filter(|&(_, p)| p > 0) {
870            let Some(r) = along.iter().position(|&b| b == Some(a)) else { rest += p * strides[a]; continue; };
871            let Steps::Stride { start, .. } = &mut steps[r] else { return None };
872            *start += p * strides[a];
873            if steps[r].axis_of(source, &strides) != Some(Some(a)) { return None; }
874        }
875        Some(Self { base: rest, steps, along, strides })
876    }
877    /// The axes of the copy that walk source axis `a`, in order.
878    pub fn walking(&self, a: usize) -> impl Iterator<Item = usize> + '_ { (0..self.along.len()).filter(move |&r| self.along[r] == Some(a)) }
879    /// What the entries walking source axis `a`, of `n` items, read as windows, such as a copy's frames and the positions in each window: where the windows start, and each entry's length and move in positions. Each entry must be a `Stride` that steps forward by a whole number of positions. An axis that no entry walks reads position 0. `None` for any other entries, and where the windows don't lie inside the axis.
880    pub fn read(&self, a: usize, n: usize) -> Option<(usize, Vec<[usize; 2]>)> {
881        let unit = self.strides[a];
882        if unit == 0 { return None; }
883        let (mut first, mut last, mut read) = (0, 0, Vec::new());
884        for r in self.walking(a) {
885            let Steps::Stride { start, len, step } = self.steps[r] else { return None };
886            if step <= 0 || step as usize % unit != 0 || start % unit != 0 || len == 0 { return None; }
887            let moves = step as usize / unit;
888            (first, last) = (first + start / unit, last + (len - 1) * moves);
889            read.push([len, moves]);
890        }
891        (first + last < n).then_some((first, read))
892    }
893    /// The leading windows read by this copy, as `[number, width, step]` per axis. The copy's axes must be frames, windows, then unchanged source axes. Each frame starts at 0.
894    pub fn windows(&self, source: &[usize]) -> Option<Vec<[usize; 3]>> {
895        let s = self.steps.len().checked_sub(source.len())?;
896        if self.base != 0 || s == 0 || 2 * s > self.steps.len() { return None; }
897        if (s..source.len())
898            .any(|a| !matches!(self.steps[s + a], Steps::Stride { start: 0, len, step } if len == source[a] && step == self.strides[a] as isize)) {
899                return None;
900            }
901        (0..s).map(|a| {
902            if self.along[a] != Some(a) || self.along[s + a] != Some(a) { return None; }
903            match self.read(a, source[a])? { (0, read) => match read[..] { [[n, moves], [w, 1]] => Some([n, w, moves]), _ => None }, _ => None }
904        }).collect()
905    }
906    /// The transpose of unit-step leading windows as an axis order and copies to sum. The copies read the reordered windows back to the source, with zero at missing positions.
907    pub fn transpose_windows(&self, source: &[usize]) -> Option<(Vec<usize>, Vec<Mapping>)> {
908        let windows = self.windows(source)?;
909        if windows.iter().any(|&[_, _, step]| step != 1) { return None; }
910        let s = windows.len();
911        // Move the shorter of each frame/window pair first, so each term reads the remaining axes forwards.
912        let cells: Vec<usize> = windows.iter().enumerate().map(|(a, &[n, w, _])| if n < w { a } else { s + a }).collect();
913        let order: Vec<usize> = cells.iter().copied().chain((0..self.steps.len()).filter(|r| !cells.contains(r))).collect();
914        let mut to = vec![0; order.len()];
915        for (&r, stride) in order.iter().zip(strides(&order.iter().map(|&r| self.steps[r].len()).collect::<Vec<_>>())) { to[r] = stride; }
916        let mut terms = vec![(0, vec![])];
917        for (a, &len) in source.iter().enumerate() {
918            let entries = if a < s {
919                let cell = cells[a];
920                let frame = if cell == a { s + a } else { a };
921                (0..self.steps[cell].len()).map(|c| (c * to[cell], Steps::clipped(-(c as i128), len, self.steps[frame].len(), to[frame]))).collect()
922            } else { vec![(0, Steps::along(len, to[s + a]))] };
923            terms = terms.into_iter().flat_map(|(base, steps): (usize, Vec<Steps>)| {
924                entries.iter().map(move |(offset, step)| (base + offset, steps.iter().cloned().chain([step.clone()]).collect()))
925            }).collect();
926        }
927        Some((order, terms.into_iter().map(|(base, steps)| Mapping::Steps { base, steps }).collect()))
928    }
929}
930
931/// A fold along one or two axes of a copy whose axes walk windows of its source, run as sliding folds over the source's rows without building the windows ([`crate::reduce::sliding`] and [`crate::reduce::sliding_2d`]). [`Walk`] classifies the copy. Each folded axis is a window's axis that moves one position at a time along its source axis. The copy's other axes along that source axis move by 1, where the windows slide, or by whole windows, where they tile it, as pooling's do. basedpl folds its pending copies of windows this way, and the Python API its views of windows.
932#[derive(Clone, Debug)]
933pub struct Slide {
934    /// How the copy walks its source.
935    pub walk: Walk,
936    /// The offsets of the source items that the folds read: for one fold, the rows that it reads along its source axis, or the whole source where a source axis before the fold's holds more than one item, and for two folds the whole source.
937    pub region: std::ops::Range<usize>,
938    /// The region as rows, with each fold along a middle axis: `[outer, len, inner]` for one folded axis, where windows that tile fold one tile to a row, and `[outer, len1, mid, len2, inner]` for two, in the order of their source axes.
939    pub dims: Vec<usize>,
940    /// Each fold's window width, in the order of `dims`.
941    pub widths: Vec<usize>,
942    /// The positions between the starts of each fold's windows, in the order of `dims`: 1 where they slide, and the width where they tile.
943    pub every: Vec<usize>,
944    /// The tables through which the copy without its folded axes reads the folds' results.
945    pub tables: Vec<Steps>,
946}
947
948impl Slide {
949    /// The sliding folds along `axes`, one axis or two, of the copy of `shape` that `tables` read from a source of shape `source`. Two folds need windows along two source axes, each covering its whole axis, and a last source axis after both of at least two items, as [`crate::reduce::sliding_2d`] takes them. `None` for any other copy: one that reads fills, one with an axis of more than one item that walks no source axis, one that reads a source axis that none of its axes walks at a position other than 0, one whose folded axis doesn't move by 1 along its source axis or is the only axis along it, and one whose other axes along a folded axis's source axis neither move by 1 nor tile it. `None` also for a window narrower than two items.
950    pub fn new(tables: &[Steps], shape: &[usize], source: &[usize], axes: &[usize]) -> Option<Self> {
951        let walk = Walk::new(tables, shape, source)?;
952        // An axis of the copy that walks no source axis holds one item.
953        if walk.base != 0 || walk.steps.iter().zip(&walk.along).any(|(s, a)| a.is_none() && s.len() != 1) { return None; }
954        let reads = (0..source.len()).map(|a| walk.read(a, source[a])).collect::<Option<Vec<_>>>()?;
955        let moves: Vec<usize> = walk.steps.iter().zip(&walk.along).map(|(s, a)| match (s, a) {
956            (Steps::Stride { step, .. }, Some(a)) => *step as usize / walk.strides[*a], _ => 0
957        }).collect();
958        // Each fold's source axis, the first row and number of rows that it reads along that axis, the positions between its windows' starts, and its width. A source axis before the fold's means the fold reads every row of its axis.
959        let mut lanes = axes.iter().map(|&k| {
960            let (q, width) = (walk.along[k]?, shape[k]);
961            if width < 2 || moves[k] != 1 || walk.walking(q).count() < 2 { return None; }
962            let others: Vec<usize> = walk.walking(q).filter(|&r| r != k).map(|r| moves[r]).collect();
963            let every = if others.iter().all(|&m| m == 1) { 1 } else if others.iter().all(|&m| m % width == 0) { width } else { return None; };
964            let (first, read) = &reads[q];
965            let len = 1 + read.iter().map(|&[n, m]| (n - 1) * m).sum::<usize>();
966            let (lo, len) = if crate::layout::count(&source[..q]) == 1 { (*first, len) } else { (0, source[q]) };
967            ((first - lo) % every == 0 && len % every == 0).then_some([q, lo, len, every, width])
968        }).collect::<Option<Vec<_>>>()?;
969        lanes.sort_unstable();
970        let count = |range: std::ops::Range<usize>| crate::layout::count(&source[range]);
971        let (region, dims) = match lanes[..] {
972            [[q, lo, len, every, width]] => {
973                let (outer, inner) = (count(0..q), walk.strides[q]);
974                let region = if outer == 1 { lo * inner..(lo + len) * inner } else { 0..count(0..source.len()) };
975                (region, if every == 1 { vec![outer, len, inner] } else { vec![outer * len / width, width, inner] })
976            }
977            [[q1, lo1, len1, ..], [q2, lo2, len2, ..]] => {
978                let whole = lo1 == 0 && lo2 == 0 && len1 == source[q1] && len2 == source[q2];
979                if q1 == q2 || !whole { return None; }
980                let dims = vec![count(0..q1), source[q1], count(q1 + 1..q2), source[q2], count(q2 + 1..source.len())];
981                if dims[4] < 2 { return None; }
982                (0..count(0..source.len()), dims)
983            }
984            _ => return None,
985        };
986        let mut folded = source.to_vec();
987        for &[q, _, len, every, width] in &lanes { folded[q] = (len - width) / every + 1; }
988        let folded_strides = strides(&folded);
989        let mut started = vec![false; folded.len()];
990        // Each table moves by its axis's moves, in starts along a folded axis, and the first of them along each source axis starts where the copy does.
991        let tables = (0..shape.len()).filter(|r| !axes.contains(r)).map(|r| match walk.along[r] {
992            Some(a) => {
993                let (every, lo) = lanes.iter().find(|l| l[0] == a).map_or((1, 0), |l| (l[3], l[1]));
994                let start = if std::mem::replace(&mut started[a], true) { 0 } else { (reads[a].0 - lo) / every };
995                Steps::Stride { start: start * folded_strides[a], len: shape[r], step: (moves[r] / every * folded_strides[a]) as isize }
996            }
997            None => Steps::Stride { start: 0, len: shape[r], step: 0 },
998        }).collect();
999        Some(Self { region, dims, widths: lanes.iter().map(|l| l[4]).collect(), every: lanes.iter().map(|l| l[3]).collect(), tables, walk })
1000    }
1001}
1002
1003/// The row-major shape of a source of `len` items whose axes the entries of `tables` step along, for a copy of `shape`, for a caller that holds a source's storage without its shape, as a Python view does. It has one axis for each distinct stride, and each axis holds as many items as the next larger stride spans. `None` where a stride doesn't divide the next, or `len` isn't a whole number of the largest.
1004pub fn walked_source(tables: &[Steps], shape: &[usize], len: usize) -> Option<Vec<usize>> {
1005    let (_, steps) = Steps::per_axis(tables, shape)?;
1006    let mut units: Vec<usize> = steps.iter().filter_map(|s| match *s {
1007        Steps::Stride { len, step, .. } if len > 1 && step != 0 => Some(step.unsigned_abs()), _ => None
1008    }).collect();
1009    units.push(1);
1010    units.sort_unstable();
1011    units.dedup();
1012    let largest = *units.last()?;
1013    if len % largest != 0 || units.windows(2).any(|w| w[1] % w[0] != 0) { return None; }
1014    Some(std::iter::once(len / largest).chain(units.windows(2).rev().map(|w| w[1] / w[0])).collect())
1015}
1016
1017/// The lengths of `shape` transposed by `axes`, which sends axis `i` to result axis `axes[i]`. Axes sent to the same result axis select their diagonal, at the shortest of their lengths. Every result axis from 0 to the largest destination needs a source.
1018pub fn transposed(shape: &[usize], axes: &[usize]) -> Result<Vec<usize>, Error> {
1019    let rank = axes.iter().max().map_or(0, |n| n + 1);
1020    if (0..rank).any(|r| !axes.contains(&r)) { return Err(Error::new(Kind::Rank, "transpose axes must be consecutive from 0")); }
1021    if axes.len() != shape.len() { return Err(Error::new(Kind::Length, "transpose needs one destination for each axis")); }
1022    let (mut lengths, mut seen) = (vec![0usize; rank], 0u128);
1023    for (&axis, &len) in axes.iter().zip(shape) {
1024        lengths[axis] = if seen >> axis & 1 == 1 { lengths[axis].min(len) } else { len };
1025        seen |= 1 << axis;
1026    }
1027    Ok(lengths)
1028}
1029
1030/// The steps that read `shape` transposed by `axes`, one entry for each result axis (`transpose`).
1031pub fn transpose_steps(shape: &[usize], axes: &[usize]) -> Result<Vec<Steps>, Error> {
1032    let lengths = transposed(shape, axes)?;
1033    let mut stride = vec![0usize; lengths.len()];
1034    for (&axis, step) in axes.iter().zip(strides(shape)) { stride[axis] = stride[axis].saturating_add(step); }
1035    Ok(lengths.iter().zip(&stride).map(|(&len, &stride)| Steps::along(len, stride)).collect())
1036}
1037
1038/// Sends axis `i` of `layout` to result axis `axes[i]`, at the lengths that `transposed` gives. A result axis keeps the keys and name its source axes agree on.
1039pub fn transpose(layout: &Layout, axes: &[usize]) -> Result<Plan, Error> {
1040    let source = layout.shape();
1041    let shape = transposed(source, axes)?;
1042    let sources: Vec<Vec<_>> = (0..shape.len()).map(|a| (0..axes.len()).filter(|&i| axes[i] == a).collect()).collect();
1043    let steps = transpose_steps(source, axes)?;
1044    Ok(Plan { layout: regrouped(layout, shape, &sources, &[])?, steps })
1045}
1046
1047/// Where diagonal `offset` of a matrix of `rows` by `cols` starts, as a row and a column, and its length. A positive offset starts further along the columns. A negative one starts further down the rows.
1048pub fn diagonal_span(offset: i64, [rows, cols]: [usize; 2]) -> (usize, usize, usize) {
1049    let (row, col) = if offset < 0 { (offset.unsigned_abs().min(rows as u64) as usize, 0) } else { (0, (offset as u64).min(cols as u64) as usize) };
1050    (row, col, (rows - row).min(cols - col))
1051}
1052
1053/// A diagonal of two distinct axes, moved to the last result axis. Positive offsets start further along `axes[1]`; negative ones along `axes[0]`.
1054pub fn diagonal(layout: &Layout, axes: [usize; 2], offset: i64) -> Result<Plan, Error> {
1055    let [a, b] = axes;
1056    let shape = layout.shape();
1057    if a == b || a >= shape.len() || b >= shape.len() { return Err(Error::new(Kind::Index, "a diagonal needs two distinct axes")); }
1058    let (row, col, n) = diagonal_span(offset, [shape[a], shape[b]]);
1059    let selected = layout.select(a, (row..row + n).map(Some))?.select(b, (col..col + n).map(Some))?;
1060    let remaining: Vec<_> = (0..shape.len()).filter(|&i| i != a && i != b).collect();
1061    let result: Vec<_> = remaining.iter().map(|&i| shape[i]).chain([n]).collect();
1062    let sources: Vec<_> = remaining.iter().map(|&i| vec![i]).chain([vec![a, b]]).collect();
1063    let stride = strides(shape);
1064    let mut steps: Vec<_> = remaining.iter().map(|&i| Steps::along(shape[i], stride[i])).collect();
1065    steps.push(Steps::Stride { start: row * stride[a] + col * stride[b], len: n, step: (stride[a] + stride[b]) as isize });
1066    Ok(Plan { layout: regrouped(&selected, result, &sources, &[])?, steps })
1067}
1068
1069/// A layout of `shape` whose axis `i` comes from `layout`'s axes `sources[i]`, as a transpose sends them. An axis with no sources is new, as a split or a merge makes it. A result axis keeps the keys its sources agree on when each source has its length. It keeps the name they agree on, unless `names[i]` gives one.
1070pub fn regrouped(layout: &Layout, shape: Vec<usize>, sources: &[Vec<usize>], names: &[Option<Arc<str>>]) -> Result<Layout, Error> {
1071    if sources.len() != shape.len() || names.len() > shape.len() { return Err(Error::new(Kind::Length, "each result axis needs one list of sources")); }
1072    let lengths = layout.shape();
1073    if sources.iter().flatten().any(|&a| a >= lengths.len()) { return Err(Error::new(Kind::Rank, "a source axis is outside array rank")); }
1074    let keys = sources.iter().zip(&shape).map(|(s, &len)| if s.iter().all(|&a| lengths[a] == len) { agreed(s.iter().map(|&a| layout.keys(a))) } else { None })
1075        .collect();
1076    let names = sources.iter().enumerate().map(|(i, s)| names.get(i).and_then(Clone::clone).or_else(|| agreed(s.iter().map(|&a| layout.name(a))))).collect();
1077    Layout::from(shape).with_keys(keys)?.with_names(names)
1078}
1079
1080/// The lengths of `left` and `right` joined along `axis`: their lengths along it added. They must have the same rank, and the same lengths on their other axes.
1081pub fn catenated(left: &[usize], right: &[usize], axis: usize) -> Result<Vec<usize>, Error> {
1082    let rank = left.len();
1083    if right.len() != rank { return Err(Error::new(Kind::Rank, "catenate ranks differ")); }
1084    if axis >= rank { return Err(Error::new(Kind::Domain, "catenate axis is outside result rank")); }
1085    if (0..rank).any(|a| a != axis && left[a] != right[a]) { return Err(Error::new(Kind::Length, "catenate frames differ")); }
1086    let mut shape = left.to_vec();
1087    shape[axis] = left[axis].saturating_add(right[axis]);
1088    Ok(shape)
1089}
1090
1091/// The layout of `left` and `right` joined along `axis`, with the lengths that `catenated` gives. The joined axis has both sides' keys, with no key where a side has none, and each other axis has either side's keys. A name survives where both sides agree.
1092pub fn catenate(left: &Layout, right: &Layout, axis: usize) -> Result<Layout, Error> {
1093    let shape = catenated(left.shape(), right.shape(), axis)?;
1094    generated_len(&shape)?;
1095    let names = |k: Option<&Arc<Keys>>, n: usize| k.map_or_else(|| vec![None; n], |k| k.names().to_vec());
1096    let keys = (0..shape.len()).map(|a| match (left.keys(a), right.keys(a)) {
1097        (x, y) if a != axis => Ok(x.or(y).cloned()), (None, None) => Ok(None),
1098        (x, y) => Keys::partial([names(x, left.shape()[a]), names(y, right.shape()[a])].concat()).map(Some)
1099    }).collect::<Result<_, Error>>()?;
1100    let names = (0..shape.len()).map(|a| agreed([left.name(a), right.name(a)])).collect();
1101    Ok(Layout::from(shape).with_keys(keys)?.inherit_names(names))
1102}