1use crate::hashing::{HashMap, HashMapExt, HashSet, HashSetExt};
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(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<T: Copy, K: Eq + Hash>(x: &[T], key: impl Fn(T) -> K) -> Vec<bool> {
595 let mut seen = HashSet::with_capacity(x.len());
596 x.iter().map(|&v| seen.insert(key(v))).collect()
597}
598
599#[inline(never)]
601pub fn grouped(groups: &[usize], n: usize) -> Groups { grouped_by(groups, n, |g| g) }
602
603#[derive(Clone, Debug, PartialEq)]
605pub struct Groups {
606 pub sizes: Vec<usize>,
608 buffers: Vec<Vec<u32>>,
609}
610
611impl Groups {
612 pub fn each(&self) -> impl Iterator<Item = &[u32]> {
614 let (mut buffers, mut rest) = (self.buffers.iter(), &[][..]);
615 self.sizes.iter().map(move |&n| {
616 while rest.len() < n { rest = buffers.next().expect("a buffer holds each group") }
617 let (these, more) = rest.split_at(n);
618 rest = more;
619 these
620 })
621 }
622
623 pub fn into_positions(mut self) -> Vec<u32> { if self.buffers.len() == 1 { return self.buffers.pop().expect("one buffer"); } self.buffers.concat() }
625 pub fn into_buffers(self) -> Vec<Vec<u32>> { self.buffers }
627}
628
629#[derive(Clone, Debug, PartialEq)]
631pub struct Grouped {
632 pub firsts: Vec<usize>,
634 pub groups: Groups,
636}
637
638fn numbered_table<K: Key>(x: &[K]) -> Option<(i64, Vec<u32>, Vec<usize>)> {
640 checked_len(x);
641 let (min, mut table) = table(x, u32::MAX)?;
642 let firsts = first_positions(x, min, &mut table, None);
643 for (g, &f) in firsts.iter().enumerate() { table[x[f].key().wrapping_sub(min) as usize] = g as u32; }
644 Some((min, table, firsts))
645}
646
647#[inline(never)]
649pub fn grouped_keys<K: Key>(x: &[K]) -> Option<Grouped> {
650 let (min, table, firsts) = numbered_table(x)?;
651 let table = &table[..];
652 Some(if let [lower, _] = *table { grouped_classes(x, firsts, move |k| k.key().wrapping_sub(min) as usize ^ lower as usize) }
654 else { grouped_classes(x, firsts, move |k| table[k.key().wrapping_sub(min) as usize] as usize) })
655}
656
657#[inline(never)]
659pub fn numbered_keys<K: Key>(x: &[K]) -> Option<Numbered> {
660 let (min, table, firsts) = numbered_table(x)?;
661 let (at, table) = (move |k: K| k.key().wrapping_sub(min) as usize, &table[..]);
662 let numbers = with_width!(Width::below(firsts.len()), T => {
663 let mut numbers: Vec<T> = Vec::new();
664 if let [lower, _] = *table {
666 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));
667 } else {
668 crate::pairs(crate::Arg::Many(x), crate::Arg::One(()), x.len(), &mut numbers, |k, ()| (num_traits::AsPrimitive::<T>::as_(table[at(k)]), true));
669 }
670 int_buf(numbers)
671 });
672 let counts = with_ints!(numbers.view(), |numbers| ClassNumbers { count: firsts.len(), numbers }.counts());
673 Some(Numbered { firsts, numbers, counts })
674}
675
676#[inline(always)]
678pub fn grouped_by<T: Copy + Sync>(items: &[T], n: usize, group: impl Fn(T) -> usize + Sync) -> Groups {
679 let len = checked_len(items);
680 if n <= 1 { return Groups { sizes: vec![len; n], buffers: vec![(0..len as u32).collect()] }; }
681 if n == 2 { return two_groups(items, group, false); }
682 let (positions, sizes) = if n <= 256 && n * 16 <= len { parted::<4, T>(items, n, group) } else { parted::<1, T>(items, n, group) };
684 Groups { sizes, buffers: vec![positions] }
685}
686
687fn checked_len<T>(items: &[T]) -> usize { assert!(u32::try_from(items.len()).is_ok(), "grouped positions hold at most u32::MAX items"); items.len() }
689
690#[inline(always)]
692fn grouped_classes<T: Copy + Sync>(items: &[T], firsts: Vec<usize>, class: impl Fn(T) -> usize + Sync) -> Grouped {
693 checked_len(items);
694 let groups = if firsts.len() == 2 { two_groups(items, class, true) } else { grouped_by(items, firsts.len(), class) };
695 Grouped { firsts, groups }
696}
697
698fn group_bytes<T: Copy + Sync>(items: &[T], group: impl Fn(T) -> usize + Sync) -> (Vec<u8>, usize) {
700 let mut groups = Vec::with_capacity(items.len());
701 crate::pairs(crate::Arg::Many(items), crate::Arg::One(()), items.len(), &mut groups, |k, ()| (group(k) as u8, true));
702 let ones = count_eq(&groups, 1);
703 (groups, ones)
704}
705
706fn two_groups<T: Copy + Sync>(items: &[T], group: impl Fn(T) -> usize + Sync, apart: bool) -> Groups {
708 let (groups, ones) = group_bytes(items, group);
709 let zeros = groups.len() - ones;
710 let buffers = if apart {
711 let (mut d0, mut d1) = (Vec::new(), Vec::new());
712 unsafe { crate::slice::appended(&mut d0, zeros, |a| crate::slice::appended(&mut d1, ones, |b| split2(&groups, a, b))) };
713 vec![d0, d1]
714 } else {
715 let mut d = Vec::new();
716 unsafe {
717 crate::slice::appended(&mut d, groups.len(), |d| {
718 let (a, b) = d.split_at_mut(zeros);
719 split2(&groups, a, b)
720 })
721 };
722 vec![d]
723 };
724 Groups { sizes: vec![zeros, ones], buffers }
725}
726
727fn split2(groups: &[u8], d0: &mut [MaybeUninit<u32>], d1: &mut [MaybeUninit<u32>]) {
729 let (blocks, rest) = groups.as_chunks::<8>();
730 let (mut k0, mut k1) = (0, 0);
731 for (b, block) in blocks.iter().enumerate() {
732 let (ones, start) = (crate::gather::mask_bits(block), (8 * b) as u32);
733 k1 = crate::gather::selected(d1, k1, ones, |i| start + u32::from(i));
734 k0 = crate::gather::selected(d0, k0, !ones, |i| start + u32::from(i));
735 }
736 let (start, ones) = ((8 * blocks.len()) as u32, crate::gather::mask_bits(rest));
737 k1 = crate::gather::selected(d1, k1, ones, |i| start + u32::from(i));
738 k0 = crate::gather::selected(d0, k0, !ones & ((1u16 << rest.len()) - 1) as u8, |i| start + u32::from(i));
739 assert_eq!((k0, k1), (d0.len(), d1.len()), "each group must be 0 or 1");
740}
741
742fn parted<const P: usize, T: Copy>(items: &[T], n: usize, group: impl Fn(T) -> usize) -> (Vec<u32>, Vec<usize>) {
744 let (len, part) = (items.len(), items.len() / P);
745 let parts: [&[T]; P] = std::array::from_fn(|q| &items[q * part..][..part]);
746 let rest = &items[P * part..];
747 let (mut positions, mut sizes) = (vec![0u32; len], vec![0; n]);
748 let mut next = vec![[0u32; P]; n];
749 let (slots, next) = (&mut positions[..], &mut next[..]);
751 for j in 0..part { for q in 0..P { next[group(parts[q][j])][q] += 1; } }
752 for &item in rest { next[group(item)][P - 1] += 1; }
753 let mut start = 0;
755 for (places, size) in next.iter_mut().zip(&mut sizes) {
756 for place in places.iter_mut() {
757 *size += *place as usize;
758 (*place, start) = (start, start + *place);
759 }
760 }
761 let mut place = move |i: usize, item: T, q: usize| {
762 let p = &mut next[group(item)][q];
763 let at = *p;
764 slots[at as usize] = i as u32;
765 *p = at + 1;
766 };
767 for j in 0..part { for q in 0..P { place(q * part + j, parts[q][j], q); } }
768 for (i, &item) in rest.iter().enumerate() { place(P * part + i, item, P - 1); }
769 (positions, sizes)
770}
771
772const RADIX_FROM: usize = 256;
774
775const fn value_radix_from(bytes: usize) -> usize { if bytes > 4 { 1 << 17 } else { RADIX_FROM } }
777
778#[inline(never)]
780pub fn sort_rows<K: RadixKey>(n: usize, columns: Vec<Vec<K>>) -> Vec<u32> {
781 if n < RADIX_FROM {
782 let mut order: Vec<u32> = (0..n as u32).collect();
783 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));
784 return order;
785 }
786 if let [keys] = &columns[..] && let Some(order) = counted(keys) { return order; }
787 let mut order: Vec<u32> = (0..n as u32).collect();
788 if columns.is_empty() { return order; }
789 let mut spare = (vec![K::default(); n], vec![0; n]);
790 for (c, mut keys) in columns.into_iter().rev().enumerate() { radix_column(&mut keys, &mut order, &mut spare, c == 0); }
791 order
792}
793
794fn radix_column<K: RadixKey>(keys: &mut Vec<K>, order: &mut Vec<u32>, spare: &mut (Vec<K>, Vec<u32>), first: bool) {
796 if !first {
797 for (key, &i) in spare.0.iter_mut().zip(order.iter()) { *key = keys[i as usize]; }
798 std::mem::swap(keys, &mut spare.0);
799 }
800 radix(keys, order, (&mut spare.0, &mut spare.1));
801}
802
803fn counted<K: RadixKey>(keys: &[K]) -> Option<Vec<u32>> {
805 let n = keys.len();
806 let size = |(lo, hi): (K, K)| table_size(hi.word() - lo.word(), n);
807 size(K::extremes(&keys[..n.min(256)]))?;
808 let (lo, hi) = K::extremes(keys);
809 let range = size((lo, hi))?;
810 let lo = lo.word();
811 Some(grouped_by(keys, range, move |k| (k.word() - lo) as usize).into_positions())
812}
813
814pub(crate) fn sorted<K: RadixKey>(mut keys: Vec<K>) -> (Vec<K>, Vec<u32>) {
816 let n = keys.len();
817 let mut order: Vec<u32> = (0..n as u32).collect();
818 if n < RADIX_FROM { order.sort_by_key(|&i| keys[i as usize]); return (order.iter().map(|&i| keys[i as usize]).collect(), order); }
819 radix(&mut keys, &mut order, (&mut vec![K::default(); n], &mut vec![0; n]));
820 (keys, order)
821}
822
823fn byte_counts<K: RadixKey>(keys: &[K]) -> [[u32; 256]; 8] {
825 let mut counts = [[0u32; 256]; 8];
826 for &k in keys { for (b, count) in counts[..K::BYTES].iter_mut().enumerate() { count[k.byte(b)] += 1; } }
827 counts
828}
829
830fn scattered<K: RadixKey, P: Copy>(keys: &[K], carried: &[P], b: usize, next: &mut [usize; 256], dest: &mut [K], places: &mut [P]) {
832 let carried = &carried[..keys.len()];
833 for j in 0..keys.len() {
834 let place = &mut next[keys[j].byte(b)];
835 (dest[*place], places[*place]) = (keys[j], carried[j]);
836 *place += 1;
837 }
838}
839
840fn radix<K: RadixKey, P: Copy + Send + Sync>(keys: &mut [K], payload: &mut [P], spare: (&mut [K], &mut [P])) {
842 let n = keys.len();
843 if n > RADIX_PIECE && crate::parallel::threads() > 1 { return radix_pieces(keys, payload, spare); }
844 let counts = byte_counts(keys);
845 let counts = &counts[..K::BYTES];
846 let (mut from, mut to, mut passes) = ((keys, payload), spare, 0);
847 for (b, count) in counts.iter().enumerate() {
848 if count.contains(&(n as u32)) { continue; }
849 let (mut next, mut total) = ([0usize; 256], 0);
850 for (place, &c) in next.iter_mut().zip(count) { (*place, total) = (total, total + c as usize) }
851 let (source, carried, dest, places) = (&from.0[..n], &from.1[..n], &mut to.0[..n], &mut to.1[..n]);
852 let heavy = source[0].byte(b);
854 if count[heavy] as usize >= n - n / 64 {
855 for (keys, carried) in source.chunks(64).zip(carried.chunks(64)) {
856 if keys.iter().fold(true, |all, k| all & (k.byte(b) == heavy)) {
857 let at = next[heavy];
858 next[heavy] = at + keys.len();
859 dest[at..next[heavy]].copy_from_slice(keys);
860 places[at..next[heavy]].copy_from_slice(carried);
861 } else { scattered(keys, carried, b, &mut next, dest, places) }
862 }
863 } else { scattered(source, carried, b, &mut next, dest, places) }
864 (from, to, passes) = (to, from, passes + 1);
865 }
866 settled(passes, from, to);
867}
868
869const RADIX_PIECE: usize = 1 << 16;
871
872fn settled<K: Copy, P: Copy>(passes: usize, from: (&mut [K], &mut [P]), to: (&mut [K], &mut [P])) {
874 if passes % 2 == 1 {
875 to.0.copy_from_slice(from.0);
876 to.1.copy_from_slice(from.1);
877 }
878}
879
880fn radix_pieces<K: RadixKey, P: Copy + Send + Sync>(keys: &mut [K], payload: &mut [P], spare: (&mut [K], &mut [P])) {
882 use crate::parallel::pieces_of;
883 struct Runs<'a, K, P> { keys: Vec<std::slice::IterMut<'a, K>>, places: Vec<std::slice::IterMut<'a, P>> }
885 let n = keys.len();
886 let count = (n / RADIX_PIECE).min(crate::parallel::threads());
887 let size = n.div_ceil(count);
888 let piece = move |p: usize| p * size..n.min((p + 1) * size);
889 let mut every = vec![[[0u32; 256]; 8]; count];
890 pieces_of(&mut every, 1, usize::MAX, &|p, c| { c[0] = byte_counts(&keys[piece(p)]); true });
891 let (mut from, mut to, mut passes) = ((keys, payload), spare, 0);
892 for b in 0..K::BYTES {
893 let total = (0..256).map(|d| every.iter().map(|c| c[b][d] as usize).sum::<usize>());
894 if total.into_iter().any(|t| t == n) { continue; }
895 let mut counts: Vec<[u32; 256]> = every.iter().map(|c| c[b]).collect();
896 if passes > 0 {
897 pieces_of(&mut counts, 1, usize::MAX, &|p, c| {
898 c[0] = [0; 256];
899 for &k in &from.0[piece(p)] { c[0][k.byte(b)] += 1; }
900 true
901 });
902 }
903 {
904 let (mut dest, mut places) = (&mut to.0[..n], &mut to.1[..n]);
905 let mut runs: Vec<Runs<'_, K, P>> = (0..count).map(|_| Runs { keys: Vec::with_capacity(256), places: Vec::with_capacity(256) }).collect();
906 for d in 0..256 {
907 for (run, c) in runs.iter_mut().zip(&counts) {
908 let (k, rest) = std::mem::take(&mut dest).split_at_mut(c[d] as usize);
909 let (o, others) = std::mem::take(&mut places).split_at_mut(c[d] as usize);
910 (dest, places) = (rest, others);
911 run.keys.push(k.iter_mut());
912 run.places.push(o.iter_mut());
913 }
914 }
915 let (source, carried) = (&from.0[..n], &from.1[..n]);
916 pieces_of(&mut runs, 1, usize::MAX, &|p, run| {
917 let run = &mut run[0];
918 let range = piece(p);
919 let (source, carried) = (&source[range.clone()], &carried[range]);
920 for j in 0..source.len() {
921 let d = source[j].byte(b);
922 *run.keys[d].next().expect("a run holds each key that the counts gave its piece") = source[j];
923 *run.places[d].next().expect("a run holds each key that the counts gave its piece") = carried[j];
924 }
925 true
926 });
927 }
928 (from, to, passes) = (to, from, passes + 1);
929 }
930 settled(passes, from, to);
931}
932
933pub struct Chains { ends: HashMap<u64, (u32, u32)>, next: Vec<u32> }
935
936impl Chains {
937 pub fn new(len: usize) -> Self { Self { ends: HashMap::with_capacity(len), next: vec![u32::MAX; len] } }
939 pub fn push(&mut self, key: u64, i: usize) {
941 let i = i as u32;
942 match self.ends.entry(key) {
943 Entry::Occupied(mut e) => {
944 let (_, last) = e.get_mut();
945 self.next[*last as usize] = i;
946 *last = i;
947 }
948 Entry::Vacant(e) => { e.insert((i, i)); }
949 }
950 }
951 pub fn push_new(&mut self, key: u64, i: usize, copy: impl Fn(usize) -> bool) {
953 match self.ends.entry(key) {
954 Entry::Occupied(mut e) => {
955 let (first, last) = e.get_mut();
956 if links(&self.next, *first, u32::MAX).any(copy) { return; }
957 self.next[*last as usize] = i as u32;
958 *last = i as u32;
959 }
960 Entry::Vacant(e) => { e.insert((i as u32, i as u32)); }
961 }
962 }
963 #[inline]
965 pub fn chain(&self, key: u64) -> Links<'_> { links(&self.next, self.first(key), u32::MAX) }
966 #[inline]
968 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))) }
969 #[inline]
971 fn first(&self, key: u64) -> u32 { self.ends.get(&key).map_or(u32::MAX, |&(first, _)| first) }
972}
973
974#[inline]
976fn links(next: &[u32], at: u32, then: u32) -> Links<'_> { Links { next, at, then } }
977
978pub struct Links<'a> { next: &'a [u32], at: u32, then: u32 }
980
981impl Iterator for Links<'_> {
982 type Item = usize;
983 #[inline]
984 fn next(&mut self) -> Option<usize> {
985 if self.at == u32::MAX {
986 if self.then == u32::MAX { return None; }
987 (self.at, self.then) = (self.then, u32::MAX);
988 }
989 let i = self.at;
990 self.at = self.next[i as usize];
991 Some(i as usize)
992 }
993}
994
995macro_rules! keys_of {
997 ($view:expr, |$x:ident| $body:expr) => {
998 match $view {
999 View::Booleans($x) => Some($body),
1000 View::Characters($x) => Some($body),
1001 View::Integers(ints) => Some(with_ints!(ints, |$x| $body)),
1002 _ => None,
1003 }
1004 };
1005}
1006
1007macro_rules! key_pairs {
1009 ($a:expr, $b:expr, |$x:ident, $y:ident| $body:expr) => {
1010 match ($a, $b) {
1011 (View::Booleans($x), View::Booleans($y)) => Some($body),
1012 (View::Characters($x), View::Characters($y)) => Some($body),
1013 (View::Integers($x), View::Integers($y)) => Some(with_int_pair!($x, $y, |$x, $y| $body)),
1014 _ => None,
1015 }
1016 };
1017}
1018
1019macro_rules! keys_or_floats {
1021 ($x:expr, $t:expr, $keys:path, $floats:path) => {{
1022 if let Some(found) = keys_of!($x, |x| $keys(x)) {
1023 return Some(Ok(found));
1024 }
1025 let View::Floats(x) = $x else { return None };
1026 Some(with_floats!(x, |x| crate::with_tolerance!($t, |t| $floats(t, x))))
1027 }};
1028}
1029
1030pub fn view_classes(x: View<'_>, t: Tolerance) -> Option<Result<Vec<u32>, Error>> { keys_or_floats!(x, t, classes, tolerant::classes) }
1032
1033pub fn view_firsts(x: View<'_>, t: Tolerance) -> Option<Result<Vec<bool>, Error>> { keys_or_floats!(x, t, firsts, tolerant::firsts) }
1035
1036pub fn view_class_counts(x: View<'_>, t: Tolerance) -> Option<Result<Vec<usize>, Error>> {
1038 if let Some(counts) = keys_of!(x, |x| class_counts(x)) { return Some(Ok(counts)); }
1039 Some(view_classes(x, t)?.map(|classes| numbered(&classes).counts))
1040}
1041
1042pub fn view_grouped(x: View<'_>, t: Tolerance) -> Option<Result<Grouped, Error>> {
1044 if let Some(Some(grouped)) = keys_of!(x, |x| grouped_keys(x)) { return Some(Ok(grouped)); }
1045 Some(view_classes(x, t)?.map(|classes| numbered(&classes).grouped()))
1046}
1047
1048pub fn view_numbered(x: View<'_>, t: Tolerance) -> Option<Result<Numbered, Error>> {
1050 if let Some(Some(numbered)) = keys_of!(x, |x| numbered_keys(x)) { return Some(Ok(numbered)); }
1051 Some(view_classes(x, t)?.map(|classes| numbered(&classes)))
1052}
1053
1054pub fn view_index_of(x: View<'_>, y: View<'_>, t: Tolerance) -> Option<Result<Vec<u32>, Error>> {
1056 if let (DType::Float(w), true) = (x.dtype(), x.dtype() == y.dtype()) {
1057 return with_float_width!(w, F => Some(crate::with_tolerance!(t, |t| tolerant::index_of(t, F::of(x)?, F::of(y)?))));
1058 }
1059 if let (View::Integers(x), View::Floats(Floats::F64(y))) = (x, y) { return with_ints!(x, |x| whole_needles(x, y, t).map(Ok)); }
1060 key_pairs!(x, y, |x, y| Ok(index_of(x, y)))
1061}
1062
1063pub fn view_matches(x: View<'_>, y: View<'_>, t: Tolerance) -> Option<bool> {
1065 use tolerant::{Tolerant, TolerantComplex};
1066 if x.len() != y.len() { return Some(false); }
1067 let whole = |v: View<'_>, i: usize| match v {
1069 View::Booleans(v) => i64::from(v[i]), View::Integers(v) => v.get(i), _ => unreachable!("Booleans or integers")
1070 };
1071 Some(match (x, y) {
1072 (View::Booleans(x), View::Booleans(y)) => x == y,
1073 (View::Integers(x), View::Integers(y)) => x == y,
1074 (x @ (View::Booleans(_) | View::Integers(_)), y @ (View::Booleans(_) | View::Integers(_))) => (0..x.len()).all(|i| whole(x, i) == whole(y, i)),
1075 (n @ (View::Booleans(_) | View::Integers(_)), View::Floats(f)) | (View::Floats(f), n @ (View::Booleans(_) | View::Integers(_))) => {
1076 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()))))
1077 }
1078 (View::Floats(x), View::Floats(y)) => {
1079 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())))))
1080 }
1081 (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)))),
1082 (View::Characters(x), View::Characters(y)) => x == y,
1083 _ => return None,
1084 })
1085}
1086
1087pub fn view_contains(x: View<'_>, y: View<'_>, cells: usize, t: Tolerance) -> Option<Result<bool, Error>> {
1089 if y.len() == 1 { if let Some(found) = view_index_of(x, y, t) { return Some(found.map(|found| (found[0] as usize) < cells)); } }
1090 for i in 0..cells {
1091 if i.is_multiple_of(crate::BETWEEN_CHECKS) { if let Err(e) = check() { return Some(Err(e)); } }
1092 if view_matches(x.slice(i * y.len()..(i + 1) * y.len()), y, t)? { return Some(Ok(true)); }
1093 }
1094 Some(Ok(false))
1095}
1096
1097pub fn member_layout(layout: &Layout, query: &Layout) -> Result<Option<Layout>, Error> {
1099 let rank = layout.shape().len();
1100 if rank == 0 { return Err(Error::new(Kind::Rank, "a rank-0 array is not a container")); }
1101 let cell = layout.axes(1..rank);
1102 Ok(cell.matches(query).then_some(cell))
1103}
1104
1105pub fn view_member(x: &Layout, xs: View<'_>, query: &Layout, ys: View<'_>, t: Tolerance) -> Result<bool, Error> {
1107 let Some(cell) = member_layout(x, query)? else { return Ok(false) };
1108 let Some(dtype) = xs.dtype().common(ys.dtype()) else { return Ok(false) };
1109 let haystack = ArrayView::compact(x, xs).cast(dtype)?;
1110 let axes: Vec<usize> = (0..cell.shape().len()).collect();
1111 let mapping = Mapping::new(query, &axes, &cell, ys.len());
1112 let needle = ArrayView::compact(query, ys).remapped(mapping, cell.count()).cast(dtype)?;
1113 view_contains(haystack.view(), needle.view(), x.shape()[0], t).ok_or_else(|| Error::new(Kind::Domain, "membership can't compare these kinds"))?
1114}
1115
1116fn whole_needles<K: Key>(x: &[K], y: &[f64], t: Tolerance) -> Option<Vec<u32>> {
1118 let whole: Vec<i64> = y.iter().map(|&v| tolerant::integer_of(t, v)).collect::<Option<_>>()?;
1119 let mut found = index_of(x, &whole.iter().filter_map(|&n| K::from_key(n)).collect::<Vec<_>>()).into_iter();
1120 Some(whole.iter().map(|&n| K::from_key(n).and_then(|_| found.next()).unwrap_or(x.len() as u32)).collect())
1121}
1122
1123#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1125pub enum Set { Union, Intersection, Difference }
1126
1127pub fn view_set(x: View<'_>, y: View<'_>, how: Set, t: Tolerance) -> Option<Result<Buffer, Error>> {
1129 if x.dtype() != y.dtype() { return None; }
1130 let mut joined = Buffer::with_capacity(x.dtype(), if how == Set::Union { x.len() + y.len() } else { 0 });
1131 let values = if how == Set::Union {
1132 joined.extend_from(x, 0..x.len());
1133 joined.extend_from(y, 0..y.len());
1134 joined.view()
1135 } else { x };
1136 let firsts = match view_firsts(values, t)? { Ok(firsts) => firsts, Err(e) => return Some(Err(e)) };
1137 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)) } };
1138 let positions: Vec<u32> = firsts.iter().enumerate().filter_map(|(i, &first)| {
1139 (first && matches.as_ref().is_none_or(|found| ((found[i] as usize) < y.len()) == (how == Set::Intersection))).then_some(i as u32)
1140 }).collect();
1141 Some(Ok(Buffer::gathered(values, &positions)))
1142}
1143
1144pub fn view_bins(boundaries: View<'_>, values: View<'_>) -> Option<Vec<u32>> {
1146 if let (DType::Float(w), true) = (boundaries.dtype(), boundaries.dtype() == values.dtype()) {
1147 return with_float_width!(w, F => Some(bins(&tolerant::order_keys(F::of(boundaries)?), &tolerant::order_keys(F::of(values)?))));
1148 }
1149 key_pairs!(boundaries, values, |x, y| bins(x, y))
1150}
1151
1152pub fn view_bins_below(boundaries: View<'_>, values: View<'_>) -> Option<Vec<u32>> {
1154 if boundaries.dtype() != values.dtype() { return None; }
1155 let (mut a, mut v) = (view_order_keys(boundaries)?, view_order_keys(values)?);
1156 let n = a.len() as u32;
1157 a.reverse();
1158 for k in a.iter_mut().chain(&mut v) { *k = !*k; }
1159 Some(bins(&a, &v).into_iter().map(|p| n - p).collect())
1160}
1161
1162macro_rules! ordered_view {
1164 ($view:expr, |$x:ident| $body:expr) => {
1165 match $view {
1166 View::Booleans($x) => Some($body),
1167 View::Characters($x) => Some($body),
1168 View::Integers(ints) => Some(with_ints!(ints, |$x| $body)),
1169 View::Floats(floats) => Some(with_floats!(floats, |$x| $body)),
1170 View::Complexes(_) => None,
1171 }
1172 };
1173 ($view:expr, |$x:ident, $key:ident| $body:expr) => {
1174 match $view {
1175 View::Booleans($x) => {
1176 let $key = SortKey::sort_key;
1177 Some($body)
1178 }
1179 View::Characters($x) => {
1180 let $key = SortKey::sort_key;
1181 Some($body)
1182 }
1183 View::Integers(ints) => Some(with_ints!(ints, |$x| {
1184 let $key = SortKey::sort_key;
1185 $body
1186 })),
1187 View::Floats(floats) => Some(with_floats!(floats, |$x| {
1188 let $key = Float::order_key;
1189 $body
1190 })),
1191 View::Complexes(_) => None,
1192 }
1193 };
1194}
1195
1196pub fn view_order_keys(x: View<'_>) -> Option<Vec<u64>> { ordered_view!(x, |x, key| x.iter().map(|&v| key(v).word()).collect()) }
1198
1199pub fn view_grade(x: View<'_>, down: bool) -> Option<Vec<u32>> {
1201 ordered_view!(x, |x, key| sort_rows(x.len(), vec![order_column(x, x.len(), 0, 1, down, key, |k| k)]))
1202}
1203
1204pub fn view_lexgrade<'a>(n: usize, columns: impl Iterator<Item = (View<'a>, usize, usize)>, down: bool) -> Option<Vec<u32>> {
1206 let columns: Vec<_> = columns.collect();
1207 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)])); }
1208 if n < RADIX_FROM {
1209 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())))
1211 .collect::<Option<Vec<_>>>()?;
1212 return Some(sort_rows(n, columns));
1213 }
1214 let (mut order, mut positions) = ((0..n as u32).collect::<Vec<_>>(), vec![0; n]);
1215 for (c, (view, start, step)) in columns.into_iter().rev().enumerate() {
1216 ordered_view!(view, |x, key| {
1217 let mut keys = order_column(x, n, start, step, down, key, |k| k);
1218 let mut spare = (vec![Default::default(); n], std::mem::take(&mut positions));
1219 radix_column(&mut keys, &mut order, &mut spare, c == 0);
1220 positions = spare.1;
1221 })?;
1222 }
1223 Some(order)
1224}
1225
1226pub fn view_lexgrade_along(columns: &[(View<'_>, usize, usize)], dims @ [outer, len, inner]: [usize; 3], down: bool) -> Result<Option<Vec<i64>>, Error> {
1228 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);
1229 if outer * inner == 1 { return Ok(lane(0, 0).map(|g| g.into_iter().map(i64::from).collect())); }
1230 if columns.iter().any(|(c, ..)| matches!(c, View::Complexes(_))) { return Ok(None); }
1231 Ok(Some(lanes::<i64, ()>(dims, len, |o, i, out, ()| {
1232 for (p, g) in out.iter_mut().zip(lane(o, i).expect("keys that order")) { *p = g.into(); }
1233 })?))
1234}
1235
1236fn order_column<T: Copy + Sync, K: RadixKey, O: Send>(
1238 x: &[T], n: usize, start: usize, step: usize, down: bool, key: impl Fn(T) -> K + Sync, convert: impl Fn(K) -> O + Sync) -> Vec<O> {
1239 let key = move |v| convert(if down { !key(v) } else { key(v) });
1240 if step != 1 { return (0..n).map(|i| key(x[start + i * step])).collect(); }
1241 let mut keys = Vec::with_capacity(n);
1242 crate::pairs(crate::Arg::Many(&x[start..start + n]), crate::Arg::One(()), n, &mut keys, move |v, ()| (key(v), true));
1243 keys
1244}
1245
1246fn spread_into<T: Element + Key>(x: &[T], out: &mut [T], down: bool) -> bool {
1248 if x.len() < 256 { return false; }
1249 let Some((min, counts)) = counts(x) else { return false };
1250 let mut at = 0;
1251 let mut write = |(k, &c): (usize, &usize)| {
1252 let v = T::from_key(min + k as i64).unwrap_or(T::FILL);
1253 if at + 4 <= out.len() { out[at..at + 4].fill(v); if c > 4 { out[at + 4..at + c].fill(v); } }
1254 else { out[at..at + c].fill(v); }
1255 at += c;
1256 };
1257 if down { counts.iter().enumerate().rev().for_each(&mut write); }
1258 else { counts.iter().enumerate().for_each(&mut write); }
1259 true
1260}
1261
1262const VECTORS: bool = cfg!(any(target_arch = "aarch64", target_feature = "avx2", target_feature = "simd128"));
1264const VECTOR_BYTES: usize = if cfg!(target_feature = "avx2") { 32 } else { 16 };
1266const HELD_ROWS: usize = if cfg!(target_arch = "aarch64") { 16 } else { 8 };
1268const NETWORK_BYTES: usize = 32 << 10;
1270const ANY_FILL_BYTES: usize = 4 << 10;
1272
1273pub(crate) trait Network: Sized { fn networked(_x: &[Self], _out: &mut [Self], _down: bool, _scratch: &mut Vec<Self>) -> bool { false } }
1275impl Network for half::f16 {}
1276impl Network for f32 {
1277 fn networked(x: &[f32], out: &mut [f32], down: bool, scratch: &mut Vec<f32>) -> bool { networked::<_, { VECTOR_BYTES / 4 }>(x, out, down, scratch) }
1278}
1279impl Network for f64 {
1280 fn networked(x: &[f64], out: &mut [f64], down: bool, scratch: &mut Vec<f64>) -> bool { networked::<_, { VECTOR_BYTES / 8 }>(x, out, down, scratch) }
1281}
1282
1283fn networked<F: Native, const W: usize>(x: &[F], out: &mut [F], down: bool, scratch: &mut Vec<F>) -> bool {
1285 let padded = x.len().next_power_of_two() * size_of::<F>();
1286 if !VECTORS || x.len() < W * HELD_ROWS || padded > NETWORK_BYTES || padded > ANY_FILL_BYTES && 4 * size_of_val(x) < 3 * padded { return false; }
1287 scratch.clear();
1288 let real = x.len() - x.iter().filter(|v| v.is_nan()).count();
1289 if real == x.len() { scratch.extend_from_slice(x) }
1290 else { scratch.extend(x.iter().copied().filter(|v| !v.is_nan())) }
1291 scratch.resize(real.next_power_of_two().max(W * HELD_ROWS), F::infinity());
1292 let rows = scratch.as_chunks_mut::<W>().0;
1293 network(rows);
1294 for (j, lane) in out[..real].chunks_mut(rows.len()).enumerate() { for (o, row) in lane.iter_mut().zip(&*rows) { *o = row[j]; } }
1295 if real < x.len() {
1296 let mut nans: Vec<F::OrderKey> = x.iter().filter(|v| v.is_nan()).map(|v| v.total_key()).collect();
1297 nans.sort_unstable();
1298 for (o, &k) in out[real..].iter_mut().zip(&nans) { *o = F::from_order_key(k); }
1299 }
1300 let sorted = &mut out[..real];
1301 let (zeros, positive) = (sorted.partition_point(|&v| v < F::zero()), sorted.partition_point(|&v| v <= F::zero()));
1302 let negative = zeros + sorted[zeros..positive].iter().filter(|v| v.is_sign_negative()).count();
1303 sorted[zeros..negative].fill(-F::zero());
1304 sorted[negative..positive].fill(F::zero());
1305 if down { out.reverse(); }
1306 true
1307}
1308
1309fn network<F: Native, const W: usize>(rows: &mut [[F; W]]) {
1311 for block in rows.as_chunks_mut::<HELD_ROWS>().0 { sorted_block(block) }
1312 let mut b = 2 * HELD_ROWS;
1313 while b <= rows.len() {
1314 for part in rows.chunks_exact_mut(b) { flip(part) }
1315 cleaned(rows, b / 4);
1316 b *= 2;
1317 }
1318 lanes_stage::<F, W, 1>(rows);
1319 if W >= 4 { lanes_stage::<F, W, 2>(rows) }
1320 if W >= 8 { lanes_stage::<F, W, 3>(rows) }
1321}
1322
1323#[inline(always)]
1325fn sorted_block<F: Native, const W: usize>(block: &mut [[F; W]; HELD_ROWS]) {
1326 let mut v = *block;
1327 held_step::<F, W, 2, true>(&mut v);
1328 held_step::<F, W, 4, true>(&mut v);
1329 half_cleans(&mut v, 2);
1330 held_step::<F, W, 8, true>(&mut v);
1331 half_cleans(&mut v, 4);
1332 if HELD_ROWS == 16 {
1333 held_step::<F, W, 16, true>(&mut v);
1334 half_cleans(&mut v, 8);
1335 }
1336 *block = v;
1337}
1338
1339fn cleaned<F: Native, const W: usize>(rows: &mut [[F; W]], mut d: usize) {
1341 while d >= HELD_ROWS { for part in rows.chunks_exact_mut(2 * d) { half_clean(part) } d /= 2 }
1342 for block in rows.as_chunks_mut::<HELD_ROWS>().0 {
1343 let mut v = *block;
1344 half_cleans(&mut v, HELD_ROWS);
1345 *block = v;
1346 }
1347}
1348
1349#[inline(always)]
1351fn half_cleans<F: Native, const W: usize>(v: &mut [[F; W]; HELD_ROWS], span: usize) {
1352 if span >= 16 { held_step::<F, W, 16, false>(v) }
1353 if span >= 8 { held_step::<F, W, 8, false>(v) }
1354 if span >= 4 { held_step::<F, W, 4, false>(v) }
1355 held_step::<F, W, 2, false>(v);
1356}
1357
1358#[inline(always)]
1360fn held_step<F: Native, const W: usize, const SPAN: usize, const FLIP: bool>(v: &mut [[F; W]; HELD_ROWS]) {
1361 for base in (0..HELD_ROWS).step_by(SPAN) {
1362 for r in 0..SPAN / 2 {
1363 let (i, k) = (base + r, if FLIP { base + SPAN - 1 - r } else { base + SPAN / 2 + r });
1364 (v[i], v[k]) = F::ordered_lanes(v[i], v[k]);
1365 }
1366 }
1367}
1368
1369fn lanes_stage<F: Native, const W: usize, const T: u32>(rows: &mut [[F; W]]) {
1371 let (m, h) = ((1 << T) - 1, 1 << (T - 1));
1372 let half = rows.len() / 2;
1373 let (lo, hi) = rows.split_at_mut(half);
1374 for (a, c) in lo.iter_mut().zip(hi.iter_mut().rev()) {
1375 let (s, l) = F::ordered_lanes(*a, std::array::from_fn(|j| c[j ^ m]));
1376 *a = std::array::from_fn(|j| if j & h == 0 { s[j] } else { l[j] });
1377 *c = std::array::from_fn(|j| if j & h == 0 { s[j ^ m] } else { l[j ^ m] });
1378 }
1379 for u in (0..T - 1).rev() {
1380 let g = 1 << u;
1381 for v in rows.iter_mut() {
1382 let (s, l) = F::ordered_lanes::<W>(std::array::from_fn(|j| v[j & !g]), std::array::from_fn(|j| v[j | g]));
1383 *v = std::array::from_fn(|j| if j & g == 0 { s[j] } else { l[j] });
1384 }
1385 }
1386 cleaned(rows, half);
1387}
1388
1389#[inline(always)]
1391fn flip<F: Native, const W: usize>(part: &mut [[F; W]]) {
1392 let (lo, hi) = part.split_at_mut(part.len() / 2);
1393 for (a, c) in lo.iter_mut().zip(hi.iter_mut().rev()) { (*a, *c) = F::ordered_lanes(*a, *c) }
1394}
1395
1396#[inline(always)]
1398fn half_clean<F: Native, const W: usize>(part: &mut [[F; W]]) {
1399 let (lo, hi) = part.split_at_mut(part.len() / 2);
1400 for (a, c) in lo.iter_mut().zip(hi) { (*a, *c) = F::ordered_lanes(*a, *c) }
1401}
1402
1403pub fn view_sort(x: View<'_>, down: bool) -> Result<Option<Buffer>, Error> { view_sort_along(x, [1, x.len(), 1], down) }
1405
1406pub fn view_sort_along(x: View<'_>, dims: [usize; 3], down: bool) -> Result<Option<Buffer>, Error> {
1408 fn run<T: SortKey>(x: &[T], dims @ [outer, len, inner]: [usize; 3], down: bool) -> Result<Buffer, Error> {
1409 let sort = |lane: &[T], out: &mut [T], (keys, spare, scratch): &mut (Vec<T::Key>, Vec<T::Key>, Vec<T>)| {
1410 if T::sorted_into(lane, out, down, scratch) { return; }
1411 T::keys_into(lane, keys, down);
1412 let n = keys.len();
1413 if n < value_radix_from(<T::Key as RadixKey>::BYTES) { keys.sort_unstable(); } else {
1414 spare.resize(n, Default::default());
1415 radix(keys, &mut vec![(); n], (spare, &mut vec![(); n]));
1416 }
1417 T::write_back(keys, out, down);
1418 };
1419 Ok(T::buffer(if outer * inner == 1 {
1420 let mut out = vec![T::FILL; len];
1421 sort(x, &mut out, &mut Default::default());
1422 out
1423 } else { map_lanes::<_, _, (Vec<T::Key>, Vec<T::Key>, Vec<T>)>(x, dims, len, |lane, out, work| sort(lane, out, work))? }))
1424 }
1425 ordered_view!(x, |x| run(x, dims, down)).transpose()
1426}
1427
1428pub fn select<T>(x: &mut [T], ranks: &[usize], compare: impl Fn(&T, &T) -> Ordering + Copy) {
1430 fn within<T>(x: &mut [T], ranks: &[usize], base: usize, compare: impl Fn(&T, &T) -> Ordering + Copy) {
1431 if ranks.is_empty() { return; }
1432 let middle = ranks.len() / 2;
1433 let rank = ranks[middle] - base;
1434 let (before, _, after) = x.select_nth_unstable_by(rank, compare);
1435 within(before, &ranks[..middle], base, compare);
1436 within(after, &ranks[middle + 1..], base + rank + 1, compare);
1437 }
1438 within(x, ranks, 0, compare);
1439}
1440
1441pub fn view_argpartition(x: View<'_>, rank: usize) -> Option<Vec<u32>> {
1443 let keys = view_order_keys(x)?;
1444 let mut out: Vec<u32> = (0..x.len() as u32).collect();
1445 select(&mut out, &[rank], |&a, &b| keys[a as usize].cmp(&keys[b as usize]));
1446 Some(out)
1447}
1448
1449pub fn view_partition_along(x: View<'_>, dims: [usize; 3], rank: usize, indices: bool) -> Result<Option<Buffer>, Error> {
1451 fn run<T: SortKey>(x: &[T], dims: [usize; 3], rank: usize, indices: bool) -> Result<Buffer, Error> {
1452 Ok(if indices {
1453 i64::buffer(map_lanes::<_, _, Vec<T::Key>>(x, dims, dims[1], |lane, out, keys| {
1454 T::keys_into(lane, keys, false);
1455 for (i, p) in out.iter_mut().enumerate() { *p = i as i64; }
1456 select(out, &[rank], |a, b| keys[*a as usize].cmp(&keys[*b as usize]));
1457 })?)
1458 } else {
1459 T::buffer(map_lanes::<_, _, Vec<T::Key>>(x, dims, dims[1], |lane, out, keys| {
1460 T::keys_into(lane, keys, false);
1461 select(keys, &[rank], |a, b| a.cmp(b));
1462 T::write_back(keys, out, false);
1463 })?)
1464 })
1465 }
1466 ordered_view!(x, |x| run(x, dims, rank, indices)).transpose()
1467}
1468
1469const BLOCK: usize = 64;
1471
1472fn ranked<T: Copy + PartialOrd, const BACK: bool>(x: &[T], key: impl Fn(T) -> u64, down: bool, range: std::ops::Range<usize>) -> Vec<u32> {
1474 if range.is_empty() { return Vec::new(); }
1475 let (n, m, flip) = (x.len(), range.end, if down { u64::MAX } else { 0 });
1476 let at = |j: usize| if BACK { n - 1 - j } else { j };
1477 let place = |t: u32| (if BACK { !t } else { t }) as usize;
1478 let pair = |p: usize| (key(x[p]) ^ flip, if BACK { !(p as u32) } else { p as u32 });
1479 let mut kept: Vec<(u64, u32)> = (0..m).map(|j| pair(at(j))).collect();
1480 if m * BLOCK < n {
1482 let mut last = *kept.iter().max().expect("range.end is at least 1");
1484 let (room, mut compacted) = (2 * m, m);
1485 let mut j = m;
1486 'scan: while j < n {
1487 let len = BLOCK.min(n - j);
1488 let start = if BACK { n - j - len } else { j };
1489 let limit = x[place(last.1)];
1490 let below = <&[T; BLOCK]>::try_from(&x[start..start + len]).map_or(true, |block| {
1491 if down { block.iter().fold(false, |below, &v| below | !(v <= limit)) } else { block.iter().fold(false, |below, &v| below | !(v >= limit)) }
1492 });
1493 if below {
1494 for i in j..j + len {
1495 let p = pair(at(i));
1496 if p.0 >= last.0 { continue; }
1497 kept.push(p);
1498 if kept.len() < room { continue; }
1499 if i - compacted < room && i >= 4 * room { kept.extend((i + 1..n).map(|i| pair(at(i)))); break 'scan; }
1501 kept.select_nth_unstable(m - 1);
1502 kept.truncate(m);
1503 last = kept[m - 1];
1504 compacted = i;
1505 }
1506 }
1507 j += len;
1508 }
1509 } else { kept.extend((m..n).map(|j| pair(at(j)))); }
1510 let mut ranks = vec![range.start, range.end - 1];
1511 ranks.dedup();
1512 select(&mut kept, &ranks, Ord::cmp);
1513 let run = &mut kept[range];
1514 run.sort_unstable();
1515 run.iter().map(|&(_, t)| place(t) as u32).collect()
1516}
1517
1518pub fn view_selected(x: View<'_>, range: std::ops::Range<usize>, down: bool) -> Option<Vec<u32>> {
1520 let n = x.len();
1521 ordered_view!(x, |x, key| if range.end <= n - range.start {
1522 ranked::<_, false>(x, |v| key(v).word(), down, range)
1523 } else {
1524 let mut run = ranked::<_, true>(x, |v| key(v).word(), !down, n - range.end..n - range.start);
1525 run.reverse();
1526 run
1527 })
1528}
1529
1530pub fn view_argextremes(x: View<'_>, [outer, len, inner]: [usize; 3], down: bool) -> Option<Vec<i64>> {
1532 fn extremes<T: Copy + Sync, K: Ord>(x: &[T], [outer, len, inner]: [usize; 3], down: bool, key: impl Fn(T) -> K + Sync) -> Vec<i64> {
1533 let mut out = Vec::with_capacity(outer * inner);
1534 crate::parallel::extended(&mut out, outer * inner, (crate::CHUNK / len.max(1)).max(1), &|range, slots| {
1535 unsafe {
1536 slots.fill(range.len(), |slots| {
1537 for (lane, slot) in range.zip(slots) {
1538 let start = lane / inner * len * inner + lane % inner;
1539 let (mut best, mut at) = (key(x[start]), 0);
1540 for j in 1..len {
1541 let k = key(x[start + j * inner]);
1542 if if down { k > best } else { k < best } {
1543 best = k;
1544 at = j;
1545 }
1546 }
1547 slot.write(at as i64);
1548 }
1549 true
1550 })
1551 };
1552 true
1553 });
1554 out
1555 }
1556 if len == 0 { return None; }
1557 ordered_view!(x, |x, key| extremes(x, [outer, len, inner], down, key))
1558}
1559
1560struct Quantiles { points: Vec<(usize, usize, f64)>, ranks: Vec<usize> }
1562
1563impl Quantiles {
1564 fn new(len: usize, probabilities: &[f64]) -> Result<Self, Error> {
1566 if !probabilities.iter().all(|q| (0.0..=1.0).contains(q)) {
1567 return Err(Error::new(Kind::Domain, "quantile probabilities must be between zero and one"));
1568 }
1569 let points: Vec<_> = probabilities.iter().map(|&q| {
1570 let p = q * len.saturating_sub(1) as f64;
1571 (p.floor() as usize, p.ceil() as usize, p.fract())
1572 }).collect();
1573 let mut ranks: Vec<_> = points.iter().flat_map(|&(lo, hi, _)| [lo, hi]).collect();
1574 ranks.sort_unstable();
1575 ranks.dedup();
1576 Ok(Self { points, ranks })
1577 }
1578
1579 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) {
1580 if x.is_empty() { out.fill(F::nan()); return; }
1581 select(x, &self.ranks, compare);
1582 for (out, &(lo, hi, weight)) in out.iter_mut().zip(&self.points) {
1583 let a = read(x[lo]);
1584 *out = F::from_working(if lo == hi { a }
1585 else { let weight = F::Working::narrow(weight); a * (F::Working::narrow(1.0) - weight) + read(x[hi]) * weight });
1586 }
1587 }
1588}
1589
1590pub fn quantiles<F: Float>(x: &mut [F], probabilities: &[f64]) -> Result<Vec<F>, Error> {
1592 let mut out = vec![F::nan(); probabilities.len()];
1593 Quantiles::new(x.len(), probabilities)?.apply(x, &mut out, |a, b| a.order_key().cmp(&b.order_key()), Float::working);
1594 Ok(out)
1595}
1596
1597pub fn view_quantiles_along(x: View<'_>, dims @ [outer, _, inner]: [usize; 3], probabilities: &[f64]) -> Result<Option<Buffer>, Error> {
1599 use crate::buffer::Source;
1600 fn run<T: SortKey, F: Element + Float + Default>(
1601 x: &[T], dims: [usize; 3], plan: &Quantiles, read: impl Fn(T) -> F::Working + Sync) -> Result<Vec<F>, Error> {
1602 map_lanes::<_, _, Vec<T::Key>>(x, dims, plan.points.len(), |lane, out, keys| {
1603 T::keys_into(lane, keys, false);
1604 plan.apply(keys, out, |a, b| a.cmp(b), |k| read(T::from_sort_key(k)));
1605 })
1606 }
1607 fn floats<F: Element + Float + Default + Network>(x: &[F], dims: [usize; 3], plan: &Quantiles) -> Result<Buffer, Error> {
1608 Ok(F::buffer(run::<_, F>(x, dims, plan, Float::working)?))
1609 }
1610 let plan = Quantiles::new(dims[1], probabilities)?;
1611 let items = match x {
1612 View::Booleans(x) => f64::buffer(run::<_, f64>(x, dims, &plan, Source::<f64>::read)?),
1613 View::Integers(ints) => with_ints!(ints, |x| f64::buffer(run::<_, f64>(x, dims, &plan, Source::<f64>::read)?)),
1614 View::Floats(x) => with_floats!(x, |x| floats(x, dims, &plan)?), _ => return Ok(None)
1615 };
1616 let count = probabilities.len();
1617 Ok(Some(if count > 1 && outer > 1 {
1618 let plan = crate::structure::transpose(&crate::layout::Layout::from([outer, count, inner]), &[1, 0, 2]).expect("valid quantile axes");
1619 Buffer::walked(items.view(), 0, &plan.steps, items.len())
1620 } else { items }))
1621}
1622
1623fn map_lanes<T: Copy + Sync, U: Element + Default, S: Default>(
1625 x: &[T], dims @ [_, len, inner]: [usize; 3], width: usize, map: impl Fn(&[T], &mut [U], &mut S) + Sync) -> Result<Vec<U>, Error> {
1626 lanes::<U, (Vec<T>, S)>(dims, width, |o, i, out, (input, work)| {
1627 let input = if inner == 1 { &x[o * len..(o + 1) * len] } else {
1628 input.clear();
1629 input.extend((0..len).map(|j| x[(o * len + j) * inner + i]));
1630 &input[..]
1631 };
1632 map(input, out, work);
1633 })
1634}
1635
1636fn lanes<U: Element + Default, S: Default>(
1638 [outer, _, inner]: [usize; 3], width: usize, fill: impl Fn(usize, usize, &mut [U], &mut S) + Sync) -> Result<Vec<U>, Error> {
1639 if outer * inner * width == 0 { return Ok(Vec::new()); }
1640 let ((out,), _) = crate::parallel::stack::<(Vec<U>,), S>(outer * inner, [width], &|lane, (out,), work| {
1641 fill(lane / inner, lane % inner, out, work);
1642 true
1643 })?;
1644 if inner == 1 || width == 1 { return Ok(out); }
1645 let mut transposed = Vec::with_capacity(out.len());
1646 transposed.resize(out.len(), U::FILL);
1647 for o in 0..outer {
1648 crate::linalg::transpose_into(
1649 &out[o * inner * width..][..inner * width], width, &mut transposed[o * inner * width..][..inner * width], inner, [inner, width], |x| x);
1650 }
1651 Ok(transposed)
1652}