1use crate::{
5 Arg, buffer::{Element, kept}, error::{Error, Kind}, gather::{Same, Steps, Writes, cycled, walk_plan},
6 layout::{Alignment, Layout, count, generated_len, hidden, names_conflict, strides}
7};
8use std::{borrow::Cow, ops::Range, sync::{Arc, OnceLock}};
9
10#[derive(Clone, Debug, PartialEq, Eq)]
12pub enum Mapping {
13 Single,
15 Linear(usize),
17 Tiled(usize),
19 Steps { base: usize, steps: Vec<Steps> },
21}
22
23#[derive(Clone, Debug, PartialEq, Eq)]
25pub enum Missing { Default, Input(ReadPlan), Scalar(usize), Invalid }
26
27impl Missing {
28 pub fn prototype(dtype: crate::buffer::DType, sources: &mut Vec<Binding<'_>>) -> Self {
30 if dtype != crate::buffer::DType::Character { return Self::Default; }
31 sources.push(Binding::full(crate::buffer::View::Characters(&[<char as crate::buffer::Element>::FILL])));
32 Self::Scalar(sources.len() - 1)
33 }
34}
35
36#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub enum Run { Source { source: usize, first: usize, stride: isize }, Index { first: i64, step: i64 }, Fill, Invalid }
39
40impl Run {
41 pub fn positions(self, len: usize) -> impl Iterator<Item = usize> {
43 let (first, stride, len) = if let Self::Source { first, stride, .. } = self { (first, stride, len) } else { (0, 0, 0) };
44 (0..len).map(move |j| first.wrapping_add_signed(stride * j as isize))
45 }
46}
47
48pub trait RunSink {
50 fn run(&mut self, at: usize, len: usize, run: Run) -> bool;
52 fn walk(&mut self, _at: usize, _input: &ReadPlan, _map: &Mapping, _missing: &ReadPlan, _range: Range<usize>) -> Option<bool> { None }
54}
55
56impl<F: FnMut(usize, usize, Run) -> bool> RunSink for F { fn run(&mut self, at: usize, len: usize, run: Run) -> bool { self(at, len, run) } }
57
58#[derive(Clone, Debug, PartialEq, Eq)]
60pub struct ReadPart { pub end: usize, pub plan: ReadPlan }
61
62#[derive(Clone, Debug, PartialEq, Eq)]
64pub enum ReadRoute {
65 Input(usize), Index { base: i64, step: i64 }, Map { input: ReadPlan, map: Mapping, missing: ReadPlan }, Join(Arc<[ReadPart]>), Default, Invalid
66}
67
68#[derive(Clone, Debug, PartialEq, Eq)]
69pub struct ReadPlan { route: Arc<ReadRoute>, items: usize }
70
71impl ReadPlan {
72 pub fn input(source: usize, items: usize) -> Self { Self { route: Arc::new(ReadRoute::Input(source)), items } }
73 pub fn fill(items: usize) -> Self { Self { route: Arc::new(ReadRoute::Default), items } }
74 pub fn index(base: i64, step: i64, items: usize) -> Self { Self { route: Arc::new(ReadRoute::Index { base, step }), items } }
75 fn fallback(missing: Missing, items: usize) -> Self {
76 match missing {
77 Missing::Input(plan) => plan, Missing::Scalar(source) => Self::input(source, 1).remapped(Mapping::Single, items, Missing::Invalid),
78 Missing::Default => Self::fill(items), Missing::Invalid => Self { route: Arc::new(ReadRoute::Invalid), items }
79 }
80 }
81 pub fn copy(source: usize, source_items: usize, map: Mapping, items: usize, missing: Missing) -> Self {
82 Self::input(source, source_items).remapped(map, items, missing)
83 }
84 pub fn joined(rows: usize, parts: impl IntoIterator<Item = Self>) -> Result<Self, Error> {
86 let too_large = || Error::new(Kind::Length, "joined source is too large");
87 let (mut joined, mut width) = (Vec::new(), 0usize);
88 for plan in parts {
89 let share = if rows == 0 { 0 } else { plan.len() / rows };
90 if share * rows != plan.len() { return Err(Error::new(Kind::Length, "a joined part doesn't fill its rows")); }
91 width = width.checked_add(share).ok_or_else(too_large)?;
92 joined.push(ReadPart { end: width, plan });
93 }
94 Ok(Self { route: Arc::new(ReadRoute::Join(joined.into())), items: width.checked_mul(rows).ok_or_else(too_large)? })
95 }
96 pub fn len(&self) -> usize { self.items }
97 pub fn is_empty(&self) -> bool { self.items == 0 }
98 pub fn route(&self) -> &ReadRoute { &self.route }
99 pub fn at(&self, item: usize) -> Run {
101 let mut found = Run::Invalid;
102 if item < self.items { self.runs(item..item + 1, &mut |_, _, run| { found = run; false }); }
103 found
104 }
105 pub fn runs(&self, range: Range<usize>, sink: &mut dyn RunSink) -> bool { self.runs_through(range.start, 1, range.len(), range.start, sink) }
107 fn runs_through(&self, first: usize, stride: isize, len: usize, at: usize, sink: &mut dyn RunSink) -> bool {
109 if len == 0 { return true; }
110 match self.route() {
111 &ReadRoute::Input(source) => sink.run(at, len, Run::Source { source, first, stride }),
112 &ReadRoute::Index { base, step } => {
113 sink.run(at, len, Run::Index { first: base.wrapping_add(step.wrapping_mul(first as i64)), step: step.wrapping_mul(stride as i64) })
114 }
115 ReadRoute::Default => sink.run(at, len, Run::Fill),
116 ReadRoute::Invalid => sink.run(at, len, Run::Invalid),
117 ReadRoute::Join(parts) => {
118 let mut j = 0;
119 while j < len {
120 let Some((plan, item, before, after)) = joined_at(parts, first.wrapping_add_signed(stride * j as isize)) else {
121 return sink.run(at + j, len - j, Run::Invalid);
122 };
123 let n = (len - j).min(match stride.cmp(&0) {
124 std::cmp::Ordering::Greater => after.div_ceil(stride as usize), std::cmp::Ordering::Less => before / stride.unsigned_abs() + 1,
125 std::cmp::Ordering::Equal => len - j
126 });
127 if !plan.runs_through(item, stride, n, at + j, sink) { return false; }
128 j += n;
129 }
130 true
131 }
132 ReadRoute::Map { input, map, missing } => {
133 if input.is_empty() { return missing.runs_through(first, stride, len, at, sink); }
134 match stride {
135 1 => {
136 if let Some(go) = sink.walk(at, input, map, missing, first..first + len) { return go; }
137 let mut go = true;
138 map.spans(first..first + len, input.len(), &mut |p, n, run| {
139 go = match run {
140 Some((q, s)) => input.runs_through(q, s, n, at + p - first, sink), None => missing.runs_through(p, 1, n, at + p - first, sink)
141 };
142 go
143 });
144 go
145 }
146 0 => match map.get(first).filter(|&q| q < input.len()) {
147 Some(q) => input.runs_through(q, 0, len, at, sink), None => missing.runs_through(first, 0, len, at, sink)
148 },
149 _ => (0..len).all(|j| {
150 let p = first.wrapping_add_signed(stride * j as isize);
151 match map.get(p).filter(|&q| q < input.len()) {
152 Some(q) => input.runs_through(q, 1, 1, at + j, sink), None => missing.runs_through(p, 1, 1, at + j, sink)
153 }
154 }),
155 }
156 }
157 }
158 }
159 pub fn mapped(input: ReadPlan, map: Mapping, missing: ReadPlan, items: usize) -> Self {
160 if map.in_order(input.items, &[items]) { return input; }
161 if let Some(plan) = Self::single_map(&input, &map, items) { return plan; }
162 Self { route: Arc::new(ReadRoute::Map { input, map, missing }), items }
163 }
164 fn single_map(input: &ReadPlan, map: &Mapping, items: usize) -> Option<Self> {
166 let ReadRoute::Map { input: inner, map: first, missing } = input.route() else { return None };
167 let frame = map.shape(items);
168 if map.bounds(&frame)?.end > input.items { return None; }
169 let misses = match missing.route() { ReadRoute::Default => Missing::Default, ReadRoute::Invalid => Missing::Invalid, _ => return None };
170 let composed = first.compose(map, &input.shape_hint(), &frame)?;
171 (composed.footprint() <= first.footprint() + map.footprint()).then(|| Self::mapped(inner.clone(), composed, Self::fallback(misses, items), items))
172 }
173 pub fn remapped(&self, map: Mapping, items: usize, missing: Missing) -> Self { Self::mapped(self.clone(), map, Self::fallback(missing, items), items) }
174 pub fn try_for_each_source<E>(&self, visit: &mut impl FnMut(usize, usize) -> Result<(), E>) -> Result<(), E> {
176 match self.route() {
177 ReadRoute::Input(source) => visit(*source, self.items),
178 ReadRoute::Map { input, missing, .. } => { input.try_for_each_source(visit)?; missing.try_for_each_source(visit) }
179 ReadRoute::Join(parts) => { for part in &**parts { part.plan.try_for_each_source(visit)?; } Ok(()) }
180 ReadRoute::Index { .. } | ReadRoute::Default | ReadRoute::Invalid => Ok(()),
181 }
182 }
183 pub fn as_codes(&self, character: bool) -> Cow<'_, Self> {
185 if !character { return Cow::Borrowed(self); }
186 Cow::Owned(self.defaulted(<char as Element>::FILL as i64))
187 }
188 fn defaulted(&self, value: i64) -> Self {
190 match self.route() {
191 ReadRoute::Map { input, map, missing } => Self::mapped(input.defaulted(value), map.clone(), missing.defaulted(value), self.items),
192 ReadRoute::Join(parts) => {
193 let parts: Vec<_> = parts.iter().map(|part| ReadPart { end: part.end, plan: part.plan.defaulted(value) }).collect();
194 Self { route: Arc::new(ReadRoute::Join(parts.into())), items: self.items }
195 }
196 ReadRoute::Default => Self::index(value, 0, self.items),
197 ReadRoute::Input(_) | ReadRoute::Index { .. } | ReadRoute::Invalid => self.clone(),
198 }
199 }
200 pub fn contiguous(&self, range: Range<usize>) -> Option<(usize, Range<usize>)> {
201 let (source, start, stride) = self.run(&range)?;
202 (stride == 1 || range.len() == 1).then(|| (source, start..start + range.len()))
203 }
204 pub fn shape_hint(&self) -> Vec<usize> { match self.route() { ReadRoute::Map { map, .. } => map.shape(self.items), _ => vec![self.items] } }
205 pub fn positional(&self, shape: &[usize]) -> Option<(usize, Mapping)> {
207 if count(shape) != self.items { return None; }
208 match self.route() {
209 ReadRoute::Input(source) => Some((*source, Mapping::Linear(1))),
210 ReadRoute::Map { input, map, .. } => {
211 if input.is_empty() || map.bounds(shape)?.end > input.len() { return None; }
212 let inner_shape = input.shape_hint();
213 let (source, inner) = input.positional(&inner_shape)?;
214 Some((source, inner.compose(map, &inner_shape, shape)?))
215 }
216 ReadRoute::Join(parts) if parts.len() == 1 => parts[0].plan.positional(shape),
217 _ => None,
218 }
219 }
220 pub fn affine(&self, shape: &[usize]) -> Option<(usize, usize, Vec<isize>)> {
222 let (source, map) = self.positional(shape)?;
223 let (base, strides) = map.affine(shape)?;
224 Some((source, base, strides))
225 }
226 pub fn inverse(&self, source: usize, shape: &[usize], items: usize) -> Option<Mapping> {
228 if self.items != items { return None; }
229 let (found, map) = self.positional(shape)?;
230 if found != source { return None; }
231 if items == 0 { return Some(Mapping::Linear(1)); }
232 let (base, forward) = map.affine(shape)?;
233 let mut axes: Vec<_> = (0..shape.len()).filter(|&axis| shape[axis] > 1).collect();
234 axes.sort_unstable_by_key(|&axis| forward[axis].unsigned_abs());
235 let (mut covered, mut first) = (1usize, base as i128);
236 for &axis in &axes {
237 if forward[axis].unsigned_abs() != covered { return None; }
238 covered = covered.checked_mul(shape[axis])?;
239 if forward[axis] < 0 { first += forward[axis] as i128 * (shape[axis] - 1) as i128; }
240 }
241 if covered != items || first != 0 { return None; }
242 let strides = crate::layout::strides(shape);
243 let steps = axes.into_iter().rev().map(|axis| {
244 let stride = isize::try_from(strides[axis]).ok()?;
245 let reverse = forward[axis] < 0;
246 Some(Steps::Stride { start: if reverse { (shape[axis] - 1) * strides[axis] } else { 0 }, len: shape[axis],
247 step: if reverse { -stride } else { stride } })
248 }).collect::<Option<Vec<_>>>()?;
249 Some(Mapping::from(steps))
250 }
251 fn run(&self, range: &Range<usize>) -> Option<(usize, usize, isize)> {
253 if range.is_empty() || range.end > self.items { return None; }
254 let mut found = None;
255 self.runs(range.clone(), &mut |_, len, run| {
256 if let Run::Source { source, first, stride } = run && len == range.len() && (len == 1 || matches!(stride, 0 | 1)) {
257 found = Some((source, first, if len == 1 { 1 } else { stride }));
258 }
259 false
260 });
261 found
262 }
263}
264
265fn affine_bounds(base: usize, strides: &[isize], shape: &[usize]) -> Option<Range<usize>> {
266 if shape.len() != strides.len() { return None; }
267 if shape.contains(&0) { return Some(0..0); }
268 let (mut lo, mut hi) = (base as i128, base as i128);
269 for (&len, &stride) in shape.iter().zip(strides) {
270 let delta = (stride as i128).checked_mul(len.saturating_sub(1) as i128)?;
271 if delta < 0 { lo = lo.checked_add(delta)?; } else { hi = hi.checked_add(delta)?; }
272 }
273 Some(usize::try_from(lo).ok()?..usize::try_from(hi).ok()?.checked_add(1)?)
274}
275
276fn row_spans(
278 start: usize, axis: &Steps, j0: usize, n: usize, items: usize, at: usize, f: &mut dyn FnMut(usize, usize, Option<(usize, isize)>) -> bool) -> bool {
279 if let Steps::Stride { start: offset, step, .. } = *axis {
280 let first = start + offset.wrapping_add_signed(step * j0 as isize);
281 return match step.cmp(&0) {
282 std::cmp::Ordering::Greater => {
284 let below = if first >= items { 0 } else { n.min((items - 1 - first) / step as usize + 1) };
285 (below == 0 || f(at, below, Some((first, step)))) && (below == n || f(at + below, n - below, None))
286 }
287 std::cmp::Ordering::Equal => f(at, n, (first < items).then_some((first, 0))),
288 std::cmp::Ordering::Less => {
289 let past = if first < items { 0 } else { n.min((first - items) / step.unsigned_abs() + 1) };
290 (past == 0 || f(at, past, None)) && (past == n || f(at + past, n - past, Some((first.wrapping_add_signed(step * past as isize), step))))
291 }
292 };
293 }
294 let position = |j: usize| axis.get(j).map(|o| start + o).filter(|&p| p < items);
295 let (mut j, end) = (j0, j0 + n);
296 while j < end {
297 let first = position(j);
298 let step = match (first, (j + 1 < end).then(|| position(j + 1)).flatten()) { (Some(p), Some(q)) => q as isize - p as isize, _ => 0 };
299 let mut k = j + 1;
300 while k < end && match (first, position(k)) {
301 (None, None) => true, (Some(p), Some(q)) => q as isize - p as isize == step * (k - j) as isize, _ => false
302 } { k += 1; }
303 if !f(at + j - j0, k - j, first.map(|p| (p, step))) { return false; }
304 j = k;
305 }
306 true
307}
308
309#[derive(Clone, Copy)]
311pub struct Binding<'a> { pub view: crate::buffer::View<'a>, pub start: usize, pub items: usize }
312
313impl<'a> Binding<'a> {
314 pub fn full(view: crate::buffer::View<'a>) -> Self { Self { items: view.len(), view, start: 0 } }
315 pub fn as_codes(self) -> Self { Self { view: crate::characters::as_codes(self.view), ..self } }
316 pub fn borrow(&self, range: Range<usize>) -> Option<crate::buffer::View<'a>> {
317 if range.start > range.end || range.end > self.items { return None; }
318 let local = range.start.checked_sub(self.start)?..range.end.checked_sub(self.start)?;
319 (local.end <= self.view.len()).then(|| self.view.slice(local))
320 }
321 pub fn affine(&self, base: usize, strides: &[isize], shape: &[usize]) -> Option<(crate::buffer::View<'a>, usize)> {
322 self.borrow(affine_bounds(base, strides, shape)?)?;
323 Some((self.view, base.checked_sub(self.start)?))
324 }
325 pub fn window(view: crate::buffer::View<'a>, start: usize, items: usize) -> Result<Self, Error> {
326 if start.checked_add(view.len()).is_none_or(|end| end > items) { return Err(Error::new(Kind::Length, "read window exceeds its producer")); }
327 Ok(Self { view, start, items })
328 }
329}
330
331#[derive(Clone, Copy)]
333pub struct Read<'a> { pub sources: &'a [Binding<'a>], pub plan: &'a ReadPlan }
334
335pub struct Rows<'a, T> { pub items: &'a [T], pub first: usize, pub step: isize, pub run: bool }
337
338impl<'s> Read<'s> {
339 fn bindings(&self) -> Result<(), Error> {
341 self.plan.try_for_each_source(&mut |source, count| {
342 let binding = self.sources.get(source).ok_or_else(|| Error::new(Kind::Length, "missing read source"))?;
343 if binding.items != count { return Err(Error::new(Kind::Length, "read source changed length")); }
344 if binding.start.checked_add(binding.view.len()).is_none_or(|end| end > count) {
345 return Err(Error::new(Kind::Length, "read window exceeds its producer"));
346 }
347 Ok(())
348 })
349 }
350 pub fn positional(&self, shape: &[usize]) -> Option<(crate::buffer::View<'_>, Mapping)> {
351 let (source, mapping) = self.plan.positional(shape)?;
352 let binding = self.sources.get(source)?;
353 binding.borrow(mapping.bounds(shape)?)?;
354 let (base, steps) = mapping.entries(shape)?;
355 Some((binding.view, Mapping::Steps { base: base.checked_sub(binding.start)?, steps }))
356 }
357 fn range(&self, range: &Range<usize>) -> Result<(), Error> {
358 if range.start > range.end || range.end > self.plan.len() { return Err(Error::new(Kind::Length, "read range exceeds its result")); }
359 self.bindings()
360 }
361 fn run_view(&self, range: &Range<usize>) -> Option<(crate::buffer::View<'s>, usize)> {
363 let (source, start, stride) = self.plan.run(range)?;
364 let binding = self.sources.get(source)?;
365 let end = start.checked_add(if stride == 0 { 1 } else { range.len() })?;
366 Some((binding.borrow(start..end)?, stride as usize))
367 }
368 fn storage_run(&self, range: &Range<usize>, dtype: crate::buffer::DType) -> Option<(crate::buffer::View<'_>, usize)> {
370 self.run_view(range).filter(|(view, _)| view.dtype() == dtype)
371 }
372 pub fn part_view<'a>(
374 &'a self, range: Range<usize>, dtype: crate::buffer::DType, checked: bool, scratch: &'a mut crate::buffer::Buffer
375 ) -> Result<crate::buffer::View<'a>, Error> {
376 self.range(&range)?;
377 if let Some((view, stride)) = self.storage_run(&range, dtype) { if stride == 1 || range.len() == 1 { return Ok(view); } }
378 self.read_into(range, dtype, checked, scratch)?;
379 Ok(scratch.view())
380 }
381 pub fn part<'a, T: crate::buffer::Cast>(&'a self, range: Range<usize>, checked: bool, scratch: &'a mut Vec<T>) -> Result<Arg<'a, T>, Error> {
382 self.range(&range)?;
383 if range.is_empty() { return Ok(Arg::Many(&[])); }
384 if let Some((view, stride)) = self.run_view(&range) {
385 if view.dtype() != T::DTYPE {
386 if stride == 0 { return Ok(Arg::One(kept(T::cast(view, 0), checked)?)); }
387 cast_run(view, stride, range.len(), checked, scratch)?;
388 return Ok(Arg::Many(scratch));
389 }
390 let items = T::of(view).expect("a typed read selected its dtype");
391 return Ok(if stride == 0 { Arg::One(items[0]) } else { Arg::Many(items) });
392 }
393 scratch.clear();
394 self.read_onto(range, checked, scratch)?;
395 Ok(Arg::Many(scratch))
396 }
397 pub fn in_place<T: crate::buffer::Cast>(&self, range: &Range<usize>) -> bool {
399 self.run_view(range).is_some_and(|(view, stride)| stride == 0 || view.dtype() == T::DTYPE)
400 }
401 pub fn rows<T: crate::buffer::Cast>(&self, shape: [usize; 2]) -> Option<Rows<'s, T>> {
403 self.bindings().ok()?;
404 let (source, base, strides) = self.plan.affine(&shape)?;
405 let (view, first) = self.sources.get(source)?.affine(base, &strides, &shape)?;
406 let &[step, inner @ (0 | 1)] = &strides[..] else { return None };
407 Some(Rows { items: T::of(view)?, first, step, run: inner == 1 })
408 }
409 fn read_onto<T: crate::buffer::Cast>(&self, range: Range<usize>, checked: bool, output: &mut Vec<T>) -> Result<(), Error> {
411 output.reserve(range.len());
412 if self.join_onto(&range, checked, output)? { return Ok(()); }
413 let mut onto = Onto { read: self, checked, output, error: None };
414 self.plan.runs(range, &mut onto);
415 onto.error.map_or(Ok(()), Err)
416 }
417 fn join_onto<T: crate::buffer::Cast>(&self, range: &Range<usize>, checked: bool, output: &mut Vec<T>) -> Result<bool, Error> {
419 let ReadRoute::Join(parts) = self.plan.route() else { return Ok(false) };
420 let width = parts.last().map_or(0, |part| part.end);
421 if width == 0 || range.start % width != 0 || range.len() % width != 0 { return Ok(false); }
422 let (first, rows) = (range.start / width, range.len() / width);
423 let (mut lanes, mut start) = (Vec::with_capacity(parts.len()), 0);
424 for part in parts.iter() {
425 let n = part.end - std::mem::replace(&mut start, part.end);
426 let Some((view, stride)) = (Read { sources: self.sources, plan: &part.plan }).run_view(&(first * n..(first + rows) * n)) else { return Ok(false) };
427 lanes.push((match (T::of(view), stride) {
428 (Some(items), 1) => Cow::Borrowed(items),
429 (_, 0) => Cow::Owned(vec![kept(T::cast(view, 0), checked)?; rows * n]),
430 _ => {
431 let mut items = Vec::new();
432 T::extend_cast(view, rows * n, checked, &mut items)?;
433 Cow::Owned(items)
434 }
435 }, n));
436 }
437 crate::gather::interleave(output, &lanes.iter().map(|(items, n)| (&items[..], *n)).collect::<Vec<_>>(), rows);
438 Ok(true)
439 }
440 fn run_onto<T: crate::buffer::Cast>(
442 &self, source: usize, first: usize, stride: isize, len: usize, checked: bool, output: &mut Vec<T>) -> Result<(), Error> {
443 if len == 0 { return Ok(()); }
444 let outside = || Error::new(Kind::Length, "read reaches outside its bound window");
445 let last = first.checked_add_signed(stride * (len as isize - 1)).ok_or_else(outside)?;
446 let low = first.min(last);
447 let view = self.sources[source].borrow(low..first.max(last) + 1).ok_or_else(outside)?;
448 if stride == 0 {
449 let item = kept(T::cast(view, 0), checked)?;
450 output.resize(output.len() + len, item);
451 return Ok(());
452 }
453 if stride == 1 {
454 if let Some(items) = T::of(view) { output.extend_from_slice(items); return Ok(()); }
455 return T::extend_cast(view, len, checked, output);
456 }
457 let walked = Walked::from(Mapping::Steps { base: 0, steps: vec![Steps::Stride { start: first - low, len, step: stride }] });
458 if let Some(items) = T::of(view) { copy_walked(&walked, items, 0, len, T::FILL, output); return Ok(()); }
459 with_dtype!(view.dtype(), S => {
460 let mut items = Vec::with_capacity(len);
461 copy_walked(&walked, S::of(view).expect("the view has its dtype"), 0, len, S::FILL, &mut items);
462 T::extend_cast(S::view(&items), len, checked, output)
463 })
464 }
465 fn walkable(&self, input: &ReadPlan, map: &Mapping, end: usize) -> Option<&'s Binding<'s>> {
467 let &ReadRoute::Input(source) = input.route() else { return None };
468 let binding = self.sources.get(source)?;
469 let len = binding.view.len();
470 (binding.start == 0 && len == input.len() && match *map {
471 Mapping::Single => true, Mapping::Linear(r) => r > 0 && end.saturating_sub(1) / r < len, Mapping::Tiled(m) => m > 0 && m <= len,
472 Mapping::Steps { base, ref steps } => Steps::furthest(base, &steps[..]).is_none_or(|furthest| furthest < len)
473 }).then_some(binding)
474 }
475 pub fn mapped_source(&self) -> Option<(crate::buffer::View<'s>, &'s Mapping)> {
477 let ReadRoute::Map { input, map, .. } = self.plan.route() else { return None };
478 if map.fills() { return None; }
479 Some((self.walkable(input, map, self.plan.len())?.view, map))
480 }
481 fn walk_onto<T: crate::buffer::Cast>(
483 &self, input: &ReadPlan, map: &Mapping, missing: &ReadPlan, range: Range<usize>, checked: bool, output: &mut Vec<T>) -> Result<bool, Error> {
484 let Some(binding) = self.walkable(input, map, range.end) else { return Ok(false) };
485 let Some(fill) = Filled::of(missing).filter(|&fill| fill != Filled::Never || !map.fills()) else { return Ok(false) };
486 let walked = Walked::from(map);
487 if let Some(items) = T::of(binding.view) {
488 let Some(fill) = fill.item::<T>(self.sources) else { return Ok(false) };
489 copy_walked(&walked, items, range.start, range.len(), fill, output);
490 return Ok(true);
491 }
492 with_dtype!(binding.view.dtype(), S => {
493 let Some(fill) = fill.item::<S>(self.sources) else { return Ok(false) };
494 let mut items = Vec::with_capacity(range.len());
495 copy_walked(&walked, S::of(binding.view).expect("the view has its dtype"), range.start, range.len(), fill, &mut items);
496 T::extend_cast(S::view(&items), items.len(), checked, output)?;
497 Ok(true)
498 })
499 }
500 pub fn read_into(&self, range: Range<usize>, dtype: crate::buffer::DType, checked: bool, output: &mut crate::buffer::Buffer) -> Result<(), Error> {
501 self.range(&range)?;
502 if output.dtype() != dtype { *output = crate::buffer::Buffer::with_capacity(dtype, range.len()); }
503 with_dtype!(dtype, T => {
504 let items = T::vec_mut(output).expect("output has the selected dtype");
505 if let Some((view, stride)) = self.run_view(&range) {
506 if view.dtype() != dtype { return cast_run(view, stride, range.len(), checked, items); }
507 let source = T::of(view).expect("a typed read selected its dtype");
508 items.clear();
509 if stride == 0 { items.resize(range.len(), source[0]); } else { items.extend_from_slice(source); }
510 return Ok(());
511 }
512 items.clear();
513 self.read_onto(range, checked, items)
514 })
515 }
516}
517
518struct Onto<'r, 'a, T> { read: &'r Read<'a>, checked: bool, output: &'r mut Vec<T>, error: Option<Error> }
520
521impl<T: crate::buffer::Cast> RunSink for Onto<'_, '_, T> {
522 fn run(&mut self, _at: usize, len: usize, run: Run) -> bool {
523 let read = match run {
524 Run::Source { source, first, stride } => self.read.run_onto(source, first, stride, len, self.checked, self.output),
525 Run::Index { first, step } => {
526 let values: Vec<i64> = (0..len).map(|j| first.wrapping_add(step.wrapping_mul(j as i64))).collect();
527 T::extend_cast(crate::buffer::View::Integers(crate::buffer::Ints::I64(&values)), len, self.checked, self.output)
528 }
529 Run::Fill => { self.output.resize(self.output.len() + len, T::FILL); Ok(()) }
530 Run::Invalid => Err(Error::new(Kind::Index, "read position is outside its source")),
531 };
532 read.map_err(|e| self.error = Some(e)).is_ok()
533 }
534 fn walk(&mut self, _at: usize, input: &ReadPlan, map: &Mapping, missing: &ReadPlan, range: Range<usize>) -> Option<bool> {
535 match self.read.walk_onto(input, map, missing, range, self.checked, self.output) {
536 Ok(true) => Some(true), Ok(false) => None, Err(e) => { self.error = Some(e); Some(false) }
537 }
538 }
539}
540
541fn joined_at(parts: &[ReadPart], position: usize) -> Option<(&ReadPlan, usize, usize, usize)> {
543 let width = parts.last().map_or(0, |part| part.end);
544 if width == 0 { return None; }
545 let (row, column) = (position / width, position % width);
546 let k = parts.partition_point(|part| part.end <= column);
547 let start = if k == 0 { 0 } else { parts[k - 1].end };
548 let part = parts.get(k)?;
549 Some((&part.plan, row * (part.end - start) + column - start, column - start, part.end - column))
550}
551fn cast_run<T: crate::buffer::Cast>(view: crate::buffer::View<'_>, stride: usize, len: usize, checked: bool, output: &mut Vec<T>) -> Result<(), Error> {
553 output.clear();
554 if stride == 0 { output.resize(len, kept(T::cast(view, 0), checked)?); return Ok(()); }
555 T::extend_cast(view, len, checked, output)
556}
557
558#[derive(Clone, Copy, PartialEq)]
560enum Filled { Fill, Value(i64), Source(usize), Never }
561
562impl Filled {
563 fn of(missing: &ReadPlan) -> Option<Self> {
564 match missing.route() {
565 ReadRoute::Default => Some(Self::Fill), ReadRoute::Invalid => Some(Self::Never), &ReadRoute::Index { base, step: 0 } => Some(Self::Value(base)),
566 &ReadRoute::Input(source) if missing.len() == 1 => Some(Self::Source(source)),
567 ReadRoute::Map { input, map: Mapping::Single, .. } => match *input.route() { ReadRoute::Input(source) => Some(Self::Source(source)), _ => None },
568 _ => None
569 }
570 }
571 fn item<A: crate::buffer::Cast>(self, sources: &[Binding<'_>]) -> Option<A> {
573 match self {
574 Self::Fill | Self::Never => Some(A::FILL), Self::Value(value) => A::cast(crate::buffer::View::Integers(crate::buffer::Ints::I64(&[value])), 0).ok(),
575 Self::Source(source) => A::of(sources.get(source).filter(|binding| binding.start == 0)?.view)?.first().copied()
576 }
577 }
578}
579
580#[derive(Clone, Copy, Debug, PartialEq, Eq)]
582pub enum Period<'a> {
583 Repeats(usize),
585 Cycles(usize),
587 Walks { len: usize, base: usize, steps: &'a [Steps] },
589}
590
591impl Period<'_> {
592 pub fn tiling(&self, run: usize) -> Option<(usize, usize)> {
594 match *self {
595 Self::Cycles(p) | Self::Walks { len: p, .. } if p < LONG_CYCLE => Some((p, run)), Self::Walks { len, .. } => Some((len, 0)),
596 Self::Cycles(_) | Self::Repeats(_) => None
597 }
598 }
599}
600
601#[derive(Clone, Debug, PartialEq, Eq)]
603pub struct Block {
604 pub range: Range<usize>,
606 pub dims: Vec<usize>,
608 pub axes: Vec<Option<(usize, bool)>>,
610}
611
612impl Block {
613 pub fn fold(&self, axis: Option<usize>) -> ([usize; 3], Option<Vec<Steps>>) {
615 let Some(axis) = axis else { return ([1, self.range.len(), 1], None) };
616 let folded = self.axes[axis].map(|(b, _)| b);
617 let dims = &self.dims;
618 let around = match folded { Some(f) => [count(&dims[..f]), dims[f], count(&dims[f + 1..])], None => [self.range.len(), 1, 1] };
619 let rest: Vec<usize> = dims.iter().enumerate().filter(|&(b, _)| Some(b) != folded).map(|(_, &n)| n).collect();
621 let stride = strides(&rest);
622 let steps: Vec<Steps> = self.axes.iter().enumerate().filter(|&(a, _)| a != axis).map(|(_, kept)| match *kept {
623 None => Steps::along(1, 0),
624 Some((b, back)) => {
625 let b = b - usize::from(folded.is_some_and(|f| b > f));
626 let (len, step) = (rest[b], stride[b] as isize);
627 if back { Steps::Stride { start: (len - 1) * stride[b], len, step: -step } } else { Steps::Stride { start: 0, len, step } }
628 }
629 }).collect();
630 (around, (!Steps::in_order(&steps)).then_some(steps))
631 }
632}
633
634impl From<Vec<Steps>> for Mapping {
636 #[inline]
637 fn from(tables: Vec<Steps>) -> Self { Self::Steps { base: 0, steps: tables.into() } }
638}
639
640impl From<Writes<'_>> for Mapping {
642 fn from(writes: Writes<'_>) -> Self {
643 match writes {
644 Writes::Offsets(offsets) => Self::from(vec![Steps::table(offsets.to_vec())]),
645 Writes::Walk { base, steps } => Self::Steps { base, steps: steps.to_vec().into() },
646 }
647 }
648}
649
650#[derive(Clone, Debug)]
652pub struct Walked<'a> { mapping: Cow<'a, Mapping>, walk: OnceLock<(usize, Vec<Cow<'static, Steps>>)>, cycle: OnceLock<Option<usize>> }
653
654impl<'a> Walked<'a> {
655 pub fn new(mapping: Cow<'a, Mapping>) -> Self { Self { mapping, walk: OnceLock::new(), cycle: OnceLock::new() } }
657 #[inline]
659 fn walk(&self) -> (usize, &[Cow<'static, Steps>]) {
660 let (offset, walk) = self.walk.get_or_init(|| match &*self.mapping {
661 Mapping::Steps { steps, .. } => {
662 let (offset, walk) = Steps::simplify(0, &steps[..]);
663 (offset, walk.into_iter().map(|s| Cow::Owned(s.into_owned())).collect())
664 }
665 _ => (0, Vec::new()),
666 });
667 (*offset, walk)
668 }
669 pub fn cycle(&self) -> Option<usize> { *self.cycle.get_or_init(|| match self.mapping.period() { Some(Period::Cycles(p)) => Some(p), _ => None }) }
671}
672
673impl<'a> From<&'a Mapping> for Walked<'a> { fn from(mapping: &'a Mapping) -> Self { Self::new(Cow::Borrowed(mapping)) } }
674impl From<Mapping> for Walked<'static> { fn from(mapping: Mapping) -> Self { Self::new(Cow::Owned(mapping)) } }
675impl std::ops::Deref for Walked<'_> {
676 type Target = Mapping;
677 #[inline]
678 fn deref(&self) -> &Mapping { &self.mapping }
679}
680impl PartialEq for Walked<'_> { fn eq(&self, other: &Self) -> bool { self.mapping == other.mapping } }
681
682impl Mapping {
683 pub fn tabled(&self) -> bool { matches!(self, Self::Steps { steps, .. } if steps.iter().any(|s| matches!(s, Steps::Table(_)))) }
685 pub fn as_writes(&self) -> Option<Writes<'_>> { if let Self::Steps { base, steps } = self { Some(Writes::Walk { base: *base, steps }) } else { None } }
687 pub fn entries(&self, shape: &[usize]) -> Option<(usize, Vec<Steps>)> {
689 if shape.contains(&0) { return Some((0, shape.iter().map(|&n| Steps::along(n, 0)).collect())); }
690 if let Self::Steps { base, steps } = self { let (offset, steps) = Steps::per_axis(steps, shape)?; return Some((base + offset, steps)); }
691 if let Self::Linear(period) | Self::Tiled(period) = *self {
692 let n = count(shape);
693 if period < n && period != 1 && !shape.iter().rev().scan(1, |p, &n| {
694 *p *= n;
695 Some(*p)
696 }).any(|p| p == period) { return None; }
697 }
698 let base = self.get(0)?;
699 let steps = shape.iter().zip(strides(shape))
700 .map(|(&len, stride)| Some(Steps::Stride { start: 0, len, step: if len > 1 { self.get(stride)? as isize - base as isize } else { 0 } }))
701 .collect::<Option<_>>()?;
702 Some((base, steps))
703 }
704 pub fn through(&self, outer: &[Steps], shape: &[usize]) -> Option<Self> {
706 if outer.iter().any(|s| s.len() == 0) { return Some(Self::Linear(1)); }
707 let (base, inner) = self.entries(shape)?;
708 let (base, steps) = Steps::through(outer, shape, base, &inner)?;
709 Some(Self::Steps { base, steps })
710 }
711 pub fn split(&self, shape: &[usize]) -> Self { self.entries(shape).map_or_else(|| self.clone(), |(base, steps)| Self::Steps { base, steps }) }
713 pub fn to_walk(&self, shape: &[usize]) -> (usize, Cow<'_, [Steps]>) {
714 if let Self::Steps { base, steps } = self { return (*base, Cow::Borrowed(steps)); }
715 let (base, steps) = match self.entries(shape) {
716 Some((base, steps)) => (base, steps), None => (0, vec![Steps::table((0..count(shape)).map(|i| self.get(i).unwrap_or(Steps::FILL)).collect())])
717 };
718 (base, Cow::Owned(steps))
719 }
720 fn shape(&self, items: usize) -> Vec<usize> {
722 match self { Self::Steps { steps, .. } if Steps::count(steps) == items => steps.iter().map(Steps::len).collect(), _ => vec![items] }
723 }
724 pub fn in_order(&self, len: usize, shape: &[usize]) -> bool {
726 len == count(shape) && (matches!(self, Self::Linear(1)) || { let (base, steps) = self.to_walk(shape); Self::walk_in_order(base, &steps, len) })
727 }
728 pub fn walk_in_order(base: usize, steps: &[Steps], len: usize) -> bool { base == 0 && Steps::count(steps) == len && Steps::in_order(steps) }
730 pub fn tables(&self, shape: &[usize]) -> Vec<Steps> { let (base, steps) = self.to_walk(shape); Steps::with_base(base, steps.into_owned()) }
732 pub fn footprint(&self) -> usize { if let Self::Steps { steps, .. } = self { steps.iter().map(Steps::footprint).sum() } else { 0 } }
734 pub fn compose(&self, outer: &Self, shape: &[usize], frame: &[usize]) -> Option<Self> {
736 if frame.contains(&0) { return Some(Self::Linear(1)); }
737 if matches!(self, Self::Linear(1)) { return Some(outer.clone()); }
738 if matches!(outer, Self::Linear(1)) { return Some(self.clone()); }
739 let (base, steps) = match outer { Self::Steps { base, steps } => (*base, steps.to_vec()), _ => outer.entries(frame)? };
740 Some(self.through(&Steps::with_base(base, steps), shape)?)
741 }
742 pub fn composed(&self, outer: &Self, shape: &[usize], frame: &[usize], table: bool) -> Option<Self> {
744 if let Some(mapping) = self.compose(outer, shape, frame) { return Some(mapping); }
745 if !table { return None; }
746 let walked = |map: &Self, shape: &[usize]| {
747 let (base, steps) = map.to_walk(shape);
748 Steps::combined(base, &steps[..])
749 };
750 let inner = (count(shape) <= count(frame)).then(|| walked(self, shape));
751 let at = |j: usize| match &inner {
752 _ if j == Steps::FILL => j, Some(listed) => listed.get(j).copied().unwrap_or(Steps::FILL), None => self.get(j).unwrap_or(Steps::FILL)
753 };
754 Some(Self::Steps { base: 0, steps: vec![Steps::table(walked(outer, frame).into_iter().map(at).collect())].into() })
755 }
756 fn bounds(&self, shape: &[usize]) -> Option<Range<usize>> { let (base, strides) = self.affine(shape)?; affine_bounds(base, &strides, shape) }
757 pub fn affine(&self, shape: &[usize]) -> Option<(usize, Vec<isize>)> {
759 if !shape.contains(&0) && let Self::Steps { steps, .. } = self && steps.iter().any(|s| s.len() > 1 && s.affine().is_none()) { return None; }
761 let (base, steps) = self.entries(shape)?;
762 let (mut first, mut strides) = (base, Vec::with_capacity(steps.len()));
763 for s in &steps {
764 let (start, step) = s.affine()?;
765 first = first.wrapping_add(start);
766 strides.push(step);
767 }
768 Some((first, strides))
769 }
770 pub fn block(&self, shape: &[usize]) -> Option<Block> {
772 if shape.contains(&0) { return None; }
773 let (first, strides) = self.affine(shape)?;
774 let mut walked: Vec<usize> = (0..shape.len()).filter(|&a| shape[a] > 1).collect();
775 walked.sort_by_key(|&a| std::cmp::Reverse(strides[a].unsigned_abs()));
776 let mut size = 1;
777 for &a in walked.iter().rev() {
778 if strides[a].unsigned_abs() != size { return None; }
779 size *= shape[a];
780 }
781 let start = walked.iter().filter(|&&a| strides[a] < 0).try_fold(first, |s, &a| s.checked_sub(strides[a].unsigned_abs() * (shape[a] - 1)))?;
782 let mut axes = vec![None; shape.len()];
783 for (b, &a) in walked.iter().enumerate() { axes[a] = Some((b, strides[a] < 0)); }
784 Some(Block { range: start..start + size, dims: walked.iter().map(|&a| shape[a]).collect(), axes })
785 }
786 #[inline]
788 pub fn index(&self, i: usize) -> usize { self.get(i).expect("positional mapping") }
789 #[inline]
791 pub fn run(&self, start: usize, n: usize) -> Option<(usize, isize)> {
792 let Self::Steps { base, steps } = self else { return None };
793 Steps::run(*base, steps, start, n)
794 }
795 #[inline]
798 pub fn get(&self, i: usize) -> Option<usize> {
799 match self {
800 Self::Linear(0) | Self::Tiled(0) => None, Self::Single => Some(0), Self::Linear(1) => Some(i), Self::Linear(repeat) => Some(i / repeat),
801 Self::Tiled(size) => Some(i % size), Self::Steps { base, steps } => Self::reached(*base, steps, i)
802 }
803 }
804 #[inline(never)]
805 fn reached(base: usize, steps: &[Steps], i: usize) -> Option<usize> { Steps::reach(base, steps, i) }
806 pub fn contract(left: &Layout, laxis: usize, right: &Layout, raxis: usize) -> Result<Self, Error> {
808 match (left.keys(laxis), right.keys(raxis)) {
809 (Some(x), Some(y)) if x != y => Ok(Self::Steps { base: 0, steps: vec![Steps::table(y.positions(x, false)?)].into() }), _ => Ok(Self::Linear(1))
810 }
811 }
812 pub fn new(input: &Layout, axes: &[usize], result: &Layout, count: usize) -> Self {
814 let alignments = axes.iter().enumerate().map(|(a, &b)| match (input.keys(a), result.keys(b)) {
815 (Some(src), Some(dst)) => src.aligned(dst), _ => Alignment::Same
816 }).collect();
817 Self::aligned(input, axes, result, count, alignments)
818 }
819 fn aligned(input: &Layout, axes: &[usize], result: &Layout, count: usize, alignments: Vec<Alignment>) -> Self {
821 if count == 0 { return Self::Single; }
822 let (shape, output) = (input.shape(), result.shape());
823 if alignments.iter().all(|a| matches!(a, Alignment::Same)) {
824 let size = input.count();
825 if size == 1 { return Self::Single; }
826 let strides = strides(output);
827 let (mut source, mut linear, mut tiled, repeat) = (1, true, true, count / size);
828 for (&len, &axis) in shape.iter().zip(axes).rev() {
829 if len != 1 {
830 linear &= strides[axis] == source * repeat;
831 tiled &= strides[axis] == source;
832 }
833 source *= len;
834 }
835 if linear { return Self::Linear(repeat); }
836 if tiled { return Self::Tiled(size); }
837 }
838 let mut steps: Vec<Steps> = output.iter().map(|&len| Steps::Stride { start: 0, len, step: 0 }).collect();
839 let mut source = 1;
840 for ((&len, &b), alignment) in shape.iter().zip(axes).zip(alignments).rev() {
841 steps[b] = match alignment {
842 Alignment::Same => Steps::Stride { start: 0, len: output[b], step: if len == 1 { 0 } else { source as isize } },
843 Alignment::Leading => Steps::clipped(0, output[b], len, source),
844 Alignment::Table(t) => {
846 let table = Steps::table(t.into_iter().map(|p| p.map_or(Steps::FILL, |p| p * source)).collect());
847 table.even().unwrap_or(table)
848 }
849 };
850 source *= len;
851 }
852 Self::Steps { base: 0, steps: steps.into() }
853 }
854 pub fn line_block(&self, size: usize) -> usize {
856 match self.inner() { Some(inner) if inner.skips() => (inner.len() * (LINE / size).max(1)).min(LINE_BLOCK), _ => 0 }
857 }
858 pub fn period(&self) -> Option<Period<'_>> {
860 if let Self::Tiled(m) = *self { return Some(Period::Cycles(m)); }
861 if let Some(r) = self.repeats() { return Some(Period::Repeats(r)); }
862 let Self::Steps { base, steps } = self else { return None };
863 if let [Steps::Padded { start: 0, len, size, stride: 1, edge: crate::structure::Edge::Wrap }] = steps[..] && *base == 0 && size >= 2 && len > size {
864 return Some(Period::Cycles(size));
865 }
866 let k = steps.iter().position(|s| !matches!(*s, Steps::Stride { step: 0, .. }))?;
867 let (outer, inner) = steps.split_at(k);
868 let p = inner.iter().map(Steps::len).product::<usize>();
869 if p < 2 || outer.iter().all(|s| s.len() < 2) || self.fills() { return None; }
870 let base = base + outer.iter().map(|s| if let Steps::Stride { start, .. } = *s { start } else { 0 }).sum::<usize>();
871 Some(if base == 0 && Steps::in_order(inner) { Period::Cycles(p) } else { Period::Walks { len: p, base, steps: inner } })
872 }
873 pub fn copies(&self, block: usize) -> bool { match self { Self::Steps { steps, .. } => Steps::contiguous(steps) < block, _ => false } }
875 pub fn tiles_whole(&self, block: usize) -> bool {
877 let Self::Steps { steps, .. } = self else { return false };
878 let (_, steps) = Steps::simplify(0, steps);
879 let [.., outer, inner] = &steps[..] else { return false };
880 Steps::transposing(outer, inner).is_some_and(|(_, _, n, _)| n >= block)
881 }
882 pub fn regular(&self) -> Option<(usize, isize)> {
884 match self { Self::Linear(r) => Some((*r, 0)), Self::Steps { steps, .. } => Some(Steps::regular(steps)), _ => None }
885 }
886 pub fn repeats(&self) -> Option<usize> { self.regular().and_then(|(len, step)| (step == 0 && len > 1).then_some(len)) }
888 pub fn runs(&self, positions: usize) -> Option<(usize, isize, Vec<usize>)> {
890 let (base, steps) = match self {
891 Self::Linear(r) => return Some((*r, 0, (0..positions.div_ceil(*r)).collect())), Self::Steps { base, steps } => (base, steps), _ => return None
892 };
893 if self.fills() { return None; }
894 let (len, stride) = Steps::regular(steps);
895 let (mut width, mut outer) = (1, steps.to_vec());
897 for axis in outer.iter_mut().rev() {
898 if width == len { break; }
899 if axis.len() == 1 { continue; }
900 let Steps::Stride { start, len: n, step } = *axis else { return None };
901 *axis = Steps::Stride { start, len: 1, step };
902 width *= n;
903 }
904 Some((len, stride, Steps::combined(*base, &outer)))
905 }
906 pub fn fills(&self) -> bool { matches!(self, Self::Steps { steps, .. } if steps.iter().any(Steps::fills)) }
908 pub fn spans(&self, range: Range<usize>, items: usize, f: &mut dyn FnMut(usize, usize, Option<(usize, isize)>) -> bool) {
910 let mut at = range.start;
911 match *self {
912 Self::Single => { f(at, range.len(), (items > 0).then_some((0, 0))); }
913 Self::Linear(0) | Self::Tiled(0) => { f(at, range.len(), None); }
914 Self::Linear(r) => {
915 while at < range.end {
916 let (k, end) = (at / r, range.end.min((at / r + 1) * r));
917 if !f(at, end - at, (k < items).then_some((k, 0))) { return; }
918 at = end;
919 }
920 }
921 Self::Tiled(m) => {
922 while at < range.end {
923 let (p, end) = (at % m, range.end.min(at - at % m + m));
924 let below = items.saturating_sub(p).min(end - at);
925 if below > 0 && !f(at, below, Some((p, 1))) { return; }
926 if end > at + below && !f(at + below, end - at - below, None) { return; }
927 at = end;
928 }
929 }
930 Self::Steps { base, ref steps } => {
931 let Some((inner, outer)) = steps.split_last() else { f(at, range.len(), (base < items).then_some((base, 0))); return; };
932 if let Some((first, step)) = Steps::run(base, steps, at, range.len())
934 && first.max(first.wrapping_add_signed(step * (range.len() as isize - 1))) < items { f(at, range.len(), Some((first, step))); return; }
935 let width = inner.len();
936 while at < range.end && width > 0 {
937 let (row, j) = (at / width, at % width);
938 let n = (width - j).min(range.end - at);
939 let go = match Steps::reach(base, outer, row) { Some(start) => row_spans(start, inner, j, n, items, at, f), None => f(at, n, None) };
940 if !go { return; }
941 at += n;
942 }
943 }
944 }
945 }
946
947 fn inner(&self) -> Option<&Steps> { let Self::Steps { steps, .. } = self else { return None }; steps.iter().rev().find(|s| s.len() > 1) }
949}
950
951pub const LISTED: usize = 128;
953pub fn small_read(read: usize, len: usize) -> bool { read.saturating_mul(LISTED) <= len }
955
956pub fn paired_length(x: usize, y: usize) -> Result<usize, Error> {
958 if x == y || y == 1 { Ok(x) } else if x == 1 { Ok(y) } else { Err(Error::new(Kind::Length, format!("lengths {x} and {y} don't agree"))) }
959}
960
961#[derive(Clone)]
963pub struct Agreement<L = Layout> { pub layout: L, pub len: usize, pub left: Mapping, pub right: Mapping }
964
965impl Agreement {
966 pub fn new(left: &Layout, right: &Layout) -> Result<Self, Error> {
968 if left.plain() && right.plain() && (left.shape() == right.shape() || left.shape().is_empty() || right.shape().is_empty()) {
970 let shape = if left.shape().is_empty() { right.shape() } else { left.shape() };
971 let len = generated_len(shape)?;
972 let map = |l: &Layout| if len > 1 && l.shape() == shape { Mapping::Linear(1) } else { Mapping::Single };
973 return Ok(Self { layout: Layout::from(shape.to_vec()), len, left: map(left), right: map(right) });
974 }
975 let [xa, ya] = Self::axis_pairs(left, right);
976 Self::mapped(left, right, &xa, &ya, false)
977 }
978 pub fn axis_pairs(left: &Layout, right: &Layout) -> [Vec<usize>; 2] {
979 let (mut ya, nx) = (Self::paired_axes(left, right, true), left.shape().len());
980 if ya.len() > nx && !ya[..nx].iter().copied().eq(0..nx) && let Some(xa) = (0..nx).map(|a| ya.iter().position(|&b| b == a)).collect::<Option<Vec<_>>>() {
981 ya.iter_mut().enumerate().for_each(|(b, a)| *a = b);
982 return [xa, ya];
983 }
984 [(0..nx).collect(), ya]
985 }
986 pub fn onto(left: &Layout, right: &Layout) -> Result<Self, Error> {
988 if left.plain() && right.plain() { return Self::new(left, right); }
989 Self::mapped(left, right, &(0..left.shape().len()).collect::<Vec<_>>(), &Self::paired_axes(left, right, true), true)
990 }
991 pub fn paired_axes(left: &Layout, right: &Layout, positional: bool) -> Vec<usize> {
993 let (nx, ny) = (left.shape().len(), right.shape().len());
994 let mut ya: Vec<_> = (0..ny).map(|b| right.name(b).and_then(|n| (0..nx).find(|&a| left.name(a) == Some(n)))).collect();
995 if positional {
996 let remaining_left: Vec<_> = (0..nx).filter(|a| !ya.contains(&Some(*a))).collect();
997 let remaining_right: Vec<_> = (0..ny).filter(|&b| ya[b].is_none()).collect();
998 for (a, b) in remaining_left.into_iter().zip(remaining_right) {
999 if !names_conflict(left.name(a).map(|n| &**n), right.name(b).map(|n| &**n)) { ya[b] = Some(a); }
1000 }
1001 }
1002 let mut next = nx;
1003 ya.into_iter().map(|a| { a.unwrap_or_else(|| { next += 1; next - 1 }) }).collect()
1004 }
1005 pub fn paired(left: &Layout, right: &Layout, axes: &[usize]) -> Result<Self, Error> {
1007 Self::mapped(left, right, &(0..left.shape().len()).collect::<Vec<_>>(), axes, false)
1008 }
1009 pub fn with_axes(left: &Layout, right: &Layout, axes: &[usize]) -> Result<Self, Error> {
1011 if left.shape().len() < right.shape().len() { Self::mapped(left, right, axes, &(0..right.shape().len()).collect::<Vec<_>>(), false) }
1012 else { Self::mapped(left, right, &(0..left.shape().len()).collect::<Vec<_>>(), axes, false) }
1013 }
1014 pub fn outer(left: &Layout, right: &Layout) -> Result<Self, Error> {
1016 let nx = left.shape().len();
1017 Self::mapped(left, right, &(0..nx).collect::<Vec<_>>(), &(nx..nx + right.shape().len()).collect::<Vec<_>>(), false)
1018 }
1019 pub fn mapped(left: &Layout, right: &Layout, xa: &[usize], ya: &[usize], left_keys: bool) -> Result<Self, Error> {
1020 if xa.len() != left.shape().len() || ya.len() != right.shape().len() { return Err(Error::new(Kind::Length, "axis maps must match input ranks")); }
1021 let mut shape = vec![1; xa.iter().chain(ya).max().map_or(0, |a| a + 1)];
1022 let mut keys = vec![None; shape.len()];
1023 let mut names = vec![None; shape.len()];
1024 let (mut lt, mut rt) = (vec![Alignment::Same; left.shape().len()], vec![Alignment::Same; right.shape().len()]);
1025 for axis in 0..shape.len() {
1026 let (x, y) = (xa.iter().position(|&a| a == axis), ya.iter().position(|&a| a == axis));
1027 let (nx, ny) = (x.map_or(1, |a| left.shape()[a]), y.map_or(1, |a| right.shape()[a]));
1028 let (kx, ky) = (x.and_then(|a| left.keys(a)), y.and_then(|a| right.keys(a)));
1029 let (lx, ly) = (x.and_then(|a| left.name(a)), y.and_then(|a| right.name(a)));
1030 names[axis] = match (lx, ly) { (Some(l), Some(r)) if hidden(l) && !hidden(r) => Some(r), _ => lx.or(ly) }.cloned();
1031 if let (Some(kx), Some(ky), Some(a), Some(b)) = (kx, ky, x, y) {
1032 let key = if left_keys { rt[b] = ky.aligned(kx); kx.clone() } else {
1033 let (key, [l, r]) = kx.union(ky)?;
1034 (lt[a], rt[b]) = (l, r);
1035 key
1036 };
1037 shape[axis] = key.len();
1038 keys[axis] = Some(key);
1039 } else {
1040 let n = paired_length(nx, ny)?;
1041 shape[axis] = n;
1042 keys[axis] = kx.filter(|_| nx == n).or(ky.filter(|_| ny == n)).cloned();
1043 }
1044 }
1045 let len = generated_len(&shape)?;
1046 let layout = Layout::from(shape).with_keys(keys)?.inherit_names(names);
1047 let (x, y) = (Mapping::aligned(left, xa, &layout, len, lt), Mapping::aligned(right, ya, &layout, len, rt));
1048 Ok(Self { layout, len, left: x, right: y })
1049 }
1050}
1051
1052impl<L> Agreement<L> {
1053 pub fn map_layout<M>(self, f: impl FnOnce(L) -> M) -> Agreement<M> {
1055 Agreement { layout: f(self.layout), len: self.len, left: self.left, right: self.right }
1056 }
1057}
1058
1059pub struct Frames { pub layout: Layout, pub count: usize, cells: Vec<Option<Vec<usize>>> }
1061
1062impl Frames {
1063 pub fn new(args: &[(&Layout, usize)]) -> Result<Self, Error> {
1065 let frames = args.iter().map(|&(l, core)| {
1066 let rank = l.shape().len();
1067 if rank < core { return Err(Error::new(Kind::Rank, format!("an argument needs at least {core} axes, not {rank}"))); }
1068 Ok(l.axes(0..rank - core))
1069 }).collect::<Result<Vec<_>, _>>()?;
1070 let Some((first, rest)) = frames.split_first() else { return Ok(Self { layout: Layout::EMPTY, count: 1, cells: Vec::new() }) };
1071 let (mut layout, mut chains) = (first.clone(), vec![Vec::<Mapping>::new()]);
1072 for frame in rest {
1073 let a = Agreement::new(&layout, frame)?;
1074 for chain in &mut chains { chain.push(a.left.clone()); }
1075 chains.push(vec![a.right]);
1076 layout = a.layout;
1077 }
1078 let count = generated_len(layout.shape())?;
1079 let cells = chains.iter().map(|chain| {
1081 if chain.iter().all(|m| *m == Mapping::Linear(1)) { return Some(None); }
1082 (0..count).map(|i| chain.iter().rev().try_fold(i, |j, m| m.get(j))).collect::<Option<Vec<_>>>().map(Some)
1083 }).collect::<Option<_>>().ok_or_else(|| Error::new(Kind::Length, "keyed frames need a cell that an argument lacks"))?;
1084 Ok(Self { layout, count, cells })
1085 }
1086
1087 #[inline]
1089 pub fn one() -> Self { Self { layout: Layout::EMPTY, count: 1, cells: Vec::new() } }
1090
1091 #[inline]
1093 pub fn cell(&self, k: usize, i: usize) -> usize { match self.cells.get(k) { Some(Some(c)) => c[i], _ => i } }
1094
1095 #[inline]
1097 pub fn items(&self, k: usize, i: usize, per: usize) -> std::ops::Range<usize> { let c = self.cell(k, i); c * per..(c + 1) * per }
1098}
1099
1100const LONG_CYCLE: usize = 256;
1102pub(crate) const LINE: usize = 128;
1104const LINE_BLOCK: usize = 1 << 16;
1106
1107fn single<A: Copy>(mapping: &Mapping, data: &[A], start: usize, n: usize) -> Option<A> {
1109 match *mapping {
1110 Mapping::Single => Some(data[0]), Mapping::Linear(r) if r != 1 && n <= r - start % r => Some(data[start / r]),
1111 Mapping::Steps { .. } => mapping.run(start, n).and_then(|(o, step)| (step == 0).then(|| data[o])), _ => None
1112 }
1113}
1114
1115pub fn expanded_items<'a, A: Element>(mapping: &Walked, data: &'a [A], start: usize, n: usize, buf: &'a mut Vec<A>) -> &'a [A] {
1117 if let Some(x) = single(mapping, data, start, n) {
1118 buf.clear();
1119 buf.resize(n, x);
1120 return buf;
1121 }
1122 match expanded(mapping, data, start, n, buf) { Arg::Many(items) => items, Arg::One(_) => unreachable!("`single` gives every repeated item") }
1123}
1124
1125pub fn in_place<'a, A: Copy>(mapping: &Walked, data: &'a [A], start: usize, n: usize) -> Option<Arg<'a, A>> {
1127 if let Mapping::Steps { .. } = **mapping {
1128 return match mapping.run(start, n) {
1130 Some((o, 0)) => Some(Arg::One(data[o])), Some((o, 1)) => Some(Arg::Many(&data[o..o + n])),
1131 _ => mapping.cycle().filter(|&p| start % p + n <= p).map(|p| Arg::Many(&data[start % p..start % p + n]))
1132 };
1133 }
1134 if let Some(x) = single(mapping, data, start, n) { return Some(Arg::One(x)); }
1135 match **mapping {
1136 Mapping::Linear(1) => Some(Arg::Many(&data[start..start + n])),
1137 Mapping::Tiled(m) if start % m + n <= data.len() => Some(Arg::Many(&data[start % m..start % m + n])), _ => None
1138 }
1139}
1140
1141pub fn expanded<'a, A: Element>(mapping: &Walked, data: &'a [A], start: usize, n: usize, buf: &'a mut Vec<A>) -> Arg<'a, A> {
1143 if let Some(arg) = in_place(mapping, data, start, n) { return arg; }
1144 Arg::Many(copied(mapping, data, start, n, buf))
1145}
1146
1147pub fn tiled_cycle<A: Element>(data: &[A], mapping: &Mapping, run: usize, buf: &mut Vec<A>) -> Option<Mapping> {
1149 let period = mapping.period()?;
1150 let (p, cover) = period.tiling(run)?;
1151 match period {
1152 Period::Walks { base, steps, .. } => { copied(&Walked::from(Mapping::Steps { base, steps: steps.to_vec().into() }), data, 0, p, buf); }
1153 _ => {
1154 buf.clear();
1155 buf.extend_from_slice(&data[..p]);
1156 }
1157 }
1158 buf.truncate(p);
1159 if p < cover { *buf = buf.repeat(cover.div_ceil(p) + 1); }
1160 Some(Mapping::Tiled(p))
1161}
1162
1163pub fn copied<'a, A: Element>(mapping: &Walked, data: &[A], start: usize, n: usize, buf: &'a mut Vec<A>) -> &'a [A] {
1165 buf.clear();
1166 let reason = copy_walked(mapping, data, start, n, A::FILL, buf);
1167 crate::events::declined(crate::events::Route::InPlace, reason, n);
1168 &buf[..n]
1169}
1170
1171#[inline(never)]
1173fn copy_walked<A: Element>(mapping: &Walked, data: &[A], start: usize, n: usize, fill: A, buf: &mut Vec<A>) -> &'static str {
1174 let at = buf.len();
1175 match **mapping {
1176 _ if n == 0 => "no items",
1177 Mapping::Single => { buf.resize(at + n, data[0]); "a repeated item" }
1178 Mapping::Linear(r) => {
1179 let (k0, end) = (start / r, start + n);
1181 buf.resize(at + ((k0 + 1) * r).min(end) - start, data[k0]);
1182 if buf.len() - at < n {
1183 let (k1, k2) = (k0 + 1, end / r);
1184 crate::gather::repeat(buf, k2 - k1, &[r as i64], (k2 - k1) * r, |j| data[k1 + j]);
1185 if end % r != 0 { buf.resize(at + n, data[k2]) }
1186 }
1187 "repeats of each item"
1188 }
1189 Mapping::Tiled(m) => { cycled(buf, &data[..m], start, n); "a cycle that wraps" }
1190 Mapping::Steps { base, .. } => {
1191 let (offset, plan) = mapping.walk();
1192 walk_plan(buf, data, base + offset, plan, fill, start, n, Same);
1193 "steps other than a forward run"
1194 }
1195 }
1196}