1use 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#[derive(Clone, Debug)]
11pub struct Plan { pub layout: Layout, pub steps: Vec<Steps> }
12
13#[derive(Clone, Copy, Debug)]
15pub struct Count { pub negative: bool, pub magnitude: usize }
16
17impl Count {
18 pub fn of(n: i64) -> Self { Count { negative: n < 0, magnitude: usize::try_from(n.unsigned_abs()).unwrap_or(usize::MAX) } }
20}
21
22#[derive(Clone, Debug, PartialEq)]
24pub enum IndexPart<T> { Axis(T), All, NewAxis, Rest }
25impl<T> IndexPart<T> {
26 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 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 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 pub fn axis(&self) -> Option<&T> { match self { IndexPart::Axis(t) => Some(t), _ => None } }
49}
50
51pub 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
80pub 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
103pub 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#[derive(Clone, Debug)]
113pub struct Routes { pub selectors: Vec<Selector>, pub sources: Vec<Vec<usize>>, pub one: bool }
114
115pub 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
151pub 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#[derive(Clone, Debug)]
164pub enum Selector { At(usize), Range { start: usize, len: usize, step: isize }, All, Indices { layout: Layout, positions: Table }, NewAxis }
165
166pub 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 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 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 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
214struct 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
249fn 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
258fn 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 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 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
328fn 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
337pub 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
343pub 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
368pub 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
374pub 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
380pub 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
387pub 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
399pub 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
413pub fn index_table(values: Arc<dyn IndexSource>, size: usize, t: crate::tolerant::Tolerance) -> Result<Table, Error> {
415 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 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#[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
446pub 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
457pub const NOT_WHOLE: &str = "index must be a whole number";
459
460fn not_whole() -> Error { Error::new(Kind::Domain, NOT_WHOLE) }
462
463macro_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
475pub 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
478pub 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
486pub 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
493pub enum Edge { Fill, Clamp, Wrap, Reflect, Symmetric }
494
495#[derive(Debug, PartialEq)]
497pub enum Edged { Positions(Vec<usize>), Floats(Vec<f64>) }
498
499impl Edge {
500 pub const MODES: [(&'static str, Edge); 5] =
502 [("fill", Edge::Fill), ("clamp", Edge::Clamp), ("wrap", Edge::Wrap), ("reflect", Edge::Reflect), ("symmetric", Edge::Symmetric)];
503 pub fn named(name: &str) -> Option<Edge> { Self::MODES.iter().find(|(n, _)| *n == name).map(|&(_, e)| e) }
505 pub fn name(self) -> &'static str { Self::MODES.iter().find(|&&(_, e)| e == self).expect("every mode has a name").0 }
507 #[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 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 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
535pub fn repeat(layout: &Layout, axis: usize, counts: &[i64]) -> Result<Plan, Error> { replicated(layout, axis, counts, false) }
537
538pub fn expand(layout: &Layout, axis: usize, counts: &[i64]) -> Result<Plan, Error> { replicated(layout, axis, counts, true) }
540
541pub 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
562fn 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 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
589fn 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
594fn 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 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
636pub 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
642pub 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
649pub 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
660pub 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
671fn 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
681pub 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
688pub 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
701pub 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
709pub 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
723pub 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 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
749pub struct WindowAxis { pub size: usize, pub step: usize, pub edge: Option<Edge> }
750
751impl WindowAxis {
752 pub fn padded(&self) -> bool { self.edge.is_some() }
754}
755
756impl WindowAxis {
757 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
761pub 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
777pub 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
784pub 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#[derive(Clone, Debug)]
798pub struct Windows { shape: Vec<usize>, axes: Vec<WindowAxis>, lengths: Vec<usize> }
799
800impl Windows {
801 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 pub fn frame(&self) -> &[usize] { &self.lengths[..self.axes.len()] }
818 pub fn cell(&self) -> &[usize] { &self.lengths[self.axes.len()..] }
820 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 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 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 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#[derive(Clone, Debug)]
850pub struct Walk {
851 pub base: usize,
853 pub steps: Vec<Steps>,
855 pub along: Vec<Option<usize>>,
857 pub strides: Vec<usize>,
859}
860
861impl Walk {
862 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 pub fn walking(&self, a: usize) -> impl Iterator<Item = usize> + '_ { (0..self.along.len()).filter(move |&r| self.along[r] == Some(a)) }
879 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 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 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 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#[derive(Clone, Debug)]
933pub struct Slide {
934 pub walk: Walk,
936 pub region: std::ops::Range<usize>,
938 pub dims: Vec<usize>,
940 pub widths: Vec<usize>,
942 pub every: Vec<usize>,
944 pub tables: Vec<Steps>,
946}
947
948impl Slide {
949 pub fn new(tables: &[Steps], shape: &[usize], source: &[usize], axes: &[usize]) -> Option<Self> {
951 let walk = Walk::new(tables, shape, source)?;
952 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 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 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
1003pub 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
1017pub 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
1030pub 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
1038pub 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
1047pub 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
1053pub 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
1069pub 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
1080pub 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
1091pub 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}