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(Arc<ReadMap>), Join(Arc<[ReadPart]>), Default, Invalid
66}
67
68#[derive(Debug, PartialEq, Eq)]
69pub struct ReadMap { pub input: ReadPlan, pub map: Mapping, pub missing: ReadPlan }
70
71#[derive(Clone, Debug, PartialEq, Eq)]
72pub struct ReadPlan { route: ReadRoute, items: usize }
73
74impl ReadPlan {
75 pub fn input(source: usize, items: usize) -> Self { Self { route: ReadRoute::Input(source), items } }
76 pub fn fill(items: usize) -> Self { Self { route: ReadRoute::Default, items } }
77 pub fn index(base: i64, step: i64, items: usize) -> Self { Self { route: ReadRoute::Index { base, step }, items } }
78 fn fallback(missing: Missing, items: usize) -> Self {
79 match missing {
80 Missing::Input(plan) => plan.remapped(Mapping::Linear(1), items, Missing::Invalid),
81 Missing::Scalar(source) => Self::input(source, 1).remapped(Mapping::Single, items, Missing::Invalid),
82 Missing::Default => Self::fill(items), Missing::Invalid => Self { route: ReadRoute::Invalid, items }
83 }
84 }
85 pub fn copy(source: usize, source_items: usize, map: Mapping, items: usize, missing: Missing) -> Self {
86 Self::input(source, source_items).remapped(map, items, missing)
87 }
88 pub fn joined(rows: usize, parts: impl IntoIterator<Item = Self>) -> Result<Self, Error> {
90 let too_large = || Error::new(Kind::Length, "joined source is too large");
91 let (mut joined, mut width) = (Vec::new(), 0usize);
92 for plan in parts {
93 let share = if rows == 0 { 0 } else { plan.len() / rows };
94 if share * rows != plan.len() { return Err(Error::new(Kind::Length, "a joined part doesn't fill its rows")); }
95 width = width.checked_add(share).ok_or_else(too_large)?;
96 joined.push(ReadPart { end: width, plan });
97 }
98 Ok(Self { route: ReadRoute::Join(joined.into()), items: width.checked_mul(rows).ok_or_else(too_large)? })
99 }
100 pub fn len(&self) -> usize { self.items }
101 pub fn is_empty(&self) -> bool { self.items == 0 }
102 pub fn route(&self) -> &ReadRoute { &self.route }
103 pub fn at(&self, item: usize) -> Run {
105 let mut found = Run::Invalid;
106 if item < self.items { self.runs(item..item + 1, &mut |_, _, run| { found = run; false }); }
107 found
108 }
109 pub fn runs(&self, range: Range<usize>, sink: &mut dyn RunSink) -> bool { self.runs_through(range.start, 1, range.len(), range.start, sink) }
111 fn runs_through(&self, first: usize, stride: isize, len: usize, at: usize, sink: &mut dyn RunSink) -> bool {
113 if len == 0 { return true; }
114 match self.route() {
115 &ReadRoute::Input(source) => sink.run(at, len, Run::Source { source, first, stride }),
116 &ReadRoute::Index { base, step } => {
117 sink.run(at, len, Run::Index { first: base.wrapping_add(step.wrapping_mul(first as i64)), step: step.wrapping_mul(stride as i64) })
118 }
119 ReadRoute::Default => sink.run(at, len, Run::Fill),
120 ReadRoute::Invalid => sink.run(at, len, Run::Invalid),
121 ReadRoute::Join(parts) => {
122 let mut j = 0;
123 while j < len {
124 let Some((plan, item, before, after)) = joined_at(parts, first.wrapping_add_signed(stride * j as isize)) else {
125 return sink.run(at + j, len - j, Run::Invalid);
126 };
127 let n = (len - j).min(match stride.cmp(&0) {
128 std::cmp::Ordering::Greater => after.div_ceil(stride as usize), std::cmp::Ordering::Less => before / stride.unsigned_abs() + 1,
129 std::cmp::Ordering::Equal => len - j
130 });
131 if !plan.runs_through(item, stride, n, at + j, sink) { return false; }
132 j += n;
133 }
134 true
135 }
136 ReadRoute::Map(read) => {
137 let ReadMap { input, map, missing } = read.as_ref();
138 if input.is_empty() { return missing.runs_through(first, stride, len, at, sink); }
139 match stride {
140 1 => {
141 if let Some(go) = sink.walk(at, input, map, missing, first..first + len) { return go; }
142 let mut go = true;
143 map.spans(first..first + len, input.len(), &mut |p, n, run| {
144 go = match run {
145 Some((q, s)) => input.runs_through(q, s, n, at + p - first, sink), None => missing.runs_through(p, 1, n, at + p - first, sink)
146 };
147 go
148 });
149 go
150 }
151 0 => match map.get(first).filter(|&q| q < input.len()) {
152 Some(q) => input.runs_through(q, 0, len, at, sink), None => missing.runs_through(first, 0, len, at, sink)
153 },
154 _ => (0..len).all(|j| {
155 let p = first.wrapping_add_signed(stride * j as isize);
156 match map.get(p).filter(|&q| q < input.len()) {
157 Some(q) => input.runs_through(q, 1, 1, at + j, sink), None => missing.runs_through(p, 1, 1, at + j, sink)
158 }
159 }),
160 }
161 }
162 }
163 }
164 fn mapped(input: ReadPlan, map: Mapping, missing: Missing, items: usize) -> Self {
165 if map.in_order(input.items, &[items]) { return input; }
166 if let Some(plan) = Self::single_map(&input, &map, items) { return plan; }
167 Self { route: ReadRoute::Map(Arc::new(ReadMap { input, map, missing: Self::fallback(missing, items) })), items }
168 }
169 fn single_map(input: &ReadPlan, map: &Mapping, items: usize) -> Option<Self> {
171 let ReadRoute::Map(read) = input.route() else { return None };
172 let ReadMap { input: inner, map: first, missing } = read.as_ref();
173 let frame = map.shape(items);
174 if map.bounds(&frame)?.end > input.items { return None; }
175 let misses = match missing.route() { ReadRoute::Default => Missing::Default, ReadRoute::Invalid => Missing::Invalid, _ => return None };
176 let composed = first.compose(map, &input.shape_hint(), &frame)?;
177 (composed.footprint() <= first.footprint() + map.footprint()).then(|| Self::mapped(inner.clone(), composed, misses, items))
178 }
179 pub fn remapped(&self, map: Mapping, items: usize, missing: Missing) -> Self { Self::mapped(self.clone(), map, missing, items) }
180 pub fn try_for_each_source<E>(&self, visit: &mut impl FnMut(usize, usize) -> Result<(), E>) -> Result<(), E> {
182 match self.route() {
183 ReadRoute::Input(source) => visit(*source, self.items),
184 ReadRoute::Map(read) => { read.input.try_for_each_source(visit)?; read.missing.try_for_each_source(visit) }
185 ReadRoute::Join(parts) => { for part in &**parts { part.plan.try_for_each_source(visit)?; } Ok(()) }
186 ReadRoute::Index { .. } | ReadRoute::Default | ReadRoute::Invalid => Ok(()),
187 }
188 }
189 pub fn as_codes(&self, character: bool) -> Cow<'_, Self> {
191 if !character { return Cow::Borrowed(self); }
192 Cow::Owned(self.defaulted(<char as Element>::FILL as i64))
193 }
194 fn defaulted(&self, value: i64) -> Self {
196 match self.route() {
197 ReadRoute::Map(read) => Self::mapped(read.input.defaulted(value), read.map.clone(), Missing::Input(read.missing.defaulted(value)), self.items),
198 ReadRoute::Join(parts) => {
199 let parts: Vec<_> = parts.iter().map(|part| ReadPart { end: part.end, plan: part.plan.defaulted(value) }).collect();
200 Self { route: ReadRoute::Join(parts.into()), items: self.items }
201 }
202 ReadRoute::Default => Self::index(value, 0, self.items),
203 ReadRoute::Input(_) | ReadRoute::Index { .. } | ReadRoute::Invalid => self.clone(),
204 }
205 }
206 pub fn contiguous(&self, range: Range<usize>) -> Option<(usize, Range<usize>)> {
207 let (source, start, stride) = self.run(&range)?;
208 (stride == 1 || range.len() == 1).then(|| (source, start..start + range.len()))
209 }
210 pub fn shape_hint(&self) -> Vec<usize> { match self.route() { ReadRoute::Map(read) => read.map.shape(self.items), _ => vec![self.items] } }
211 pub fn positional(&self, shape: &[usize]) -> Option<(usize, Mapping)> {
213 if count(shape) != self.items { return None; }
214 match self.route() {
215 ReadRoute::Input(source) => Some((*source, Mapping::Linear(1))),
216 ReadRoute::Map(read) => {
217 let ReadMap { input, map, .. } = read.as_ref();
218 if input.is_empty() || map.bounds(shape)?.end > input.len() { return None; }
219 let inner_shape = input.shape_hint();
220 let (source, inner) = input.positional(&inner_shape)?;
221 Some((source, inner.compose(map, &inner_shape, shape)?))
222 }
223 ReadRoute::Join(parts) if parts.len() == 1 => parts[0].plan.positional(shape),
224 _ => None,
225 }
226 }
227 fn affine_into(&self, shape: &[usize], strides: &mut [isize]) -> Option<(usize, usize, Range<usize>)> {
228 if count(shape) != self.items || shape.len() != strides.len() { return None; }
229 let (source, map, items) = match self.route() {
230 ReadRoute::Input(source) => (*source, &Mapping::Linear(1), self.items),
231 ReadRoute::Map(read) => {
232 let ReadMap { input, map, .. } = read.as_ref();
233 let &ReadRoute::Input(source) = input.route() else { return None };
234 (source, map, input.len())
235 }
236 ReadRoute::Join(parts) if parts.len() == 1 => return parts[0].plan.affine_into(shape, strides),
237 _ => return None,
238 };
239 let base = map.affine_axes(shape, &mut |k, step| strides[k] = step)?;
240 let bounds = affine_bounds(base, strides, shape)?;
241 (bounds.end <= items).then_some((source, base, bounds))
242 }
243 pub fn inverse(&self, source: usize, shape: &[usize], items: usize) -> Option<Mapping> {
245 if self.items != items { return None; }
246 let (found, map) = self.positional(shape)?;
247 if found != source { return None; }
248 if items == 0 { return Some(Mapping::Linear(1)); }
249 let (base, forward) = map.affine(shape)?;
250 let mut axes: Vec<_> = (0..shape.len()).filter(|&axis| shape[axis] > 1).collect();
251 axes.sort_unstable_by_key(|&axis| forward[axis].unsigned_abs());
252 let (mut covered, mut first) = (1usize, base as i128);
253 for &axis in &axes {
254 if forward[axis].unsigned_abs() != covered { return None; }
255 covered = covered.checked_mul(shape[axis])?;
256 if forward[axis] < 0 { first += forward[axis] as i128 * (shape[axis] - 1) as i128; }
257 }
258 if covered != items || first != 0 { return None; }
259 let strides = crate::layout::strides(shape);
260 let steps = axes.into_iter().rev().map(|axis| {
261 let stride = isize::try_from(strides[axis]).ok()?;
262 let reverse = forward[axis] < 0;
263 Some(Steps::Stride { start: if reverse { (shape[axis] - 1) * strides[axis] } else { 0 }, len: shape[axis],
264 step: if reverse { -stride } else { stride } })
265 }).collect::<Option<Vec<_>>>()?;
266 Some(Mapping::from(steps))
267 }
268 fn run(&self, range: &Range<usize>) -> Option<(usize, usize, isize)> {
270 if range.is_empty() || range.end > self.items { return None; }
271 let n = range.len();
272 let (source, first, stride) = match self.route() {
273 &ReadRoute::Input(source) => (source, range.start, 1),
274 ReadRoute::Map(read) => {
275 let ReadMap { input, map, missing } = read.as_ref();
276 if input.is_empty() { return missing.run(range); }
277 let (first, stride) = map.run(range.start, n)?;
278 let stride = if n == 1 { 1 } else { stride };
279 if !matches!(stride, 0 | 1) { return None; }
280 if first >= input.len() { return missing.run(range); }
281 let end = first.checked_add(if stride == 0 { 1 } else { n })?;
282 let (source, at, step) = input.run(&(first..end))?;
283 (source, at, if stride == 0 { 0 } else { step })
284 }
285 ReadRoute::Join(parts) => {
286 let (part, at, _, after) = joined_at(parts, range.start)?;
287 if n > after { return None; }
288 part.run(&(at..at.checked_add(n)?))?
289 }
290 ReadRoute::Index { .. } | ReadRoute::Default | ReadRoute::Invalid => return None
291 };
292 Some((source, first, if n == 1 { 1 } else { stride }))
293 }
294}
295
296fn affine_bounds(base: usize, strides: &[isize], shape: &[usize]) -> Option<Range<usize>> {
297 if shape.len() != strides.len() { return None; }
298 if shape.contains(&0) { return Some(0..0); }
299 let (mut lo, mut hi) = (base as i128, base as i128);
300 for (&len, &stride) in shape.iter().zip(strides) {
301 let delta = (stride as i128).checked_mul(len.saturating_sub(1) as i128)?;
302 if delta < 0 { lo = lo.checked_add(delta)?; } else { hi = hi.checked_add(delta)?; }
303 }
304 Some(usize::try_from(lo).ok()?..usize::try_from(hi).ok()?.checked_add(1)?)
305}
306
307fn row_spans(
309 start: usize, axis: &Steps, j0: usize, n: usize, items: usize, at: usize, f: &mut dyn FnMut(usize, usize, Option<(usize, isize)>) -> bool) -> bool {
310 if let Steps::Stride { start: offset, step, .. } = *axis {
311 let first = start + offset.wrapping_add_signed(step * j0 as isize);
312 return match step.cmp(&0) {
313 std::cmp::Ordering::Greater => {
315 let below = if first >= items { 0 } else { n.min((items - 1 - first) / step as usize + 1) };
316 (below == 0 || f(at, below, Some((first, step)))) && (below == n || f(at + below, n - below, None))
317 }
318 std::cmp::Ordering::Equal => f(at, n, (first < items).then_some((first, 0))),
319 std::cmp::Ordering::Less => {
320 let past = if first < items { 0 } else { n.min((first - items) / step.unsigned_abs() + 1) };
321 (past == 0 || f(at, past, None)) && (past == n || f(at + past, n - past, Some((first.wrapping_add_signed(step * past as isize), step))))
322 }
323 };
324 }
325 let position = |j: usize| axis.get(j).map(|o| start + o).filter(|&p| p < items);
326 let (mut j, end) = (j0, j0 + n);
327 while j < end {
328 let first = position(j);
329 let step = match (first, (j + 1 < end).then(|| position(j + 1)).flatten()) { (Some(p), Some(q)) => q as isize - p as isize, _ => 0 };
330 let mut k = j + 1;
331 while k < end && match (first, position(k)) {
332 (None, None) => true, (Some(p), Some(q)) => q as isize - p as isize == step * (k - j) as isize, _ => false
333 } { k += 1; }
334 if !f(at + j - j0, k - j, first.map(|p| (p, step))) { return false; }
335 j = k;
336 }
337 true
338}
339
340#[derive(Clone, Copy)]
342pub struct Binding<'a> { pub view: crate::buffer::View<'a>, pub start: usize, pub items: usize }
343
344impl<'a> Binding<'a> {
345 pub fn full(view: crate::buffer::View<'a>) -> Self { Self { items: view.len(), view, start: 0 } }
346 pub fn as_codes(self) -> Self { Self { view: crate::characters::as_codes(self.view), ..self } }
347 pub fn borrow(&self, range: Range<usize>) -> Option<crate::buffer::View<'a>> {
348 if range.start > range.end || range.end > self.items { return None; }
349 let local = range.start.checked_sub(self.start)?..range.end.checked_sub(self.start)?;
350 (local.end <= self.view.len()).then(|| self.view.slice(local))
351 }
352 pub fn affine(&self, base: usize, strides: &[isize], shape: &[usize]) -> Option<(crate::buffer::View<'a>, usize)> {
353 self.borrow(affine_bounds(base, strides, shape)?)?;
354 Some((self.view, base.checked_sub(self.start)?))
355 }
356 pub fn window(view: crate::buffer::View<'a>, start: usize, items: usize) -> Result<Self, Error> {
357 if start.checked_add(view.len()).is_none_or(|end| end > items) { return Err(Error::new(Kind::Length, "read window exceeds its producer")); }
358 Ok(Self { view, start, items })
359 }
360}
361
362#[derive(Clone, Copy)]
364pub struct Read<'a> { pub sources: &'a [Binding<'a>], pub plan: &'a ReadPlan }
365
366pub struct Rows<'a, T> { pub items: &'a [T], pub first: usize, pub step: isize, pub run: bool }
368
369impl<'s> Read<'s> {
370 fn bindings(&self) -> Result<(), Error> {
372 self.plan.try_for_each_source(&mut |source, count| {
373 let binding = self.sources.get(source).ok_or_else(|| Error::new(Kind::Length, "missing read source"))?;
374 if binding.items != count { return Err(Error::new(Kind::Length, "read source changed length")); }
375 if binding.start.checked_add(binding.view.len()).is_none_or(|end| end > count) {
376 return Err(Error::new(Kind::Length, "read window exceeds its producer"));
377 }
378 Ok(())
379 })
380 }
381 pub fn positional(&self, shape: &[usize]) -> Option<(crate::buffer::View<'_>, Mapping)> {
382 let (source, mapping) = self.plan.positional(shape)?;
383 let binding = self.sources.get(source)?;
384 binding.borrow(mapping.bounds(shape)?)?;
385 let (base, steps) = mapping.entries(shape)?;
386 Some((binding.view, Mapping::Steps { base: base.checked_sub(binding.start)?, steps }))
387 }
388 fn range(&self, range: &Range<usize>) -> Result<(), Error> {
389 if range.start > range.end || range.end > self.plan.len() { return Err(Error::new(Kind::Length, "read range exceeds its result")); }
390 self.bindings()
391 }
392 fn run_view(&self, range: &Range<usize>) -> Option<(crate::buffer::View<'s>, usize)> {
394 let (source, start, stride) = self.plan.run(range)?;
395 let binding = self.sources.get(source)?;
396 let end = start.checked_add(if stride == 0 { 1 } else { range.len() })?;
397 Some((binding.borrow(start..end)?, stride as usize))
398 }
399 pub fn part_view<'a>(
401 &'a self, range: Range<usize>, dtype: crate::buffer::DType, checked: bool, scratch: &'a mut crate::buffer::Buffer
402 ) -> Result<crate::buffer::View<'a>, Error> {
403 self.range(&range)?;
404 let run = self.run_view(&range);
405 if let Some((view, stride)) = run && view.dtype() == dtype && (stride == 1 || range.len() == 1) { return Ok(view); }
406 self.write_into(range, dtype, checked, scratch, run)?;
407 Ok(scratch.view())
408 }
409
410 pub fn argument<'a>(
412 self, range: Range<usize>, checked: bool, scratch: &'a mut crate::buffer::Buffer
413 ) -> Result<(crate::buffer::View<'a>, usize), Error> where 's: 'a {
414 self.range(&range)?;
415 if let Some((view, _)) = self.run_view(&range) { return Ok((view, 0)); }
416 let mut dtype = None;
417 self.plan.try_for_each_source(&mut |source, _| {
418 let current = self.sources[source].view.dtype();
419 if dtype.is_some_and(|dtype| dtype != current) { return Err(Error::new(Kind::Domain, "read sources have different storage dtypes")); }
420 dtype = Some(current);
421 Ok(())
422 })?;
423 let dtype = dtype.ok_or_else(|| Error::new(Kind::Domain, "read has no storage dtype"))?;
424 let copied = range.len();
425 self.write_into(range, dtype, checked, scratch, None)?;
426 Ok((scratch.view(), copied))
427 }
428 pub fn borrowed<T: crate::buffer::Cast>(&self, range: Range<usize>, checked: bool) -> Result<Option<Arg<'s, T>>, Error> {
430 self.range(&range)?;
431 if range.is_empty() { return Ok(Some(Arg::Many(&[]))); }
432 match self.run_view(&range) { Some((view, _)) => borrowed_as(view, checked), None => Ok(None) }
433 }
434 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> {
435 self.part_counted(range, checked, scratch).map(|(items, _)| items)
436 }
437 pub fn part_counted<'a, T: crate::buffer::Cast>(
439 &'a self, range: Range<usize>, checked: bool, scratch: &'a mut Vec<T>
440 ) -> Result<(Arg<'a, T>, usize), Error> {
441 self.range(&range)?;
442 if range.is_empty() { return Ok((Arg::Many(&[]), 0)); }
443 if let Some((view, stride)) = self.run_view(&range) {
444 if let Some(items) = borrowed_as::<T>(view, checked)? { return Ok((items, 0)); }
445 cast_run(view, stride, range.len(), checked, scratch)?;
446 return Ok((Arg::Many(scratch), range.len()));
447 }
448 scratch.clear();
449 let copied = range.len();
450 self.read_onto(range, checked, scratch)?;
451 Ok((Arg::Many(scratch), copied))
452 }
453 pub fn in_place<T: crate::buffer::Cast>(&self, range: &Range<usize>) -> bool {
455 self.run_view(range).is_some_and(|(view, _)| view.len() == 1 || view.dtype() == T::DTYPE)
456 }
457 pub fn rows<T: crate::buffer::Cast>(&self, shape: [usize; 2]) -> Option<Rows<'s, T>> {
459 self.bindings().ok()?;
460 let mut strides = [0; 2];
461 let (source, base, bounds) = self.plan.affine_into(&shape, &mut strides)?;
462 let binding = self.sources.get(source)?;
463 binding.borrow(bounds)?;
464 let [step, inner @ (0 | 1)] = strides else { return None };
465 Some(Rows { items: T::of(binding.view)?, first: base.checked_sub(binding.start)?, step, run: inner == 1 })
466 }
467 fn read_onto<T: crate::buffer::Cast>(&self, range: Range<usize>, checked: bool, output: &mut Vec<T>) -> Result<(), Error> {
469 output.reserve(range.len());
470 if self.join_onto(&range, checked, output)? { return Ok(()); }
471 let mut onto = Onto { read: self, checked, output, error: None };
472 self.plan.runs(range, &mut onto);
473 onto.error.map_or(Ok(()), Err)
474 }
475 fn join_onto<T: crate::buffer::Cast>(&self, range: &Range<usize>, checked: bool, output: &mut Vec<T>) -> Result<bool, Error> {
477 let ReadRoute::Join(parts) = self.plan.route() else { return Ok(false) };
478 let width = parts.last().map_or(0, |part| part.end);
479 if width == 0 || range.start % width != 0 || range.len() % width != 0 { return Ok(false); }
480 let (first, rows) = (range.start / width, range.len() / width);
481 let (mut lanes, mut start) = (Vec::with_capacity(parts.len()), 0);
482 for part in parts.iter() {
483 let n = part.end - std::mem::replace(&mut start, part.end);
484 let Some((view, stride)) = (Read { sources: self.sources, plan: &part.plan }).run_view(&(first * n..(first + rows) * n)) else { return Ok(false) };
485 lanes.push((match (T::of(view), stride) {
486 (Some(items), 1) => Cow::Borrowed(items),
487 (_, 0) => Cow::Owned(vec![kept(T::cast(view, 0), checked)?; rows * n]),
488 _ => {
489 let mut items = Vec::new();
490 T::extend_cast(view, rows * n, checked, &mut items)?;
491 Cow::Owned(items)
492 }
493 }, n));
494 }
495 crate::gather::interleave(output, &lanes.iter().map(|(items, n)| (&items[..], *n)).collect::<Vec<_>>(), rows);
496 Ok(true)
497 }
498 fn run_onto<T: crate::buffer::Cast>(
500 &self, source: usize, first: usize, stride: isize, len: usize, checked: bool, output: &mut Vec<T>) -> Result<(), Error> {
501 if len == 0 { return Ok(()); }
502 let outside = || Error::new(Kind::Length, "read reaches outside its bound window");
503 let last = first.checked_add_signed(stride * (len as isize - 1)).ok_or_else(outside)?;
504 let low = first.min(last);
505 let view = self.sources[source].borrow(low..first.max(last) + 1).ok_or_else(outside)?;
506 if stride == 0 {
507 let item = kept(T::cast(view, 0), checked)?;
508 output.resize(output.len() + len, item);
509 return Ok(());
510 }
511 if stride == 1 {
512 if let Some(items) = T::of(view) { output.extend_from_slice(items); return Ok(()); }
513 return T::extend_cast(view, len, checked, output);
514 }
515 let walked = Walked::from(Mapping::Steps { base: 0, steps: vec![Steps::Stride { start: first - low, len, step: stride }] });
516 if let Some(items) = T::of(view) { copy_walked(&walked, items, 0, len, T::FILL, output); return Ok(()); }
517 with_dtype!(view.dtype(), S => {
518 let mut items = Vec::with_capacity(len);
519 copy_walked(&walked, S::of(view).expect("the view has its dtype"), 0, len, S::FILL, &mut items);
520 T::extend_cast(S::view(&items), len, checked, output)
521 })
522 }
523 fn walkable(&self, input: &ReadPlan, map: &Mapping, end: usize) -> Option<&'s Binding<'s>> {
525 let &ReadRoute::Input(source) = input.route() else { return None };
526 let binding = self.sources.get(source)?;
527 let len = binding.view.len();
528 (binding.start == 0 && len == input.len() && match *map {
529 Mapping::Single => true, Mapping::Linear(r) => r > 0 && end.saturating_sub(1) / r < len, Mapping::Tiled(m) => m > 0 && m <= len,
530 Mapping::Steps { base, ref steps } => Steps::furthest(base, &steps[..]).is_none_or(|furthest| furthest < len)
531 }).then_some(binding)
532 }
533 pub fn mapped_source(&self) -> Option<(crate::buffer::View<'s>, &'s Mapping)> {
535 let ReadRoute::Map(read) = self.plan.route() else { return None };
536 let ReadMap { input, map, .. } = read.as_ref();
537 if map.fills() { return None; }
538 Some((self.walkable(input, map, self.plan.len())?.view, map))
539 }
540 fn walk_onto<T: crate::buffer::Cast>(
542 &self, input: &ReadPlan, map: &Mapping, missing: &ReadPlan, range: Range<usize>, checked: bool, output: &mut Vec<T>) -> Result<bool, Error> {
543 let Some(binding) = self.walkable(input, map, range.end) else { return Ok(false) };
544 let Some(fill) = Filled::of(missing).filter(|&fill| fill != Filled::Never || !map.fills()) else { return Ok(false) };
545 let walked = Walked::from(map);
546 if let Some(items) = T::of(binding.view) {
547 let Some(fill) = fill.item::<T>(self.sources) else { return Ok(false) };
548 copy_walked(&walked, items, range.start, range.len(), fill, output);
549 return Ok(true);
550 }
551 with_dtype!(binding.view.dtype(), S => {
552 let Some(fill) = fill.item::<S>(self.sources) else { return Ok(false) };
553 let mut items = Vec::with_capacity(range.len());
554 copy_walked(&walked, S::of(binding.view).expect("the view has its dtype"), range.start, range.len(), fill, &mut items);
555 T::extend_cast(S::view(&items), items.len(), checked, output)?;
556 Ok(true)
557 })
558 }
559 pub fn read_into(&self, range: Range<usize>, dtype: crate::buffer::DType, checked: bool, output: &mut crate::buffer::Buffer) -> Result<(), Error> {
560 self.range(&range)?;
561 let run = self.run_view(&range);
562 self.write_into(range, dtype, checked, output, run)
563 }
564 fn write_into(
565 &self, range: Range<usize>, dtype: crate::buffer::DType, checked: bool,
566 output: &mut crate::buffer::Buffer, run: Option<(crate::buffer::View<'s>, usize)>
567 ) -> Result<(), Error> {
568 if output.dtype() != dtype { *output = crate::buffer::Buffer::with_capacity(dtype, range.len()); }
569 with_dtype!(dtype, T => {
570 let items = T::vec_mut(output).expect("output has the selected dtype");
571 if let Some((view, stride)) = run {
572 if view.dtype() != dtype { return cast_run(view, stride, range.len(), checked, items); }
573 let source = T::of(view).expect("a typed read selected its dtype");
574 items.clear();
575 if stride == 0 { items.resize(range.len(), source[0]); } else { items.extend_from_slice(source); }
576 return Ok(());
577 }
578 items.clear();
579 self.read_onto(range, checked, items)
580 })
581 }
582}
583
584struct Onto<'r, 'a, T> { read: &'r Read<'a>, checked: bool, output: &'r mut Vec<T>, error: Option<Error> }
586
587impl<T: crate::buffer::Cast> RunSink for Onto<'_, '_, T> {
588 fn run(&mut self, _at: usize, len: usize, run: Run) -> bool {
589 let read = match run {
590 Run::Source { source, first, stride } => self.read.run_onto(source, first, stride, len, self.checked, self.output),
591 Run::Index { first, step } => {
592 let values: Vec<i64> = (0..len).map(|j| first.wrapping_add(step.wrapping_mul(j as i64))).collect();
593 T::extend_cast(crate::buffer::View::Integers(crate::buffer::Ints::I64(&values)), len, self.checked, self.output)
594 }
595 Run::Fill => { self.output.resize(self.output.len() + len, T::FILL); Ok(()) }
596 Run::Invalid => Err(Error::new(Kind::Index, "read position is outside its source")),
597 };
598 read.map_err(|e| self.error = Some(e)).is_ok()
599 }
600 fn walk(&mut self, _at: usize, input: &ReadPlan, map: &Mapping, missing: &ReadPlan, range: Range<usize>) -> Option<bool> {
601 match self.read.walk_onto(input, map, missing, range, self.checked, self.output) {
602 Ok(true) => Some(true), Ok(false) => None, Err(e) => { self.error = Some(e); Some(false) }
603 }
604 }
605}
606
607fn joined_at(parts: &[ReadPart], position: usize) -> Option<(&ReadPlan, usize, usize, usize)> {
609 let width = parts.last().map_or(0, |part| part.end);
610 if width == 0 { return None; }
611 let (row, column) = (position / width, position % width);
612 let k = parts.partition_point(|part| part.end <= column);
613 let start = if k == 0 { 0 } else { parts[k - 1].end };
614 let part = parts.get(k)?;
615 Some((&part.plan, row * (part.end - start) + column - start, column - start, part.end - column))
616}
617fn borrowed_as<'a, T: crate::buffer::Cast>(view: crate::buffer::View<'a>, checked: bool) -> Result<Option<Arg<'a, T>>, Error> {
618 if let Some(items) = T::of(view) { return Ok(Some(Arg::of(items))); }
619 if view.len() == 1 { return Ok(Some(Arg::One(kept(T::cast(view, 0), checked)?))); }
620 Ok(None)
621}
622
623fn cast_run<T: crate::buffer::Cast>(view: crate::buffer::View<'_>, stride: usize, len: usize, checked: bool, output: &mut Vec<T>) -> Result<(), Error> {
625 output.clear();
626 if stride == 0 { output.resize(len, kept(T::cast(view, 0), checked)?); return Ok(()); }
627 T::extend_cast(view, len, checked, output)
628}
629
630#[derive(Clone, Copy, PartialEq)]
632enum Filled { Fill, Value(i64), Source(usize), Never }
633
634impl Filled {
635 fn of(missing: &ReadPlan) -> Option<Self> {
636 match missing.route() {
637 ReadRoute::Default => Some(Self::Fill), ReadRoute::Invalid => Some(Self::Never), &ReadRoute::Index { base, step: 0 } => Some(Self::Value(base)),
638 &ReadRoute::Input(source) if missing.len() == 1 => Some(Self::Source(source)),
639 ReadRoute::Map(read) if matches!(&read.map, Mapping::Single) => match *read.input.route() { ReadRoute::Input(source) => Some(Self::Source(source)), _ => None },
640 _ => None
641 }
642 }
643 fn item<A: crate::buffer::Cast>(self, sources: &[Binding<'_>]) -> Option<A> {
645 match self {
646 Self::Fill | Self::Never => Some(A::FILL), Self::Value(value) => A::cast(crate::buffer::View::Integers(crate::buffer::Ints::I64(&[value])), 0).ok(),
647 Self::Source(source) => A::of(sources.get(source).filter(|binding| binding.start == 0)?.view)?.first().copied()
648 }
649 }
650}
651
652#[derive(Clone, Copy, Debug, PartialEq, Eq)]
654pub enum Period<'a> {
655 Repeats(usize),
657 Cycles(usize),
659 Walks { len: usize, base: usize, steps: &'a [Steps] },
661}
662
663impl Period<'_> {
664 pub fn tiling(&self, run: usize) -> Option<(usize, usize)> {
666 match *self {
667 Self::Cycles(p) | Self::Walks { len: p, .. } if p > 0 && p < LONG_CYCLE =>
668 Some((p, if p < run { p * (run.div_ceil(p) + 1) } else { p })), Self::Walks { len, .. } => Some((len, len)),
669 Self::Cycles(_) | Self::Repeats(_) => None
670 }
671 }
672}
673
674#[derive(Clone, Debug, PartialEq, Eq)]
676pub struct Block {
677 pub range: Range<usize>,
679 pub dims: Vec<usize>,
681 pub axes: Vec<Option<(usize, bool)>>,
683}
684
685impl Block {
686 pub fn fold(&self, axis: Option<usize>) -> ([usize; 3], Option<Vec<Steps>>) {
688 let Some(axis) = axis else { return ([1, self.range.len(), 1], None) };
689 let folded = self.axes[axis].map(|(b, _)| b);
690 let dims = &self.dims;
691 let around = match folded { Some(f) => [count(&dims[..f]), dims[f], count(&dims[f + 1..])], None => [self.range.len(), 1, 1] };
692 let rest: Vec<usize> = dims.iter().enumerate().filter(|&(b, _)| Some(b) != folded).map(|(_, &n)| n).collect();
694 let stride = strides(&rest);
695 let steps: Vec<Steps> = self.axes.iter().enumerate().filter(|&(a, _)| a != axis).map(|(_, kept)| match *kept {
696 None => Steps::along(1, 0),
697 Some((b, back)) => {
698 let b = b - usize::from(folded.is_some_and(|f| b > f));
699 let (len, step) = (rest[b], stride[b] as isize);
700 if back { Steps::Stride { start: (len - 1) * stride[b], len, step: -step } } else { Steps::Stride { start: 0, len, step } }
701 }
702 }).collect();
703 (around, (!Steps::in_order(&steps)).then_some(steps))
704 }
705}
706
707impl From<Vec<Steps>> for Mapping {
709 #[inline]
710 fn from(tables: Vec<Steps>) -> Self { Self::Steps { base: 0, steps: tables.into() } }
711}
712
713impl From<Writes<'_>> for Mapping {
715 fn from(writes: Writes<'_>) -> Self {
716 match writes {
717 Writes::Offsets(offsets) => Self::from(vec![Steps::table(offsets.to_vec())]),
718 Writes::Walk { base, steps } => Self::Steps { base, steps: steps.to_vec().into() },
719 }
720 }
721}
722
723#[derive(Clone, Debug)]
725pub struct Walked<'a> { mapping: Cow<'a, Mapping>, walk: OnceLock<(usize, Vec<Cow<'static, Steps>>)>, cycle: OnceLock<Option<usize>> }
726
727impl<'a> Walked<'a> {
728 pub fn new(mapping: Cow<'a, Mapping>) -> Self { Self { mapping, walk: OnceLock::new(), cycle: OnceLock::new() } }
730 #[inline]
732 fn walk(&self) -> (usize, &[Cow<'static, Steps>]) {
733 let (offset, walk) = self.walk.get_or_init(|| match &*self.mapping {
734 Mapping::Steps { steps, .. } => {
735 let (offset, walk) = Steps::simplify(0, &steps[..]);
736 (offset, walk.into_iter().map(|s| Cow::Owned(s.into_owned())).collect())
737 }
738 _ => (0, Vec::new()),
739 });
740 (*offset, walk)
741 }
742 pub fn cycle(&self) -> Option<usize> { *self.cycle.get_or_init(|| match self.mapping.period() { Some(Period::Cycles(p)) => Some(p), _ => None }) }
744
745 pub fn gathered(&self, input: crate::buffer::View<'_>, cover: usize) -> Option<crate::buffer::Buffer> {
748 with_dtype!(input.dtype(), A => {
749 let data = A::of(input).expect("the view has its dtype");
750 let mut items = Vec::new();
751 if tiled_cycle(data, self, cover, &mut items).is_none() {
752 let Mapping::Steps { steps, .. } = &**self else { return None; };
753 let n = Steps::count(steps);
754 items.reserve(n);
755 expanded_items(self, data, 0, n, &mut items);
756 }
757 Some(A::buffer(items))
758 })
759 }
760}
761
762impl<'a> From<&'a Mapping> for Walked<'a> { fn from(mapping: &'a Mapping) -> Self { Self::new(Cow::Borrowed(mapping)) } }
763impl From<Mapping> for Walked<'static> { fn from(mapping: Mapping) -> Self { Self::new(Cow::Owned(mapping)) } }
764impl std::ops::Deref for Walked<'_> {
765 type Target = Mapping;
766 #[inline]
767 fn deref(&self) -> &Mapping { &self.mapping }
768}
769impl PartialEq for Walked<'_> { fn eq(&self, other: &Self) -> bool { self.mapping == other.mapping } }
770
771impl Mapping {
772 pub fn tabled(&self) -> bool { matches!(self, Self::Steps { steps, .. } if steps.iter().any(|s| matches!(s, Steps::Table(_)))) }
774 pub fn as_writes(&self) -> Option<Writes<'_>> { if let Self::Steps { base, steps } = self { Some(Writes::Walk { base: *base, steps }) } else { None } }
776 fn affine_axes(&self, shape: &[usize], axis: &mut dyn FnMut(usize, isize)) -> Option<usize> {
778 if shape.contains(&0) { for k in 0..shape.len() { axis(k, 0); } return Some(0); }
779 if let Self::Steps { base, steps } = self {
780 if steps.iter().map(Steps::len).ne(shape.iter().copied()) { return None; }
781 return steps.iter().enumerate().try_fold(*base, |at, (k, s)| {
782 let (start, step) = if s.len() == 1 { (s.get(0)?, 0) } else { s.affine()? };
783 axis(k, step);
784 Some(at.wrapping_add(start))
785 });
786 }
787 if let Self::Linear(period) | Self::Tiled(period) = *self {
788 if period < count(shape) && period != 1 && !shape.iter().rev().scan(1, |p, &n| { *p *= n; Some(*p) }).any(|p| p == period) { return None; }
789 }
790 let (base, mut stride) = (self.get(0)?, count(shape));
791 for (k, &len) in shape.iter().enumerate() {
792 stride /= len;
793 axis(k, if len > 1 { self.get(stride)? as isize - base as isize } else { 0 });
794 }
795 Some(base)
796 }
797
798 pub fn entries(&self, shape: &[usize]) -> Option<(usize, Vec<Steps>)> {
800 if shape.contains(&0) { return Some((0, shape.iter().map(|&n| Steps::along(n, 0)).collect())); }
801 if let Self::Steps { base, steps } = self { let (offset, steps) = Steps::per_axis(steps, shape)?; return Some((base + offset, steps)); }
802 let mut steps = Vec::with_capacity(shape.len());
803 let base = self.affine_axes(shape, &mut |k, step| steps.push(Steps::Stride { start: 0, len: shape[k], step }))?;
804 Some((base, steps))
805 }
806 pub fn through(&self, outer: &[Steps], shape: &[usize]) -> Option<Self> {
808 if outer.iter().any(|s| s.len() == 0) { return Some(Self::Linear(1)); }
809 let (base, inner) = self.entries(shape)?;
810 let (base, steps) = Steps::through(outer, shape, base, &inner)?;
811 Some(Self::Steps { base, steps })
812 }
813 pub fn split(&self, shape: &[usize]) -> Self { self.entries(shape).map_or_else(|| self.clone(), |(base, steps)| Self::Steps { base, steps }) }
815 pub fn to_walk(&self, shape: &[usize]) -> (usize, Cow<'_, [Steps]>) {
816 if let Self::Steps { base, steps } = self { return (*base, Cow::Borrowed(steps)); }
817 let (base, steps) = match self.entries(shape) {
818 Some((base, steps)) => (base, steps), None => (0, vec![Steps::table((0..count(shape)).map(|i| self.get(i).unwrap_or(Steps::FILL)).collect())])
819 };
820 (base, Cow::Owned(steps))
821 }
822 fn shape(&self, items: usize) -> Vec<usize> {
824 match self { Self::Steps { steps, .. } if Steps::count(steps) == items => steps.iter().map(Steps::len).collect(), _ => vec![items] }
825 }
826 pub fn in_order(&self, len: usize, shape: &[usize]) -> bool {
828 len == count(shape) && (matches!(self, Self::Linear(1)) || { let (base, steps) = self.to_walk(shape); Self::walk_in_order(base, &steps, len) })
829 }
830 pub fn walk_in_order(base: usize, steps: &[Steps], len: usize) -> bool { base == 0 && Steps::count(steps) == len && Steps::in_order(steps) }
832 pub fn tables(&self, shape: &[usize]) -> Vec<Steps> { let (base, steps) = self.to_walk(shape); Steps::with_base(base, steps.into_owned()) }
834 pub fn footprint(&self) -> usize { if let Self::Steps { steps, .. } = self { steps.iter().map(Steps::footprint).sum() } else { 0 } }
836 pub fn compose(&self, outer: &Self, shape: &[usize], frame: &[usize]) -> Option<Self> {
838 if frame.contains(&0) { return Some(Self::Linear(1)); }
839 if matches!(self, Self::Linear(1)) { return Some(outer.clone()); }
840 if matches!(outer, Self::Linear(1)) { return Some(self.clone()); }
841 let (base, steps) = match outer { Self::Steps { base, steps } => (*base, steps.to_vec()), _ => outer.entries(frame)? };
842 Some(self.through(&Steps::with_base(base, steps), shape)?)
843 }
844 pub fn composed(&self, outer: &Self, shape: &[usize], frame: &[usize], table: bool) -> Option<Self> {
846 if let Some(mapping) = self.compose(outer, shape, frame) { return Some(mapping); }
847 if !table { return None; }
848 let walked = |map: &Self, shape: &[usize]| {
849 let (base, steps) = map.to_walk(shape);
850 Steps::combined(base, &steps[..])
851 };
852 let inner = (count(shape) <= count(frame)).then(|| walked(self, shape));
853 let at = |j: usize| match &inner {
854 _ if j == Steps::FILL => j, Some(listed) => listed.get(j).copied().unwrap_or(Steps::FILL), None => self.get(j).unwrap_or(Steps::FILL)
855 };
856 Some(Self::Steps { base: 0, steps: vec![Steps::table(walked(outer, frame).into_iter().map(at).collect())].into() })
857 }
858 fn bounds(&self, shape: &[usize]) -> Option<Range<usize>> { let (base, strides) = self.affine(shape)?; affine_bounds(base, &strides, shape) }
859 pub fn affine(&self, shape: &[usize]) -> Option<(usize, Vec<isize>)> {
861 if !shape.contains(&0) && let Self::Steps { steps, .. } = self && steps.iter().any(|s| s.len() > 1 && s.affine().is_none()) { return None; }
863 let (base, steps) = self.entries(shape)?;
864 let (mut first, mut strides) = (base, Vec::with_capacity(steps.len()));
865 for s in &steps {
866 let (start, step) = s.affine()?;
867 first = first.wrapping_add(start);
868 strides.push(step);
869 }
870 Some((first, strides))
871 }
872 pub fn block(&self, shape: &[usize]) -> Option<Block> {
874 if shape.contains(&0) { return None; }
875 let (first, strides) = self.affine(shape)?;
876 let mut walked: Vec<usize> = (0..shape.len()).filter(|&a| shape[a] > 1).collect();
877 walked.sort_by_key(|&a| std::cmp::Reverse(strides[a].unsigned_abs()));
878 let mut size = 1;
879 for &a in walked.iter().rev() {
880 if strides[a].unsigned_abs() != size { return None; }
881 size *= shape[a];
882 }
883 let start = walked.iter().filter(|&&a| strides[a] < 0).try_fold(first, |s, &a| s.checked_sub(strides[a].unsigned_abs() * (shape[a] - 1)))?;
884 let mut axes = vec![None; shape.len()];
885 for (b, &a) in walked.iter().enumerate() { axes[a] = Some((b, strides[a] < 0)); }
886 Some(Block { range: start..start + size, dims: walked.iter().map(|&a| shape[a]).collect(), axes })
887 }
888 #[inline]
890 pub fn index(&self, i: usize) -> usize { self.get(i).expect("positional mapping") }
891 #[inline]
893 pub fn run(&self, start: usize, n: usize) -> Option<(usize, isize)> {
894 if n == 0 { return None; }
895 match *self {
896 Self::Single => Some((0, 0)), Self::Linear(0) | Self::Tiled(0) => None,
897 Self::Linear(1) => Some((start, 1)),
898 Self::Linear(r) if n <= r - start % r => Some((start / r, 0)),
899 Self::Tiled(p) if n <= p - start % p => Some((start % p, 1)),
900 Self::Steps { base, ref steps } => Steps::run(base, steps, start, n),
901 _ => None
902 }
903 }
904 #[inline]
907 pub fn get(&self, i: usize) -> Option<usize> {
908 match self {
909 Self::Linear(0) | Self::Tiled(0) => None, Self::Single => Some(0), Self::Linear(1) => Some(i), Self::Linear(repeat) => Some(i / repeat),
910 Self::Tiled(size) => Some(i % size), Self::Steps { base, steps } => Self::reached(*base, steps, i)
911 }
912 }
913 #[inline(never)]
914 fn reached(base: usize, steps: &[Steps], i: usize) -> Option<usize> { Steps::reach(base, steps, i) }
915 pub fn contract(left: &Layout, laxis: usize, right: &Layout, raxis: usize) -> Result<Self, Error> {
917 match (left.keys(laxis), right.keys(raxis)) {
918 (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))
919 }
920 }
921 pub fn new(input: &Layout, axes: &[usize], result: &Layout, count: usize) -> Self {
923 let alignments = axes.iter().enumerate().map(|(a, &b)| match (input.keys(a), result.keys(b)) {
924 (Some(src), Some(dst)) => src.aligned(dst), _ => Alignment::Same
925 }).collect();
926 Self::aligned(input, axes, result, count, alignments)
927 }
928 fn aligned(input: &Layout, axes: &[usize], result: &Layout, count: usize, alignments: Vec<Alignment>) -> Self {
930 if count == 0 { return Self::Single; }
931 let (shape, output) = (input.shape(), result.shape());
932 if alignments.iter().all(|a| matches!(a, Alignment::Same)) {
933 let size = input.count();
934 if size == 1 { return Self::Single; }
935 let strides = strides(output);
936 let (mut source, mut linear, mut tiled, repeat) = (1, true, true, count / size);
937 for (&len, &axis) in shape.iter().zip(axes).rev() {
938 if len != 1 {
939 linear &= strides[axis] == source * repeat;
940 tiled &= strides[axis] == source;
941 }
942 source *= len;
943 }
944 if linear { return Self::Linear(repeat); }
945 if tiled { return Self::Tiled(size); }
946 }
947 let mut steps: Vec<Steps> = output.iter().map(|&len| Steps::Stride { start: 0, len, step: 0 }).collect();
948 let mut source = 1;
949 for ((&len, &b), alignment) in shape.iter().zip(axes).zip(alignments).rev() {
950 steps[b] = match alignment {
951 Alignment::Same => Steps::Stride { start: 0, len: output[b], step: if len == 1 { 0 } else { source as isize } },
952 Alignment::Leading => Steps::clipped(0, output[b], len, source),
953 Alignment::Table(t) => {
955 let table = Steps::table(t.into_iter().map(|p| p.map_or(Steps::FILL, |p| p * source)).collect());
956 table.even().unwrap_or(table)
957 }
958 };
959 source *= len;
960 }
961 Self::Steps { base: 0, steps: steps.into() }
962 }
963 pub fn line_block(&self, size: usize) -> usize {
965 match self.inner() { Some(inner) if inner.skips() => (inner.len() * (LINE / size).max(1)).min(LINE_BLOCK), _ => 0 }
966 }
967 pub fn period(&self) -> Option<Period<'_>> {
969 if let Self::Tiled(m) = *self { return Some(Period::Cycles(m)); }
970 if let Some(r) = self.repeats() { return Some(Period::Repeats(r)); }
971 let Self::Steps { base, steps } = self else { return None };
972 if let [Steps::Padded { start: 0, len, size, stride: 1, edge: crate::structure::Edge::Wrap }] = steps[..] && *base == 0 && size >= 2 && len > size {
973 return Some(Period::Cycles(size));
974 }
975 let k = steps.iter().position(|s| !matches!(*s, Steps::Stride { step: 0, .. }))?;
976 let (outer, inner) = steps.split_at(k);
977 let p = inner.iter().map(Steps::len).product::<usize>();
978 if p < 2 || outer.iter().all(|s| s.len() < 2) || self.fills() { return None; }
979 let base = base + outer.iter().map(|s| if let Steps::Stride { start, .. } = *s { start } else { 0 }).sum::<usize>();
980 Some(if base == 0 && Steps::in_order(inner) { Period::Cycles(p) } else { Period::Walks { len: p, base, steps: inner } })
981 }
982 pub fn copies(&self, block: usize) -> bool { match self { Self::Steps { steps, .. } => Steps::contiguous(steps) < block, _ => false } }
984 pub fn tiles_whole(&self, block: usize) -> bool {
986 let Self::Steps { steps, .. } = self else { return false };
987 let (_, steps) = Steps::simplify(0, steps);
988 let [.., outer, inner] = &steps[..] else { return false };
989 Steps::transposing(outer, inner).is_some_and(|(_, _, n, _)| n >= block)
990 }
991 pub fn regular(&self) -> Option<(usize, isize)> {
993 match self { Self::Linear(r) => Some((*r, 0)), Self::Steps { steps, .. } => Some(Steps::regular(steps)), _ => None }
994 }
995 pub fn repeats(&self) -> Option<usize> { self.regular().and_then(|(len, step)| (step == 0 && len > 1).then_some(len)) }
997 pub fn runs(&self, positions: usize) -> Option<(usize, isize, Vec<usize>)> {
999 let (base, steps) = match self {
1000 Self::Linear(r) => return Some((*r, 0, (0..positions.div_ceil(*r)).collect())), Self::Steps { base, steps } => (base, steps), _ => return None
1001 };
1002 if self.fills() { return None; }
1003 let (len, stride) = Steps::regular(steps);
1004 let (mut width, mut outer) = (1, steps.to_vec());
1006 for axis in outer.iter_mut().rev() {
1007 if width == len { break; }
1008 if axis.len() == 1 { continue; }
1009 let Steps::Stride { start, len: n, step } = *axis else { return None };
1010 *axis = Steps::Stride { start, len: 1, step };
1011 width *= n;
1012 }
1013 Some((len, stride, Steps::combined(*base, &outer)))
1014 }
1015 pub fn fills(&self) -> bool { matches!(self, Self::Steps { steps, .. } if steps.iter().any(Steps::fills)) }
1017 pub fn spans(&self, range: Range<usize>, items: usize, f: &mut dyn FnMut(usize, usize, Option<(usize, isize)>) -> bool) {
1019 let mut at = range.start;
1020 match *self {
1021 Self::Single => { f(at, range.len(), (items > 0).then_some((0, 0))); }
1022 Self::Linear(0) | Self::Tiled(0) => { f(at, range.len(), None); }
1023 Self::Linear(1) => {
1024 let below = items.saturating_sub(at).min(range.len());
1025 if below > 0 && !f(at, below, Some((at, 1))) { return; }
1026 if below < range.len() { f(at + below, range.len() - below, None); }
1027 }
1028 Self::Linear(r) => {
1029 while at < range.end {
1030 let (k, end) = (at / r, range.end.min((at / r + 1) * r));
1031 if !f(at, end - at, (k < items).then_some((k, 0))) { return; }
1032 at = end;
1033 }
1034 }
1035 Self::Tiled(m) => {
1036 while at < range.end {
1037 let (p, end) = (at % m, range.end.min(at - at % m + m));
1038 let below = items.saturating_sub(p).min(end - at);
1039 if below > 0 && !f(at, below, Some((p, 1))) { return; }
1040 if end > at + below && !f(at + below, end - at - below, None) { return; }
1041 at = end;
1042 }
1043 }
1044 Self::Steps { base, ref steps } => {
1045 let Some((inner, outer)) = steps.split_last() else { f(at, range.len(), (base < items).then_some((base, 0))); return; };
1046 if let Some((first, step)) = Steps::run(base, steps, at, range.len())
1048 && first.max(first.wrapping_add_signed(step * (range.len() as isize - 1))) < items { f(at, range.len(), Some((first, step))); return; }
1049 let width = inner.len();
1050 while at < range.end && width > 0 {
1051 let (row, j) = (at / width, at % width);
1052 let n = (width - j).min(range.end - at);
1053 let go = match Steps::reach(base, outer, row) { Some(start) => row_spans(start, inner, j, n, items, at, f), None => f(at, n, None) };
1054 if !go { return; }
1055 at += n;
1056 }
1057 }
1058 }
1059 }
1060
1061 fn inner(&self) -> Option<&Steps> { let Self::Steps { steps, .. } = self else { return None }; steps.iter().rev().find(|s| s.len() > 1) }
1063}
1064
1065pub const LISTED: usize = 128;
1067pub fn small_read(read: usize, len: usize) -> bool { read.saturating_mul(LISTED) <= len }
1069
1070pub fn paired_length(x: usize, y: usize) -> Result<usize, Error> {
1072 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"))) }
1073}
1074
1075#[derive(Clone)]
1077pub struct Agreement<L = Layout> { pub layout: L, pub len: usize, pub left: Mapping, pub right: Mapping }
1078
1079impl Agreement {
1080 pub fn new(left: &Layout, right: &Layout) -> Result<Self, Error> {
1082 if left.plain() && right.plain() && (left.shape() == right.shape() || left.shape().is_empty() || right.shape().is_empty()) {
1084 let shape = if left.shape().is_empty() { right.shape() } else { left.shape() };
1085 let len = generated_len(shape)?;
1086 let map = |l: &Layout| if len > 1 && l.shape() == shape { Mapping::Linear(1) } else { Mapping::Single };
1087 return Ok(Self { layout: Layout::from(shape), len, left: map(left), right: map(right) });
1088 }
1089 let [xa, ya] = Self::axis_pairs(left, right);
1090 Self::mapped(left, right, &xa, &ya, false)
1091 }
1092 pub fn axis_pairs(left: &Layout, right: &Layout) -> [Vec<usize>; 2] {
1093 let (mut ya, nx) = (Self::paired_axes(left, right, true), left.shape().len());
1094 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<_>>>() {
1095 ya.iter_mut().enumerate().for_each(|(b, a)| *a = b);
1096 return [xa, ya];
1097 }
1098 [(0..nx).collect(), ya]
1099 }
1100 pub fn onto(left: &Layout, right: &Layout) -> Result<Self, Error> {
1102 if left.plain() && right.plain() { return Self::new(left, right); }
1103 Self::mapped(left, right, &(0..left.shape().len()).collect::<Vec<_>>(), &Self::paired_axes(left, right, true), true)
1104 }
1105 pub fn paired_axes(left: &Layout, right: &Layout, positional: bool) -> Vec<usize> {
1107 let (nx, ny) = (left.shape().len(), right.shape().len());
1108 let mut ya: Vec<_> = (0..ny).map(|b| right.name(b).and_then(|n| (0..nx).find(|&a| left.name(a) == Some(n)))).collect();
1109 if positional {
1110 let remaining_left: Vec<_> = (0..nx).filter(|a| !ya.contains(&Some(*a))).collect();
1111 let remaining_right: Vec<_> = (0..ny).filter(|&b| ya[b].is_none()).collect();
1112 for (a, b) in remaining_left.into_iter().zip(remaining_right) {
1113 if !names_conflict(left.name(a).map(|n| &**n), right.name(b).map(|n| &**n)) { ya[b] = Some(a); }
1114 }
1115 }
1116 let mut next = nx;
1117 ya.into_iter().map(|a| { a.unwrap_or_else(|| { next += 1; next - 1 }) }).collect()
1118 }
1119 pub fn paired(left: &Layout, right: &Layout, axes: &[usize]) -> Result<Self, Error> {
1121 Self::mapped(left, right, &(0..left.shape().len()).collect::<Vec<_>>(), axes, false)
1122 }
1123 pub fn with_axes(left: &Layout, right: &Layout, axes: &[usize]) -> Result<Self, Error> {
1125 if left.shape().len() < right.shape().len() { Self::mapped(left, right, axes, &(0..right.shape().len()).collect::<Vec<_>>(), false) }
1126 else { Self::mapped(left, right, &(0..left.shape().len()).collect::<Vec<_>>(), axes, false) }
1127 }
1128 pub fn outer(left: &Layout, right: &Layout) -> Result<Self, Error> {
1130 let nx = left.shape().len();
1131 Self::mapped(left, right, &(0..nx).collect::<Vec<_>>(), &(nx..nx + right.shape().len()).collect::<Vec<_>>(), false)
1132 }
1133 pub fn mapped(left: &Layout, right: &Layout, xa: &[usize], ya: &[usize], left_keys: bool) -> Result<Self, Error> {
1134 if xa.len() != left.shape().len() || ya.len() != right.shape().len() { return Err(Error::new(Kind::Length, "axis maps must match input ranks")); }
1135 let mut shape = vec![1; xa.iter().chain(ya).max().map_or(0, |a| a + 1)];
1136 let mut keys = vec![None; shape.len()];
1137 let mut names = vec![None; shape.len()];
1138 let (mut lt, mut rt) = (vec![Alignment::Same; left.shape().len()], vec![Alignment::Same; right.shape().len()]);
1139 for axis in 0..shape.len() {
1140 let (x, y) = (xa.iter().position(|&a| a == axis), ya.iter().position(|&a| a == axis));
1141 let (nx, ny) = (x.map_or(1, |a| left.shape()[a]), y.map_or(1, |a| right.shape()[a]));
1142 let (kx, ky) = (x.and_then(|a| left.keys(a)), y.and_then(|a| right.keys(a)));
1143 let (lx, ly) = (x.and_then(|a| left.name(a)), y.and_then(|a| right.name(a)));
1144 names[axis] = match (lx, ly) { (Some(l), Some(r)) if hidden(l) && !hidden(r) => Some(r), _ => lx.or(ly) }.cloned();
1145 if let (Some(kx), Some(ky), Some(a), Some(b)) = (kx, ky, x, y) {
1146 let key = if left_keys { rt[b] = ky.aligned(kx); kx.clone() } else {
1147 let (key, [l, r]) = kx.union(ky)?;
1148 (lt[a], rt[b]) = (l, r);
1149 key
1150 };
1151 shape[axis] = key.len();
1152 keys[axis] = Some(key);
1153 } else {
1154 let n = paired_length(nx, ny)?;
1155 shape[axis] = n;
1156 keys[axis] = kx.filter(|_| nx == n).or(ky.filter(|_| ny == n)).cloned();
1157 }
1158 }
1159 let len = generated_len(&shape)?;
1160 let layout = Layout::from(shape).with_keys(keys)?.inherit_names(names);
1161 let (x, y) = (Mapping::aligned(left, xa, &layout, len, lt), Mapping::aligned(right, ya, &layout, len, rt));
1162 Ok(Self { layout, len, left: x, right: y })
1163 }
1164}
1165
1166impl<L> Agreement<L> {
1167 pub fn map_layout<M>(self, f: impl FnOnce(L) -> M) -> Agreement<M> {
1169 Agreement { layout: f(self.layout), len: self.len, left: self.left, right: self.right }
1170 }
1171}
1172
1173pub struct Frames { pub layout: Layout, pub count: usize, cells: Vec<Option<Vec<usize>>> }
1175
1176impl Frames {
1177 pub fn new(args: &[(&Layout, usize)]) -> Result<Self, Error> {
1179 let frames = args.iter().map(|&(l, core)| {
1180 let rank = l.shape().len();
1181 if rank < core { return Err(Error::new(Kind::Rank, format!("an argument needs at least {core} axes, not {rank}"))); }
1182 Ok(l.axes(0..rank - core))
1183 }).collect::<Result<Vec<_>, _>>()?;
1184 let Some((first, rest)) = frames.split_first() else { return Ok(Self { layout: Layout::EMPTY, count: 1, cells: Vec::new() }) };
1185 let (mut layout, mut chains) = (first.clone(), vec![Vec::<Mapping>::new()]);
1186 for frame in rest {
1187 let a = Agreement::new(&layout, frame)?;
1188 for chain in &mut chains { chain.push(a.left.clone()); }
1189 chains.push(vec![a.right]);
1190 layout = a.layout;
1191 }
1192 let count = generated_len(layout.shape())?;
1193 let cells = chains.iter().map(|chain| {
1195 if chain.iter().all(|m| *m == Mapping::Linear(1)) { return Some(None); }
1196 (0..count).map(|i| chain.iter().rev().try_fold(i, |j, m| m.get(j))).collect::<Option<Vec<_>>>().map(Some)
1197 }).collect::<Option<_>>().ok_or_else(|| Error::new(Kind::Length, "keyed frames need a cell that an argument lacks"))?;
1198 Ok(Self { layout, count, cells })
1199 }
1200
1201 #[inline]
1203 pub fn one() -> Self { Self { layout: Layout::EMPTY, count: 1, cells: Vec::new() } }
1204
1205 #[inline]
1207 pub fn cell(&self, k: usize, i: usize) -> usize { match self.cells.get(k) { Some(Some(c)) => c[i], _ => i } }
1208
1209 #[inline]
1211 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 }
1212}
1213
1214const LONG_CYCLE: usize = 256;
1216pub(crate) const LINE: usize = 128;
1218const LINE_BLOCK: usize = 1 << 16;
1220
1221fn single<A: Copy>(mapping: &Mapping, data: &[A], start: usize, n: usize) -> Option<A> {
1223 mapping.run(start, n).and_then(|(o, step)| (step == 0).then(|| data[o]))
1224}
1225
1226pub fn expanded_items<'a, A: Element>(mapping: &Walked, data: &'a [A], start: usize, n: usize, buf: &'a mut Vec<A>) -> &'a [A] {
1228 if let Some(x) = single(mapping, data, start, n) {
1229 buf.clear();
1230 buf.resize(n, x);
1231 return buf;
1232 }
1233 match expanded(mapping, data, start, n, buf) { Arg::Many(items) => items, Arg::One(_) => unreachable!("`single` gives every repeated item") }
1234}
1235
1236pub fn in_place<'a, A: Copy>(mapping: &Walked, data: &'a [A], start: usize, n: usize) -> Option<Arg<'a, A>> {
1238 match mapping.run(start, n) {
1239 Some((o, 0)) => return data.get(o).copied().map(Arg::One),
1240 Some((o, 1)) => return data.get(o..o.checked_add(n)?).map(Arg::Many),
1241 _ => {}
1242 }
1243 let period = match **mapping {
1245 Mapping::Tiled(p) if p > 0 => p,
1246 Mapping::Steps { .. } => mapping.cycle().filter(|&p| p > 0 && n <= p - start % p)?,
1247 _ => return None
1248 };
1249 let from = start % period;
1250 data.get(from..from.checked_add(n)?).map(Arg::Many)
1251}
1252
1253pub fn expanded<'a, A: Element>(mapping: &Walked, data: &'a [A], start: usize, n: usize, buf: &'a mut Vec<A>) -> Arg<'a, A> {
1255 if let Some(arg) = in_place(mapping, data, start, n) { return arg; }
1256 Arg::Many(copied(mapping, data, start, n, buf))
1257}
1258
1259pub fn tiled_cycle<A: Element>(data: &[A], mapping: &Mapping, run: usize, buf: &mut Vec<A>) -> Option<Mapping> {
1261 let period = mapping.period()?;
1262 let (p, n) = period.tiling(run)?;
1263 match period {
1264 Period::Walks { base, steps, .. } => { copied(&Walked::from(Mapping::Steps { base, steps: steps.to_vec().into() }), data, 0, p, buf); }
1265 _ => {
1266 buf.clear();
1267 buf.extend_from_slice(&data[..p]);
1268 }
1269 }
1270 buf.truncate(p);
1271 if p < n { *buf = buf.repeat(n / p); }
1272 Some(Mapping::Tiled(p))
1273}
1274
1275pub fn copied<'a, A: Element>(mapping: &Walked, data: &[A], start: usize, n: usize, buf: &'a mut Vec<A>) -> &'a [A] {
1277 buf.clear();
1278 let reason = copy_walked(mapping, data, start, n, A::FILL, buf);
1279 crate::events::declined(crate::events::Route::InPlace, reason, n);
1280 &buf[..n]
1281}
1282
1283#[inline(never)]
1285fn copy_walked<A: Element>(mapping: &Walked, data: &[A], start: usize, n: usize, fill: A, buf: &mut Vec<A>) -> &'static str {
1286 let at = buf.len();
1287 match **mapping {
1288 _ if n == 0 => "no items",
1289 Mapping::Single => { buf.resize(at + n, data[0]); "a repeated item" }
1290 Mapping::Linear(r) => {
1291 let (k0, end) = (start / r, start + n);
1293 buf.resize(at + ((k0 + 1) * r).min(end) - start, data[k0]);
1294 if buf.len() - at < n {
1295 let (k1, k2) = (k0 + 1, end / r);
1296 crate::gather::repeat(buf, k2 - k1, &[r as i64], (k2 - k1) * r, |j| data[k1 + j]);
1297 if end % r != 0 { buf.resize(at + n, data[k2]) }
1298 }
1299 "repeats of each item"
1300 }
1301 Mapping::Tiled(m) => { cycled(buf, &data[..m], start, n); "a cycle that wraps" }
1302 Mapping::Steps { base, .. } => {
1303 let (offset, plan) = mapping.walk();
1304 walk_plan(buf, data, base + offset, plan, fill, start, n, Same);
1305 "steps other than a forward run"
1306 }
1307 }
1308}