1use crate::hashing::{HashMap, HashMapExt};
5use crate::{
6 Float, Native, agreement::Mapping, array::ArrayView, buffer::{Buffer, DType, Element, Floats, IntBuf, View, Width}, error::{Error, Kind}, interrupt::check,
7 layout::{Layout, count}, tolerant::{self, Tolerance}
8};
9use hashbrown::hash_map::Entry;
10use std::cmp::Ordering;
11use std::{convert::Infallible, hash::Hash, mem::MaybeUninit, ops::Not};
12
13pub trait Key: Copy + Send + Sync + Eq + Ord + Hash {
15 fn key(self) -> i64;
16 fn from_key(key: i64) -> Option<Self>;
18 fn range(x: &[Self]) -> (i64, i64);
20}
21#[inline]
23fn extreme_keys<K: Key>(x: &[K]) -> (i64, i64) where crate::reduce::Min: crate::reduce::Reduction<K, Total = K>,
24 crate::reduce::Max: crate::reduce::Reduction<K, Total = K> { let (lo, hi) = crate::reduce::extremes(x); (lo.key(), hi.key()) }
25macro_rules! keys {
26 ($($t:ty),+) => {$(
27 impl Key for $t {
28 #[inline]
29 fn key(self) -> i64 { self as i64 }
30 #[inline]
31 fn from_key(key: i64) -> Option<Self> { ((key as $t).key() == key).then_some(key as $t) }
32 #[inline]
33 fn range(x: &[Self]) -> (i64, i64) { extreme_keys(x) }
34 }
35 )+};
36}
37keys!(u8, u16, u32, u64, i16, i32, i64);
38impl Key for bool {
39 #[inline]
40 fn key(self) -> i64 { self.into() }
41 #[inline]
42 fn from_key(key: i64) -> Option<Self> { matches!(key, 0 | 1).then_some(key == 1) }
43 #[inline]
44 fn range(x: &[Self]) -> (i64, i64) { extreme_keys(x) }
45}
46impl Key for char {
47 #[inline]
48 fn key(self) -> i64 { i64::from(u32::from(self)) }
49 #[inline]
50 fn from_key(key: i64) -> Option<Self> { char::from_u32(u32::try_from(key).ok()?) }
51 #[inline]
52 fn range(x: &[Self]) -> (i64, i64) { let (lo, hi) = crate::reduce::extremes_of(x, |c: char| c as u32); (lo.into(), hi.into()) }
53}
54
55pub trait RadixKey: Copy + Send + Sync + Ord + Default + Not<Output = Self> {
57 const BYTES: usize;
58 fn byte(self, b: usize) -> usize;
59 fn word(self) -> u64;
61 fn extremes(x: &[Self]) -> (Self, Self);
63}
64macro_rules! radix_keys {
65 ($($t:ty),+) => {$(
66 impl RadixKey for $t {
67 const BYTES: usize = size_of::<$t>();
68 #[inline]
69 fn byte(self, b: usize) -> usize { (self >> (8 * b)) as usize & 255 }
70 #[inline]
71 fn word(self) -> u64 { self as u64 }
72 #[inline]
73 fn extremes(x: &[Self]) -> (Self, Self) { crate::reduce::extremes(x) }
74 }
75 )+};
76}
77radix_keys!(u8, u16, u32, u64);
78
79pub(crate) trait SortKey: Element + Default {
81 type Key: RadixKey;
82 fn sort_key(self) -> Self::Key;
83 fn from_sort_key(k: Self::Key) -> Self;
84 fn sorted_into(_x: &[Self], _out: &mut [Self], _down: bool, _scratch: &mut Vec<Self>) -> bool { false }
86 fn keys_into(x: &[Self], keys: &mut Vec<Self::Key>, down: bool) {
88 keys.clear();
89 keys.extend(x.iter().map(|&v| if down { !v.sort_key() } else { v.sort_key() }));
90 }
91 fn write_back(keys: &[Self::Key], out: &mut [Self], down: bool) {
93 for (o, &k) in out.iter_mut().zip(keys) { *o = Self::from_sort_key(if down { !k } else { k }); }
94 }
95}
96impl<F: Float + Element + Default + Network> SortKey for F {
97 type Key = F::OrderKey;
98 #[inline]
99 fn sort_key(self) -> F::OrderKey { self.total_key() }
100 #[inline]
101 fn from_sort_key(k: F::OrderKey) -> Self { F::from_order_key(k) }
102 fn sorted_into(x: &[Self], out: &mut [Self], down: bool, scratch: &mut Vec<Self>) -> bool { F::networked(x, out, down, scratch) }
103}
104macro_rules! sort_keys {
105 ($($t:ty => $key:ty, $to:expr, $back:expr);+) => {$(
106 impl SortKey for $t {
107 type Key = $key;
108 #[inline]
109 fn sort_key(self) -> $key { ($to)(self) }
110 #[inline]
111 fn from_sort_key(k: $key) -> Self { ($back)(k) }
112 fn sorted_into(x: &[Self], out: &mut [Self], down: bool, _: &mut Vec<Self>) -> bool { spread_into(x, out, down) }
113 }
114 )+};
115}
116sort_keys!(i16 => u16, |x: i16| x as u16 ^ 1 << 15, |k: u16| (k ^ 1 << 15) as i16; i32 => u32, |x: i32| x as u32 ^ 1 << 31, |k: u32| (k ^ 1 << 31) as i32;
117 i64 => u64, |x: i64| x as u64 ^ 1 << 63, |k: u64| (k ^ 1 << 63) as i64; u8 => u8, |x| x, |k| k; bool => u8, u8::from, |k| k != 0; char => u32, u32::from, |k| char::from_u32(k).unwrap_or_default());
118
119pub const SCANNED: usize = 16;
121
122const FIRST_BLOCK: usize = 64;
124
125#[inline(always)]
127pub(crate) fn first<T: Copy>(x: &[T], hit: impl Fn(T) -> bool) -> Option<usize> {
128 let (blocks, rest) = x.as_chunks::<FIRST_BLOCK>();
129 for (b, block) in blocks.iter().enumerate() {
130 if block.iter().fold(false, |any, &h| any | hit(h)) { return block.iter().position(|&h| hit(h)).map(|p| b * FIRST_BLOCK + p); }
131 }
132 rest.iter().position(|&h| hit(h)).map(|p| blocks.len() * FIRST_BLOCK + p)
133}
134
135#[inline(never)]
137pub fn index_of<K: Key>(x: &[K], y: &[K]) -> Vec<u32> {
138 let miss = x.len() as u32;
139 if y.len() <= SCANNED { return y.iter().map(|&n| first(x, |h| h == n).map_or(miss, |p| p as u32)).collect(); }
140 let Some((min, mut table)) = table(x, u32::MAX) else { return first_matches(x, y) };
141 let first = write_first(x, min, &mut table, None, |_, _| ());
142 let at = move |n: K| usize::try_from(n.key().wrapping_sub(min)).ok().and_then(|i| first.get(i));
144 let mut found = Vec::with_capacity(y.len());
145 look_up(&mut found, y, move |n| at(n).map_or(miss, |&f| f.min(miss)));
146 found
147}
148
149pub(crate) fn look_up<K: Copy + Sync>(found: &mut Vec<u32>, y: &[K], look: impl Fn(K) -> u32 + Sync) {
151 crate::parallel::extended(found, y.len(), crate::BETWEEN_CHECKS, &|r, slots| unsafe {
152 slots.fill(r.len(), |slots| {
153 let needles = &y[r];
154 let slots = &mut slots[..needles.len()];
155 for i in 0..needles.len() { slots[i].write(look(needles[i])); }
156 true
157 })
158 });
159}
160
161#[inline(never)]
163pub fn classes<K: Key>(x: &[K]) -> Vec<u32> {
164 let Some((min, mut table)) = table(x, u32::MAX) else { return first_classes(x, |n| n) };
165 let first = write_first(x, min, &mut table, None, |_, _| ());
166 x.iter().map(move |&n| first[n.key().wrapping_sub(min) as usize]).collect()
167}
168
169#[inline(never)]
171pub fn firsts<K: Key>(x: &[K]) -> Vec<bool> {
172 let Some((min, mut table)) = table(x, u32::MAX) else { return first_mask(&first_classes(x, |n| n)) };
173 let mut mask = vec![false; x.len()];
174 write_first(x, min, &mut table, None, |p, new| mask[p] |= new);
175 mask
176}
177
178#[inline]
180fn count_eq<K: Key>(x: &[K], v: K) -> usize { x.chunks(192).map(|c| usize::from(c.iter().map(|&k| u8::from(k == v)).sum::<u8>())).sum() }
181
182pub fn counts<K: Key>(x: &[K]) -> Option<(i64, Vec<usize>)> {
184 let (min, mut counts) = table(x, 0)?;
185 counts_into(x, min, &mut counts);
186 Some((min, counts))
187}
188
189#[inline(never)]
191pub fn counts_into<K: Key>(x: &[K], min: i64, counts: &mut [usize]) {
192 let n = counts.len();
193 let at = move |k: &K| k.key().wrapping_sub(min) as usize;
194 if n <= 4 {
195 for (g, count) in counts.iter_mut().enumerate().skip(1) { *count = K::from_key(min + g as i64).map_or(0, |v| count_eq(x, v)); }
196 counts[0] = x.len() - counts[1..].iter().sum::<usize>();
197 return;
198 }
199 if n >= 64 { for k in x { counts[at(k)] += 1 } return; }
200 let mut parts = vec![0; 4 * n];
201 let (chunks, rest) = x.as_chunks::<4>();
202 for chunk in chunks { for (p, k) in chunk.iter().enumerate() { parts[p * n + at(k)] += 1 } }
203 for k in rest { parts[at(k)] += 1 }
204 for (g, count) in counts.iter_mut().enumerate() { *count = (0..4).map(|p| parts[p * n + g]).sum() }
205}
206
207#[derive(Clone, Debug)]
209pub struct Numbered {
210 pub firsts: Vec<usize>,
212 pub numbers: IntBuf,
214 pub counts: Vec<usize>,
216}
217
218impl Numbered {
219 pub fn grouped(self) -> Grouped {
221 let Numbered { firsts, numbers, .. } = self;
222 with_ints!(numbers.view(), |numbers| grouped_classes(numbers, firsts, class))
223 }
224}
225
226#[derive(Clone, Copy, Debug)]
228pub struct ClassNumbers<'a, N> {
229 pub count: usize,
231 pub numbers: &'a [N],
233}
234
235impl<N: Key> ClassNumbers<'_, N> {
236 pub fn counts(self) -> Vec<usize> {
238 let mut counts = vec![0; self.count];
239 counts_into(self.numbers, 0, &mut counts);
240 counts
241 }
242
243 #[inline]
245 fn folded<T: Copy, R: crate::reduce::Reduction<T>, const PARTS: usize>(self, data: &[T]) -> (Vec<R::Total>, bool) {
246 let (mut totals, mut fits) = (vec![R::identity(); self.count * PARTS], true);
247 for (k, (&c, &x)) in self.numbers.iter().zip(data).enumerate() {
248 let i = class(c) * PARTS + k % PARTS;
249 let ok;
250 (totals[i], ok) = R::checked(totals[i], x);
251 fits &= ok;
252 }
253 if PARTS == 1 { return (totals, fits); }
254 let merged = totals.chunks_exact(PARTS).map(|t| {
255 t[1..].iter().fold(t[0], |a, &b| {
256 let (m, ok) = R::merge(a, b);
257 fits &= ok;
258 m
259 })
260 }).collect();
261 (merged, fits)
262 }
263}
264
265#[inline]
267fn class<N: Key>(k: N) -> usize { k.key() as usize }
268
269pub trait Grouping {
271 fn len(&self) -> usize;
273 fn fold<T: Copy, R: crate::reduce::Reduction<T>>(&self, data: &[T]) -> (Vec<R::Total>, bool);
275 fn scan<T: Copy, R: crate::reduce::Reduction<T>>(&self, data: &[T]) -> (Vec<R::Total>, bool);
277}
278
279impl<N: Key> Grouping for ClassNumbers<'_, N> {
280 fn len(&self) -> usize { self.count }
281 #[inline]
283 fn fold<T: Copy, R: crate::reduce::Reduction<T>>(&self, data: &[T]) -> (Vec<R::Total>, bool) {
284 if self.count <= R::SPLIT { self.folded::<T, R, 4>(data) } else { self.folded::<T, R, 1>(data) }
285 }
286 fn scan<T: Copy, R: crate::reduce::Reduction<T>>(&self, data: &[T]) -> (Vec<R::Total>, bool) {
287 let mut at = crate::layout::exclusive_sums(self.counts());
288 let (mut running, mut totals, mut fits) = (vec![R::identity(); self.count], vec![R::identity(); data.len()], true);
289 for (&k, &x) in self.numbers.iter().zip(data) {
290 let k = class(k);
291 let (t, ok) = R::checked(running[k], x);
292 fits &= ok;
293 (running[k], totals[at[k]]) = (t, t);
294 at[k] += 1;
295 }
296 (totals, fits)
297 }
298}
299
300#[derive(Clone, Copy, Debug)]
302pub struct Ranges<'a> {
303 pub starts: &'a [usize],
305 pub lens: &'a [usize],
307}
308
309impl Grouping for Ranges<'_> {
310 fn len(&self) -> usize { self.starts.len() }
311 fn fold<T: Copy, R: crate::reduce::Reduction<T>>(&self, data: &[T]) -> (Vec<R::Total>, bool) {
312 let mut fits = true;
313 let totals = self.starts.iter().zip(self.lens).map(|(&s, &n)| {
314 let (t, ok) = R::slice(&data[s..s + n]);
315 fits &= ok;
316 t
317 }).collect();
318 (totals, fits)
319 }
320 fn scan<T: Copy, R: crate::reduce::Reduction<T>>(&self, data: &[T]) -> (Vec<R::Total>, bool) {
321 let mut totals = Vec::with_capacity(self.lens.iter().sum());
322 let fits = self.starts.iter().zip(self.lens).fold(true, |fits, (&s, &n)| R::scan(&data[s..s + n], None, &mut totals) & fits);
323 (totals, fits)
324 }
325}
326
327pub fn numbered(classes: &[u32]) -> Numbered {
329 let (mut number, mut firsts) = (vec![0u32; classes.len()], Vec::new());
330 for (i, &c) in classes.iter().enumerate() {
331 if c as usize == i {
332 number[i] = firsts.len() as u32;
333 firsts.push(i);
334 }
335 }
336 let mut counts = vec![0; firsts.len()];
337 let numbers = with_width!(Width::below(firsts.len()), T => int_buf(classes.iter().map(|&c| {
338 let k = number[c as usize];
339 counts[k as usize] += 1;
340 num_traits::AsPrimitive::<T>::as_(k)
341 }).collect::<Vec<T>>()));
342 Numbered { firsts, numbers, counts }
343}
344
345fn int_buf<T: Element>(items: Vec<T>) -> IntBuf {
347 match T::buffer(items) { Buffer::Integers(b) => b, _ => unreachable!("integer items make an integer buffer") }
348}
349
350pub fn dense_offsets(offsets: &[usize], size: usize) -> bool { size <= 8 * offsets.len() }
352
353pub fn occurrences(offsets: &[usize], size: usize) -> Option<(Vec<u32>, usize)> {
355 let next = |seen: &mut u32| { *seen += 1; *seen - 1 };
356 let counts: Vec<u32> = if dense_offsets(offsets, size) {
357 if offsets.len() <= size {
358 let (mut taken, mut repeated) = (vec![0u64; size.div_ceil(64)], 0);
359 for &o in offsets {
360 let (word, bit) = (&mut taken[o / 64], 1 << (o % 64));
361 repeated |= *word & bit;
362 *word |= bit;
363 }
364 if repeated == 0 { return None; }
365 }
366 let mut seen = vec![0u32; size];
367 offsets.iter().map(|&o| next(&mut seen[o])).collect()
368 } else { let mut seen = HashMap::new(); offsets.iter().map(|&o| next(seen.entry(o).or_insert(0))).collect() };
369 let rounds = counts.iter().fold(0, |most, &c| most.max(c + 1));
370 (rounds > 1).then_some((counts, rounds as usize))
371}
372
373pub fn inverse_permutation<T: Copy + Default + 'static>(order: impl ExactSizeIterator<Item = usize>) -> Vec<T> where usize: num_traits::AsPrimitive<T> {
375 let mut ranks = vec![T::default(); order.len()];
376 for (rank, position) in order.enumerate() { ranks[position] = num_traits::AsPrimitive::as_(rank); }
377 ranks
378}
379
380#[inline]
382pub fn find_with<E>(pattern_shape: &[usize], shape: &[usize], mut same: impl FnMut(usize, usize) -> Result<bool, E>) -> Result<Vec<bool>, E> {
383 let mut out = vec![false; count(shape)];
384 if pattern_shape.len() > shape.len() { return Ok(out); }
385 let mut pattern = vec![1; shape.len() - pattern_shape.len()];
386 pattern.extend_from_slice(pattern_shape);
387 let strides = crate::layout::strides(shape);
388 let offsets: Vec<usize> = (0..count(&pattern)).map(|i| crate::layout::coordinates(i, &pattern).map(|(a, c)| c * strides[a]).sum()).collect();
389 for (flat, matched) in out.iter_mut().enumerate() {
390 if crate::layout::coordinates(flat, shape).any(|(a, c)| pattern[a] > shape[a] - c) { continue; }
391 *matched = true;
392 for (i, &offset) in offsets.iter().enumerate() { if !same(i, flat + offset)? { *matched = false; break; } }
393 }
394 Ok(out)
395}
396
397pub fn view_find(pattern: View<'_>, pattern_shape: &[usize], target: View<'_>, shape: &[usize], t: Tolerance) -> Option<Vec<bool>> {
399 #[inline]
400 fn find(pattern_shape: &[usize], shape: &[usize], mut same: impl FnMut(usize, usize) -> bool) -> Vec<bool> {
401 let Ok(found) = find_with(pattern_shape, shape, |i, j| Ok::<_, Infallible>(same(i, j)));
402 found
403 }
404 use tolerant::Tolerant;
405 Some(match (pattern, target) {
406 (View::Booleans(x), View::Booleans(y)) => find(pattern_shape, shape, |i, j| x[i] == y[j]),
407 (View::Integers(x), View::Integers(y)) => find(pattern_shape, shape, |i, j| x.get(i) == y.get(j)),
408 (View::Characters(x), View::Characters(y)) => find(pattern_shape, shape, |i, j| x[i] == y[j]),
409 (View::Floats(x), View::Floats(y)) => crate::with_tolerance!(t, |t| find(pattern_shape, shape, |i, j| t.matches(x.get(i), y.get(j)))), _ => return None
410 })
411}
412
413#[inline(never)]
415pub fn class_counts<K: Key>(x: &[K]) -> Vec<usize> {
416 let Some((min, counts)) = counts(x) else { return numbered(&classes(x)).counts };
417 let mut first = vec![u32::MAX; counts.len()];
418 first_positions(x, min, &mut first, Some(counts.iter().filter(|&&c| c > 0).count())).into_iter().map(|f| counts[x[f].key().wrapping_sub(min) as usize])
419 .collect()
420}
421
422#[inline(never)]
424pub fn bins<K: Key>(boundaries: &[K], values: &[K]) -> Vec<u32> {
425 const LANES: usize = 16;
426 let n = boundaries.len();
427 let (Some(&first), Some(&last)) = (boundaries.first(), boundaries.last()) else { return vec![0; values.len()] };
428 let at = move |k: K| k.key().wrapping_sub(first.key()) as u64 as usize;
429 if let Some(range) = table_size(last.key().wrapping_sub(first.key()) as u64, n).filter(|_| boundaries.is_sorted()) {
430 let mut below = vec![0u32; range];
431 for (i, &b) in boundaries.iter().enumerate() { below[at(b)] = i as u32 + 1 }
433 let mut total = 0;
434 for b in &mut below { total = total.max(*b); *b = total }
435 return values.iter().map(move |&v| {
436 if v < first { 0 } else if v > last { n as u32 } else { below[at(v)] }
437 }).collect();
438 }
439 let search = move |v: &[K], out: &mut [u32]| {
440 let mut base = [0; LANES];
441 let mut len = n;
442 while len > 1 {
443 let half = len / 2;
444 for (b, &x) in base.iter_mut().zip(v) { *b = core::hint::select_unpredictable(boundaries[*b + half] <= x, *b + half, *b) }
446 len -= half;
447 }
448 for ((o, &b), &x) in out.iter_mut().zip(&base).zip(v) { *o = (b + usize::from(boundaries[b] <= x)) as u32 }
449 };
450 let mut out = vec![0; values.len()];
451 let mut start = 0;
452 while start < values.len() {
453 let end = values.len().min(start + LANES);
454 search(&values[start..end], &mut out[start..end]);
455 start = end;
456 }
457 out
458}
459
460fn table<K: Key, E: Copy>(x: &[K], empty: E) -> Option<(i64, Vec<E>)> {
462 if x.is_empty() { return None; }
463 let (min, max) = K::range(x);
464 table_size(max.abs_diff(min), x.len()).map(|range| (min, vec![empty; range]))
465}
466
467fn table_size(span: u64, len: usize) -> Option<usize> { usize::try_from(span).ok()?.checked_add(1).filter(|&r| r <= len.saturating_mul(2).max(1024)) }
469
470fn write_first<'t, K: Key>(x: &[K], min: i64, first: &'t mut [u32], present: Option<usize>, mut found: impl FnMut(usize, bool)) -> &'t [u32] {
472 let (mut empty, mut stall, mut counted) = match present { Some(h) => (h, 0, true), None => (first.len(), 0, first.len() > 128) };
473 for (b, block) in x.chunks(64).enumerate() {
474 let before = empty;
475 for (j, &n) in block.iter().enumerate() {
476 let entry = &mut first[n.key().wrapping_sub(min) as usize];
477 let new = *entry == u32::MAX;
478 empty -= usize::from(new);
479 *entry = if new { (b * 64 + j) as u32 } else { *entry };
480 found(b * 64 + j, new);
481 }
482 stall = if empty == before { stall + 1 } else { 0 };
483 if !counted && stall == STALL {
485 counted = true;
486 empty -= first.len() - held(x, min);
487 }
488 if empty == 0 { break; }
489 }
490 first
491}
492
493const STALL: usize = 4;
495
496fn held<K: Key>(x: &[K], min: i64) -> usize {
498 let mut seen = [false; 128];
499 for k in x { seen[k.key().wrapping_sub(min) as usize & 127] = true }
500 seen.iter().filter(|&&s| s).count()
501}
502
503fn first_positions<K: Key>(x: &[K], min: i64, first: &mut [u32], present: Option<usize>) -> Vec<usize> {
505 let (mut positions, mut count) = (vec![0; first.len().min(x.len()) + 1], 0);
507 write_first(x, min, first, present, |p, new| {
508 positions[count] = p;
509 count += usize::from(new);
510 });
511 positions.truncate(count);
512 positions
513}
514
515#[inline(never)]
517pub fn first_matches<K: Copy + Sync + Eq + Hash>(x: &[K], y: &[K]) -> Vec<u32> {
518 let (miss, mut first) = (x.len() as u32, HashMap::with_capacity(x.len()));
519 for (i, &n) in x.iter().enumerate().rev() { first.insert(n, i as u32); }
520 let (first, mut found) = (&first, Vec::with_capacity(y.len()));
521 look_up(&mut found, y, move |n| first.get(&n).map_or(miss, |&f| f));
522 found
523}
524
525pub fn first_classes<T: Copy, K: Eq + Hash>(x: &[T], key: impl Fn(T) -> K) -> Vec<u32> {
527 let mut first = HashMap::with_capacity(x.len());
528 x.iter().enumerate().map(|(i, &v)| *first.entry(key(v)).or_insert(i as u32)).collect()
529}
530
531#[derive(Clone, Copy)]
533pub struct Exact<'a> { shape: Option<&'a [usize]>, items: View<'a> }
534
535impl<'a> Exact<'a> {
536 pub fn new(shape: Option<&'a [usize]>, items: View<'a>) -> Option<Self> {
538 matches!(items, View::Booleans(_) | View::Integers(_) | View::Characters(_)).then_some(Self { shape, items })
539 }
540
541 fn item(self, i: usize) -> (bool, i64) {
543 match self.items {
544 View::Characters(x) => (true, x[i].key()), View::Booleans(x) => (false, x[i].key()), View::Integers(x) => (false, x.get(i)),
545 View::Floats(_) | View::Complexes(_) => unreachable!("an exact cell holds no floats")
546 }
547 }
548}
549
550impl PartialEq for Exact<'_> {
551 fn eq(&self, other: &Self) -> bool {
552 fn equal<T: PartialEq>(x: &[T], y: &[T]) -> bool { x.len() == y.len() && x.iter().zip(y).all(|(a, b)| a == b) }
554 let shapes = match (self.shape, other.shape) { (Some(x), Some(y)) => equal(x, y), (x, y) => x.is_none() && y.is_none() };
555 shapes && match (self.items, other.items) {
556 (View::Characters(x), View::Characters(y)) => equal(x, y), (x, y) => view_matches(x, y, Tolerance::EXACT) == Some(true)
557 }
558 }
559}
560
561impl Eq for Exact<'_> {}
562
563impl Hash for Exact<'_> {
564 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
565 match self.items {
566 View::Characters(x) => crate::characters::codes(x).hash(state),
567 View::Booleans(x) => { for &b in x { state.write_i64(b.into()) } }
568 View::Integers(x) => with_ints!(x, |x| for &i in x {
569 state.write_i64(i.into())
570 }),
571 View::Floats(_) | View::Complexes(_) => unreachable!("an exact cell holds no floats"),
572 }
573 }
574}
575
576impl Ord for Exact<'_> {
577 fn cmp(&self, other: &Self) -> Ordering {
578 let (x, y) = match (self.shape, other.shape) {
579 (None, Some(_)) => return Ordering::Less, (Some(_), None) => return Ordering::Greater, (x, y) => (x.unwrap_or_default(), y.unwrap_or_default())
580 };
581 if x.len() != y.len() { return x.len().cmp(&y.len()); }
582 for i in 0..self.items.len().min(other.items.len()) {
583 let ((xc, xk), (yc, yk)) = (self.item(i), other.item(i));
584 let order = crate::characters::order_of_kinds(xc, yc).then(xk.cmp(&yk));
585 if order.is_ne() { return order; }
586 }
587 self.items.len().cmp(&other.items.len()).then_with(|| x.cmp(y))
588 }
589}
590
591impl PartialOrd for Exact<'_> { fn partial_cmp(&self, other: &Self) -> Option<Ordering> { Some(self.cmp(other)) } }
592
593pub(crate) fn first_mask(classes: &[u32]) -> Vec<bool> { classes.iter().enumerate().map(|(i, &c)| c as usize == i).collect() }
595
596#[inline(never)]
598pub fn grouped(groups: &[usize], n: usize) -> Groups { grouped_by(groups, n, |g| g) }
599
600#[derive(Clone, Debug, PartialEq)]
602pub struct Groups {
603 pub sizes: Vec<usize>,
605 buffers: Vec<Vec<u32>>,
606}
607
608impl Groups {
609 pub fn each(&self) -> impl Iterator<Item = &[u32]> {
611 let (mut buffers, mut rest) = (self.buffers.iter(), &[][..]);
612 self.sizes.iter().map(move |&n| {
613 while rest.len() < n { rest = buffers.next().expect("a buffer holds each group") }
614 let (these, more) = rest.split_at(n);
615 rest = more;
616 these
617 })
618 }
619
620 pub fn into_positions(mut self) -> Vec<u32> { if self.buffers.len() == 1 { return self.buffers.pop().expect("one buffer"); } self.buffers.concat() }
622 pub fn into_buffers(self) -> Vec<Vec<u32>> { self.buffers }
624}
625
626#[derive(Clone, Debug, PartialEq)]
628pub struct Grouped {
629 pub firsts: Vec<usize>,
631 pub groups: Groups,
633}
634
635fn numbered_table<K: Key>(x: &[K]) -> Option<(i64, Vec<u32>, Vec<usize>)> {
637 checked_len(x);
638 let (min, mut table) = table(x, u32::MAX)?;
639 let firsts = first_positions(x, min, &mut table, None);
640 for (g, &f) in firsts.iter().enumerate() { table[x[f].key().wrapping_sub(min) as usize] = g as u32; }
641 Some((min, table, firsts))
642}
643
644#[inline(never)]
646pub fn grouped_keys<K: Key>(x: &[K]) -> Option<Grouped> {
647 let (min, table, firsts) = numbered_table(x)?;
648 let table = &table[..];
649 Some(if let [lower, _] = *table { grouped_classes(x, firsts, move |k| k.key().wrapping_sub(min) as usize ^ lower as usize) }
651 else { grouped_classes(x, firsts, move |k| table[k.key().wrapping_sub(min) as usize] as usize) })
652}
653
654#[inline(never)]
656pub fn numbered_keys<K: Key>(x: &[K]) -> Option<Numbered> {
657 let (min, table, firsts) = numbered_table(x)?;
658 let (at, table) = (move |k: K| k.key().wrapping_sub(min) as usize, &table[..]);
659 let numbers = with_width!(Width::below(firsts.len()), T => {
660 let mut numbers: Vec<T> = Vec::new();
661 if let [lower, _] = *table {
663 crate::pairs(crate::Arg::Many(x), crate::Arg::One(()), x.len(), &mut numbers, |k, ()| (num_traits::AsPrimitive::<T>::as_(at(k) ^ lower as usize), true));
664 } else {
665 crate::pairs(crate::Arg::Many(x), crate::Arg::One(()), x.len(), &mut numbers, |k, ()| (num_traits::AsPrimitive::<T>::as_(table[at(k)]), true));
666 }
667 int_buf(numbers)
668 });
669 let counts = with_ints!(numbers.view(), |numbers| ClassNumbers { count: firsts.len(), numbers }.counts());
670 Some(Numbered { firsts, numbers, counts })
671}
672
673#[inline(always)]
675pub fn grouped_by<T: Copy + Sync>(items: &[T], n: usize, group: impl Fn(T) -> usize + Sync) -> Groups {
676 let len = checked_len(items);
677 if n <= 1 { return Groups { sizes: vec![len; n], buffers: vec![(0..len as u32).collect()] }; }
678 if n == 2 { return two_groups(items, group, false); }
679 let (positions, sizes) = if n <= 256 && n * 16 <= len { parted::<4, T>(items, n, group) } else { parted::<1, T>(items, n, group) };
681 Groups { sizes, buffers: vec![positions] }
682}
683
684fn checked_len<T>(items: &[T]) -> usize { assert!(u32::try_from(items.len()).is_ok(), "grouped positions hold at most u32::MAX items"); items.len() }
686
687#[inline(always)]
689fn grouped_classes<T: Copy + Sync>(items: &[T], firsts: Vec<usize>, class: impl Fn(T) -> usize + Sync) -> Grouped {
690 checked_len(items);
691 let groups = if firsts.len() == 2 { two_groups(items, class, true) } else { grouped_by(items, firsts.len(), class) };
692 Grouped { firsts, groups }
693}
694
695fn group_bytes<T: Copy + Sync>(items: &[T], group: impl Fn(T) -> usize + Sync) -> (Vec<u8>, usize) {
697 let mut groups = Vec::with_capacity(items.len());
698 crate::pairs(crate::Arg::Many(items), crate::Arg::One(()), items.len(), &mut groups, |k, ()| (group(k) as u8, true));
699 let ones = count_eq(&groups, 1);
700 (groups, ones)
701}
702
703fn two_groups<T: Copy + Sync>(items: &[T], group: impl Fn(T) -> usize + Sync, apart: bool) -> Groups {
705 let (groups, ones) = group_bytes(items, group);
706 let zeros = groups.len() - ones;
707 let buffers = if apart {
708 let (mut d0, mut d1) = (Vec::new(), Vec::new());
709 unsafe { crate::slice::appended(&mut d0, zeros, |a| crate::slice::appended(&mut d1, ones, |b| split2(&groups, a, b))) };
710 vec![d0, d1]
711 } else {
712 let mut d = Vec::new();
713 unsafe {
714 crate::slice::appended(&mut d, groups.len(), |d| {
715 let (a, b) = d.split_at_mut(zeros);
716 split2(&groups, a, b)
717 })
718 };
719 vec![d]
720 };
721 Groups { sizes: vec![zeros, ones], buffers }
722}
723
724fn split2(groups: &[u8], d0: &mut [MaybeUninit<u32>], d1: &mut [MaybeUninit<u32>]) {
726 let (blocks, rest) = groups.as_chunks::<8>();
727 let (mut k0, mut k1) = (0, 0);
728 for (b, block) in blocks.iter().enumerate() {
729 let (ones, start) = (crate::gather::mask_bits(block), (8 * b) as u32);
730 k1 = crate::gather::selected(d1, k1, ones, |i| start + u32::from(i));
731 k0 = crate::gather::selected(d0, k0, !ones, |i| start + u32::from(i));
732 }
733 let (start, ones) = ((8 * blocks.len()) as u32, crate::gather::mask_bits(rest));
734 k1 = crate::gather::selected(d1, k1, ones, |i| start + u32::from(i));
735 k0 = crate::gather::selected(d0, k0, !ones & ((1u16 << rest.len()) - 1) as u8, |i| start + u32::from(i));
736 assert_eq!((k0, k1), (d0.len(), d1.len()), "each group must be 0 or 1");
737}
738
739fn parted<const P: usize, T: Copy>(items: &[T], n: usize, group: impl Fn(T) -> usize) -> (Vec<u32>, Vec<usize>) {
741 let (len, part) = (items.len(), items.len() / P);
742 let parts: [&[T]; P] = std::array::from_fn(|q| &items[q * part..][..part]);
743 let rest = &items[P * part..];
744 let (mut positions, mut sizes) = (vec![0u32; len], vec![0; n]);
745 let mut next = vec![[0u32; P]; n];
746 let (slots, next) = (&mut positions[..], &mut next[..]);
748 for j in 0..part { for q in 0..P { next[group(parts[q][j])][q] += 1; } }
749 for &item in rest { next[group(item)][P - 1] += 1; }
750 let mut start = 0;
752 for (places, size) in next.iter_mut().zip(&mut sizes) {
753 for place in places.iter_mut() {
754 *size += *place as usize;
755 (*place, start) = (start, start + *place);
756 }
757 }
758 let mut place = move |i: usize, item: T, q: usize| {
759 let p = &mut next[group(item)][q];
760 let at = *p;
761 slots[at as usize] = i as u32;
762 *p = at + 1;
763 };
764 for j in 0..part { for q in 0..P { place(q * part + j, parts[q][j], q); } }
765 for (i, &item) in rest.iter().enumerate() { place(P * part + i, item, P - 1); }
766 (positions, sizes)
767}
768
769const RADIX_FROM: usize = 256;
771
772const fn value_radix_from(bytes: usize) -> usize { if bytes > 4 { 1 << 17 } else { RADIX_FROM } }
774
775#[inline(never)]
777pub fn sort_rows<K: RadixKey>(n: usize, columns: Vec<Vec<K>>) -> Vec<u32> {
778 if n < RADIX_FROM {
779 let mut order: Vec<u32> = (0..n as u32).collect();
780 order.sort_by(|&a, &b| columns.iter().map(|c| c[a as usize].cmp(&c[b as usize])).find(|o| o.is_ne()).unwrap_or(std::cmp::Ordering::Equal));
781 return order;
782 }
783 if let [keys] = &columns[..] && let Some(order) = counted(keys) { return order; }
784 let mut order: Vec<u32> = (0..n as u32).collect();
785 if columns.is_empty() { return order; }
786 let mut spare = (vec![K::default(); n], vec![0; n]);
787 for (c, mut keys) in columns.into_iter().rev().enumerate() { radix_column(&mut keys, &mut order, &mut spare, c == 0); }
788 order
789}
790
791fn radix_column<K: RadixKey>(keys: &mut Vec<K>, order: &mut Vec<u32>, spare: &mut (Vec<K>, Vec<u32>), first: bool) {
793 if !first {
794 for (key, &i) in spare.0.iter_mut().zip(order.iter()) { *key = keys[i as usize]; }
795 std::mem::swap(keys, &mut spare.0);
796 }
797 radix(keys, order, (&mut spare.0, &mut spare.1));
798}
799
800fn counted<K: RadixKey>(keys: &[K]) -> Option<Vec<u32>> {
802 let n = keys.len();
803 let size = |(lo, hi): (K, K)| table_size(hi.word() - lo.word(), n);
804 size(K::extremes(&keys[..n.min(256)]))?;
805 let (lo, hi) = K::extremes(keys);
806 let range = size((lo, hi))?;
807 let lo = lo.word();
808 Some(grouped_by(keys, range, move |k| (k.word() - lo) as usize).into_positions())
809}
810
811pub(crate) fn sorted<K: RadixKey>(mut keys: Vec<K>) -> (Vec<K>, Vec<u32>) {
813 let n = keys.len();
814 let mut order: Vec<u32> = (0..n as u32).collect();
815 if n < RADIX_FROM { order.sort_by_key(|&i| keys[i as usize]); return (order.iter().map(|&i| keys[i as usize]).collect(), order); }
816 radix(&mut keys, &mut order, (&mut vec![K::default(); n], &mut vec![0; n]));
817 (keys, order)
818}
819
820fn byte_counts<K: RadixKey>(keys: &[K]) -> [[u32; 256]; 8] {
822 let mut counts = [[0u32; 256]; 8];
823 for &k in keys { for (b, count) in counts[..K::BYTES].iter_mut().enumerate() { count[k.byte(b)] += 1; } }
824 counts
825}
826
827fn scattered<K: RadixKey, P: Copy>(keys: &[K], carried: &[P], b: usize, next: &mut [usize; 256], dest: &mut [K], places: &mut [P]) {
829 let carried = &carried[..keys.len()];
830 for j in 0..keys.len() {
831 let place = &mut next[keys[j].byte(b)];
832 (dest[*place], places[*place]) = (keys[j], carried[j]);
833 *place += 1;
834 }
835}
836
837fn radix<K: RadixKey, P: Copy + Send + Sync>(keys: &mut [K], payload: &mut [P], spare: (&mut [K], &mut [P])) {
839 let n = keys.len();
840 if n > RADIX_PIECE && crate::parallel::threads() > 1 { return radix_pieces(keys, payload, spare); }
841 let counts = byte_counts(keys);
842 let counts = &counts[..K::BYTES];
843 let (mut from, mut to, mut passes) = ((keys, payload), spare, 0);
844 for (b, count) in counts.iter().enumerate() {
845 if count.contains(&(n as u32)) { continue; }
846 let (mut next, mut total) = ([0usize; 256], 0);
847 for (place, &c) in next.iter_mut().zip(count) { (*place, total) = (total, total + c as usize) }
848 let (source, carried, dest, places) = (&from.0[..n], &from.1[..n], &mut to.0[..n], &mut to.1[..n]);
849 let heavy = source[0].byte(b);
851 if count[heavy] as usize >= n - n / 64 {
852 for (keys, carried) in source.chunks(64).zip(carried.chunks(64)) {
853 if keys.iter().fold(true, |all, k| all & (k.byte(b) == heavy)) {
854 let at = next[heavy];
855 next[heavy] = at + keys.len();
856 dest[at..next[heavy]].copy_from_slice(keys);
857 places[at..next[heavy]].copy_from_slice(carried);
858 } else { scattered(keys, carried, b, &mut next, dest, places) }
859 }
860 } else { scattered(source, carried, b, &mut next, dest, places) }
861 (from, to, passes) = (to, from, passes + 1);
862 }
863 settled(passes, from, to);
864}
865
866const RADIX_PIECE: usize = 1 << 16;
868
869fn settled<K: Copy, P: Copy>(passes: usize, from: (&mut [K], &mut [P]), to: (&mut [K], &mut [P])) {
871 if passes % 2 == 1 {
872 to.0.copy_from_slice(from.0);
873 to.1.copy_from_slice(from.1);
874 }
875}
876
877fn radix_pieces<K: RadixKey, P: Copy + Send + Sync>(keys: &mut [K], payload: &mut [P], spare: (&mut [K], &mut [P])) {
879 use crate::parallel::pieces_of;
880 struct Runs<'a, K, P> { keys: Vec<std::slice::IterMut<'a, K>>, places: Vec<std::slice::IterMut<'a, P>> }
882 let n = keys.len();
883 let count = (n / RADIX_PIECE).min(crate::parallel::threads());
884 let size = n.div_ceil(count);
885 let piece = move |p: usize| p * size..n.min((p + 1) * size);
886 let mut every = vec![[[0u32; 256]; 8]; count];
887 pieces_of(&mut every, 1, usize::MAX, &|p, c| { c[0] = byte_counts(&keys[piece(p)]); true });
888 let (mut from, mut to, mut passes) = ((keys, payload), spare, 0);
889 for b in 0..K::BYTES {
890 let total = (0..256).map(|d| every.iter().map(|c| c[b][d] as usize).sum::<usize>());
891 if total.into_iter().any(|t| t == n) { continue; }
892 let mut counts: Vec<[u32; 256]> = every.iter().map(|c| c[b]).collect();
893 if passes > 0 {
894 pieces_of(&mut counts, 1, usize::MAX, &|p, c| {
895 c[0] = [0; 256];
896 for &k in &from.0[piece(p)] { c[0][k.byte(b)] += 1; }
897 true
898 });
899 }
900 {
901 let (mut dest, mut places) = (&mut to.0[..n], &mut to.1[..n]);
902 let mut runs: Vec<Runs<'_, K, P>> = (0..count).map(|_| Runs { keys: Vec::with_capacity(256), places: Vec::with_capacity(256) }).collect();
903 for d in 0..256 {
904 for (run, c) in runs.iter_mut().zip(&counts) {
905 let (k, rest) = std::mem::take(&mut dest).split_at_mut(c[d] as usize);
906 let (o, others) = std::mem::take(&mut places).split_at_mut(c[d] as usize);
907 (dest, places) = (rest, others);
908 run.keys.push(k.iter_mut());
909 run.places.push(o.iter_mut());
910 }
911 }
912 let (source, carried) = (&from.0[..n], &from.1[..n]);
913 pieces_of(&mut runs, 1, usize::MAX, &|p, run| {
914 let run = &mut run[0];
915 let range = piece(p);
916 let (source, carried) = (&source[range.clone()], &carried[range]);
917 for j in 0..source.len() {
918 let d = source[j].byte(b);
919 *run.keys[d].next().expect("a run holds each key that the counts gave its piece") = source[j];
920 *run.places[d].next().expect("a run holds each key that the counts gave its piece") = carried[j];
921 }
922 true
923 });
924 }
925 (from, to, passes) = (to, from, passes + 1);
926 }
927 settled(passes, from, to);
928}
929
930pub struct Chains { ends: HashMap<u64, (u32, u32)>, next: Vec<u32> }
932
933impl Chains {
934 pub fn new(len: usize) -> Self { Self { ends: HashMap::with_capacity(len), next: vec![u32::MAX; len] } }
936 pub fn push(&mut self, key: u64, i: usize) {
938 let i = i as u32;
939 match self.ends.entry(key) {
940 Entry::Occupied(mut e) => {
941 let (_, last) = e.get_mut();
942 self.next[*last as usize] = i;
943 *last = i;
944 }
945 Entry::Vacant(e) => { e.insert((i, i)); }
946 }
947 }
948 pub fn push_new(&mut self, key: u64, i: usize, copy: impl Fn(usize) -> bool) {
950 match self.ends.entry(key) {
951 Entry::Occupied(mut e) => {
952 let (first, last) = e.get_mut();
953 if links(&self.next, *first, u32::MAX).any(copy) { return; }
954 self.next[*last as usize] = i as u32;
955 *last = i as u32;
956 }
957 Entry::Vacant(e) => { e.insert((i as u32, i as u32)); }
958 }
959 }
960 #[inline]
962 pub fn chain(&self, key: u64) -> Links<'_> { links(&self.next, self.first(key), u32::MAX) }
963 #[inline]
965 pub fn chains(&self, key: u64, then: Option<u64>) -> Links<'_> { links(&self.next, self.first(key), then.map_or(u32::MAX, |k| self.first(k))) }
966 #[inline]
968 fn first(&self, key: u64) -> u32 { self.ends.get(&key).map_or(u32::MAX, |&(first, _)| first) }
969}
970
971#[inline]
973fn links(next: &[u32], at: u32, then: u32) -> Links<'_> { Links { next, at, then } }
974
975pub struct Links<'a> { next: &'a [u32], at: u32, then: u32 }
977
978impl Iterator for Links<'_> {
979 type Item = usize;
980 #[inline]
981 fn next(&mut self) -> Option<usize> {
982 if self.at == u32::MAX {
983 if self.then == u32::MAX { return None; }
984 (self.at, self.then) = (self.then, u32::MAX);
985 }
986 let i = self.at;
987 self.at = self.next[i as usize];
988 Some(i as usize)
989 }
990}
991
992macro_rules! keys_of {
994 ($view:expr, |$x:ident| $body:expr) => {
995 match $view {
996 View::Booleans($x) => Some($body),
997 View::Characters($x) => Some($body),
998 View::Integers(ints) => Some(with_ints!(ints, |$x| $body)),
999 _ => None,
1000 }
1001 };
1002}
1003
1004macro_rules! key_pairs {
1006 ($a:expr, $b:expr, |$x:ident, $y:ident| $body:expr) => {
1007 match ($a, $b) {
1008 (View::Booleans($x), View::Booleans($y)) => Some($body),
1009 (View::Characters($x), View::Characters($y)) => Some($body),
1010 (View::Integers($x), View::Integers($y)) => Some(with_int_pair!($x, $y, |$x, $y| $body)),
1011 _ => None,
1012 }
1013 };
1014}
1015
1016macro_rules! keys_or_floats {
1018 ($x:expr, $t:expr, $keys:path, $floats:path) => {{
1019 if let Some(found) = keys_of!($x, |x| $keys(x)) {
1020 return Some(Ok(found));
1021 }
1022 let View::Floats(x) = $x else { return None };
1023 Some(with_floats!(x, |x| crate::with_tolerance!($t, |t| $floats(t, x))))
1024 }};
1025}
1026
1027pub fn view_classes(x: View<'_>, t: Tolerance) -> Option<Result<Vec<u32>, Error>> { keys_or_floats!(x, t, classes, tolerant::classes) }
1029
1030pub fn view_firsts(x: View<'_>, t: Tolerance) -> Option<Result<Vec<bool>, Error>> { keys_or_floats!(x, t, firsts, tolerant::firsts) }
1032
1033pub fn view_class_counts(x: View<'_>, t: Tolerance) -> Option<Result<Vec<usize>, Error>> {
1035 if let Some(counts) = keys_of!(x, |x| class_counts(x)) { return Some(Ok(counts)); }
1036 Some(view_classes(x, t)?.map(|classes| numbered(&classes).counts))
1037}
1038
1039pub fn view_grouped(x: View<'_>, t: Tolerance) -> Option<Result<Grouped, Error>> {
1041 if let Some(Some(grouped)) = keys_of!(x, |x| grouped_keys(x)) { return Some(Ok(grouped)); }
1042 Some(view_classes(x, t)?.map(|classes| numbered(&classes).grouped()))
1043}
1044
1045pub fn view_numbered(x: View<'_>, t: Tolerance) -> Option<Result<Numbered, Error>> {
1047 if let Some(Some(numbered)) = keys_of!(x, |x| numbered_keys(x)) { return Some(Ok(numbered)); }
1048 Some(view_classes(x, t)?.map(|classes| numbered(&classes)))
1049}
1050
1051pub fn view_index_of(x: View<'_>, y: View<'_>, t: Tolerance) -> Option<Result<Vec<u32>, Error>> {
1053 if let (DType::Float(w), true) = (x.dtype(), x.dtype() == y.dtype()) {
1054 return with_float_width!(w, F => Some(crate::with_tolerance!(t, |t| tolerant::index_of(t, F::of(x)?, F::of(y)?))));
1055 }
1056 if let (View::Integers(x), View::Floats(Floats::F64(y))) = (x, y) { return with_ints!(x, |x| whole_needles(x, y, t).map(Ok)); }
1057 key_pairs!(x, y, |x, y| Ok(index_of(x, y)))
1058}
1059
1060pub fn view_matches(x: View<'_>, y: View<'_>, t: Tolerance) -> Option<bool> {
1062 use tolerant::{Tolerant, TolerantComplex};
1063 if x.len() != y.len() { return Some(false); }
1064 let whole = |v: View<'_>, i: usize| match v {
1066 View::Booleans(v) => i64::from(v[i]), View::Integers(v) => v.get(i), _ => unreachable!("Booleans or integers")
1067 };
1068 Some(match (x, y) {
1069 (View::Booleans(x), View::Booleans(y)) => x == y,
1070 (View::Integers(x), View::Integers(y)) => x == y,
1071 (x @ (View::Booleans(_) | View::Integers(_)), y @ (View::Booleans(_) | View::Integers(_))) => (0..x.len()).all(|i| whole(x, i) == whole(y, i)),
1072 (n @ (View::Booleans(_) | View::Integers(_)), View::Floats(f)) | (View::Floats(f), n @ (View::Booleans(_) | View::Integers(_))) => {
1073 crate::with_tolerance!(t, |t| with_floats!(f, |f| f.iter().enumerate().all(|(i, &f)| t.matches::<f64>(whole(n, i) as f64, f.into()))))
1074 }
1075 (View::Floats(x), View::Floats(y)) => {
1076 crate::with_tolerance!(t, |t| with_floats!(x, |x| with_floats!(y, |y| x.iter().zip(y).all(|(&x, &y)| t.matches::<f64>(x.into(), y.into())))))
1077 }
1078 (View::Complexes(x), View::Complexes(y)) => crate::with_tolerance!(t, |t| (0..x.len()).all(|i| t.complex_equal(x.get(i), y.get(i)))),
1079 (View::Characters(x), View::Characters(y)) => x == y,
1080 _ => return None,
1081 })
1082}
1083
1084pub fn view_contains(x: View<'_>, y: View<'_>, cells: usize, t: Tolerance) -> Option<Result<bool, Error>> {
1086 if y.len() == 1 { if let Some(found) = view_index_of(x, y, t) { return Some(found.map(|found| (found[0] as usize) < cells)); } }
1087 for i in 0..cells {
1088 if i.is_multiple_of(crate::BETWEEN_CHECKS) { if let Err(e) = check() { return Some(Err(e)); } }
1089 if view_matches(x.slice(i * y.len()..(i + 1) * y.len()), y, t)? { return Some(Ok(true)); }
1090 }
1091 Some(Ok(false))
1092}
1093
1094pub fn member_layout(layout: &Layout, query: &Layout) -> Result<Option<Layout>, Error> {
1096 let rank = layout.shape().len();
1097 if rank == 0 { return Err(Error::new(Kind::Rank, "a rank-0 array is not a container")); }
1098 let cell = layout.axes(1..rank);
1099 Ok(cell.matches(query).then_some(cell))
1100}
1101
1102pub fn view_member(x: &Layout, xs: View<'_>, query: &Layout, ys: View<'_>, t: Tolerance) -> Result<bool, Error> {
1104 let Some(cell) = member_layout(x, query)? else { return Ok(false) };
1105 let Some(dtype) = xs.dtype().common(ys.dtype()) else { return Ok(false) };
1106 let haystack = ArrayView::compact(x, xs).cast(dtype)?;
1107 let axes: Vec<usize> = (0..cell.shape().len()).collect();
1108 let mapping = Mapping::new(query, &axes, &cell, ys.len());
1109 let needle = ArrayView::compact(query, ys).remapped(mapping, cell.count()).cast(dtype)?;
1110 view_contains(haystack.view(), needle.view(), x.shape()[0], t).ok_or_else(|| Error::new(Kind::Domain, "membership can't compare these kinds"))?
1111}
1112
1113fn whole_needles<K: Key>(x: &[K], y: &[f64], t: Tolerance) -> Option<Vec<u32>> {
1115 let whole: Vec<i64> = y.iter().map(|&v| tolerant::integer_of(t, v)).collect::<Option<_>>()?;
1116 let mut found = index_of(x, &whole.iter().filter_map(|&n| K::from_key(n)).collect::<Vec<_>>()).into_iter();
1117 Some(whole.iter().map(|&n| K::from_key(n).and_then(|_| found.next()).unwrap_or(x.len() as u32)).collect())
1118}
1119
1120#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1122pub enum Set { Union, Intersection, Difference }
1123
1124pub fn view_set(x: View<'_>, y: View<'_>, how: Set, t: Tolerance) -> Option<Result<Buffer, Error>> {
1126 if x.dtype() != y.dtype() { return None; }
1127 let mut joined = Buffer::with_capacity(x.dtype(), if how == Set::Union { x.len() + y.len() } else { 0 });
1128 let values = if how == Set::Union {
1129 joined.extend_from(x, 0..x.len());
1130 joined.extend_from(y, 0..y.len());
1131 joined.view()
1132 } else { x };
1133 let firsts = match view_firsts(values, t)? { Ok(firsts) => firsts, Err(e) => return Some(Err(e)) };
1134 let matches = if how == Set::Union { None } else { match view_index_of(y, values, t)? { Ok(found) => Some(found), Err(e) => return Some(Err(e)) } };
1135 let positions: Vec<u32> = firsts.iter().enumerate().filter_map(|(i, &first)| {
1136 (first && matches.as_ref().is_none_or(|found| ((found[i] as usize) < y.len()) == (how == Set::Intersection))).then_some(i as u32)
1137 }).collect();
1138 Some(Ok(Buffer::gathered(values, &positions)))
1139}
1140
1141pub fn view_bins(boundaries: View<'_>, values: View<'_>) -> Option<Vec<u32>> {
1143 if let (DType::Float(w), true) = (boundaries.dtype(), boundaries.dtype() == values.dtype()) {
1144 return with_float_width!(w, F => Some(bins(&tolerant::order_keys(F::of(boundaries)?), &tolerant::order_keys(F::of(values)?))));
1145 }
1146 key_pairs!(boundaries, values, |x, y| bins(x, y))
1147}
1148
1149pub fn view_bins_below(boundaries: View<'_>, values: View<'_>) -> Option<Vec<u32>> {
1151 if boundaries.dtype() != values.dtype() { return None; }
1152 let (mut a, mut v) = (view_order_keys(boundaries)?, view_order_keys(values)?);
1153 let n = a.len() as u32;
1154 a.reverse();
1155 for k in a.iter_mut().chain(&mut v) { *k = !*k; }
1156 Some(bins(&a, &v).into_iter().map(|p| n - p).collect())
1157}
1158
1159macro_rules! ordered_view {
1161 ($view:expr, |$x:ident| $body:expr) => {
1162 match $view {
1163 View::Booleans($x) => Some($body),
1164 View::Characters($x) => Some($body),
1165 View::Integers(ints) => Some(with_ints!(ints, |$x| $body)),
1166 View::Floats(floats) => Some(with_floats!(floats, |$x| $body)),
1167 View::Complexes(_) => None,
1168 }
1169 };
1170 ($view:expr, |$x:ident, $key:ident| $body:expr) => {
1171 match $view {
1172 View::Booleans($x) => {
1173 let $key = SortKey::sort_key;
1174 Some($body)
1175 }
1176 View::Characters($x) => {
1177 let $key = SortKey::sort_key;
1178 Some($body)
1179 }
1180 View::Integers(ints) => Some(with_ints!(ints, |$x| {
1181 let $key = SortKey::sort_key;
1182 $body
1183 })),
1184 View::Floats(floats) => Some(with_floats!(floats, |$x| {
1185 let $key = Float::order_key;
1186 $body
1187 })),
1188 View::Complexes(_) => None,
1189 }
1190 };
1191}
1192
1193pub fn view_order_keys(x: View<'_>) -> Option<Vec<u64>> { ordered_view!(x, |x, key| x.iter().map(|&v| key(v).word()).collect()) }
1195
1196pub fn view_grade(x: View<'_>, down: bool) -> Option<Vec<u32>> {
1198 ordered_view!(x, |x, key| sort_rows(x.len(), vec![order_column(x, x.len(), 0, 1, down, key, |k| k)]))
1199}
1200
1201pub fn view_lexgrade<'a>(n: usize, columns: impl Iterator<Item = (View<'a>, usize, usize)>, down: bool) -> Option<Vec<u32>> {
1203 let columns: Vec<_> = columns.collect();
1204 if let [(view, start, step)] = columns[..] { return ordered_view!(view, |x, key| sort_rows(n, vec![order_column(x, n, start, step, down, key, |k| k)])); }
1205 if n < RADIX_FROM {
1206 let columns = columns.into_iter().map(|(view, start, step)| ordered_view!(view, |x, key| order_column(x, n, start, step, down, key, |k| k.word())))
1208 .collect::<Option<Vec<_>>>()?;
1209 return Some(sort_rows(n, columns));
1210 }
1211 let (mut order, mut positions) = ((0..n as u32).collect::<Vec<_>>(), vec![0; n]);
1212 for (c, (view, start, step)) in columns.into_iter().rev().enumerate() {
1213 ordered_view!(view, |x, key| {
1214 let mut keys = order_column(x, n, start, step, down, key, |k| k);
1215 let mut spare = (vec![Default::default(); n], std::mem::take(&mut positions));
1216 radix_column(&mut keys, &mut order, &mut spare, c == 0);
1217 positions = spare.1;
1218 })?;
1219 }
1220 Some(order)
1221}
1222
1223pub fn view_lexgrade_along(columns: &[(View<'_>, usize, usize)], dims @ [outer, len, inner]: [usize; 3], down: bool) -> Result<Option<Vec<i64>>, Error> {
1225 let lane = |o: usize, i: usize| view_lexgrade(len, columns.iter().map(|&(c, start, step)| (c, start + (o * len * inner + i) * step, inner * step)), down);
1226 if outer * inner == 1 { return Ok(lane(0, 0).map(|g| g.into_iter().map(i64::from).collect())); }
1227 if columns.iter().any(|(c, ..)| matches!(c, View::Complexes(_))) { return Ok(None); }
1228 Ok(Some(lanes::<i64, ()>(dims, len, |o, i, out, ()| {
1229 for (p, g) in out.iter_mut().zip(lane(o, i).expect("keys that order")) { *p = g.into(); }
1230 })?))
1231}
1232
1233fn order_column<T: Copy + Sync, K: RadixKey, O: Send>(
1235 x: &[T], n: usize, start: usize, step: usize, down: bool, key: impl Fn(T) -> K + Sync, convert: impl Fn(K) -> O + Sync) -> Vec<O> {
1236 let key = move |v| convert(if down { !key(v) } else { key(v) });
1237 if step != 1 { return (0..n).map(|i| key(x[start + i * step])).collect(); }
1238 let mut keys = Vec::with_capacity(n);
1239 crate::pairs(crate::Arg::Many(&x[start..start + n]), crate::Arg::One(()), n, &mut keys, move |v, ()| (key(v), true));
1240 keys
1241}
1242
1243fn spread_into<T: Element + Key>(x: &[T], out: &mut [T], down: bool) -> bool {
1245 if x.len() < 256 { return false; }
1246 let Some((min, counts)) = counts(x) else { return false };
1247 let mut at = 0;
1248 let mut write = |(k, &c): (usize, &usize)| {
1249 let v = T::from_key(min + k as i64).unwrap_or(T::FILL);
1250 if at + 4 <= out.len() { out[at..at + 4].fill(v); if c > 4 { out[at + 4..at + c].fill(v); } }
1251 else { out[at..at + c].fill(v); }
1252 at += c;
1253 };
1254 if down { counts.iter().enumerate().rev().for_each(&mut write); }
1255 else { counts.iter().enumerate().for_each(&mut write); }
1256 true
1257}
1258
1259const VECTORS: bool = cfg!(any(target_arch = "aarch64", target_feature = "avx2", target_feature = "simd128"));
1261const VECTOR_BYTES: usize = if cfg!(target_feature = "avx2") { 32 } else { 16 };
1263const HELD_ROWS: usize = if cfg!(target_arch = "aarch64") { 16 } else { 8 };
1265const NETWORK_BYTES: usize = 32 << 10;
1267const ANY_FILL_BYTES: usize = 4 << 10;
1269
1270pub(crate) trait Network: Sized { fn networked(_x: &[Self], _out: &mut [Self], _down: bool, _scratch: &mut Vec<Self>) -> bool { false } }
1272impl Network for half::f16 {}
1273impl Network for f32 {
1274 fn networked(x: &[f32], out: &mut [f32], down: bool, scratch: &mut Vec<f32>) -> bool { networked::<_, { VECTOR_BYTES / 4 }>(x, out, down, scratch) }
1275}
1276impl Network for f64 {
1277 fn networked(x: &[f64], out: &mut [f64], down: bool, scratch: &mut Vec<f64>) -> bool { networked::<_, { VECTOR_BYTES / 8 }>(x, out, down, scratch) }
1278}
1279
1280fn networked<F: Native, const W: usize>(x: &[F], out: &mut [F], down: bool, scratch: &mut Vec<F>) -> bool {
1282 let padded = x.len().next_power_of_two() * size_of::<F>();
1283 if !VECTORS || x.len() < W * HELD_ROWS || padded > NETWORK_BYTES || padded > ANY_FILL_BYTES && 4 * size_of_val(x) < 3 * padded { return false; }
1284 scratch.clear();
1285 let real = x.len() - x.iter().filter(|v| v.is_nan()).count();
1286 if real == x.len() { scratch.extend_from_slice(x) }
1287 else { scratch.extend(x.iter().copied().filter(|v| !v.is_nan())) }
1288 scratch.resize(real.next_power_of_two().max(W * HELD_ROWS), F::infinity());
1289 let rows = scratch.as_chunks_mut::<W>().0;
1290 network(rows);
1291 for (j, lane) in out[..real].chunks_mut(rows.len()).enumerate() { for (o, row) in lane.iter_mut().zip(&*rows) { *o = row[j]; } }
1292 if real < x.len() {
1293 let mut nans: Vec<F::OrderKey> = x.iter().filter(|v| v.is_nan()).map(|v| v.total_key()).collect();
1294 nans.sort_unstable();
1295 for (o, &k) in out[real..].iter_mut().zip(&nans) { *o = F::from_order_key(k); }
1296 }
1297 let sorted = &mut out[..real];
1298 let (zeros, positive) = (sorted.partition_point(|&v| v < F::zero()), sorted.partition_point(|&v| v <= F::zero()));
1299 let negative = zeros + sorted[zeros..positive].iter().filter(|v| v.is_sign_negative()).count();
1300 sorted[zeros..negative].fill(-F::zero());
1301 sorted[negative..positive].fill(F::zero());
1302 if down { out.reverse(); }
1303 true
1304}
1305
1306fn network<F: Native, const W: usize>(rows: &mut [[F; W]]) {
1308 for block in rows.as_chunks_mut::<HELD_ROWS>().0 { sorted_block(block) }
1309 let mut b = 2 * HELD_ROWS;
1310 while b <= rows.len() {
1311 for part in rows.chunks_exact_mut(b) { flip(part) }
1312 cleaned(rows, b / 4);
1313 b *= 2;
1314 }
1315 lanes_stage::<F, W, 1>(rows);
1316 if W >= 4 { lanes_stage::<F, W, 2>(rows) }
1317 if W >= 8 { lanes_stage::<F, W, 3>(rows) }
1318}
1319
1320#[inline(always)]
1322fn sorted_block<F: Native, const W: usize>(block: &mut [[F; W]; HELD_ROWS]) {
1323 let mut v = *block;
1324 held_step::<F, W, 2, true>(&mut v);
1325 held_step::<F, W, 4, true>(&mut v);
1326 half_cleans(&mut v, 2);
1327 held_step::<F, W, 8, true>(&mut v);
1328 half_cleans(&mut v, 4);
1329 if HELD_ROWS == 16 {
1330 held_step::<F, W, 16, true>(&mut v);
1331 half_cleans(&mut v, 8);
1332 }
1333 *block = v;
1334}
1335
1336fn cleaned<F: Native, const W: usize>(rows: &mut [[F; W]], mut d: usize) {
1338 while d >= HELD_ROWS { for part in rows.chunks_exact_mut(2 * d) { half_clean(part) } d /= 2 }
1339 for block in rows.as_chunks_mut::<HELD_ROWS>().0 {
1340 let mut v = *block;
1341 half_cleans(&mut v, HELD_ROWS);
1342 *block = v;
1343 }
1344}
1345
1346#[inline(always)]
1348fn half_cleans<F: Native, const W: usize>(v: &mut [[F; W]; HELD_ROWS], span: usize) {
1349 if span >= 16 { held_step::<F, W, 16, false>(v) }
1350 if span >= 8 { held_step::<F, W, 8, false>(v) }
1351 if span >= 4 { held_step::<F, W, 4, false>(v) }
1352 held_step::<F, W, 2, false>(v);
1353}
1354
1355#[inline(always)]
1357fn held_step<F: Native, const W: usize, const SPAN: usize, const FLIP: bool>(v: &mut [[F; W]; HELD_ROWS]) {
1358 for base in (0..HELD_ROWS).step_by(SPAN) {
1359 for r in 0..SPAN / 2 {
1360 let (i, k) = (base + r, if FLIP { base + SPAN - 1 - r } else { base + SPAN / 2 + r });
1361 (v[i], v[k]) = F::ordered_lanes(v[i], v[k]);
1362 }
1363 }
1364}
1365
1366fn lanes_stage<F: Native, const W: usize, const T: u32>(rows: &mut [[F; W]]) {
1368 let (m, h) = ((1 << T) - 1, 1 << (T - 1));
1369 let half = rows.len() / 2;
1370 let (lo, hi) = rows.split_at_mut(half);
1371 for (a, c) in lo.iter_mut().zip(hi.iter_mut().rev()) {
1372 let (s, l) = F::ordered_lanes(*a, std::array::from_fn(|j| c[j ^ m]));
1373 *a = std::array::from_fn(|j| if j & h == 0 { s[j] } else { l[j] });
1374 *c = std::array::from_fn(|j| if j & h == 0 { s[j ^ m] } else { l[j ^ m] });
1375 }
1376 for u in (0..T - 1).rev() {
1377 let g = 1 << u;
1378 for v in rows.iter_mut() {
1379 let (s, l) = F::ordered_lanes::<W>(std::array::from_fn(|j| v[j & !g]), std::array::from_fn(|j| v[j | g]));
1380 *v = std::array::from_fn(|j| if j & g == 0 { s[j] } else { l[j] });
1381 }
1382 }
1383 cleaned(rows, half);
1384}
1385
1386#[inline(always)]
1388fn flip<F: Native, const W: usize>(part: &mut [[F; W]]) {
1389 let (lo, hi) = part.split_at_mut(part.len() / 2);
1390 for (a, c) in lo.iter_mut().zip(hi.iter_mut().rev()) { (*a, *c) = F::ordered_lanes(*a, *c) }
1391}
1392
1393#[inline(always)]
1395fn half_clean<F: Native, const W: usize>(part: &mut [[F; W]]) {
1396 let (lo, hi) = part.split_at_mut(part.len() / 2);
1397 for (a, c) in lo.iter_mut().zip(hi) { (*a, *c) = F::ordered_lanes(*a, *c) }
1398}
1399
1400pub fn view_sort(x: View<'_>, down: bool) -> Result<Option<Buffer>, Error> { view_sort_along(x, [1, x.len(), 1], down) }
1402
1403pub fn view_sort_along(x: View<'_>, dims: [usize; 3], down: bool) -> Result<Option<Buffer>, Error> {
1405 fn run<T: SortKey>(x: &[T], dims @ [outer, len, inner]: [usize; 3], down: bool) -> Result<Buffer, Error> {
1406 let sort = |lane: &[T], out: &mut [T], (keys, spare, scratch): &mut (Vec<T::Key>, Vec<T::Key>, Vec<T>)| {
1407 if T::sorted_into(lane, out, down, scratch) { return; }
1408 T::keys_into(lane, keys, down);
1409 let n = keys.len();
1410 if n < value_radix_from(<T::Key as RadixKey>::BYTES) { keys.sort_unstable(); } else {
1411 spare.resize(n, Default::default());
1412 radix(keys, &mut vec![(); n], (spare, &mut vec![(); n]));
1413 }
1414 T::write_back(keys, out, down);
1415 };
1416 Ok(T::buffer(if outer * inner == 1 {
1417 let mut out = vec![T::FILL; len];
1418 sort(x, &mut out, &mut Default::default());
1419 out
1420 } else { map_lanes::<_, _, (Vec<T::Key>, Vec<T::Key>, Vec<T>)>(x, dims, len, |lane, out, work| sort(lane, out, work))? }))
1421 }
1422 ordered_view!(x, |x| run(x, dims, down)).transpose()
1423}
1424
1425pub fn select<T>(x: &mut [T], ranks: &[usize], compare: impl Fn(&T, &T) -> Ordering + Copy) {
1427 fn within<T>(x: &mut [T], ranks: &[usize], base: usize, compare: impl Fn(&T, &T) -> Ordering + Copy) {
1428 if ranks.is_empty() { return; }
1429 let middle = ranks.len() / 2;
1430 let rank = ranks[middle] - base;
1431 let (before, _, after) = x.select_nth_unstable_by(rank, compare);
1432 within(before, &ranks[..middle], base, compare);
1433 within(after, &ranks[middle + 1..], base + rank + 1, compare);
1434 }
1435 within(x, ranks, 0, compare);
1436}
1437
1438pub fn view_argpartition(x: View<'_>, rank: usize) -> Option<Vec<u32>> {
1440 let keys = view_order_keys(x)?;
1441 let mut out: Vec<u32> = (0..x.len() as u32).collect();
1442 select(&mut out, &[rank], |&a, &b| keys[a as usize].cmp(&keys[b as usize]));
1443 Some(out)
1444}
1445
1446pub fn view_partition_along(x: View<'_>, dims: [usize; 3], rank: usize, indices: bool) -> Result<Option<Buffer>, Error> {
1448 fn run<T: SortKey>(x: &[T], dims: [usize; 3], rank: usize, indices: bool) -> Result<Buffer, Error> {
1449 Ok(if indices {
1450 i64::buffer(map_lanes::<_, _, Vec<T::Key>>(x, dims, dims[1], |lane, out, keys| {
1451 T::keys_into(lane, keys, false);
1452 for (i, p) in out.iter_mut().enumerate() { *p = i as i64; }
1453 select(out, &[rank], |a, b| keys[*a as usize].cmp(&keys[*b as usize]));
1454 })?)
1455 } else {
1456 T::buffer(map_lanes::<_, _, Vec<T::Key>>(x, dims, dims[1], |lane, out, keys| {
1457 T::keys_into(lane, keys, false);
1458 select(keys, &[rank], |a, b| a.cmp(b));
1459 T::write_back(keys, out, false);
1460 })?)
1461 })
1462 }
1463 ordered_view!(x, |x| run(x, dims, rank, indices)).transpose()
1464}
1465
1466const BLOCK: usize = 64;
1468
1469fn ranked<T: Copy + PartialOrd, const BACK: bool>(x: &[T], key: impl Fn(T) -> u64, down: bool, range: std::ops::Range<usize>) -> Vec<u32> {
1471 if range.is_empty() { return Vec::new(); }
1472 let (n, m, flip) = (x.len(), range.end, if down { u64::MAX } else { 0 });
1473 let at = |j: usize| if BACK { n - 1 - j } else { j };
1474 let place = |t: u32| (if BACK { !t } else { t }) as usize;
1475 let pair = |p: usize| (key(x[p]) ^ flip, if BACK { !(p as u32) } else { p as u32 });
1476 let mut kept: Vec<(u64, u32)> = (0..m).map(|j| pair(at(j))).collect();
1477 if m * BLOCK < n {
1479 let mut last = *kept.iter().max().expect("range.end is at least 1");
1481 let (room, mut compacted) = (2 * m, m);
1482 let mut j = m;
1483 'scan: while j < n {
1484 let len = BLOCK.min(n - j);
1485 let start = if BACK { n - j - len } else { j };
1486 let limit = x[place(last.1)];
1487 let below = <&[T; BLOCK]>::try_from(&x[start..start + len]).map_or(true, |block| {
1488 if down { block.iter().fold(false, |below, &v| below | !(v <= limit)) } else { block.iter().fold(false, |below, &v| below | !(v >= limit)) }
1489 });
1490 if below {
1491 for i in j..j + len {
1492 let p = pair(at(i));
1493 if p.0 >= last.0 { continue; }
1494 kept.push(p);
1495 if kept.len() < room { continue; }
1496 if i - compacted < room && i >= 4 * room { kept.extend((i + 1..n).map(|i| pair(at(i)))); break 'scan; }
1498 kept.select_nth_unstable(m - 1);
1499 kept.truncate(m);
1500 last = kept[m - 1];
1501 compacted = i;
1502 }
1503 }
1504 j += len;
1505 }
1506 } else { kept.extend((m..n).map(|j| pair(at(j)))); }
1507 let mut ranks = vec![range.start, range.end - 1];
1508 ranks.dedup();
1509 select(&mut kept, &ranks, Ord::cmp);
1510 let run = &mut kept[range];
1511 run.sort_unstable();
1512 run.iter().map(|&(_, t)| place(t) as u32).collect()
1513}
1514
1515pub fn view_selected(x: View<'_>, range: std::ops::Range<usize>, down: bool) -> Option<Vec<u32>> {
1517 let n = x.len();
1518 ordered_view!(x, |x, key| if range.end <= n - range.start {
1519 ranked::<_, false>(x, |v| key(v).word(), down, range)
1520 } else {
1521 let mut run = ranked::<_, true>(x, |v| key(v).word(), !down, n - range.end..n - range.start);
1522 run.reverse();
1523 run
1524 })
1525}
1526
1527pub fn view_argextremes(x: View<'_>, [outer, len, inner]: [usize; 3], down: bool) -> Option<Vec<i64>> {
1529 fn extremes<T: Copy + Sync, K: Ord>(x: &[T], [outer, len, inner]: [usize; 3], down: bool, key: impl Fn(T) -> K + Sync) -> Vec<i64> {
1530 let mut out = Vec::with_capacity(outer * inner);
1531 crate::parallel::extended(&mut out, outer * inner, (crate::CHUNK / len.max(1)).max(1), &|range, slots| {
1532 unsafe {
1533 slots.fill(range.len(), |slots| {
1534 for (lane, slot) in range.zip(slots) {
1535 let start = lane / inner * len * inner + lane % inner;
1536 let (mut best, mut at) = (key(x[start]), 0);
1537 for j in 1..len {
1538 let k = key(x[start + j * inner]);
1539 if if down { k > best } else { k < best } {
1540 best = k;
1541 at = j;
1542 }
1543 }
1544 slot.write(at as i64);
1545 }
1546 true
1547 })
1548 };
1549 true
1550 });
1551 out
1552 }
1553 if len == 0 { return None; }
1554 ordered_view!(x, |x, key| extremes(x, [outer, len, inner], down, key))
1555}
1556
1557struct Quantiles { points: Vec<(usize, usize, f64)>, ranks: Vec<usize> }
1559
1560impl Quantiles {
1561 fn new(len: usize, probabilities: &[f64]) -> Result<Self, Error> {
1563 if !probabilities.iter().all(|q| (0.0..=1.0).contains(q)) {
1564 return Err(Error::new(Kind::Domain, "quantile probabilities must be between zero and one"));
1565 }
1566 let points: Vec<_> = probabilities.iter().map(|&q| {
1567 let p = q * len.saturating_sub(1) as f64;
1568 (p.floor() as usize, p.ceil() as usize, p.fract())
1569 }).collect();
1570 let mut ranks: Vec<_> = points.iter().flat_map(|&(lo, hi, _)| [lo, hi]).collect();
1571 ranks.sort_unstable();
1572 ranks.dedup();
1573 Ok(Self { points, ranks })
1574 }
1575
1576 fn apply<T: Copy, F: Float>(&self, x: &mut [T], out: &mut [F], compare: impl Fn(&T, &T) -> Ordering + Copy, read: impl Fn(T) -> F::Working) {
1577 if x.is_empty() { out.fill(F::nan()); return; }
1578 select(x, &self.ranks, compare);
1579 for (out, &(lo, hi, weight)) in out.iter_mut().zip(&self.points) {
1580 let a = read(x[lo]);
1581 *out = F::from_working(if lo == hi { a }
1582 else { let weight = F::Working::narrow(weight); a * (F::Working::narrow(1.0) - weight) + read(x[hi]) * weight });
1583 }
1584 }
1585}
1586
1587pub fn quantiles<F: Float>(x: &mut [F], probabilities: &[f64]) -> Result<Vec<F>, Error> {
1589 let mut out = vec![F::nan(); probabilities.len()];
1590 Quantiles::new(x.len(), probabilities)?.apply(x, &mut out, |a, b| a.order_key().cmp(&b.order_key()), Float::working);
1591 Ok(out)
1592}
1593
1594pub fn view_quantiles_along(x: View<'_>, dims @ [outer, _, inner]: [usize; 3], probabilities: &[f64]) -> Result<Option<Buffer>, Error> {
1596 use crate::buffer::Source;
1597 fn run<T: SortKey, F: Element + Float + Default>(
1598 x: &[T], dims: [usize; 3], plan: &Quantiles, read: impl Fn(T) -> F::Working + Sync) -> Result<Vec<F>, Error> {
1599 map_lanes::<_, _, Vec<T::Key>>(x, dims, plan.points.len(), |lane, out, keys| {
1600 T::keys_into(lane, keys, false);
1601 plan.apply(keys, out, |a, b| a.cmp(b), |k| read(T::from_sort_key(k)));
1602 })
1603 }
1604 fn floats<F: Element + Float + Default + Network>(x: &[F], dims: [usize; 3], plan: &Quantiles) -> Result<Buffer, Error> {
1605 Ok(F::buffer(run::<_, F>(x, dims, plan, Float::working)?))
1606 }
1607 let plan = Quantiles::new(dims[1], probabilities)?;
1608 let items = match x {
1609 View::Booleans(x) => f64::buffer(run::<_, f64>(x, dims, &plan, Source::<f64>::read)?),
1610 View::Integers(ints) => with_ints!(ints, |x| f64::buffer(run::<_, f64>(x, dims, &plan, Source::<f64>::read)?)),
1611 View::Floats(x) => with_floats!(x, |x| floats(x, dims, &plan)?), _ => return Ok(None)
1612 };
1613 let count = probabilities.len();
1614 Ok(Some(if count > 1 && outer > 1 {
1615 let plan = crate::structure::transpose(&crate::layout::Layout::from([outer, count, inner]), &[1, 0, 2]).expect("valid quantile axes");
1616 Buffer::walked(items.view(), 0, &plan.steps, items.len())
1617 } else { items }))
1618}
1619
1620fn map_lanes<T: Copy + Sync, U: Element + Default, S: Default>(
1622 x: &[T], dims @ [_, len, inner]: [usize; 3], width: usize, map: impl Fn(&[T], &mut [U], &mut S) + Sync) -> Result<Vec<U>, Error> {
1623 lanes::<U, (Vec<T>, S)>(dims, width, |o, i, out, (input, work)| {
1624 let input = if inner == 1 { &x[o * len..(o + 1) * len] } else {
1625 input.clear();
1626 input.extend((0..len).map(|j| x[(o * len + j) * inner + i]));
1627 &input[..]
1628 };
1629 map(input, out, work);
1630 })
1631}
1632
1633fn lanes<U: Element + Default, S: Default>(
1635 [outer, _, inner]: [usize; 3], width: usize, fill: impl Fn(usize, usize, &mut [U], &mut S) + Sync) -> Result<Vec<U>, Error> {
1636 if outer * inner * width == 0 { return Ok(Vec::new()); }
1637 let ((out,), _) = crate::parallel::stack::<(Vec<U>,), S>(outer * inner, [width], &|lane, (out,), work| {
1638 fill(lane / inner, lane % inner, out, work);
1639 true
1640 })?;
1641 if inner == 1 || width == 1 { return Ok(out); }
1642 let mut transposed = Vec::with_capacity(out.len());
1643 transposed.resize(out.len(), U::FILL);
1644 for o in 0..outer {
1645 crate::linalg::transpose_into(
1646 &out[o * inner * width..][..inner * width], width, &mut transposed[o * inner * width..][..inner * width], inner, [inner, width], |x| x);
1647 }
1648 Ok(transposed)
1649}