1use std::collections::{HashMap, HashSet};
8use std::sync::Arc;
9
10use crate::array::{Array, Buf, Data, Layout};
11use crate::complex::{self as cx, Cx};
12use crate::dtype::DType;
13use crate::error::{Error, ErrorKind, Result, Span};
14use crate::exact::{self, Ext, Rat};
15use crate::fmt::FmtOpts;
16use crate::frontend::{ComplexOrder, Rules};
17use crate::par;
18use crate::simd::multiversioned;
19
20pub const RANK_INF: i64 = i64::MAX;
22
23#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub enum Agreement {
27 LeadingPrefix,
29 ExactOrScalar,
31}
32
33#[derive(Clone, Copy, Debug, PartialEq)]
43pub struct Tol {
44 pub ct: f64,
46 pub by_smaller: bool,
48}
49
50impl Tol {
51 pub const EXACT: Tol = Tol { ct: 0.0, by_smaller: true };
53 pub const J: Tol = Tol { ct: 5.684_341_886_080_802e-14, by_smaller: true };
55 pub const APL: Tol = Tol { ct: 1e-13, by_smaller: false };
57
58 #[inline(always)]
60 pub fn eq(self, a: f64, b: f64) -> bool {
61 if a == b {
62 return true;
63 }
64 let s = if self.by_smaller {
67 a.abs().min(b.abs())
68 } else {
69 a.abs().max(b.abs())
70 };
71 (a - b).abs() < self.ct * s
72 }
73
74 #[inline(always)]
80 pub fn is_j(self) -> bool {
81 self.by_smaller
82 }
83
84 #[inline(always)]
90 pub fn is_zero(self, y: f64) -> bool {
91 self.is_j() && y.abs() < self.ct
92 }
93
94 #[inline(always)]
96 pub fn lt(self, a: f64, b: f64) -> bool {
97 a < b && !self.eq(a, b)
98 }
99
100 #[inline(always)]
102 pub fn le(self, a: f64, b: f64) -> bool {
103 a <= b || self.eq(a, b)
104 }
105
106 #[inline]
110 pub fn eq_cx(self, a: Cx, b: Cx) -> bool {
111 if a == b {
112 return true;
113 }
114 let (ma, mb) = (cx::abs(a), cx::abs(b));
115 let s = if self.by_smaller { ma.min(mb) } else { ma.max(mb) };
116 cx::abs(cx::sub(a, b)) < self.ct * s
117 }
118
119 #[inline(always)]
122 pub fn floor(self, y: f64) -> f64 {
123 let c = y.ceil();
124 if self.eq(y, c) {
125 c
126 } else {
127 y.floor()
128 }
129 }
130
131 #[inline(always)]
134 pub fn ceil(self, y: f64) -> f64 {
135 let f = y.floor();
136 if self.eq(y, f) {
137 f
138 } else {
139 y.ceil()
140 }
141 }
142}
143
144#[derive(Clone, Copy, Debug)]
149pub struct EvalCfg {
150 pub agreement: Agreement,
151 pub fmt: FmtOpts,
152 pub tol: Tol,
155 pub rules: Rules,
158}
159
160impl EvalCfg {
161 pub(crate) fn pure<R>(self, f: impl FnOnce(&mut Ctx<'_>) -> R) -> R {
166 let mut sink = |_: &str| debug_assert!(false, "a pure verb wrote to the output sink");
167 let mut env = Env::new(Vec::new());
168 f(&mut Ctx { cfg: self, out: &mut sink, inp: None, env: &mut env, device: None })
169 }
170}
171
172pub const RECURSION_LIMIT: usize = 64;
181
182pub struct Env {
190 globals: HashMap<String, Array>,
191 frames: Vec<HashMap<String, Array>>,
192 running: Vec<std::sync::Arc<crate::ir::ExplicitDef>>,
195 verbs: HashMap<String, Verb>,
196 args: Vec<Array>,
197}
198
199impl Env {
200 pub fn new(args: Vec<Array>) -> Env {
201 Env {
202 globals: HashMap::new(),
203 frames: Vec::new(),
204 running: Vec::new(),
205 verbs: HashMap::new(),
206 args,
207 }
208 }
209
210 pub fn get(&self, name: &str) -> Option<Array> {
211 if let Some(frame) = self.frames.last() && let Some(v) = frame.get(name) {
212 return Some(v.clone());
213 }
214 self.globals.get(name).cloned()
215 }
216
217 pub fn assign(&mut self, name: String, value: Array, scope: crate::ir::Scope) {
218 if scope == crate::ir::Scope::LocalDefault && self.get(&name).is_some() {
219 return;
220 }
221 let target = match (scope, self.frames.last_mut()) {
222 (crate::ir::Scope::Local | crate::ir::Scope::LocalDefault, Some(frame)) => frame,
223 _ => &mut self.globals,
224 };
225 target.insert(name, value);
226 }
227
228 pub fn define(&mut self, name: String, verb: Verb) {
229 self.verbs.insert(name, verb);
230 }
231
232 pub fn undefine(&mut self, name: &str) {
233 self.verbs.remove(name);
234 }
235
236 pub fn verb(&self, name: &str) -> Option<&Verb> {
237 self.verbs.get(name)
238 }
239
240 pub fn arg(&self, i: usize) -> Result<Array> {
241 self.args
242 .get(i)
243 .cloned()
244 .ok_or_else(|| Error::internal("a parameter was read where none is bound"))
245 }
246
247 pub fn enter(
249 &mut self,
250 frame: HashMap<String, Array>,
251 def: std::sync::Arc<crate::ir::ExplicitDef>,
252 span: Span,
253 ) -> Result<()> {
254 if self.frames.len() >= RECURSION_LIMIT {
255 return Err(Error::new(
256 ErrorKind::Domain,
257 format!("explicit definitions called each other more than {RECURSION_LIMIT} deep"),
258 Some(span),
259 )
260 .note("a definition that recurses needs a case that stops"));
261 }
262 self.frames.push(frame);
263 self.running.push(def);
264 Ok(())
265 }
266
267 pub fn leave(&mut self) -> HashMap<String, Array> {
269 self.running.pop();
270 self.frames.pop().unwrap_or_default()
271 }
272
273 pub fn current_def(&self) -> Option<std::sync::Arc<crate::ir::ExplicitDef>> {
275 self.running.last().cloned()
276 }
277}
278
279pub type InputFn<'a> = Option<&'a mut dyn FnMut() -> Option<String>>;
287
288pub fn reborrow_input<'s, 'a: 's>(inp: &'s mut InputFn<'a>) -> InputFn<'s> {
292 match inp {
293 Some(f) => Some(&mut **f),
294 None => None,
295 }
296}
297
298pub struct Ctx<'a> {
300 pub cfg: EvalCfg,
301 pub out: &'a mut dyn FnMut(&str),
304 pub inp: InputFn<'a>,
308 pub env: &'a mut Env,
310 pub device: Option<&'a crate::device::Device>,
313}
314
315const MAX_NESTING: usize = 400;
324
325thread_local! {
326 static NESTING: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
327}
328
329pub(crate) fn check_nesting(depth: usize, span: Span) -> Result<()> {
331 if depth > MAX_NESTING {
332 return Err(Error::new(
333 ErrorKind::Limit,
334 format!("this program nests more than {MAX_NESTING} applications deep"),
335 Some(span),
336 ));
337 }
338 Ok(())
339}
340
341pub(crate) struct Nesting;
343
344impl Nesting {
345 pub(crate) fn enter(span: Span) -> Result<Nesting> {
347 let depth = NESTING.with(|c| {
348 let d = c.get() + 1;
349 c.set(d);
350 d
351 });
352 if depth > MAX_NESTING {
353 NESTING.with(|c| c.set(c.get() - 1));
354 return Err(Error::new(
355 ErrorKind::Limit,
356 format!("this program nests more than {MAX_NESTING} applications deep"),
357 Some(span),
358 ));
359 }
360 Ok(Nesting)
361 }
362}
363
364impl Drop for Nesting {
365 fn drop(&mut self) {
366 NESTING.with(|c| c.set(c.get().saturating_sub(1)));
367 }
368}
369
370impl Ctx<'_> {
371 fn with_tol<R>(&mut self, tol: Tol, f: impl FnOnce(&mut Ctx<'_>) -> R) -> R {
373 let cfg = EvalCfg { tol, ..self.cfg };
374 f(&mut Ctx {
375 cfg,
376 out: &mut *self.out,
377 inp: reborrow_input(&mut self.inp),
378 env: &mut *self.env,
379 device: self.device,
380 })
381 }
382
383 pub(crate) fn read_line(&mut self, span: Span) -> Result<String> {
389 let Some(read) = self.inp.as_deref_mut() else {
390 return Err(Error::new(
391 ErrorKind::Value,
392 "this expression reads input, and this run has no input source attached",
393 Some(span),
394 )
395 .note("attach one with Program::run_io (Rust), input= (Python), or jay_run_io (C)"));
396 };
397 read().ok_or_else(|| {
398 Error::new(ErrorKind::Value, "the input has ended: there is no line to read", Some(span))
399 })
400 }
401}
402
403#[derive(Clone, Copy, Debug, PartialEq, Eq)]
405pub enum ScalarMonad {
406 Conj,
408 Neg,
409 Signum,
410 Recip,
411 Sqrt,
412 Exp,
413 Abs,
414 Floor,
415 Ceil,
416 Not,
418 OneMinus,
420 Inc,
422 Dec,
424 Double,
426 Halve,
428 Square,
430 Ln,
432 Pi,
434 Factorial,
437 Imaginary,
439 Polar,
441}
442
443#[derive(Clone, Copy, Debug, PartialEq, Eq)]
445pub enum ScalarDyad {
446 Add,
447 Sub,
448 Mul,
449 DivJ,
451 DivApl,
453 Min,
454 Max,
455 Pow,
456 Residue,
458 Eq,
459 Ne,
460 Lt,
461 Le,
462 Gt,
463 Ge,
464 Lcm,
466 Gcd,
468 Log,
470 Root,
472 Circle,
476 Binomial,
479 MakeComplex,
481 PolarBy,
483}
484
485#[derive(Clone, Copy, Debug, PartialEq, Eq)]
487pub enum Enclose {
488 Always,
490 ExceptSimpleScalar,
492}
493
494#[derive(Clone, Copy, Debug, PartialEq, Eq)]
496pub enum MonadOp {
497 Scalar(ScalarMonad),
498 ShapeOf,
500 Tally,
502 Ravel,
504 TransposeAxes,
506 Catalogue,
509 AtomicRep,
513 RazeIn,
516 Head,
518 Behead,
520 Tail,
522 Curtail,
524 Reverse,
526 Nub,
528 GradeUp { origin: i64 },
530 GradeDown { origin: i64 },
532 IotaJ,
534 IotaApl { origin: i64 },
536 Echo,
538 ReadStream,
542 TypeCode,
544 Same,
546 Format,
550 DecodeBits,
552 EncodeBits,
556 Itemize,
558 TableOf,
561 Enclose(Enclose),
563 Open,
567 Raze,
569 First,
572 Enlist,
574 Depth,
577 Indices { origin: i64, boxed_coords: bool },
581 Steps,
583 ToExact,
586 NthPrime,
588 PrimeFactors,
590 MatrixInverse,
592 Roll { origin: i64, fixed: bool, float_at_zero: bool },
597 ComplexParts { polar: bool },
601 SelfClassify,
604 NubSieve,
607 Unicode { pass_chars: bool },
611 Symbols,
615 Words,
617 Nest,
620 LevelOf,
623 MapPaths,
626 PolyRoots,
630 PolyDeriv,
633 AnagramIndex,
635 CycleForm,
638 Split,
641 Execute { apl: bool },
646 NotYet(&'static str),
648 None,
650}
651
652#[derive(Clone, Copy, Debug, PartialEq, Eq)]
654pub enum DyadOp {
655 Scalar(ScalarDyad),
656 Reshape,
659 Take,
661 Drop,
663 Right,
665 Left,
667 Rotate,
669 AppendLeading,
671 AppendLast,
673 IndexOf { origin: i64 },
676 MemberJ,
678 MemberApl,
680 From,
683 Match,
685 NotMatch,
687 GradeSelect { down: bool },
689 Copy,
692 Decode,
695 Encode,
699 CollateGrade { down: bool, origin: i64 },
702 TransposeJ,
705 TransposeApl,
708 DecodeApl,
711 EncodeApl,
714 Laminate,
716 Link,
719 Strand,
721 IntervalIndex { offset: i64, closed: bool },
725 IndexOfLast { origin: i64 },
727 MatrixDivide,
729 PartitionEnclose,
732 Squad { origin: i64 },
734 SelectAxis { axis: usize, rank: usize, origin: i64 },
738 Fetch,
740 PolyEval,
743 PolyIntegral,
746 FormatSpec,
749 FormatSpecJ,
754 ParseNumbers,
757 SequentialMachine,
759 TruthTable(u8),
763 ExactForm,
767 Deal { origin: i64, fixed: bool },
770 Boolean(BoolDyad),
773 Less,
775 Union,
777 Intersect,
779 AnagramFrom,
782 Permute,
785 FindSeq,
788 UnicodeForm,
791 PrimeMeta,
795 PrimeExponents,
798 SymbolForm,
801 Pick { origin: i64 },
803 Expand,
806 WriteStream,
810 NotYet(&'static str),
811 None,
812}
813
814#[derive(Clone, Copy, Debug, PartialEq, Eq)]
816pub enum BoolDyad {
817 Nor,
819 Nand,
821}
822
823#[derive(Clone, Copy, Debug, PartialEq, Eq)]
826pub struct Prim {
827 pub name: &'static str,
828 pub monad: MonadOp,
829 pub dyad: DyadOp,
830 pub ranks: [i64; 3],
831}
832
833#[derive(Clone, Copy, Debug, PartialEq, Eq)]
837pub enum WindowKind {
838 Prefix,
841 Suffix,
844 Scan,
848}
849
850#[derive(Clone, Debug, PartialEq, Eq)]
852pub enum Power {
853 Times(u64),
855 Converge,
857 Each(Vec<u64>),
860 ConvergeTrace,
862}
863
864const CONVERGE_LIMIT: usize = 1 << 20;
866
867pub type MemoCache = Arc<std::sync::Mutex<HashMap<Vec<u64>, Array>>>;
871
872#[derive(Clone, Debug)]
876pub enum Verb {
877 Prim(Prim),
878 Rank(Box<Verb>, [i64; 3]),
880 Reduce(Box<Verb>),
882 Windowed(Box<Verb>, WindowKind),
885 Commute(Box<Verb>),
887 PowerN(Box<Verb>, Power),
889 Fork(Box<Verb>, Box<Verb>, Box<Verb>),
891 NounFork(Array, Box<Verb>, Box<Verb>),
893 Hook(Box<Verb>, Box<Verb>),
895 Atop(Box<Verb>, Box<Verb>),
897 Compose(Box<Verb>, Box<Verb>),
900 BondLeft(Array, Box<Verb>),
903 BondRight(Box<Verb>, Array),
905 Each(Box<Verb>, Enclose),
908 Fit(Box<Verb>, f64),
910 Amend(Array),
912 AmendVerb(Box<Verb>),
915 ShiftFill(Array),
918 Memo(Box<Verb>, MemoCache),
922 Level { u: Box<Verb>, level: i64, spread: bool },
926 Characteristics(Box<Verb>),
929 Before(Box<Verb>, Box<Verb>),
932 UserDerived { def: Box<Verb>, alpha: Box<Verb>, omega: Option<Box<Verb>> },
936 KeyPairs(Box<Verb>),
941 Key(Box<Verb>),
944 Cut(Box<Verb>, i64),
946 PowerV(Box<Verb>, Box<Verb>),
948 PowerUntil(Box<Verb>, Box<Verb>),
950 AlongAxis(Box<Verb>, usize),
954 Explicit(Arc<crate::ir::ExplicitDef>),
958 SelfRef,
961 Named(String),
964 WithObverse(Box<Verb>, Box<Verb>),
967 Agenda(Vec<Verb>, Box<Verb>),
970 Adverse(Box<Verb>, Box<Verb>),
974 Hypergeometric { num: Vec<crate::complex::Cx>, den: Vec<crate::complex::Cx> },
977 Beside(Box<Verb>, Box<Verb>),
981 Stencil(Box<Verb>, Vec<i64>),
985 Evoke(Vec<Verb>, i64),
991 InnerProduct { u: Box<Verb>, v: Box<Verb>, apl: bool },
998}
999
1000impl Verb {
1001 pub fn ranks(&self) -> [i64; 3] {
1003 match self {
1004 Verb::Prim(p) => p.ranks,
1005 Verb::Rank(_, r) => *r,
1006 Verb::Windowed(_, WindowKind::Prefix) => [RANK_INF, 0, RANK_INF],
1009 Verb::Each(..) => [0, 0, 0],
1010 Verb::Fit(v, _) => v.ranks(),
1011 Verb::Amend(_)
1014 | Verb::AmendVerb(_)
1015 | Verb::ShiftFill(_)
1016 | Verb::Level { .. }
1017 | Verb::Characteristics(_)
1018 | Verb::UserDerived { .. }
1019 | Verb::KeyPairs(_)
1020 | Verb::Key(_)
1021 | Verb::Cut(..)
1022 | Verb::PowerV(..)
1023 | Verb::PowerUntil(..)
1024 | Verb::AlongAxis(..) => [RANK_INF, RANK_INF, RANK_INF],
1025 Verb::Memo(v, _) => v.ranks(),
1026 Verb::WithObverse(v, _) | Verb::Adverse(v, _) => v.ranks(),
1027 Verb::Beside(..) => [RANK_INF, RANK_INF, RANK_INF],
1028 Verb::Hypergeometric { .. } => [0, 0, 0],
1030 Verb::InnerProduct { .. } => [2, RANK_INF, RANK_INF],
1033 _ => [RANK_INF, RANK_INF, RANK_INF],
1034 }
1035 }
1036
1037 pub fn name(&self) -> String {
1039 match self {
1040 Verb::Prim(p) => p.name.to_string(),
1041 Verb::Rank(v, r) => format!("{}\"{}", v.name(), rank_str(*r)),
1042 Verb::Reduce(v) => format!("{}/", v.name()),
1043 Verb::Windowed(v, WindowKind::Suffix) => format!("{}\\.", v.name()),
1044 Verb::Windowed(v, _) => format!("{}\\", v.name()),
1045 Verb::Commute(v) => format!("{}~", v.name()),
1046 Verb::PowerN(v, Power::Converge) => format!("{}^:_", v.name()),
1047 Verb::PowerN(v, Power::Times(n)) => format!("{}^:{n}", v.name()),
1048 Verb::PowerN(v, Power::Each(_)) => format!("{}^:n", v.name()),
1049 Verb::PowerN(v, Power::ConvergeTrace) => format!("{}^:a:", v.name()),
1050 Verb::Fork(f, g, h) => format!("({} {} {})", f.name(), g.name(), h.name()),
1051 Verb::NounFork(_, g, h) => format!("(n {} {})", g.name(), h.name()),
1052 Verb::Hook(f, g) => format!("({} {})", f.name(), g.name()),
1053 Verb::Atop(f, g) => format!("({}@:{})", f.name(), g.name()),
1054 Verb::Compose(f, g) => format!("({}&:{})", f.name(), g.name()),
1055 Verb::BondLeft(_, v) => format!("(n&{})", v.name()),
1056 Verb::BondRight(v, _) => format!("({}&n)", v.name()),
1057 Verb::Each(v, Enclose::Always) => format!("({}&.>)", v.name()),
1058 Verb::Each(v, _) => format!("({}¨)", v.name()),
1059 Verb::Fit(v, n) => format!("{}!.{n}", v.name()),
1060 Verb::Amend(_) => "(m})".to_string(),
1061 Verb::AmendVerb(v) => format!("({}}})", v.name()),
1062 Verb::ShiftFill(_) => "|.!.n".to_string(),
1063 Verb::Characteristics(v) => format!("{} b.", v.name()),
1064 Verb::Before(f, g) => format!("({}⍛{})", f.name(), g.name()),
1065 Verb::KeyPairs(v) => format!("{}⌸", v.name()),
1066 Verb::UserDerived { def, alpha, omega } => match omega {
1067 Some(g) => format!("({} {} {})", alpha.name(), def.name(), g.name()),
1068 None => format!("({} {})", alpha.name(), def.name()),
1069 },
1070 Verb::Memo(v, _) => format!("{} M.", v.name()),
1071 Verb::Level { u, level, spread } => {
1072 format!("{} {} {level}", u.name(), if *spread { "S:" } else { "L:" })
1073 }
1074 Verb::Key(v) => format!("{}/.", v.name()),
1075 Verb::Cut(v, n) => format!("{};.{n}", v.name()),
1076 Verb::PowerV(v, w) => format!("{}^:{}", v.name(), w.name()),
1077 Verb::PowerUntil(v, w) => format!("{}⍣{}", v.name(), w.name()),
1078 Verb::AlongAxis(v, k) => format!("{}[{k}]", v.name()),
1079 Verb::Explicit(d) => d.name.clone(),
1080 Verb::SelfRef => "$:".to_string(),
1081 Verb::Named(n) => n.clone(),
1082 Verb::WithObverse(v, w) => format!("({}:.{})", v.name(), w.name()),
1083 Verb::Adverse(v, w) => format!("({}::{})", v.name(), w.name()),
1084 Verb::Beside(f, g) => format!("({}∘{})", f.name(), g.name()),
1085 Verb::Hypergeometric { num, den } => {
1086 format!("({} H. {})", cx_list(num), cx_list(den))
1087 }
1088 Verb::Agenda(vs, w) => {
1089 let names: Vec<String> = vs.iter().map(Verb::name).collect();
1090 format!("({}@.{})", names.join("`"), w.name())
1091 }
1092 Verb::Evoke(vs, n) => {
1093 let names: Vec<String> = vs.iter().map(Verb::name).collect();
1094 format!("({}`:{n})", names.join("`"))
1095 }
1096 Verb::Stencil(u, w) => {
1097 let sizes: Vec<String> = w.iter().map(i64::to_string).collect();
1098 format!("({}⌺{})", u.name(), sizes.join(" "))
1099 }
1100 Verb::InnerProduct { u, v, .. } => format!("({} . {})", u.name(), v.name()),
1101 }
1102 }
1103
1104 pub fn uses_tolerance(&self) -> bool {
1109 match self {
1110 Verb::Prim(p) => {
1111 matches!(
1112 p.monad,
1113 MonadOp::Scalar(ScalarMonad::Floor)
1114 | MonadOp::Scalar(ScalarMonad::Ceil)
1115 | MonadOp::Nub
1116 ) || matches!(
1117 p.dyad,
1118 DyadOp::Scalar(
1119 ScalarDyad::Eq
1120 | ScalarDyad::Ne
1121 | ScalarDyad::Lt
1122 | ScalarDyad::Le
1123 | ScalarDyad::Gt
1124 | ScalarDyad::Ge
1125 ) | DyadOp::Match
1126 | DyadOp::NotMatch
1127 | DyadOp::MemberJ
1128 | DyadOp::MemberApl
1129 | DyadOp::IndexOf { .. }
1130 | DyadOp::IndexOfLast { .. }
1131 )
1132 }
1133 Verb::Rank(v, _)
1134 | Verb::Reduce(v)
1135 | Verb::Windowed(v, _)
1136 | Verb::Commute(v)
1137 | Verb::PowerN(v, _)
1138 | Verb::BondLeft(_, v)
1139 | Verb::BondRight(v, _)
1140 | Verb::Each(v, _)
1141 | Verb::Fit(v, _)
1142 | Verb::Key(v)
1143 | Verb::Cut(v, _)
1144 | Verb::AlongAxis(v, _) => v.uses_tolerance(),
1145 Verb::PowerV(v, w) | Verb::PowerUntil(v, w) => {
1146 v.uses_tolerance() || w.uses_tolerance()
1147 }
1148 Verb::Amend(_)
1151 | Verb::AmendVerb(_)
1152 | Verb::ShiftFill(_)
1153 | Verb::Characteristics(_)
1154 | Verb::Explicit(_)
1155 | Verb::SelfRef
1156 | Verb::Named(_)
1157 | Verb::Hypergeometric { .. } => false,
1158 Verb::Memo(v, _) | Verb::Level { u: v, .. } => v.uses_tolerance(),
1159 Verb::WithObverse(v, _) => v.uses_tolerance(),
1160 Verb::Adverse(v, w) | Verb::Beside(v, w) | Verb::Before(v, w) => {
1161 v.uses_tolerance() || w.uses_tolerance()
1162 }
1163 Verb::KeyPairs(v) => v.uses_tolerance(),
1164 Verb::UserDerived { def, alpha, omega } => {
1165 def.uses_tolerance()
1166 || alpha.uses_tolerance()
1167 || omega.as_ref().is_some_and(|g| g.uses_tolerance())
1168 }
1169 Verb::Agenda(vs, w) => {
1170 w.uses_tolerance() || vs.iter().any(Verb::uses_tolerance)
1171 }
1172 Verb::Evoke(vs, _) => vs.iter().any(Verb::uses_tolerance),
1173 Verb::Stencil(u, _) => u.uses_tolerance(),
1174 Verb::InnerProduct { u, v, .. } => u.uses_tolerance() || v.uses_tolerance(),
1175 Verb::Fork(f, g, h) => {
1176 f.uses_tolerance() || g.uses_tolerance() || h.uses_tolerance()
1177 }
1178 Verb::NounFork(_, g, h)
1179 | Verb::Hook(g, h)
1180 | Verb::Atop(g, h)
1181 | Verb::Compose(g, h) => g.uses_tolerance() || h.uses_tolerance(),
1182 }
1183 }
1184
1185 pub fn is_pure(&self) -> bool {
1190 match self {
1191 Verb::Prim(p) => {
1194 !matches!(
1195 p.monad,
1196 MonadOp::Echo | MonadOp::Roll { .. } | MonadOp::ReadStream
1197 ) && !matches!(p.dyad, DyadOp::Deal { .. } | DyadOp::WriteStream)
1198 }
1199 Verb::Rank(v, _)
1200 | Verb::Reduce(v)
1201 | Verb::Windowed(v, _)
1202 | Verb::Commute(v)
1203 | Verb::PowerN(v, _) => v.is_pure(),
1204 Verb::Fork(f, g, h) => f.is_pure() && g.is_pure() && h.is_pure(),
1205 Verb::NounFork(_, g, h)
1206 | Verb::Hook(g, h)
1207 | Verb::Atop(g, h)
1208 | Verb::Compose(g, h) => g.is_pure() && h.is_pure(),
1209 Verb::BondLeft(_, v) | Verb::BondRight(v, _) | Verb::Each(v, _) | Verb::Fit(v, _) => {
1210 v.is_pure()
1211 }
1212 Verb::Key(v) | Verb::Cut(v, _) | Verb::AlongAxis(v, _) => v.is_pure(),
1213 Verb::Hypergeometric { .. } => true,
1214 Verb::PowerV(v, w) | Verb::PowerUntil(v, w) => v.is_pure() && w.is_pure(),
1215 Verb::WithObverse(v, _) => v.is_pure(),
1216 Verb::Adverse(v, w) | Verb::Beside(v, w) | Verb::Before(v, w) => {
1217 v.is_pure() && w.is_pure()
1218 }
1219 Verb::KeyPairs(v) => v.is_pure(),
1220 Verb::UserDerived { .. } => false,
1223 Verb::Agenda(vs, w) => w.is_pure() && vs.iter().all(Verb::is_pure),
1224 Verb::Evoke(vs, _) => vs.iter().all(Verb::is_pure),
1225 Verb::Stencil(u, _) => u.is_pure(),
1226 Verb::InnerProduct { u, v, .. } => u.is_pure() && v.is_pure(),
1227 Verb::Amend(_) | Verb::ShiftFill(_) | Verb::Characteristics(_) => true,
1228 Verb::AmendVerb(v) | Verb::Level { u: v, .. } => v.is_pure(),
1229 Verb::Memo(v, _) => v.is_pure(),
1233 Verb::Explicit(_) | Verb::SelfRef | Verb::Named(_) => false,
1238 }
1239 }
1240
1241 pub fn monad(&self, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1247 let _depth = Nesting::enter(span)?;
1248 if y.is_row_major() {
1249 return self.monad_rows(y, ctx, span);
1250 }
1251 match self.monad_columns(y, ctx, span) {
1252 Some(r) => r,
1253 None => self.monad_rows(&y.to_row_major(), ctx, span),
1254 }
1255 }
1256
1257 fn monad_rows(&self, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1260 debug_assert!(y.is_row_major());
1261 match self {
1262 Verb::Prim(p) => {
1263 if let MonadOp::Scalar(op) = p.monad {
1266 return scalar_monad(op, y, ctx.cfg, span);
1267 }
1268 if p.monad == MonadOp::Open && is_mixed_simple(y) {
1272 return Ok(y.clone());
1273 }
1274 let frame_rank = y.rank() - effective_rank(p.ranks[0], y.rank());
1275 if frame_rank == 0 {
1276 return monad_op(p, y, ctx, span);
1277 }
1278 let frame = y.shape[..frame_rank].to_vec();
1279 let n: usize = frame.iter().product();
1280 let cells = each_cell(n, y.count(), self.is_pure(), ctx, |i, c| {
1281 monad_op(p, &y.cell_at(frame_rank, i), c, span)
1282 })?;
1283 assemble(&frame, cells, span)
1284 }
1285 Verb::Rank(v, r) => {
1286 let frame_rank = y.rank() - effective_rank(r[0], y.rank());
1287 if frame_rank == 0 {
1288 return v.monad(y, ctx, span);
1291 }
1292 if let Some(a) = reduce_vector_cells(v, y, frame_rank) {
1295 return Ok(a);
1296 }
1297 let frame = y.shape[..frame_rank].to_vec();
1298 let n: usize = frame.iter().product();
1299 let cells = each_cell(n, y.count(), self.is_pure(), ctx, |i, c| {
1300 v.monad(&y.cell_at(frame_rank, i), c, span)
1301 })?;
1302 assemble(&frame, cells, span)
1303 }
1304 Verb::Reduce(v) => reduce(v, y, ctx, span),
1305 Verb::Windowed(v, kind) => {
1306 runs(v, y, *kind == WindowKind::Suffix, ctx, span)
1307 }
1308 Verb::Commute(v) => v.dyad(y, y, ctx, span),
1309 Verb::PowerN(v, p) => power(v, p.clone(), None, y, ctx, span),
1310 Verb::Fork(f, g, h) => {
1311 let l = f.monad(y, ctx, span)?;
1312 let r = h.monad(y, ctx, span)?;
1313 g.dyad(&l, &r, ctx, span)
1314 }
1315 Verb::NounFork(n, g, h) => {
1316 let r = h.monad(y, ctx, span)?;
1317 g.dyad(n, &r, ctx, span)
1318 }
1319 Verb::Hook(f, g) => {
1320 let r = g.monad(y, ctx, span)?;
1321 f.dyad(y, &r, ctx, span)
1322 }
1323 Verb::Atop(f, g) | Verb::Compose(f, g) => {
1324 let r = g.monad(y, ctx, span)?;
1325 f.monad(&r, ctx, span)
1326 }
1327 Verb::BondLeft(m, v) => v.dyad(m, y, ctx, span),
1328 Verb::BondRight(v, n) => v.dyad(y, n, ctx, span),
1329 Verb::Each(u, rule) => {
1330 let n = y.count();
1331 let cells = each_cell(n, n, self.is_pure(), ctx, |i, c| {
1332 let opened = open_cell(&atom(y, i));
1333 Ok(enclose(&u.monad(&opened, c, span)?, *rule))
1334 })?;
1335 assemble(&y.shape, cells, span)
1336 }
1337 Verb::Fit(v, n) => {
1338 let tol = Tol { ct: *n, ..ctx.cfg.tol };
1339 ctx.with_tol(tol, |c| v.monad(y, c, span))
1340 }
1341 Verb::Amend(m) => {
1343 if m.rank() != 0 || y.rank() > 1 {
1344 return Err(Error::new(
1345 ErrorKind::Rank,
1346 "selecting with m} takes one index into a list",
1347 Some(span),
1348 ));
1349 }
1350 from_index(m, y, span)
1351 }
1352 Verb::AmendVerb(u) => {
1354 let m = u.monad(y, ctx, span)?;
1355 Verb::Amend(m).monad(y, ctx, span)
1356 }
1357 Verb::ShiftFill(fill) => shift_fill(&Array::scalar_i64(-1), y, fill, span),
1360 Verb::Memo(u, cache) => memoised(u, cache, None, y, ctx, span),
1361 Verb::Characteristics(u) => characteristics(u, y, span),
1362 Verb::Before(f, g) => {
1363 let l = f.monad(y, ctx, span)?;
1364 g.dyad(&l, y, ctx, span)
1365 }
1366 Verb::KeyPairs(u) => key_pairs(u, y, None, ctx, span),
1367 Verb::UserDerived { def, alpha, omega } => {
1368 with_operands(alpha, omega.as_deref(), ctx, |c| def.monad(y, c, span))
1369 }
1370 Verb::Level { u, level, spread } => {
1371 at_level(u, *level, *spread, y, ctx, span)
1372 }
1373 Verb::Key(u) => oblique(u, y, ctx, span),
1374 Verb::Cut(u, n) => cut(u, None, y, *n, ctx, span),
1375 Verb::PowerV(u, v) => power_v(u, v, None, y, ctx, span),
1376 Verb::PowerUntil(u, v) => power_until(u, v, y, ctx, span),
1377 Verb::AlongAxis(u, k) => along_axis(u, None, y, *k, ctx, span),
1378 Verb::Explicit(d) => crate::ir::call_explicit(d, None, y, ctx, span),
1379 Verb::SelfRef => {
1380 let d = self_ref(ctx, span)?;
1381 crate::ir::call_explicit(&d, None, y, ctx, span)
1382 }
1383 Verb::Named(n) => named_verb(ctx, n, span)?.monad(y, ctx, span),
1384 Verb::WithObverse(v, _) => v.monad(y, ctx, span),
1385 Verb::Adverse(v, w) => match v.monad(y, ctx, span) {
1386 Err(e) if e.kind != ErrorKind::NotYet => w.monad(y, ctx, span),
1387 other => other,
1388 },
1389 Verb::Beside(f, g) => {
1390 let r = g.monad(y, ctx, span)?;
1391 f.monad(&r, ctx, span)
1392 }
1393 Verb::Hypergeometric { num, den } => hypergeometric(num, den, y, span),
1394 Verb::Agenda(vs, w) => {
1395 agenda_pick(vs, w, None, y, ctx, span)?.monad(y, ctx, span)
1396 }
1397 Verb::Evoke(vs, n) => evoke(vs, *n, None, y, ctx, span),
1398 Verb::Stencil(u, w) => stencil(u, w, y, ctx, span),
1399 Verb::InnerProduct { u, v, apl } => determinant(u, v, *apl, y, ctx, span),
1400 }
1401 }
1402
1403 fn monad_columns(&self, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Option<Result<Array>> {
1412 debug_assert!(!y.is_row_major());
1413 match self {
1414 Verb::Prim(p) => match p.monad {
1415 MonadOp::Scalar(op) => Some(scalar_monad(op, y, ctx.cfg, span)),
1418 MonadOp::ShapeOf | MonadOp::Tally => Some(monad_op(p, y, ctx, span)),
1420 MonadOp::TransposeAxes => Some(Ok(transpose_axes(y))),
1424 _ => None,
1425 },
1426 Verb::Reduce(v) => reduce_columns(v, y).map(Ok),
1429 Verb::Rank(v, r) => {
1432 if y.rank() != effective_rank(r[0], y.rank()) + 1 {
1433 return None;
1434 }
1435 reduce_rows_columns(v, y).map(Ok)
1436 }
1437 _ => None,
1438 }
1439 }
1440
1441 pub fn dyad(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1447 let _depth = Nesting::enter(span)?;
1448 if x.is_row_major() && y.is_row_major() {
1449 return self.dyad_rows(x, y, ctx, span);
1450 }
1451 if let Some(layout) = self.elementwise_layout(x, y) {
1452 return Ok(self.dyad_rows(x, y, ctx, span)?.with_layout(layout));
1453 }
1454 self.dyad_rows(&x.to_row_major(), &y.to_row_major(), ctx, span)
1455 }
1456
1457 fn elementwise_layout(&self, x: &Array, y: &Array) -> Option<Layout> {
1462 let Verb::Prim(p) = self else { return None };
1463 if !matches!(p.dyad, DyadOp::Scalar(_)) {
1464 return None;
1465 }
1466 if x.rank() == 0 {
1467 return Some(y.layout());
1468 }
1469 if y.rank() == 0 {
1470 return Some(x.layout());
1471 }
1472 (x.shape == y.shape && x.layout() == y.layout()).then(|| x.layout())
1473 }
1474
1475 fn dyad_rows(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1479 match self {
1480 Verb::Prim(_) | Verb::Rank(_, _) | Verb::Each(..) => {
1481 self.dyad_ranked(x, y, ctx, span)
1482 }
1483 Verb::Windowed(_, WindowKind::Prefix) => self.dyad_ranked(x, y, ctx, span),
1485 Verb::Windowed(u, WindowKind::Suffix) => outfix(u, x, y, ctx, span),
1488 Verb::Windowed(_, WindowKind::Scan) => {
1489 Err(Error::not_yet("dyadic scan (x f\\ y)", span))
1490 }
1491 Verb::Commute(v) => v.dyad(y, x, ctx, span),
1492 Verb::PowerN(v, p) => power(v, p.clone(), Some(x), y, ctx, span),
1493 Verb::Reduce(v) => table(v, x, y, ctx, span),
1495 Verb::Fork(f, g, h) => {
1496 let l = f.dyad(x, y, ctx, span)?;
1497 let r = h.dyad(x, y, ctx, span)?;
1498 g.dyad(&l, &r, ctx, span)
1499 }
1500 Verb::NounFork(n, g, h) => {
1501 let r = h.dyad(x, y, ctx, span)?;
1502 g.dyad(n, &r, ctx, span)
1503 }
1504 Verb::Hook(f, g) => {
1505 let r = g.monad(y, ctx, span)?;
1506 f.dyad(x, &r, ctx, span)
1507 }
1508 Verb::Atop(f, g) => {
1509 let r = g.dyad(x, y, ctx, span)?;
1510 f.monad(&r, ctx, span)
1511 }
1512 Verb::Compose(f, g) => {
1513 let l = g.monad(x, ctx, span)?;
1514 let r = g.monad(y, ctx, span)?;
1515 f.dyad(&l, &r, ctx, span)
1516 }
1517 Verb::Fit(v, n) => {
1518 let tol = Tol { ct: *n, ..ctx.cfg.tol };
1519 ctx.with_tol(tol, |c| v.dyad(x, y, c, span))
1520 }
1521 Verb::Amend(m) => amend(m, x, y, span),
1522 Verb::AmendVerb(u) => {
1524 let m = u.dyad(x, y, ctx, span)?;
1525 amend(&m, x, y, span)
1526 }
1527 Verb::ShiftFill(fill) => shift_fill(x, y, fill, span),
1528 Verb::Memo(u, cache) => memoised(u, cache, Some(x), y, ctx, span),
1529 Verb::Characteristics(_) => {
1530 Err(Error::domain("u b. has no dyadic meaning", span))
1531 }
1532 Verb::Before(f, g) => {
1533 let l = f.monad(x, ctx, span)?;
1534 g.dyad(&l, y, ctx, span)
1535 }
1536 Verb::KeyPairs(u) => key_pairs(u, x, Some(y), ctx, span),
1537 Verb::UserDerived { def, alpha, omega } => {
1538 with_operands(alpha, omega.as_deref(), ctx, |c| def.dyad(x, y, c, span))
1539 }
1540 Verb::Level { u, level, spread } => {
1541 at_level_dyad(u, *level, *spread, x, y, ctx, span)
1542 }
1543 Verb::Key(u) => key(u, x, y, ctx, span),
1544 Verb::Cut(u, n) => cut(u, Some(x), y, *n, ctx, span),
1545 Verb::PowerV(u, v) => power_v(u, v, Some(x), y, ctx, span),
1546 Verb::PowerUntil(..) => {
1547 Err(Error::not_yet("dyadic power with a function operand (x f⍣g y)", span))
1548 }
1549 Verb::AlongAxis(u, k) => along_axis(u, Some(x), y, *k, ctx, span),
1550 Verb::Explicit(d) => crate::ir::call_explicit(d, Some(x), y, ctx, span),
1551 Verb::SelfRef => {
1552 let d = self_ref(ctx, span)?;
1553 crate::ir::call_explicit(&d, Some(x), y, ctx, span)
1554 }
1555 Verb::Named(n) => named_verb(ctx, n, span)?.dyad(x, y, ctx, span),
1556 Verb::WithObverse(v, _) => v.dyad(x, y, ctx, span),
1557 Verb::Adverse(v, w) => match v.dyad(x, y, ctx, span) {
1558 Err(e) if e.kind != ErrorKind::NotYet => w.dyad(x, y, ctx, span),
1559 other => other,
1560 },
1561 Verb::Beside(f, g) => {
1562 let r = g.monad(y, ctx, span)?;
1563 f.dyad(x, &r, ctx, span)
1564 }
1565 Verb::Hypergeometric { .. } => {
1566 Err(Error::domain("m H. n has no dyadic meaning", span))
1567 }
1568 Verb::Agenda(vs, w) => {
1569 agenda_pick(vs, w, Some(x), y, ctx, span)?.dyad(x, y, ctx, span)
1570 }
1571 Verb::Evoke(vs, n) => evoke(vs, *n, Some(x), y, ctx, span),
1572 Verb::InnerProduct { u, v, apl } => inner_product(u, v, *apl, x, y, ctx, span),
1573 Verb::Stencil(..) => {
1574 Err(Error::domain("f⌺w has no dyadic meaning", span))
1575 }
1576 Verb::BondLeft(..) | Verb::BondRight(..) => {
1578 Err(Error::domain(format!("{} has no dyadic meaning", self.name()), span))
1579 }
1580 }
1581 }
1582
1583 fn dyad_ranked(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1585 let ranks = self.ranks();
1586 let er_l = effective_rank(ranks[1], x.rank());
1587 let er_r = effective_rank(ranks[2], y.rank());
1588 if er_l == 0 && er_r == 0 {
1589 if let Some(op) = self.scalar_dyad_op() {
1592 return scalar_dyad(op, x, y, ctx.cfg, span);
1593 }
1594 }
1595 let fxl = x.rank() - er_l;
1596 let fyl = y.rank() - er_r;
1597 let p = agree(&x.shape[..fxl], &y.shape[..fyl], &x.shape, &y.shape, ctx.cfg.agreement, span)?;
1598 if p.frame.is_empty() {
1599 return self.dyad_cell(x, y, ctx, span);
1600 }
1601 let work = x.count().max(y.count());
1602 let cells = each_cell(p.n, work, self.is_pure(), ctx, |i, c| {
1603 let xc = x.cell_at(fxl, i / p.x_div);
1604 let yc = y.cell_at(fyl, i / p.y_div);
1605 self.dyad_cell(&xc, &yc, c, span)
1606 })?;
1607 assemble(&p.frame, cells, span)
1608 }
1609
1610 fn dyad_cell(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1612 match self {
1613 Verb::Prim(p) if p.dyad == DyadOp::WriteStream => {
1616 stream_number(y, 2, "1!:2 writes", span)?;
1617 (ctx.out)(&format!("{}\n", crate::fmt::format_array(x, &ctx.cfg.fmt)));
1618 Ok(x.clone())
1619 }
1620 Verb::Prim(p) => dyad_op(p, x, y, ctx.cfg, span),
1621 Verb::Rank(v, _) => v.dyad(x, y, ctx, span),
1622 Verb::Windowed(v, _) => infix(v, x, y, ctx, span),
1623 Verb::Each(u, rule) => {
1624 let r = u.dyad(&open_cell(x), &open_cell(y), ctx, span)?;
1625 Ok(enclose(&r, *rule))
1626 }
1627 _ => Err(Error::internal("dyad_cell on a verb without cell ranks")),
1628 }
1629 }
1630
1631 fn scalar_dyad_op(&self) -> Option<ScalarDyad> {
1634 match self {
1635 Verb::Prim(p) => match p.dyad {
1636 DyadOp::Scalar(op) => Some(op),
1637 _ => None,
1638 },
1639 Verb::Rank(v, _) => v.scalar_dyad_op(),
1640 _ => None,
1641 }
1642 }
1643}
1644
1645pub fn effective_rank(r: i64, arg_rank: usize) -> usize {
1648 if r >= 0 {
1649 (r as usize).min(arg_rank)
1650 } else {
1651 arg_rank.saturating_sub(r.unsigned_abs() as usize)
1652 }
1653}
1654
1655fn self_ref(ctx: &Ctx<'_>, span: Span) -> Result<Arc<crate::ir::ExplicitDef>> {
1665 ctx.env.current_def().ok_or_else(|| {
1666 Error::new(
1667 ErrorKind::Value,
1668 "self-reference outside an explicit definition",
1669 Some(span),
1670 )
1671 })
1672}
1673
1674fn named_verb(ctx: &Ctx<'_>, name: &str, span: Span) -> Result<Verb> {
1676 ctx.env.verb(name).cloned().ok_or_else(|| {
1677 Error::new(ErrorKind::Value, format!("undefined verb: {name}"), Some(span))
1678 })
1679}
1680
1681fn each_cell<F>(
1682 n: usize,
1683 work: usize,
1684 pure: bool,
1685 ctx: &mut Ctx<'_>,
1686 f: F,
1687) -> Result<Vec<Array>>
1688where
1689 F: Fn(usize, &mut Ctx<'_>) -> Result<Array> + Sync + Send,
1690{
1691 if pure && n > 1 && par::worth_it(work) {
1692 let cfg = ctx.cfg;
1693 return par::map_indexed(n, |i| cfg.pure(|c| f(i, c))).into_iter().collect();
1694 }
1695 (0..n).map(|i| f(i, ctx)).collect()
1696}
1697
1698fn one_rank(r: i64) -> String {
1701 if r == RANK_INF { "_".to_string() } else { r.to_string() }
1702}
1703
1704fn rank_str(r: [i64; 3]) -> String {
1707 if r[0] == r[1] && r[1] == r[2] {
1708 one_rank(r[0])
1709 } else {
1710 format!("{} {} {}", one_rank(r[0]), one_rank(r[1]), one_rank(r[2]))
1711 }
1712}
1713
1714fn show_shape(shape: &[usize]) -> String {
1716 if shape.is_empty() {
1717 return "(scalar)".to_string();
1718 }
1719 shape.iter().map(|n| n.to_string()).collect::<Vec<_>>().join(" ")
1720}
1721
1722fn strides(shape: &[usize]) -> Vec<usize> {
1726 let mut s = vec![1usize; shape.len()];
1727 for k in (0..shape.len().saturating_sub(1)).rev() {
1728 s[k] = s[k + 1] * shape[k + 1];
1729 }
1730 s
1731}
1732
1733fn odometer(coord: &mut [usize], shape: &[usize]) {
1735 for k in (0..coord.len()).rev() {
1736 coord[k] += 1;
1737 if coord[k] < shape[k] {
1738 return;
1739 }
1740 coord[k] = 0;
1741 }
1742}
1743
1744fn push_elem(dst: &mut Data, src: &Data, i: usize) {
1746 match (dst, src) {
1747 (Data::Bool(a), Data::Bool(b)) => a.push(b[i]),
1748 (Data::I64(a), Data::I64(b)) => a.push(b[i]),
1749 (Data::Ext(a), Data::Ext(b)) => a.push(b[i].clone()),
1750 (Data::Rat(a), Data::Rat(b)) => a.push(b[i].clone()),
1751 (Data::F64(a), Data::F64(b)) => a.push(b[i]),
1752 (Data::Complex(a), Data::Complex(b)) => a.push(b[i]),
1753 (Data::Char(a), Data::Char(b)) => a.push(b[i]),
1754 (Data::Symbol(a), Data::Symbol(b)) => a.push(b[i]),
1755 (Data::Box(a), Data::Box(b)) => a.push(b[i].clone()),
1756 _ => debug_assert!(false, "push_elem across dtypes"),
1757 }
1758}
1759
1760fn fill_data(dtype: DType, n: usize) -> Data {
1762 let mut d = Data::empty(dtype);
1763 for _ in 0..n {
1764 d.push_fill();
1765 }
1766 d
1767}
1768
1769struct Pairing {
1774 frame: Vec<usize>,
1775 n: usize,
1776 x_div: usize,
1777 y_div: usize,
1778}
1779
1780fn frame_mismatch(
1781 xs: &[usize],
1782 ys: &[usize],
1783 fx: &[usize],
1784 fy: &[usize],
1785 axis: usize,
1786 span: Span,
1787) -> Error {
1788 let kind = if fx.len() == 1 && fy.len() == 1 { ErrorKind::Length } else { ErrorKind::Shape };
1791 let note = if axis < fx.len() && axis < fy.len() {
1792 format!("frames first differ at axis {axis}: {} vs {}", fx[axis], fy[axis])
1793 } else {
1794 format!(
1795 "frames have different numbers of axes: {} vs {}, diverging at axis {axis}",
1796 fx.len(),
1797 fy.len()
1798 )
1799 };
1800 Error::new(
1801 kind,
1802 format!(
1803 "arguments do not agree: left shape {}, right shape {}",
1804 show_shape(xs),
1805 show_shape(ys)
1806 ),
1807 Some(span),
1808 )
1809 .note(note)
1810}
1811
1812fn agree(
1815 fx: &[usize],
1816 fy: &[usize],
1817 xs: &[usize],
1818 ys: &[usize],
1819 mode: Agreement,
1820 span: Span,
1821) -> Result<Pairing> {
1822 let common = fx.len().min(fy.len());
1823 match mode {
1824 Agreement::LeadingPrefix => {
1825 for i in 0..common {
1826 if fx[i] != fy[i] {
1827 return Err(frame_mismatch(xs, ys, fx, fy, i, span));
1828 }
1829 }
1830 let (long, short) = if fx.len() >= fy.len() { (fx, fy) } else { (fy, fx) };
1831 let n: usize = long.iter().product();
1832 let surplus: usize = long[short.len()..].iter().product();
1833 let (x_div, y_div) =
1834 if fx.len() >= fy.len() { (1, surplus.max(1)) } else { (surplus.max(1), 1) };
1835 Ok(Pairing { frame: long.to_vec(), n, x_div, y_div })
1836 }
1837 Agreement::ExactOrScalar => {
1838 if fx == fy {
1839 let n: usize = fx.iter().product();
1840 return Ok(Pairing { frame: fx.to_vec(), n, x_div: 1, y_div: 1 });
1841 }
1842 let one = |f: &[usize]| f.iter().product::<usize>() == 1;
1849 if fx.is_empty() || (one(fx) && !fy.is_empty()) {
1850 let n: usize = fy.iter().product();
1851 return Ok(Pairing { frame: fy.to_vec(), n, x_div: n.max(1), y_div: 1 });
1852 }
1853 if fy.is_empty() || one(fy) {
1854 let n: usize = fx.iter().product();
1855 return Ok(Pairing { frame: fx.to_vec(), n, x_div: 1, y_div: n.max(1) });
1856 }
1857 let axis = (0..common).find(|&i| fx[i] != fy[i]).unwrap_or(common);
1858 Err(frame_mismatch(xs, ys, fx, fy, axis, span))
1859 }
1860 }
1861}
1862
1863fn assemble(frame: &[usize], cells: Vec<Array>, span: Span) -> Result<Array> {
1867 if cells.is_empty() {
1868 return Ok(Array::new(frame.to_vec(), Data::empty(DType::I64)));
1871 }
1872 let mut dt = cells[0].dtype();
1873 for c in &cells[1..] {
1874 dt = DType::promote(dt, c.dtype()).ok_or_else(|| {
1875 let boxed = dt == DType::Box || c.dtype() == DType::Box;
1876 let what = if boxed {
1877 "cannot frame boxed and unboxed results into one array"
1878 } else {
1879 "cannot frame character and numeric results into one array"
1880 };
1881 Error::new(ErrorKind::Type, what, Some(span))
1882 })?;
1883 }
1884 let widen = |c: &Array| -> Result<Data> {
1885 c.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening while framing"))
1886 };
1887
1888 if cells[1..].iter().all(|c| c.shape == cells[0].shape) {
1889 let mut data = Data::empty(dt);
1890 for c in &cells {
1891 if c.dtype() == dt {
1892 data.extend_from(&c.data);
1893 } else {
1894 data.extend_from(&widen(c)?);
1895 }
1896 }
1897 let mut shape = frame.to_vec();
1898 shape.extend_from_slice(&cells[0].shape);
1899 return Ok(Array::new(shape, data));
1900 }
1901
1902 let crank = cells.iter().map(|c| c.rank()).max().unwrap_or(0);
1905 let padded: Vec<Vec<usize>> = cells
1906 .iter()
1907 .map(|c| {
1908 let mut s = vec![1usize; crank - c.rank()];
1909 s.extend_from_slice(&c.shape);
1910 s
1911 })
1912 .collect();
1913 let mut common = vec![0usize; crank];
1914 for s in &padded {
1915 for k in 0..crank {
1916 common[k] = common[k].max(s[k]);
1917 }
1918 }
1919 let cell_n: usize = common.iter().product();
1920 let mut data = Data::empty(dt);
1921 for (c, ps) in cells.iter().zip(&padded) {
1922 let cd = if c.dtype() == dt { c.data.clone() } else { widen(c)? };
1923 let st = strides(ps);
1924 let mut coord = vec![0usize; crank];
1925 for _ in 0..cell_n {
1926 let mut idx = 0usize;
1927 let mut inside = true;
1928 for k in 0..crank {
1929 if coord[k] >= ps[k] {
1930 inside = false;
1931 break;
1932 }
1933 idx += coord[k] * st[k];
1934 }
1935 if inside {
1936 push_elem(&mut data, &cd, idx);
1937 } else {
1938 data.push_fill();
1939 }
1940 odometer(&mut coord, &common);
1941 }
1942 }
1943 let mut shape = frame.to_vec();
1944 shape.extend_from_slice(&common);
1945 Ok(Array::new(shape, data))
1946}
1947
1948fn atom(a: &Array, i: usize) -> Array {
1953 debug_assert!(a.is_row_major(), "an atom out of a column-major buffer");
1954 Array::new(Vec::new(), a.data.slice(i, i + 1))
1955}
1956
1957fn enclose(y: &Array, rule: Enclose) -> Array {
1959 if rule == Enclose::ExceptSimpleScalar && y.rank() == 0 && y.dtype() != DType::Box {
1960 return y.clone();
1961 }
1962 Array::boxed(y.clone())
1963}
1964
1965fn open_cell(y: &Array) -> Array {
1968 match &y.data {
1969 Data::Box(v) if !v.is_empty() => v[0].clone(),
1970 _ => y.clone(),
1971 }
1972}
1973
1974fn first(y: &Array) -> Array {
1977 if y.count() == 0 {
1978 let mut d = Data::empty(y.dtype());
1979 d.push_fill();
1980 return open_cell(&Array::new(Vec::new(), d));
1981 }
1982 open_cell(&atom(y, 0))
1983}
1984
1985fn depth(y: &Array) -> i64 {
1987 match &y.data {
1988 Data::Box(v) => 1 + v.iter().map(depth).max().unwrap_or(0),
1989 _ => i64::from(y.rank() > 0),
1990 }
1991}
1992
1993fn leaves(a: &Array, out: &mut Vec<Array>) {
1995 match &a.data {
1996 Data::Box(v) => {
1997 for b in v.iter() {
1998 leaves(b, out);
1999 }
2000 }
2001 _ => out.push(a.clone()),
2002 }
2003}
2004
2005fn enlist(y: &Array, span: Span) -> Result<Array> {
2007 let mut parts = Vec::new();
2008 leaves(y, &mut parts);
2009 let mut dt = None;
2012 for p in parts.iter().filter(|p| p.count() > 0) {
2013 dt = Some(match dt {
2014 None => p.dtype(),
2015 Some(t) => DType::promote(t, p.dtype()).ok_or_else(|| {
2016 Error::new(
2017 ErrorKind::Type,
2018 "cannot enlist character and numeric data into one vector",
2019 Some(span),
2020 )
2021 })?,
2022 });
2023 }
2024 let dt = dt.unwrap_or(DType::I64);
2025 let mut data = Data::empty(dt);
2026 for p in &parts {
2027 let cast = p.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening in enlist"))?;
2028 data.extend_from(&cast);
2029 }
2030 Ok(Array::new(vec![data.len()], data))
2031}
2032
2033fn spread(a: &Array, shape: &[usize]) -> Array {
2035 let n: usize = shape.iter().product();
2036 let mut data = Data::empty(a.dtype());
2037 for _ in 0..n {
2038 push_elem(&mut data, &a.data, 0);
2039 }
2040 Array::new(shape.to_vec(), data)
2041}
2042
2043fn wider_shape(a: &[usize], b: &[usize]) -> Vec<usize> {
2046 let r = a.len().max(b.len());
2047 let pad = |s: &[usize]| {
2048 let mut v = vec![1usize; r - s.len()];
2049 v.extend_from_slice(s);
2050 v
2051 };
2052 let (pa, pb) = (pad(a), pad(b));
2053 (0..r).map(|k| pa[k].max(pb[k])).collect()
2054}
2055
2056fn raze(y: &Array, span: Span) -> Result<Array> {
2061 let opened: Vec<Array> = (0..y.count()).map(|i| open_cell(&atom(y, i))).collect();
2062 let mut common: Option<Vec<usize>> = None;
2063 for a in opened.iter().filter(|a| a.rank() > 0) {
2064 common = Some(match common {
2065 None => a.shape[1..].to_vec(),
2066 Some(c) => wider_shape(&c, &a.shape[1..]),
2067 });
2068 }
2069 let common = common.unwrap_or_default();
2070 let mut cells: Vec<Array> = Vec::new();
2071 for a in &opened {
2072 if a.rank() == 0 {
2073 cells.push(spread(a, &common));
2074 continue;
2075 }
2076 for i in 0..a.items() {
2077 cells.push(a.item(i));
2078 }
2079 }
2080 if cells.is_empty() {
2081 return Ok(Array::new(vec![0], Data::empty(DType::I64)));
2082 }
2083 let n = cells.len();
2084 assemble(&[n], cells, span)
2085}
2086
2087fn link(x: &Array, y: &Array, span: Span) -> Result<Array> {
2090 let head = Array::boxed(x.clone());
2091 let tail = if y.dtype() == DType::Box { y.clone() } else { Array::boxed(y.clone()) };
2092 catenate(&head, &tail, true, false, span)
2093}
2094
2095fn nest_like(a: &Array, other: &Array) -> Array {
2098 if a.dtype() == DType::Box || other.dtype() != DType::Box {
2099 return a.clone();
2100 }
2101 let cells: Vec<Array> = (0..a.count()).map(|i| atom(a, i)).collect();
2102 Array::new(a.shape.clone(), Data::Box(cells.into()))
2103}
2104
2105fn box_items(y: &Array) -> Array {
2107 if y.dtype() == DType::Box {
2108 return y.clone();
2109 }
2110 let n = y.items();
2111 let boxes: Vec<Array> = (0..n).map(|i| item_or_self(y, i)).collect();
2112 Array::new(vec![n], Data::Box(boxes.into()))
2113}
2114
2115fn strand(x: &Array, y: &Array, span: Span) -> Result<Array> {
2119 let item = enclose(x, Enclose::ExceptSimpleScalar);
2120 let one = |a: &Array| Array::new(vec![1], a.data.clone());
2121 if item.dtype() != DType::Box
2126 && y.dtype() != DType::Box
2127 && DType::promote(item.dtype(), y.dtype()).is_some()
2128 {
2129 return catenate(&one(&item), y, true, false, span);
2130 }
2131 let head = if item.dtype() == DType::Box { item } else { Array::boxed(item) };
2132 catenate(&one(&head), &box_items(y), true, false, span)
2133}
2134
2135fn char_arith(span: Span) -> Error {
2138 Error::new(ErrorKind::Type, "cannot do arithmetic on characters", Some(span))
2139}
2140
2141fn symbol_arith(span: Span) -> Error {
2142 Error::new(
2143 ErrorKind::Type,
2144 "cannot do arithmetic on symbols; `5 s:` gives their names back",
2145 Some(span),
2146 )
2147}
2148
2149fn box_arith(span: Span) -> Error {
2150 Error::new(
2151 ErrorKind::Type,
2152 "cannot do arithmetic on boxed values; open them first (J `>`, APL `⊃`)",
2153 Some(span),
2154 )
2155}
2156
2157fn wrong_type(d: DType, span: Span) -> Error {
2160 match d {
2161 DType::Box => box_arith(span),
2162 DType::Symbol => symbol_arith(span),
2163 _ => char_arith(span),
2164 }
2165}
2166
2167fn borrow_i64<'a>(d: &'a Data, tmp: &'a mut Vec<i64>) -> &'a [i64] {
2173 match d {
2174 Data::I64(v) => v,
2175 Data::Bool(v) => {
2176 *tmp = par::map(v, |&b| b as i64);
2177 &tmp[..]
2178 }
2179 _ => &[],
2181 }
2182}
2183
2184fn borrow_f64<'a>(d: &'a Data, tmp: &'a mut Vec<f64>) -> &'a [f64] {
2186 match d {
2187 Data::F64(v) => v,
2188 Data::I64(v) => {
2189 *tmp = par::map(v, |&x| x as f64);
2190 &tmp[..]
2191 }
2192 Data::Bool(v) => {
2193 *tmp = par::map(v, |&x| x as f64);
2194 &tmp[..]
2195 }
2196 Data::Ext(v) => {
2197 *tmp = par::map(v, exact::ext_to_f64);
2198 &tmp[..]
2199 }
2200 Data::Rat(v) => {
2201 *tmp = par::map(v, Rat::to_f64);
2202 &tmp[..]
2203 }
2204 _ => &[],
2205 }
2206}
2207
2208fn borrow_cx<'a>(d: &'a Data, tmp: &'a mut Vec<Cx>) -> &'a [Cx] {
2210 match d {
2211 Data::Complex(v) => v,
2212 Data::Ext(v) => {
2213 *tmp = par::map(v, |x| [exact::ext_to_f64(x), 0.0]);
2214 &tmp[..]
2215 }
2216 Data::Rat(v) => {
2217 *tmp = par::map(v, |x| [x.to_f64(), 0.0]);
2218 &tmp[..]
2219 }
2220 Data::F64(v) => {
2221 *tmp = par::map(v, |&x| [x, 0.0]);
2222 &tmp[..]
2223 }
2224 Data::I64(v) => {
2225 *tmp = par::map(v, |&x| [x as f64, 0.0]);
2226 &tmp[..]
2227 }
2228 Data::Bool(v) => {
2229 *tmp = v.iter().map(|&x| [x as f64, 0.0]).collect();
2230 &tmp[..]
2231 }
2232 _ => &[],
2233 }
2234}
2235
2236pub(crate) trait Widen<T>: Copy + Send + Sync {
2245 fn widen(self) -> T;
2246}
2247
2248macro_rules! widens {
2249 ($($from:ty => $to:ty : |$v:ident| $e:expr;)*) => {
2250 $(impl Widen<$to> for $from {
2251 #[inline(always)]
2252 fn widen(self) -> $to {
2253 let $v = self;
2254 $e
2255 }
2256 })*
2257 };
2258}
2259
2260widens! {
2261 u8 => i64: |v| v as i64;
2262 i64 => i64: |v| v;
2263 u8 => f64: |v| v as f64;
2264 i64 => f64: |v| v as f64;
2265 f64 => f64: |v| v;
2266 u8 => Cx: |v| [v as f64, 0.0];
2267 i64 => Cx: |v| [v as f64, 0.0];
2268 f64 => Cx: |v| [v, 0.0];
2269 Cx => Cx: |v| v;
2270}
2271
2272macro_rules! i64_source {
2275 ($d:expr, $tmp:ident, $s:ident, $body:expr) => {
2276 match $d {
2277 Data::I64(v) => {
2278 let $s: &[i64] = v;
2279 $body
2280 }
2281 Data::Bool(v) => {
2282 let $s: &[u8] = v;
2283 $body
2284 }
2285 other => {
2286 let $s: &[i64] = borrow_i64(other, &mut $tmp);
2287 $body
2288 }
2289 }
2290 };
2291}
2292
2293macro_rules! f64_source {
2296 ($d:expr, $tmp:ident, $s:ident, $body:expr) => {
2297 match $d {
2298 Data::F64(v) => {
2299 let $s: &[f64] = v;
2300 $body
2301 }
2302 Data::I64(v) => {
2303 let $s: &[i64] = v;
2304 $body
2305 }
2306 Data::Bool(v) => {
2307 let $s: &[u8] = v;
2308 $body
2309 }
2310 other => {
2311 let $s: &[f64] = borrow_f64(other, &mut $tmp);
2312 $body
2313 }
2314 }
2315 };
2316}
2317
2318macro_rules! cx_source {
2320 ($d:expr, $tmp:ident, $s:ident, $body:expr) => {
2321 match $d {
2322 Data::Complex(v) => {
2323 let $s: &[Cx] = v;
2324 $body
2325 }
2326 Data::F64(v) => {
2327 let $s: &[f64] = v;
2328 $body
2329 }
2330 Data::I64(v) => {
2331 let $s: &[i64] = v;
2332 $body
2333 }
2334 Data::Bool(v) => {
2335 let $s: &[u8] = v;
2336 $body
2337 }
2338 other => {
2339 let $s: &[Cx] = borrow_cx(other, &mut $tmp);
2340 $body
2341 }
2342 }
2343 };
2344}
2345
2346fn as_f64<'a>(d: &'a Data, tmp: &'a mut Vec<f64>, span: Span) -> Result<&'a [f64]> {
2348 if !d.dtype().is_numeric() {
2349 return Err(wrong_type(d.dtype(), span));
2350 }
2351 Ok(borrow_f64(d, tmp))
2352}
2353
2354fn arith_type(a: DType, b: DType, span: Span) -> Result<DType> {
2356 if a == DType::Box || b == DType::Box {
2357 return Err(box_arith(span));
2358 }
2359 if a == DType::Symbol || b == DType::Symbol {
2360 return Err(symbol_arith(span));
2361 }
2362 match DType::promote(a, b) {
2363 Some(DType::Char) => Err(char_arith(span)),
2364 None => Err(Error::new(
2365 ErrorKind::Type,
2366 "cannot mix character and numeric data",
2367 Some(span),
2368 )),
2369 Some(DType::Bool) => Ok(DType::I64),
2370 Some(t) => Ok(t),
2371 }
2372}
2373
2374#[allow(clippy::too_many_arguments)]
2387#[inline]
2388fn zip_chunk<A, B, U, F>(
2389 xs: &[A],
2390 xoff: usize,
2391 xdiv: usize,
2392 ys: &[B],
2393 yoff: usize,
2394 ydiv: usize,
2395 start: usize,
2396 out: &mut [U],
2397 mut f: F,
2398) -> bool
2399where
2400 A: Copy,
2401 B: Copy,
2402 F: FnMut(A, B, &mut U) -> bool,
2403{
2404 let len = out.len();
2405 if len == 0 {
2406 return true;
2407 }
2408 let last = start + len - 1;
2409 let one_x = xdiv > 1 && start / xdiv == last / xdiv;
2410 let one_y = ydiv > 1 && start / ydiv == last / ydiv;
2411 if xdiv == 1 && ydiv == 1 {
2412 let xc = &xs[xoff + start..xoff + start + len];
2413 let yc = &ys[yoff + start..yoff + start + len];
2414 for ((slot, &a), &b) in out.iter_mut().zip(xc).zip(yc) {
2415 if !f(a, b, slot) {
2416 return false;
2417 }
2418 }
2419 } else if xdiv == 1 && one_y {
2420 let b = ys[yoff + start / ydiv];
2421 let xc = &xs[xoff + start..xoff + start + len];
2422 for (slot, &a) in out.iter_mut().zip(xc) {
2423 if !f(a, b, slot) {
2424 return false;
2425 }
2426 }
2427 } else if one_x && ydiv == 1 {
2428 let a = xs[xoff + start / xdiv];
2429 let yc = &ys[yoff + start..yoff + start + len];
2430 for (slot, &b) in out.iter_mut().zip(yc) {
2431 if !f(a, b, slot) {
2432 return false;
2433 }
2434 }
2435 } else {
2436 for (k, slot) in out.iter_mut().enumerate() {
2437 let i = start + k;
2438 if !f(xs[xoff + i / xdiv], ys[yoff + i / ydiv], slot) {
2439 return false;
2440 }
2441 }
2442 }
2443 true
2444}
2445
2446const LANCZOS: [f64; 9] = [
2450 0.999_999_999_999_809_9,
2451 676.520_368_121_885_1,
2452 -1_259.139_216_722_402_8,
2453 771.323_428_777_653_1,
2454 -176.615_029_162_140_6,
2455 12.507_343_278_686_905,
2456 -0.138_571_095_265_720_12,
2457 9.984_369_578_019_572e-6,
2458 1.505_632_735_149_311_6e-7,
2459];
2460
2461fn gamma(x: f64) -> f64 {
2465 use std::f64::consts::PI;
2466 if x < 0.5 {
2467 return PI / ((PI * x).sin() * gamma(1.0 - x));
2469 }
2470 let z = x - 1.0;
2471 let mut a = LANCZOS[0];
2472 for (i, &c) in LANCZOS.iter().enumerate().skip(1) {
2473 a += c / (z + i as f64);
2474 }
2475 let t = z + 7.5;
2476 (2.0 * PI).sqrt() * t.powf(z + 0.5) * (-t).exp() * a
2477}
2478
2479fn factorial(y: f64) -> f64 {
2482 if y.fract() == 0.0 && y.abs() < 1e17 {
2483 let n = y as i64;
2484 if n < 0 {
2485 return if n % 2 == -1 { f64::INFINITY } else { f64::NEG_INFINITY };
2487 }
2488 if n > 170 {
2489 return f64::INFINITY;
2490 }
2491 let mut c = 1.0f64;
2492 for i in 2..=n {
2493 c *= i as f64;
2494 }
2495 return c;
2496 }
2497 gamma(y + 1.0)
2498}
2499
2500const BINOMIAL_PRODUCT_LIMIT: i64 = 4096;
2503
2504fn binomial_product(x: i64, y: f64) -> f64 {
2508 let mut c = 1.0f64;
2509 for i in 1..=x {
2510 c = c * (y - i as f64 + 1.0) / i as f64;
2511 if c == 0.0 {
2512 break;
2513 }
2514 }
2515 c
2516}
2517
2518fn binomial_i64(x: i64, y: i64) -> Option<i64> {
2522 if x < 0 {
2523 if y >= 0 || y < x {
2526 return Some(0);
2527 }
2528 let v = binomial_exact(-y - 1, -x - 1)?;
2529 return if (y - x) % 2 == 0 { Some(v) } else { v.checked_neg() };
2530 }
2531 binomial_exact(x, y)
2532}
2533
2534fn binomial_exact(x: i64, y: i64) -> Option<i64> {
2537 if x > BINOMIAL_PRODUCT_LIMIT {
2538 return None;
2539 }
2540 let mut c: i128 = 1;
2541 for i in 1..=x as i128 {
2542 c = c.checked_mul(y as i128 - i + 1)? / i;
2543 if c == 0 {
2544 break;
2545 }
2546 }
2547 i64::try_from(c).ok()
2548}
2549
2550fn binomial(x: f64, y: f64) -> f64 {
2552 if x.fract() == 0.0 && x.abs() < 1e17 {
2553 let xi = x as i64;
2554 if xi < 0 {
2555 if y.fract() == 0.0 && y < 0.0 && y >= x {
2556 let sign = if (y as i64 - xi) % 2 == 0 { 1.0 } else { -1.0 };
2557 return sign * binomial_product(-y as i64 - 1, -x - 1.0);
2558 }
2559 return 0.0;
2560 }
2561 if xi <= BINOMIAL_PRODUCT_LIMIT {
2562 return binomial_product(xi, y);
2563 }
2564 }
2565 gamma(y + 1.0) / (gamma(x + 1.0) * gamma(y - x + 1.0))
2566}
2567
2568#[inline]
2571fn i64_op(op: ScalarDyad, a: i64, b: i64) -> Option<i64> {
2572 use ScalarDyad::*;
2573 Some(match op {
2574 Add => a.checked_add(b)?,
2575 Sub => a.checked_sub(b)?,
2576 Mul => a.checked_mul(b)?,
2577 Min => a.min(b),
2578 Max => a.max(b),
2579 Residue => {
2580 if a == 0 {
2581 b
2582 } else {
2583 let mut r = b.wrapping_rem(a);
2585 if r != 0 && (r < 0) != (a < 0) {
2586 r += a;
2587 }
2588 r
2589 }
2590 }
2591 Pow => {
2592 if b < 0 {
2593 return None;
2594 }
2595 a.checked_pow(u32::try_from(b).ok()?)?
2596 }
2597 Binomial => binomial_i64(a, b)?,
2598 _ => return None,
2599 })
2600}
2601
2602#[inline]
2604fn f64_op(op: ScalarDyad, a: f64, b: f64, span: Span) -> Result<f64> {
2605 use ScalarDyad::*;
2606 Ok(match op {
2607 Add => a + b,
2608 Sub => a - b,
2609 Mul => a * b,
2610 Min => a.min(b),
2611 Max => a.max(b),
2612 DivJ => {
2613 if b == 0.0 {
2614 if a == 0.0 { 0.0 } else { f64::INFINITY.copysign(a) }
2615 } else {
2616 a / b
2617 }
2618 }
2619 DivApl => {
2620 if b == 0.0 {
2621 if a == 0.0 {
2622 1.0
2623 } else {
2624 return Err(Error::domain("division by zero", span));
2625 }
2626 } else {
2627 a / b
2628 }
2629 }
2630 Pow => {
2631 if a == 0.0 && b == 0.0 {
2632 1.0
2633 } else {
2634 a.powf(b)
2635 }
2636 }
2637 Residue => {
2638 if a.is_infinite() {
2643 if b == 0.0 || (b > 0.0) == (a > 0.0) { b } else { a }
2644 } else if a == 0.0 {
2645 b
2646 } else {
2647 b - a * (b / a).floor()
2648 }
2649 }
2650 Log => {
2651 if a < 0.0 || b < 0.0 {
2652 return Err(Error::not_yet("complex numbers", span));
2653 }
2654 b.ln() / a.ln()
2655 }
2656 Root => {
2657 if b < 0.0 {
2658 return Err(Error::not_yet("complex numbers", span));
2659 }
2660 b.powf(1.0 / a)
2661 }
2662 Circle => return circle(a, b, span),
2663 Binomial => binomial(a, b),
2664 _ => return Err(Error::internal("non-arithmetic op in the float path")),
2665 })
2666}
2667
2668#[inline]
2672fn escapes_reals(op: ScalarDyad, a: f64, b: f64) -> bool {
2673 use ScalarDyad::*;
2674 match op {
2675 Pow => a < 0.0 && b.fract() != 0.0,
2677 Log => a < 0.0 || b < 0.0,
2678 Root => b < 0.0,
2679 Circle => circle_escapes(a, b),
2680 _ => false,
2681 }
2682}
2683
2684#[inline]
2687fn circle_escapes(k: f64, y: f64) -> bool {
2688 if k.fract() != 0.0 {
2689 return false;
2690 }
2691 match k as i64 {
2692 0 | -1 | -2 | -7 => y.abs() > 1.0,
2693 -4 => y.abs() < 1.0,
2694 -6 => y < 1.0,
2695 8 | -8 | -11 | -12 => true,
2698 _ => false,
2699 }
2700}
2701
2702#[inline]
2712fn circle(k: f64, y: f64, span: Span) -> Result<f64> {
2713 if k.fract() != 0.0 {
2714 return Err(Error::domain("the circle function needs an integer left argument", span));
2715 }
2716 let complex = || Error::internal("a circle function left the reals on the real path");
2717 Ok(match k as i64 {
2718 0 => {
2719 if y.abs() > 1.0 {
2720 return Err(complex());
2721 }
2722 (1.0 - y * y).max(0.0).sqrt()
2723 }
2724 1 => y.sin(),
2725 2 => y.cos(),
2726 3 => y.tan(),
2727 4 => (1.0 + y * y).sqrt(),
2728 5 => y.sinh(),
2729 6 => y.cosh(),
2730 7 => y.tanh(),
2731 -1 => {
2732 if y.abs() > 1.0 {
2733 return Err(complex());
2734 }
2735 y.asin()
2736 }
2737 -2 => {
2738 if y.abs() > 1.0 {
2739 return Err(complex());
2740 }
2741 y.acos()
2742 }
2743 -3 => y.atan(),
2744 -4 => {
2745 if y.abs() < 1.0 {
2746 return Err(complex());
2747 }
2748 y.signum() * (y * y - 1.0).max(0.0).sqrt()
2750 }
2751 -5 => y.asinh(),
2752 -6 => {
2753 if y < 1.0 {
2754 return Err(complex());
2755 }
2756 y.acosh()
2757 }
2758 -7 => {
2759 if y.abs() > 1.0 {
2760 return Err(complex());
2761 }
2762 y.atanh()
2763 }
2764 9 | -9 | -10 => y,
2766 10 => y.abs(),
2767 11 => 0.0,
2768 12 => {
2769 if y < 0.0 {
2770 std::f64::consts::PI
2771 } else {
2772 0.0
2773 }
2774 }
2775 8 | -8 | -11 | -12 => return Err(complex()),
2776 _ => {
2777 return Err(Error::domain(
2778 "the circle functions run from _12 to 12",
2779 span,
2780 ));
2781 }
2782 })
2783}
2784
2785#[inline]
2787fn cx_op(op: ScalarDyad, a: Cx, b: Cx, span: Span) -> Result<Cx> {
2788 use ScalarDyad::*;
2789 Ok(match op {
2790 Add => cx::add(a, b),
2791 Sub => cx::sub(a, b),
2792 Mul => cx::mul(a, b),
2793 DivJ => cx::div(a, b),
2794 DivApl => {
2795 if b == cx::ZERO {
2796 if a == cx::ZERO {
2797 cx::ONE
2798 } else {
2799 return Err(Error::domain("division by zero", span));
2800 }
2801 } else {
2802 cx::div(a, b)
2803 }
2804 }
2805 Pow => cx::pow(a, b),
2806 Log => cx::log(a, b),
2807 Root => cx::root(a, b),
2808 Residue => cx::residue(a, b),
2809 Lcm => cx::lcm(a, b),
2810 Gcd => cx::gcd(a, b),
2811 MakeComplex => cx::add(a, cx::mul(cx::I, b)),
2812 PolarBy => cx::mul(a, cx::exp(cx::mul(cx::I, b))),
2813 Circle => {
2814 if a[1] != 0.0 || a[0].fract() != 0.0 {
2815 return Err(Error::domain(
2816 "the circle function needs an integer left argument",
2817 span,
2818 ));
2819 }
2820 cx::circle(a[0] as i64, b).ok_or_else(|| {
2821 Error::domain("the circle functions run from _12 to 12", span)
2822 })?
2823 }
2824 Min | Max => return Err(no_complex_order(span)),
2825 Binomial => {
2826 return Err(Error::not_yet("the binomial function on complex numbers", span));
2827 }
2828 Eq | Ne | Lt | Le | Gt | Ge => {
2829 return Err(Error::internal("a comparison in the complex arithmetic path"));
2830 }
2831 })
2832}
2833
2834fn no_complex_order(span: Span) -> Error {
2837 Error::new(
2838 ErrorKind::Domain,
2839 "complex numbers have no order; only equality (=, ~:) applies to them",
2840 Some(span),
2841 )
2842}
2843
2844#[allow(clippy::too_many_arguments)]
2845#[inline(always)]
2846fn dyad_cx_chunk_body<A: Widen<Cx>, B: Widen<Cx>>(
2847 op: ScalarDyad,
2848 xs: &[A],
2849 xoff: usize,
2850 xdiv: usize,
2851 ys: &[B],
2852 yoff: usize,
2853 ydiv: usize,
2854 start: usize,
2855 out: &mut [Cx],
2856 span: Span,
2857) -> Result<()> {
2858 use ScalarDyad::*;
2859 macro_rules! plain {
2862 ($step:expr) => {{
2863 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut Cx| {
2864 *slot = $step(a.widen(), b.widen());
2865 true
2866 });
2867 return Ok(());
2868 }};
2869 }
2870 match op {
2871 Add => plain!(cx::add),
2872 Sub => plain!(cx::sub),
2873 Mul => plain!(cx::mul),
2874 DivJ => plain!(cx::div),
2875 _ => {}
2876 }
2877 let mut err = None;
2878 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut Cx| {
2879 match cx_op(op, a.widen(), b.widen(), span) {
2880 Ok(v) => {
2881 *slot = v;
2882 true
2883 }
2884 Err(e) => {
2885 err = Some(e);
2886 false
2887 }
2888 }
2889 });
2890 match err {
2891 Some(e) => Err(e),
2892 None => Ok(()),
2893 }
2894}
2895
2896multiversioned! {
2897 #[allow(clippy::too_many_arguments)]
2900 fn dyad_cx_chunk[A: Widen<Cx>, B: Widen<Cx>](
2901 op: ScalarDyad,
2902 xs: &[A],
2903 xoff: usize,
2904 xdiv: usize,
2905 ys: &[B],
2906 yoff: usize,
2907 ydiv: usize,
2908 start: usize,
2909 out: &mut [Cx],
2910 span: Span,
2911 ) -> Result<()> = dyad_cx_chunk_body;
2912}
2913
2914#[allow(clippy::too_many_arguments)]
2915fn dyad_cx<A: Widen<Cx>, B: Widen<Cx>>(
2916 op: ScalarDyad,
2917 xs: &[A],
2918 xoff: usize,
2919 xdiv: usize,
2920 ys: &[B],
2921 yoff: usize,
2922 ydiv: usize,
2923 n: usize,
2924 span: Span,
2925) -> Result<Vec<Cx>> {
2926 par::try_fill(n, |start, part| {
2927 dyad_cx_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part, span)
2928 })
2929}
2930
2931#[allow(clippy::too_many_arguments)]
2940fn complex_dyad_data(
2941 op: ScalarDyad,
2942 x: &Data,
2943 xoff: usize,
2944 xdiv: usize,
2945 y: &Data,
2946 yoff: usize,
2947 ydiv: usize,
2948 n: usize,
2949 span: Span,
2950) -> Result<Data> {
2951 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2952 macro_rules! pass {
2953 ($xs:expr, $ys:expr) => {
2954 Data::Complex(dyad_cx(op, $xs, xoff, xdiv, $ys, yoff, ydiv, n, span)?.into())
2955 };
2956 }
2957 Ok(match (x, y) {
2958 (Data::Complex(a), _) => {
2959 let xs: &[Cx] = a;
2960 cx_source!(y, ty, ys, pass!(xs, ys))
2961 }
2962 (_, Data::Complex(b)) => {
2963 let ys: &[Cx] = b;
2964 cx_source!(x, tx, xs, pass!(xs, ys))
2965 }
2966 _ => pass!(borrow_cx(x, &mut tx), borrow_cx(y, &mut ty)),
2967 })
2968}
2969
2970fn circle_reads_a_part(x: &Data, xoff: usize, xdiv: usize, n: usize) -> bool {
2975 if x.dtype() == DType::Complex {
2976 return false;
2978 }
2979 let mut tmp = Vec::new();
2980 let xs = borrow_f64(x, &mut tmp);
2981 (0..n).all(|i| {
2982 let k = xs[xoff + i / xdiv];
2983 k.fract() == 0.0 && (9.0..=12.0).contains(&k)
2984 })
2985}
2986
2987#[allow(clippy::too_many_arguments)]
2990fn pass_leaves_reals(
2991 op: ScalarDyad,
2992 x: &Data,
2993 xoff: usize,
2994 xdiv: usize,
2995 y: &Data,
2996 yoff: usize,
2997 ydiv: usize,
2998 n: usize,
2999) -> bool {
3000 use ScalarDyad::*;
3001 if !matches!(op, Pow | Log | Root | Circle) {
3002 return false;
3003 }
3004 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3005 let xs = borrow_f64(x, &mut tx);
3006 let ys = borrow_f64(y, &mut ty);
3007 (0..n).any(|i| escapes_reals(op, xs[xoff + i / xdiv], ys[yoff + i / ydiv]))
3008}
3009
3010#[allow(clippy::too_many_arguments)]
3011#[inline(always)]
3012fn dyad_i64_chunk_body<A: Widen<i64>, B: Widen<i64>>(
3013 op: ScalarDyad,
3014 xs: &[A],
3015 xoff: usize,
3016 xdiv: usize,
3017 ys: &[B],
3018 yoff: usize,
3019 ydiv: usize,
3020 start: usize,
3021 out: &mut [i64],
3022) -> bool {
3023 use ScalarDyad::*;
3024 macro_rules! overflowing {
3028 ($m:ident) => {{
3029 let mut over = false;
3030 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
3031 let (v, o) = i64::$m(a.widen(), b.widen());
3032 *slot = v;
3033 over |= o;
3034 true
3035 });
3036 !over
3037 }};
3038 }
3039 macro_rules! plain {
3040 ($step:expr) => {{
3041 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
3042 *slot = $step(a.widen(), b.widen());
3043 true
3044 })
3045 }};
3046 }
3047 match op {
3048 Add => overflowing!(overflowing_add),
3049 Sub => overflowing!(overflowing_sub),
3050 Mul => overflowing!(overflowing_mul),
3051 Min => plain!(i64::min),
3052 Max => plain!(i64::max),
3053 _ => zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
3054 match i64_op(op, a.widen(), b.widen()) {
3055 Some(v) => {
3056 *slot = v;
3057 true
3058 }
3059 None => false,
3060 }
3061 }),
3062 }
3063}
3064
3065multiversioned! {
3066 #[allow(clippy::too_many_arguments)]
3073 fn dyad_i64_chunk[A: Widen<i64>, B: Widen<i64>](
3074 op: ScalarDyad,
3075 xs: &[A],
3076 xoff: usize,
3077 xdiv: usize,
3078 ys: &[B],
3079 yoff: usize,
3080 ydiv: usize,
3081 start: usize,
3082 out: &mut [i64],
3083 ) -> bool = dyad_i64_chunk_body;
3084}
3085
3086#[allow(clippy::too_many_arguments)]
3088fn dyad_i64<A: Widen<i64>, B: Widen<i64>>(
3089 op: ScalarDyad,
3090 xs: &[A],
3091 xoff: usize,
3092 xdiv: usize,
3093 ys: &[B],
3094 yoff: usize,
3095 ydiv: usize,
3096 n: usize,
3097) -> Option<Vec<i64>> {
3098 let (out, ok) = par::fill(n, |start, part| {
3099 dyad_i64_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part)
3100 });
3101 ok.then_some(out)
3102}
3103
3104#[allow(clippy::too_many_arguments)]
3107fn int_dyad_data(
3108 op: ScalarDyad,
3109 x: &Data,
3110 xoff: usize,
3111 xdiv: usize,
3112 y: &Data,
3113 yoff: usize,
3114 ydiv: usize,
3115 n: usize,
3116) -> Option<Data> {
3117 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3118 let out = i64_source!(x, tx, xs, {
3119 i64_source!(y, ty, ys, dyad_i64(op, xs, xoff, xdiv, ys, yoff, ydiv, n))
3120 })?;
3121 Some(Data::I64(out.into()))
3122}
3123
3124#[allow(clippy::too_many_arguments)]
3125#[inline(always)]
3126fn dyad_f64_chunk_body<A: Widen<f64>, B: Widen<f64>>(
3127 op: ScalarDyad,
3128 xs: &[A],
3129 xoff: usize,
3130 xdiv: usize,
3131 ys: &[B],
3132 yoff: usize,
3133 ydiv: usize,
3134 start: usize,
3135 out: &mut [f64],
3136 span: Span,
3137) -> Result<()> {
3138 use ScalarDyad::*;
3139 macro_rules! plain {
3142 ($step:expr) => {{
3143 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut f64| {
3144 *slot = $step(a.widen(), b.widen());
3145 true
3146 });
3147 return Ok(());
3148 }};
3149 }
3150 match op {
3151 Add => plain!(|a: f64, b: f64| a + b),
3152 Sub => plain!(|a: f64, b: f64| a - b),
3153 Mul => plain!(|a: f64, b: f64| a * b),
3154 Min => plain!(f64::min),
3155 Max => plain!(f64::max),
3156 _ => {}
3157 }
3158 let mut err = None;
3159 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut f64| {
3160 match f64_op(op, a.widen(), b.widen(), span) {
3161 Ok(v) => {
3162 *slot = v;
3163 true
3164 }
3165 Err(e) => {
3166 err = Some(e);
3167 false
3168 }
3169 }
3170 });
3171 match err {
3172 Some(e) => Err(e),
3173 None => Ok(()),
3174 }
3175}
3176
3177multiversioned! {
3178 #[allow(clippy::too_many_arguments)]
3181 fn dyad_f64_chunk[A: Widen<f64>, B: Widen<f64>](
3182 op: ScalarDyad,
3183 xs: &[A],
3184 xoff: usize,
3185 xdiv: usize,
3186 ys: &[B],
3187 yoff: usize,
3188 ydiv: usize,
3189 start: usize,
3190 out: &mut [f64],
3191 span: Span,
3192 ) -> Result<()> = dyad_f64_chunk_body;
3193}
3194
3195#[allow(clippy::too_many_arguments)]
3196fn dyad_f64<A: Widen<f64>, B: Widen<f64>>(
3197 op: ScalarDyad,
3198 xs: &[A],
3199 xoff: usize,
3200 xdiv: usize,
3201 ys: &[B],
3202 yoff: usize,
3203 ydiv: usize,
3204 n: usize,
3205 span: Span,
3206) -> Result<Vec<f64>> {
3207 par::try_fill(n, |start, part| {
3208 dyad_f64_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part, span)
3209 })
3210}
3211
3212#[allow(clippy::too_many_arguments)]
3214fn float_dyad_data(
3215 op: ScalarDyad,
3216 x: &Data,
3217 xoff: usize,
3218 xdiv: usize,
3219 y: &Data,
3220 yoff: usize,
3221 ydiv: usize,
3222 n: usize,
3223 span: Span,
3224) -> Result<Data> {
3225 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3226 let out = f64_source!(x, tx, xs, {
3227 f64_source!(y, ty, ys, dyad_f64(op, xs, xoff, xdiv, ys, yoff, ydiv, n, span)?)
3228 });
3229 Ok(Data::F64(out.into()))
3230}
3231
3232fn crossed_types(a: DType, b: DType) -> bool {
3236 let class = |d: DType| match d {
3237 DType::Box => 3,
3238 DType::Symbol => 2,
3239 DType::Char => 1,
3240 _ => 0,
3241 };
3242 class(a) != class(b)
3243}
3244
3245#[allow(clippy::too_many_arguments)]
3248fn symbol_min_max(
3249 op: ScalarDyad,
3250 x: &Data,
3251 xoff: usize,
3252 xdiv: usize,
3253 y: &Data,
3254 yoff: usize,
3255 ydiv: usize,
3256 n: usize,
3257 span: Span,
3258) -> Result<Data> {
3259 let (Data::Symbol(a), Data::Symbol(b)) = (x, y) else {
3260 return Err(symbol_arith(span));
3261 };
3262 let down = op == ScalarDyad::Min;
3263 let (out, _) = par::fill(n, |start, part: &mut [crate::symbol::Id]| {
3264 zip_chunk(a, xoff, xdiv, b, yoff, ydiv, start, part, |p, q, slot| {
3265 *slot = if crate::symbol::cmp(p, q).is_le() == down { p } else { q };
3266 true
3267 })
3268 });
3269 Ok(Data::Symbol(out.into()))
3270}
3271
3272#[allow(clippy::too_many_arguments)]
3273fn compare_data(
3274 op: ScalarDyad,
3275 x: &Data,
3276 xoff: usize,
3277 xdiv: usize,
3278 y: &Data,
3279 yoff: usize,
3280 ydiv: usize,
3281 n: usize,
3282 tol: Tol,
3283 span: Span,
3284) -> Result<Data> {
3285 use ScalarDyad::*;
3286 let (dx, dy) = (x.dtype(), y.dtype());
3287 let equality = matches!(op, Eq | Ne);
3288 let boxed = dx == DType::Box || dy == DType::Box;
3294 if equality && crossed_types(dx, dy) && (!boxed || tol.is_j()) {
3295 let unequal = op == Ne;
3296 return Ok(Data::Bool(vec![u8::from(unequal); n].into()));
3297 }
3298 if boxed {
3299 if !equality {
3302 return Err(box_arith(span));
3303 }
3304 let (Data::Box(a), Data::Box(b)) = (x, y) else {
3305 return Err(Error::not_yet("a scalar function inside a nested array", span));
3308 };
3309 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
3310 for (k, slot) in part.iter_mut().enumerate() {
3311 let i = start + k;
3312 let e = arrays_match(&a[xoff + i / xdiv], &b[yoff + i / ydiv], tol);
3313 *slot = u8::from(if op == Eq { e } else { !e });
3314 }
3315 true
3316 });
3317 return Ok(Data::Bool(out.into()));
3318 }
3319 if dx == DType::Symbol || dy == DType::Symbol {
3320 if dx != dy {
3323 return Err(Error::new(
3324 ErrorKind::Type,
3325 "cannot compare a symbol with data that is not a symbol",
3326 Some(span),
3327 ));
3328 }
3329 let (Data::Symbol(a), Data::Symbol(b)) = (x, y) else {
3330 return Err(Error::internal("symbol comparison on non-symbol data"));
3331 };
3332 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
3334 zip_chunk(a, xoff, xdiv, b, yoff, ydiv, start, part, |p, q, slot| {
3335 *slot = u8::from(match op {
3336 Eq => p == q,
3337 Ne => p != q,
3338 _ => {
3339 let o = crate::symbol::cmp(p, q);
3340 match op {
3341 Lt => o.is_lt(),
3342 Le => o.is_le(),
3343 Gt => o.is_gt(),
3344 _ => o.is_ge(),
3345 }
3346 }
3347 });
3348 true
3349 })
3350 });
3351 return Ok(Data::Bool(out.into()));
3352 }
3353 if dx == DType::Char || dy == DType::Char {
3354 if dx != dy {
3355 return Err(Error::new(
3356 ErrorKind::Type,
3357 "cannot compare character and numeric data",
3358 Some(span),
3359 ));
3360 }
3361 if !equality {
3362 return Err(Error::new(
3363 ErrorKind::Type,
3364 "cannot order character data; only equality applies",
3365 Some(span),
3366 ));
3367 }
3368 let (Data::Char(a), Data::Char(b)) = (x, y) else {
3369 return Err(Error::internal("character comparison on non-character data"));
3370 };
3371 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
3372 zip_chunk(a, xoff, xdiv, b, yoff, ydiv, start, part, |p, q, slot| {
3373 let e = p == q;
3374 *slot = if op == Eq { e as u8 } else { !e as u8 };
3375 true
3376 })
3377 });
3378 return Ok(Data::Bool(out.into()));
3379 }
3380 if DType::promote(dx, dy).is_some_and(DType::is_exact)
3381 && let Some(d) = exact_compare_data(op, x, xoff, xdiv, y, yoff, ydiv, n)
3382 {
3383 return Ok(d);
3384 }
3385 if dx == DType::Complex || dy == DType::Complex {
3386 if !equality {
3387 return Err(no_complex_order(span));
3388 }
3389 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3390 let out = cx_source!(x, tx, xs, {
3391 cx_source!(y, ty, ys, {
3392 par::fill(n, |start, part: &mut [u8]| {
3393 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
3394 let e = tol.eq_cx(a.widen(), b.widen());
3395 *slot = if op == Eq { e as u8 } else { !e as u8 };
3396 true
3397 })
3398 })
3399 .0
3400 })
3401 });
3402 return Ok(Data::Bool(out.into()));
3403 }
3404 let out = if DType::promote(dx, dy) == Some(DType::F64) {
3407 let (mut tx, mut ty) = (Vec::<f64>::new(), Vec::<f64>::new());
3408 f64_source!(x, tx, xs, {
3409 f64_source!(y, ty, ys, {
3410 par::fill(n, |start, part: &mut [u8]| {
3411 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
3412 *slot = tol_cmp(op, a.widen(), b.widen(), tol) as u8;
3413 true
3414 })
3415 })
3416 .0
3417 })
3418 })
3419 } else {
3420 let (mut tx, mut ty) = (Vec::<i64>::new(), Vec::<i64>::new());
3421 i64_source!(x, tx, xs, {
3422 i64_source!(y, ty, ys, {
3423 par::fill(n, |start, part: &mut [u8]| {
3424 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
3425 let (a, b): (i64, i64) = (a.widen(), b.widen());
3426 *slot = cmp_result(op, Some(i64::cmp(&a, &b))) as u8;
3427 true
3428 })
3429 })
3430 .0
3431 })
3432 })
3433 };
3434 Ok(Data::Bool(out.into()))
3435}
3436
3437#[inline(always)]
3439pub(crate) fn tol_cmp(op: ScalarDyad, a: f64, b: f64, tol: Tol) -> bool {
3440 use ScalarDyad::*;
3441 match op {
3442 Eq => tol.eq(a, b),
3443 Ne => !tol.eq(a, b),
3444 Lt => tol.lt(a, b),
3445 Le => tol.le(a, b),
3446 Gt => tol.lt(b, a),
3447 Ge => tol.le(b, a),
3448 _ => false,
3449 }
3450}
3451
3452fn cmp_result(op: ScalarDyad, ord: Option<std::cmp::Ordering>) -> bool {
3454 use std::cmp::Ordering::*;
3455 use ScalarDyad::*;
3456 match ord {
3457 None => matches!(op, Ne),
3458 Some(o) => match op {
3459 Eq => o == Equal,
3460 Ne => o != Equal,
3461 Lt => o == Less,
3462 Le => o != Greater,
3463 Gt => o == Greater,
3464 Ge => o != Less,
3465 _ => false,
3466 },
3467 }
3468}
3469
3470fn gcd_i128(a: i128, b: i128) -> i128 {
3472 let (mut a, mut b) = (a.abs(), b.abs());
3473 while b != 0 {
3474 let t = a % b;
3475 a = b;
3476 b = t;
3477 }
3478 a
3479}
3480
3481fn decimal_parts(v: f64) -> Option<(i128, u32)> {
3485 if !v.is_finite() {
3486 return None;
3487 }
3488 let text = format!("{v:e}");
3489 let (mantissa, exponent) = text.split_once('e')?;
3490 let exponent: i32 = exponent.parse().ok()?;
3491 let (whole, fraction) = mantissa.split_once('.').unwrap_or((mantissa, ""));
3492 let mut digits: i128 = format!("{whole}{fraction}").parse().ok()?;
3493 let mut scale = fraction.len() as i32 - exponent;
3494 while scale < 0 {
3497 digits = digits.checked_mul(10)?;
3498 scale += 1;
3499 }
3500 (scale <= 34).then_some((digits, scale as u32))
3503}
3504
3505fn gcd_decimal(a: f64, b: f64) -> Option<f64> {
3510 let (pa, sa) = decimal_parts(a)?;
3511 let (pb, sb) = decimal_parts(b)?;
3512 let scale = sa.max(sb);
3513 let lift = |p: i128, s: u32| 10i128.checked_pow(scale - s).and_then(|k| p.checked_mul(k));
3514 let g = gcd_i128(lift(pa, sa)?, lift(pb, sb)?);
3515 format!("{g}e-{scale}").parse().ok()
3518}
3519
3520fn gcd_f64(a: f64, b: f64, tol: Tol) -> Option<f64> {
3526 let (mut a, mut b) = (a.abs(), b.abs());
3527 if !a.is_finite() || !b.is_finite() {
3528 return None;
3529 }
3530 for _ in 0..1000 {
3533 if b == 0.0 {
3534 return Some(a);
3535 }
3536 if a == 0.0 {
3537 return Some(b);
3538 }
3539 let q = a / b;
3544 let mut k = q.floor();
3545 if tol.eq(q, k + 1.0) {
3546 k += 1.0;
3547 }
3548 let mut r = a - b * k;
3549 if r <= 0.0 || tol.eq(r, b) {
3550 r = 0.0;
3551 }
3552 a = b;
3553 b = r;
3554 }
3555 Some(a)
3556}
3557
3558#[allow(clippy::too_many_arguments)]
3562fn real_lcm_gcd(
3563 op: ScalarDyad,
3564 xs: &[f64],
3565 xoff: usize,
3566 xdiv: usize,
3567 ys: &[f64],
3568 yoff: usize,
3569 ydiv: usize,
3570 n: usize,
3571 tol: Tol,
3572 span: Span,
3573) -> Result<Data> {
3574 let mut out = vec![0.0f64; n];
3575 let mut ok = true;
3576 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, 0, &mut out, |a, b, slot| {
3577 let Some(g) = gcd_decimal(a, b).or_else(|| gcd_f64(a, b, tol)) else {
3578 ok = false;
3579 return false;
3580 };
3581 *slot = if op == ScalarDyad::Gcd {
3582 g
3583 } else if g == 0.0 {
3584 0.0
3585 } else {
3586 a / g * b
3587 };
3588 true
3589 });
3590 if !ok {
3591 return Err(Error::domain("LCM/GCD needs finite values", span));
3592 }
3593 Ok(Data::F64(out.into()))
3594}
3595
3596#[allow(clippy::too_many_arguments)]
3600fn lcm_gcd_data(
3601 op: ScalarDyad,
3602 x: &Data,
3603 xoff: usize,
3604 xdiv: usize,
3605 y: &Data,
3606 yoff: usize,
3607 ydiv: usize,
3608 n: usize,
3609 tol: Tol,
3610 span: Span,
3611) -> Result<Data> {
3612 let t = arith_type(x.dtype(), y.dtype(), span)?;
3613 if t == DType::Complex {
3614 return complex_dyad_data(op, x, xoff, xdiv, y, yoff, ydiv, n, span);
3616 }
3617 if t.is_exact()
3618 && let Some(d) = exact_dyad_data(op, t, x, xoff, xdiv, y, yoff, ydiv, n, span)?
3619 {
3620 return Ok(d);
3621 }
3622 let both_bool = x.dtype() == DType::Bool && y.dtype() == DType::Bool;
3623 let float = t == DType::F64;
3624 let (xs, ys) = if float {
3625 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3626 let xf = borrow_f64(x, &mut tx);
3627 let yf = borrow_f64(y, &mut ty);
3628 let integral = |v: &[f64]| v.iter().all(|&a| a.fract() == 0.0 && fits_i64(a));
3629 if !integral(xf) || !integral(yf) {
3630 return real_lcm_gcd(op, xf, xoff, xdiv, yf, yoff, ydiv, n, tol, span);
3631 }
3632 (
3633 xf.iter().map(|&a| a as i64).collect::<Vec<_>>(),
3634 yf.iter().map(|&a| a as i64).collect::<Vec<_>>(),
3635 )
3636 } else {
3637 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3638 (borrow_i64(x, &mut tx).to_vec(), borrow_i64(y, &mut ty).to_vec())
3639 };
3640 let (out, fits) = par::fill(n, |start, part: &mut [i128]| {
3643 let mut fits = true;
3644 zip_chunk(&xs, xoff, xdiv, &ys, yoff, ydiv, start, part, |a, b, slot| {
3645 let (a, b) = (a as i128, b as i128);
3646 let g = gcd_i128(a, b);
3647 let v = if op == ScalarDyad::Gcd {
3648 g
3649 } else if g == 0 {
3650 0
3651 } else {
3652 a / g * b
3653 };
3654 fits &= i64::try_from(v).is_ok();
3655 *slot = v;
3656 true
3657 });
3658 fits
3659 });
3660 if !fits || float {
3661 return Ok(Data::F64(par::map(&out, |&v| v as f64).into()));
3662 }
3663 if both_bool {
3664 return Ok(Data::Bool(par::map(&out, |&v| v as u8).into()));
3665 }
3666 Ok(Data::I64(par::map(&out, |&v| v as i64).into()))
3667}
3668
3669fn to_rat_vec(d: &Data) -> Option<Vec<Rat>> {
3674 Some(match d {
3675 Data::Bool(v) => v.iter().map(|&b| Rat::from_int(Ext::from(b))).collect(),
3676 Data::I64(v) => v.iter().map(|&x| Rat::from_int(Ext::from(x))).collect(),
3677 Data::Ext(v) => v.iter().map(|x| Rat::from_int(x.clone())).collect(),
3678 Data::Rat(v) => v.to_vec(),
3679 Data::F64(_) | Data::Complex(_) | Data::Char(_) | Data::Symbol(_) | Data::Box(_) => {
3680 return None;
3681 }
3682 })
3683}
3684
3685fn rat_window(d: &Data, off: usize, div: usize, n: usize) -> Option<Vec<Rat>> {
3693 if n == 0 {
3694 return Some(Vec::new());
3695 }
3696 let end = off + (n - 1) / div + 1;
3697 if off == 0 && end == d.len() {
3698 return to_rat_vec(d);
3699 }
3700 to_rat_vec(&d.slice(off, end))
3701}
3702
3703fn exact_data(t: DType, out: Vec<Rat>) -> Data {
3711 if t == DType::Ext && out.iter().all(Rat::is_integer) {
3712 return Data::Ext(out.iter().map(|r| r.to_int().expect("whole")).collect());
3713 }
3714 Data::Rat(out.into())
3715}
3716
3717fn too_large(span: Span) -> Error {
3719 Error::domain(
3720 format!(
3721 "the exact result needs more than {} bits; use floats for a value this large",
3722 exact::MAX_BITS
3723 ),
3724 span,
3725 )
3726}
3727
3728fn exact_pow(a: &Rat, b: &Rat, span: Span) -> Result<Option<Rat>> {
3731 let Some(e) = b.to_int().as_ref().and_then(exact::ext_to_i64) else {
3732 return Ok(None);
3733 };
3734 if let Some(v) = a.pow(e) {
3735 return Ok(Some(v));
3736 }
3737 if a.is_zero() && e < 0 { Ok(None) } else { Err(too_large(span)) }
3739}
3740
3741#[allow(clippy::too_many_arguments)]
3745fn exact_dyad_data(
3746 op: ScalarDyad,
3747 t: DType,
3748 x: &Data,
3749 xoff: usize,
3750 xdiv: usize,
3751 y: &Data,
3752 yoff: usize,
3753 ydiv: usize,
3754 n: usize,
3755 span: Span,
3756) -> Result<Option<Data>> {
3757 use ScalarDyad::*;
3758 let (Some(xs), Some(ys)) = (rat_window(x, xoff, xdiv, n), rat_window(y, yoff, ydiv, n))
3759 else {
3760 return Ok(None);
3761 };
3762 let mut out = Vec::with_capacity(n);
3763 for i in 0..n {
3764 let a = &xs[i / xdiv];
3765 let b = &ys[i / ydiv];
3766 let v = match op {
3767 Add => a.add(b),
3768 Sub => a.sub(b),
3769 Mul => a.mul(b),
3770 DivJ | DivApl => match a.div(b) {
3772 Some(v) => v,
3773 None => return Ok(None),
3774 },
3775 Min => a.min(b).clone(),
3776 Max => a.max(b).clone(),
3777 Residue => exact::rat_residue(a, b),
3778 Gcd => exact::rat_gcd(a, b),
3779 Lcm => exact::rat_lcm(a, b),
3780 Pow => match exact_pow(a, b, span)? {
3781 Some(v) => v,
3782 None => return Ok(None),
3783 },
3784 Binomial => match (a.to_int(), b.to_int()) {
3785 (Some(k), Some(m)) => match exact::ext_binomial(&k, &m) {
3786 Some(v) => Rat::from_int(v),
3787 None => return Ok(None),
3788 },
3789 _ => return Ok(None),
3790 },
3791 Root if t == DType::Ext => {
3794 let (Some(k), Some(m)) = (a.to_int(), b.to_int()) else {
3795 return Ok(None);
3796 };
3797 let Some(k) = exact::ext_to_i64(&k).and_then(|k| u32::try_from(k).ok()) else {
3798 return Ok(None);
3799 };
3800 match exact::exact_root(k, &m) {
3801 Some(v) => Rat::from_int(v),
3802 None => return Ok(None),
3803 }
3804 }
3805 Root | Log | Circle | MakeComplex | PolarBy => return Ok(None),
3806 Eq | Ne | Lt | Le | Gt | Ge => return Ok(None),
3808 };
3809 out.push(v);
3810 }
3811 Ok(Some(exact_data(t, out)))
3812}
3813
3814fn exact_monad(op: ScalarMonad, y: &Array) -> Option<Array> {
3817 use ScalarMonad::*;
3818 let v = to_rat_vec(&y.data)?;
3819 let shape = y.shape.clone();
3820 if matches!(op, Floor | Ceil | Signum) {
3823 let out: Vec<Ext> = v
3824 .iter()
3825 .map(|r| match op {
3826 Floor => r.floor(),
3827 Ceil => r.ceil(),
3828 _ => r.signum(),
3829 })
3830 .collect();
3831 return Some(Array::new(shape, Data::Ext(out.into())).with_layout(y.layout()));
3832 }
3833 let two = Rat::from_int(Ext::from(2));
3834 let mut out = Vec::with_capacity(v.len());
3835 for r in &v {
3836 let value = match op {
3837 Conj => r.clone(),
3838 Neg => r.neg(),
3839 Abs => r.abs(),
3840 Recip => r.recip()?,
3841 Inc => r.add(&Rat::one()),
3842 Dec => r.sub(&Rat::one()),
3843 OneMinus => Rat::one().sub(r),
3844 Double => r.add(r),
3845 Halve => r.div(&two).expect("two is not zero"),
3846 Square => r.mul(r),
3847 Sqrt => r.sqrt()?,
3848 Factorial => Rat::from_int(r.to_int().as_ref().and_then(exact::ext_factorial)?),
3849 Exp | Ln | Pi | Imaginary | Polar | Not => return None,
3852 Floor | Ceil | Signum => unreachable!("handled above"),
3853 };
3854 out.push(value);
3855 }
3856 Some(Array::new(shape, exact_data(y.dtype(), out)).with_layout(y.layout()))
3857}
3858
3859fn to_exact(y: &Array, span: Span) -> Result<Array> {
3864 let data = match &y.data {
3865 Data::Ext(_) | Data::Rat(_) => return Ok(y.clone()),
3866 Data::Bool(v) => Data::Ext(v.iter().map(|&b| Ext::from(b)).collect()),
3867 Data::I64(v) => Data::Ext(v.iter().map(|&x| Ext::from(x)).collect()),
3868 Data::F64(v) => {
3869 let mut out = Vec::with_capacity(v.len());
3870 for &x in v.iter() {
3871 out.push(exact::f64_to_rat(x).ok_or_else(|| {
3872 Error::domain("an infinity has no exact value", span)
3873 })?);
3874 }
3875 exact_data(DType::Ext, out)
3876 }
3877 Data::Complex(_) | Data::Char(_) | Data::Symbol(_) | Data::Box(_) => {
3878 return Err(Error::domain(
3879 format!("x: needs real numbers, not {} data", y.dtype().name()),
3880 span,
3881 ));
3882 }
3883 };
3884 Ok(Array::new(y.shape.clone(), data).with_layout(y.layout()))
3885}
3886
3887fn from_exact(y: &Array) -> Array {
3890 let shape = y.shape.clone();
3891 match &y.data {
3892 Data::Ext(v) => match v.iter().map(exact::ext_to_i64).collect::<Option<Vec<i64>>>() {
3893 Some(out) => Array::new(shape, Data::I64(out.into())).with_layout(y.layout()),
3894 None => Array::new(shape, Data::F64(v.iter().map(exact::ext_to_f64).collect()))
3895 .with_layout(y.layout()),
3896 },
3897 Data::Rat(v) => Array::new(shape, Data::F64(v.iter().map(Rat::to_f64).collect()))
3898 .with_layout(y.layout()),
3899 _ => y.clone(),
3900 }
3901}
3902
3903fn exact_form(x: &Array, y: &Array, span: Span) -> Result<Array> {
3905 match one_whole(x, "the form x: converts to", span)? {
3906 1 => {
3907 let e = to_exact(y, span)?;
3908 e.cast(DType::Rat).ok_or_else(|| Error::internal("an exact value has no rational form"))
3909 }
3910 2 => {
3911 let e = to_exact(y, span)?;
3912 let v = to_rat_vec(&e.data).ok_or_else(|| Error::internal("x: gave an inexact value"))?;
3913 let mut out = Vec::with_capacity(2 * v.len());
3914 for r in &v {
3915 out.push(r.numer().clone());
3916 out.push(r.denom().clone());
3917 }
3918 let mut shape = y.shape.clone();
3919 shape.push(2);
3920 Ok(Array::new(shape, Data::Ext(out.into())))
3921 }
3922 -1 => Ok(from_exact(y)),
3923 -2 => {
3925 if !y.dtype().is_numeric() {
3926 return Err(Error::domain(
3927 format!("x: needs real numbers, not {} data", y.dtype().name()),
3928 span,
3929 ));
3930 }
3931 Ok(y.clone())
3932 }
3933 n => Err(Error::domain(
3934 format!("x: converts to form 1, 2, _1 or _2, not {n}"),
3935 span,
3936 )),
3937 }
3938}
3939
3940#[allow(clippy::too_many_arguments)]
3944fn exact_compare_data(
3945 op: ScalarDyad,
3946 x: &Data,
3947 xoff: usize,
3948 xdiv: usize,
3949 y: &Data,
3950 yoff: usize,
3951 ydiv: usize,
3952 n: usize,
3953) -> Option<Data> {
3954 let (xs, ys) = (rat_window(x, xoff, xdiv, n)?, rat_window(y, yoff, ydiv, n)?);
3955 let out: Vec<u8> = (0..n)
3956 .map(|i| {
3957 let ord = xs[i / xdiv].cmp(&ys[i / ydiv]);
3958 cmp_result(op, Some(ord)) as u8
3959 })
3960 .collect();
3961 Some(Data::Bool(out.into()))
3962}
3963
3964#[allow(clippy::too_many_arguments)]
3968fn scalar_dyad_data(
3969 op: ScalarDyad,
3970 x: &Data,
3971 xoff: usize,
3972 xdiv: usize,
3973 y: &Data,
3974 yoff: usize,
3975 ydiv: usize,
3976 n: usize,
3977 tol: Tol,
3978 span: Span,
3979) -> Result<Data> {
3980 use ScalarDyad::*;
3981 if x.dtype() == DType::Symbol || y.dtype() == DType::Symbol {
3982 match op {
3983 Eq | Ne | Lt | Le | Gt | Ge => {}
3985 Min | Max => {
3988 return symbol_min_max(op, x, xoff, xdiv, y, yoff, ydiv, n, span);
3989 }
3990 _ => return Err(symbol_arith(span)),
3991 }
3992 }
3993 if matches!(op, Eq | Ne | Lt | Le | Gt | Ge) {
3994 return compare_data(op, x, xoff, xdiv, y, yoff, ydiv, n, tol, span);
3995 }
3996 if matches!(op, Lcm | Gcd) {
3997 return lcm_gcd_data(op, x, xoff, xdiv, y, yoff, ydiv, n, tol, span);
3998 }
3999 let t = arith_type(x.dtype(), y.dtype(), span)?;
4000 if t.is_exact()
4001 && let Some(d) = exact_dyad_data(op, t, x, xoff, xdiv, y, yoff, ydiv, n, span)?
4002 {
4003 return Ok(d);
4004 }
4005 if t == DType::I64 && !matches!(op, DivJ | DivApl | Log | Root | Circle) {
4007 if let Some(d) = int_dyad_data(op, x, xoff, xdiv, y, yoff, ydiv, n) {
4010 return Ok(d);
4011 }
4012 }
4014 if t == DType::Complex
4015 || matches!(op, MakeComplex | PolarBy)
4016 || pass_leaves_reals(op, x, xoff, xdiv, y, yoff, ydiv, n)
4017 {
4018 let data = complex_dyad_data(op, x, xoff, xdiv, y, yoff, ydiv, n, span)?;
4019 if op == Circle && circle_reads_a_part(x, xoff, xdiv, n) && let Data::Complex(v) = &data {
4020 return Ok(Data::F64(v.iter().map(|z| z[0]).collect()));
4021 }
4022 return Ok(data);
4023 }
4024 float_dyad_data(op, x, xoff, xdiv, y, yoff, ydiv, n, span)
4025}
4026
4027fn frame_pervaded(frame: Vec<usize>, cells: Vec<Array>, span: Span) -> Result<Array> {
4032 if cells.iter().all(|c| c.rank() == 0 && c.dtype() != DType::Box) {
4033 return assemble(&frame, cells, span);
4034 }
4035 let boxes: Vec<Array> = cells.into_iter().collect();
4036 Ok(Array::new(frame, Data::Box(boxes.into())))
4037}
4038
4039fn pervade_dyad(
4045 op: ScalarDyad,
4046 x: &Array,
4047 y: &Array,
4048 cfg: EvalCfg,
4049 span: Span,
4050) -> Result<Array> {
4051 let p = agree(&x.shape, &y.shape, &x.shape, &y.shape, cfg.agreement, span)?;
4052 if p.n == 0 {
4053 return Ok(Array::new(p.frame, Data::empty(DType::Box)));
4054 }
4055 let (xr, yr) = (x.to_row_major(), y.to_row_major());
4056 let mut cells = Vec::with_capacity(p.n);
4057 for i in 0..p.n {
4058 let a = open_cell(&atom(&xr, i / p.x_div));
4059 let b = open_cell(&atom(&yr, i / p.y_div));
4060 cells.push(scalar_dyad(op, &a, &b, cfg, span)?);
4061 }
4062 frame_pervaded(p.frame, cells, span)
4063}
4064
4065fn pervade_monad(op: ScalarMonad, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
4067 if y.count() == 0 {
4068 return Ok(Array::new(y.shape.clone(), Data::empty(DType::Box)));
4069 }
4070 let yr = y.to_row_major();
4071 let mut cells = Vec::with_capacity(y.count());
4072 for i in 0..y.count() {
4073 let a = open_cell(&atom(&yr, i));
4074 cells.push(scalar_monad(op, &a, cfg, span)?);
4075 }
4076 frame_pervaded(y.shape.clone(), cells, span)
4077}
4078
4079fn scalar_dyad(
4080 op: ScalarDyad,
4081 x: &Array,
4082 y: &Array,
4083 cfg: EvalCfg,
4084 span: Span,
4085) -> Result<Array> {
4086 if cfg.rules.lang == crate::Lang::Apl
4087 && (x.dtype() == DType::Box || y.dtype() == DType::Box)
4088 {
4089 return pervade_dyad(op, x, y, cfg, span);
4090 }
4091 let p = agree(&x.shape, &y.shape, &x.shape, &y.shape, cfg.agreement, span)?;
4092 if p.n == 0 {
4096 return Ok(Array::new(p.frame, Data::empty(empty_result_type(x, y))));
4097 }
4098 let data =
4099 scalar_dyad_data(op, &x.data, 0, p.x_div, &y.data, 0, p.y_div, p.n, cfg.tol, span)?;
4100 Ok(Array::new(p.frame, data))
4101}
4102
4103fn empty_result_type(x: &Array, y: &Array) -> DType {
4106 for a in [x, y] {
4107 if a.dtype().is_numeric() {
4108 return a.dtype();
4109 }
4110 }
4111 DType::I64
4112}
4113
4114fn fits_i64(v: f64) -> bool {
4116 v.is_finite() && v >= i64::MIN as f64 && v < i64::MAX as f64
4117}
4118
4119fn monad_leaves_reals(op: ScalarMonad, d: &Data) -> bool {
4121 use ScalarMonad::*;
4122 match op {
4123 Imaginary | Polar => d.dtype().is_numeric(),
4125 Sqrt | Ln => match d {
4126 Data::I64(v) => par::any(v, |&x| x < 0),
4127 Data::F64(v) => par::any(v, |&x| x < 0.0),
4128 Data::Ext(v) => v.iter().any(|x| x.sign() == num_bigint::Sign::Minus),
4129 Data::Rat(v) => v.iter().any(|x| x < &Rat::zero()),
4130 _ => false,
4131 },
4132 _ => false,
4133 }
4134}
4135
4136fn complex_monad(op: ScalarMonad, y: &Array, span: Span) -> Result<Array> {
4138 use ScalarMonad::*;
4139 let mut tmp = Vec::new();
4140 let v = borrow_cx(&y.data, &mut tmp);
4141 if y.count() > 0 && v.is_empty() {
4142 return Err(wrong_type(y.dtype(), span));
4143 }
4144 let data = match op {
4145 Abs => Data::F64(par::map(v, |&z| cx::abs(z)).into()),
4147 Not => return Err(Error::domain("logical negation needs values of 0 or 1", span)),
4148 Factorial => {
4149 return Err(Error::not_yet("the factorial of a complex number", span));
4150 }
4151 _ => {
4152 let step: fn(Cx) -> Cx = match op {
4153 Conj => cx::conj,
4154 Neg => cx::neg,
4155 Signum => cx::signum,
4156 Recip => cx::recip,
4157 Sqrt => cx::sqrt,
4158 Exp => cx::exp,
4159 Ln => cx::ln,
4160 Floor => cx::floor,
4161 Ceil => cx::ceil,
4162 OneMinus => |z| cx::sub(cx::ONE, z),
4163 Inc => |z| cx::add(z, cx::ONE),
4164 Dec => |z| cx::sub(z, cx::ONE),
4165 Double => |z| cx::add(z, z),
4166 Halve => |z| [z[0] / 2.0, z[1] / 2.0],
4167 Square => |z| cx::mul(z, z),
4168 Pi => |z| [std::f64::consts::PI * z[0], std::f64::consts::PI * z[1]],
4169 Imaginary => |z| cx::mul(cx::I, z),
4170 Polar => |z| cx::exp(cx::mul(cx::I, z)),
4171 Abs | Not | Factorial => unreachable!("handled above"),
4172 };
4173 Data::Complex(par::map(v, |&z| step(z)).into())
4174 }
4175 };
4176 Ok(Array::new(y.shape.clone(), data).with_layout(y.layout()))
4177}
4178
4179fn scalar_monad(op: ScalarMonad, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
4181 use ScalarMonad::*;
4182 if cfg.rules.lang == crate::Lang::Apl && y.dtype() == DType::Box {
4183 return pervade_monad(op, y, cfg, span);
4184 }
4185 let tol = cfg.tol;
4186 let d = &y.data;
4187 if y.count() == 0 && !d.dtype().is_numeric() {
4190 return Ok(Array::new(y.shape.clone(), Data::empty(DType::I64)));
4191 }
4192 if d.dtype() == DType::Complex || monad_leaves_reals(op, d) {
4193 return complex_monad(op, y, span);
4194 }
4195 if d.dtype().is_exact() && let Some(a) = exact_monad(op, y) {
4196 return Ok(a);
4197 }
4198 let mut tmp = Vec::new();
4202 let data = match op {
4203 Conj if d.dtype().is_numeric() => d.clone(),
4205 Conj => return Err(wrong_type(d.dtype(), span)),
4206 Imaginary | Polar => return Err(Error::internal("a complex monad on the real path")),
4209 Neg => match d {
4210 Data::Bool(v) => Data::I64(par::map(v, |&b| -(b as i64)).into()),
4211 Data::I64(v) => match par::try_map(v, i64::checked_neg) {
4212 Some(out) => Data::I64(out.into()),
4213 None => Data::F64(par::map(v, |&x| -(x as f64)).into()),
4214 },
4215 Data::F64(v) => Data::F64(par::map(v, |&x| -x).into()),
4216 _ => return Err(wrong_type(d.dtype(), span)),
4217 },
4218 Signum => match d {
4219 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64).into()),
4220 Data::I64(v) => Data::I64(par::map(v, |&x| x.signum()).into()),
4221 Data::F64(v) => Data::F64(
4224 par::map(v, |&x| {
4225 if tol.is_zero(x) {
4226 0.0
4227 } else if x > 0.0 {
4228 1.0
4229 } else if x < 0.0 {
4230 -1.0
4231 } else {
4232 0.0
4233 }
4234 })
4235 .into(),
4236 ),
4237 _ => return Err(wrong_type(d.dtype(), span)),
4238 },
4239 Recip => {
4240 let v = as_f64(d, &mut tmp, span)?;
4244 Data::F64(par::map(v, |&x| if x == 0.0 { f64::INFINITY } else { 1.0 / x }).into())
4245 }
4246 Sqrt => {
4247 let v = as_f64(d, &mut tmp, span)?;
4249 Data::F64(par::map(v, |&x| x.sqrt()).into())
4250 }
4251 Exp => {
4252 let v = as_f64(d, &mut tmp, span)?;
4253 Data::F64(par::map(v, |&x| x.exp()).into())
4254 }
4255 Abs => match d {
4256 Data::Bool(_) => d.clone(),
4257 Data::I64(v) => match par::try_map(v, i64::checked_abs) {
4258 Some(out) => Data::I64(out.into()),
4259 None => Data::F64(par::map(v, |&x| (x as f64).abs()).into()),
4260 },
4261 Data::F64(v) => Data::F64(par::map(v, |&x| x.abs()).into()),
4262 _ => return Err(wrong_type(d.dtype(), span)),
4263 },
4264 Floor | Ceil => match d {
4265 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64).into()),
4266 Data::I64(_) => d.clone(),
4267 Data::F64(v) => {
4268 let round = |x: f64| if op == Floor { tol.floor(x) } else { tol.ceil(x) };
4269 match par::try_map(v, |x| {
4271 let r = round(x);
4272 fits_i64(r).then_some(r as i64)
4273 }) {
4274 Some(out) => Data::I64(out.into()),
4275 None => Data::F64(par::map(v, |&x| round(x)).into()),
4276 }
4277 }
4278 _ => return Err(wrong_type(d.dtype(), span)),
4279 },
4280 Inc | Dec => {
4281 let step = if op == Inc { 1i64 } else { -1 };
4282 match d {
4283 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64 + step).into()),
4284 Data::I64(v) => match par::try_map(v, |x: i64| x.checked_add(step)) {
4285 Some(out) => Data::I64(out.into()),
4286 None => Data::F64(par::map(v, |&x| x as f64 + step as f64).into()),
4287 },
4288 Data::F64(v) => Data::F64(par::map(v, |&x| x + step as f64).into()),
4289 _ => return Err(wrong_type(d.dtype(), span)),
4290 }
4291 }
4292 Double | Square => match d {
4293 Data::Bool(v) => {
4294 Data::I64(par::map(v, |&b| if op == Double { 2 * b as i64 } else { b as i64 }).into())
4295 }
4296 Data::I64(v) => {
4297 let f = |x: i64| if op == Double { x.checked_mul(2) } else { x.checked_mul(x) };
4298 match par::try_map(v, f) {
4299 Some(out) => Data::I64(out.into()),
4300 None => Data::F64(
4301 par::map(v, |&x| {
4302 let x = x as f64;
4303 if op == Double { x + x } else { x * x }
4304 })
4305 .into(),
4306 ),
4307 }
4308 }
4309 Data::F64(v) => {
4310 Data::F64(par::map(v, |&x| if op == Double { x + x } else { x * x }).into())
4311 }
4312 _ => return Err(wrong_type(d.dtype(), span)),
4313 },
4314 Halve => {
4315 let v = as_f64(d, &mut tmp, span)?;
4316 Data::F64(par::map(v, |&x| x / 2.0).into())
4317 }
4318 Pi => {
4319 let v = as_f64(d, &mut tmp, span)?;
4320 Data::F64(par::map(v, |&x| std::f64::consts::PI * x).into())
4321 }
4322 Factorial => {
4323 let v = as_f64(d, &mut tmp, span)?;
4324 Data::F64(par::map(v, |&x| factorial(x)).into())
4325 }
4326 Ln => {
4327 let v = as_f64(d, &mut tmp, span)?;
4329 Data::F64(par::map(v, |&x| x.ln()).into())
4331 }
4332 OneMinus => match d {
4333 Data::Bool(v) => Data::Bool(par::map(v, |&b| 1 - b).into()),
4334 Data::I64(v) => match par::try_map(v, |x: i64| 1i64.checked_sub(x)) {
4335 Some(out) => Data::I64(out.into()),
4336 None => Data::F64(par::map(v, |&x| 1.0 - x as f64).into()),
4337 },
4338 Data::F64(v) => Data::F64(par::map(v, |&x| 1.0 - x).into()),
4339 _ => return Err(wrong_type(d.dtype(), span)),
4340 },
4341 Not => {
4342 let bad = || Error::domain("logical negation needs values of 0 or 1", span);
4343 match d {
4344 Data::Bool(v) => Data::Bool(par::map(v, |&b| 1 - b).into()),
4345 Data::I64(v) => {
4346 let out = par::try_map(v, |x: i64| match x {
4347 0 => Some(1u8),
4348 1 => Some(0u8),
4349 _ => None,
4350 })
4351 .ok_or_else(bad)?;
4352 Data::Bool(out.into())
4353 }
4354 Data::F64(v) => {
4355 let out = par::try_map(v, |x: f64| {
4356 if x == 0.0 {
4357 Some(1u8)
4358 } else if x == 1.0 {
4359 Some(0u8)
4360 } else {
4361 None
4362 }
4363 })
4364 .ok_or_else(bad)?;
4365 Data::Bool(out.into())
4366 }
4367 _ => return Err(bad()),
4368 }
4369 }
4370 };
4371 Ok(Array::new(y.shape.clone(), data).with_layout(y.layout()))
4372}
4373
4374fn transpose_axes(y: &Array) -> Array {
4384 if y.rank() < 2 {
4385 return y.clone();
4386 }
4387 let out_shape: Vec<usize> = y.shape.iter().rev().copied().collect();
4388 let flipped = match y.layout() {
4389 Layout::RowMajor => Layout::ColMajor,
4390 Layout::ColMajor => Layout::RowMajor,
4391 };
4392 Array::new(out_shape, y.data.clone()).with_layout(flipped)
4393}
4394
4395fn iota_j(y: &Array, span: Span) -> Result<Array> {
4398 if y.rank() > 1 {
4399 return Err(Error::new(
4400 ErrorKind::Rank,
4401 "index generator needs a scalar or vector argument",
4402 Some(span),
4403 ));
4404 }
4405 let dims = y
4406 .to_i64_vec()
4407 .ok_or_else(|| Error::domain("index generator needs integer lengths", span))?;
4408 let shape: Vec<usize> = dims.iter().map(|d| d.unsigned_abs() as usize).collect();
4409 let n = crate::limits::elements(&shape, span)?;
4410 let st = strides(&shape);
4411 let mut out = Vec::with_capacity(n);
4412 let mut coord = vec![0usize; shape.len()];
4413 for _ in 0..n {
4414 let mut v = 0usize;
4415 for k in 0..shape.len() {
4416 let c = if dims[k] < 0 { shape[k] - 1 - coord[k] } else { coord[k] };
4417 v += c * st[k];
4418 }
4419 out.push(v as i64);
4420 odometer(&mut coord, &shape);
4421 }
4422 let data = Data::I64(out.into());
4423 let data = if y.dtype() == DType::Ext {
4426 data.cast(DType::Ext).ok_or_else(|| Error::internal("integers have no extended form"))?
4427 } else {
4428 data
4429 };
4430 Ok(Array::new(shape, data))
4431}
4432
4433fn head(y: &Array) -> Array {
4435 if y.rank() == 0 {
4436 return y.clone();
4437 }
4438 if y.items() == 0 {
4439 let cell_shape = y.shape[1..].to_vec();
4440 let n: usize = cell_shape.iter().product();
4441 return Array::new(cell_shape, fill_data(y.dtype(), n));
4442 }
4443 y.item(0)
4444}
4445
4446fn behead(y: &Array, span: Span) -> Result<Array> {
4447 if y.rank() == 0 {
4448 return Err(Error::domain("cannot drop the first item of a scalar", span));
4449 }
4450 if y.items() == 0 {
4451 return Ok(y.clone());
4452 }
4453 let m = y.item_size();
4454 let mut shape = y.shape.clone();
4455 shape[0] -= 1;
4456 Ok(Array::new(shape, y.data.slice(m, y.count())))
4457}
4458
4459fn tail(y: &Array) -> Array {
4461 if y.rank() == 0 {
4462 return y.clone();
4463 }
4464 let n = y.items();
4465 if n == 0 {
4466 let cell_shape = y.shape[1..].to_vec();
4467 let m: usize = cell_shape.iter().product();
4468 return Array::new(cell_shape, fill_data(y.dtype(), m));
4469 }
4470 y.item(n - 1)
4471}
4472
4473fn curtail(y: &Array) -> Array {
4475 if y.rank() == 0 {
4476 return Array::empty(y.dtype());
4477 }
4478 let n = y.items();
4479 if n == 0 {
4480 return y.clone();
4481 }
4482 let m = y.item_size();
4483 let mut shape = y.shape.clone();
4484 shape[0] = n - 1;
4485 Array::new(shape, y.data.slice(0, (n - 1) * m))
4486}
4487
4488fn reverse(y: &Array) -> Array {
4490 if y.rank() == 0 {
4491 return y.clone();
4492 }
4493 let n = y.items();
4494 let m = y.item_size();
4495 let mut data = Data::empty(y.dtype());
4496 for i in (0..n).rev() {
4497 for k in 0..m {
4498 push_elem(&mut data, &y.data, i * m + k);
4499 }
4500 }
4501 Array::new(y.shape.clone(), data)
4502}
4503
4504fn rotate(x: &Array, y: &Array, span: Span) -> Result<Array> {
4507 let counts = axis_counts(x, "rotate", span)?;
4508 if y.rank() == 0 {
4509 return Ok(y.clone());
4510 }
4511 if counts.len() > y.rank() {
4512 return Err(Error::new(
4513 ErrorKind::Length,
4514 format!(
4515 "rotate has {} amounts for an argument of rank {}",
4516 counts.len(),
4517 y.rank()
4518 ),
4519 Some(span),
4520 ));
4521 }
4522 let st = strides(&y.shape);
4523 let n = y.count();
4524 let r = y.rank();
4525 let mut data = Data::empty(y.dtype());
4526 let mut coord = vec![0usize; r];
4527 for _ in 0..n {
4528 let mut idx = 0usize;
4529 for k in 0..r {
4530 let len = y.shape[k] as i64;
4532 let s = counts.get(k).copied().unwrap_or(0);
4533 idx += (coord[k] as i64 + s).rem_euclid(len) as usize * st[k];
4534 }
4535 push_elem(&mut data, &y.data, idx);
4536 odometer(&mut coord, &y.shape);
4537 }
4538 Ok(Array::new(y.shape.clone(), data))
4539}
4540
4541fn elem_key(d: &Data, i: usize) -> u64 {
4544 match d {
4545 Data::Bool(v) => v[i] as u64,
4546 Data::I64(v) => v[i] as u64,
4547 Data::F64(v) => {
4548 let x = v[i];
4549 if x == 0.0 { 0 } else { x.to_bits() }
4550 }
4551 Data::Complex(v) => cx_key(v[i]),
4552 Data::Char(v) => v[i] as u64,
4553 Data::Symbol(v) => v[i] as u64,
4555 Data::Ext(_) | Data::Rat(_) | Data::Box(_) => 0,
4558 }
4559}
4560
4561fn num_key(d: &Data, i: usize) -> u64 {
4564 match d {
4565 Data::Bool(v) => (v[i] as f64).to_bits(),
4566 Data::I64(v) => (v[i] as f64).to_bits(),
4567 Data::F64(v) => {
4568 let x = v[i];
4569 if x == 0.0 { 0.0f64.to_bits() } else { x.to_bits() }
4570 }
4571 Data::Complex(v) => cx_key(v[i]),
4572 Data::Char(v) => v[i] as u64,
4573 Data::Symbol(v) => v[i] as u64,
4574 Data::Ext(_) | Data::Rat(_) | Data::Box(_) => 0,
4576 }
4577}
4578
4579fn cx_key(z: Cx) -> u64 {
4581 let bits = |x: f64| if x == 0.0 { 0u64 } else { x.to_bits() };
4582 bits(z[0]) ^ bits(z[1]).rotate_left(32)
4583}
4584
4585fn nub(y: &Array, tol: Tol) -> Array {
4587 if y.rank() == 0 {
4588 return Array::new(vec![1], y.data.clone());
4589 }
4590 let n = y.items();
4591 let m = y.item_size();
4592 let mut keep = Vec::new();
4593 if y.dtype() == DType::Box || y.dtype().is_exact() {
4594 for i in 0..n {
4597 if !keep.iter().any(|&j| arrays_match(&y.item(i), &y.item(j), tol)) {
4598 keep.push(i);
4599 }
4600 }
4601 } else if y.dtype() == DType::F64 && tol.ct != 0.0 {
4602 let mut tv = Vec::new();
4605 let v = borrow_f64(&y.data, &mut tv);
4606 for i in 0..n {
4607 if !keep.iter().any(|&j| (0..m).all(|k| tol.eq(v[i * m + k], v[j * m + k]))) {
4608 keep.push(i);
4609 }
4610 }
4611 } else {
4612 let mut seen: HashSet<Vec<u64>> = HashSet::with_capacity(n);
4613 for i in 0..n {
4614 let key: Vec<u64> = (0..m).map(|k| elem_key(&y.data, i * m + k)).collect();
4615 if seen.insert(key) {
4616 keep.push(i);
4617 }
4618 }
4619 }
4620 let mut data = Data::empty(y.dtype());
4621 for &i in &keep {
4622 for k in 0..m {
4623 push_elem(&mut data, &y.data, i * m + k);
4624 }
4625 }
4626 let mut shape = y.shape.clone();
4627 shape[0] = keep.len();
4628 Array::new(shape, data)
4629}
4630
4631#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4636enum Tao {
4637 J,
4638 Apl2,
4639}
4640
4641impl Tao {
4642 fn of(rules: Rules) -> Tao {
4643 match rules.lang {
4644 crate::Lang::J => Tao::J,
4645 crate::Lang::Apl => Tao::Apl2,
4648 }
4649 }
4650
4651 fn class(self, dt: DType) -> u8 {
4657 match self {
4658 Tao::J => match dt {
4659 DType::Symbol => 1,
4660 DType::Char => 2,
4661 DType::Box => 3,
4662 _ => 0,
4663 },
4664 Tao::Apl2 => match dt {
4665 DType::Char | DType::Symbol => 0,
4666 DType::Box => 2,
4667 _ => 1,
4668 },
4669 }
4670 }
4671}
4672
4673fn cmp_items_total(x: &Array, y: &Array, tao: Tao) -> std::cmp::Ordering {
4688 use std::cmp::Ordering::Equal;
4689 match tao {
4690 Tao::J => {
4691 let class = |a: &Array| if a.count() == 0 { 0 } else { tao.class(a.dtype()) };
4692 class(x)
4693 .cmp(&class(y))
4694 .then_with(|| x.rank().cmp(&y.rank()))
4695 .then_with(|| x.shape.iter().rev().cmp(y.shape.iter().rev()))
4696 .then_with(|| cmp_atoms(x, y, tao))
4697 }
4698 Tao::Apl2 => x
4699 .rank()
4700 .cmp(&y.rank())
4701 .then_with(|| x.shape.iter().cmp(y.shape.iter()))
4702 .then_with(|| cmp_atoms(x, y, tao))
4703 .then_with(|| {
4704 if x.count() == 0 {
4705 tao.class(x.dtype()).cmp(&tao.class(y.dtype()))
4706 } else {
4707 Equal
4708 }
4709 }),
4710 }
4711}
4712
4713fn cmp_atoms(x: &Array, y: &Array, tao: Tao) -> std::cmp::Ordering {
4716 use std::cmp::Ordering::Equal;
4717 let n = x.count();
4718 if n == 0 {
4719 return Equal;
4720 }
4721 let (xr, yr) = (x.to_row_major(), y.to_row_major());
4722 let (dx, dy) = (xr.row_major_data(), yr.row_major_data());
4723 let opened = |d: &Data, i: usize| -> Array {
4724 match d {
4725 Data::Box(v) => v[i].clone(),
4726 _ => {
4727 let mut one = Data::empty(d.dtype());
4728 push_elem(&mut one, d, i);
4729 Array::new(vec![], one)
4730 }
4731 }
4732 };
4733 if matches!(dx, Data::Box(_)) || matches!(dy, Data::Box(_)) {
4734 return (0..n)
4735 .map(|i| cmp_items_total(&opened(dx, i), &opened(dy, i), tao))
4736 .find(|o| *o != Equal)
4737 .unwrap_or(Equal);
4738 }
4739 let classes = tao.class(dx.dtype()).cmp(&tao.class(dy.dtype()));
4741 if classes != Equal {
4742 return classes;
4743 }
4744 match (dx, dy) {
4745 (Data::Char(a), Data::Char(b)) => a[..n].cmp(&b[..n]),
4746 _ => cmp_numbers(dx, dy, n),
4747 }
4748}
4749
4750fn cmp_numbers(dx: &Data, dy: &Data, n: usize) -> std::cmp::Ordering {
4754 use std::cmp::Ordering::Equal;
4755 let seek = |f: &dyn Fn(usize) -> std::cmp::Ordering| {
4756 (0..n).map(f).find(|o| *o != Equal).unwrap_or(Equal)
4757 };
4758 match DType::promote(dx.dtype(), dy.dtype()) {
4759 Some(DType::Complex) => {
4760 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4761 let (a, b) = (borrow_cx(dx, &mut ta), borrow_cx(dy, &mut tb));
4762 seek(&|k| {
4763 a[k][0]
4764 .partial_cmp(&b[k][0])
4765 .unwrap_or(Equal)
4766 .then_with(|| a[k][1].partial_cmp(&b[k][1]).unwrap_or(Equal))
4767 })
4768 }
4769 Some(DType::F64) => {
4770 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4771 let (a, b) = (borrow_f64(dx, &mut ta), borrow_f64(dy, &mut tb));
4772 seek(&|k| a[k].partial_cmp(&b[k]).unwrap_or(Equal))
4773 }
4774 Some(t) if t.is_exact() => match (to_rat_vec(dx), to_rat_vec(dy)) {
4775 (Some(a), Some(b)) => seek(&|k| a[k].cmp(&b[k])),
4776 _ => Equal,
4777 },
4778 None => Equal,
4780 Some(_) => {
4781 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4782 let (a, b) = (borrow_i64(dx, &mut ta), borrow_i64(dy, &mut tb));
4783 seek(&|k| a[k].cmp(&b[k]))
4784 }
4785 }
4786}
4787
4788fn cmp_items(d: &Data, i: usize, j: usize, m: usize, tao: Tao) -> std::cmp::Ordering {
4792 use std::cmp::Ordering::Equal;
4793 let (a, b) = (i * m, j * m);
4794 let ord = |k: usize| match d {
4795 Data::Bool(v) => v[a + k].cmp(&v[b + k]),
4796 Data::I64(v) => v[a + k].cmp(&v[b + k]),
4797 Data::F64(v) => v[a + k].partial_cmp(&v[b + k]).unwrap_or(Equal),
4798 Data::Complex(v) => v[a + k][0]
4805 .partial_cmp(&v[b + k][0])
4806 .unwrap_or(Equal)
4807 .then_with(|| v[a + k][1].partial_cmp(&v[b + k][1]).unwrap_or(Equal)),
4808 Data::Char(v) => v[a + k].cmp(&v[b + k]),
4809 Data::Symbol(v) => crate::symbol::cmp(v[a + k], v[b + k]),
4812 Data::Ext(v) => v[a + k].cmp(&v[b + k]),
4815 Data::Rat(v) => v[a + k].cmp(&v[b + k]),
4816 Data::Box(v) => cmp_items_total(&v[a + k], &v[b + k], tao),
4819 };
4820 (0..m).map(ord).find(|o| *o != Equal).unwrap_or(Equal)
4821}
4822
4823fn grade_order(y: &Array, down: bool, tao: Tao) -> Vec<usize> {
4825 if y.rank() == 0 {
4826 return vec![0];
4827 }
4828 let n = y.items();
4829 let m = y.item_size();
4830 let mut idx: Vec<usize> = (0..n).collect();
4831 if down {
4834 idx.sort_by(|&a, &b| cmp_items(&y.data, b, a, m, tao));
4835 } else {
4836 idx.sort_by(|&a, &b| cmp_items(&y.data, a, b, m, tao));
4837 }
4838 idx
4839}
4840
4841fn collate_grade(x: &Array, y: &Array, down: bool, origin: i64, span: Span) -> Result<Array> {
4846 let chars_of = |a: &Array| -> Result<Vec<char>> {
4847 match a.row_major_data() {
4848 Data::Char(v) => Ok(v.as_slice().to_vec()),
4849 _ => Err(Error::domain("a collating grade takes characters", span)),
4850 }
4851 };
4852 let (xs, ys) = (chars_of(x)?, chars_of(y)?);
4853 let xshape = if x.rank() == 0 { vec![1] } else { x.shape.clone() };
4854 let width = xshape.len();
4855 let absent: Vec<usize> = xshape.iter().rev().copied().collect();
4859 let mut keys: std::collections::HashMap<char, Vec<usize>> =
4860 std::collections::HashMap::new();
4861 let xst = strides(&xshape);
4862 for (i, &c) in xs.iter().enumerate() {
4863 keys.entry(c).or_insert_with(|| {
4864 (0..width).map(|a| (i / xst[a]) % xshape[a]).rev().collect()
4865 });
4866 }
4867 let key_of = |c: char| keys.get(&c).unwrap_or(&absent).clone();
4868 let n = if y.rank() == 0 { 1 } else { y.items() };
4869 let m = if n == 0 { 0 } else { ys.len() / n };
4870 let item_keys: Vec<Vec<usize>> = (0..n)
4871 .map(|i| ys[i * m..(i + 1) * m].iter().flat_map(|&c| key_of(c)).collect())
4872 .collect();
4873 let mut idx: Vec<usize> = (0..n).collect();
4874 if down {
4875 idx.sort_by(|&a, &b| item_keys[b].cmp(&item_keys[a]));
4876 } else {
4877 idx.sort_by(|&a, &b| item_keys[a].cmp(&item_keys[b]));
4878 }
4879 Ok(Array::from_i64(idx.into_iter().map(|i| origin + i as i64).collect()))
4880}
4881
4882fn atomic_rep(y: &Array, ctx: &Ctx<'_>, span: Span) -> Result<Array> {
4886 let name = match y.as_boxes() {
4887 Some([b]) if y.rank() == 0 => crate::gerund::text_of(b),
4888 _ => None,
4889 };
4890 let Some(name) = name else {
4891 return Err(Error::domain("5!:1 takes a boxed name", span));
4892 };
4893 if let Some(v) = ctx.env.verb(&name) {
4894 let ar = crate::gerund::verb_ar(v).ok_or_else(|| {
4895 Error::not_yet(
4896 format!("the atomic representation of {}", v.name()),
4897 span,
4898 )
4899 })?;
4900 return Ok(Array::boxed(ar.to_array()));
4901 }
4902 match ctx.env.get(&name) {
4903 Some(a) => Ok(Array::boxed(crate::gerund::Ar::Noun(a).to_array())),
4904 None => Err(Error::new(
4905 ErrorKind::Value,
4906 format!("undefined name: {name}"),
4907 Some(span),
4908 )),
4909 }
4910}
4911
4912fn catalogue(y: &Array, span: Span) -> Result<Array> {
4916 let items = if y.rank() == 0 { vec![y.clone()] } else { y.cells(1) };
4917 let opened: Vec<Array> = items
4919 .iter()
4920 .map(|it| match it.as_boxes() {
4921 Some(bs) if it.rank() == 0 => bs[0].clone(),
4922 _ => it.clone(),
4923 })
4924 .collect();
4925 let mut shape: Vec<usize> = Vec::new();
4926 for o in &opened {
4927 shape.extend_from_slice(&o.shape);
4928 }
4929 let total: usize = shape.iter().product();
4930 let mut out = Vec::with_capacity(total);
4931 let mut coord = vec![0usize; shape.len()];
4932 for _ in 0..total {
4933 let mut at = 0usize;
4934 let mut picks = Vec::with_capacity(opened.len());
4935 for o in &opened {
4936 let st = strides(&o.shape);
4937 let idx: usize = (0..o.rank()).map(|a| coord[at + a] * st[a]).sum();
4938 at += o.rank();
4939 let mut data = Data::empty(o.dtype());
4940 push_elem(&mut data, o.row_major_data(), idx);
4941 picks.push(Array::new(vec![], data));
4942 }
4943 out.push(assemble(&[picks.len()], picks, span)?);
4944 odometer(&mut coord, &shape);
4945 }
4946 Ok(Array::new(shape, Data::Box(out.into())))
4947}
4948
4949fn raze_in(y: &Array, tol: Tol, span: Span) -> Result<Array> {
4952 let all = raze(y, span)?;
4953 let n = if all.rank() == 0 { 1 } else { all.items() };
4954 let elements: Vec<Array> = (0..y.count())
4955 .map(|i| {
4956 let mut data = Data::empty(y.dtype());
4957 push_elem(&mut data, y.row_major_data(), i);
4958 let one = Array::new(vec![], data);
4959 match one.as_boxes() {
4960 Some(bs) => bs[0].clone(),
4961 None => one,
4962 }
4963 })
4964 .collect();
4965 let mut out = Vec::with_capacity(elements.len() * n);
4966 for e in &elements {
4967 let row = member_j(&all, e, tol);
4968 out.extend_from_slice(row.to_i64_vec().unwrap_or_default().as_slice());
4969 }
4970 let mut shape = y.shape.clone();
4971 shape.push(n);
4972 Ok(Array::new(shape, Data::Bool(out.into_iter().map(|v| v as u8).collect::<Vec<u8>>().into())))
4973}
4974
4975fn select_items(y: &Array, order: &[usize]) -> Array {
4977 let m = y.item_size();
4978 let mut data = Data::empty(y.dtype());
4979 for &i in order {
4980 for k in 0..m {
4981 push_elem(&mut data, &y.data, i * m + k);
4982 }
4983 }
4984 let mut shape = y.shape.clone();
4985 shape[0] = order.len();
4986 Array::new(shape, data)
4987}
4988
4989fn check_gradable(y: &Array, rules: Rules, span: Span) -> Result<()> {
4994 if y.dtype() == DType::Complex && rules.complex_order != ComplexOrder::RealThenImaginary {
4995 return Err(Error::not_yet("grading complex values by magnitude and angle", span));
4996 }
4997 Ok(())
4998}
4999
5000fn grade_select(x: &Array, y: &Array, down: bool, rules: Rules, span: Span) -> Result<Array> {
5004 check_gradable(y, rules, span)?;
5005 let order = grade_order(y, down, Tao::of(rules));
5006 if x.rank() == 0 {
5007 return Ok(x.clone());
5008 }
5009 if let Some(&past) = order.iter().find(|&&i| i >= x.items()) {
5010 return Err(Error::domain(
5011 format!("index {past} is out of range: the argument has {} items", x.items()),
5012 span,
5013 ));
5014 }
5015 Ok(select_items(x, &order))
5016}
5017
5018pub(crate) fn arrays_match(x: &Array, y: &Array, tol: Tol) -> bool {
5021 if x.shape != y.shape {
5022 return false;
5023 }
5024 if x.layout() != y.layout() {
5027 return arrays_match(&x.to_row_major(), &y.to_row_major(), tol);
5028 }
5029 if x.count() == 0 {
5032 return true;
5033 }
5034 if let (Data::Box(a), Data::Box(b)) = (&x.data, &y.data) {
5035 return a.iter().zip(b.iter()).all(|(p, q)| arrays_match(p, q, tol));
5036 }
5037 let (dx, dy) = (x.dtype(), y.dtype());
5038 match DType::promote(dx, dy) {
5039 None => false,
5040 Some(DType::Char) => match (&x.data, &y.data) {
5041 (Data::Char(a), Data::Char(b)) => a.as_slice() == b.as_slice(),
5042 _ => false,
5043 },
5044 Some(DType::Symbol) => match (&x.data, &y.data) {
5047 (Data::Symbol(a), Data::Symbol(b)) => a.as_slice() == b.as_slice(),
5048 _ => false,
5049 },
5050 Some(DType::F64) => {
5051 let (mut ta, mut tb) = (Vec::new(), Vec::new());
5052 let a = borrow_f64(&x.data, &mut ta);
5053 let b = borrow_f64(&y.data, &mut tb);
5054 a.iter().zip(b).all(|(p, q)| tol.eq(*p, *q))
5055 }
5056 Some(DType::Complex) => {
5057 let (mut ta, mut tb) = (Vec::new(), Vec::new());
5058 let a = borrow_cx(&x.data, &mut ta);
5059 let b = borrow_cx(&y.data, &mut tb);
5060 a.iter().zip(b).all(|(p, q)| tol.eq_cx(*p, *q))
5061 }
5062 Some(t) if t.is_exact() => match (to_rat_vec(&x.data), to_rat_vec(&y.data)) {
5063 (Some(a), Some(b)) => a == b,
5064 _ => false,
5065 },
5066 Some(_) => {
5067 let (mut ta, mut tb) = (Vec::new(), Vec::new());
5068 let a = borrow_i64(&x.data, &mut ta);
5069 let b = borrow_i64(&y.data, &mut tb);
5070 a.iter().zip(b).all(|(p, q)| p == q)
5071 }
5072 }
5073}
5074
5075fn item_or_self(a: &Array, i: usize) -> Array {
5077 if a.rank() == 0 { a.clone() } else { a.item(i) }
5078}
5079
5080fn member_j(x: &Array, y: &Array, tol: Tol) -> Array {
5083 let cell_rank = y.rank().saturating_sub(1).min(x.rank());
5084 let frame_rank = x.rank() - cell_rank;
5085 let frame: Vec<usize> = x.shape[..frame_rank].to_vec();
5086 let nf: usize = frame.iter().product();
5087 let items = y.items();
5088 let mut out = Vec::with_capacity(nf);
5089 for i in 0..nf {
5090 let cell = x.cell_at(frame_rank, i);
5091 out.push((0..items).any(|j| arrays_match(&cell, &item_or_self(y, j), tol)) as u8);
5092 }
5093 Array::new(frame, Data::Bool(out.into()))
5094}
5095
5096fn member_apl(x: &Array, y: &Array, tol: Tol) -> Array {
5098 let n = x.count();
5099 if x.dtype() == DType::Box
5100 || y.dtype() == DType::Box
5101 || x.dtype().is_exact()
5102 || y.dtype().is_exact()
5103 {
5104 let opened = |a: &Array, i: usize| -> Array {
5110 let e = atom(a, i);
5111 match e.as_boxes() {
5112 Some([b]) if b.rank() == 0 && b.dtype() != DType::Box => b.clone(),
5113 _ => e,
5114 }
5115 };
5116 let out: Vec<u8> = (0..n)
5117 .map(|i| {
5118 let e = opened(x, i);
5119 u8::from((0..y.count()).any(|j| arrays_match(&e, &opened(y, j), tol)))
5120 })
5121 .collect();
5122 return Array::new(x.shape.clone(), Data::Bool(out.into()));
5123 }
5124 if x.dtype() != y.dtype()
5125 && [x.dtype(), y.dtype()].iter().any(|&d| matches!(d, DType::Char | DType::Symbol))
5126 {
5127 return Array::new(x.shape.clone(), Data::Bool(vec![0u8; n].into()));
5128 }
5129 if tol.ct != 0.0
5130 && (x.dtype() == DType::F64 || y.dtype() == DType::F64)
5131 && x.dtype() != DType::Char
5132 {
5133 let (mut tx, mut ty) = (Vec::new(), Vec::new());
5135 let xs = borrow_f64(&x.data, &mut tx);
5136 let ys = borrow_f64(&y.data, &mut ty);
5137 let out: Vec<u8> =
5138 xs.iter().map(|a| ys.iter().any(|b| tol.eq(*a, *b)) as u8).collect();
5139 return Array::new(x.shape.clone(), Data::Bool(out.into()));
5140 }
5141 let seen: HashSet<u64> = (0..y.count()).map(|i| num_key(&y.data, i)).collect();
5142 let out: Vec<u8> =
5143 (0..n).map(|i| seen.contains(&num_key(&x.data, i)) as u8).collect();
5144 Array::new(x.shape.clone(), Data::Bool(out.into()))
5145}
5146
5147fn index_of(x: &Array, y: &Array, origin: i64, tol: Tol) -> Array {
5149 let cell_rank = x.rank().saturating_sub(1).min(y.rank());
5150 let frame_rank = y.rank() - cell_rank;
5151 let frame: Vec<usize> = y.shape[..frame_rank].to_vec();
5152 let nf: usize = frame.iter().product();
5153 let items = x.items();
5154 let mut out = Vec::with_capacity(nf);
5155 for i in 0..nf {
5156 let cell = y.cell_at(frame_rank, i);
5157 let at = (0..items)
5158 .find(|&j| arrays_match(&cell, &item_or_self(x, j), tol))
5159 .unwrap_or(items);
5160 out.push(origin + at as i64);
5161 }
5162 Array::new(frame, Data::I64(out.into()))
5163}
5164
5165fn from_index(x: &Array, y: &Array, span: Span) -> Result<Array> {
5167 if let Some(spec) = x.as_boxes().and_then(<[Array]>::first) {
5170 let spec = index_spec(spec, y, span)?;
5171 return Ok(select_spec(&spec, y));
5172 }
5173 let idx = x
5174 .to_i64_vec()
5175 .ok_or_else(|| Error::domain("index must be an integer", span))?;
5176 let Some(&i) = idx.first() else {
5177 return Err(Error::internal("from_index with no index"));
5178 };
5179 let n = y.items() as i64;
5180 let k = if i < 0 { i + n } else { i };
5181 if k < 0 || k >= n {
5182 return Err(Error::domain(
5183 format!("index {i} is out of range: the argument has {n} items"),
5184 span,
5185 ));
5186 }
5187 Ok(item_or_self(y, k as usize))
5188}
5189
5190fn cat_promote(a: &Array, other: &Array, rank: usize, axis: usize, span: Span) -> Result<Array> {
5194 if a.rank() == rank {
5195 return Ok(a.clone());
5196 }
5197 if a.rank() == 0 {
5198 let mut shape =
5199 if other.rank() == rank { other.shape.clone() } else { vec![1usize; rank] };
5200 shape[axis] = 1;
5201 let n: usize = shape.iter().product();
5202 let mut data = Data::empty(a.dtype());
5203 for _ in 0..n {
5204 push_elem(&mut data, &a.data, 0);
5205 }
5206 return Ok(Array::new(shape, data));
5207 }
5208 if a.rank() + 1 == rank {
5209 let mut shape = a.shape.clone();
5210 shape.insert(axis, 1);
5211 return Ok(Array::new(shape, a.data.clone()));
5212 }
5213 Err(Error::new(
5214 ErrorKind::Rank,
5215 format!("cannot catenate rank {} with rank {}", a.rank(), other.rank()),
5216 Some(span),
5217 ))
5218}
5219
5220pub(crate) fn catenate(
5222 x: &Array,
5223 y: &Array,
5224 leading: bool,
5225 fill: bool,
5226 span: Span,
5227) -> Result<Array> {
5228 let rank = x.rank().max(y.rank()).max(1);
5229 let axis = if leading { 0 } else { rank - 1 };
5230 let xa = cat_promote(x, y, rank, axis, span)?;
5231 let ya = cat_promote(y, x, rank, axis, span)?;
5232 let mut ragged = false;
5236 let want: Vec<i64> = (0..rank)
5237 .map(|k| {
5238 ragged |= k != axis && xa.shape[k] != ya.shape[k];
5239 xa.shape[k].max(ya.shape[k]) as i64
5240 })
5241 .collect();
5242 if ragged && !fill {
5243 return Err(Error::new(
5244 ErrorKind::Length,
5245 format!(
5246 "cannot catenate: left shape {}, right shape {}",
5247 show_shape(&xa.shape),
5248 show_shape(&ya.shape)
5249 ),
5250 Some(span),
5251 ));
5252 }
5253 let (xa, ya) = if ragged {
5254 let fit = |a: &Array| -> Result<Array> {
5255 let mut to = want.clone();
5256 to[axis] = a.shape[axis] as i64;
5257 take(&Array::from_i64(to), a, false, false, span)
5258 };
5259 (fit(&xa)?, fit(&ya)?)
5260 } else {
5261 (xa, ya)
5262 };
5263 let (xa, ya) = if !fill && (xa.dtype() == DType::Box) != (ya.dtype() == DType::Box) {
5267 (nest_like(&xa, &ya), nest_like(&ya, &xa))
5268 } else {
5269 (xa, ya)
5270 };
5271 let dt = DType::promote(xa.dtype(), ya.dtype()).ok_or_else(|| {
5272 let boxed = xa.dtype() == DType::Box || ya.dtype() == DType::Box;
5273 let what = if boxed {
5274 "cannot catenate boxed and unboxed data; box the other side first"
5275 } else {
5276 "cannot catenate character and numeric data"
5277 };
5278 Error::new(ErrorKind::Type, what, Some(span))
5279 })?;
5280 let widen = |a: &Array| -> Result<Data> {
5281 if a.dtype() == dt {
5282 Ok(a.data.clone())
5283 } else {
5284 a.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening in catenate"))
5285 }
5286 };
5287 let xd = widen(&xa)?;
5288 let yd = widen(&ya)?;
5289 let outer: usize = xa.shape[..axis].iter().product();
5290 let ix: usize = xa.shape[axis..].iter().product();
5291 let iy: usize = ya.shape[axis..].iter().product();
5292 let mut data = Data::empty(dt);
5293 for o in 0..outer {
5294 for k in 0..ix {
5295 push_elem(&mut data, &xd, o * ix + k);
5296 }
5297 for k in 0..iy {
5298 push_elem(&mut data, &yd, o * iy + k);
5299 }
5300 }
5301 let mut shape = xa.shape.clone();
5302 shape[axis] = xa.shape[axis] + ya.shape[axis];
5303 Ok(Array::new(shape, data))
5304}
5305
5306fn copy_items(x: &Array, y: &Array, apl: bool, span: Span) -> Result<Array> {
5314 let counts = x
5315 .to_i64_vec()
5316 .ok_or_else(|| Error::domain("replication counts must be integers", span))?;
5317 if !apl && counts.iter().any(|&c| c < 0) {
5318 return Err(Error::domain("replication counts must be nonnegative", span));
5319 }
5320 let one_item = apl && x.rank() > 0 && y.rank() > 0 && y.items() == 1 && counts.len() != 1;
5324 let scalar_y = y.rank() == 0 || one_item;
5325 let m = y.item_size();
5326 let n = if x.rank() == 0 || !scalar_y { y.items() } else { counts.len() };
5327 let per = if x.rank() == 0 { vec![counts[0]; n] } else { counts };
5328 if per.len() != n {
5329 return Err(Error::new(
5330 ErrorKind::Length,
5331 format!("{} replication count(s) for {n} item(s)", per.len()),
5332 Some(span),
5333 ));
5334 }
5335 let items: u128 = per.iter().map(|&c| c.unsigned_abs() as u128).sum();
5338 let total = crate::limits::count(items * m.max(1) as u128, span)? / m.max(1);
5339 let mut data = Data::empty(y.dtype());
5340 for (i, &c) in per.iter().enumerate() {
5341 let src = if scalar_y { 0 } else { i };
5343 for _ in 0..c.unsigned_abs() {
5344 for k in 0..m {
5345 if c < 0 {
5346 data.push_fill();
5347 } else {
5348 push_elem(&mut data, &y.data, src * m + k);
5349 }
5350 }
5351 }
5352 }
5353 let mut shape = if y.rank() == 0 { vec![1] } else { y.shape.clone() };
5356 shape[0] = total;
5357 Ok(Array::new(shape, data))
5358}
5359
5360fn format_chars(y: &Array, opts: &FmtOpts) -> Array {
5369 if y.dtype() == DType::Char {
5370 return y.clone();
5371 }
5372 if y.count() == 0 {
5374 return Array::new(y.shape.clone(), Data::empty(DType::Char));
5375 }
5376 let text = crate::fmt::format_array(y, opts);
5377 if y.dtype() == DType::Box {
5378 let lines = text.lines().filter(|l| !l.is_empty()).count();
5383 let rows: usize =
5384 if y.rank() == 0 { 1 } else { y.shape[..y.rank() - 1].iter().product() };
5385 if lines != rows {
5386 return text_planes(&text, &y.shape[..y.rank().saturating_sub(2)]);
5387 }
5388 }
5389 if y.rank() < 2 {
5390 let chars: Vec<char> = text.chars().collect();
5391 return Array::new(vec![chars.len()], Data::Char(chars.into()));
5392 }
5393 let lines: Vec<&str> = text.lines().filter(|l| !l.is_empty()).collect();
5396 let width = lines.iter().map(|l| l.chars().count()).max().unwrap_or(0);
5397 let mut chars: Vec<char> = Vec::with_capacity(lines.len() * width);
5398 for line in &lines {
5399 chars.extend(line.chars());
5400 chars.resize(chars.len() + width - line.chars().count(), ' ');
5401 }
5402 let mut shape = y.shape[..y.rank() - 1].to_vec();
5405 shape.push(width);
5406 debug_assert_eq!(lines.len(), shape[..shape.len() - 1].iter().product::<usize>());
5407 Array::new(shape, Data::Char(chars.into()))
5408}
5409
5410fn text_planes(text: &str, frame: &[usize]) -> Array {
5413 let lines: Vec<&str> = text.lines().filter(|l| !l.is_empty()).collect();
5414 let width = lines.iter().map(|l| l.chars().count()).max().unwrap_or(0);
5415 let planes: usize = frame.iter().product::<usize>().max(1);
5416 let per = lines.len() / planes;
5417 let mut chars: Vec<char> = Vec::with_capacity(lines.len() * width);
5418 for line in &lines {
5419 chars.extend(line.chars());
5420 chars.resize(chars.len() + width - line.chars().count(), ' ');
5421 }
5422 let mut shape = frame.to_vec();
5423 shape.push(per);
5424 shape.push(width);
5425 Array::new(shape, Data::Char(chars.into()))
5426}
5427
5428fn digits_of(a: &Array, what: &str, span: Span) -> Result<Vec<f64>> {
5430 a.to_f64_vec().ok_or_else(|| Error::domain(format!("{what} needs numeric data"), span))
5431}
5432
5433fn narrow(values: Vec<f64>, integral: bool) -> Data {
5437 if integral && values.iter().all(|&v| v.fract() == 0.0 && fits_i64(v)) {
5438 return Data::I64(values.iter().map(|&v| v as i64).collect::<Vec<_>>().into());
5439 }
5440 Data::F64(values.into())
5441}
5442
5443fn is_integral(a: &Array) -> bool {
5445 !matches!(a.dtype(), DType::F64 | DType::Rat | DType::Char | DType::Symbol)
5446}
5447
5448fn decode_exact(x: Option<&Array>, y: &Array) -> Option<Array> {
5454 let yr = y.to_row_major();
5455 let digits = to_rat_vec(&yr.data)?;
5456 let two = Rat::from_int(Ext::from(2));
5457 let radix: Vec<Rat> = match x {
5458 None => vec![two; digits.len()],
5459 Some(x) => {
5460 let r = to_rat_vec(&x.to_row_major().data)?;
5461 match r.len() {
5462 1 => vec![r[0].clone(); digits.len()],
5463 n if n == digits.len() => r,
5464 _ => return None,
5465 }
5466 }
5467 };
5468 let mut acc = Rat::from_int(Ext::from(0));
5469 for (d, b) in digits.iter().zip(&radix) {
5470 acc = acc.mul(b).add(d);
5471 }
5472 let exact_in = |a: &Array| matches!(a.dtype(), DType::Ext | DType::Rat);
5473 if exact_in(y) || x.is_some_and(exact_in) {
5474 return Some(Array::new(Vec::new(), exact_data(DType::Ext, vec![acc])));
5475 }
5476 let whole = acc.to_int()?;
5479 Some(Array::scalar_i64(exact::ext_to_i64(&whole)?))
5480}
5481
5482fn decode(x: Option<&Array>, y: &Array, span: Span) -> Result<Array> {
5485 if let Some(exact) = decode_exact(x, y) {
5486 return Ok(exact);
5487 }
5488 let digits = digits_of(y, "decode", span)?;
5489 let radix: Vec<f64> = match x {
5490 None => vec![2.0; digits.len()],
5491 Some(x) => {
5492 let r = digits_of(x, "decode", span)?;
5493 match r.len() {
5494 1 => vec![r[0]; digits.len()],
5495 n if n == digits.len() => r,
5496 n => {
5497 return Err(Error::new(
5498 ErrorKind::Length,
5499 format!("{n} radices for {} digits", digits.len()),
5500 Some(span),
5501 ));
5502 }
5503 }
5504 }
5505 };
5506 let mut acc = 0.0f64;
5507 for (d, b) in digits.iter().zip(&radix) {
5508 acc = acc * b + d;
5509 }
5510 let integral = is_integral(y) && x.is_none_or(is_integral);
5511 Ok(Array::new(vec![], narrow(vec![acc], integral)))
5512}
5513
5514fn decode_apl(x: &Array, y: &Array, span: Span) -> Result<Array> {
5518 let digits = digits_of(y, "decode", span)?;
5519 let radices = digits_of(x, "decode", span)?;
5520 let k = if y.rank() == 0 { 1 } else { y.shape[0] };
5525 let n: usize = if y.rank() == 0 { 1 } else { y.shape[1..].iter().product() };
5526 let (rows, width) = match x.rank() {
5527 0 => (1usize, 0usize),
5528 r => (x.shape[..r - 1].iter().product(), x.shape[r - 1]),
5529 };
5530 if width != 0 && width != k {
5531 return Err(Error::new(
5532 ErrorKind::Length,
5533 format!("{width} radices for {k} digits"),
5534 Some(span),
5535 ));
5536 }
5537 let per_row = if x.rank() > 0 && width == 0 { 0 } else { k };
5541 let mut out = vec![0.0f64; rows * n];
5542 for i in 0..rows {
5543 for j in 0..n {
5544 let mut acc = 0.0f64;
5545 for d in 0..per_row {
5546 let b = if width == 0 { radices[0] } else { radices[i * width + d] };
5547 acc = acc * b + digits[d * n + j];
5548 }
5549 out[i * n + j] = acc;
5550 }
5551 }
5552 let mut shape: Vec<usize> = if x.rank() == 0 {
5553 Vec::new()
5554 } else {
5555 x.shape[..x.rank() - 1].to_vec()
5556 };
5557 if y.rank() > 0 {
5558 shape.extend_from_slice(&y.shape[1..]);
5559 }
5560 let integral = is_integral(y) && is_integral(x);
5561 Ok(Array::new(shape, narrow(out, integral)))
5562}
5563
5564fn encode_apl(x: &Array, y: &Array, span: Span) -> Result<Array> {
5567 let radices = digits_of(x, "encode", span)?;
5568 let values = digits_of(y, "encode", span)?;
5569 let k = if x.rank() == 0 { 1 } else { x.shape[0] };
5570 let frames = if k == 0 { 0 } else { radices.len() / k };
5571 let n = values.len();
5572 let mut out = vec![0.0f64; k * frames * n];
5573 let mut radix = vec![0.0f64; k];
5574 let mut cell = vec![0.0f64; k];
5575 for p in 0..frames {
5576 for (i, r) in radix.iter_mut().enumerate() {
5577 *r = radices[i * frames + p];
5578 }
5579 for (j, &v) in values.iter().enumerate() {
5580 encode_one(&radix, v, &mut cell);
5581 for i in 0..k {
5582 out[(i * frames + p) * n + j] = cell[i];
5583 }
5584 }
5585 }
5586 let mut shape = x.shape.clone();
5587 shape.extend_from_slice(&y.shape);
5588 Ok(Array::new(shape, narrow(out, is_integral(x) && is_integral(y))))
5589}
5590
5591fn bit_width(values: &[f64], span: Span) -> Result<usize> {
5594 if values.is_empty() {
5596 return Ok(0);
5597 }
5598 let mut m = 0.0f64;
5599 for &v in values {
5600 if !v.is_finite() {
5601 return Err(Error::domain("cannot encode an infinite value", span));
5602 }
5603 m = m.max(v.abs());
5604 }
5605 let whole = m.floor();
5606 if whole >= 1e15 {
5607 return Err(Error::domain("the value is too large to encode in binary", span));
5608 }
5609 let mut w = 1usize;
5610 let mut n = whole as i64;
5611 while n > 1 {
5612 n /= 2;
5613 w += 1;
5614 }
5615 Ok(w)
5616}
5617
5618fn encode_one(radix: &[f64], v: f64, out: &mut [f64]) {
5622 let mut rem = v;
5623 for i in (0..radix.len()).rev() {
5624 let b = radix[i];
5625 if b == 0.0 {
5626 out[i] = rem;
5627 rem = 0.0;
5628 } else {
5629 let r = rem - b * (rem / b).floor();
5630 out[i] = r;
5631 rem = (rem - r) / b;
5632 }
5633 }
5634}
5635
5636fn encode(x: &Array, y: &Array, span: Span) -> Result<Array> {
5640 let radix = digits_of(x, "encode", span)?;
5641 let values = digits_of(y, "encode", span)?;
5642 let k = radix.len();
5643 let n = values.len();
5644 let mut out = vec![0.0f64; k * n];
5645 let mut cell = vec![0.0f64; k];
5646 for (j, &v) in values.iter().enumerate() {
5647 encode_one(&radix, v, &mut cell);
5648 for i in 0..k {
5649 out[i * n + j] = cell[i];
5650 }
5651 }
5652 let mut shape = if x.rank() == 0 { Vec::new() } else { vec![k] };
5655 shape.extend_from_slice(&y.shape);
5656 Ok(Array::new(shape, narrow(out, is_integral(x) && is_integral(y))))
5657}
5658
5659fn encode_bits(y: &Array, span: Span) -> Result<Array> {
5661 let values = digits_of(y, "encode", span)?;
5662 let k = bit_width(&values, span)?;
5663 let radix = vec![2.0; k];
5664 let mut out = vec![0.0f64; values.len() * k];
5665 for (j, &v) in values.iter().enumerate() {
5666 encode_one(&radix, v, &mut out[j * k..(j + 1) * k]);
5667 }
5668 let mut shape = y.shape.clone();
5669 shape.push(k);
5670 Ok(Array::new(shape, narrow(out, is_integral(y))))
5671}
5672
5673fn laminate(x: &Array, y: &Array, span: Span) -> Result<Array> {
5678 let spread = |a: &Array, other: &Array| -> Array {
5679 if a.rank() != 0 {
5680 return a.clone();
5681 }
5682 let shape = if other.rank() == 0 { vec![1] } else { other.shape.clone() };
5683 let n: usize = shape.iter().product();
5684 let mut data = Data::empty(a.dtype());
5685 for _ in 0..n {
5686 push_elem(&mut data, &a.data, 0);
5687 }
5688 Array::new(shape, data)
5689 };
5690 assemble(&[2], vec![spread(x, y), spread(y, x)], span)
5691}
5692
5693fn table_of(y: &Array) -> Array {
5695 let shape = match y.rank() {
5696 0 => vec![1, 1],
5697 _ => vec![y.items(), y.item_size()],
5698 };
5699 Array::new(shape, y.data.clone())
5700}
5701
5702fn table(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
5708 let ranks = u.ranks();
5709 let fxl = x.rank() - effective_rank(ranks[1], x.rank());
5710 let fyl = y.rank() - effective_rank(ranks[2], y.rank());
5711 let mut frame = x.shape[..fxl].to_vec();
5712 frame.extend_from_slice(&y.shape[..fyl]);
5713 let nx: usize = x.shape[..fxl].iter().product();
5714 let ny: usize = y.shape[..fyl].iter().product();
5715 let n = nx * ny;
5716 if n == 0 {
5717 return assemble(&frame, Vec::new(), span);
5718 }
5719 if frame.is_empty() {
5720 return u.dyad(x, y, ctx, span);
5721 }
5722 let work = x.count().max(y.count()).max(n);
5723 let cells = each_cell(n, work, u.is_pure(), ctx, |i, c| {
5724 u.dyad(&x.cell_at(fxl, i / ny), &y.cell_at(fyl, i % ny), c, span)
5725 })?;
5726 assemble(&frame, cells, span)
5727}
5728
5729pub(crate) fn with_origin(v: &Verb, origin: i64) -> Option<Verb> {
5736 match v {
5737 Verb::Prim(p) => {
5738 let mut out = *p;
5739 let mut changed = false;
5740 out.monad = match p.monad {
5741 MonadOp::GradeUp { .. } => {
5742 changed = true;
5743 MonadOp::GradeUp { origin }
5744 }
5745 MonadOp::GradeDown { .. } => {
5746 changed = true;
5747 MonadOp::GradeDown { origin }
5748 }
5749 MonadOp::IotaApl { .. } => {
5750 changed = true;
5751 MonadOp::IotaApl { origin }
5752 }
5753 MonadOp::Indices { boxed_coords, .. } => {
5754 changed = true;
5755 MonadOp::Indices { origin, boxed_coords }
5756 }
5757 MonadOp::Roll { fixed, float_at_zero, .. } => {
5758 changed = true;
5759 MonadOp::Roll { origin, fixed, float_at_zero }
5760 }
5761 other => other,
5762 };
5763 out.dyad = match p.dyad {
5764 DyadOp::IndexOf { .. } => {
5765 changed = true;
5766 DyadOp::IndexOf { origin }
5767 }
5768 DyadOp::IndexOfLast { .. } => {
5769 changed = true;
5770 DyadOp::IndexOfLast { origin }
5771 }
5772 DyadOp::CollateGrade { down, .. } => {
5773 changed = true;
5774 DyadOp::CollateGrade { down, origin }
5775 }
5776 DyadOp::Squad { .. } => {
5777 changed = true;
5778 DyadOp::Squad { origin }
5779 }
5780 DyadOp::Pick { .. } => {
5781 changed = true;
5782 DyadOp::Pick { origin }
5783 }
5784 DyadOp::SelectAxis { axis, rank, .. } => {
5785 changed = true;
5786 DyadOp::SelectAxis { axis, rank, origin }
5787 }
5788 DyadOp::Deal { fixed, .. } => {
5789 changed = true;
5790 DyadOp::Deal { origin, fixed }
5791 }
5792 other => other,
5793 };
5794 changed.then_some(Verb::Prim(out))
5795 }
5796 Verb::Rank(u, r) => Some(Verb::Rank(Box::new(with_origin(u, origin)?), *r)),
5797 Verb::Reduce(u) => Some(Verb::Reduce(Box::new(with_origin(u, origin)?))),
5798 Verb::Windowed(u, k) => Some(Verb::Windowed(Box::new(with_origin(u, origin)?), *k)),
5799 Verb::Commute(u) => Some(Verb::Commute(Box::new(with_origin(u, origin)?))),
5800 Verb::Each(u, e) => Some(Verb::Each(Box::new(with_origin(u, origin)?), *e)),
5801 Verb::Fit(u, n) => Some(Verb::Fit(Box::new(with_origin(u, origin)?), *n)),
5802 Verb::AlongAxis(u, k) => Some(Verb::AlongAxis(Box::new(with_origin(u, origin)?), *k)),
5803 _ => None,
5804 }
5805}
5806
5807fn scalar_dyad_of(v: &Verb) -> Option<ScalarDyad> {
5812 match v {
5813 Verb::Prim(p) => match p.dyad {
5814 DyadOp::Scalar(op) => Some(op),
5815 _ => None,
5816 },
5817 _ => None,
5818 }
5819}
5820
5821fn folds_with(u: &Verb, op: ScalarDyad) -> bool {
5824 matches!(u, Verb::Reduce(inner) if scalar_dyad_of(inner) == Some(op))
5825}
5826
5827fn inner_product(
5834 u: &Verb,
5835 v: &Verb,
5836 apl: bool,
5837 x: &Array,
5838 y: &Array,
5839 ctx: &mut Ctx<'_>,
5840 span: Span,
5841) -> Result<Array> {
5842 if let Some(a) = matrix_product(u, v, x, y, span) {
5843 return Ok(a);
5844 }
5845 if apl && scalar_dyad_of(v).is_none() {
5848 return apl_inner_product(u, v, x, y, ctx, span);
5849 }
5850 if !apl {
5851 return inner_cells(u, v, x, y, ctx, span);
5852 }
5853 let saved = ctx.cfg.agreement;
5859 ctx.cfg.agreement = Agreement::LeadingPrefix;
5860 let out = inner_cells(u, v, x, y, ctx, span);
5861 ctx.cfg.agreement = saved;
5862 out
5863}
5864
5865fn inner_cells(
5868 u: &Verb,
5869 v: &Verb,
5870 x: &Array,
5871 y: &Array,
5872 ctx: &mut Ctx<'_>,
5873 span: Span,
5874) -> Result<Array> {
5875 let cell_rank = effective_rank(v.ranks()[1].max(1), x.rank());
5876 let frame_rank = x.rank() - cell_rank;
5877 if frame_rank == 0 {
5878 let inner = v.dyad(x, y, ctx, span)?;
5879 return u.monad(&inner, ctx, span);
5880 }
5881 let frame = x.shape[..frame_rank].to_vec();
5882 let n: usize = frame.iter().product();
5883 if n == 0 {
5884 return assemble(&frame, Vec::new(), span);
5885 }
5886 let work = x.count().max(y.count());
5887 let pure = u.is_pure() && v.is_pure();
5888 let cells = each_cell(n, work, pure, ctx, |i, c| {
5889 let inner = v.dyad(&x.cell_at(frame_rank, i), y, c, span)?;
5890 u.monad(&inner, c, span)
5891 })?;
5892 assemble(&frame, cells, span)
5893}
5894
5895fn apl_inner_product(
5900 u: &Verb,
5901 v: &Verb,
5902 x: &Array,
5903 y: &Array,
5904 ctx: &mut Ctx<'_>,
5905 span: Span,
5906) -> Result<Array> {
5907 let k = match (x.rank(), y.rank()) {
5910 (0, 0) => 1,
5911 (0, _) => y.shape[0],
5912 _ => x.shape[x.rank() - 1],
5913 };
5914 if x.rank() > 0 && y.rank() > 0 && x.shape[x.rank() - 1] != y.shape[0] {
5915 return Err(Error::new(
5916 ErrorKind::Length,
5917 format!("inner product over {} and {} elements", x.shape[x.rank() - 1], y.shape[0]),
5918 Some(span),
5919 ));
5920 }
5921 let lead: &[usize] = if x.rank() > 0 { &x.shape[..x.rank() - 1] } else { &[] };
5922 let trail: &[usize] = if y.rank() > 0 { &y.shape[1..] } else { &[] };
5923 let rows: usize = lead.iter().product();
5924 let cols: usize = trail.iter().product();
5925 let mut frame = lead.to_vec();
5926 frame.extend_from_slice(trail);
5927 let n = rows * cols;
5928 if n == 0 {
5929 return assemble(&frame, Vec::new(), span);
5930 }
5931 let vector = |d: &Data, at: &dyn Fn(usize) -> usize| {
5932 let mut out = Data::empty(d.dtype());
5933 for t in 0..k {
5934 out.push_from(d, at(t));
5935 }
5936 Array::new(vec![k], out)
5937 };
5938 let pure = u.is_pure() && v.is_pure();
5939 let cells = each_cell(n, x.count().max(y.count()), pure, ctx, |i, c| {
5940 let (r, col) = (i / cols, i % cols);
5941 let left = vector(&x.data, &|t| if x.rank() > 0 { r * k + t } else { 0 });
5942 let right = vector(&y.data, &|t| if y.rank() > 0 { t * cols + col } else { 0 });
5943 let inner = v.dyad(&left, &right, c, span)?;
5944 u.monad(&inner, c, span)
5945 })?;
5946 assemble(&frame, cells, span)
5947}
5948
5949fn matrix_product(u: &Verb, v: &Verb, x: &Array, y: &Array, span: Span) -> Option<Array> {
5955 if !folds_with(u, ScalarDyad::Add) || scalar_dyad_of(v) != Some(ScalarDyad::Mul) {
5956 return None;
5957 }
5958 if x.rank() == 0 || y.rank() == 0 {
5959 return None;
5960 }
5961 let k = x.shape[x.rank() - 1];
5962 if k != y.shape[0] {
5963 return None;
5964 }
5965 let rows: usize = x.shape[..x.rank() - 1].iter().product();
5966 let cols: usize = y.shape[1..].iter().product();
5967 let mut shape = x.shape[..x.rank() - 1].to_vec();
5968 shape.extend_from_slice(&y.shape[1..]);
5969 if crate::limits::elements(&shape, span).is_err() {
5970 return None;
5971 }
5972 let whole = matches!(x.dtype(), DType::Bool | DType::I64)
5973 && matches!(y.dtype(), DType::Bool | DType::I64);
5974 if whole
5975 && let (Some(xs), Some(ys)) = (x.to_i64_vec(), y.to_i64_vec())
5976 && let Some(out) = matmul_whole(&xs, &ys, rows, k, cols)
5977 {
5978 return Some(Array::new(shape, Data::I64(out.into())));
5979 }
5980 let (xs, ys) = (x.to_f64_vec()?, y.to_f64_vec()?);
5981 let out = par::fill_rows(rows, cols, rows * k * cols, |r0, part| {
5982 matmul_f64(&xs, &ys, k, cols, r0, part);
5983 });
5984 Some(Array::new(shape, Data::F64(out.into())))
5985}
5986
5987const MATMUL_BLOCK: usize = 128;
5991
5992#[inline(always)]
5993fn matmul_f64_body(xs: &[f64], ys: &[f64], k: usize, n: usize, r0: usize, out: &mut [f64]) {
5994 if n == 0 {
5995 return;
5996 }
5997 let rows = out.len() / n;
5998 for k0 in (0..k).step_by(MATMUL_BLOCK) {
5999 let k1 = (k0 + MATMUL_BLOCK).min(k);
6000 for r in 0..rows {
6001 let left = &xs[(r0 + r) * k..(r0 + r + 1) * k];
6002 let dst = &mut out[r * n..(r + 1) * n];
6003 for (t, &a) in left.iter().enumerate().take(k1).skip(k0) {
6004 let row = &ys[t * n..(t + 1) * n];
6005 for (o, &b) in dst.iter_mut().zip(row) {
6006 *o += a * b;
6007 }
6008 }
6009 }
6010 }
6011}
6012
6013multiversioned! {
6014 fn matmul_f64(
6018 xs: &[f64],
6019 ys: &[f64],
6020 k: usize,
6021 n: usize,
6022 r0: usize,
6023 out: &mut [f64],
6024 ) -> () = matmul_f64_body;
6025}
6026
6027#[inline(always)]
6028fn matmul_i64_body(xs: &[i64], ys: &[i64], k: usize, n: usize, r0: usize, out: &mut [i64]) {
6029 if n == 0 {
6030 return;
6031 }
6032 let rows = out.len() / n;
6033 for k0 in (0..k).step_by(MATMUL_BLOCK) {
6034 let k1 = (k0 + MATMUL_BLOCK).min(k);
6035 for r in 0..rows {
6036 let left = &xs[(r0 + r) * k..(r0 + r + 1) * k];
6037 let dst = &mut out[r * n..(r + 1) * n];
6038 for (t, &a) in left.iter().enumerate().take(k1).skip(k0) {
6039 let row = &ys[t * n..(t + 1) * n];
6040 for (o, &b) in dst.iter_mut().zip(row) {
6041 *o = o.wrapping_add(a.wrapping_mul(b));
6042 }
6043 }
6044 }
6045 }
6046}
6047
6048multiversioned! {
6049 fn matmul_i64(
6053 xs: &[i64],
6054 ys: &[i64],
6055 k: usize,
6056 n: usize,
6057 r0: usize,
6058 out: &mut [i64],
6059 ) -> () = matmul_i64_body;
6060}
6061
6062fn matmul_whole(xs: &[i64], ys: &[i64], rows: usize, k: usize, n: usize) -> Option<Vec<i64>> {
6066 let bound = |v: &[i64]| v.iter().map(|&a| (a as i128).abs()).max().unwrap_or(0);
6071 if bound(xs).saturating_mul(bound(ys)).saturating_mul(k as i128) <= i64::MAX as i128 {
6072 return Some(par::fill_rows(rows, n, rows * k * n, |r0, part| {
6073 matmul_i64(xs, ys, k, n, r0, part);
6074 }));
6075 }
6076 let mut out = vec![0i64; rows * n];
6077 for r in 0..rows {
6078 let left = &xs[r * k..(r + 1) * k];
6079 let dst = &mut out[r * n..(r + 1) * n];
6080 for (t, &a) in left.iter().enumerate() {
6081 for (o, &b) in dst.iter_mut().zip(&ys[t * n..(t + 1) * n]) {
6082 *o = a.checked_mul(b).and_then(|p| o.checked_add(p))?;
6083 }
6084 }
6085 }
6086 Some(out)
6087}
6088
6089const DETERMINANT_MINORS_MAX: usize = 16;
6094
6095fn determinant(
6101 u: &Verb,
6102 v: &Verb,
6103 apl: bool,
6104 y: &Array,
6105 ctx: &mut Ctx<'_>,
6106 span: Span,
6107) -> Result<Array> {
6108 if apl {
6109 return Err(Error::domain("an inner product has no monadic meaning in APL", span));
6110 }
6111 if y.rank() > 2 {
6115 let frame = y.shape[..y.rank() - 2].to_vec();
6116 let n: usize = frame.iter().product();
6117 let pure = u.is_pure() && v.is_pure();
6118 let cells = each_cell(n, y.count(), pure, ctx, |i, c| {
6119 determinant(u, v, apl, &y.cell_at(y.rank() - 2, i), c, span)
6120 })?;
6121 return assemble(&frame, cells, span);
6122 }
6123 let rows = y.items();
6124 let cols = y.item_size();
6125 if folds_with(u, ScalarDyad::Sub)
6126 && scalar_dyad_of(v) == Some(ScalarDyad::Mul)
6127 && rows == cols
6128 && rows >= 3
6129 && matches!(y.dtype(), DType::Bool | DType::I64 | DType::F64)
6130 && let Some(values) = y.to_f64_vec()
6131 {
6132 return Ok(Array::scalar_f64(determinant_lu(values, rows)));
6133 }
6134 if rows > DETERMINANT_MINORS_MAX {
6135 return Err(Error::not_yet(
6136 format!(
6137 "a determinant of more than {DETERMINANT_MINORS_MAX} rows by minors \
6138 (only -/ . * over machine numbers has a direct method)"
6139 ),
6140 span,
6141 ));
6142 }
6143 let mut seen: HashMap<u64, Array> = HashMap::new();
6144 let all = if rows == 64 { u64::MAX } else { (1u64 << rows) - 1 };
6145 minors(u, v, y, cols, rows, all, &mut seen, ctx, span)
6146}
6147
6148#[allow(clippy::too_many_arguments)]
6152fn minors(
6153 u: &Verb,
6154 v: &Verb,
6155 y: &Array,
6156 cols: usize,
6157 rows: usize,
6158 left: u64,
6159 seen: &mut HashMap<u64, Array>,
6160 ctx: &mut Ctx<'_>,
6161 span: Span,
6162) -> Result<Array> {
6163 if let Some(a) = seen.get(&left) {
6164 return Ok(a.clone());
6165 }
6166 let column = rows - left.count_ones() as usize;
6169 let value = if column >= cols {
6170 let data = reduce_identity(v, 1).ok_or_else(|| {
6171 Error::not_yet(
6172 format!("the identity element of {} (a determinant with no columns)", v.name()),
6173 span,
6174 )
6175 })?;
6176 Array::new(Vec::new(), data)
6177 } else if left == 0 {
6178 u.monad(&Array::new(vec![0], Data::empty(DType::I64)), ctx, span)?
6179 } else {
6180 let mut terms = Vec::with_capacity(left.count_ones() as usize);
6181 for r in 0..rows {
6182 if left & (1 << r) == 0 {
6183 continue;
6184 }
6185 let minor = minors(u, v, y, cols, rows, left & !(1 << r), seen, ctx, span)?;
6186 let head = Array::new(Vec::new(), y.data.slice(r * cols + column, r * cols + column + 1));
6187 terms.push(v.dyad(&head, &minor, ctx, span)?);
6188 }
6189 let n = terms.len();
6190 u.monad(&assemble(&[n], terms, span)?, ctx, span)?
6191 };
6192 seen.insert(left, value.clone());
6193 Ok(value)
6194}
6195
6196fn determinant_lu(mut a: Vec<f64>, n: usize) -> f64 {
6201 let mut det = 1.0f64;
6202 for c in 0..n {
6203 let mut pivot = c;
6204 for r in c + 1..n {
6205 if a[r * n + c].abs() > a[pivot * n + c].abs() {
6206 pivot = r;
6207 }
6208 }
6209 if a[pivot * n + c] == 0.0 {
6210 return 0.0;
6211 }
6212 if pivot != c {
6213 for j in 0..n {
6214 a.swap(c * n + j, pivot * n + j);
6215 }
6216 det = -det;
6217 }
6218 let head = a[c * n + c];
6219 det *= head;
6220 for r in c + 1..n {
6221 let factor = a[r * n + c] / head;
6222 if factor == 0.0 {
6223 continue;
6224 }
6225 for j in c..n {
6226 a[r * n + j] -= factor * a[c * n + j];
6227 }
6228 }
6229 }
6230 det
6231}
6232
6233fn monad_op(p: &Prim, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
6235 match p.monad {
6236 MonadOp::Scalar(op) => scalar_monad(op, y, ctx.cfg, span),
6237 MonadOp::ShapeOf => {
6238 Ok(carry_exact(Array::from_i64(y.shape.iter().map(|&n| n as i64).collect()), y))
6239 }
6240 MonadOp::Tally => Ok(carry_exact(Array::scalar_i64(y.items() as i64), y)),
6241 MonadOp::Ravel => Ok(Array::new(vec![y.count()], y.data.clone())),
6242 MonadOp::TransposeAxes => Ok(transpose_axes(y)),
6243 MonadOp::Head => Ok(head(y)),
6244 MonadOp::Behead => behead(y, span),
6245 MonadOp::Tail => Ok(tail(y)),
6246 MonadOp::Curtail => Ok(curtail(y)),
6247 MonadOp::Reverse => Ok(reverse(y)),
6248 MonadOp::Nub => Ok(nub(y, ctx.cfg.tol)),
6251 MonadOp::GradeUp { origin } | MonadOp::GradeDown { origin } => {
6252 check_gradable(y, ctx.cfg.rules, span)?;
6253 if ctx.cfg.rules.lang == crate::Lang::Apl && y.rank() == 0 {
6256 return Err(Error::domain("a grade needs an array, not a scalar", span));
6257 }
6258 let down = matches!(p.monad, MonadOp::GradeDown { .. });
6259 let order = grade_order(y, down, Tao::of(ctx.cfg.rules));
6260 Ok(Array::from_i64(order.iter().map(|&i| origin + i as i64).collect()))
6261 }
6262 MonadOp::IotaJ => iota_j(y, span),
6263 MonadOp::IotaApl { origin } => iota_apl(y, origin, span),
6264 MonadOp::Echo => {
6265 (ctx.out)(&format!("{}\n", crate::fmt::format_array(y, &ctx.cfg.fmt)));
6266 Ok(Array::empty(DType::I64))
6267 }
6268 MonadOp::ReadStream => {
6269 stream_number(y, 1, "1!:1 reads", span)?;
6270 let line = ctx.read_line(span)?;
6271 Ok(Array::from_chars(line.chars().collect()))
6272 }
6273 MonadOp::TypeCode => Ok(Array::scalar_i64(type_code(y))),
6274 MonadOp::Same => Ok(y.clone()),
6275 MonadOp::Format => Ok(format_chars(y, &ctx.cfg.fmt)),
6276 MonadOp::DecodeBits => decode(None, y, span).map(|r| carry_exact(r, y)),
6277 MonadOp::EncodeBits => encode_bits(y, span).map(|r| carry_exact(r, y)),
6278 MonadOp::Itemize => {
6279 let mut shape = vec![1usize];
6280 shape.extend_from_slice(&y.shape);
6281 Ok(Array::new(shape, y.data.clone()))
6282 }
6283 MonadOp::TableOf => Ok(table_of(y)),
6284 MonadOp::Enclose(rule) => Ok(enclose(y, rule)),
6285 MonadOp::Open => Ok(open_cell(y)),
6286 MonadOp::Raze => raze(y, span),
6287 MonadOp::Catalogue => catalogue(y, span),
6288 MonadOp::AtomicRep => atomic_rep(y, ctx, span),
6289 MonadOp::RazeIn => raze_in(y, ctx.cfg.tol, span),
6290 MonadOp::First => Ok(first(y)),
6291 MonadOp::Enlist => enlist(y, span),
6292 MonadOp::Depth => Ok(Array::scalar_i64(depth(y))),
6293 MonadOp::Indices { origin, boxed_coords } => {
6294 where_indices(y, origin, boxed_coords, span)
6295 }
6296 MonadOp::Steps => steps(y, span),
6297 MonadOp::ToExact => to_exact(y, span),
6298 MonadOp::NthPrime => {
6299 let n = y
6300 .to_i64_vec()
6301 .ok_or_else(|| Error::domain("the prime index must be an integer", span))?;
6302 let v = n.first().copied().unwrap_or(0);
6303 Ok(carry_exact(Array::scalar_i64(nth_prime(v, span)?), y))
6304 }
6305 MonadOp::PrimeFactors => {
6306 let n = y
6307 .to_i64_vec()
6308 .ok_or_else(|| Error::domain("prime factors need an integer", span))?;
6309 let v = n.first().copied().unwrap_or(0);
6310 Ok(carry_exact(Array::from_i64(prime_factors(v, span)?), y))
6311 }
6312 MonadOp::MatrixInverse => matrix_inverse(y, span),
6313 MonadOp::Roll { origin, fixed, float_at_zero } => {
6314 roll(y, origin, fixed, float_at_zero, span)
6315 }
6316 MonadOp::ComplexParts { polar } => complex_parts(y, polar, span),
6317 MonadOp::SelfClassify => Ok(self_classify(y, ctx.cfg.tol)),
6318 MonadOp::NubSieve => Ok(nub_sieve(y, ctx.cfg.tol)),
6319 MonadOp::Unicode { pass_chars } => unicode(y, pass_chars, span),
6320 MonadOp::Symbols => to_symbols(y, span),
6321 MonadOp::Words => words(y, span),
6322 MonadOp::LevelOf => Ok(Array::scalar_i64(boxing_level(y))),
6323 MonadOp::MapPaths => Ok(map_paths(y)),
6324 MonadOp::Nest => Ok(nest(y)),
6325 MonadOp::PolyRoots => poly_roots(y, span),
6326 MonadOp::PolyDeriv => poly_deriv(y, span),
6327 MonadOp::AnagramIndex => anagram_index(y, ctx.cfg.rules, span),
6328 MonadOp::CycleForm => cycle_form(y, span),
6329 MonadOp::Split => Ok(split_items(y)),
6330 MonadOp::Execute { apl } => execute(y, apl, ctx, span),
6331 MonadOp::NotYet(what) => Err(Error::not_yet(what, span)),
6332 MonadOp::None => {
6333 Err(Error::domain(format!("{} has no monadic meaning", p.name), span))
6334 }
6335 }
6336}
6337
6338fn complex_parts(y: &Array, polar: bool, span: Span) -> Result<Array> {
6342 let Some(v) = y.to_complex_vec() else {
6343 return Err(wrong_type(y.dtype(), span));
6344 };
6345 let z = v.first().copied().unwrap_or(cx::ZERO);
6346 let pair = if polar { vec![cx::abs(z), cx::arg(z)] } else { vec![z[0], z[1]] };
6347 Ok(Array::from_f64(pair))
6348}
6349
6350fn axis_counts(x: &Array, what: &str, span: Span) -> Result<Vec<i64>> {
6351 if x.rank() > 1 {
6352 return Err(Error::new(
6353 ErrorKind::Rank,
6354 format!("{what} needs a scalar or vector left argument"),
6355 Some(span),
6356 ));
6357 }
6358 if x.count() == 0 {
6361 return Ok(Vec::new());
6362 }
6363 x.to_i64_vec()
6364 .ok_or_else(|| Error::domain(format!("{what} needs integer lengths"), span))
6365}
6366
6367fn reshape(x: &Array, y: &Array, by_items: bool, span: Span) -> Result<Array> {
6378 let dims = axis_counts(x, "reshape", span)?;
6379 if dims.iter().any(|&d| d < 0) {
6380 return Err(Error::domain("reshape lengths must be nonnegative", span));
6381 }
6382 let mut shape: Vec<usize> = dims.iter().map(|&d| d as usize).collect();
6383 let (unit, src) = if by_items {
6385 let item_shape = if y.rank() == 0 { &[][..] } else { &y.shape[1..] };
6386 shape.extend_from_slice(item_shape);
6387 (item_shape.iter().product::<usize>(), y.items().max(usize::from(y.rank() == 0)))
6388 } else {
6389 (1, y.count())
6390 };
6391 let n = crate::limits::elements(&shape, span)?;
6392 let mut data = Data::empty(y.dtype());
6393 if n > 0 && src == 0 {
6394 if by_items {
6395 return Err(Error::new(ErrorKind::Length, "reshape of an empty array", Some(span)));
6396 }
6397 return Ok(Array::new(shape, fill_data(y.dtype(), n)));
6398 }
6399 if y.is_row_major() && n <= unit.saturating_mul(src) && n <= y.data.len() {
6405 return Ok(Array::new(shape, y.data.slice(0, n)));
6406 }
6407 for i in 0..n {
6408 push_elem(&mut data, &y.data, (i / unit) % src * unit + i % unit);
6409 }
6410 Ok(Array::new(shape, data))
6411}
6412
6413fn leading_run(y: &Array, counts: &[i64], drop: bool) -> Option<Array> {
6417 if y.rank() == 0 || counts.is_empty() {
6418 return None;
6419 }
6420 let trailing_untouched = counts[1..].iter().enumerate().all(|(a, &c)| {
6424 if drop { c == 0 } else { c.unsigned_abs() as usize == y.shape[a + 1] }
6425 });
6426 if !trailing_untouched {
6427 return None;
6428 }
6429 let n = y.items();
6430 let k = counts[0];
6431 let a = k.unsigned_abs() as usize;
6432 let (lo, keep) = if drop {
6433 let a = a.min(n);
6434 if k >= 0 { (a, n - a) } else { (0, n - a) }
6435 } else {
6436 if a > n {
6438 return None;
6439 }
6440 if k >= 0 { (0, a) } else { (n - a, a) }
6441 };
6442 Some(section(y, lo, lo + keep))
6443}
6444
6445fn count_rank(verb: &str, counts: usize, rank: usize, span: Span) -> Error {
6450 Error::new(
6451 ErrorKind::Length,
6452 format!("{counts} {verb} counts for a rank-{rank} argument"),
6453 Some(span),
6454 )
6455}
6456
6457fn take(x: &Array, y: &Array, prototype_fill: bool, apl: bool, span: Span) -> Result<Array> {
6458 let counts = axis_counts(x, "take", span)?;
6459 let fill = if prototype_fill { prototype_of(y) } else { None };
6463 let promoted;
6464 let base = if y.rank() == 0 {
6467 promoted = Array::new(vec![1; counts.len()], y.data.clone());
6468 &promoted
6469 } else {
6470 y
6471 };
6472 let wrong = if apl {
6474 counts.len() != base.rank()
6475 } else {
6476 counts.len() > base.rank() || (counts.is_empty() && base.rank() > 0)
6477 };
6478 if wrong {
6479 return Err(count_rank("take", counts.len(), base.rank(), span));
6480 }
6481 if let Some(run) = leading_run(base, &counts, false) {
6482 return Ok(run);
6483 }
6484 let mut out_shape = base.shape.clone();
6485 for (a, &k) in counts.iter().enumerate() {
6486 out_shape[a] = k.unsigned_abs() as usize;
6487 }
6488 let n = crate::limits::elements(&out_shape, span)?;
6489 let st = strides(&base.shape);
6490 let mut data = Data::empty(base.dtype());
6491 let mut coord = vec![0usize; out_shape.len()];
6492 for _ in 0..n {
6493 let mut idx = 0usize;
6494 let mut inside = true;
6495 for a in 0..out_shape.len() {
6496 let len = base.shape[a] as i64;
6497 let c = coord[a] as i64;
6498 let s = match counts.get(a) {
6501 Some(&k) if k < 0 => c + len - k.unsigned_abs() as i64,
6502 _ => c,
6503 };
6504 if s < 0 || s >= len {
6505 inside = false;
6506 break;
6507 }
6508 idx += s as usize * st[a];
6509 }
6510 if inside {
6511 push_elem(&mut data, &base.data, idx);
6512 } else if let (Data::Box(v), Some(p)) = (&mut data, &fill) {
6513 v.push(p.clone());
6514 } else {
6515 data.push_fill();
6516 }
6517 odometer(&mut coord, &out_shape);
6518 }
6519 Ok(Array::new(out_shape, data))
6520}
6521
6522fn prototype_of(y: &Array) -> Option<Array> {
6526 fn zeroed(a: &Array) -> Array {
6527 if let Some(items) = a.as_boxes() {
6528 let inner: Vec<Array> = items.iter().map(zeroed).collect();
6529 return Array::new(a.shape.clone(), Data::Box(inner.into()));
6530 }
6531 let dtype = match a.dtype() {
6532 DType::Char | DType::Symbol => a.dtype(),
6533 _ => DType::I64,
6534 };
6535 Array::new(a.shape.clone(), fill_data(dtype, a.count()))
6536 }
6537 let first = y.as_boxes()?.first()?;
6538 Some(zeroed(first))
6539}
6540
6541fn drop_(x: &Array, y: &Array, apl: bool, span: Span) -> Result<Array> {
6542 let counts = axis_counts(x, "drop", span)?;
6543 let promoted;
6544 let base = if y.rank() == 0 {
6545 promoted = Array::new(vec![1; counts.len()], y.data.clone());
6546 &promoted
6547 } else {
6548 y
6549 };
6550 let wrong =
6551 if apl { counts.len() != base.rank() } else { counts.len() > base.rank() };
6552 if wrong {
6553 return Err(count_rank("drop", counts.len(), base.rank(), span));
6554 }
6555 if let Some(run) = leading_run(base, &counts, true) {
6556 return Ok(run);
6557 }
6558 let mut out_shape = base.shape.clone();
6559 let mut offset = vec![0usize; base.rank()];
6560 for (a, &k) in counts.iter().enumerate() {
6561 let len = base.shape[a];
6562 let d = (k.unsigned_abs() as usize).min(len);
6563 out_shape[a] = len - d;
6564 if k > 0 {
6565 offset[a] = d;
6566 }
6567 }
6568 let n: usize = out_shape.iter().product();
6569 let st = strides(&base.shape);
6570 let mut data = Data::empty(base.dtype());
6571 let mut coord = vec![0usize; out_shape.len()];
6572 for _ in 0..n {
6573 let idx: usize = (0..out_shape.len()).map(|a| (coord[a] + offset[a]) * st[a]).sum();
6574 push_elem(&mut data, &base.data, idx);
6575 odometer(&mut coord, &out_shape);
6576 }
6577 Ok(Array::new(out_shape, data))
6578}
6579
6580fn dyad_op(p: &Prim, x: &Array, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
6582 let tol = cfg.tol;
6583 match p.dyad {
6584 DyadOp::Scalar(op) => scalar_dyad(op, x, y, cfg, span),
6587 DyadOp::Reshape => reshape(x, y, cfg.agreement == Agreement::LeadingPrefix, span),
6588 DyadOp::Take => {
6589 let apl = cfg.rules.lang == crate::Lang::Apl;
6590 take(x, y, cfg.agreement == Agreement::ExactOrScalar, apl, span)
6591 }
6592 DyadOp::Drop => drop_(x, y, cfg.rules.lang == crate::Lang::Apl, span),
6593 DyadOp::Right => Ok(y.clone()),
6594 DyadOp::Left => Ok(x.clone()),
6595 DyadOp::Rotate => rotate(x, y, span),
6596 DyadOp::AppendLeading => {
6599 catenate(x, y, true, cfg.agreement == Agreement::LeadingPrefix, span)
6600 }
6601 DyadOp::AppendLast => {
6602 catenate(x, y, false, cfg.agreement == Agreement::LeadingPrefix, span)
6603 }
6604 DyadOp::IndexOf { origin } => Ok(index_of(x, y, origin, tol)),
6605 DyadOp::MemberJ => Ok(member_j(x, y, tol)),
6606 DyadOp::MemberApl => Ok(member_apl(x, y, tol)),
6607 DyadOp::From => from_index(x, y, span),
6608 DyadOp::Match => {
6609 let empties_differ = cfg.rules.lang == crate::Lang::Apl
6613 && x.count() == 0
6614 && y.count() == 0
6615 && (x.dtype() == DType::Char) != (y.dtype() == DType::Char);
6616 Ok(Array::scalar_bool(!empties_differ && arrays_match(x, y, tol)))
6617 }
6618 DyadOp::NotMatch => Ok(Array::scalar_bool(!arrays_match(x, y, tol))),
6619 DyadOp::GradeSelect { down } => grade_select(x, y, down, cfg.rules, span),
6620 DyadOp::Copy => copy_items(x, y, cfg.agreement == Agreement::ExactOrScalar, span),
6621 DyadOp::CollateGrade { down, origin } => collate_grade(x, y, down, origin, span),
6622 DyadOp::TransposeJ => transpose_j(x, y, span),
6623 DyadOp::TransposeApl => transpose_apl(x, y, cfg.rules.origin, span),
6624 DyadOp::DecodeApl => decode_apl(x, y, span).map(|r| carry_exact2(r, x, y)),
6625 DyadOp::EncodeApl => encode_apl(x, y, span).map(|r| carry_exact2(r, x, y)),
6626 DyadOp::Decode => decode(Some(x), y, span).map(|r| carry_exact2(r, x, y)),
6627 DyadOp::Encode => encode(x, y, span).map(|r| carry_exact2(r, x, y)),
6628 DyadOp::Laminate => laminate(x, y, span),
6629 DyadOp::Link => link(x, y, span),
6630 DyadOp::Strand => strand(x, y, span),
6631 DyadOp::IntervalIndex { offset, closed } => {
6632 interval_index(x, y, offset, closed, tol, span)
6633 }
6634 DyadOp::IndexOfLast { origin } => Ok(index_of_last(x, y, origin, tol)),
6635 DyadOp::MatrixDivide => matrix_divide(x, y, span),
6636 DyadOp::PartitionEnclose => partition_enclose(x, y, span),
6637 DyadOp::Squad { origin } => squad(x, y, origin, span),
6638 DyadOp::SelectAxis { axis, rank, origin } => {
6639 select_axis(x, y, axis, rank, origin, span)
6640 }
6641 DyadOp::Fetch => fetch(x, y, span),
6642 DyadOp::PolyEval => poly_eval(x, y, span),
6643 DyadOp::PolyIntegral => poly_integral(x, y, span),
6644 DyadOp::TruthTable(m) => truth_table(m, x, y, span),
6645 DyadOp::FormatSpec => format_spec(x, y, &cfg.fmt, span),
6646 DyadOp::FormatSpecJ => format_spec_j(x, y, &cfg.fmt, span),
6647 DyadOp::ParseNumbers => parse_numbers(x, y, span),
6648 DyadOp::SequentialMachine => sequential_machine(x, y, span),
6649 DyadOp::Deal { origin, fixed } => deal(x, y, origin, fixed, span),
6650 DyadOp::ExactForm => exact_form(x, y, span),
6651 DyadOp::Boolean(op) => bool_dyad(op, x, y, cfg, span),
6652 DyadOp::Less => {
6653 set_rank(cfg, "without", x, y, span)?;
6654 Ok(set_less(x, y, tol))
6655 }
6656 DyadOp::Union => {
6657 set_rank(cfg, "union", x, y, span)?;
6658 union_items(x, y, tol, span)
6659 }
6660 DyadOp::Intersect => {
6661 set_rank(cfg, "intersection", x, y, span)?;
6662 Ok(intersect_items(x, y, tol))
6663 }
6664 DyadOp::AnagramFrom => anagram_from(x, y, span),
6665 DyadOp::Permute => permute(x, y, span),
6666 DyadOp::FindSeq => {
6667 find_seq(x, y, tol, cfg.rules.lang == crate::Lang::Apl, span)
6668 }
6669 DyadOp::UnicodeForm => unicode_form(x, y, span),
6670 DyadOp::SymbolForm => symbol_form(x, y, span),
6671 DyadOp::PrimeMeta => prime_meta(x, y, span).map(|r| carry_exact2(r, x, y)),
6672 DyadOp::PrimeExponents => prime_exponents(x, y, span).map(|r| carry_exact2(r, x, y)),
6673 DyadOp::Pick { origin } => pick(x, y, origin, span),
6674 DyadOp::Expand => expand(x, y, span),
6675 DyadOp::WriteStream => Err(Error::internal("1!:2 reached the pure dyad dispatcher")),
6678 DyadOp::NotYet(what) => Err(Error::not_yet(what, span)),
6679 DyadOp::None => Err(Error::domain(format!("{} has no dyadic meaning", p.name), span)),
6680 }
6681}
6682
6683fn reduce_identity(v: &Verb, n: usize) -> Option<Data> {
6695 let Verb::Prim(p) = v else { return None };
6696 let DyadOp::Scalar(op) = p.dyad else { return None };
6697 let ints = |k: i64| Data::I64(vec![k; n].into());
6698 let bits = |k: u8| Data::Bool(vec![k; n].into());
6699 Some(match op {
6700 ScalarDyad::Add | ScalarDyad::Sub | ScalarDyad::Gcd | ScalarDyad::Residue => ints(0),
6701 ScalarDyad::Mul
6702 | ScalarDyad::DivJ
6703 | ScalarDyad::DivApl
6704 | ScalarDyad::Pow
6705 | ScalarDyad::Lcm
6706 | ScalarDyad::Root
6707 | ScalarDyad::Binomial => ints(1),
6708 ScalarDyad::Min => Data::F64(vec![f64::INFINITY; n].into()),
6709 ScalarDyad::Max => Data::F64(vec![f64::NEG_INFINITY; n].into()),
6710 ScalarDyad::Eq | ScalarDyad::Le | ScalarDyad::Ge => bits(1),
6711 ScalarDyad::Ne | ScalarDyad::Lt | ScalarDyad::Gt => bits(0),
6712 ScalarDyad::MakeComplex | ScalarDyad::PolarBy => return None,
6715 ScalarDyad::Log | ScalarDyad::Circle => return None,
6718 })
6719}
6720
6721fn is_associative(op: ScalarDyad) -> bool {
6727 use ScalarDyad::*;
6728 matches!(op, Add | Mul | Min | Max)
6729}
6730
6731#[inline(always)]
6732fn fold_range_body<S, T, F>(
6733 v: &[S],
6734 m: usize,
6735 lo: usize,
6736 hi: usize,
6737 j0: usize,
6738 acc: &mut [T],
6739 step: &F,
6740) -> bool
6741where
6742 S: Widen<T>,
6743 T: Copy,
6744 F: Fn(T, T) -> (T, bool),
6745{
6746 let w = acc.len();
6747 let base = (hi - 1) * m + j0;
6748 for (slot, &x) in acc.iter_mut().zip(&v[base..base + w]) {
6749 *slot = x.widen();
6750 }
6751 let mut over = false;
6754 for i in (lo..hi - 1).rev() {
6755 let row = &v[i * m + j0..i * m + j0 + w];
6756 for (slot, &x) in acc.iter_mut().zip(row) {
6757 let (r, o) = step(x.widen(), *slot);
6758 *slot = r;
6759 over |= o;
6760 }
6761 }
6762 !over
6763}
6764
6765multiversioned! {
6766 #[allow(clippy::too_many_arguments)]
6767 fn fold_range_vectorised[S: Widen<T>, T: Copy, F: Fn(T, T) -> (T, bool)](
6768 v: &[S],
6769 m: usize,
6770 lo: usize,
6771 hi: usize,
6772 j0: usize,
6773 acc: &mut [T],
6774 step: &F,
6775 ) -> bool = fold_range_body;
6776}
6777
6778const VECTOR_COLUMNS: usize = 16;
6786
6787#[allow(clippy::too_many_arguments)]
6800#[inline]
6801fn fold_range<S, T, F>(
6802 v: &[S],
6803 m: usize,
6804 lo: usize,
6805 hi: usize,
6806 j0: usize,
6807 acc: &mut [T],
6808 step: &F,
6809) -> bool
6810where
6811 S: Widen<T>,
6812 T: Copy,
6813 F: Fn(T, T) -> (T, bool),
6814{
6815 if acc.len() < VECTOR_COLUMNS {
6816 fold_range_body(v, m, lo, hi, j0, acc, step)
6817 } else {
6818 fold_range_vectorised(v, m, lo, hi, j0, acc, step)
6819 }
6820}
6821
6822const FOLD_LANES: usize = 8;
6833
6834const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
6839
6840#[inline(always)]
6845fn fold_lanes_body<S, T, F>(v: &[S], step: &F) -> Option<T>
6846where
6847 S: Widen<T>,
6848 T: Copy,
6849 F: Fn(T, T) -> (T, bool),
6850{
6851 let n = v.len();
6852 let mut over = false;
6853 if n < MIN_LANE_WORK {
6854 let mut acc = v[n - 1].widen();
6855 for &x in v[..n - 1].iter().rev() {
6856 let (r, o) = step(x.widen(), acc);
6857 acc = r;
6858 over |= o;
6859 }
6860 return (!over).then_some(acc);
6861 }
6862 let rows = n / FOLD_LANES;
6865 let head = n - rows * FOLD_LANES;
6866 let last = head + (rows - 1) * FOLD_LANES;
6867 let mut acc = [v[last].widen(); FOLD_LANES];
6868 for (slot, &x) in acc.iter_mut().zip(&v[last..last + FOLD_LANES]) {
6869 *slot = x.widen();
6870 }
6871 for r in (0..rows - 1).rev() {
6872 let row = &v[head + r * FOLD_LANES..head + (r + 1) * FOLD_LANES];
6873 for (slot, &x) in acc.iter_mut().zip(row) {
6874 let (r, o) = step(x.widen(), *slot);
6875 *slot = r;
6876 over |= o;
6877 }
6878 }
6879 let mut a = acc[FOLD_LANES - 1];
6880 for &x in acc[..FOLD_LANES - 1].iter().rev() {
6881 let (r, o) = step(x, a);
6882 a = r;
6883 over |= o;
6884 }
6885 for &x in v[..head].iter().rev() {
6886 let (r, o) = step(x.widen(), a);
6887 a = r;
6888 over |= o;
6889 }
6890 (!over).then_some(a)
6891}
6892
6893multiversioned! {
6894 fn fold_lanes_vectorised[S: Widen<T>, T: Copy, F: Fn(T, T) -> (T, bool)](
6895 v: &[S],
6896 step: &F,
6897 ) -> Option<T> = fold_lanes_body;
6898}
6899
6900#[inline]
6903fn fold_lanes<S, T, F>(v: &[S], step: &F) -> Option<T>
6904where
6905 S: Widen<T>,
6906 T: Copy,
6907 F: Fn(T, T) -> (T, bool),
6908{
6909 if v.len() < MIN_LANE_WORK {
6910 fold_lanes_body(v, step)
6911 } else {
6912 fold_lanes_vectorised(v, step)
6913 }
6914}
6915
6916fn fold_flat<S, T, F>(v: &[S], n: usize, assoc: bool, step: &F) -> Option<T>
6919where
6920 S: Widen<T>,
6921 T: Copy + Send + Sync,
6922 F: Fn(T, T) -> (T, bool) + Sync + Send,
6923{
6924 if assoc {
6925 return par::try_fold_chunks(
6926 &v[..n],
6927 |part| fold_lanes(part, step),
6928 |a, b| {
6929 let (r, o) = step(a, b);
6930 (!o).then_some(r)
6931 },
6932 );
6933 }
6934 let mut acc = v[n - 1].widen();
6935 let mut over = false;
6936 for &x in v[..n - 1].iter().rev() {
6937 let (r, o) = step(x.widen(), acc);
6938 acc = r;
6939 over |= o;
6940 }
6941 (!over).then_some(acc)
6942}
6943
6944fn fold_items<S, T, F>(v: &[S], n: usize, m: usize, assoc: bool, step: F) -> Option<Vec<T>>
6955where
6956 S: Widen<T>,
6957 T: Copy + Default + Send + Sync,
6958 F: Fn(T, T) -> (T, bool) + Sync + Send,
6959{
6960 if m >= par::WIDE_ITEM {
6961 let (out, ok) = par::fill_wide(m, n * m, |j0, acc: &mut [T]| {
6962 fold_range(v, m, 0, n, j0, acc, &step)
6963 });
6964 return ok.then_some(out);
6965 }
6966 if m == 1 {
6967 return fold_flat(v, n, assoc, &step).map(|x| vec![x]);
6968 }
6969 let chunks = if assoc { par::chunks(n, n * m) } else { 1 };
6970 if chunks < 2 {
6971 let mut acc = vec![T::default(); m];
6972 return fold_range(v, m, 0, n, 0, &mut acc, &step).then_some(acc);
6973 }
6974 let per = n.div_ceil(chunks);
6975 let parts = par::map_indexed(n.div_ceil(per), |c| {
6976 let mut acc = vec![T::default(); m];
6977 let ok = fold_range(v, m, c * per, ((c + 1) * per).min(n), 0, &mut acc, &step);
6978 ok.then_some(acc)
6979 });
6980 let mut it = parts.into_iter().rev();
6983 let mut acc = it.next()??;
6984 for part in it {
6985 let part = part?;
6986 let mut over = false;
6987 for (slot, &x) in acc.iter_mut().zip(&part) {
6988 let (r, o) = step(x, *slot);
6989 *slot = r;
6990 over |= o;
6991 }
6992 if over {
6993 return None;
6994 }
6995 }
6996 Some(acc)
6997}
6998
6999fn fold_i64<S: Widen<i64>>(op: ScalarDyad, v: &[S], n: usize, m: usize) -> Option<Vec<i64>> {
7002 use ScalarDyad::*;
7003 let assoc = is_associative(op);
7004 match op {
7005 Add => fold_items(v, n, m, assoc, i64::overflowing_add),
7006 Sub => fold_items(v, n, m, assoc, i64::overflowing_sub),
7007 Mul => fold_items(v, n, m, assoc, i64::overflowing_mul),
7008 Min => fold_items(v, n, m, assoc, |a: i64, b: i64| (a.min(b), false)),
7009 Max => fold_items(v, n, m, assoc, |a: i64, b: i64| (a.max(b), false)),
7010 _ => None,
7011 }
7012}
7013
7014fn fold_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize) -> Option<Vec<Cx>> {
7015 use ScalarDyad::*;
7016 let assoc = is_associative(op);
7017 match op {
7018 Add => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::add(a, b), false)),
7019 Sub => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::sub(a, b), false)),
7020 Mul => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::mul(a, b), false)),
7021 _ => None,
7023 }
7024}
7025
7026fn fold_f64(op: ScalarDyad, v: &[f64], n: usize, m: usize) -> Option<Vec<f64>> {
7027 use ScalarDyad::*;
7028 let assoc = is_associative(op);
7029 match op {
7030 Add => fold_items(v, n, m, assoc, |a: f64, b: f64| (a + b, false)),
7031 Sub => fold_items(v, n, m, assoc, |a: f64, b: f64| (a - b, false)),
7032 Mul => fold_items(v, n, m, assoc, |a: f64, b: f64| (a * b, false)),
7033 Min => fold_items(v, n, m, assoc, |a: f64, b: f64| (a.min(b), false)),
7034 Max => fold_items(v, n, m, assoc, |a: f64, b: f64| (a.max(b), false)),
7035 _ => None,
7036 }
7037}
7038
7039fn reduce_typed(op: ScalarDyad, d: &Data, n: usize, m: usize) -> Option<Data> {
7043 use ScalarDyad::*;
7044 if !matches!(op, Add | Sub | Mul | Min | Max) {
7047 return None;
7048 }
7049 match d {
7050 Data::F64(v) => Some(Data::F64(fold_f64(op, v, n, m)?.into())),
7051 Data::Complex(v) => Some(Data::Complex(fold_cx(op, v, n, m)?.into())),
7052 Data::I64(v) => Some(Data::I64(fold_i64(op, v, n, m)?.into())),
7053 Data::Bool(v) => Some(Data::I64(fold_i64(op, v.as_slice(), n, m)?.into())),
7057 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
7060 }
7061}
7062
7063#[inline(always)]
7069fn fold_runs_body<S, T, F>(v: &[S], start: usize, m: usize, out: &mut [T], step: &F) -> bool
7070where
7071 S: Widen<T>,
7072 T: Copy,
7073 F: Fn(T, T) -> (T, bool),
7074{
7075 let mut over = false;
7076 for (k, slot) in out.iter_mut().enumerate() {
7077 let run = &v[(start + k) * m..(start + k + 1) * m];
7078 let mut acc = run[m - 1].widen();
7079 for &x in run[..m - 1].iter().rev() {
7080 let (r, o) = step(x.widen(), acc);
7081 acc = r;
7082 over |= o;
7083 }
7084 *slot = acc;
7085 }
7086 !over
7087}
7088
7089multiversioned! {
7090 fn fold_runs_vectorised[S: Widen<T>, T: Copy, F: Fn(T, T) -> (T, bool)](
7091 v: &[S],
7092 start: usize,
7093 m: usize,
7094 out: &mut [T],
7095 step: &F,
7096 ) -> bool = fold_runs_body;
7097}
7098
7099fn fold_runs<S, T, F>(v: &[S], n: usize, m: usize, step: F) -> Option<Vec<T>>
7102where
7103 S: Widen<T>,
7104 T: Copy + Default + Send + Sync,
7105 F: Fn(T, T) -> (T, bool) + Sync + Send,
7106{
7107 let wide = m >= VECTOR_COLUMNS;
7111 let (out, ok) = par::fill_wide(n, n * m, |start, part: &mut [T]| {
7112 if wide {
7113 fold_runs_vectorised(v, start, m, part, &step)
7114 } else {
7115 fold_runs_body(v, start, m, part, &step)
7116 }
7117 });
7118 ok.then_some(out)
7119}
7120
7121fn fold_runs_data(op: ScalarDyad, d: &Data, n: usize, m: usize) -> Option<Data> {
7122 use ScalarDyad::*;
7123 match d {
7124 Data::F64(v) => Some(Data::F64(
7125 match op {
7126 Add => fold_runs(v, n, m, |a: f64, b: f64| (a + b, false)),
7127 Sub => fold_runs(v, n, m, |a: f64, b: f64| (a - b, false)),
7128 Mul => fold_runs(v, n, m, |a: f64, b: f64| (a * b, false)),
7129 Min => fold_runs(v, n, m, |a: f64, b: f64| (a.min(b), false)),
7130 Max => fold_runs(v, n, m, |a: f64, b: f64| (a.max(b), false)),
7131 _ => None,
7132 }?
7133 .into(),
7134 )),
7135 Data::I64(v) => Some(Data::I64(fold_runs_i64(op, v.as_slice(), n, m)?.into())),
7136 Data::Complex(v) => Some(Data::Complex(
7138 match op {
7139 Add => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::add(a, b), false)),
7140 Sub => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::sub(a, b), false)),
7141 Mul => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::mul(a, b), false)),
7142 _ => None,
7143 }?
7144 .into(),
7145 )),
7146 Data::Bool(v) => Some(Data::I64(fold_runs_i64(op, v.as_slice(), n, m)?.into())),
7148 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
7149 }
7150}
7151
7152fn fold_runs_i64<S: Widen<i64>>(op: ScalarDyad, v: &[S], n: usize, m: usize) -> Option<Vec<i64>> {
7154 use ScalarDyad::*;
7155 match op {
7156 Add => fold_runs(v, n, m, i64::overflowing_add),
7157 Sub => fold_runs(v, n, m, i64::overflowing_sub),
7158 Mul => fold_runs(v, n, m, i64::overflowing_mul),
7159 Min => fold_runs(v, n, m, |a: i64, b: i64| (a.min(b), false)),
7160 Max => fold_runs(v, n, m, |a: i64, b: i64| (a.max(b), false)),
7161 _ => None,
7162 }
7163}
7164
7165fn run_slices<T: Clone>(b: &Buf<T>, runs: usize, len: usize) -> Vec<&[T]> {
7181 if let Some(parts) = b.parts() && parts.len() == runs && parts.iter().all(|p| p.len() == len) {
7182 return parts.iter().map(Buf::as_slice).collect();
7183 }
7184 let flat = b.as_slice();
7185 (0..runs).map(|c| &flat[c * len..(c + 1) * len]).collect()
7186}
7187
7188fn fold_columns<S, T, F>(cols: &[&[S]], len: usize, assoc: bool, step: F) -> Option<Vec<T>>
7197where
7198 S: Widen<T>,
7199 T: Copy + Default + Send + Sync,
7200 F: Fn(T, T) -> (T, bool) + Sync + Send,
7201{
7202 if par::worth_it(len) {
7207 let mut out = Vec::with_capacity(cols.len());
7208 for c in cols {
7209 out.push(fold_flat(c, len, assoc, &step)?);
7210 }
7211 return Some(out);
7212 }
7213 let (out, ok) = par::fill_wide(cols.len(), cols.len() * len, |start, part: &mut [T]| {
7214 let mut ok = true;
7215 for (k, slot) in part.iter_mut().enumerate() {
7216 match fold_flat(cols[start + k], len, assoc, &step) {
7217 Some(v) => *slot = v,
7218 None => ok = false,
7219 }
7220 }
7221 ok
7222 });
7223 ok.then_some(out)
7224}
7225
7226fn fold_columns_data(op: ScalarDyad, d: &Data, runs: usize, len: usize) -> Option<Data> {
7228 use ScalarDyad::*;
7229 if !matches!(op, Add | Sub | Mul | Min | Max) {
7230 return None;
7231 }
7232 let assoc = is_associative(op);
7233 macro_rules! by {
7234 ($v:expr, $add:expr, $sub:expr, $mul:expr, $min:expr, $max:expr) => {{
7235 let cols = run_slices($v, runs, len);
7236 match op {
7237 Add => fold_columns(&cols, len, assoc, $add),
7238 Sub => fold_columns(&cols, len, assoc, $sub),
7239 Mul => fold_columns(&cols, len, assoc, $mul),
7240 Min => fold_columns(&cols, len, assoc, $min),
7241 Max => fold_columns(&cols, len, assoc, $max),
7242 _ => None,
7243 }?
7244 }};
7245 }
7246 match d {
7247 Data::F64(v) => Some(Data::F64(
7248 by!(
7249 v,
7250 |a: f64, b: f64| (a + b, false),
7251 |a: f64, b: f64| (a - b, false),
7252 |a: f64, b: f64| (a * b, false),
7253 |a: f64, b: f64| (a.min(b), false),
7254 |a: f64, b: f64| (a.max(b), false)
7255 )
7256 .into(),
7257 )),
7258 Data::I64(v) => Some(Data::I64(
7259 by!(
7260 v,
7261 i64::overflowing_add,
7262 i64::overflowing_sub,
7263 i64::overflowing_mul,
7264 |a: i64, b: i64| (a.min(b), false),
7265 |a: i64, b: i64| (a.max(b), false)
7266 )
7267 .into(),
7268 )),
7269 Data::Complex(v) => {
7270 if !matches!(op, Add | Sub | Mul) {
7271 return None;
7272 }
7273 Some(Data::Complex(
7274 by!(
7275 v,
7276 |a: Cx, b: Cx| (cx::add(a, b), false),
7277 |a: Cx, b: Cx| (cx::sub(a, b), false),
7278 |a: Cx, b: Cx| (cx::mul(a, b), false),
7279 |_: Cx, _: Cx| unreachable!("refused above"),
7280 |_: Cx, _: Cx| unreachable!("refused above")
7281 )
7282 .into(),
7283 ))
7284 }
7285 Data::Bool(v) => Some(Data::I64(
7289 by!(
7290 v,
7291 i64::overflowing_add,
7292 i64::overflowing_sub,
7293 i64::overflowing_mul,
7294 |a: i64, b: i64| (a.min(b), false),
7295 |a: i64, b: i64| (a.max(b), false)
7296 )
7297 .into(),
7298 )),
7299 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
7300 }
7301}
7302
7303fn reduce_columns(v: &Verb, y: &Array) -> Option<Array> {
7308 let Verb::Prim(p) = v else { return None };
7309 let DyadOp::Scalar(op) = p.dyad else { return None };
7310 if !y.dtype().is_numeric() {
7311 return None;
7312 }
7313 let n = y.shape[0];
7314 let m: usize = y.shape[1..].iter().product();
7315 if n == 0 || m == 0 {
7318 return None;
7319 }
7320 let shape = y.shape[1..].to_vec();
7321 if n == 1 {
7324 return Some(Array::col_major(shape, y.data.clone()));
7325 }
7326 let data = fold_columns_data(op, &y.data, m, n)?;
7327 Some(Array::col_major(shape, data))
7328}
7329
7330fn fold_across<S, T, F>(cols: &[&[S]], rows: usize, step: F) -> Option<Vec<T>>
7333where
7334 S: Widen<T>,
7335 T: Copy + Default + Send + Sync,
7336 F: Fn(T, T) -> (T, bool) + Sync + Send,
7337{
7338 let (last, rest) = cols.split_last()?;
7339 let (out, ok) = par::fill(rows, |start, part: &mut [T]| {
7340 let mut over = false;
7341 for (k, slot) in part.iter_mut().enumerate() {
7342 let i = start + k;
7343 let mut acc = last[i].widen();
7344 for c in rest.iter().rev() {
7345 let (r, o) = step(c[i].widen(), acc);
7346 acc = r;
7347 over |= o;
7348 }
7349 *slot = acc;
7350 }
7351 !over
7352 });
7353 ok.then_some(out)
7354}
7355
7356fn fold_across_data(op: ScalarDyad, d: &Data, rows: usize, cols: usize) -> Option<Data> {
7357 use ScalarDyad::*;
7358 if !matches!(op, Add | Sub | Mul | Min | Max) {
7359 return None;
7360 }
7361 macro_rules! by {
7362 ($v:expr, $add:expr, $sub:expr, $mul:expr, $min:expr, $max:expr) => {{
7363 let parts = run_slices($v, cols, rows);
7364 match op {
7365 Add => fold_across(&parts, rows, $add),
7366 Sub => fold_across(&parts, rows, $sub),
7367 Mul => fold_across(&parts, rows, $mul),
7368 Min => fold_across(&parts, rows, $min),
7369 Max => fold_across(&parts, rows, $max),
7370 _ => None,
7371 }?
7372 }};
7373 }
7374 match d {
7375 Data::F64(v) => Some(Data::F64(
7376 by!(
7377 v,
7378 |a: f64, b: f64| (a + b, false),
7379 |a: f64, b: f64| (a - b, false),
7380 |a: f64, b: f64| (a * b, false),
7381 |a: f64, b: f64| (a.min(b), false),
7382 |a: f64, b: f64| (a.max(b), false)
7383 )
7384 .into(),
7385 )),
7386 Data::I64(v) => Some(Data::I64(
7387 by!(
7388 v,
7389 i64::overflowing_add,
7390 i64::overflowing_sub,
7391 i64::overflowing_mul,
7392 |a: i64, b: i64| (a.min(b), false),
7393 |a: i64, b: i64| (a.max(b), false)
7394 )
7395 .into(),
7396 )),
7397 Data::Complex(v) => {
7398 if !matches!(op, Add | Sub | Mul) {
7399 return None;
7400 }
7401 Some(Data::Complex(
7402 by!(
7403 v,
7404 |a: Cx, b: Cx| (cx::add(a, b), false),
7405 |a: Cx, b: Cx| (cx::sub(a, b), false),
7406 |a: Cx, b: Cx| (cx::mul(a, b), false),
7407 |_: Cx, _: Cx| unreachable!("refused above"),
7408 |_: Cx, _: Cx| unreachable!("refused above")
7409 )
7410 .into(),
7411 ))
7412 }
7413 Data::Bool(v) => Some(Data::I64(
7415 by!(
7416 v,
7417 i64::overflowing_add,
7418 i64::overflowing_sub,
7419 i64::overflowing_mul,
7420 |a: i64, b: i64| (a.min(b), false),
7421 |a: i64, b: i64| (a.max(b), false)
7422 )
7423 .into(),
7424 )),
7425 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
7426 }
7427}
7428
7429fn reduce_rows_columns(u: &Verb, y: &Array) -> Option<Array> {
7432 let Verb::Reduce(inner) = u else { return None };
7433 let Verb::Prim(p) = &**inner else { return None };
7434 let DyadOp::Scalar(op) = p.dyad else { return None };
7435 if y.rank() != 2 || !y.dtype().is_numeric() {
7438 return None;
7439 }
7440 let (rows, cols) = (y.shape[0], y.shape[1]);
7441 if rows == 0 || cols == 0 {
7444 return None;
7445 }
7446 if cols == 1 {
7447 return Some(Array::new(vec![rows], y.data.clone()));
7449 }
7450 let data = fold_across_data(op, &y.data, rows, cols)?;
7451 Some(Array::new(vec![rows], data))
7452}
7453
7454fn reduce_vector_cells(u: &Verb, y: &Array, frame_rank: usize) -> Option<Array> {
7463 let Verb::Reduce(inner) = u else { return None };
7464 let Verb::Prim(p) = &**inner else { return None };
7465 let DyadOp::Scalar(op) = p.dyad else { return None };
7466 if y.rank() != frame_rank + 1 || !y.dtype().is_numeric() {
7469 return None;
7470 }
7471 let m = y.shape[frame_rank];
7472 if m == 0 {
7475 return None;
7476 }
7477 use ScalarDyad::{Add, Max, Min, Mul, Sub};
7478 if !matches!(op, Add | Sub | Mul | Min | Max) {
7479 return None;
7480 }
7481 let frame = y.shape[..frame_rank].to_vec();
7482 if m == 1 {
7483 return Some(Array::new(frame, y.data.clone()));
7486 }
7487 let n: usize = frame.iter().product();
7488 let data = fold_runs_data(op, &y.data, n, m)?;
7489 Some(Array::new(frame, data))
7490}
7491
7492fn reduce(v: &Verb, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
7494 if y.rank() == 0 {
7495 return Ok(y.clone());
7496 }
7497 let n = y.items();
7498 if n == 1 {
7499 return Ok(y.item(0));
7500 }
7501 let cell_shape = y.shape[1..].to_vec();
7502 let m: usize = cell_shape.iter().product();
7503 if n == 0 {
7504 if matches!(v, Verb::Prim(p) if matches!(p.dyad, DyadOp::AppendLeading | DyadOp::AppendLast))
7507 {
7508 return Ok(Array::new(vec![0], Data::empty(y.dtype())));
7509 }
7510 return match reduce_identity(v, m) {
7511 Some(d) => Ok(Array::new(cell_shape, d)),
7512 None => Err(Error::domain(
7513 format!("empty reduction has no identity for {}", v.name()),
7514 span,
7515 )),
7516 };
7517 }
7518 if y.dtype().is_numeric() && let Verb::Prim(p) = v && let DyadOp::Scalar(op) = p.dyad {
7519 if let Some(d) = reduce_typed(op, y.row_major_data(), n, m) {
7524 return Ok(Array::new(cell_shape, d));
7525 }
7526 let mut acc = y.data.slice((n - 1) * m, n * m);
7529 for i in (0..n - 1).rev() {
7530 acc =
7531 scalar_dyad_data(op, &y.data, i * m, 1, &acc, 0, 1, m, ctx.cfg.tol, span)?;
7532 }
7533 return Ok(Array::new(cell_shape, acc));
7534 }
7535 let mut acc = y.item(n - 1);
7536 for i in (0..n - 1).rev() {
7537 acc = v.dyad(&y.item(i), &acc, ctx, span)?;
7538 }
7539 Ok(acc)
7540}
7541
7542fn folded_op(u: &Verb) -> Option<ScalarDyad> {
7549 let Verb::Reduce(inner) = u else { return None };
7550 let Verb::Prim(p) = &**inner else { return None };
7551 match p.dyad {
7552 DyadOp::Scalar(op) => Some(op),
7553 _ => None,
7554 }
7555}
7556
7557fn section(y: &Array, lo: usize, hi: usize) -> Array {
7559 let m = y.item_size();
7560 let mut shape = y.shape.clone();
7561 shape[0] = hi - lo;
7562 Array::new(shape, y.data.slice(lo * m, hi * m))
7563}
7564
7565fn as_items(y: &Array) -> Option<Array> {
7568 (y.rank() == 0).then(|| Array::new(vec![1], y.data.clone()))
7569}
7570
7571#[inline(always)]
7572fn scan_flat_body<S, T, F>(v: &[S], n: usize, m: usize, back: bool, step: F) -> Option<Vec<T>>
7573where
7574 S: Widen<T>,
7575 T: Copy + Default,
7576 F: Fn(T, T) -> (T, bool),
7577{
7578 if m == 1 {
7579 let mut out = vec![T::default(); n];
7582 let mut over = false;
7583 if back {
7584 let mut acc = v[n - 1].widen();
7585 out[n - 1] = acc;
7586 for (slot, &x) in out[..n - 1].iter_mut().zip(&v[..n - 1]).rev() {
7587 let (r, o) = step(x.widen(), acc);
7588 acc = r;
7589 over |= o;
7590 *slot = acc;
7591 }
7592 } else {
7593 let mut acc = v[0].widen();
7594 out[0] = acc;
7595 for (slot, &x) in out[1..n].iter_mut().zip(&v[1..n]) {
7596 let (r, o) = step(acc, x.widen());
7597 acc = r;
7598 over |= o;
7599 *slot = acc;
7600 }
7601 }
7602 return (!over).then_some(out);
7603 }
7604 let mut out = vec![T::default(); n * m];
7605 let mut acc = vec![T::default(); m];
7606 let mut over = false;
7607 if back {
7608 for (slot, &x) in acc.iter_mut().zip(&v[(n - 1) * m..n * m]) {
7609 *slot = x.widen();
7610 }
7611 out[(n - 1) * m..n * m].copy_from_slice(&acc);
7612 for i in (0..n - 1).rev() {
7613 for (j, slot) in acc.iter_mut().enumerate() {
7614 let (r, o) = step(v[i * m + j].widen(), *slot);
7615 *slot = r;
7616 over |= o;
7617 }
7618 out[i * m..i * m + m].copy_from_slice(&acc);
7619 }
7620 } else {
7621 for (slot, &x) in acc.iter_mut().zip(&v[..m]) {
7622 *slot = x.widen();
7623 }
7624 out[..m].copy_from_slice(&acc);
7625 for i in 1..n {
7626 for (j, slot) in acc.iter_mut().enumerate() {
7627 let (r, o) = step(*slot, v[i * m + j].widen());
7628 *slot = r;
7629 over |= o;
7630 }
7631 out[i * m..i * m + m].copy_from_slice(&acc);
7632 }
7633 }
7634 (!over).then_some(out)
7635}
7636
7637multiversioned! {
7638 fn scan_flat_vectorised[S: Widen<T>, T: Copy + Default, F: Fn(T, T) -> (T, bool)](
7639 v: &[S],
7640 n: usize,
7641 m: usize,
7642 back: bool,
7643 step: F,
7644 ) -> Option<Vec<T>> = scan_flat_body;
7645}
7646
7647fn scan_flat<S, T, F>(v: &[S], n: usize, m: usize, back: bool, step: F) -> Option<Vec<T>>
7659where
7660 S: Widen<T>,
7661 T: Copy + Default,
7662 F: Fn(T, T) -> (T, bool),
7663{
7664 if m < VECTOR_COLUMNS {
7665 scan_flat_body(v, n, m, back, step)
7666 } else {
7667 scan_flat_vectorised(v, n, m, back, step)
7668 }
7669}
7670
7671fn scan_i64<S: Widen<i64>>(
7672 op: ScalarDyad,
7673 v: &[S],
7674 n: usize,
7675 m: usize,
7676 back: bool,
7677) -> Option<Vec<i64>> {
7678 use ScalarDyad::*;
7679 match op {
7680 Add => scan_flat(v, n, m, back, i64::overflowing_add),
7681 Sub => scan_flat(v, n, m, back, i64::overflowing_sub),
7682 Mul => scan_flat(v, n, m, back, i64::overflowing_mul),
7683 Min => scan_flat(v, n, m, back, |a: i64, b: i64| (a.min(b), false)),
7684 Max => scan_flat(v, n, m, back, |a: i64, b: i64| (a.max(b), false)),
7685 _ => None,
7686 }
7687}
7688
7689fn scan_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize, back: bool) -> Option<Vec<Cx>> {
7690 use ScalarDyad::*;
7691 match op {
7692 Add => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::add(a, b), false)),
7693 Sub => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::sub(a, b), false)),
7694 Mul => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::mul(a, b), false)),
7695 _ => None,
7696 }
7697}
7698
7699fn scan_f64<S: Widen<f64>>(
7700 op: ScalarDyad,
7701 v: &[S],
7702 n: usize,
7703 m: usize,
7704 back: bool,
7705) -> Option<Vec<f64>> {
7706 use ScalarDyad::*;
7707 match op {
7708 Add => scan_flat(v, n, m, back, |a: f64, b: f64| (a + b, false)),
7709 Sub => scan_flat(v, n, m, back, |a: f64, b: f64| (a - b, false)),
7710 Mul => scan_flat(v, n, m, back, |a: f64, b: f64| (a * b, false)),
7711 Min => scan_flat(v, n, m, back, |a: f64, b: f64| (a.min(b), false)),
7712 Max => scan_flat(v, n, m, back, |a: f64, b: f64| (a.max(b), false)),
7713 _ => None,
7714 }
7715}
7716
7717fn scan_typed(op: ScalarDyad, d: &Data, n: usize, m: usize, back: bool) -> Option<Data> {
7721 use ScalarDyad::*;
7722 if !matches!(op, Add | Sub | Mul | Min | Max) {
7723 return None;
7724 }
7725 fn ints<S: Widen<i64> + Widen<f64>>(
7729 op: ScalarDyad,
7730 v: &[S],
7731 n: usize,
7732 m: usize,
7733 back: bool,
7734 ) -> Data {
7735 match scan_i64(op, v, n, m, back) {
7736 Some(out) => Data::I64(out.into()),
7737 None => Data::F64(
7738 scan_f64(op, v, n, m, back).expect("the float scan cannot overflow").into(),
7739 ),
7740 }
7741 }
7742 match d {
7743 Data::F64(v) => Some(Data::F64(scan_f64(op, v.as_slice(), n, m, back)?.into())),
7744 Data::Complex(v) => Some(Data::Complex(scan_cx(op, v, n, m, back)?.into())),
7745 Data::I64(v) => Some(ints(op, v.as_slice(), n, m, back)),
7746 Data::Bool(v) => Some(ints(op, v.as_slice(), n, m, back)),
7747 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
7750 }
7751}
7752
7753fn affine_step(u: &Verb) -> Option<&Array> {
7760 fn prim(v: &Verb, want: DyadOp, ranks: [i64; 3]) -> bool {
7764 matches!(v, Verb::Prim(p) if p.dyad == want && p.ranks == ranks)
7765 }
7766 const ATOMS: [i64; 3] = [0, 0, 0];
7767 const WHOLE: [i64; 3] = [RANK_INF; 3];
7768 fn scaled(v: &Verb) -> Option<&Array> {
7770 let Verb::NounFork(c, g, h) = v else { return None };
7771 let noun = c.rank() == 0
7772 && matches!(c.dtype(), DType::Bool | DType::I64 | DType::F64 | DType::Complex);
7773 let tree = prim(g, DyadOp::Scalar(ScalarDyad::Mul), ATOMS)
7774 && prim(h, DyadOp::Right, WHOLE);
7775 (noun && tree).then_some(c)
7776 }
7777 let Verb::Fork(f, g, h) = u else { return None };
7778 if !prim(g, DyadOp::Scalar(ScalarDyad::Add), ATOMS) {
7779 return None;
7780 }
7781 if prim(f, DyadOp::Left, WHOLE) {
7782 scaled(h)
7783 } else if prim(h, DyadOp::Left, WHOLE) {
7784 scaled(f)
7785 } else {
7786 None
7787 }
7788}
7789
7790struct Ring<T> {
7793 add: fn(T, T) -> T,
7794 mul: fn(T, T) -> T,
7795 one: T,
7796 finite: fn(T) -> bool,
7797}
7798
7799fn affine_flat<T>(v: &[T], c: T, n: usize, m: usize, back: bool, r: &Ring<T>) -> Option<Vec<T>>
7805where
7806 T: Copy + Default,
7807{
7808 let (add, mul) = (r.add, r.mul);
7809 let mut out = vec![T::default(); n * m];
7810 if back {
7811 out[(n - 1) * m..].copy_from_slice(&v[(n - 1) * m..n * m]);
7812 for i in (0..n - 1).rev() {
7813 for j in 0..m {
7814 out[i * m + j] = add(v[i * m + j], mul(c, out[(i + 1) * m + j]));
7815 }
7816 }
7817 } else {
7818 out[..m].copy_from_slice(&v[..m]);
7819 let mut pow = r.one;
7820 for i in 1..n {
7821 pow = mul(pow, c);
7822 if !(r.finite)(pow) {
7823 return None;
7824 }
7825 for j in 0..m {
7826 out[i * m + j] = add(out[(i - 1) * m + j], mul(pow, v[i * m + j]));
7827 }
7828 }
7829 }
7830 Some(out)
7831}
7832
7833fn affine_scan(c: &Array, y: &Array, back: bool) -> Option<Data> {
7843 let (n, m) = (y.items(), y.item_size());
7844 let machine = |t: DType| matches!(t, DType::Bool | DType::I64 | DType::F64 | DType::Complex);
7845 if n == 0 || !machine(c.dtype()) || !machine(y.dtype()) {
7846 return None;
7847 }
7848 match DType::promote(c.dtype(), y.dtype())? {
7849 DType::F64 => {
7850 let (mut tc, mut tv) = (Vec::new(), Vec::new());
7851 let k = *borrow_f64(&c.data, &mut tc).first()?;
7852 let v = borrow_f64(y.row_major_data(), &mut tv);
7853 let r = Ring { add: |a, b| a + b, mul: |a, b| a * b, one: 1.0, finite: f64::is_finite };
7854 Some(Data::F64(affine_flat(v, k, n, m, back, &r)?.into()))
7855 }
7856 DType::Complex => {
7857 let (mut tc, mut tv) = (Vec::new(), Vec::new());
7858 let k = *borrow_cx(&c.data, &mut tc).first()?;
7859 let v = borrow_cx(y.row_major_data(), &mut tv);
7860 let finite = |z: Cx| z[0].is_finite() && z[1].is_finite();
7861 let r = Ring { add: cx::add, mul: cx::mul, one: [1.0, 0.0], finite };
7862 Some(Data::Complex(affine_flat(v, k, n, m, back, &r)?.into()))
7863 }
7864 _ => None,
7865 }
7866}
7867
7868fn window_fold<S, T, F>(v: &[S], n: usize, m: usize, w: usize, step: F) -> Option<Vec<T>>
7883where
7884 S: Widen<T>,
7885 T: Copy + Default + Send + Sync,
7886 F: Fn(T, T) -> (T, bool) + Sync + Send,
7887{
7888 debug_assert!(w >= 1 && n >= w);
7889 if m == 1 {
7890 return window_fold_flat(v, n, w, step);
7891 }
7892 let count = n - w + 1;
7893 let mut out = vec![T::default(); count * m];
7894 let mut pre = vec![T::default(); w * m];
7897 let mut suf = vec![T::default(); w * m];
7898 let mut prev = vec![T::default(); w * m];
7899 let mut over = false;
7900 for b in 0..n.div_ceil(w) {
7901 let bs = b * w;
7902 let be = ((b + 1) * w).min(n);
7903 for (slot, &x) in pre[..m].iter_mut().zip(&v[bs * m..bs * m + m]) {
7904 *slot = x.widen();
7905 }
7906 for i in 1..be - bs {
7907 let (o, p) = (i * m, (i - 1) * m);
7908 for j in 0..m {
7909 let (r, f) = step(pre[p + j], v[(bs + i) * m + j].widen());
7910 pre[o + j] = r;
7911 over |= f;
7912 }
7913 }
7914 for e in bs.max(w - 1)..be {
7917 let i = e + 1 - w;
7918 let (oo, po) = (i * m, (e - bs) * m);
7919 if i == bs {
7920 out[oo..oo + m].copy_from_slice(&pre[po..po + m]);
7921 } else {
7922 let so = (i + w - bs) * m;
7923 for j in 0..m {
7924 let (r, f) = step(prev[so + j], pre[po + j]);
7925 out[oo + j] = r;
7926 over |= f;
7927 }
7928 }
7929 }
7930 let last = be - 1 - bs;
7931 for (slot, &x) in suf[last * m..last * m + m]
7932 .iter_mut()
7933 .zip(&v[(be - 1) * m..be * m])
7934 {
7935 *slot = x.widen();
7936 }
7937 for i in (0..last).rev() {
7938 let (o, p) = (i * m, (i + 1) * m);
7939 for j in 0..m {
7940 let (r, f) = step(v[(bs + i) * m + j].widen(), suf[p + j]);
7941 suf[o + j] = r;
7942 over |= f;
7943 }
7944 }
7945 std::mem::swap(&mut prev, &mut suf);
7946 }
7947 (!over).then_some(out)
7948}
7949
7950fn window_fold_flat<S, T, F>(v: &[S], n: usize, w: usize, step: F) -> Option<Vec<T>>
7960where
7961 S: Widen<T>,
7962 T: Copy + Default + Send + Sync,
7963 F: Fn(T, T) -> (T, bool) + Sync + Send,
7964{
7965 let (out, ok) = par::fill(n - w + 1, |lo, part: &mut [T]| {
7966 window_fold_range(v, n, w, lo, part, &step)
7967 });
7968 ok.then_some(out)
7969}
7970
7971#[inline(always)]
7972fn window_fold_range_body<S, T, F>(
7973 v: &[S],
7974 n: usize,
7975 w: usize,
7976 lo: usize,
7977 out: &mut [T],
7978 step: &F,
7979) -> bool
7980where
7981 S: Widen<T>,
7982 T: Copy + Default,
7983 F: Fn(T, T) -> (T, bool),
7984{
7985 if out.is_empty() {
7986 return true;
7987 }
7988 let hi = lo + out.len();
7989 let mut pre = vec![T::default(); w];
7990 let mut suf = vec![T::default(); w];
7991 let mut prev = vec![T::default(); w];
7992 let mut over = false;
7993 let mut bs = lo / w * w;
7994 while bs < n && bs <= hi + w - 2 {
7996 let block = &v[bs..(bs + w).min(n)];
7997 let lb = block.len();
7998 let mut acc = block[0].widen();
7999 pre[0] = acc;
8000 for (slot, &x) in pre[1..lb].iter_mut().zip(&block[1..]) {
8001 let (r, o) = step(acc, x.widen());
8002 acc = r;
8003 over |= o;
8004 *slot = acc;
8005 }
8006 for e in bs.max(lo + w - 1)..(bs + lb).min(hi + w - 1) {
8011 let i = e + 1 - w;
8012 out[i - lo] = if i == bs {
8013 pre[e - bs]
8014 } else {
8015 let (r, o) = step(prev[i + w - bs], pre[e - bs]);
8016 over |= o;
8017 r
8018 };
8019 }
8020 let mut acc = block[lb - 1].widen();
8021 suf[lb - 1] = acc;
8022 for (slot, &x) in suf[..lb - 1].iter_mut().zip(&block[..lb - 1]).rev() {
8023 let (r, o) = step(x.widen(), acc);
8024 acc = r;
8025 over |= o;
8026 *slot = acc;
8027 }
8028 std::mem::swap(&mut prev, &mut suf);
8029 bs += w;
8030 }
8031 !over
8032}
8033
8034multiversioned! {
8035 fn window_fold_range[S: Widen<T>, T: Copy + Default, F: Fn(T, T) -> (T, bool)](
8040 v: &[S],
8041 n: usize,
8042 w: usize,
8043 lo: usize,
8044 out: &mut [T],
8045 step: &F,
8046 ) -> bool = window_fold_range_body;
8047}
8048
8049pub(crate) fn windows_into<S, T, F>(v: &[S], w: usize, lo: usize, out: &mut [T], step: &F) -> bool
8058where
8059 S: Widen<T>,
8060 T: Copy + Default,
8061 F: Fn(T, T) -> (T, bool),
8062{
8063 window_fold_range(v, v.len(), w, lo, out, step)
8064}
8065
8066fn window_i64<S: Widen<i64>>(
8067 op: ScalarDyad,
8068 v: &[S],
8069 n: usize,
8070 m: usize,
8071 w: usize,
8072) -> Option<Vec<i64>> {
8073 use ScalarDyad::*;
8074 match op {
8075 Add => window_fold(v, n, m, w, i64::overflowing_add),
8076 Mul => window_fold(v, n, m, w, i64::overflowing_mul),
8077 Min => window_fold(v, n, m, w, |a: i64, b: i64| (a.min(b), false)),
8078 Max => window_fold(v, n, m, w, |a: i64, b: i64| (a.max(b), false)),
8079 _ => None,
8080 }
8081}
8082
8083fn window_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize, w: usize) -> Option<Vec<Cx>> {
8084 use ScalarDyad::*;
8085 match op {
8086 Add => window_fold(v, n, m, w, |a: Cx, b: Cx| (cx::add(a, b), false)),
8087 Mul => window_fold(v, n, m, w, |a: Cx, b: Cx| (cx::mul(a, b), false)),
8088 _ => None,
8089 }
8090}
8091
8092fn window_f64<S: Widen<f64>>(
8093 op: ScalarDyad,
8094 v: &[S],
8095 n: usize,
8096 m: usize,
8097 w: usize,
8098) -> Option<Vec<f64>> {
8099 use ScalarDyad::*;
8100 match op {
8101 Add => window_fold(v, n, m, w, |a: f64, b: f64| (a + b, false)),
8102 Mul => window_fold(v, n, m, w, |a: f64, b: f64| (a * b, false)),
8103 Min => window_fold(v, n, m, w, |a: f64, b: f64| (a.min(b), false)),
8104 Max => window_fold(v, n, m, w, |a: f64, b: f64| (a.max(b), false)),
8105 _ => None,
8106 }
8107}
8108
8109fn window_typed(op: ScalarDyad, d: &Data, n: usize, m: usize, w: usize) -> Option<Data> {
8113 use ScalarDyad::*;
8114 if !matches!(op, Add | Mul | Min | Max) {
8115 return None;
8116 }
8117 fn ints<S: Widen<i64> + Widen<f64>>(
8120 op: ScalarDyad,
8121 v: &[S],
8122 n: usize,
8123 m: usize,
8124 w: usize,
8125 ) -> Data {
8126 match window_i64(op, v, n, m, w) {
8127 Some(out) => Data::I64(out.into()),
8128 None => {
8129 Data::F64(window_f64(op, v, n, m, w).expect("the float fold cannot overflow").into())
8130 }
8131 }
8132 }
8133 match d {
8134 Data::F64(v) => Some(Data::F64(window_f64(op, v.as_slice(), n, m, w)?.into())),
8135 Data::Complex(v) => Some(Data::Complex(window_cx(op, v, n, m, w)?.into())),
8136 Data::I64(v) => Some(ints(op, v.as_slice(), n, m, w)),
8137 Data::Bool(v) => Some(ints(op, v.as_slice(), n, m, w)),
8138 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
8141 }
8142}
8143
8144fn runs(u: &Verb, y: &Array, back: bool, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
8146 let promoted = as_items(y);
8147 let base = promoted.as_ref().unwrap_or(y);
8148 let n = base.items();
8149 let m = base.item_size();
8150 if n > 0 && base.dtype().is_numeric() && let Some(op) = folded_op(u) {
8151 if (back || is_associative(op))
8154 && let Some(d) = scan_typed(op, base.row_major_data(), n, m, back)
8155 {
8156 return Ok(Array::new(base.shape.clone(), d));
8157 }
8158 }
8159 if n > 0 && let Verb::Reduce(inner) = u {
8160 if base.dtype().is_numeric()
8161 && let Some(c) = affine_step(inner)
8162 && let Some(d) = affine_scan(c, base, back)
8163 {
8164 return Ok(Array::new(base.shape.clone(), d));
8165 }
8166 if back && u.is_pure() {
8173 let mut acc = base.item(n - 1);
8174 let mut cells = Vec::with_capacity(n);
8175 cells.push(acc.clone());
8176 for i in (0..n - 1).rev() {
8177 acc = inner.dyad(&base.item(i), &acc, ctx, span)?;
8178 cells.push(acc.clone());
8179 }
8180 cells.reverse();
8181 return assemble(&[n], cells, span);
8182 }
8183 }
8184 let cells = each_cell(n, n * m, u.is_pure(), ctx, |i, c| {
8185 let part = if back { section(base, i, n) } else { section(base, 0, i + 1) };
8186 u.monad(&part, c, span)
8187 })?;
8188 assemble(&[n], cells, span)
8189}
8190
8191fn empty_windows(u: &Verb, y: &Array, w: usize, ctx: &mut Ctx<'_>, span: Span) -> Array {
8196 let m = y.item_size();
8197 if u.is_pure() && let Some(cells) = w.checked_mul(m).filter(|&s| s <= 1 << 20) {
8198 let mut shape = y.shape.clone();
8199 shape[0] = w;
8200 let probe = Array::new(shape, fill_data(y.dtype(), cells));
8201 if let Ok(cell) = u.monad(&probe, ctx, span) {
8202 let mut shape = vec![0usize];
8203 shape.extend_from_slice(&cell.shape);
8204 return Array::new(shape, Data::empty(cell.dtype()));
8205 }
8206 }
8207 Array::new(vec![0], Data::empty(DType::I64))
8208}
8209
8210fn window_size(x: &Array, span: Span) -> Result<i64> {
8212 let v = x
8213 .to_i64_vec()
8214 .ok_or_else(|| Error::domain("the window size must be an integer", span))?;
8215 match v.as_slice() {
8216 [k] => Ok(*k),
8217 _ => Err(Error::new(
8218 ErrorKind::Length,
8219 "the window size must be a single number",
8220 Some(span),
8221 )),
8222 }
8223}
8224
8225fn infix(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
8233 let k = window_size(x, span)?;
8234 let promoted = as_items(y);
8235 let base = promoted.as_ref().unwrap_or(y);
8236 let n = base.items();
8237 let m = base.item_size();
8238 if k < 0 {
8239 let w = k.unsigned_abs() as usize;
8240 let count = n.div_ceil(w);
8241 let cells = each_cell(count, n * m, u.is_pure(), ctx, |i, c| {
8242 u.monad(§ion(base, i * w, ((i + 1) * w).min(n)), c, span)
8243 })?;
8244 return assemble(&[count], cells, span);
8245 }
8246 let w = k as usize;
8247 if n < w {
8248 return Ok(empty_windows(u, base, w, ctx, span));
8249 }
8250 let count = n - w + 1;
8251 if w > 0 && base.dtype().is_numeric()
8252 && let Some(op) = folded_op(u) && let Some(d) = window_typed(op, &base.data, n, m, w)
8253 {
8254 let mut shape = base.shape.clone();
8255 shape[0] = count;
8256 return Ok(Array::new(shape, d));
8257 }
8258 let work = count.saturating_mul(w).saturating_mul(m);
8259 let cells = each_cell(count, work, u.is_pure(), ctx, |i, c| {
8260 u.monad(§ion(base, i, i + w), c, span)
8261 })?;
8262 assemble(&[count], cells, span)
8263}
8264
8265fn power(
8268 u: &Verb,
8269 p: Power,
8270 x: Option<&Array>,
8271 y: &Array,
8272 ctx: &mut Ctx<'_>,
8273 span: Span,
8274) -> Result<Array> {
8275 let step = |acc: &Array, c: &mut Ctx<'_>| match x {
8276 Some(x) => u.dyad(x, acc, c, span),
8277 None => u.monad(acc, c, span),
8278 };
8279 match p {
8280 Power::Times(n) => {
8281 let mut acc = y.clone();
8282 for _ in 0..n {
8283 acc = step(&acc, ctx)?;
8284 }
8285 Ok(acc)
8286 }
8287 Power::Converge => {
8288 let mut acc = y.clone();
8289 for _ in 0..CONVERGE_LIMIT {
8290 let next = step(&acc, ctx)?;
8291 if arrays_match(&next, &acc, ctx.cfg.tol) {
8292 return Ok(next);
8293 }
8294 acc = next;
8295 }
8296 Err(Error::domain("the iteration did not converge", span))
8297 }
8298 Power::Each(ref counts) => {
8302 let mut acc = y.clone();
8303 let mut done = 0u64;
8304 let mut order: Vec<usize> = (0..counts.len()).collect();
8305 order.sort_by_key(|&i| counts[i]);
8306 let mut cells: Vec<Option<Array>> = vec![None; counts.len()];
8307 for i in order {
8308 while done < counts[i] {
8309 acc = step(&acc, ctx)?;
8310 done += 1;
8311 }
8312 cells[i] = Some(acc.clone());
8313 }
8314 let cells: Vec<Array> = cells.into_iter().map(|c| c.expect("every count filled")).collect();
8315 assemble(&[cells.len()], cells, span)
8316 }
8317 Power::ConvergeTrace => {
8318 let mut acc = y.clone();
8319 let mut cells = vec![acc.clone()];
8320 for _ in 0..CONVERGE_LIMIT {
8321 let next = step(&acc, ctx)?;
8322 if arrays_match(&next, &acc, ctx.cfg.tol) {
8323 return assemble(&[cells.len()], cells, span);
8324 }
8325 cells.push(next.clone());
8326 acc = next;
8327 }
8328 Err(Error::domain("the iteration did not converge", span))
8329 }
8330 }
8331}
8332
8333fn power_v(
8336 u: &Verb,
8337 v: &Verb,
8338 x: Option<&Array>,
8339 y: &Array,
8340 ctx: &mut Ctx<'_>,
8341 span: Span,
8342) -> Result<Array> {
8343 let count = match x {
8344 Some(x) => v.dyad(x, y, ctx, span)?,
8345 None => v.monad(y, ctx, span)?,
8346 };
8347 let n = count
8348 .to_i64_vec()
8349 .ok_or_else(|| Error::domain("the power count must be an integer", span))?;
8350 if n.len() != 1 {
8351 return Err(Error::not_yet("a list of power counts (u^:v with several)", span));
8352 }
8353 let n = n[0];
8354 if n < 0 {
8355 return Err(Error::not_yet("a negative power (the verb's inverse)", span));
8356 }
8357 power(u, Power::Times(n as u64), x, y, ctx, span)
8358}
8359
8360fn power_until(
8362 u: &Verb,
8363 test: &Verb,
8364 y: &Array,
8365 ctx: &mut Ctx<'_>,
8366 span: Span,
8367) -> Result<Array> {
8368 let mut acc = y.clone();
8369 for _ in 0..CONVERGE_LIMIT {
8370 let next = u.monad(&acc, ctx, span)?;
8371 let done = test.dyad(&next, &acc, ctx, span)?;
8372 let stop = done
8373 .to_f64_vec()
8374 .ok_or_else(|| Error::domain("the ⍣ test must answer with numbers", span))?;
8375 if !stop.is_empty() && stop.iter().all(|&v| v != 0.0) {
8376 return Ok(next);
8377 }
8378 acc = next;
8379 }
8380 Err(Error::domain("the iteration did not converge", span))
8381}
8382
8383fn along_axis(
8390 u: &Verb,
8391 x: Option<&Array>,
8392 y: &Array,
8393 k: usize,
8394 ctx: &mut Ctx<'_>,
8395 span: Span,
8396) -> Result<Array> {
8397 if k >= y.rank().max(1) {
8398 return Err(Error::new(
8399 ErrorKind::Rank,
8400 format!("axis {k} does not exist on an argument of rank {}", y.rank()),
8401 Some(span),
8402 ));
8403 }
8404 let moved = axis_to_front(y, k);
8405 let r = moved.rank();
8406 let out = match x {
8407 Some(x) => u.dyad(x, &moved, ctx, span)?,
8408 None => u.monad(&moved, ctx, span)?,
8409 };
8410 if out.rank() == r {
8411 return Ok(front_to_axis(&out, k));
8412 }
8413 Ok(out)
8414}
8415
8416fn where_indices(y: &Array, origin: i64, boxed: bool, span: Span) -> Result<Array> {
8424 let counts = y
8425 .to_i64_vec()
8426 .ok_or_else(|| Error::domain("indices needs non-negative integers", span))?;
8427 if counts.iter().any(|&c| c < 0) {
8428 return Err(Error::domain("indices needs non-negative integers", span));
8429 }
8430 if !boxed || y.rank() < 2 {
8431 let mut out = Vec::new();
8432 for (i, &c) in counts.iter().enumerate() {
8433 for _ in 0..c {
8434 out.push(origin + i as i64);
8435 }
8436 }
8437 return Ok(Array::from_i64(out));
8438 }
8439 let r = y.rank();
8440 let mut coord = vec![0usize; r];
8441 let mut out: Vec<Array> = Vec::new();
8442 for &c in &counts {
8443 if c > 0 {
8444 let point =
8445 Array::from_i64(coord.iter().map(|&k| origin + k as i64).collect::<Vec<_>>());
8446 for _ in 0..c {
8447 out.push(point.clone());
8448 }
8449 }
8450 odometer(&mut coord, &y.shape);
8451 }
8452 Ok(Array::new(vec![out.len()], Data::Box(out.into())))
8453}
8454
8455fn interval_index(
8461 x: &Array,
8462 y: &Array,
8463 offset: i64,
8464 closed: bool,
8465 tol: Tol,
8466 span: Span,
8467) -> Result<Array> {
8468 if !x.dtype().is_numeric() || !y.dtype().is_numeric() {
8471 return ordered_interval_index(x, y, offset, closed, span);
8472 }
8473 let bounds = x
8474 .to_f64_vec()
8475 .ok_or_else(|| Error::domain("interval index needs numeric bounds", span))?;
8476 let vals = y
8477 .to_f64_vec()
8478 .ok_or_else(|| Error::domain("interval index needs numeric values", span))?;
8479 let out: Vec<i64> = vals
8480 .iter()
8481 .map(|&v| {
8482 let count =
8485 bounds.iter().filter(|&&b| if closed { !tol.lt(v, b) } else { tol.lt(b, v) });
8486 offset + count.count() as i64
8487 })
8488 .collect();
8489 Ok(Array::new(y.shape.clone(), Data::I64(out.into())))
8490}
8491
8492fn ordered_interval_index(
8496 x: &Array,
8497 y: &Array,
8498 offset: i64,
8499 closed: bool,
8500 span: Span,
8501) -> Result<Array> {
8502 let (xr, yr) = (x.to_row_major(), y.to_row_major());
8503 let (bounds, vals) = (&xr.data, &yr.data);
8504 let cmp = |i: usize, j: usize| -> Option<std::cmp::Ordering> {
8505 match (bounds, vals) {
8506 (Data::Char(p), Data::Char(q)) => Some(p[i].cmp(&q[j])),
8507 (Data::Symbol(p), Data::Symbol(q)) => Some(crate::symbol::cmp(p[i], q[j])),
8508 _ => None,
8509 }
8510 };
8511 let mut out = Vec::with_capacity(y.count());
8512 for j in 0..y.count() {
8513 let mut count = 0i64;
8514 for i in 0..x.count() {
8515 let ord = cmp(i, j).ok_or_else(|| {
8516 Error::domain(
8517 format!(
8518 "interval index compares {} bounds with {} values",
8519 x.dtype().name(),
8520 y.dtype().name()
8521 ),
8522 span,
8523 )
8524 })?;
8525 count += i64::from(if closed { ord.is_le() } else { ord.is_lt() });
8527 }
8528 out.push(offset + count);
8529 }
8530 Ok(Array::new(y.shape.clone(), Data::I64(out.into())))
8531}
8532
8533fn steps(y: &Array, span: Span) -> Result<Array> {
8536 let vals = y.to_f64_vec().ok_or_else(|| Error::domain("steps needs a number", span))?;
8537 let v = match vals.first() {
8538 Some(&v) if v.is_finite() => v,
8539 _ => return Err(Error::domain("steps needs a finite number", span)),
8540 };
8541 let n = (2.0 * v.abs()).floor();
8542 if n > 1e7 {
8543 return Err(Error::domain("steps would produce too many items", span));
8544 }
8545 let n = n as i64 + 1;
8546 let step = if v < 0.0 { -1.0 } else { 1.0 };
8547 let start = -v;
8548 if v.fract() == 0.0 {
8549 let start = start as i64;
8550 let step = step as i64;
8551 return Ok(Array::from_i64((0..n).map(|k| start + k * step).collect()));
8552 }
8553 Ok(Array::from_f64((0..n).map(|k| start + k as f64 * step).collect()))
8554}
8555
8556fn index_of_last(x: &Array, y: &Array, origin: i64, tol: Tol) -> Array {
8558 let cell_rank = x.rank().saturating_sub(1).min(y.rank());
8559 let frame_rank = y.rank() - cell_rank;
8560 let frame: Vec<usize> = y.shape[..frame_rank].to_vec();
8561 let nf: usize = frame.iter().product();
8562 let items = x.items();
8563 let mut out = Vec::with_capacity(nf);
8564 for i in 0..nf {
8565 let cell = y.cell_at(frame_rank, i);
8566 let at = (0..items)
8567 .rev()
8568 .find(|&j| arrays_match(&cell, &item_or_self(x, j), tol))
8569 .unwrap_or(items);
8570 out.push(origin + at as i64);
8571 }
8572 Array::new(frame, Data::I64(out.into()))
8573}
8574
8575fn roll(
8582 y: &Array,
8583 origin: i64,
8584 fixed: bool,
8585 float_at_zero: bool,
8586 span: Span,
8587) -> Result<Array> {
8588 let bounds = y
8589 .to_i64_vec()
8590 .ok_or_else(|| Error::domain("roll needs whole numbers", span))?;
8591 if bounds.iter().any(|&b| b < 0) {
8592 return Err(Error::domain("roll needs non-negative numbers", span));
8593 }
8594 if !float_at_zero && bounds.contains(&0) {
8595 return Err(Error::domain("? 0 has no value: the range is empty", span));
8596 }
8597 let any_zero = bounds.contains(&0);
8599 crate::rng::with(fixed, |g| {
8600 if any_zero {
8601 let out: Vec<f64> = bounds
8602 .iter()
8603 .map(|&b| {
8604 if b == 0 {
8605 g.unit()
8606 } else {
8607 (origin + g.below(b as u64) as i64) as f64
8608 }
8609 })
8610 .collect();
8611 return Ok(Array::new(y.shape.clone(), Data::F64(out.into())));
8612 }
8613 let out: Vec<i64> =
8614 bounds.iter().map(|&b| origin + g.below(b as u64) as i64).collect();
8615 Ok(Array::new(y.shape.clone(), Data::I64(out.into())))
8616 })
8617}
8618
8619fn deal(x: &Array, y: &Array, origin: i64, fixed: bool, span: Span) -> Result<Array> {
8621 let want = one_whole(x, "the count dealt", span)?;
8622 let from = one_whole(y, "the range dealt from", span)?;
8623 if want < 0 || from < 0 {
8624 return Err(Error::domain("deal needs non-negative numbers", span));
8625 }
8626 if want > from {
8627 return Err(Error::domain(
8628 format!("cannot deal {want} distinct value(s) from {from}"),
8629 span,
8630 ));
8631 }
8632 if want == 0 {
8633 return Ok(Array::from_i64(Vec::new()));
8634 }
8635 let drawn = crate::rng::with(fixed, |g| g.deal(want as usize, from as u64));
8636 Ok(Array::from_i64(drawn.into_iter().map(|v| v + origin).collect()))
8637}
8638
8639fn one_whole(a: &Array, what: &str, span: Span) -> Result<i64> {
8641 let v = a
8642 .to_i64_vec()
8643 .ok_or_else(|| Error::domain(format!("{what} must be a whole number"), span))?;
8644 match v[..] {
8645 [n] => Ok(n),
8646 _ => Err(Error::new(
8647 ErrorKind::Rank,
8648 format!("{what} must be one number"),
8649 Some(span),
8650 )),
8651 }
8652}
8653
8654fn nth_prime(n: i64, span: Span) -> Result<i64> {
8658 if n < 0 {
8659 return Err(Error::domain("the prime index must not be negative", span));
8660 }
8661 const LIMIT: i64 = 5_000_000;
8662 if n >= LIMIT {
8663 return Err(Error::domain(
8664 format!("prime index {n} is beyond the {LIMIT}th prime"),
8665 span,
8666 ));
8667 }
8668 let k = (n + 1) as f64;
8671 let bound = if n < 6 { 15.0 } else { k * (k.ln() + k.ln().ln()) };
8672 let bound = bound.ceil() as usize + 1;
8673 let mut sieve = vec![true; bound + 1];
8674 sieve[0] = false;
8675 if bound >= 1 {
8676 sieve[1] = false;
8677 }
8678 let mut p = 2usize;
8679 while p * p <= bound {
8680 if sieve[p] {
8681 let mut q = p * p;
8682 while q <= bound {
8683 sieve[q] = false;
8684 q += p;
8685 }
8686 }
8687 p += 1;
8688 }
8689 let mut seen = 0i64;
8690 for (v, &is_p) in sieve.iter().enumerate() {
8691 if is_p {
8692 if seen == n {
8693 return Ok(v as i64);
8694 }
8695 seen += 1;
8696 }
8697 }
8698 Err(Error::internal("the prime sieve was too small"))
8699}
8700
8701fn prime_factors(n: i64, span: Span) -> Result<Vec<i64>> {
8703 if n < 1 {
8704 return Err(Error::domain("prime factors need a positive integer", span));
8705 }
8706 let mut out = Vec::new();
8707 let mut m = n;
8708 let mut d = 2i64;
8709 while d.saturating_mul(d) <= m {
8710 while m % d == 0 {
8711 out.push(d);
8712 m /= d;
8713 }
8714 d += if d == 2 { 1 } else { 2 };
8715 }
8716 if m > 1 {
8717 out.push(m);
8718 }
8719 Ok(out)
8720}
8721
8722fn lstsq(a: &[f64], m: usize, n: usize, b: &[f64], k: usize) -> Option<Vec<f64>> {
8730 let mut r = a.to_vec();
8732 let mut c = b.to_vec();
8733 let at = |i: usize, j: usize, w: usize| i * w + j;
8734 let scale = a.iter().fold(0.0f64, |acc, v| acc.max(v.abs()));
8735 if scale == 0.0 {
8736 return None;
8737 }
8738 for j in 0..n {
8739 let norm = (j..m).map(|i| r[at(i, j, n)] * r[at(i, j, n)]).sum::<f64>().sqrt();
8741 if norm <= 1e-13 * scale {
8742 return None;
8743 }
8744 let alpha = if r[at(j, j, n)] > 0.0 { -norm } else { norm };
8745 let mut v = vec![0.0f64; m];
8746 for i in j..m {
8747 v[i] = r[at(i, j, n)];
8748 }
8749 v[j] -= alpha;
8750 let vnorm2: f64 = (j..m).map(|i| v[i] * v[i]).sum();
8751 if vnorm2 > 0.0 {
8752 for col in j..n {
8753 let dot: f64 = (j..m).map(|i| v[i] * r[at(i, col, n)]).sum();
8754 let f = 2.0 * dot / vnorm2;
8755 for i in j..m {
8756 r[at(i, col, n)] -= f * v[i];
8757 }
8758 }
8759 for col in 0..k {
8760 let dot: f64 = (j..m).map(|i| v[i] * c[at(i, col, k)]).sum();
8761 let f = 2.0 * dot / vnorm2;
8762 for i in j..m {
8763 c[at(i, col, k)] -= f * v[i];
8764 }
8765 }
8766 }
8767 }
8768 let mut x = vec![0.0f64; n * k];
8770 for col in 0..k {
8771 for i in (0..n).rev() {
8772 let mut acc = c[at(i, col, k)];
8773 for j in i + 1..n {
8774 acc -= r[at(i, j, n)] * x[at(j, col, k)];
8775 }
8776 let d = r[at(i, i, n)];
8777 if d.abs() <= 1e-13 * scale {
8778 return None;
8779 }
8780 x[at(i, col, k)] = acc / d;
8781 }
8782 }
8783 Some(x)
8784}
8785
8786fn as_matrix(a: &Array, span: Span) -> Result<(Vec<f64>, usize, usize)> {
8789 let v = a
8790 .to_f64_vec()
8791 .ok_or_else(|| Error::domain("matrix division needs numeric data", span))?;
8792 match a.rank() {
8793 0 => Ok((v, 1, 1)),
8794 1 => {
8795 let m = a.shape[0];
8796 Ok((v, m, 1))
8797 }
8798 2 => Ok((v, a.shape[0], a.shape[1])),
8799 _ => Err(Error::new(
8800 ErrorKind::Rank,
8801 "matrix division needs an argument of rank 2 or less",
8802 Some(span),
8803 )),
8804 }
8805}
8806
8807fn matrix_inverse(y: &Array, span: Span) -> Result<Array> {
8811 let (a, m, n) = as_matrix(y, span)?;
8812 if m < n {
8813 return Err(Error::new(
8814 ErrorKind::Length,
8815 format!("cannot invert a {m} by {n} matrix: it has more columns than rows"),
8816 Some(span),
8817 ));
8818 }
8819 let mut eye = vec![0.0f64; m * m];
8820 for i in 0..m {
8821 eye[i * m + i] = 1.0;
8822 }
8823 let x = lstsq(&a, m, n, &eye, m)
8824 .ok_or_else(|| Error::domain("the matrix is singular", span))?;
8825 let shape = if y.rank() == 2 { vec![n, m] } else { y.shape.clone() };
8828 Ok(Array::new(shape, Data::F64(x.into())))
8829}
8830
8831fn matrix_divide(x: &Array, y: &Array, span: Span) -> Result<Array> {
8833 let (a, m, n) = as_matrix(y, span)?;
8834 let (b, bm, k) = as_matrix(x, span)?;
8835 if bm != m {
8836 return Err(Error::new(
8837 ErrorKind::Length,
8838 format!("the system has {m} rows but the right-hand side has {bm}"),
8839 Some(span),
8840 ));
8841 }
8842 if m < n {
8843 return Err(Error::new(
8844 ErrorKind::Length,
8845 format!("the {m} by {n} system is underdetermined"),
8846 Some(span),
8847 ));
8848 }
8849 let sol = lstsq(&a, m, n, &b, k)
8850 .ok_or_else(|| Error::domain("the system is singular", span))?;
8851 let shape = if x.rank() == 2 { vec![n, k] } else { vec![n] };
8854 Ok(Array::new(shape, Data::F64(sol.into())))
8855}
8856
8857fn squad(x: &Array, y: &Array, origin: i64, span: Span) -> Result<Array> {
8861 if x.rank() > 1 {
8862 return Err(Error::new(
8863 ErrorKind::Rank,
8864 "the index of ⌷ must be a scalar or a vector",
8865 Some(span),
8866 ));
8867 }
8868 let items: Vec<Array> = if x.rank() == 0 { vec![x.clone()] } else { x.cells(1) };
8871 if items.len() != y.rank() {
8872 return Err(Error::new(
8873 ErrorKind::Rank,
8874 format!("{} index(es) for an argument of rank {}", items.len(), y.rank()),
8875 Some(span),
8876 ));
8877 }
8878 let mut specs = Vec::with_capacity(items.len());
8879 let mut shape = Vec::new();
8880 for (k, item) in items.iter().enumerate() {
8881 let spec = match item.as_boxes() {
8882 Some(bs) if item.rank() == 0 => bs[0].clone(),
8883 _ => item.clone(),
8884 };
8885 let idx = spec
8886 .to_i64_vec()
8887 .ok_or_else(|| Error::domain("index must be an integer", span))?;
8888 for &i in &idx {
8889 let j = i - origin;
8890 if j < 0 || j as usize >= y.shape[k] {
8891 return Err(Error::domain(
8892 format!("index {i} is out of range on axis {k}"),
8893 span,
8894 ));
8895 }
8896 }
8897 shape.extend_from_slice(&spec.shape);
8898 specs.push((spec.shape.clone(), idx));
8899 }
8900 let y = y.to_row_major();
8901 let st = strides(&y.shape);
8902 let total: usize = shape.iter().product();
8903 let mut data = Data::empty(y.dtype());
8904 let mut coord = vec![0usize; shape.len()];
8905 for _ in 0..total {
8906 let mut at = 0usize;
8907 let mut used = 0usize;
8908 for (k, (sshape, idx)) in specs.iter().enumerate() {
8909 let sst = strides(sshape);
8910 let pick: usize = (0..sshape.len()).map(|a| coord[used + a] * sst[a]).sum();
8911 used += sshape.len();
8912 at += (idx[pick] - origin) as usize * st[k];
8913 }
8914 push_elem(&mut data, y.row_major_data(), at);
8915 odometer(&mut coord, &shape);
8916 }
8917 Ok(Array::new(shape, data))
8918}
8919
8920fn select_axis(
8927 x: &Array,
8928 y: &Array,
8929 axis: usize,
8930 rank: usize,
8931 origin: i64,
8932 span: Span,
8933) -> Result<Array> {
8934 if rank != 0 && y.rank() != rank {
8935 return Err(Error::new(
8936 ErrorKind::Rank,
8937 format!("{rank} index slot(s) for an argument of rank {}", y.rank()),
8938 Some(span),
8939 ));
8940 }
8941 if axis >= y.rank() {
8942 return Err(Error::new(
8943 ErrorKind::Rank,
8944 format!("axis {axis} does not exist on an argument of rank {}", y.rank()),
8945 Some(span),
8946 ));
8947 }
8948 let idx = x
8949 .to_i64_vec()
8950 .ok_or_else(|| Error::domain("index must be an integer", span))?;
8951 let len = y.shape[axis];
8952 let mut picks = Vec::with_capacity(idx.len());
8953 for &i in &idx {
8954 let j = i - origin;
8955 if j < 0 || j as usize >= len {
8956 return Err(Error::domain(
8957 format!("index {i} is out of range: axis {axis} has {len} items"),
8958 span,
8959 ));
8960 }
8961 picks.push(j as usize);
8962 }
8963 let mut shape = Vec::with_capacity(y.rank() + x.rank());
8964 shape.extend_from_slice(&y.shape[..axis]);
8965 shape.extend_from_slice(&x.shape);
8966 shape.extend_from_slice(&y.shape[axis + 1..]);
8967 let outer: usize = y.shape[..axis].iter().product();
8968 let inner: usize = y.shape[axis + 1..].iter().product();
8969 let mut data = Data::empty(y.dtype());
8970 for o in 0..outer {
8971 for &p in &picks {
8972 let base = (o * len + p) * inner;
8973 for e in 0..inner {
8974 push_elem(&mut data, &y.data, base + e);
8975 }
8976 }
8977 }
8978 Ok(Array::new(shape, data))
8979}
8980
8981fn amend(m: &Array, x: &Array, y: &Array, span: Span) -> Result<Array> {
8985 if y.rank() == 0 {
8986 return Err(Error::new(ErrorKind::Rank, "cannot amend a scalar", Some(span)));
8987 }
8988 if let Some(spec) = m.as_boxes().and_then(<[Array]>::first) {
8990 let spec = index_spec(spec, y, span)?;
8991 return amend_spec(&spec, x, y, span);
8992 }
8993 let idx = m
8994 .to_i64_vec()
8995 .ok_or_else(|| Error::domain("amend indices must be integers", span))?;
8996 let items = y.items() as i64;
8997 let mut at = Vec::with_capacity(idx.len());
8998 for &i in &idx {
8999 let k = if i < 0 { i + items } else { i };
9000 if k < 0 || k >= items {
9001 return Err(Error::domain(
9002 format!("index {i} is out of range: the argument has {items} items"),
9003 span,
9004 ));
9005 }
9006 at.push(k as usize);
9007 }
9008 let cell = y.item_size();
9009 let per_index = if x.count() == cell {
9010 false
9011 } else if x.count() == cell * at.len() {
9012 true
9013 } else {
9014 return Err(Error::new(
9015 ErrorKind::Length,
9016 format!(
9017 "cannot amend {} item(s) of {} element(s) each with {} element(s)",
9018 at.len(),
9019 cell,
9020 x.count()
9021 ),
9022 Some(span),
9023 ));
9024 };
9025 let Some(t) = DType::promote(x.dtype(), y.dtype()) else {
9028 return Err(Error::new(
9029 ErrorKind::Type,
9030 "the replacement and the argument hold different kinds of value",
9031 Some(span),
9032 ));
9033 };
9034 let (Some(src), Some(base)) = (x.data.cast(t), y.data.cast(t)) else {
9035 return Err(Error::new(
9036 ErrorKind::Type,
9037 "the replacement and the argument hold different kinds of value",
9038 Some(span),
9039 ));
9040 };
9041 let mut data = Data::empty(t);
9043 let mut plan: Vec<Option<usize>> = vec![None; y.items()];
9044 for (n, &k) in at.iter().enumerate() {
9045 plan[k] = Some(if per_index { n } else { 0 });
9046 }
9047 for (i, slot) in plan.iter().enumerate() {
9048 match slot {
9049 Some(n) => {
9050 for e in 0..cell {
9051 push_elem(&mut data, &src, n * cell + e);
9052 }
9053 }
9054 None => {
9055 for e in 0..cell {
9056 push_elem(&mut data, &base, i * cell + e);
9057 }
9058 }
9059 }
9060 }
9061 Ok(Array::new(y.shape.clone(), data))
9062}
9063
9064fn fetch(x: &Array, y: &Array, span: Span) -> Result<Array> {
9069 let steps: Vec<Array> = match x.as_boxes() {
9070 Some(bs) => bs.to_vec(),
9071 None => vec![x.clone()],
9072 };
9073 let mut cur = y.clone();
9074 for step in steps {
9075 let idx = if step.count() == 0 {
9078 Vec::new()
9079 } else {
9080 step.to_i64_vec()
9081 .ok_or_else(|| Error::domain("a fetch path holds integers", span))?
9082 };
9083 let base =
9085 if cur.rank() == 0 { Array::new(vec![1], cur.data.clone()) } else { cur.clone() };
9086 if idx.len() > base.rank() {
9087 return Err(Error::new(
9088 ErrorKind::Length,
9089 format!(
9090 "a path step of {} index(es) into a value of rank {}",
9091 idx.len(),
9092 cur.rank()
9093 ),
9094 Some(span),
9095 ));
9096 }
9097 let at = cell_index(&base, &idx, span)?;
9098 cur = open_cell(&base.cell_at(idx.len(), at));
9099 }
9100 Ok(cur)
9101}
9102
9103fn cell_index(y: &Array, idx: &[i64], span: Span) -> Result<usize> {
9105 let mut at = 0usize;
9106 for (k, &i) in idx.iter().enumerate() {
9107 let len = y.shape[k] as i64;
9108 let j = if i < 0 { i + len } else { i };
9109 if j < 0 || j >= len {
9110 return Err(Error::domain(
9111 format!("index {i} is out of range: axis {k} has {len} items"),
9112 span,
9113 ));
9114 }
9115 at = at * y.shape[k] + j as usize;
9116 }
9117 Ok(at)
9118}
9119
9120fn partition_enclose(x: &Array, y: &Array, span: Span) -> Result<Array> {
9129 if y.rank() > 1 {
9132 let last = y.shape[y.rank() - 1];
9133 let rows = y.count() / last.max(1);
9134 let mut cells: Vec<Array> = Vec::new();
9135 let mut width = None;
9136 for r in 0..rows {
9137 let row = Array::new(vec![last], y.data.slice(r * last, (r + 1) * last));
9138 let parts = partition_enclose(x, &row, span)?;
9139 let n = parts.count();
9140 if *width.get_or_insert(n) != n {
9141 return Err(Error::internal("partitions of unequal count"));
9142 }
9143 match parts.data {
9144 Data::Box(v) => cells.extend(v.as_slice().iter().cloned()),
9145 _ => return Err(Error::internal("a partition is boxed")),
9146 }
9147 }
9148 let mut shape = y.shape[..y.rank() - 1].to_vec();
9149 shape.push(width.unwrap_or(0));
9150 return Ok(Array::new(shape, Data::Box(cells.into())));
9151 }
9152 if y.rank() == 0 {
9153 return Err(Error::new(
9154 ErrorKind::Rank,
9155 "partitioned enclose needs an array to partition",
9156 Some(span),
9157 ));
9158 }
9159 let flags = x
9160 .to_i64_vec()
9161 .ok_or_else(|| Error::domain("partition flags must be integers", span))?;
9162 if flags.iter().any(|&f| f < 0) {
9163 return Err(Error::domain("partition flags must not be negative", span));
9164 }
9165 if flags.len() != y.shape[0] {
9166 return Err(Error::new(
9167 ErrorKind::Length,
9168 format!("{} flag(s) for {} item(s)", flags.len(), y.shape[0]),
9169 Some(span),
9170 ));
9171 }
9172 let mut parts: Vec<Array> = Vec::new();
9173 let mut cur: Option<Data> = None;
9174 let mut prev = 0i64;
9175 for (i, &f) in flags.iter().enumerate() {
9176 if f > prev {
9177 if let Some(d) = cur.take() {
9178 parts.push(Array::new(vec![d.len()], d));
9179 }
9180 cur = Some(Data::empty(y.dtype()));
9181 }
9182 prev = f;
9183 if f == 0 {
9184 continue;
9185 }
9186 if let Some(d) = cur.as_mut() {
9187 push_elem(d, &y.data, i);
9188 }
9189 }
9190 if let Some(d) = cur.take() {
9191 parts.push(Array::new(vec![d.len()], d));
9192 }
9193 Ok(Array::new(vec![parts.len()], Data::Box(parts.into())))
9194}
9195
9196fn key(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
9199 let keys = if x.rank() == 0 { Array::new(vec![1], x.data.clone()) } else { x.clone() };
9200 let n = keys.items();
9201 if n != y.items() && !(y.rank() == 0 && n == 1) {
9202 return Err(Error::new(
9203 ErrorKind::Length,
9204 format!("{n} key(s) for {} item(s)", y.items()),
9205 Some(span),
9206 ));
9207 }
9208 let groups = group_positions(&keys, ctx.cfg.tol);
9209 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
9210 let mut cells = Vec::with_capacity(groups.len());
9211 for (_, at) in &groups {
9212 cells.push(u.monad(&select_items(&items, at), ctx, span)?);
9213 }
9214 assemble(&[groups.len()], cells, span)
9215}
9216
9217fn oblique(u: &Verb, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
9220 if y.rank() < 2 {
9221 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
9222 let n = items.items();
9223 let mut cells = Vec::with_capacity(n);
9224 for i in 0..n {
9225 cells.push(u.monad(&select_items(&items, &[i]), ctx, span)?);
9226 }
9227 return assemble(&[n], cells, span);
9228 }
9229 if y.rank() > 2 {
9230 return Err(Error::not_yet("oblique (u/.) on a rank-3 or higher argument", span));
9231 }
9232 let (rows, cols) = (y.shape[0], y.shape[1]);
9233 let mut cells = Vec::with_capacity(rows + cols - 1);
9234 for d in 0..rows + cols - 1 {
9235 let mut data = Data::empty(y.dtype());
9236 let mut len = 0usize;
9237 for i in 0..rows {
9238 if d >= i && d - i < cols {
9239 push_elem(&mut data, &y.data, i * cols + (d - i));
9240 len += 1;
9241 }
9242 }
9243 cells.push(u.monad(&Array::new(vec![len], data), ctx, span)?);
9244 }
9245 assemble(&[rows + cols - 1], cells, span)
9246}
9247
9248fn cut_ranges(frets: &[bool], mode: i64) -> Vec<(usize, usize)> {
9256 let n = frets.len();
9257 let mut out = Vec::new();
9258 if mode.abs() == 1 {
9259 let mut start: Option<usize> = None;
9260 for (i, &fret) in frets.iter().enumerate() {
9261 if fret {
9262 if let Some(s) = start {
9263 out.push((s, i));
9264 }
9265 start = Some(i);
9266 }
9267 }
9268 if let Some(s) = start {
9269 out.push((s, n));
9270 }
9271 if mode < 0 {
9272 return out.into_iter().map(|(s, e)| (s + 1, e)).collect();
9273 }
9274 } else {
9275 let mut start = 0usize;
9276 for (i, &fret) in frets.iter().enumerate() {
9277 if fret {
9278 out.push((start, i + 1));
9279 start = i + 1;
9280 }
9281 }
9282 if mode < 0 {
9283 return out.into_iter().map(|(s, e)| (s, e - 1)).collect();
9284 }
9285 }
9286 out
9287}
9288
9289fn cut(
9291 u: &Verb,
9292 x: Option<&Array>,
9293 y: &Array,
9294 mode: i64,
9295 ctx: &mut Ctx<'_>,
9296 span: Span,
9297) -> Result<Array> {
9298 if mode == 0 {
9299 let Some(x) = x else {
9300 return u.monad(&reverse_all_axes(y), ctx, span);
9301 };
9302 let (origin, size) = rectangle(x, span)?;
9303 let origin = origin.unwrap_or_else(|| vec![0; size.len()]);
9304 return u.monad(&subarray(y, &origin, &size, span)?, ctx, span);
9305 }
9306 if mode.abs() == 3 {
9307 let Some(x) = x else {
9308 return Err(Error::not_yet("monadic tessellation (u;.3 y)", span));
9309 };
9310 return tessellate(u, x, y, mode < 0, ctx, span);
9311 }
9312 if !matches!(mode, 1 | -1 | 2 | -2) {
9313 return Err(Error::not_yet(format!("cut (u;.{mode})"), span));
9314 }
9315 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
9316 let n = items.items();
9317 let tol = ctx.cfg.tol;
9318 let frets: Vec<bool> = match x {
9319 Some(x) => {
9320 let flags = x
9321 .to_i64_vec()
9322 .ok_or_else(|| Error::domain("cut frets must be integers", span))?;
9323 if let Some(&bad) = flags.iter().find(|&&f| f != 0 && f != 1) {
9326 return Err(Error::domain(format!("{bad} is not a fret: a fret is 0 or 1"), span));
9327 }
9328 if x.rank() == 0 {
9331 vec![flags[0] != 0; n]
9332 } else {
9333 if flags.len() != n {
9334 return Err(Error::new(
9335 ErrorKind::Length,
9336 format!("{} fret(s) for {n} item(s)", flags.len()),
9337 Some(span),
9338 ));
9339 }
9340 flags.iter().map(|&f| f != 0).collect()
9341 }
9342 }
9343 None => {
9344 if n == 0 {
9346 Vec::new()
9347 } else {
9348 let at = if mode.abs() == 1 { 0 } else { n - 1 };
9349 let mark = items.item(at);
9350 (0..n).map(|i| arrays_match(&items.item(i), &mark, tol)).collect()
9351 }
9352 }
9353 };
9354 let ranges = cut_ranges(&frets, mode);
9355 let mut cells = Vec::with_capacity(ranges.len());
9356 for (s, e) in &ranges {
9357 cells.push(u.monad(§ion(&items, *s, *e), ctx, span)?);
9358 }
9359 assemble(&[ranges.len()], cells, span)
9360}
9361
9362fn rectangle(x: &Array, span: Span) -> Result<(Option<Vec<i64>>, Vec<i64>)> {
9365 let values = x
9366 .to_i64_vec()
9367 .ok_or_else(|| Error::domain("a cut rectangle is whole numbers", span))?;
9368 match x.rank() {
9369 0 | 1 => Ok((None, values)),
9370 2 if x.shape[0] == 2 => {
9371 let n = x.shape[1];
9372 Ok((Some(values[..n].to_vec()), values[n..].to_vec()))
9373 }
9374 _ => Err(Error::new(
9375 ErrorKind::Rank,
9376 "a cut rectangle is a vector of sizes, or two rows of origins and sizes",
9377 Some(span),
9378 )),
9379 }
9380}
9381
9382fn subarray(y: &Array, origin: &[i64], size: &[i64], span: Span) -> Result<Array> {
9386 if origin.len() > y.rank() {
9387 return Err(Error::new(
9388 ErrorKind::Rank,
9389 format!("a cut of {} axis/axes into a rank-{} value", origin.len(), y.rank()),
9390 Some(span),
9391 ));
9392 }
9393 let r = y.rank();
9394 let st = strides(&y.shape);
9395 let mut shape = y.shape.clone();
9396 let mut start = vec![0i64; r];
9397 let mut step = vec![1i64; r];
9398 for k in 0..origin.len() {
9399 let len = size[k].unsigned_abs() as usize;
9400 let from = if origin[k] < 0 { origin[k] + y.shape[k] as i64 } else { origin[k] };
9401 if from < 0 || from + len as i64 > y.shape[k] as i64 {
9402 return Err(Error::domain(
9403 format!("a cut of {len} from {from} leaves axis {k} of {}", y.shape[k]),
9404 span,
9405 ));
9406 }
9407 shape[k] = len;
9408 if size[k] < 0 {
9409 start[k] = from + len as i64 - 1;
9410 step[k] = -1;
9411 } else {
9412 start[k] = from;
9413 }
9414 }
9415 Ok(gather(y, &shape, &start, &step, &st))
9416}
9417
9418fn gather(y: &Array, shape: &[usize], start: &[i64], step: &[i64], st: &[usize]) -> Array {
9420 let n: usize = shape.iter().product();
9421 let mut data = Data::empty(y.dtype());
9422 let mut coord = vec![0usize; shape.len()];
9423 for _ in 0..n {
9424 let idx: usize = (0..shape.len())
9425 .map(|k| (start[k] + step[k] * coord[k] as i64) as usize * st[k])
9426 .sum();
9427 push_elem(&mut data, &y.data, idx);
9428 odometer(&mut coord, shape);
9429 }
9430 Array::new(shape.to_vec(), data)
9431}
9432
9433fn tessellate(
9437 u: &Verb,
9438 x: &Array,
9439 y: &Array,
9440 complete: bool,
9441 ctx: &mut Ctx<'_>,
9442 span: Span,
9443) -> Result<Array> {
9444 let (movement, size) = rectangle(x, span)?;
9446 if size.iter().any(|&s| s < 0) && movement.is_none() {
9451 return Err(Error::not_yet(
9452 "a negative block size without a movement row (x u;.3 y)",
9453 span,
9454 ));
9455 }
9456 let movement = movement.unwrap_or_else(|| vec![1; size.len()]);
9457 if size.len() > y.rank() {
9458 return Err(Error::new(
9459 ErrorKind::Rank,
9460 format!("a tessellation of {} axis/axes into a rank-{} value", size.len(), y.rank()),
9461 Some(span),
9462 ));
9463 }
9464 let mut frame = Vec::with_capacity(size.len());
9465 for k in 0..size.len() {
9466 let (len, step, block) = (y.shape[k] as i64, movement[k], size[k].abs());
9467 if step <= 0 {
9468 return Err(Error::domain("a tessellation moves by a positive step", span));
9469 }
9470 let count = if complete {
9471 if len < block { 0 } else { (len - block) / step + 1 }
9472 } else {
9473 (len + step - 1) / step
9474 };
9475 frame.push(count as usize);
9476 }
9477 let total: usize = frame.iter().product();
9478 let mut cells = Vec::with_capacity(total);
9479 let mut coord = vec![0usize; frame.len()];
9480 for _ in 0..total {
9481 let origin: Vec<i64> = (0..frame.len()).map(|k| coord[k] as i64 * movement[k]).collect();
9482 let block: Vec<i64> = (0..frame.len())
9485 .map(|k| {
9486 let len = size[k].abs().min(y.shape[k] as i64 - origin[k]);
9487 if size[k] < 0 { -len } else { len }
9488 })
9489 .collect();
9490 cells.push(u.monad(&subarray(y, &origin, &block, span)?, ctx, span)?);
9491 odometer(&mut coord, &frame);
9492 }
9493 assemble(&frame, cells, span)
9494}
9495
9496fn reverse_all_axes(y: &Array) -> Array {
9498 if y.rank() == 0 {
9499 return y.clone();
9500 }
9501 let st = strides(&y.shape);
9502 let n = y.count();
9503 let r = y.rank();
9504 let mut data = Data::empty(y.dtype());
9505 let mut coord = vec![0usize; r];
9506 for _ in 0..n {
9507 let idx: usize = (0..r).map(|k| (y.shape[k] - 1 - coord[k]) * st[k]).sum();
9508 push_elem(&mut data, &y.data, idx);
9509 odometer(&mut coord, &y.shape);
9510 }
9511 Array::new(y.shape.clone(), data)
9512}
9513
9514fn axis_to_front(y: &Array, k: usize) -> Array {
9518 if k == 0 || y.rank() < 2 {
9519 return y.clone();
9520 }
9521 let r = y.rank();
9522 let src: Vec<usize> = std::iter::once(k).chain((0..r).filter(|&a| a != k)).collect();
9523 permute_axes(y, &src)
9524}
9525
9526fn front_to_axis(y: &Array, k: usize) -> Array {
9528 if k == 0 || y.rank() < 2 {
9529 return y.clone();
9530 }
9531 let r = y.rank();
9532 let mut src = Vec::with_capacity(r);
9535 for a in 0..r {
9536 src.push(match a.cmp(&k) {
9537 std::cmp::Ordering::Less => a + 1,
9538 std::cmp::Ordering::Equal => 0,
9539 std::cmp::Ordering::Greater => a,
9540 });
9541 }
9542 permute_axes(y, &src)
9543}
9544
9545fn transpose_to(y: &Array, dest: &[usize], span: Span) -> Result<Array> {
9550 let rank_out = dest.iter().copied().max().map_or(0, |m| m + 1);
9551 let mut out_shape = vec![usize::MAX; rank_out];
9552 for (a, &d) in dest.iter().enumerate() {
9553 out_shape[d] = out_shape[d].min(y.shape[a]);
9554 }
9555 if out_shape.contains(&usize::MAX) {
9556 return Err(Error::new(
9557 ErrorKind::Domain,
9558 "a transpose must name every axis of the result",
9559 Some(span),
9560 ));
9561 }
9562 let y = y.to_row_major();
9563 let st = strides(&y.shape);
9564 let n: usize = out_shape.iter().product();
9565 let mut data = Data::empty(y.dtype());
9566 let mut coord = vec![0usize; rank_out];
9567 for _ in 0..n {
9568 let idx: usize = dest.iter().enumerate().map(|(a, &d)| coord[d] * st[a]).sum();
9569 push_elem(&mut data, &y.data, idx);
9570 odometer(&mut coord, &out_shape);
9571 }
9572 Ok(Array::new(out_shape, data))
9573}
9574
9575fn transpose_apl(x: &Array, y: &Array, io: i64, span: Span) -> Result<Array> {
9578 let axes = x
9579 .to_i64_vec()
9580 .ok_or_else(|| Error::domain("a transpose is given whole numbers", span))?;
9581 if axes.len() != y.rank() {
9582 return Err(Error::new(
9583 ErrorKind::Length,
9584 format!("{} axes for a rank-{} value", axes.len(), y.rank()),
9585 Some(span),
9586 ));
9587 }
9588 let mut dest = Vec::with_capacity(axes.len());
9589 for a in axes {
9590 let d = a - io;
9591 if d < 0 || d as usize >= y.rank() {
9592 return Err(Error::new(
9593 ErrorKind::Domain,
9594 format!("axis {a} is outside a rank-{} value", y.rank()),
9595 Some(span),
9596 ));
9597 }
9598 dest.push(d as usize);
9599 }
9600 transpose_to(y, &dest, span)
9601}
9602
9603fn transpose_j(x: &Array, y: &Array, span: Span) -> Result<Array> {
9607 let groups: Vec<Vec<i64>> = match x.as_boxes() {
9608 Some(bs) => bs
9609 .iter()
9610 .map(|b| {
9611 b.to_i64_vec().ok_or_else(|| {
9612 Error::domain("a transpose is given whole numbers", span)
9613 })
9614 })
9615 .collect::<Result<Vec<_>>>()?,
9616 None => x
9617 .to_i64_vec()
9618 .ok_or_else(|| Error::domain("a transpose is given whole numbers", span))?
9619 .into_iter()
9620 .map(|a| vec![a])
9621 .collect(),
9622 };
9623 let r = y.rank();
9624 let mut group_of = vec![None; r];
9627 for (g, axes) in groups.iter().enumerate() {
9628 for &a in axes {
9629 let k = if a < 0 { a + r as i64 } else { a };
9630 if k < 0 || k as usize >= r {
9631 return Err(Error::new(
9632 ErrorKind::Domain,
9633 format!("axis {a} is outside a rank-{r} value"),
9634 Some(span),
9635 ));
9636 }
9637 if group_of[k as usize].is_some() {
9638 return Err(Error::new(
9639 ErrorKind::Domain,
9640 format!("axis {a} is named twice in a transpose"),
9641 Some(span),
9642 ));
9643 }
9644 group_of[k as usize] = Some(g);
9645 }
9646 }
9647 let leading = group_of.iter().filter(|g| g.is_none()).count();
9648 let mut dest = vec![0usize; r];
9649 let mut next = 0;
9650 for a in 0..r {
9651 match group_of[a] {
9652 None => {
9653 dest[a] = next;
9654 next += 1;
9655 }
9656 Some(g) => dest[a] = leading + g,
9657 }
9658 }
9659 transpose_to(y, &dest, span)
9660}
9661
9662fn permute_axes(y: &Array, src: &[usize]) -> Array {
9664 let st = strides(&y.shape);
9665 let out_shape: Vec<usize> = src.iter().map(|&a| y.shape[a]).collect();
9666 let n = y.count();
9667 let mut data = Data::empty(y.dtype());
9668 let mut coord = vec![0usize; src.len()];
9669 for _ in 0..n {
9670 let idx: usize = (0..src.len()).map(|a| coord[a] * st[src[a]]).sum();
9671 push_elem(&mut data, &y.data, idx);
9672 odometer(&mut coord, &out_shape);
9673 }
9674 Array::new(out_shape, data)
9675}
9676
9677struct Spec {
9681 width: usize,
9683 cells: Vec<Vec<usize>>,
9685 shape: Vec<usize>,
9688}
9689
9690fn axis_position(v: i64, len: usize, span: Span) -> Result<usize> {
9693 let p = if v < 0 { v + len as i64 } else { v };
9694 if p < 0 || p >= len as i64 {
9695 return Err(Error::domain(
9696 format!("index {v} is out of range: the axis has {len} element(s)"),
9697 span,
9698 ));
9699 }
9700 Ok(p as usize)
9701}
9702
9703fn index_spec(content: &Array, y: &Array, span: Span) -> Result<Spec> {
9713 let too_deep = |n: usize| {
9714 Error::new(
9715 ErrorKind::Rank,
9716 format!("an index specification of {n} axis/axes into a rank-{} value", y.rank()),
9717 Some(span),
9718 )
9719 };
9720 if let Some(items) = content.as_boxes() {
9721 if items.len() > y.rank() {
9722 return Err(too_deep(items.len()));
9723 }
9724 let mut per_axis: Vec<Vec<usize>> = Vec::with_capacity(items.len());
9725 let mut shape: Vec<usize> = Vec::new();
9726 for (k, c) in items.iter().enumerate() {
9727 let len = y.shape[k];
9728 if c.as_boxes().is_some() {
9729 let inner = open_cell(c);
9730 let excluded = inner.to_i64_vec().ok_or_else(|| {
9731 Error::domain("an index complement holds integers", span)
9732 })?;
9733 let mut dropped = vec![false; len];
9734 for v in excluded {
9735 dropped[axis_position(v, len, span)?] = true;
9736 }
9737 let kept: Vec<usize> = (0..len).filter(|i| !dropped[*i]).collect();
9738 shape.push(kept.len());
9739 per_axis.push(kept);
9740 } else {
9741 let idx = c
9742 .to_i64_vec()
9743 .ok_or_else(|| Error::domain("an index holds integers", span))?;
9744 let mut positions = Vec::with_capacity(idx.len());
9745 for v in idx {
9746 positions.push(axis_position(v, len, span)?);
9747 }
9748 shape.extend_from_slice(&c.shape);
9749 per_axis.push(positions);
9750 }
9751 }
9752 let mut cells: Vec<Vec<usize>> = vec![Vec::new()];
9754 for positions in &per_axis {
9755 let mut next = Vec::with_capacity(cells.len() * positions.len());
9756 for prefix in &cells {
9757 for &p in positions {
9758 let mut cell = prefix.clone();
9759 cell.push(p);
9760 next.push(cell);
9761 }
9762 }
9763 cells = next;
9764 }
9765 return Ok(Spec { width: per_axis.len(), cells, shape });
9766 }
9767 let idx = content
9768 .to_i64_vec()
9769 .ok_or_else(|| Error::domain("an index specification holds integers", span))?;
9770 let rank = content.rank();
9771 let width = if rank == 0 { 1 } else { content.shape[rank - 1] };
9772 if width > y.rank() {
9773 return Err(too_deep(width));
9774 }
9775 let shape: Vec<usize> = if rank == 0 { Vec::new() } else { content.shape[..rank - 1].to_vec() };
9776 let count: usize = shape.iter().product();
9777 let mut cells: Vec<Vec<usize>> = Vec::new();
9778 if width == 0 {
9779 cells.resize(count, Vec::new());
9780 } else {
9781 for chunk in idx.chunks(width) {
9782 let mut cell = Vec::with_capacity(width);
9783 for (k, &v) in chunk.iter().enumerate() {
9784 cell.push(axis_position(v, y.shape[k], span)?);
9785 }
9786 cells.push(cell);
9787 }
9788 }
9789 Ok(Spec { width, cells, shape })
9790}
9791
9792fn spec_offset(st: &[usize], cell: &[usize]) -> usize {
9794 cell.iter().enumerate().map(|(k, &p)| p * st[k]).sum()
9795}
9796
9797fn select_spec(spec: &Spec, y: &Array) -> Array {
9799 let st = strides(&y.shape);
9800 let size: usize = y.shape[spec.width..].iter().product();
9801 let mut data = Data::empty(y.dtype());
9802 for cell in &spec.cells {
9803 let base = spec_offset(&st, cell);
9804 for e in 0..size {
9805 push_elem(&mut data, &y.data, base + e);
9806 }
9807 }
9808 let mut shape = spec.shape.clone();
9809 shape.extend_from_slice(&y.shape[spec.width..]);
9810 Array::new(shape, data)
9811}
9812
9813fn amend_spec(spec: &Spec, x: &Array, y: &Array, span: Span) -> Result<Array> {
9816 let size: usize = y.shape[spec.width..].iter().product();
9817 let per_cell = if x.count() == size {
9818 false
9819 } else if x.count() == size * spec.cells.len() {
9820 true
9821 } else {
9822 return Err(Error::new(
9823 ErrorKind::Length,
9824 format!(
9825 "cannot amend {} cell(s) of {size} element(s) each with {} element(s)",
9826 spec.cells.len(),
9827 x.count()
9828 ),
9829 Some(span),
9830 ));
9831 };
9832 let mismatch = || {
9833 Error::new(
9834 ErrorKind::Type,
9835 "the replacement and the argument hold different kinds of value",
9836 Some(span),
9837 )
9838 };
9839 let t = DType::promote(x.dtype(), y.dtype()).ok_or_else(mismatch)?;
9840 let (Some(src), Some(base)) = (x.data.cast(t), y.data.cast(t)) else {
9841 return Err(mismatch());
9842 };
9843 let st = strides(&y.shape);
9844 let mut plan: Vec<Option<usize>> = vec![None; y.count()];
9845 for (n, cell) in spec.cells.iter().enumerate() {
9846 let at = spec_offset(&st, cell);
9847 for e in 0..size {
9848 plan[at + e] = Some(if per_cell { n * size + e } else { e });
9849 }
9850 }
9851 let mut data = Data::empty(t);
9852 for (i, slot) in plan.iter().enumerate() {
9853 match slot {
9854 Some(n) => push_elem(&mut data, &src, *n),
9855 None => push_elem(&mut data, &base, i),
9856 }
9857 }
9858 Ok(Array::new(y.shape.clone(), data))
9859}
9860
9861fn map_paths(y: &Array) -> Array {
9870 fn coord_of(shape: &[usize], mut i: usize) -> Array {
9871 let mut out = vec![0i64; shape.len()];
9872 for k in (0..shape.len()).rev() {
9873 out[k] = (i % shape[k]) as i64;
9874 i /= shape[k];
9875 }
9876 Array::from_i64(out)
9877 }
9878 fn go(y: &Array, prefix: &[Array]) -> Array {
9879 let Some(boxes) = y.as_boxes() else {
9880 if prefix.is_empty() {
9881 return Array::new(vec![0], Data::I64(Vec::new().into()));
9882 }
9883 return Array::new(vec![prefix.len()], Data::Box(prefix.to_vec().into()));
9884 };
9885 let cells: Vec<Array> = boxes
9886 .iter()
9887 .enumerate()
9888 .map(|(i, b)| {
9889 let mut path = prefix.to_vec();
9890 path.push(coord_of(&y.shape, i));
9891 go(b, &path)
9892 })
9893 .collect();
9894 Array::new(y.shape.clone(), Data::Box(cells.into()))
9895 }
9896 go(y, &[])
9897}
9898
9899fn shift_fill(x: &Array, y: &Array, fill: &Array, span: Span) -> Result<Array> {
9904 let counts = axis_counts(x, "shift", span)?;
9905 if y.rank() == 0 {
9906 return Ok(y.clone());
9907 }
9908 if counts.len() > y.rank() {
9909 return Err(Error::new(
9910 ErrorKind::Length,
9911 format!("shift has {} amounts for an argument of rank {}", counts.len(), y.rank()),
9912 Some(span),
9913 ));
9914 }
9915 if fill.count() != 1 {
9916 return Err(Error::new(ErrorKind::Length, "a fill is one atom", Some(span)));
9917 }
9918 let mismatch = || {
9919 Error::new(ErrorKind::Type, "the fill and the argument differ in kind", Some(span))
9920 };
9921 let t = DType::promote(y.dtype(), fill.dtype()).ok_or_else(mismatch)?;
9922 let (Some(base), Some(f)) = (y.data.cast(t), fill.data.cast(t)) else {
9923 return Err(mismatch());
9924 };
9925 let st = strides(&y.shape);
9926 let r = y.rank();
9927 let mut data = Data::empty(t);
9928 let mut coord = vec![0usize; r];
9929 for _ in 0..y.count() {
9930 let mut idx = 0usize;
9931 let mut vacated = false;
9932 for k in 0..r {
9933 let from = coord[k] as i64 + counts.get(k).copied().unwrap_or(0);
9934 if from < 0 || from >= y.shape[k] as i64 {
9935 vacated = true;
9936 break;
9937 }
9938 idx += from as usize * st[k];
9939 }
9940 if vacated {
9941 push_elem(&mut data, &f, 0);
9942 } else {
9943 push_elem(&mut data, &base, idx);
9944 }
9945 odometer(&mut coord, &y.shape);
9946 }
9947 Ok(Array::new(y.shape.clone(), data))
9948}
9949
9950fn memo_key(a: &Array, out: &mut Vec<u64>) -> bool {
9955 out.push(a.rank() as u64);
9956 out.extend(a.shape.iter().map(|&n| n as u64));
9957 out.push(a.dtype() as u64);
9958 match &a.data {
9959 Data::Ext(_) | Data::Rat(_) => false,
9960 Data::Box(items) => items.iter().all(|item| memo_key(item, out)),
9961 d => {
9962 for i in 0..d.len() {
9963 out.push(elem_key(d, i));
9964 }
9965 true
9966 }
9967 }
9968}
9969
9970fn memoised(
9972 u: &Verb,
9973 cache: &MemoCache,
9974 x: Option<&Array>,
9975 y: &Array,
9976 ctx: &mut Ctx<'_>,
9977 span: Span,
9978) -> Result<Array> {
9979 let apply = |ctx: &mut Ctx<'_>| match x {
9980 Some(x) => u.dyad(x, y, ctx, span),
9981 None => u.monad(y, ctx, span),
9982 };
9983 let mut key = vec![u64::from(x.is_some())];
9984 let keyed = x.is_none_or(|x| memo_key(x, &mut key)) && memo_key(y, &mut key);
9985 if !keyed {
9986 return apply(ctx);
9987 }
9988 if let Ok(map) = cache.lock() && let Some(hit) = map.get(&key) {
9989 return Ok(hit.clone());
9990 }
9991 let out = apply(ctx)?;
9992 if let Ok(mut map) = cache.lock() {
9993 map.insert(key, out.clone());
9994 }
9995 Ok(out)
9996}
9997
9998fn at_level(
10004 u: &Verb,
10005 level: i64,
10006 spread: bool,
10007 y: &Array,
10008 ctx: &mut Ctx<'_>,
10009 span: Span,
10010) -> Result<Array> {
10011 let n = if level < 0 { (boxing_level(y) + level).max(0) } else { level };
10013 if !spread {
10014 return map_level(u, n, y, ctx, span);
10015 }
10016 let mut cells = Vec::new();
10017 collect_level(u, n, y, ctx, span, &mut cells)?;
10018 let count = cells.len();
10019 assemble(&[count], cells, span)
10020}
10021
10022fn map_level(u: &Verb, n: i64, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
10023 let Some(boxes) = y.as_boxes().filter(|_| boxing_level(y) > n) else {
10024 return u.monad(y, ctx, span);
10025 };
10026 let boxes = boxes.to_vec();
10027 let mut cells = Vec::with_capacity(boxes.len());
10028 for b in &boxes {
10029 cells.push(map_level(u, n, b, ctx, span)?);
10030 }
10031 Ok(Array::new(y.shape.clone(), Data::Box(cells.into())))
10032}
10033
10034fn at_level_dyad(
10039 u: &Verb,
10040 level: i64,
10041 spread: bool,
10042 x: &Array,
10043 y: &Array,
10044 ctx: &mut Ctx<'_>,
10045 span: Span,
10046) -> Result<Array> {
10047 let depth = |a: &Array| if level < 0 { (boxing_level(a) + level).max(0) } else { level };
10050 let (nx, ny) = (depth(x), depth(y));
10051 if !spread {
10052 return map_level_dyad(u, nx, ny, x, y, ctx, span);
10053 }
10054 let mut cells = Vec::new();
10055 collect_level_dyad(u, nx, ny, x, y, ctx, span, &mut cells)?;
10056 let count = cells.len();
10057 assemble(&[count], cells, span)
10058}
10059
10060struct LevelPairs {
10062 left: Vec<Array>,
10063 right: Vec<Array>,
10064 shape: Vec<usize>,
10065}
10066
10067fn level_pairs(
10070 nx: i64,
10071 ny: i64,
10072 x: &Array,
10073 y: &Array,
10074 span: Span,
10075) -> Result<Option<LevelPairs>> {
10076 let bx = x.as_boxes().filter(|_| boxing_level(x) > nx);
10077 let by = y.as_boxes().filter(|_| boxing_level(y) > ny);
10078 Ok(match (bx, by) {
10079 (None, None) => None,
10080 (Some(bx), None) => {
10081 let n = bx.len();
10082 Some(LevelPairs {
10083 left: bx.to_vec(),
10084 right: vec![y.clone(); n],
10085 shape: x.shape.clone(),
10086 })
10087 }
10088 (None, Some(by)) => {
10089 let n = by.len();
10090 Some(LevelPairs {
10091 left: vec![x.clone(); n],
10092 right: by.to_vec(),
10093 shape: y.shape.clone(),
10094 })
10095 }
10096 (Some(bx), Some(by)) => {
10097 if x.shape != y.shape {
10098 return Err(Error::new(
10099 ErrorKind::Length,
10100 format!(
10101 "the levels do not agree: left shape {}, right shape {}",
10102 show_shape(&x.shape),
10103 show_shape(&y.shape)
10104 ),
10105 Some(span),
10106 ));
10107 }
10108 Some(LevelPairs { left: bx.to_vec(), right: by.to_vec(), shape: x.shape.clone() })
10109 }
10110 })
10111}
10112
10113fn map_level_dyad(
10114 u: &Verb,
10115 nx: i64,
10116 ny: i64,
10117 x: &Array,
10118 y: &Array,
10119 ctx: &mut Ctx<'_>,
10120 span: Span,
10121) -> Result<Array> {
10122 let Some(step) = level_pairs(nx, ny, x, y, span)? else {
10123 return u.dyad(x, y, ctx, span);
10124 };
10125 let mut cells = Vec::with_capacity(step.left.len());
10126 for (a, b) in step.left.iter().zip(step.right.iter()) {
10127 cells.push(map_level_dyad(u, nx, ny, a, b, ctx, span)?);
10128 }
10129 Ok(Array::new(step.shape, Data::Box(cells.into())))
10130}
10131
10132#[allow(clippy::too_many_arguments)]
10133fn collect_level_dyad(
10134 u: &Verb,
10135 nx: i64,
10136 ny: i64,
10137 x: &Array,
10138 y: &Array,
10139 ctx: &mut Ctx<'_>,
10140 span: Span,
10141 out: &mut Vec<Array>,
10142) -> Result<()> {
10143 let Some(step) = level_pairs(nx, ny, x, y, span)? else {
10144 out.push(u.dyad(x, y, ctx, span)?);
10145 return Ok(());
10146 };
10147 for (a, b) in step.left.iter().zip(step.right.iter()) {
10148 collect_level_dyad(u, nx, ny, a, b, ctx, span, out)?;
10149 }
10150 Ok(())
10151}
10152
10153fn collect_level(
10154 u: &Verb,
10155 n: i64,
10156 y: &Array,
10157 ctx: &mut Ctx<'_>,
10158 span: Span,
10159 out: &mut Vec<Array>,
10160) -> Result<()> {
10161 let Some(boxes) = y.as_boxes().filter(|_| boxing_level(y) > n) else {
10162 out.push(u.monad(y, ctx, span)?);
10163 return Ok(());
10164 };
10165 let boxes = boxes.to_vec();
10166 for b in &boxes {
10167 collect_level(u, n, b, ctx, span, out)?;
10168 }
10169 Ok(())
10170}
10171
10172fn poly_coeffs(y: &Array, span: Span) -> Result<Vec<Cx>> {
10176 let c = y
10177 .data
10178 .cast(DType::Complex)
10179 .ok_or_else(|| Error::domain("a polynomial's coefficients are numbers", span))?;
10180 match c {
10181 Data::Complex(v) => Ok(v.as_slice().to_vec()),
10182 _ => Err(Error::internal("coefficients did not cast to complex")),
10183 }
10184}
10185
10186const HYPERGEOMETRIC_TERMS: usize = 1 << 16;
10190
10191fn cx_list(v: &[Cx]) -> String {
10193 v.iter()
10194 .map(|z| if z[1] == 0.0 { format!("{}", z[0]) } else { format!("{}j{}", z[0], z[1]) })
10195 .collect::<Vec<_>>()
10196 .join(" ")
10197}
10198
10199fn hypergeometric(num: &[Cx], den: &[Cx], y: &Array, span: Span) -> Result<Array> {
10207 let (num, den) = cancel_parameters(num, den);
10208 let at = poly_coeffs(y, span)?;
10209 let mut out = Vec::with_capacity(at.len());
10210 for z in &at {
10211 out.push(hypergeometric_at(&num, &den, *z, span)?);
10212 }
10213 let mut a = complex_or_real(out);
10214 a.shape = y.shape.clone();
10215 Ok(a)
10216}
10217
10218fn cancel_parameters(num: &[Cx], den: &[Cx]) -> (Vec<Cx>, Vec<Cx>) {
10221 let mut left: Vec<Cx> = Vec::with_capacity(num.len());
10222 let mut right: Vec<Cx> = den.to_vec();
10223 for a in num {
10224 match right.iter().position(|b| b == a) {
10225 Some(i) => {
10226 right.remove(i);
10227 }
10228 None => left.push(*a),
10229 }
10230 }
10231 (left, right)
10232}
10233
10234fn hypergeometric_at(num: &[Cx], den: &[Cx], z: Cx, span: Span) -> Result<Cx> {
10235 let real = |v: &[Cx]| v.iter().all(|c| c[1] == 0.0);
10239 if z[1] == 0.0 && real(num) && real(den) {
10240 let n: Vec<f64> = num.iter().map(|c| c[0]).collect();
10241 let d: Vec<f64> = den.iter().map(|c| c[0]).collect();
10242 return Ok([hypergeometric_real(&n, &d, z[0], span)?, 0.0]);
10243 }
10244 let mut sum = cx::ONE;
10245 let mut term = cx::ONE;
10246 for k in 0..HYPERGEOMETRIC_TERMS {
10247 let kk = [k as f64, 0.0];
10248 let mut ratio = z;
10249 for a in num {
10250 ratio = cx::mul(ratio, cx::add(*a, kk));
10251 }
10252 for b in den {
10253 ratio = cx::div(ratio, cx::add(*b, kk));
10254 }
10255 term = cx::div(cx::mul(term, ratio), [k as f64 + 1.0, 0.0]);
10256 if !term[0].is_finite() || !term[1].is_finite() {
10257 return Ok(term);
10260 }
10261 let before = sum;
10262 sum = cx::add(sum, term);
10263 if sum == before {
10265 return Ok(sum);
10266 }
10267 }
10268 Err(Error::domain(
10269 format!("the hypergeometric series did not converge within {HYPERGEOMETRIC_TERMS} terms"),
10270 span,
10271 ))
10272}
10273
10274fn hypergeometric_real(num: &[f64], den: &[f64], z: f64, span: Span) -> Result<f64> {
10275 let mut sum = 1.0f64;
10276 let mut term = 1.0f64;
10277 for k in 0..HYPERGEOMETRIC_TERMS {
10278 let kk = k as f64;
10279 let mut ratio = z;
10280 for a in num {
10281 ratio *= a + kk;
10282 }
10283 for b in den {
10284 ratio /= b + kk;
10285 }
10286 term = term * ratio / (kk + 1.0);
10287 if !term.is_finite() {
10288 return Ok(term);
10289 }
10290 let before = sum;
10291 sum += term;
10292 if sum == before {
10293 return Ok(sum);
10294 }
10295 }
10296 Err(Error::domain(
10297 format!("the hypergeometric series did not converge within {HYPERGEOMETRIC_TERMS} terms"),
10298 span,
10299 ))
10300}
10301
10302fn complex_or_real(values: Vec<Cx>) -> Array {
10304 if values.iter().all(|z| z[1] == 0.0) {
10305 return Array::from_f64(values.iter().map(|z| z[0]).collect());
10306 }
10307 Array::new(vec![values.len()], Data::Complex(values.into()))
10308}
10309
10310fn poly_eval(x: &Array, y: &Array, span: Span) -> Result<Array> {
10313 let at = poly_coeffs(y, span)?;
10314 let at = at.first().copied().unwrap_or(cx::ZERO);
10315 let value = match x.as_boxes() {
10316 Some(parts) => {
10317 if parts.len() != 2 {
10318 return Err(Error::domain(
10319 "the root form of a polynomial is `multiplier ; roots`",
10320 span,
10321 ));
10322 }
10323 let multiplier = poly_coeffs(&parts[0], span)?;
10324 let mut v = multiplier.first().copied().unwrap_or(cx::ONE);
10325 for r in poly_coeffs(&parts[1], span)? {
10326 v = cx::mul(v, cx::sub(at, r));
10327 }
10328 v
10329 }
10330 None => {
10331 let c = poly_coeffs(x, span)?;
10332 let mut v = cx::ZERO;
10333 for &k in c.iter().rev() {
10334 v = cx::add(cx::mul(v, at), k);
10335 }
10336 v
10337 }
10338 };
10339 Ok(scalar_complex_or_real(value))
10340}
10341
10342fn scalar_complex_or_real(z: Cx) -> Array {
10343 if z[1] == 0.0 {
10344 return Array::scalar_f64(z[0]);
10345 }
10346 Array::new(vec![], Data::Complex(vec![z].into()))
10347}
10348
10349fn poly_roots(y: &Array, span: Span) -> Result<Array> {
10353 if let Some(parts) = y.as_boxes() {
10354 if parts.len() != 2 {
10355 return Err(Error::domain(
10356 "the root form of a polynomial is `multiplier ; roots`",
10357 span,
10358 ));
10359 }
10360 let multiplier = poly_coeffs(&parts[0], span)?;
10361 let multiplier = multiplier.first().copied().unwrap_or(cx::ONE);
10362 let mut coeffs = vec![multiplier];
10364 for r in poly_coeffs(&parts[1], span)? {
10365 let mut next = vec![cx::ZERO; coeffs.len() + 1];
10366 for (k, &c) in coeffs.iter().enumerate() {
10367 next[k + 1] = cx::add(next[k + 1], c);
10368 next[k] = cx::sub(next[k], cx::mul(c, r));
10369 }
10370 coeffs = next;
10371 }
10372 return Ok(complex_or_real(coeffs));
10373 }
10374 let mut c = poly_coeffs(y, span)?;
10375 while c.len() > 1 && c[c.len() - 1] == cx::ZERO {
10376 c.pop();
10377 }
10378 if c.iter().all(|&k| k == cx::ZERO) {
10382 let pair = vec![Array::scalar_i64(0), Array::new(vec![0], Data::empty(DType::I64))];
10383 return Ok(Array::new(vec![2], Data::Box(pair.into())));
10384 }
10385 if c.len() < 2 {
10386 return Err(Error::domain("a polynomial's roots need a coefficient of x", span));
10387 }
10388 let lead = c[c.len() - 1];
10389 let monic: Vec<Cx> = c.iter().map(|&k| cx::div(k, lead)).collect();
10390 let roots = durand_kerner(&monic);
10391 let pair = vec![scalar_complex_or_real(lead), complex_or_real(roots)];
10392 Ok(Array::new(vec![2], Data::Box(pair.into())))
10393}
10394
10395fn durand_kerner(monic: &[Cx]) -> Vec<Cx> {
10402 let d = monic.len() - 1;
10403 let seed = [0.4, 0.9];
10404 let mut z: Vec<Cx> = Vec::with_capacity(d);
10405 let mut p = cx::ONE;
10406 for _ in 0..d {
10407 z.push(p);
10408 p = cx::mul(p, seed);
10409 }
10410 let value = |monic: &[Cx], at: Cx| {
10411 let mut v = cx::ZERO;
10412 for &k in monic.iter().rev() {
10413 v = cx::add(cx::mul(v, at), k);
10414 }
10415 v
10416 };
10417 for _ in 0..500 {
10418 let mut moved: f64 = 0.0;
10419 for i in 0..d {
10420 let mut denom = cx::ONE;
10421 for j in 0..d {
10422 if i != j {
10423 denom = cx::mul(denom, cx::sub(z[i], z[j]));
10424 }
10425 }
10426 if denom == cx::ZERO {
10427 continue;
10428 }
10429 let step = cx::div(value(monic, z[i]), denom);
10430 z[i] = cx::sub(z[i], step);
10431 moved = moved.max(step[0].hypot(step[1]));
10432 }
10433 if moved < 1e-15 {
10434 break;
10435 }
10436 }
10437 for r in &mut z {
10439 if r[1].abs() < 1e-9 {
10440 r[1] = 0.0;
10441 }
10442 if r[0].abs() < 1e-12 {
10443 r[0] = 0.0;
10444 }
10445 }
10446 z.sort_by(|a, b| {
10450 let close = (a[0] - b[0]).abs() <= 1e-9 * (a[0].abs().max(b[0].abs()) + 1.0);
10451 let by_re = if close {
10452 std::cmp::Ordering::Equal
10453 } else {
10454 b[0].partial_cmp(&a[0]).unwrap_or(std::cmp::Ordering::Equal)
10455 };
10456 by_re.then(b[1].partial_cmp(&a[1]).unwrap_or(std::cmp::Ordering::Equal))
10457 });
10458 z
10459}
10460
10461fn poly_deriv(y: &Array, span: Span) -> Result<Array> {
10464 let c = poly_coeffs(y, span)?;
10465 if c.len() < 2 {
10466 return Ok(Array::from_i64(vec![0]));
10467 }
10468 let out: Vec<Cx> =
10469 c.iter().enumerate().skip(1).map(|(k, &v)| cx::mul(v, cx::from_real(k as f64))).collect();
10470 Ok(narrow_numbers(complex_or_real(out)))
10471}
10472
10473fn poly_integral(x: &Array, y: &Array, span: Span) -> Result<Array> {
10475 let c = poly_coeffs(y, span)?;
10476 let k = poly_coeffs(x, span)?;
10477 let mut out = vec![k.first().copied().unwrap_or(cx::ZERO)];
10478 for (i, &v) in c.iter().enumerate() {
10479 out.push(cx::div(v, cx::from_real((i + 1) as f64)));
10480 }
10481 Ok(narrow_numbers(complex_or_real(out)))
10482}
10483
10484fn narrow_numbers(a: Array) -> Array {
10488 let Data::F64(v) = &a.data else { return a };
10489 if v.iter().any(|x| !x.is_finite() || x.fract() != 0.0 || x.abs() > 9e15) {
10490 return a;
10491 }
10492 let values: Vec<i64> = v.iter().map(|&x| x as i64).collect();
10493 Array::new(a.shape, Data::I64(values.into()))
10494}
10495
10496fn characteristics(u: &Verb, y: &Array, span: Span) -> Result<Array> {
10500 let which = y.to_i64_vec().and_then(|v| v.first().copied());
10501 let chars = |s: String| Ok(Array::from_chars(s.chars().collect()));
10502 match which {
10503 Some(0) => {
10504 let ranks = u.ranks();
10505 Ok(Array::from_f64(
10506 ranks
10507 .iter()
10508 .map(|&r| if r == RANK_INF { f64::INFINITY } else { r as f64 })
10509 .collect(),
10510 ))
10511 }
10512 Some(-1) => match obverse(u) {
10516 Some(v) => chars(v.name()),
10517 None => Err(Error::not_yet(
10518 format!("the obverse of {} (no inverse is known)", u.name()),
10519 span,
10520 )),
10521 },
10522 Some(1) => match reduce_identity(u, 1).as_ref().map(identity_spelling) {
10523 Some(s) => chars(s),
10524 None => Err(Error::not_yet(
10525 format!("the identity function of {} (u b. 1)", u.name()),
10526 span,
10527 )),
10528 },
10529 _ => Err(Error::not_yet("a verb characteristic other than 0, 1 and _1", span)),
10530 }
10531}
10532
10533fn identity_spelling(d: &Data) -> String {
10536 let one = Array::new(Vec::new(), d.slice(0, 1));
10537 let text = crate::fmt::format_array(&one, &crate::fmt::FmtOpts::J);
10538 format!("{} $~ }}.@$", text.trim())
10539}
10540
10541fn with_operands<R>(
10544 alpha: &Verb,
10545 omega: Option<&Verb>,
10546 ctx: &mut Ctx<'_>,
10547 f: impl FnOnce(&mut Ctx<'_>) -> Result<R>,
10548) -> Result<R> {
10549 let saved = (ctx.env.verb("⍺⍺").cloned(), ctx.env.verb("⍵⍵").cloned());
10550 ctx.env.define("⍺⍺".to_string(), alpha.clone());
10551 if let Some(g) = omega {
10552 ctx.env.define("⍵⍵".to_string(), g.clone());
10553 }
10554 let out = f(ctx);
10555 match saved.0 {
10556 Some(v) => ctx.env.define("⍺⍺".to_string(), v),
10557 None => ctx.env.undefine("⍺⍺"),
10558 }
10559 match saved.1 {
10560 Some(v) => ctx.env.define("⍵⍵".to_string(), v),
10561 None => ctx.env.undefine("⍵⍵"),
10562 }
10563 out
10564}
10565
10566fn is_mixed_simple(a: &Array) -> bool {
10570 let Some(items) = a.as_boxes() else { return false };
10571 if items.is_empty() || items.iter().any(|b| b.rank() != 0 || b.dtype() == DType::Box) {
10572 return false;
10573 }
10574 let mut common = Some(items[0].dtype());
10575 for b in &items[1..] {
10576 common = common.and_then(|t| DType::promote(t, b.dtype()));
10577 }
10578 common.is_none()
10579}
10580
10581fn nest(y: &Array) -> Array {
10584 if y.dtype() == DType::Box || y.rank() == 0 {
10585 return y.clone();
10586 }
10587 Array::boxed(y.clone())
10588}
10589
10590fn key_pairs(
10594 u: &Verb,
10595 keys: &Array,
10596 values: Option<&Array>,
10597 ctx: &mut Ctx<'_>,
10598 span: Span,
10599) -> Result<Array> {
10600 let base = if keys.rank() == 0 { Array::new(vec![1], keys.data.clone()) } else { keys.clone() };
10601 let n = base.items();
10602 if let Some(v) = values && v.items() != n {
10603 return Err(Error::new(
10604 ErrorKind::Length,
10605 format!("{n} key(s) for {} item(s)", v.items()),
10606 Some(span),
10607 ));
10608 }
10609 let groups = group_positions(&base, ctx.cfg.tol);
10610 let origin = ctx.cfg.rules.origin;
10611 let mut cells = Vec::with_capacity(groups.len());
10612 for (first, at) in &groups {
10613 let key = item_or_self(&base, *first);
10614 let group = match values {
10615 Some(v) => select_items(v, at),
10616 None => Array::from_i64(at.iter().map(|&i| origin + i as i64).collect()),
10617 };
10618 let monadic = matches!(u, Verb::Explicit(d) if d.left.is_none());
10621 cells.push(if monadic {
10622 u.monad(&group, ctx, span)?
10623 } else {
10624 u.dyad(&key, &group, ctx, span)?
10625 });
10626 }
10627 let count = cells.len();
10628 assemble(&[count], cells, span)
10629}
10630
10631fn group_positions(y: &Array, tol: Tol) -> Vec<(usize, Vec<usize>)> {
10634 let n = y.items();
10635 let m = y.item_size();
10636 let hashable = match y.dtype() {
10641 DType::Box | DType::Ext | DType::Rat => false,
10642 DType::F64 | DType::Complex => tol.ct == 0.0,
10643 _ => true,
10644 };
10645 if hashable {
10646 return if m == 1 {
10647 group_by_key(n, |i| elem_key(&y.data, i))
10648 } else {
10649 group_by_key(n, |i| (0..m).map(|k| elem_key(&y.data, i * m + k)).collect::<Vec<u64>>())
10650 };
10651 }
10652 let mut keys: Vec<Array> = Vec::new();
10653 let mut groups: Vec<(usize, Vec<usize>)> = Vec::new();
10654 for i in 0..n {
10655 let item = y.item(i);
10656 match keys.iter().position(|k| arrays_match(k, &item, tol)) {
10657 Some(at) => groups[at].1.push(i),
10658 None => {
10659 keys.push(item);
10660 groups.push((i, vec![i]));
10661 }
10662 }
10663 }
10664 groups
10665}
10666
10667fn group_by_key<K, F>(n: usize, key: F) -> Vec<(usize, Vec<usize>)>
10671where
10672 K: Eq + std::hash::Hash,
10673 F: Fn(usize) -> K,
10674{
10675 use std::collections::hash_map::Entry;
10676 let mut groups: Vec<(usize, Vec<usize>)> = Vec::new();
10677 let mut at: HashMap<K, usize, KeyHash> =
10678 HashMap::with_capacity_and_hasher(n.min(1 << 16), KeyHash);
10679 for i in 0..n {
10680 match at.entry(key(i)) {
10681 Entry::Occupied(e) => groups[*e.get()].1.push(i),
10682 Entry::Vacant(e) => {
10683 e.insert(groups.len());
10684 groups.push((i, vec![i]));
10685 }
10686 }
10687 }
10688 groups
10689}
10690
10691#[derive(Clone, Copy, Default)]
10697struct KeyHash;
10698
10699impl std::hash::BuildHasher for KeyHash {
10700 type Hasher = KeyHasher;
10701 fn build_hasher(&self) -> KeyHasher {
10702 KeyHasher(0)
10703 }
10704}
10705
10706struct KeyHasher(u64);
10707
10708impl std::hash::Hasher for KeyHasher {
10709 fn finish(&self) -> u64 {
10710 let mut x = self.0;
10711 x ^= x >> 33;
10712 x = x.wrapping_mul(0xff51_afd7_ed55_8ccd);
10713 x ^ (x >> 29)
10714 }
10715 fn write(&mut self, bytes: &[u8]) {
10716 for &b in bytes {
10717 self.write_u64(b as u64);
10718 }
10719 }
10720 fn write_u64(&mut self, n: u64) {
10721 self.0 = (self.0.rotate_left(5) ^ n).wrapping_mul(0x9e37_79b9_7f4a_7c15);
10722 }
10723 fn write_usize(&mut self, n: usize) {
10724 self.write_u64(n as u64);
10725 }
10726}
10727
10728fn format_spec(x: &Array, y: &Array, fmt: &FmtOpts, span: Span) -> Result<Array> {
10734 let spec = x
10735 .to_i64_vec()
10736 .ok_or_else(|| Error::domain("a format specification is whole numbers", span))?;
10737 if y.dtype() == DType::Box {
10738 return Err(Error::not_yet("format by specification of a nested array", span));
10739 }
10740 let cols = if y.rank() == 0 { 1 } else { y.shape[y.rank() - 1] };
10741 let rows = y.count() / cols.max(1);
10742 let pairs: Vec<(Option<i64>, i64)> = match spec.len() {
10744 1 => vec![(None, spec[0]); cols],
10745 2 => vec![(Some(spec[0]), spec[1]); cols],
10746 n if n == 2 * cols => spec.chunks(2).map(|c| (Some(c[0]), c[1])).collect(),
10747 n => {
10748 return Err(Error::new(
10749 ErrorKind::Length,
10750 format!("{n} specification value(s) for {cols} column(s)"),
10751 Some(span),
10752 ));
10753 }
10754 };
10755 if pairs.iter().any(|&(w, p)| w.is_some_and(|w| w < 0) || p < 0) {
10756 return Err(Error::domain("a format width and precision are nonnegative", span));
10757 }
10758 let numbers = y.to_f64_vec();
10759 let text = |i: usize, p: i64| -> String {
10760 match (&y.data, &numbers) {
10761 (Data::Char(v), _) => v[i].to_string(),
10762 (_, Some(v)) => {
10763 let s = format!("{:.*}", p as usize, v[i]);
10764 if v[i] < 0.0 { format!("{}{}", fmt.neg, &s[1..]) } else { s }
10765 }
10766 _ => String::new(),
10767 }
10768 };
10769 if y.dtype() != DType::Char && numbers.is_none() {
10770 return Err(Error::domain("format by specification takes numbers or characters", span));
10771 }
10772 let widths: Vec<usize> = pairs
10774 .iter()
10775 .enumerate()
10776 .map(|(c, &(w, p))| match w {
10777 Some(w) => w as usize,
10778 None => {
10779 (0..rows).map(|r| text(r * cols + c, p).chars().count()).max().unwrap_or(0) + 1
10780 }
10781 })
10782 .collect();
10783 let line: usize = widths.iter().sum();
10784 let mut out: Vec<char> = Vec::with_capacity(rows * line);
10785 for r in 0..rows {
10786 for c in 0..cols {
10787 let s = text(r * cols + c, pairs[c].1);
10788 let len = s.chars().count();
10789 if len > widths[c] {
10790 return Err(Error::domain(
10791 format!("{s} does not fit a field {} wide", widths[c]),
10792 span,
10793 ));
10794 }
10795 out.extend(std::iter::repeat_n(' ', widths[c] - len));
10796 out.extend(s.chars());
10797 }
10798 }
10799 let mut shape = if y.rank() == 0 { Vec::new() } else { y.shape[..y.rank() - 1].to_vec() };
10800 shape.push(line);
10801 Ok(Array::new(shape, Data::Char(out.into())))
10802}
10803
10804fn sequential_machine(x: &Array, y: &Array, span: Span) -> Result<Array> {
10819 let Some(parts) = x.as_boxes() else {
10820 return Err(Error::domain("a sequential machine is a boxed description", span));
10821 };
10822 if x.rank() > 1 || !(2..=4).contains(&parts.len()) {
10823 return Err(Error::domain(
10824 "a sequential machine is 2 to 4 boxes: f ; s ; m ; ijrd",
10825 span,
10826 ));
10827 }
10828 let whole = |a: &Array, what: &str| -> Result<Vec<i64>> {
10829 a.to_i64_vec().ok_or_else(|| Error::domain(format!("{what} is whole numbers"), span))
10830 };
10831 let form = *whole(&parts[0], "a sequential machine's result form")?
10832 .first()
10833 .ok_or_else(|| Error::domain("a sequential machine needs a result form", span))?;
10834 if !(0..=5).contains(&form) {
10835 return Err(Error::domain(format!("{form} is not a result form of 0 to 5"), span));
10836 }
10837 let table = &parts[1];
10838 if table.rank() != 3 || table.shape[2] != 2 {
10839 return Err(Error::new(
10840 ErrorKind::Rank,
10841 "a sequential machine's transition table is shaped p q 2",
10842 Some(span),
10843 ));
10844 }
10845 let (states, classes) = (table.shape[0], table.shape[1]);
10846 let entries = whole(table, "a transition table")?;
10847 let map = parts.get(2).filter(|a| a.count() > 0);
10848 let start = match parts.get(3) {
10849 Some(a) => whole(a, "a sequential machine's starting values")?,
10850 None => Vec::new(),
10851 };
10852 let start = if start.is_empty() { vec![0, -1, 0, -1] } else { start };
10853 if start.len() != 4 {
10854 return Err(Error::new(
10855 ErrorKind::Length,
10856 "a sequential machine starts from four values: i j r d",
10857 Some(span),
10858 ));
10859 }
10860 let (mut i, mut word, mut state, ending) = (start[0], start[1], start[2], start[3]);
10861 let n = y.count() as i64;
10862
10863 let codes: Option<Vec<i64>> = match map {
10866 Some(m) => Some(whole(m, "a sequential machine's map")?),
10867 None => None,
10868 };
10869 let values: Vec<i64> = match (&y.data, &codes) {
10870 (Data::Char(v), Some(_)) => v.as_slice().iter().map(|&c| c as i64).collect(),
10871 (_, None) => y
10872 .to_i64_vec()
10873 .ok_or_else(|| Error::domain("a sequential machine over characters needs a map", span))?,
10874 _ => {
10875 return Err(Error::not_yet(
10876 "a sequential machine's map over a numeric argument (x's third box)",
10877 span,
10878 ));
10879 }
10880 };
10881 let class_at = |at: i64| -> Result<i64> {
10882 let raw = values[at as usize];
10883 let Some(m) = &codes else { return Ok(raw) };
10884 if raw < 0 || raw as usize >= m.len() {
10885 return Err(Error::new(
10886 ErrorKind::Domain,
10887 format!("{raw} is outside a map of {} entries", m.len()),
10888 Some(span),
10889 ));
10890 }
10891 Ok(m[raw as usize])
10892 };
10893
10894 let mut trace: Vec<i64> = Vec::new();
10895 let mut words: Vec<(i64, i64, i64)> = Vec::new();
10896 let mut emit = |word: i64, at: i64, place: i64| -> Result<()> {
10897 if word < 0 {
10898 return Err(Error::new(
10899 ErrorKind::Domain,
10900 "a sequential machine ended a word before one had begun",
10901 Some(span),
10902 ));
10903 }
10904 words.push((word, at - word, place));
10905 Ok(())
10906 };
10907 loop {
10908 let class = if i < n {
10909 class_at(i)?
10910 } else if ending >= 0 {
10911 ending
10912 } else {
10913 if word >= 0 {
10917 emit(word, i, classes as i64 * state)?;
10918 }
10919 break;
10920 };
10921 if state < 0 || state as usize >= states || class < 0 || class as usize >= classes {
10922 return Err(Error::new(
10923 ErrorKind::Domain,
10924 format!(
10925 "state {state} and class {class} are outside a {states} by {classes} table"
10926 ),
10927 Some(span),
10928 ));
10929 }
10930 let at = (state as usize * classes + class as usize) * 2;
10931 let (next, code) = (entries[at], entries[at + 1]);
10932 trace.extend_from_slice(&[i, word, state, class, next, code]);
10933 let place = class + classes as i64 * state;
10934 state = next;
10935 match code {
10936 0 => {}
10937 1 => word = i,
10938 2 => {
10939 emit(word, i, place)?;
10940 word = i;
10941 }
10942 3 => {
10943 emit(word, i, place)?;
10944 word = -1;
10945 }
10946 4 | 5 => {
10947 return Err(Error::not_yet(
10948 "a sequential machine's vector output (codes 4 and 5)",
10949 span,
10950 ));
10951 }
10952 6 => break,
10953 other => {
10954 return Err(Error::domain(
10955 format!("{other} is not a sequential machine output code"),
10956 span,
10957 ));
10958 }
10959 }
10960 if i >= n {
10961 break;
10962 }
10963 i += 1;
10964 }
10965 Ok(sequential_result(form, &words, &trace, y))
10966}
10967
10968fn sequential_result(form: i64, words: &[(i64, i64, i64)], trace: &[i64], y: &Array) -> Array {
10971 let piece = |&(at, len, _): &(i64, i64, i64)| {
10972 Array::new(vec![len as usize], y.data.slice(at as usize, (at + len) as usize))
10973 };
10974 match form {
10975 0 => Array::new(
10976 vec![words.len()],
10977 Data::Box(words.iter().map(piece).collect::<Vec<_>>().into()),
10978 ),
10979 1 => {
10980 let mut data = Data::empty(y.dtype());
10981 for w in words {
10982 data.extend_from(&piece(w).data);
10983 }
10984 let n = data.len();
10985 Array::new(vec![n], data)
10986 }
10987 2 => Array::new(
10988 vec![words.len(), 2],
10989 Data::I64(words.iter().flat_map(|&(at, len, _)| [at, len]).collect::<Vec<_>>().into()),
10990 ),
10991 3 => Array::from_i64(words.iter().map(|&(_, _, place)| place).collect()),
10992 4 => Array::new(
10993 vec![words.len(), 3],
10994 Data::I64(
10995 words
10996 .iter()
10997 .flat_map(|&(at, len, place)| [at, len, place])
10998 .collect::<Vec<_>>()
10999 .into(),
11000 ),
11001 ),
11002 _ => Array::new(vec![trace.len() / 6, 6], Data::I64(trace.to_vec().into())),
11003 }
11004}
11005
11006fn parse_numbers(x: &Array, y: &Array, span: Span) -> Result<Array> {
11011 let Data::Char(text) = &y.data else {
11012 return Err(Error::domain("reading numbers from text needs characters", span));
11013 };
11014 if x.count() != 1 {
11015 return Err(Error::new(
11016 ErrorKind::Rank,
11017 "the stand-in for an unreadable word is one value",
11018 Some(span),
11019 ));
11020 }
11021 let line: String = text.as_slice().iter().collect();
11022 crate::frontend::j::numbers_from_text(&line, x)
11023 .ok_or_else(|| Error::domain("the stand-in for an unreadable word is a number", span))
11024}
11025
11026fn format_field(value: f64, precision: usize, exponential: bool, neg: char) -> String {
11030 let sign = |s: String| match s.strip_prefix('-') {
11031 Some(rest) if rest.bytes().all(|b| !b.is_ascii_digit() || b == b'0') => rest.to_string(),
11034 Some(rest) => format!("{neg}{rest}"),
11035 None => s,
11036 };
11037 if !exponential {
11038 return sign(format!("{value:.precision$}"));
11039 }
11040 let text = format!("{value:.precision$e}");
11043 let (mantissa, exponent) = text.split_once('e').unwrap_or((text.as_str(), "0"));
11044 let exponent = match exponent.strip_prefix('-') {
11045 Some(rest) => format!("{neg}{rest}"),
11046 None => exponent.to_string(),
11047 };
11048 format!("{}e{exponent}", sign(mantissa.to_string()))
11049}
11050
11051fn format_spec_j(x: &Array, y: &Array, fmt: &FmtOpts, span: Span) -> Result<Array> {
11059 let Some(spec) = x.to_complex_vec() else {
11060 return Err(Error::domain("a format specification is numbers", span));
11061 };
11062 if y.dtype() == DType::Box {
11063 return Err(Error::domain("format by specification takes numbers", span));
11064 }
11065 let Some(values) = y.to_f64_vec() else {
11066 return Err(Error::domain("format by specification takes numbers", span));
11067 };
11068 let cols = if y.rank() == 0 { 1 } else { y.shape[y.rank() - 1] };
11069 let rows = if cols == 0 { 0 } else { y.count() / cols };
11070 let fields: Vec<[f64; 2]> = match spec.len() {
11071 1 => vec![spec[0]; cols],
11072 n if n == cols => spec,
11073 n => {
11074 return Err(Error::new(
11075 ErrorKind::Length,
11076 format!("{n} specification value(s) for {cols} column(s)"),
11077 Some(span),
11078 ));
11079 }
11080 };
11081 let text = |r: usize, c: usize| {
11082 let [w, d] = fields[c];
11083 format_field(values[r * cols + c], d.max(0.0) as usize, w < 0.0, fmt.neg)
11084 };
11085 let widths: Vec<usize> = (0..cols)
11088 .map(|c| {
11089 let w = fields[c][0];
11090 if w != 0.0 {
11091 return w.abs() as usize;
11092 }
11093 let wide = (0..rows).map(|r| text(r, c).chars().count()).max().unwrap_or(0);
11094 wide + usize::from(c > 0)
11095 })
11096 .collect();
11097 let line: usize = widths.iter().sum();
11098 let mut out: Vec<char> = Vec::with_capacity(rows * line);
11099 for r in 0..rows {
11100 for c in 0..cols {
11101 let s = text(r, c);
11102 let (lead, body) = match (fields[c][0] < 0.0, s.strip_prefix(fmt.neg)) {
11105 (false, _) => (String::new(), s.as_str()),
11106 (true, Some(rest)) => (fmt.neg.to_string(), rest),
11107 (true, None) => (" ".to_string(), s.as_str()),
11108 };
11109 let len = lead.chars().count() + body.chars().count();
11110 if len > widths[c] {
11111 out.extend(std::iter::repeat_n('*', widths[c]));
11112 continue;
11113 }
11114 if fields[c][0] < 0.0 {
11115 out.extend(lead.chars());
11116 out.extend(body.chars());
11117 out.extend(std::iter::repeat_n(' ', widths[c] - len));
11118 } else {
11119 out.extend(std::iter::repeat_n(' ', widths[c] - len));
11120 out.extend(body.chars());
11121 }
11122 }
11123 }
11124 let mut shape = if y.rank() == 0 { Vec::new() } else { y.shape[..y.rank() - 1].to_vec() };
11125 shape.push(line);
11126 Ok(Array::new(shape, Data::Char(out.into())))
11127}
11128
11129fn iota_apl(y: &Array, origin: i64, span: Span) -> Result<Array> {
11133 if y.rank() > 1 {
11134 return Err(Error::new(
11135 ErrorKind::Rank,
11136 "the index generator takes a shape, which is a scalar or a vector",
11137 Some(span),
11138 ));
11139 }
11140 let dims = y
11141 .to_i64_vec()
11142 .ok_or_else(|| Error::domain("index generator needs an integer argument", span))?;
11143 if dims.iter().any(|&n| n < 0) {
11144 return Err(Error::domain("index generator needs nonnegative lengths", span));
11145 }
11146 if dims.len() <= 1 {
11147 let n = dims.first().copied().unwrap_or(0);
11148 crate::limits::count(n as u128, span)?;
11149 return Ok(Array::from_i64((0..n).map(|i| origin + i).collect()));
11150 }
11151 let shape: Vec<usize> = dims.iter().map(|&n| n as usize).collect();
11152 let total = crate::limits::elements(&shape, span)?;
11153 let mut cells = Vec::with_capacity(total);
11154 let mut coord = vec![0usize; shape.len()];
11155 for _ in 0..total {
11156 cells.push(Array::from_i64(coord.iter().map(|&c| origin + c as i64).collect()));
11157 odometer(&mut coord, &shape);
11158 }
11159 Ok(Array::new(shape, Data::Box(cells.into())))
11160}
11161
11162fn carry_exact(result: Array, y: &Array) -> Array {
11168 if !matches!(y.dtype(), DType::Ext | DType::Rat) {
11169 return result;
11170 }
11171 match result.data.cast(DType::Ext) {
11172 Some(data) => Array::new(result.shape, data),
11173 None => result,
11174 }
11175}
11176
11177fn carry_exact2(result: Array, x: &Array, y: &Array) -> Array {
11178 let widened = carry_exact(result, x);
11179 carry_exact(widened, y)
11180}
11181
11182fn truth_table(m: u8, x: &Array, y: &Array, span: Span) -> Result<Array> {
11185 let table = m & 15;
11186 let bit = |a: i64, b: i64| ((table >> (3 - (2 * a + b))) & 1) as i64;
11187 let xs = x
11188 .to_i64_vec()
11189 .ok_or_else(|| Error::domain("a boolean function takes integers", span))?;
11190 let ys = y
11191 .to_i64_vec()
11192 .ok_or_else(|| Error::domain("a boolean function takes integers", span))?;
11193 let (a, b) = (xs.first().copied().unwrap_or(0), ys.first().copied().unwrap_or(0));
11194 if m < 16 {
11195 if !(0..=1).contains(&a) || !(0..=1).contains(&b) {
11196 return Err(Error::domain(
11197 format!("{m} b. takes 0 and 1; {m} b. + 16 is the same function on every bit"),
11198 span,
11199 ));
11200 }
11201 return Ok(Array::scalar_bool(bit(a, b) != 0));
11202 }
11203 let mut out = 0i64;
11204 for k in 0..64 {
11205 if bit((a >> k) & 1, (b >> k) & 1) != 0 {
11206 out |= 1i64 << k;
11207 }
11208 }
11209 Ok(Array::scalar_i64(out))
11210}
11211
11212pub fn amend_at(
11217 base: &Array,
11218 slots: &[Option<Array>],
11219 value: &Array,
11220 origin: i64,
11221 span: Span,
11222) -> Result<Array> {
11223 if slots.len() != base.rank() {
11224 return Err(Error::new(
11225 ErrorKind::Rank,
11226 format!(
11227 "indexed assignment needs one index per axis: {} slot(s) for a rank-{} value",
11228 slots.len(),
11229 base.rank()
11230 ),
11231 Some(span),
11232 ));
11233 }
11234 if !base.is_row_major() || !value.is_row_major() {
11237 let (b, v) = (base.to_row_major(), value.to_row_major());
11238 return amend_at(&b, slots, &v, origin, span);
11239 }
11240 let mut axes: Vec<Vec<usize>> = Vec::with_capacity(slots.len());
11242 let mut selected: Vec<usize> = Vec::new();
11243 for (k, slot) in slots.iter().enumerate() {
11244 let len = base.shape[k];
11245 let Some(idx) = slot else {
11246 axes.push((0..len).collect());
11247 selected.push(len);
11248 continue;
11249 };
11250 let Some(values) = idx.to_i64_vec() else {
11251 return Err(Error::new(
11252 ErrorKind::Type,
11253 "an index must be numeric",
11254 Some(span),
11255 ));
11256 };
11257 let mut positions = Vec::with_capacity(values.len());
11258 for v in values {
11259 let p = v - origin;
11260 if p < 0 || p as usize >= len {
11261 return Err(Error::new(
11262 ErrorKind::Domain,
11263 format!("index {v} is outside axis {k}, which has {len} element(s)"),
11264 Some(span),
11265 ));
11266 }
11267 positions.push(p as usize);
11268 }
11269 if idx.rank() > 0 {
11271 selected.push(positions.len());
11272 }
11273 axes.push(positions);
11274 }
11275 let count: usize = axes.iter().map(Vec::len).product();
11276 if value.rank() != 0 && (value.shape != selected || value.count() != count) {
11277 return Err(Error::new(
11278 ErrorKind::Shape,
11279 format!(
11280 "indexed assignment needs a scalar or a {} value, not a {} one",
11281 show_shape(&selected),
11282 show_shape(&value.shape)
11283 ),
11284 Some(span),
11285 ));
11286 }
11287 let dtype = DType::promote(base.dtype(), value.dtype()).ok_or_else(|| {
11290 Error::new(
11291 ErrorKind::Type,
11292 format!(
11293 "cannot put a {} value into a {} array",
11294 value.dtype().name(),
11295 base.dtype().name()
11296 ),
11297 Some(span),
11298 )
11299 })?;
11300 let mut out = base.cast(dtype).ok_or_else(|| Error::internal("promotion failed"))?;
11301 let src = value.cast(dtype).ok_or_else(|| Error::internal("promotion failed"))?;
11302 let strides = row_major_strides(&base.shape);
11303 let mut coords = vec![0usize; axes.len()];
11304 for n in 0..count {
11305 let mut rest = n;
11306 for k in (0..axes.len()).rev() {
11307 let len = axes[k].len();
11308 coords[k] = axes[k][rest % len];
11309 rest /= len;
11310 }
11311 let at: usize = coords.iter().zip(&strides).map(|(c, s)| c * s).sum();
11312 let from = if src.rank() == 0 { 0 } else { n };
11313 put_element(&mut out.data, at, &src.data, from);
11314 }
11315 Ok(out)
11316}
11317
11318fn row_major_strides(shape: &[usize]) -> Vec<usize> {
11319 let mut strides = vec![1usize; shape.len()];
11320 for k in (0..shape.len().saturating_sub(1)).rev() {
11321 strides[k] = strides[k + 1] * shape[k + 1];
11322 }
11323 strides
11324}
11325
11326fn put_element(dst: &mut Data, at: usize, src: &Data, from: usize) {
11328 match (dst, src) {
11329 (Data::Bool(d), Data::Bool(s)) => d.to_mut()[at] = s.as_slice()[from],
11330 (Data::I64(d), Data::I64(s)) => d.to_mut()[at] = s.as_slice()[from],
11331 (Data::Ext(d), Data::Ext(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
11332 (Data::Rat(d), Data::Rat(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
11333 (Data::F64(d), Data::F64(s)) => d.to_mut()[at] = s.as_slice()[from],
11334 (Data::Char(d), Data::Char(s)) => d.to_mut()[at] = s.as_slice()[from],
11335 (Data::Box(d), Data::Box(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
11336 _ => debug_assert!(false, "amend across types"),
11338 }
11339}
11340
11341fn agenda_pick(
11344 vs: &[Verb],
11345 w: &Verb,
11346 x: Option<&Array>,
11347 y: &Array,
11348 ctx: &mut Ctx<'_>,
11349 span: Span,
11350) -> Result<Verb> {
11351 let chosen = match x {
11352 None => w.monad(y, ctx, span)?,
11353 Some(x) => w.dyad(x, y, ctx, span)?,
11354 };
11355 let at = chosen
11356 .to_i64_vec()
11357 .and_then(|v| v.first().copied())
11358 .ok_or_else(|| Error::domain("an agenda index must be an integer", span))?;
11359 pick_gerund(vs, at, span)
11360}
11361
11362pub(crate) fn pick_gerund(vs: &[Verb], at: i64, span: Span) -> Result<Verb> {
11365 usize::try_from(at)
11366 .ok()
11367 .and_then(|k| vs.get(k))
11368 .cloned()
11369 .ok_or_else(|| {
11370 Error::domain(
11371 format!("agenda {at} is out of range: the gerund has {} verbs", vs.len()),
11372 span,
11373 )
11374 })
11375}
11376
11377fn evoke(
11382 vs: &[Verb],
11383 form: i64,
11384 x: Option<&Array>,
11385 y: &Array,
11386 ctx: &mut Ctx<'_>,
11387 span: Span,
11388) -> Result<Array> {
11389 if vs.is_empty() {
11390 return Err(Error::domain("an evoked gerund is empty", span));
11391 }
11392 if form == 0 {
11393 let mut cells = Vec::with_capacity(vs.len());
11394 for v in vs {
11395 cells.push(match x {
11396 None => v.monad(y, ctx, span)?,
11397 Some(x) => v.dyad(x, y, ctx, span)?,
11398 });
11399 }
11400 return assemble(&[vs.len()], cells, span);
11401 }
11402 if x.is_some() {
11403 return Err(Error::domain("m`:3 has no dyadic meaning", span));
11404 }
11405 let items = if y.rank() == 0 { vec![y.clone()] } else { y.cells(1) };
11406 let Some((last, rest)) = items.split_last() else {
11407 return Err(Error::domain("m`:3 needs an argument with items", span));
11408 };
11409 let mut acc = last.clone();
11410 for (i, item) in rest.iter().enumerate().rev() {
11411 acc = vs[i % vs.len()].dyad(item, &acc, ctx, span)?;
11412 }
11413 Ok(acc)
11414}
11415
11416fn stencil(u: &Verb, w: &[i64], y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
11421 if w.len() > y.rank() {
11422 return Err(Error::new(
11423 ErrorKind::Rank,
11424 format!("a stencil of {} axis/axes into a rank-{} value", w.len(), y.rank()),
11425 Some(span),
11426 ));
11427 }
11428 if w.iter().any(|&n| n <= 0) {
11429 return Err(Error::domain("a stencil window is a positive size", span));
11430 }
11431 let y = y.to_row_major();
11432 let k = w.len();
11433 let st = strides(&y.shape);
11434 let frame: Vec<usize> = y.shape[..k].to_vec();
11435 let mut wshape: Vec<usize> = w.iter().map(|&n| n as usize).collect();
11437 wshape.extend_from_slice(&y.shape[k..]);
11438 let inner: usize = y.shape[k..].iter().product();
11439 let total: usize = frame.iter().product();
11440 let mut cells = Vec::with_capacity(total);
11441 let mut at = vec![0usize; frame.len()];
11442 let mut coord = vec![0usize; k];
11443 for _ in 0..total {
11444 let mut data = Data::empty(y.dtype());
11445 coord.iter_mut().for_each(|c| *c = 0);
11446 let count: usize = w.iter().map(|&n| n as usize).product();
11447 for _ in 0..count {
11448 let mut base = 0usize;
11449 let mut inside = true;
11450 for a in 0..k {
11451 let off = at[a] as i64 + coord[a] as i64 - (w[a] - 1) / 2;
11452 if off < 0 || off >= y.shape[a] as i64 {
11453 inside = false;
11454 break;
11455 }
11456 base += off as usize * st[a];
11457 }
11458 for j in 0..inner {
11459 if inside {
11460 push_elem(&mut data, &y.data, base + j);
11461 } else {
11462 data.push_fill();
11463 }
11464 }
11465 odometer(&mut coord, &wshape[..k]);
11466 }
11467 cells.push(u.monad(&Array::new(wshape.clone(), data), ctx, span)?);
11468 odometer(&mut at, &frame);
11469 }
11470 assemble(&frame, cells, span)
11471}
11472
11473fn outfix(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
11477 let k = one_int(x, "an outfix width", span)?;
11478 let n = y.items() as i64;
11479 let list = as_list(y);
11480 let starts: Vec<i64> = if k < 0 {
11485 let step = -k;
11486 (0..(n + step - 1) / step).map(|i| i * step).collect()
11487 } else {
11488 (0..=(n - k)).collect()
11489 };
11490 let width = k.unsigned_abs() as usize;
11491 let mut cells = Vec::with_capacity(starts.len());
11492 for start in starts {
11493 let start = start as usize;
11494 let keep: Vec<usize> =
11495 (0..n as usize).filter(|&i| i < start || i >= start + width).collect();
11496 cells.push(u.monad(&select_items(&list, &keep), ctx, span)?);
11497 }
11498 assemble(&[cells.len()], cells, span)
11499}
11500
11501pub(crate) fn obverse(v: &Verb) -> Option<Verb> {
11513 let swap = |name: &'static str| -> Option<Verb> {
11514 crate::frontend::j::verb_named(name)
11515 };
11516 Some(match v {
11517 Verb::Prim(p) => {
11518 use ScalarMonad as SM;
11519 if matches!(
11523 p.monad,
11524 MonadOp::Scalar(SM::Conj | SM::Neg | SM::Recip | SM::OneMinus)
11525 | MonadOp::Reverse
11526 | MonadOp::TransposeAxes
11527 ) {
11528 return Some(v.clone());
11529 }
11530 if matches!(p.monad, MonadOp::Scalar(SM::Imaginary)) {
11533 return Some(Verb::Atop(Box::new(swap("-")?), Box::new(swap("j.")?)));
11534 }
11535 let by_monad: Option<&'static str> = match p.monad {
11536 MonadOp::Scalar(SM::Exp) => Some("^."),
11537 MonadOp::Scalar(SM::Ln) => Some("^"),
11538 MonadOp::Scalar(SM::Sqrt) => Some("*:"),
11539 MonadOp::Scalar(SM::Square) => Some("%:"),
11540 MonadOp::Scalar(SM::Double) => Some("-:"),
11541 MonadOp::Scalar(SM::Halve) => Some("+:"),
11542 MonadOp::Scalar(SM::Inc) => Some("<:"),
11543 MonadOp::Scalar(SM::Dec) => Some(">:"),
11544 MonadOp::Enclose(_) => Some(">"),
11545 MonadOp::Open => Some("<"),
11546 MonadOp::DecodeBits => Some("#:"),
11547 MonadOp::EncodeBits => Some("#."),
11548 _ => None,
11549 };
11550 swap(by_monad?)?
11551 }
11552 Verb::WithObverse(_, w) => (**w).clone(),
11554 Verb::Atop(f, g) => {
11556 Verb::Atop(Box::new(obverse(g)?), Box::new(obverse(f)?))
11557 }
11558 Verb::Compose(f, g) | Verb::Beside(f, g) => {
11559 Verb::Atop(Box::new(obverse(g)?), Box::new(obverse(f)?))
11560 }
11561 Verb::Rank(f, r) => Verb::Rank(Box::new(obverse(f)?), *r),
11562 Verb::Fit(f, n) => Verb::Fit(Box::new(obverse(f)?), *n),
11563 Verb::PowerN(f, Power::Times(n)) => {
11565 Verb::PowerN(Box::new(obverse(f)?), Power::Times(*n))
11566 }
11567 Verb::BondLeft(m, f) => bond_obverse(m, f, true)?,
11568 Verb::BondRight(f, n) => bond_obverse(n, f, false)?,
11569 _ => return None,
11570 })
11571}
11572
11573fn bond_obverse(n: &Array, f: &Verb, left: bool) -> Option<Verb> {
11577 let Verb::Prim(p) = f else { return None };
11578 let named = |name: &'static str| crate::frontend::j::verb_named(name);
11579 let bond = |name: &'static str, arg: &Array| -> Option<Verb> {
11580 let g = named(name)?;
11581 Some(if left {
11582 Verb::BondLeft(arg.clone(), Box::new(g))
11583 } else {
11584 Verb::BondRight(Box::new(g), arg.clone())
11585 })
11586 };
11587 use ScalarDyad as SD;
11588 let DyadOp::Scalar(op) = p.dyad else { return None };
11589 match (op, left) {
11590 (SD::Sub | SD::DivJ | SD::DivApl, true) => bond(p.name, n),
11592 (SD::Add, _) => Some(Verb::BondRight(Box::new(named("-")?), n.clone())),
11596 (SD::Mul, _) => Some(Verb::BondRight(Box::new(named("%")?), n.clone())),
11597 (SD::Sub, false) => bond("+", n),
11598 (SD::DivJ | SD::DivApl, false) => bond("*", n),
11599 (SD::Pow, false) => Some(Verb::BondLeft(n.clone(), Box::new(named("%:")?))),
11601 (SD::Pow, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^.")?))),
11602 (SD::Log, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^")?))),
11603 (SD::Root, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^")?))),
11604 _ => None,
11605 }
11606}
11607
11608fn self_classify(y: &Array, tol: Tol) -> Array {
11613 let items = if y.rank() == 0 { 1 } else { y.items() };
11614 let keys = nub(&as_list(y), tol);
11615 let rows = keys.items();
11616 let mut out = Vec::with_capacity(rows * items);
11617 for i in 0..rows {
11618 let key = item_or_self(&keys, i);
11619 for j in 0..items {
11620 out.push(arrays_match(&key, &item_or_self(y, j), tol) as u8);
11621 }
11622 }
11623 Array::new(vec![rows, items], Data::Bool(out.into()))
11624}
11625
11626fn nub_sieve(y: &Array, tol: Tol) -> Array {
11628 let items = if y.rank() == 0 { 1 } else { y.items() };
11629 let mut seen: Vec<Array> = Vec::new();
11630 let mut out = Vec::with_capacity(items);
11631 for i in 0..items {
11632 let cell = item_or_self(y, i);
11633 let fresh = !seen.iter().any(|s| arrays_match(s, &cell, tol));
11634 if fresh {
11635 seen.push(cell);
11636 }
11637 out.push(fresh as u8);
11638 }
11639 Array::new(vec![items], Data::Bool(out.into()))
11640}
11641
11642fn as_list(y: &Array) -> Array {
11644 if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() }
11645}
11646
11647fn conforming_cells(y: &Array, item_rank: usize) -> Vec<Array> {
11652 let frame_rank = y.rank().saturating_sub(item_rank);
11653 let nf: usize = y.shape[..frame_rank].iter().product();
11654 (0..nf).map(|i| y.cell_at(frame_rank, i)).collect()
11655}
11656
11657fn item_marks(y: &Array, x: &Array, tol: Tol) -> Vec<bool> {
11659 let n = if y.rank() == 0 { 1 } else { y.items() };
11660 let item_rank = y.rank().saturating_sub(1);
11661 let against = conforming_cells(x, item_rank);
11662 (0..n)
11663 .map(|i| {
11664 let cell = item_or_self(y, i);
11665 against.iter().any(|c| arrays_match(&cell, c, tol))
11666 })
11667 .collect()
11668}
11669
11670fn set_less(x: &Array, y: &Array, tol: Tol) -> Array {
11672 let xs = as_list(x);
11673 let marks = item_marks(&xs, y, tol);
11674 let keep: Vec<usize> = (0..marks.len()).filter(|&i| !marks[i]).collect();
11675 select_items(&xs, &keep)
11676}
11677
11678fn set_rank(cfg: EvalCfg, what: &str, x: &Array, y: &Array, span: Span) -> Result<()> {
11682 if cfg.rules.lang == crate::Lang::Apl && (x.rank() > 1 || y.rank() > 1) {
11683 return Err(Error::new(
11684 ErrorKind::Rank,
11685 format!("{what} takes vectors, not rank {} and rank {}", x.rank(), y.rank()),
11686 Some(span),
11687 ));
11688 }
11689 Ok(())
11690}
11691
11692fn intersect_items(x: &Array, y: &Array, tol: Tol) -> Array {
11694 let xs = as_list(x);
11695 let marks = item_marks(&xs, y, tol);
11696 let keep: Vec<usize> = (0..marks.len()).filter(|&i| marks[i]).collect();
11697 select_items(&xs, &keep)
11698}
11699
11700fn union_items(x: &Array, y: &Array, tol: Tol, span: Span) -> Result<Array> {
11703 let xs = as_list(x);
11704 let ys = as_list(y);
11705 let marks = item_marks(&ys, &xs, tol);
11706 let mut extra: Vec<usize> = Vec::new();
11707 for (i, &seen) in marks.iter().enumerate() {
11708 if seen {
11709 continue;
11710 }
11711 let cell = item_or_self(&ys, i);
11712 if !extra.iter().any(|&j| arrays_match(&item_or_self(&ys, j), &cell, tol)) {
11713 extra.push(i);
11714 }
11715 }
11716 catenate(&xs, &select_items(&ys, &extra), true, false, span)
11717}
11718
11719fn find_seq(x: &Array, y: &Array, tol: Tol, apl: bool, span: Span) -> Result<Array> {
11729 let (xr, yr) = (x.rank(), y.rank());
11730 if apl && xr > yr {
11731 return Ok(Array::new(y.shape.clone(), Data::Bool(vec![0u8; y.count()].into())));
11733 }
11734 if !apl && xr.max(1) != yr {
11735 return Err(Error::new(
11736 ErrorKind::Rank,
11737 format!("a rank-{xr} pattern in a rank-{yr} argument"),
11738 Some(span),
11739 ));
11740 }
11741 let mut pattern = vec![1usize; yr];
11742 pattern[yr - xr..].copy_from_slice(&x.shape);
11743 let n = y.count();
11744 let mut out = vec![0u8; n];
11745 let (xrm, yrm) = (x.to_row_major(), y.to_row_major());
11746 let yst = strides(&y.shape);
11747 let cells: usize = pattern.iter().product();
11748 let mut at = vec![0usize; yr];
11749 for slot in out.iter_mut() {
11750 if (0..yr).all(|a| at[a] + pattern[a] <= y.shape[a]) {
11751 let mut off = vec![0usize; yr];
11752 let mut hit = true;
11753 for k in 0..cells {
11754 let i: usize = (0..yr).map(|a| (at[a] + off[a]) * yst[a]).sum();
11755 if !arrays_match(&atom(&xrm, k), &atom(&yrm, i), tol) {
11756 hit = false;
11757 break;
11758 }
11759 odometer(&mut off, &pattern);
11760 }
11761 *slot = hit as u8;
11762 }
11763 odometer(&mut at, &y.shape);
11764 }
11765 Ok(Array::new(y.shape.clone(), Data::Bool(out.into())))
11766}
11767
11768fn bool_dyad(op: BoolDyad, x: &Array, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
11772 let bit = |a: &Array| -> Result<u8> {
11773 match a.to_i64_vec().as_deref() {
11774 Some([0]) => Ok(0),
11775 Some([1]) => Ok(1),
11776 _ => Err(Error::domain("this verb reads values of 0 or 1", span)),
11777 }
11778 };
11779 let _ = cfg;
11780 let (a, b) = (bit(x)?, bit(y)?);
11781 let v = match op {
11782 BoolDyad::Nor => u8::from(a == 0 && b == 0),
11783 BoolDyad::Nand => u8::from(a == 0 || b == 0),
11784 };
11785 Ok(Array::new(vec![], Data::Bool(vec![v].into())))
11786}
11787
11788fn item_ranks(y: &Array, rules: Rules, span: Span) -> Result<Vec<usize>> {
11794 check_gradable(y, rules, span)?;
11795 if !y.dtype().is_numeric() {
11796 return Err(Error::domain("an anagram index needs numbers", span));
11797 }
11798 let order = grade_order(&as_list(y), false, Tao::of(rules));
11799 let mut ranks = vec![0usize; order.len()];
11800 for (place, &i) in order.iter().enumerate() {
11801 ranks[i] = place;
11802 }
11803 Ok(ranks)
11804}
11805
11806fn anagram_index(y: &Array, rules: Rules, span: Span) -> Result<Array> {
11809 let ranks = item_ranks(y, rules, span)?;
11810 let n = ranks.len();
11811 let mut index: i128 = 0;
11812 for i in 0..n {
11813 let smaller = ranks[i + 1..].iter().filter(|&&r| r < ranks[i]).count() as i128;
11814 index = index
11815 .checked_mul((n - i) as i128)
11816 .and_then(|v| v.checked_add(smaller))
11817 .ok_or_else(|| Error::not_yet("an anagram index too large for an integer", span))?;
11818 }
11819 i64::try_from(index)
11820 .map(Array::scalar_i64)
11821 .map_err(|_| Error::not_yet("an anagram index too large for an integer", span))
11822}
11823
11824fn anagram_from(x: &Array, y: &Array, span: Span) -> Result<Array> {
11827 let idx = x
11828 .to_i64_vec()
11829 .ok_or_else(|| Error::domain("an anagram index must be an integer", span))?;
11830 let Some(&want) = idx.first() else {
11831 return Err(Error::internal("anagram with no index"));
11832 };
11833 let ys = as_list(y);
11834 let n = ys.items();
11835 let mut total: i128 = 1;
11836 for k in 1..=n as i128 {
11837 total = total
11838 .checked_mul(k)
11839 .ok_or_else(|| Error::not_yet("permuting more items than an integer counts", span))?;
11840 }
11841 let mut at = want as i128;
11842 if at < 0 {
11843 at += total;
11844 }
11845 if at < 0 || at >= total {
11846 return Err(Error::domain(
11847 format!("permutation {want} is out of range: {n} items have {total} of them"),
11848 span,
11849 ));
11850 }
11851 let mut pool: Vec<usize> = (0..n).collect();
11854 let mut order = Vec::with_capacity(n);
11855 let mut fact = total;
11856 for i in 0..n {
11857 fact /= (n - i) as i128;
11858 let d = (at / fact) as usize;
11859 at %= fact;
11860 order.push(pool.remove(d));
11861 }
11862 Ok(select_items(&ys, &order))
11863}
11864
11865fn cycle_form(y: &Array, span: Span) -> Result<Array> {
11869 if y.dtype() == DType::Box {
11870 let perm = cycles_to_direct(y, span)?;
11871 return Ok(Array::from_i64(perm.iter().map(|&i| i as i64).collect()));
11872 }
11873 let perm = direct_permutation(y, span)?;
11874 let mut boxes: Vec<Array> = Vec::new();
11875 let mut done = vec![false; perm.len()];
11876 for start in 0..perm.len() {
11877 if done[start] {
11878 continue;
11879 }
11880 let mut cycle = Vec::new();
11881 let mut at = start;
11882 while !done[at] {
11883 done[at] = true;
11884 cycle.push(at);
11885 at = perm[at];
11886 }
11887 let top = cycle.iter().position(|&v| v == *cycle.iter().max().unwrap()).unwrap();
11890 cycle.rotate_left(top);
11891 boxes.push(Array::boxed(Array::from_i64(
11892 cycle.iter().map(|&i| i as i64).collect(),
11893 )));
11894 }
11895 boxes.sort_by_key(|b| b.as_boxes().map(|s| s[0].to_i64_vec().unwrap()[0]).unwrap_or(0));
11896 let n = boxes.len();
11897 let inner: Vec<Array> =
11898 boxes.into_iter().map(|b| b.as_boxes().unwrap()[0].clone()).collect();
11899 Ok(Array::new(vec![n], Data::Box(inner.into())))
11900}
11901
11902fn direct_permutation(y: &Array, span: Span) -> Result<Vec<usize>> {
11904 let v = y
11905 .to_i64_vec()
11906 .ok_or_else(|| Error::domain("a permutation is a list of integers", span))?;
11907 let n = v.len();
11908 let mut seen = vec![false; n];
11909 let mut out = Vec::with_capacity(n);
11910 for &i in &v {
11911 let k = usize::try_from(i).ok().filter(|&k| k < n && !seen[k]).ok_or_else(|| {
11912 Error::domain(format!("{i} does not belong to a permutation of {n} items"), span)
11913 })?;
11914 seen[k] = true;
11915 out.push(k);
11916 }
11917 Ok(out)
11918}
11919
11920fn cycles_to_direct(y: &Array, span: Span) -> Result<Vec<usize>> {
11924 let boxes = y.as_boxes().ok_or_else(|| Error::internal("cycles from a simple array"))?;
11925 let mut cycles: Vec<Vec<usize>> = Vec::new();
11926 let mut top = 0usize;
11927 for b in boxes {
11928 let v = b
11929 .to_i64_vec()
11930 .ok_or_else(|| Error::domain("a cycle is a list of integers", span))?;
11931 let mut cycle = Vec::with_capacity(v.len());
11932 for &i in &v {
11933 let k = usize::try_from(i)
11934 .map_err(|_| Error::domain(format!("{i} is not an index"), span))?;
11935 top = top.max(k + 1);
11936 cycle.push(k);
11937 }
11938 cycles.push(cycle);
11939 }
11940 let mut perm: Vec<usize> = (0..top).collect();
11941 for cycle in &cycles {
11942 for w in 0..cycle.len() {
11943 perm[cycle[w]] = cycle[(w + 1) % cycle.len()];
11945 }
11946 }
11947 Ok(perm)
11948}
11949
11950fn permute(x: &Array, y: &Array, span: Span) -> Result<Array> {
11955 let ys = as_list(y);
11956 let n = ys.items();
11957 let cyclic = x.dtype() == DType::Box;
11958 if !cyclic && x.rank() == 0 {
11959 return Err(Error::not_yet("permuting by a single atom (x C. y)", span));
11960 }
11961 let mut perm =
11962 if cyclic { cycles_to_direct(x, span)? } else { direct_permutation(&as_list(x), span)? };
11963 if perm.len() > n {
11964 return Err(Error::new(
11965 ErrorKind::Length,
11966 format!("a permutation of {} items applied to {n}", perm.len()),
11967 Some(span),
11968 ));
11969 }
11970 if perm.len() < n {
11971 if cyclic {
11972 perm.extend(perm.len()..n);
11974 } else {
11975 let head: Vec<usize> = (perm.len()..n).collect();
11978 perm.splice(0..0, head);
11979 }
11980 }
11981 Ok(select_items(&ys, &perm))
11982}
11983
11984fn unicode(y: &Array, pass_chars: bool, span: Span) -> Result<Array> {
11990 if y.dtype() == DType::Char {
11991 if pass_chars {
11992 return Ok(y.clone());
11993 }
11994 return Ok(chars_to_codes(y));
11995 }
11996 codes_to_chars(y, span)
11997}
11998
11999fn chars_to_codes(y: &Array) -> Array {
12000 let Data::Char(v) = &y.data else { return y.clone() };
12001 Array::new(y.shape.clone(), Data::I64(v.iter().map(|&c| c as i64).collect()))
12002}
12003
12004fn codes_to_chars(y: &Array, span: Span) -> Result<Array> {
12005 let v = y
12006 .to_i64_vec()
12007 .ok_or_else(|| Error::domain("a codepoint must be an integer", span))?;
12008 let mut out = Vec::with_capacity(v.len());
12009 for &c in &v {
12010 let ch = u32::try_from(c).ok().and_then(char::from_u32).ok_or_else(|| {
12011 Error::domain(format!("{c} is not a Unicode codepoint"), span)
12012 })?;
12013 out.push(ch);
12014 }
12015 Ok(Array::new(y.shape.clone(), Data::Char(out.into())))
12016}
12017
12018fn unicode_form(x: &Array, y: &Array, span: Span) -> Result<Array> {
12021 let form = x
12022 .to_i64_vec()
12023 .ok_or_else(|| Error::domain("a conversion form is an integer", span))?
12024 .first()
12025 .copied()
12026 .unwrap_or(0);
12027 match form {
12028 3 if y.dtype() == DType::Char => Ok(chars_to_codes(y)),
12029 3 => Err(Error::domain("form 3 converts characters to codepoints", span)),
12030 10 => codes_to_chars(y, span),
12031 n => Err(Error::not_yet(format!("the byte-oriented unicode form ({n} u:)"), span)),
12032 }
12033}
12034
12035fn to_symbols(y: &Array, span: Span) -> Result<Array> {
12046 if let Some(boxes) = y.as_boxes() {
12047 let mut ids = Vec::with_capacity(boxes.len());
12048 for b in boxes {
12049 if b.rank() > 1 {
12050 return Err(Error::new(
12051 ErrorKind::Rank,
12052 "a boxed symbol name is a character list",
12053 Some(span),
12054 ));
12055 }
12056 let row_major = b.to_row_major();
12057 let Data::Char(v) = &row_major.data else {
12058 if b.count() == 0 {
12059 ids.push(crate::symbol::EMPTY);
12060 continue;
12061 }
12062 return Err(Error::domain("a symbol is made from characters", span));
12063 };
12064 ids.push(crate::symbol::intern(&v.as_slice().iter().collect::<String>()));
12065 }
12066 return Ok(Array::new(y.shape.clone(), Data::Symbol(ids.into())));
12067 }
12068 let row_major = y.to_row_major();
12069 let Data::Char(v) = &row_major.data else {
12070 return Err(Error::domain(
12071 format!("s: makes symbols from characters, not {} data", y.dtype().name()),
12072 span,
12073 ));
12074 };
12075 let chars = v.as_slice();
12076 if y.rank() >= 2 {
12077 let width = y.shape[y.rank() - 1];
12078 let mut ids = Vec::with_capacity(chars.len() / width.max(1));
12079 for row in chars.chunks(width) {
12080 let name: String = row.iter().collect();
12081 ids.push(crate::symbol::intern(name.trim_end_matches(' ')));
12082 }
12083 return Ok(Array::new(y.shape[..y.rank() - 1].to_vec(), Data::Symbol(ids.into())));
12084 }
12085 let Some((&delim, rest)) = chars.split_first() else {
12086 return Ok(Array::new(vec![0], Data::empty(DType::Symbol)));
12087 };
12088 let mut ids = Vec::new();
12089 let mut name = String::new();
12090 for &c in rest {
12091 if c == delim {
12092 ids.push(crate::symbol::intern(&name));
12093 name.clear();
12094 } else {
12095 name.push(c);
12096 }
12097 }
12098 ids.push(crate::symbol::intern(&name));
12099 Ok(Array::new(vec![ids.len()], Data::Symbol(ids.into())))
12100}
12101
12102fn symbol_form(x: &Array, y: &Array, span: Span) -> Result<Array> {
12108 let form = x
12109 .to_i64_vec()
12110 .ok_or_else(|| Error::domain("a symbol form is an integer", span))?
12111 .first()
12112 .copied()
12113 .unwrap_or(0);
12114 if !matches!(form, 4 | 5) {
12115 return Err(Error::not_yet(format!("the symbol-table form ({form} s:)"), span));
12116 }
12117 let row_major = y.to_row_major();
12118 let Data::Symbol(ids) = &row_major.data else {
12119 return Err(Error::domain(
12120 format!("{form} s: reads symbols, not {} data", y.dtype().name()),
12121 span,
12122 ));
12123 };
12124 let names = crate::symbol::names(ids.as_slice());
12125 if form == 5 {
12126 let boxes: Vec<Array> =
12127 names.iter().map(|n| Array::from_chars(n.chars().collect())).collect();
12128 return Ok(Array::new(y.shape.clone(), Data::Box(boxes.into())));
12129 }
12130 let width = names.iter().map(|n| n.chars().count()).max().unwrap_or(0);
12131 let mut out: Vec<char> = Vec::with_capacity(names.len() * width);
12132 for n in &names {
12133 out.extend(n.chars());
12134 out.resize(out.len() + width - n.chars().count(), ' ');
12135 }
12136 let mut shape = y.shape.clone();
12137 shape.push(width);
12138 Ok(Array::new(shape, Data::Char(out.into())))
12139}
12140
12141fn boxing_level(y: &Array) -> i64 {
12143 match y.as_boxes() {
12144 None => 0,
12145 Some(bs) => 1 + bs.iter().map(boxing_level).max().unwrap_or(0),
12146 }
12147}
12148
12149fn split_items(y: &Array) -> Array {
12153 if y.rank() == 0 {
12154 return Array::boxed(y.clone());
12155 }
12156 let last = y.shape[y.rank() - 1];
12157 let outer: Vec<usize> = y.shape[..y.rank() - 1].to_vec();
12158 let n: usize = outer.iter().product();
12159 let mut boxes = Vec::with_capacity(n);
12160 for i in 0..n {
12161 let mut data = Data::empty(y.dtype());
12162 for k in 0..last {
12163 push_elem(&mut data, &y.data, i * last + k);
12164 }
12165 boxes.push(Array::new(vec![last], data));
12166 }
12167 Array::new(outer, Data::Box(boxes.into()))
12168}
12169
12170fn pick(x: &Array, y: &Array, origin: i64, span: Span) -> Result<Array> {
12173 let xs = as_list(x);
12174 let mut cur = y.clone();
12175 for i in 0..xs.items() {
12176 let step = open_cell(&item_or_self(&xs, i));
12177 let idx = step
12178 .to_i64_vec()
12179 .ok_or_else(|| Error::domain("a pick path holds integers", span))?;
12180 let base =
12181 if cur.rank() == 0 { Array::new(vec![1], cur.data.clone()) } else { cur.clone() };
12182 if idx.len() > base.rank() {
12183 return Err(Error::new(
12184 ErrorKind::Length,
12185 format!(
12186 "a path step of {} index(es) into a value of rank {}",
12187 idx.len(),
12188 cur.rank()
12189 ),
12190 Some(span),
12191 ));
12192 }
12193 let zeroed: Vec<i64> = idx.iter().map(|&v| v - origin).collect();
12194 let at = cell_index(&base, &zeroed, span)?;
12195 cur = open_cell(&base.cell_at(idx.len(), at));
12196 }
12197 Ok(cur)
12198}
12199
12200fn prime_meta(x: &Array, y: &Array, span: Span) -> Result<Array> {
12205 let form = one_int(x, "a prime query", span)?;
12206 let n = one_int(y, "a prime query", span)?;
12207 match form {
12208 -1 => Ok(Array::scalar_i64(primes_below(n, span)?)),
12210 0 => Ok(Array::scalar_bool(!is_prime(n))),
12212 1 => Ok(Array::scalar_bool(is_prime(n))),
12213 2 | 3 => {
12215 let (ps, es) = factor_table(n, span)?;
12216 let k = ps.len();
12217 if form == 3 {
12218 return Ok(Array::from_i64(ps));
12219 }
12220 let mut all = ps;
12221 all.extend(es);
12222 Ok(Array::new(vec![2, k], Data::I64(all.into())))
12223 }
12224 4 => Ok(Array::scalar_i64(next_prime(n, span)?)),
12226 -4 => Ok(Array::scalar_i64(previous_prime(n, span)?)),
12227 other => Err(Error::domain(format!("{other} is not a prime query"), span)),
12228 }
12229}
12230
12231fn prime_exponents(x: &Array, y: &Array, span: Span) -> Result<Array> {
12234 let n = one_int(y, "prime exponents", span)?;
12235 let count = x.to_f64_vec().and_then(|v| v.first().copied()).unwrap_or(0.0);
12236 let (ps, es) = factor_table(n, span)?;
12237 if count == f64::NEG_INFINITY {
12238 let k = ps.len();
12239 let mut all = ps;
12240 all.extend(es);
12241 return Ok(Array::new(vec![2, k], Data::I64(all.into())));
12242 }
12243 let want = one_int(x, "prime exponents", span)?;
12244 if want < 0 {
12245 return Err(Error::not_yet(format!("the prime exponent form ({want} q:)"), span));
12246 }
12247 let mut out = Vec::with_capacity(want as usize);
12248 for i in 0..want {
12249 let p = nth_prime(i, span)?;
12250 out.push(ps.iter().position(|&q| q == p).map_or(0, |at| es[at]));
12251 }
12252 Ok(Array::from_i64(out))
12253}
12254
12255fn factor_table(n: i64, span: Span) -> Result<(Vec<i64>, Vec<i64>)> {
12257 let factors = prime_factors(n, span)?;
12258 let mut ps: Vec<i64> = Vec::new();
12259 let mut es: Vec<i64> = Vec::new();
12260 for f in factors {
12261 if ps.last() == Some(&f) {
12262 *es.last_mut().unwrap() += 1;
12263 } else {
12264 ps.push(f);
12265 es.push(1);
12266 }
12267 }
12268 Ok((ps, es))
12269}
12270
12271fn is_prime(n: i64) -> bool {
12272 if n < 2 {
12273 return false;
12274 }
12275 let mut d = 2i64;
12276 while d.saturating_mul(d) <= n {
12277 if n % d == 0 {
12278 return false;
12279 }
12280 d += 1;
12281 }
12282 true
12283}
12284
12285fn primes_below(n: i64, span: Span) -> Result<i64> {
12286 if n < 0 {
12287 return Err(Error::domain("counting the primes below a negative number", span));
12288 }
12289 Ok((2..n).filter(|&k| is_prime(k)).count() as i64)
12290}
12291
12292fn next_prime(n: i64, span: Span) -> Result<i64> {
12293 let mut k = n.checked_add(1).ok_or_else(|| Error::domain("no next prime", span))?;
12294 while !is_prime(k) {
12295 k = k.checked_add(1).ok_or_else(|| Error::domain("no next prime", span))?;
12296 }
12297 Ok(k)
12298}
12299
12300fn previous_prime(n: i64, span: Span) -> Result<i64> {
12301 let mut k = n - 1;
12302 while k >= 2 {
12303 if is_prime(k) {
12304 return Ok(k);
12305 }
12306 k -= 1;
12307 }
12308 Err(Error::domain(format!("there is no prime below {n}"), span))
12309}
12310
12311fn one_int(a: &Array, what: &str, span: Span) -> Result<i64> {
12313 a.to_i64_vec()
12314 .and_then(|v| v.first().copied())
12315 .ok_or_else(|| Error::domain(format!("{what} needs an integer"), span))
12316}
12317
12318fn expand(x: &Array, y: &Array, span: Span) -> Result<Array> {
12321 let mask = x
12322 .to_i64_vec()
12323 .ok_or_else(|| Error::domain("an expansion mask holds 0s and 1s", span))?;
12324 if mask.iter().any(|&b| b != 0 && b != 1) {
12325 return Err(Error::domain("an expansion mask holds 0s and 1s", span));
12326 }
12327 let ys = as_list(y);
12328 let taken = mask.iter().filter(|&&b| b == 1).count();
12329 let n = ys.items();
12330 let spread = n == 1 && taken != 1;
12332 if !spread && taken != n {
12333 return Err(Error::new(
12334 ErrorKind::Length,
12335 format!("an expansion mask taking {taken} item(s) over {n}"),
12336 Some(span),
12337 ));
12338 }
12339 let m = ys.item_size();
12340 let mut data = Data::empty(ys.dtype());
12341 let mut at = 0usize;
12342 for &b in &mask {
12343 if b == 1 {
12344 let from = if spread { 0 } else { at };
12345 for k in 0..m {
12346 push_elem(&mut data, &ys.data, from * m + k);
12347 }
12348 at += 1;
12349 } else {
12350 for _ in 0..m {
12351 data.push_fill();
12352 }
12353 }
12354 }
12355 let mut shape = ys.shape.clone();
12356 if shape.is_empty() {
12357 shape.push(mask.len());
12358 } else {
12359 shape[0] = mask.len();
12360 }
12361 Ok(Array::new(shape, data))
12362}
12363
12364fn execute(y: &Array, apl: bool, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
12372 let Data::Char(v) = &y.data else {
12373 return Err(Error::domain("execute reads a character list", span));
12374 };
12375 let src: String = v.iter().collect();
12376 execute_source(&src, apl, ctx, span)
12377}
12378
12379pub(crate) fn execute_source(
12382 src: &str,
12383 apl: bool,
12384 ctx: &mut Ctx<'_>,
12385 span: Span,
12386) -> Result<Array> {
12387 let lang = if apl { crate::Lang::Apl } else { crate::Lang::J };
12388 let dialect = ctx.cfg.rules.dialect();
12391 let nested = crate::compile(lang, src, &dialect).map_err(|e| nested_error(e, src, span))?;
12392 if !nested.params.is_empty() {
12393 return Err(Error::domain(
12394 "an executed string cannot take host data: `{name}` has nothing to bind to",
12395 span,
12396 ));
12397 }
12398 let mut rec = None;
12399 let (value, _) = crate::ir::run_block(&nested.stmts, None, ctx, &mut rec)
12400 .map_err(|e| nested_error(e, src, span))?;
12401 value.ok_or_else(|| Error::domain("the executed string yielded no value", span))
12402}
12403
12404fn stream_number(y: &Array, open: i64, what: &str, span: Span) -> Result<()> {
12411 let closed = || {
12412 Err(Error::sandbox(
12413 format!("{what} the standard stream {open} only; a file is outside the program"),
12414 span,
12415 ))
12416 };
12417 if matches!(y.data, Data::Box(_)) {
12418 return closed();
12419 }
12420 match y.to_i64_vec().as_deref() {
12421 Some([n]) if *n == open => Ok(()),
12422 Some([_]) => closed(),
12423 _ => Err(Error::domain(format!("{what} one stream number"), span)),
12424 }
12425}
12426
12427fn type_code(y: &Array) -> i64 {
12430 match y.dtype() {
12431 DType::Bool => 1,
12432 DType::Char => 2,
12433 DType::I64 => 4,
12434 DType::F64 => 8,
12435 DType::Complex => 16,
12436 DType::Box => 32,
12437 DType::Ext => 64,
12438 DType::Rat => 128,
12439 DType::Symbol => 65536,
12440 }
12441}
12442
12443fn nested_error(e: Error, src: &str, span: Span) -> Error {
12447 let inner = e.render(src);
12448 let mut out = Error::new(e.kind, format!("in the executed string: {}", e.msg), Some(span));
12449 out.notes.push(inner.trim_end().to_string());
12450 out
12451}
12452
12453fn words(y: &Array, span: Span) -> Result<Array> {
12459 let Data::Char(v) = &y.data else {
12460 return Err(Error::domain("words reads a character list", span));
12461 };
12462 let src: Vec<char> = v.as_slice().to_vec();
12463 let n = src.len();
12464 let mut out: Vec<Array> = Vec::new();
12465 let mut i = 0usize;
12466 let numeric_start = |k: usize| -> bool {
12467 k < n && (src[k].is_ascii_digit() || src[k] == '_')
12468 };
12469 while i < n {
12470 let c = src[i];
12471 if c == ' ' || c == '\t' {
12472 i += 1;
12473 continue;
12474 }
12475 let start = i;
12476 if c == '\'' {
12477 i += 1;
12478 loop {
12479 if i >= n {
12480 return Err(Error::parse("a word list ends inside a string", span));
12481 }
12482 if src[i] == '\'' {
12483 i += 1;
12484 if i < n && src[i] == '\'' {
12485 i += 1;
12486 continue;
12487 }
12488 break;
12489 }
12490 i += 1;
12491 }
12492 } else if c.is_ascii_alphabetic() {
12493 while i < n && (src[i].is_ascii_alphanumeric() || src[i] == '_') {
12494 i += 1;
12495 }
12496 if i < n && (src[i] == '.' || src[i] == ':') {
12497 i += 1;
12498 }
12499 if src[start..i].iter().collect::<String>() == "NB." {
12501 while i < n && src[i] != '\n' {
12502 i += 1;
12503 }
12504 }
12505 } else if numeric_start(i) {
12506 loop {
12507 while i < n && (src[i].is_ascii_alphanumeric() || src[i] == '.' || src[i] == '_')
12508 {
12509 i += 1;
12510 }
12511 let mut j = i;
12513 while j < n && src[j] == ' ' {
12514 j += 1;
12515 }
12516 if j > i && numeric_start(j) {
12517 i = j;
12518 continue;
12519 }
12520 break;
12521 }
12522 } else {
12523 i += 1;
12524 while i < n && (src[i] == '.' || src[i] == ':') {
12525 i += 1;
12526 }
12527 }
12528 out.push(Array::from_chars(src[start..i].to_vec()));
12529 }
12530 let k = out.len();
12531 Ok(Array::new(vec![k], Data::Box(out.into())))
12532}
12533
12534#[cfg(test)]
12535mod tests {
12536 use super::*;
12537
12538 macro_rules! ctx {
12540 ($name:ident, $agreement:expr) => {
12541 let mut sink = |_: &str| {};
12542 let mut env = Env::new(Vec::new());
12543 #[allow(unused_mut)]
12544 let mut $name = Ctx {
12545 cfg: EvalCfg {
12546 agreement: $agreement,
12547 fmt: FmtOpts::J,
12548 tol: Tol::J,
12549 rules: crate::frontend::Dialect::default()
12552 .rules(if $agreement == Agreement::ExactOrScalar {
12553 crate::Lang::Apl
12554 } else {
12555 crate::Lang::J
12556 })
12557 .expect("the shipped dialect is implemented"),
12558 },
12559 out: &mut sink,
12560 inp: None,
12561 env: &mut env,
12562 device: None,
12563 };
12564 };
12565 ($name:ident) => {
12566 ctx!($name, Agreement::LeadingPrefix);
12567 };
12568 }
12569
12570 fn scalar_prim(name: &'static str, monad: MonadOp, dyad: DyadOp) -> Verb {
12571 Verb::Prim(Prim { name, monad, dyad, ranks: [0, 0, 0] })
12572 }
12573
12574 fn inf_prim(name: &'static str, monad: MonadOp, dyad: DyadOp) -> Verb {
12575 Verb::Prim(Prim { name, monad, dyad, ranks: [RANK_INF, RANK_INF, RANK_INF] })
12576 }
12577
12578 fn plus() -> Verb {
12579 scalar_prim("+", MonadOp::Scalar(ScalarMonad::Conj), DyadOp::Scalar(ScalarDyad::Add))
12580 }
12581 fn minus() -> Verb {
12582 scalar_prim("-", MonadOp::Scalar(ScalarMonad::Neg), DyadOp::Scalar(ScalarDyad::Sub))
12583 }
12584 fn times() -> Verb {
12585 scalar_prim("*", MonadOp::Scalar(ScalarMonad::Signum), DyadOp::Scalar(ScalarDyad::Mul))
12586 }
12587 fn pct() -> Verb {
12588 scalar_prim("%", MonadOp::Scalar(ScalarMonad::Recip), DyadOp::Scalar(ScalarDyad::DivJ))
12589 }
12590 fn div_apl() -> Verb {
12591 scalar_prim("÷", MonadOp::Scalar(ScalarMonad::Recip), DyadOp::Scalar(ScalarDyad::DivApl))
12592 }
12593 fn floor_v() -> Verb {
12594 scalar_prim("<.", MonadOp::Scalar(ScalarMonad::Floor), DyadOp::Scalar(ScalarDyad::Min))
12595 }
12596 fn ceil_v() -> Verb {
12597 scalar_prim(">.", MonadOp::Scalar(ScalarMonad::Ceil), DyadOp::Scalar(ScalarDyad::Max))
12598 }
12599 fn pow_v() -> Verb {
12600 scalar_prim("^", MonadOp::Scalar(ScalarMonad::Exp), DyadOp::Scalar(ScalarDyad::Pow))
12601 }
12602 fn residue_v() -> Verb {
12603 scalar_prim("|", MonadOp::Scalar(ScalarMonad::Abs), DyadOp::Scalar(ScalarDyad::Residue))
12604 }
12605 fn eq_v() -> Verb {
12606 scalar_prim("=", MonadOp::None, DyadOp::Scalar(ScalarDyad::Eq))
12607 }
12608 fn lt_v() -> Verb {
12609 scalar_prim("<", MonadOp::None, DyadOp::Scalar(ScalarDyad::Lt))
12610 }
12611 fn not_v() -> Verb {
12612 scalar_prim("-.", MonadOp::Scalar(ScalarMonad::Not), DyadOp::None)
12613 }
12614 fn sqrt_v() -> Verb {
12615 scalar_prim("%:", MonadOp::Scalar(ScalarMonad::Sqrt), DyadOp::NotYet("dyadic root"))
12616 }
12617 fn dollar() -> Verb {
12618 inf_prim("$", MonadOp::ShapeOf, DyadOp::Reshape)
12619 }
12620 fn pound() -> Verb {
12621 inf_prim("#", MonadOp::Tally, DyadOp::NotYet("copy"))
12622 }
12623 fn comma() -> Verb {
12624 inf_prim(",", MonadOp::Ravel, DyadOp::NotYet("append"))
12625 }
12626 fn transpose_v() -> Verb {
12627 inf_prim("|:", MonadOp::TransposeAxes, DyadOp::NotYet("dyadic transpose"))
12628 }
12629 fn head_v() -> Verb {
12630 inf_prim("{.", MonadOp::Head, DyadOp::Take)
12631 }
12632 fn behead_v() -> Verb {
12633 inf_prim("}.", MonadOp::Behead, DyadOp::Drop)
12634 }
12635 fn iota() -> Verb {
12636 inf_prim("i.", MonadOp::IotaJ, DyadOp::NotYet("index of"))
12637 }
12638 fn iota_apl(origin: i64) -> Verb {
12639 inf_prim("⍳", MonadOp::IotaApl { origin }, DyadOp::NotYet("index of"))
12640 }
12641 fn right_v() -> Verb {
12642 inf_prim("]", MonadOp::Same, DyadOp::Right)
12643 }
12644 fn echo_v() -> Verb {
12645 inf_prim("echo", MonadOp::Echo, DyadOp::None)
12646 }
12647
12648 fn b(v: Verb) -> Box<Verb> {
12649 Box::new(v)
12650 }
12651
12652 fn mat(rows: usize, cols: usize, v: Vec<i64>) -> Array {
12653 Array::new(vec![rows, cols], Data::I64(v.into()))
12654 }
12655
12656 fn ints(a: &Array) -> Vec<i64> {
12658 a.to_row_major().as_i64_slice().expect("integer result").to_vec()
12659 }
12660
12661 fn floats(a: &Array) -> Vec<f64> {
12662 a.to_row_major().as_f64_slice().expect("float result").to_vec()
12663 }
12664
12665 fn bools(a: &Array) -> Vec<u8> {
12666 match &a.to_row_major().data {
12667 Data::Bool(v) => v.to_vec(),
12668 other => panic!("expected boolean result, got {other:?}"),
12669 }
12670 }
12671
12672 fn sp() -> Span {
12673 Span::new(0, 1)
12674 }
12675
12676 fn close(a: f64, b: f64) -> bool {
12677 (a - b).abs() < 1e-9 || (a.is_infinite() && b.is_infinite() && a.signum() == b.signum())
12678 }
12679
12680 #[test]
12683 fn names_of_primitives_and_derived_verbs() {
12684 assert_eq!(plus().name(), "+");
12685 assert_eq!(Verb::Rank(b(plus()), [1, 1, 1]).name(), "+\"1");
12686 assert_eq!(Verb::Rank(b(plus()), [0, 1, RANK_INF]).name(), "+\"0 1 _");
12687 assert_eq!(Verb::Rank(b(plus()), [RANK_INF; 3]).name(), "+\"_");
12688 assert_eq!(Verb::Reduce(b(plus())).name(), "+/");
12689 assert_eq!(Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]).name(), "+/\"1");
12690 assert_eq!(Verb::Fork(b(plus()), b(minus()), b(times())).name(), "(+ - *)");
12691 assert_eq!(
12692 Verb::NounFork(Array::scalar_i64(1), b(plus()), b(minus())).name(),
12693 "(n + -)"
12694 );
12695 assert_eq!(Verb::Hook(b(plus()), b(minus())).name(), "(+ -)");
12696 assert_eq!(Verb::Atop(b(plus()), b(minus())).name(), "(+@:-)");
12697 assert_eq!(Verb::Compose(b(plus()), b(minus())).name(), "(+&:-)");
12698 assert_eq!(Verb::BondLeft(Array::scalar_i64(1), b(plus())).name(), "(n&+)");
12699 assert_eq!(Verb::BondRight(b(plus()), Array::scalar_i64(1)).name(), "(+&n)");
12700 }
12701
12702 #[test]
12703 fn composition_applies_the_right_verb_to_both_arguments() {
12704 ctx!(c);
12705 let v = Verb::Compose(b(plus()), b(times()));
12706 let r = v.monad(&Array::from_i64(vec![-2, 0, 3]), &mut c, sp()).unwrap();
12708 assert_eq!(ints(&r), vec![-1, 0, 1]);
12709 let r = v
12710 .dyad(&Array::scalar_i64(-5), &Array::scalar_i64(7), &mut c, sp())
12711 .unwrap();
12712 assert_eq!(ints(&r), vec![0]);
12713 let bond = Verb::BondLeft(Array::scalar_i64(10), b(minus()));
12715 let r = bond.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
12716 assert_eq!(ints(&r), vec![9, 8]);
12717 let e = bond
12718 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(2), &mut c, sp())
12719 .unwrap_err();
12720 assert_eq!(e.kind, ErrorKind::Domain);
12721 let bond = Verb::BondRight(b(minus()), Array::scalar_i64(10));
12722 let r = bond.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
12723 assert_eq!(ints(&r), vec![-9, -8]);
12724 }
12725
12726 #[test]
12729 fn scalar_monad_covers_the_whole_buffer() {
12730 ctx!(c);
12731 let r = minus().monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
12732 assert_eq!(r.shape, vec![2, 3]);
12733 assert_eq!(ints(&r), vec![-1, -2, -3, -4, -5, -6]);
12734 }
12735
12736 #[test]
12737 fn leading_prefix_agreement_broadcasts_per_row() {
12738 ctx!(c);
12739 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
12740 let y = Array::from_i64(vec![10, 20]);
12741 let r = plus().dyad(&x, &y, &mut c, sp()).unwrap();
12742 assert_eq!(r.shape, vec![2, 3]);
12743 assert_eq!(ints(&r), vec![11, 12, 13, 24, 25, 26]);
12744 let r = plus().dyad(&y, &x, &mut c, sp()).unwrap();
12746 assert_eq!(ints(&r), vec![11, 12, 13, 24, 25, 26]);
12747 }
12748
12749 #[test]
12750 fn exact_or_scalar_rejects_a_prefix_frame() {
12751 ctx!(c, Agreement::ExactOrScalar);
12752 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
12753 let y = Array::from_i64(vec![10, 20]);
12754 let e = plus().dyad(&x, &y, &mut c, sp()).unwrap_err();
12755 assert_eq!(e.kind, ErrorKind::Shape);
12756 assert!(e.msg.contains("2 3"), "{}", e.msg);
12757 assert!(e.msg.contains("right shape 2"), "{}", e.msg);
12758 }
12759
12760 #[test]
12761 fn exact_or_scalar_accepts_equal_frames_and_scalars() {
12762 ctx!(c, Agreement::ExactOrScalar);
12763 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
12764 let r = plus().dyad(&x, &x, &mut c, sp()).unwrap();
12765 assert_eq!(ints(&r), vec![2, 4, 6, 8, 10, 12]);
12766 let r = plus().dyad(&Array::scalar_i64(10), &x, &mut c, sp()).unwrap();
12767 assert_eq!(r.shape, vec![2, 3]);
12768 assert_eq!(ints(&r), vec![11, 12, 13, 14, 15, 16]);
12769 let r = plus().dyad(&x, &Array::scalar_i64(10), &mut c, sp()).unwrap();
12770 assert_eq!(ints(&r), vec![11, 12, 13, 14, 15, 16]);
12771 }
12772
12773 #[test]
12774 fn vector_length_mismatch_is_a_length_error() {
12775 ctx!(c);
12776 let e = plus()
12777 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::from_i64(vec![1, 2, 3, 4, 5]), &mut c, sp())
12778 .unwrap_err();
12779 assert_eq!(e.kind, ErrorKind::Length);
12780 assert!(e.msg.contains("left shape 3"), "{}", e.msg);
12781 assert!(e.msg.contains("right shape 5"), "{}", e.msg);
12782 assert!(e.notes[0].contains("axis 0"), "{:?}", e.notes);
12783 }
12784
12785 #[test]
12786 fn diverging_matrix_frames_name_the_axis() {
12787 ctx!(c);
12788 let e = plus()
12789 .dyad(&mat(2, 3, vec![0; 6]), &mat(2, 4, vec![0; 8]), &mut c, sp())
12790 .unwrap_err();
12791 assert_eq!(e.kind, ErrorKind::Shape);
12792 assert!(e.notes[0].contains("axis 1"), "{:?}", e.notes);
12793 }
12794
12795 #[test]
12796 fn dyadic_rank_pairs_rows_with_the_whole_right_argument() {
12797 ctx!(c);
12798 let v = Verb::Rank(b(plus()), [0, 1, 1]);
12800 let r = v
12801 .dyad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &Array::from_i64(vec![10, 20, 30]), &mut c, sp())
12802 .unwrap();
12803 assert_eq!(r.shape, vec![2, 3]);
12804 assert_eq!(ints(&r), vec![11, 22, 33, 14, 25, 36]);
12805 }
12806
12807 #[test]
12808 fn surplus_frame_axes_repeat_the_shorter_frames_cells() {
12809 ctx!(c);
12810 let v = Verb::Rank(b(head_v()), [0, 0, 1]);
12813 let x = mat(2, 2, vec![1, 1, 2, 2]);
12814 let y = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
12815 let r = v.dyad(&x, &y, &mut c, sp()).unwrap();
12816 assert_eq!(r.shape, vec![2, 2, 2]);
12817 assert_eq!(ints(&r), vec![1, 0, 1, 0, 4, 5, 4, 5]);
12818 }
12819
12820 #[test]
12821 fn an_empty_frame_pairs_its_single_cell_with_every_other_cell() {
12822 ctx!(c, Agreement::ExactOrScalar);
12823 let v = Verb::Rank(b(head_v()), [0, 0, 1]);
12825 let x = Array::from_i64(vec![1, 2]);
12826 let y = Array::from_i64(vec![7, 8, 9]);
12827 let r = v.dyad(&x, &y, &mut c, sp()).unwrap();
12828 assert_eq!(r.shape, vec![2, 2]);
12829 assert_eq!(ints(&r), vec![7, 0, 7, 8]);
12830 }
12831
12832 #[test]
12833 fn negative_rank_leaves_frame_axes() {
12834 ctx!(c);
12835 let v = Verb::Rank(b(dollar()), [-1, -1, -1]);
12837 let r = v.monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
12838 assert_eq!(r.shape, vec![2, 1]);
12839 assert_eq!(ints(&r), vec![3, 3]);
12840 }
12841
12842 #[test]
12843 fn effective_rank_clamps_and_counts_back() {
12844 assert_eq!(effective_rank(0, 3), 0);
12845 assert_eq!(effective_rank(2, 1), 1);
12846 assert_eq!(effective_rank(RANK_INF, 4), 4);
12847 assert_eq!(effective_rank(-1, 3), 2);
12848 assert_eq!(effective_rank(-5, 3), 0);
12849 }
12850
12851 #[test]
12854 fn reduction_folds_right_to_left() {
12855 ctx!(c);
12856 let r = Verb::Reduce(b(minus()))
12858 .monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp())
12859 .unwrap();
12860 assert!(r.shape.is_empty());
12861 assert_eq!(ints(&r), vec![2]);
12862 }
12863
12864 #[test]
12865 fn reduction_of_one_item_and_of_a_scalar() {
12866 ctx!(c);
12867 let r = Verb::Reduce(b(plus()))
12868 .monad(&Array::from_i64(vec![7]), &mut c, sp())
12869 .unwrap();
12870 assert!(r.shape.is_empty());
12871 assert_eq!(ints(&r), vec![7]);
12872 let r = Verb::Reduce(b(plus())).monad(&Array::scalar_i64(7), &mut c, sp()).unwrap();
12873 assert_eq!(ints(&r), vec![7]);
12874 }
12875
12876 #[test]
12877 fn reduction_runs_along_the_leading_axis() {
12878 ctx!(c);
12879 let r = Verb::Reduce(b(plus()))
12880 .monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp())
12881 .unwrap();
12882 assert_eq!(r.shape, vec![3]);
12883 assert_eq!(ints(&r), vec![5, 7, 9]);
12884 }
12885
12886 #[test]
12887 fn rank_wrapped_reduction_sums_the_last_axis() {
12888 ctx!(c);
12889 let v = Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]);
12890 let r = v.monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
12891 assert_eq!(r.shape, vec![2]);
12892 assert_eq!(ints(&r), vec![6, 15]);
12893 }
12894
12895 #[test]
12896 fn empty_reduction_uses_the_identity_cell() {
12897 ctx!(c);
12898 let empty = Array::new(vec![0, 2], Data::I64(vec![].into()));
12899 let r = Verb::Reduce(b(plus())).monad(&empty, &mut c, sp()).unwrap();
12900 assert_eq!(r.shape, vec![2]);
12901 assert_eq!(ints(&r), vec![0, 0]);
12902 let r = Verb::Reduce(b(times())).monad(&empty, &mut c, sp()).unwrap();
12903 assert_eq!(ints(&r), vec![1, 1]);
12904 let r = Verb::Reduce(b(floor_v())).monad(&empty, &mut c, sp()).unwrap();
12905 assert!(floats(&r).iter().all(|&x| x == f64::INFINITY));
12906 let r = Verb::Reduce(b(ceil_v())).monad(&empty, &mut c, sp()).unwrap();
12907 assert!(floats(&r).iter().all(|&x| x == f64::NEG_INFINITY));
12908 let r = Verb::Reduce(b(minus())).monad(&empty, &mut c, sp()).unwrap();
12911 assert_eq!(ints(&r), vec![0, 0]);
12912 let r = Verb::Reduce(b(pct())).monad(&empty, &mut c, sp()).unwrap();
12913 assert_eq!(ints(&r), vec![1, 1]);
12914 let r = Verb::Reduce(b(eq_v())).monad(&empty, &mut c, sp()).unwrap();
12915 assert_eq!(bools(&r), vec![1, 1]);
12916 let r = Verb::Reduce(b(plus()))
12918 .monad(&Array::empty(DType::I64), &mut c, sp())
12919 .unwrap();
12920 assert!(r.shape.is_empty());
12921 assert_eq!(ints(&r), vec![0]);
12922 }
12923
12924 #[test]
12925 fn empty_reduction_without_an_identity_is_a_domain_error() {
12926 ctx!(c);
12927 let v = Verb::Hook(b(plus()), b(minus()));
12931 let e = Verb::Reduce(b(v)).monad(&Array::empty(DType::I64), &mut c, sp()).unwrap_err();
12932 assert_eq!(e.kind, ErrorKind::Domain);
12933 assert!(e.msg.contains("identity"), "{}", e.msg);
12934 }
12935
12936 #[test]
12937 fn reduction_with_a_non_primitive_verb_uses_the_general_fold() {
12938 ctx!(c);
12939 let v = Verb::Reduce(b(Verb::Hook(b(plus()), b(minus()))));
12941 let r = v.monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp()).unwrap();
12942 assert_eq!(ints(&r), vec![2]);
12943 }
12944
12945 #[test]
12946 fn dyadic_reduction_is_the_table() {
12947 ctx!(c);
12948 let v = Verb::Reduce(b(plus()));
12951 let r = v
12952 .dyad(&Array::scalar_i64(2), &Array::from_i64(vec![1, 2, 3]), &mut c, sp())
12953 .unwrap();
12954 assert_eq!(r.shape, vec![3]);
12955 assert_eq!(ints(&r), vec![3, 4, 5]);
12956 let r = v
12959 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::from_i64(vec![10, 20]), &mut c, sp())
12960 .unwrap();
12961 assert_eq!(r.shape, vec![3, 2]);
12962 assert_eq!(ints(&r), vec![11, 21, 12, 22, 13, 23]);
12963 let cat = Verb::Reduce(b(inf_prim(",", MonadOp::Ravel, DyadOp::AppendLeading)));
12965 let r = cat
12966 .dyad(&Array::from_i64(vec![1, 2]), &Array::from_i64(vec![3, 4]), &mut c, sp())
12967 .unwrap();
12968 assert_eq!(r.shape, vec![4]);
12969 assert_eq!(ints(&r), vec![1, 2, 3, 4]);
12970 }
12971
12972 #[test]
12975 fn integer_overflow_promotes_the_whole_result_to_float() {
12976 ctx!(c);
12977 let r = plus()
12978 .dyad(&Array::from_i64(vec![1, i64::MAX]), &Array::scalar_i64(1), &mut c, sp())
12979 .unwrap();
12980 assert_eq!(r.dtype(), DType::F64);
12981 let v = floats(&r);
12982 assert!(close(v[0], 2.0));
12983 assert!(close(v[1], i64::MAX as f64 + 1.0));
12984 let r = plus()
12986 .dyad(&Array::from_i64(vec![1, 2]), &Array::scalar_i64(1), &mut c, sp())
12987 .unwrap();
12988 assert_eq!(r.dtype(), DType::I64);
12989 }
12990
12991 #[test]
12992 fn reduction_overflow_promotes_too() {
12993 ctx!(c);
12994 let r = Verb::Reduce(b(plus()))
12995 .monad(&Array::from_i64(vec![i64::MAX, i64::MAX]), &mut c, sp())
12996 .unwrap();
12997 assert_eq!(r.dtype(), DType::F64);
12998 assert!(close(floats(&r)[0], 2.0 * i64::MAX as f64));
12999 }
13000
13001 #[test]
13002 fn booleans_widen_to_integers_in_arithmetic() {
13003 ctx!(c);
13004 let bits = Array::new(vec![3], Data::Bool(vec![1, 0, 1].into()));
13005 let r = plus().dyad(&bits, &bits, &mut c, sp()).unwrap();
13006 assert_eq!(r.dtype(), DType::I64);
13007 assert_eq!(ints(&r), vec![2, 0, 2]);
13008 }
13009
13010 #[test]
13011 fn j_division_is_float_and_survives_zero() {
13012 ctx!(c);
13013 let r = pct()
13014 .dyad(&Array::from_i64(vec![1, -1, 0, 6]), &Array::from_i64(vec![0, 0, 0, 4]), &mut c, sp())
13015 .unwrap();
13016 let v = floats(&r);
13017 assert_eq!(v[0], f64::INFINITY);
13018 assert_eq!(v[1], f64::NEG_INFINITY);
13019 assert_eq!(v[2], 0.0);
13020 assert!(close(v[3], 1.5));
13021 }
13022
13023 #[test]
13024 fn apl_division_by_zero_is_a_domain_error_except_zero_by_zero() {
13025 ctx!(c, Agreement::ExactOrScalar);
13026 let r = div_apl()
13027 .dyad(&Array::scalar_i64(0), &Array::scalar_i64(0), &mut c, sp())
13028 .unwrap();
13029 assert!(close(floats(&r)[0], 1.0));
13030 let e = div_apl()
13031 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(0), &mut c, sp())
13032 .unwrap_err();
13033 assert_eq!(e.kind, ErrorKind::Domain);
13034 assert!(e.msg.contains("division by zero"), "{}", e.msg);
13035 let r = div_apl()
13036 .dyad(&Array::scalar_i64(6), &Array::scalar_i64(4), &mut c, sp())
13037 .unwrap();
13038 assert!(close(floats(&r)[0], 1.5));
13039 }
13040
13041 #[test]
13042 fn reciprocal_of_zero_is_infinite() {
13043 ctx!(c);
13044 let r = pct().monad(&Array::from_i64(vec![0, 2]), &mut c, sp()).unwrap();
13045 let v = floats(&r);
13046 assert_eq!(v[0], f64::INFINITY);
13047 assert!(close(v[1], 0.5));
13048 }
13049
13050 #[test]
13051 fn residue_takes_the_sign_of_the_left_argument() {
13052 ctx!(c);
13053 let x = Array::from_i64(vec![3, 3, -3, -3, 0]);
13054 let y = Array::from_i64(vec![5, -5, 5, -5, 5]);
13055 let r = residue_v().dyad(&x, &y, &mut c, sp()).unwrap();
13056 assert_eq!(ints(&r), vec![2, 1, -1, -2, 5]);
13057 let r = residue_v()
13059 .dyad(&Array::from_f64(vec![2.5]), &Array::from_f64(vec![7.0]), &mut c, sp())
13060 .unwrap();
13061 assert!(close(floats(&r)[0], 2.0));
13062 }
13063
13064 #[test]
13065 fn power_stays_integral_when_it_can() {
13066 ctx!(c);
13067 let r = pow_v()
13068 .dyad(&Array::from_i64(vec![2, 0, 5]), &Array::from_i64(vec![10, 0, 1]), &mut c, sp())
13069 .unwrap();
13070 assert_eq!(r.dtype(), DType::I64);
13071 assert_eq!(ints(&r), vec![1024, 1, 5]);
13072 let r = pow_v()
13074 .dyad(&Array::from_i64(vec![2, 4]), &Array::from_i64(vec![-1, 2]), &mut c, sp())
13075 .unwrap();
13076 assert_eq!(r.dtype(), DType::F64);
13077 assert!(close(floats(&r)[0], 0.5));
13078 assert!(close(floats(&r)[1], 16.0));
13079 let r = pow_v()
13081 .dyad(&Array::scalar_i64(10), &Array::scalar_i64(30), &mut c, sp())
13082 .unwrap();
13083 assert_eq!(r.dtype(), DType::F64);
13084 }
13085
13086 #[test]
13087 fn comparisons_yield_booleans() {
13088 ctx!(c);
13089 let r = lt_v()
13090 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::scalar_i64(2), &mut c, sp())
13091 .unwrap();
13092 assert_eq!(bools(&r), vec![1, 0, 0]);
13093 let r = eq_v()
13094 .dyad(&Array::from_f64(vec![1.0, 2.0]), &Array::from_i64(vec![1, 3]), &mut c, sp())
13095 .unwrap();
13096 assert_eq!(bools(&r), vec![1, 0]);
13097 }
13098
13099 #[test]
13100 fn characters_compare_but_do_not_add() {
13101 ctx!(c);
13102 let a = Array::from_chars(vec!['a', 'b']);
13103 let bb = Array::from_chars(vec!['a', 'c']);
13104 assert_eq!(bools(&eq_v().dyad(&a, &bb, &mut c, sp()).unwrap()), vec![1, 0]);
13105 let e = plus().dyad(&a, &bb, &mut c, sp()).unwrap_err();
13106 assert_eq!(e.kind, ErrorKind::Type);
13107 assert!(e.msg.contains("characters"), "{}", e.msg);
13108 let e = lt_v().dyad(&a, &bb, &mut c, sp()).unwrap_err();
13109 assert_eq!(e.kind, ErrorKind::Type);
13110 let e = plus().dyad(&a, &Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap_err();
13111 assert_eq!(e.kind, ErrorKind::Type);
13112 assert!(e.msg.contains("character"), "{}", e.msg);
13113 let e = plus().monad(&a, &mut c, sp()).unwrap_err();
13114 assert_eq!(e.kind, ErrorKind::Type);
13115 }
13116
13117 #[test]
13118 fn floor_and_ceiling_return_integers_when_they_fit() {
13119 ctx!(c);
13120 let r = floor_v().monad(&Array::from_f64(vec![1.5, -1.5]), &mut c, sp()).unwrap();
13121 assert_eq!(r.dtype(), DType::I64);
13122 assert_eq!(ints(&r), vec![1, -2]);
13123 let r = ceil_v().monad(&Array::from_f64(vec![1.5, -1.5]), &mut c, sp()).unwrap();
13124 assert_eq!(ints(&r), vec![2, -1]);
13125 let r = floor_v().monad(&Array::from_f64(vec![1e30]), &mut c, sp()).unwrap();
13127 assert_eq!(r.dtype(), DType::F64);
13128 let r = floor_v().monad(&Array::from_i64(vec![3]), &mut c, sp()).unwrap();
13130 assert_eq!(ints(&r), vec![3]);
13131 }
13132
13133 #[test]
13134 fn logical_negation_needs_zero_or_one() {
13135 ctx!(c);
13136 let r = not_v().monad(&Array::from_i64(vec![0, 1]), &mut c, sp()).unwrap();
13137 assert_eq!(bools(&r), vec![1, 0]);
13138 let e = not_v().monad(&Array::from_i64(vec![2]), &mut c, sp()).unwrap_err();
13139 assert_eq!(e.kind, ErrorKind::Domain);
13140 }
13141
13142 #[test]
13143 fn signum_abs_and_negation_pick_their_types() {
13144 ctx!(c);
13145 let r = times().monad(&Array::from_i64(vec![-3, 0, 9]), &mut c, sp()).unwrap();
13146 assert_eq!(ints(&r), vec![-1, 0, 1]);
13147 let r = times().monad(&Array::from_f64(vec![-3.0, 0.0, 9.0]), &mut c, sp()).unwrap();
13148 assert_eq!(floats(&r), vec![-1.0, 0.0, 1.0]);
13149 let r = residue_v().monad(&Array::from_i64(vec![-3, 3]), &mut c, sp()).unwrap();
13150 assert_eq!(ints(&r), vec![3, 3]);
13151 let bits = Array::new(vec![2], Data::Bool(vec![0, 1].into()));
13152 let r = minus().monad(&bits, &mut c, sp()).unwrap();
13153 assert_eq!(r.dtype(), DType::I64);
13154 assert_eq!(ints(&r), vec![0, -1]);
13155 }
13156
13157 #[test]
13158 fn square_root_of_a_negative_number_is_complex() {
13159 ctx!(c);
13160 let r = sqrt_v().monad(&Array::from_i64(vec![9]), &mut c, sp()).unwrap();
13161 assert!(close(floats(&r)[0], 3.0));
13162 let r = sqrt_v().monad(&Array::from_i64(vec![-4]), &mut c, sp()).unwrap();
13163 assert_eq!(r.dtype(), DType::Complex);
13164 assert_eq!(r.as_complex_slice().expect("complex data"), &[[0.0, 2.0]]);
13165 }
13166
13167 #[test]
13170 fn shape_tally_and_ravel() {
13171 ctx!(c);
13172 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
13173 let r = dollar().monad(&m, &mut c, sp()).unwrap();
13174 assert_eq!(r.shape, vec![2]);
13175 assert_eq!(ints(&r), vec![2, 3]);
13176 let r = pound().monad(&m, &mut c, sp()).unwrap();
13177 assert!(r.shape.is_empty());
13178 assert_eq!(ints(&r), vec![2]);
13179 let r = pound().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap();
13181 assert_eq!(ints(&r), vec![1]);
13182 let r = comma().monad(&m, &mut c, sp()).unwrap();
13183 assert_eq!(r.shape, vec![6]);
13184 assert_eq!(ints(&r), vec![1, 2, 3, 4, 5, 6]);
13185 }
13186
13187 #[test]
13188 fn transpose_reverses_the_axes() {
13189 ctx!(c);
13190 let r = transpose_v().monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
13191 assert_eq!(r.shape, vec![3, 2]);
13192 assert_eq!(ints(&r), vec![1, 4, 2, 5, 3, 6]);
13193 let a = Array::new(vec![2, 1, 3], Data::I64(vec![1, 2, 3, 4, 5, 6].into()));
13195 let r = transpose_v().monad(&a, &mut c, sp()).unwrap();
13196 assert_eq!(r.shape, vec![3, 1, 2]);
13197 assert_eq!(ints(&r), vec![1, 4, 2, 5, 3, 6]);
13198 let v = Array::from_i64(vec![1, 2]);
13200 assert_eq!(transpose_v().monad(&v, &mut c, sp()).unwrap(), v);
13201 }
13202
13203 #[test]
13204 fn head_and_behead() {
13205 ctx!(c);
13206 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
13207 let r = head_v().monad(&m, &mut c, sp()).unwrap();
13208 assert_eq!(r.shape, vec![3]);
13209 assert_eq!(ints(&r), vec![1, 2, 3]);
13210 let r = behead_v().monad(&m, &mut c, sp()).unwrap();
13211 assert_eq!(r.shape, vec![1, 3]);
13212 assert_eq!(ints(&r), vec![4, 5, 6]);
13213 let e = Array::new(vec![0, 2], Data::I64(vec![].into()));
13215 let r = head_v().monad(&e, &mut c, sp()).unwrap();
13216 assert_eq!(r.shape, vec![2]);
13217 assert_eq!(ints(&r), vec![0, 0]);
13218 assert_eq!(behead_v().monad(&e, &mut c, sp()).unwrap(), e);
13219 assert_eq!(head_v().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap().shape, Vec::<usize>::new());
13220 let err = behead_v().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap_err();
13221 assert_eq!(err.kind, ErrorKind::Domain);
13222 }
13223
13224 #[test]
13225 fn iota_fills_a_shape_and_reverses_negative_axes() {
13226 ctx!(c);
13227 let r = iota().monad(&Array::from_i64(vec![2, 3]), &mut c, sp()).unwrap();
13228 assert_eq!(r.shape, vec![2, 3]);
13229 assert_eq!(ints(&r), vec![0, 1, 2, 3, 4, 5]);
13230 let r = iota().monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
13232 assert_eq!(r.shape, vec![3]);
13233 assert_eq!(ints(&r), vec![0, 1, 2]);
13234 let r = iota().monad(&Array::scalar_i64(-3), &mut c, sp()).unwrap();
13236 assert_eq!(ints(&r), vec![2, 1, 0]);
13237 let r = iota().monad(&Array::from_i64(vec![2, -3]), &mut c, sp()).unwrap();
13238 assert_eq!(r.shape, vec![2, 3]);
13239 assert_eq!(ints(&r), vec![2, 1, 0, 5, 4, 3]);
13240 let r = iota().monad(&Array::from_i64(vec![-2, 3]), &mut c, sp()).unwrap();
13241 assert_eq!(ints(&r), vec![3, 4, 5, 0, 1, 2]);
13242 let r = iota().monad(&Array::scalar_i64(0), &mut c, sp()).unwrap();
13244 assert_eq!(r.shape, vec![0]);
13245 assert!(ints(&r).is_empty());
13246 let e = iota().monad(&Array::from_f64(vec![1.5]), &mut c, sp()).unwrap_err();
13248 assert_eq!(e.kind, ErrorKind::Domain);
13249 let e = iota().monad(&mat(1, 1, vec![1]), &mut c, sp()).unwrap_err();
13250 assert_eq!(e.kind, ErrorKind::Rank);
13251 }
13252
13253 #[test]
13254 fn apl_iota_starts_at_the_index_origin() {
13255 ctx!(c, Agreement::ExactOrScalar);
13256 let r = iota_apl(1).monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
13257 assert_eq!(ints(&r), vec![1, 2, 3]);
13258 let r = iota_apl(0).monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
13259 assert_eq!(ints(&r), vec![0, 1, 2]);
13260 let e = iota_apl(1).monad(&Array::scalar_i64(-1), &mut c, sp()).unwrap_err();
13261 assert_eq!(e.kind, ErrorKind::Domain);
13262 let r = iota_apl(1).monad(&Array::from_i64(vec![2, 3]), &mut c, sp()).unwrap();
13265 assert_eq!(r.shape, vec![2, 3]);
13266 assert_eq!(ints(&r.as_boxes().expect("boxed")[4]), vec![2, 2]);
13267 }
13268
13269 #[test]
13270 fn reshape_cycles_the_ravel() {
13271 ctx!(c);
13272 let r = dollar()
13273 .dyad(&Array::from_i64(vec![2, 3]), &Array::from_i64(vec![1, 2]), &mut c, sp())
13274 .unwrap();
13275 assert_eq!(r.shape, vec![2, 3]);
13276 assert_eq!(ints(&r), vec![1, 2, 1, 2, 1, 2]);
13277 let r = dollar()
13279 .dyad(&Array::scalar_i64(3), &Array::from_i64(vec![7]), &mut c, sp())
13280 .unwrap();
13281 assert_eq!(r.shape, vec![3]);
13282 assert_eq!(ints(&r), vec![7, 7, 7]);
13283 let r = dollar()
13285 .dyad(&Array::scalar_i64(2), &Array::from_chars(vec!['a', 'b', 'c']), &mut c, sp())
13286 .unwrap();
13287 assert_eq!(r.dtype(), DType::Char);
13288 let e = dollar()
13290 .dyad(&Array::scalar_i64(2), &Array::empty(DType::I64), &mut c, sp())
13291 .unwrap_err();
13292 assert_eq!(e.kind, ErrorKind::Length);
13293 assert!(e.msg.contains("empty"), "{}", e.msg);
13294 let r = dollar()
13296 .dyad(&Array::scalar_i64(0), &Array::empty(DType::I64), &mut c, sp())
13297 .unwrap();
13298 assert_eq!(r.shape, vec![0]);
13299 let e = dollar()
13300 .dyad(&Array::scalar_i64(-1), &Array::from_i64(vec![1]), &mut c, sp())
13301 .unwrap_err();
13302 assert_eq!(e.kind, ErrorKind::Domain);
13303 }
13304
13305 #[test]
13306 fn take_from_both_ends_and_beyond() {
13307 ctx!(c);
13308 let v = Array::from_i64(vec![1, 2, 3, 4]);
13309 let take = |x: Array, y: &Array, c: &mut Ctx<'_>| head_v().dyad(&x, y, c, sp()).unwrap();
13310 assert_eq!(ints(&take(Array::scalar_i64(2), &v, &mut c)), vec![1, 2]);
13311 assert_eq!(ints(&take(Array::scalar_i64(-2), &v, &mut c)), vec![3, 4]);
13312 let short = Array::from_i64(vec![1, 2, 3]);
13314 assert_eq!(ints(&take(Array::scalar_i64(6), &short, &mut c)), vec![1, 2, 3, 0, 0, 0]);
13315 assert_eq!(ints(&take(Array::scalar_i64(-6), &short, &mut c)), vec![0, 0, 0, 1, 2, 3]);
13317 let r = take(Array::scalar_i64(2), &Array::scalar_i64(5), &mut c);
13319 assert_eq!(r.shape, vec![2]);
13320 assert_eq!(ints(&r), vec![5, 0]);
13321 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
13323 let r = take(Array::scalar_i64(1), &m, &mut c);
13324 assert_eq!(r.shape, vec![1, 3]);
13325 assert_eq!(ints(&r), vec![1, 2, 3]);
13326 let r = take(Array::scalar_i64(-1), &m, &mut c);
13327 assert_eq!(ints(&r), vec![4, 5, 6]);
13328 let r = take(Array::from_i64(vec![2, 2]), &m, &mut c);
13329 assert_eq!(r.shape, vec![2, 2]);
13330 assert_eq!(ints(&r), vec![1, 2, 4, 5]);
13331 let r = take(Array::from_i64(vec![3, -2]), &m, &mut c);
13332 assert_eq!(r.shape, vec![3, 2]);
13333 assert_eq!(ints(&r), vec![2, 3, 5, 6, 0, 0]);
13334 let r = head_v()
13336 .dyad(&Array::scalar_i64(3), &Array::from_chars(vec!['a']), &mut c, sp())
13337 .unwrap();
13338 assert_eq!(r.data, Data::Char(vec!['a', ' ', ' '].into()));
13339 let e = head_v()
13342 .dyad(&Array::from_i64(vec![1, 1]), &Array::from_i64(vec![1, 2]), &mut c, sp())
13343 .unwrap_err();
13344 assert_eq!(e.kind, ErrorKind::Length);
13345 let r = head_v()
13346 .dyad(&Array::from_i64(vec![1, 2]), &Array::scalar_i64(5), &mut c, sp())
13347 .unwrap();
13348 assert_eq!(r.shape, vec![1, 2]);
13349 assert_eq!(ints(&r), vec![5, 0]);
13350 }
13351
13352 #[test]
13353 fn drop_from_both_ends_and_beyond() {
13354 ctx!(c);
13355 let v = Array::from_i64(vec![1, 2, 3]);
13356 let drop = |x: Array, y: &Array, c: &mut Ctx<'_>| behead_v().dyad(&x, y, c, sp()).unwrap();
13357 assert_eq!(ints(&drop(Array::scalar_i64(1), &v, &mut c)), vec![2, 3]);
13358 assert_eq!(ints(&drop(Array::scalar_i64(-1), &v, &mut c)), vec![1, 2]);
13359 let r = drop(Array::scalar_i64(5), &v, &mut c);
13361 assert_eq!(r.shape, vec![0]);
13362 assert!(ints(&r).is_empty());
13363 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
13364 let r = drop(Array::scalar_i64(1), &m, &mut c);
13365 assert_eq!(r.shape, vec![1, 3]);
13366 assert_eq!(ints(&r), vec![4, 5, 6]);
13367 let r = drop(Array::from_i64(vec![0, -1]), &m, &mut c);
13368 assert_eq!(r.shape, vec![2, 2]);
13369 assert_eq!(ints(&r), vec![1, 2, 4, 5]);
13370 }
13371
13372 #[test]
13375 fn cells_of_unequal_shapes_are_padded_with_fills() {
13376 ctx!(c);
13377 let v = Verb::Rank(b(iota()), [0, 0, 0]);
13379 let r = v.monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp()).unwrap();
13380 assert_eq!(r.shape, vec![3, 3]);
13381 assert_eq!(ints(&r), vec![0, 0, 0, 0, 1, 0, 0, 1, 2]);
13382 }
13383
13384 #[test]
13385 fn framing_aligns_lower_rank_cells_at_the_trailing_axes() {
13386 let cells = vec![Array::from_i64(vec![1, 2]), mat(2, 2, vec![1, 2, 3, 4])];
13387 let r = assemble(&[2], cells, sp()).unwrap();
13388 assert_eq!(r.shape, vec![2, 2, 2]);
13389 assert_eq!(ints(&r), vec![1, 2, 0, 0, 1, 2, 3, 4]);
13390 }
13391
13392 #[test]
13393 fn framing_promotes_cell_types() {
13394 let cells = vec![Array::from_i64(vec![1]), Array::from_f64(vec![2.5])];
13395 let r = assemble(&[2], cells, sp()).unwrap();
13396 assert_eq!(r.dtype(), DType::F64);
13397 assert_eq!(floats(&r), vec![1.0, 2.5]);
13398 let cells = vec![Array::from_i64(vec![1]), Array::from_chars(vec!['a'])];
13400 let e = assemble(&[2], cells, sp()).unwrap_err();
13401 assert_eq!(e.kind, ErrorKind::Type);
13402 }
13403
13404 #[test]
13405 fn framing_over_an_empty_frame_yields_an_empty_result() {
13406 let r = assemble(&[0], Vec::new(), sp()).unwrap();
13407 assert_eq!(r.shape, vec![0]);
13408 assert_eq!(r.count(), 0);
13409 }
13410
13411 #[test]
13414 fn fork_applies_both_tines() {
13415 ctx!(c);
13416 let v = Verb::Fork(b(Verb::Reduce(b(plus()))), b(pct()), b(pound()));
13418 let r = v.monad(&Array::from_i64(vec![1, 2, 3, 4]), &mut c, sp()).unwrap();
13419 assert!(close(floats(&r)[0], 2.5));
13420 let v = Verb::Fork(b(minus()), b(plus()), b(plus()));
13422 let r = v
13423 .dyad(&Array::from_i64(vec![5]), &Array::from_i64(vec![3]), &mut c, sp())
13424 .unwrap();
13425 assert_eq!(ints(&r), vec![10]);
13426 }
13427
13428 #[test]
13429 fn noun_fork_supplies_a_constant_left_argument() {
13430 ctx!(c);
13431 let v = Verb::NounFork(Array::scalar_i64(10), b(minus()), b(right_v()));
13432 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
13433 assert_eq!(ints(&r), vec![9, 8]);
13434 let r = v
13435 .dyad(&Array::scalar_i64(0), &Array::from_i64(vec![1, 2]), &mut c, sp())
13436 .unwrap();
13437 assert_eq!(ints(&r), vec![9, 8]);
13438 }
13439
13440 #[test]
13441 fn hook_reuses_its_right_argument() {
13442 ctx!(c);
13443 let v = Verb::Hook(b(plus()), b(minus()));
13445 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
13446 assert_eq!(ints(&r), vec![0, 0]);
13447 let r = v
13449 .dyad(&Array::from_i64(vec![10]), &Array::from_i64(vec![3]), &mut c, sp())
13450 .unwrap();
13451 assert_eq!(ints(&r), vec![7]);
13452 }
13453
13454 #[test]
13455 fn atop_composes() {
13456 ctx!(c);
13457 let v = Verb::Atop(b(minus()), b(plus()));
13458 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
13459 assert_eq!(ints(&r), vec![-1, -2]);
13460 let r = v
13461 .dyad(&Array::from_i64(vec![1]), &Array::from_i64(vec![2]), &mut c, sp())
13462 .unwrap();
13463 assert_eq!(ints(&r), vec![-3]);
13464 }
13465
13466 #[test]
13467 fn trains_apply_to_the_whole_argument() {
13468 assert_eq!(Verb::Hook(b(plus()), b(minus())).ranks(), [RANK_INF; 3]);
13470 assert_eq!(Verb::Reduce(b(plus())).ranks(), [RANK_INF; 3]);
13471 }
13472
13473 #[test]
13476 fn absent_and_unwritten_meanings_are_reported_differently() {
13477 ctx!(c);
13478 let e = eq_v().monad(&Array::scalar_i64(1), &mut c, sp()).unwrap_err();
13479 assert_eq!(e.kind, ErrorKind::Domain);
13480 assert!(e.msg.contains("no monadic meaning"), "{}", e.msg);
13481 let e = not_v()
13482 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(1), &mut c, sp())
13483 .unwrap_err();
13484 assert_eq!(e.kind, ErrorKind::Domain);
13485 assert!(e.msg.contains("no dyadic meaning"), "{}", e.msg);
13486 let e = pound()
13487 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(1), &mut c, sp())
13488 .unwrap_err();
13489 assert_eq!(e.kind, ErrorKind::NotYet);
13490 assert!(e.msg.contains("copy"), "{}", e.msg);
13491 let _ = echo_v();
13493 }
13494
13495 fn seq_par<T: Send>(f: impl Fn() -> T + Sync + Send) -> (T, T) {
13504 (par::with_threads(1, &f), par::with_threads(4, &f))
13505 }
13506
13507 fn noise(n: usize) -> Vec<f64> {
13509 let mut x = 0x2545_f491_4f6c_dd1du64;
13510 (0..n)
13511 .map(|_| {
13512 x ^= x << 13;
13513 x ^= x >> 7;
13514 x ^= x << 17;
13515 (x >> 11) as f64 / (1u64 << 53) as f64 - 0.5
13516 })
13517 .collect()
13518 }
13519
13520 fn f64_mat(rows: usize, cols: usize) -> Array {
13521 Array::new(vec![rows, cols], Data::F64(noise(rows * cols).into()))
13522 }
13523
13524 const BIG: usize = 200_000;
13526
13527 #[test]
13528 fn an_elementwise_dyad_splits_into_the_same_result() {
13529 let x = Array::from_f64(noise(BIG));
13530 let y = Array::from_f64(noise(BIG).iter().map(|v| v + 0.25).collect());
13531 let (one, many) = seq_par(|| {
13532 ctx!(c);
13533 times().dyad(&x, &y, &mut c, sp()).unwrap()
13534 });
13535 assert_eq!(floats(&one), floats(&many));
13536 let (one, many) = seq_par(|| {
13538 ctx!(c);
13539 plus().dyad(&Array::scalar_f64(0.5), &y, &mut c, sp()).unwrap()
13540 });
13541 assert_eq!(floats(&one), floats(&many));
13542 }
13543
13544 #[test]
13545 fn an_elementwise_dyad_that_overflows_widens_the_same_way() {
13546 let mut v = vec![1i64; BIG];
13549 v[BIG - 3] = i64::MAX;
13550 let x = Array::from_i64(v);
13551 let (one, many) = seq_par(|| {
13552 ctx!(c);
13553 plus().dyad(&x, &x, &mut c, sp()).unwrap()
13554 });
13555 assert_eq!(one.dtype(), DType::F64);
13556 assert_eq!(floats(&one), floats(&many));
13557 }
13558
13559 #[test]
13560 fn an_elementwise_monad_splits_into_the_same_result() {
13561 let y = Array::from_f64(noise(BIG));
13562 for v in [minus(), sqrt_v(), floor_v(), pct()] {
13563 let (one, many) = seq_par(|| {
13564 ctx!(c);
13565 v.monad(&Array::from_f64(y.as_f64_slice().unwrap().iter().map(|x| x.abs()).collect()), &mut c, sp())
13566 .unwrap()
13567 });
13568 assert_eq!(one.data, many.data, "{}", v.name());
13569 }
13570 }
13571
13572 #[test]
13573 fn monadic_cells_run_in_parallel_and_frame_in_order() {
13574 let y = f64_mat(400, 512);
13577 let v = Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]);
13578 let (one, many) = seq_par(|| {
13579 ctx!(c);
13580 v.monad(&y, &mut c, sp()).unwrap()
13581 });
13582 assert_eq!(one.shape, vec![400]);
13583 assert_eq!(floats(&one), floats(&many));
13584 }
13585
13586 #[test]
13587 fn dyadic_cells_run_in_parallel_and_frame_in_order() {
13588 let x = f64_mat(400, 512);
13589 let y = f64_mat(400, 512);
13590 let v = Verb::Rank(b(plus()), [1, 1, 1]);
13592 let (one, many) = seq_par(|| {
13593 ctx!(c);
13594 v.dyad(&x, &y, &mut c, sp()).unwrap()
13595 });
13596 assert_eq!(one.shape, vec![400, 512]);
13597 assert_eq!(floats(&one), floats(&many));
13598 }
13599
13600 #[test]
13601 fn a_verb_that_writes_output_is_not_pure() {
13602 assert!(plus().is_pure());
13603 assert!(Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]).is_pure());
13604 assert!(!echo_v().is_pure());
13605 assert!(!Verb::Rank(b(Verb::Atop(b(echo_v()), b(plus()))), [1, 1, 1]).is_pure());
13606 }
13607
13608 #[test]
13609 fn an_impure_verb_keeps_its_cells_in_order() {
13610 let y = Array::new(vec![16, 8192], Data::I64((0..16 * 8192).collect::<Vec<i64>>().into()));
13613 let v = Verb::Rank(b(Verb::Atop(b(echo_v()), b(head_v()))), [1, 1, 1]);
13614 let mut seen: Vec<i64> = Vec::new();
13615 let mut sink = |s: &str| {
13616 if let Some(first) = s.split_whitespace().next() && let Ok(n) = first.parse::<i64>() {
13617 seen.push(n);
13618 }
13619 };
13620 let mut env = Env::new(Vec::new());
13621 let mut c = Ctx {
13622 cfg: EvalCfg {
13623 agreement: Agreement::LeadingPrefix,
13624 fmt: FmtOpts::J,
13625 tol: Tol::J,
13626 rules: Rules::default(),
13627 },
13628 out: &mut sink,
13629 inp: None,
13630 env: &mut env,
13631 device: None,
13632 };
13633 v.monad(&y, &mut c, sp()).unwrap();
13634 assert_eq!(seen, (0..16).map(|i| i * 8192).collect::<Vec<i64>>());
13635 }
13636
13637 #[test]
13638 fn a_wide_item_reduce_folds_every_column_in_order() {
13639 let y = f64_mat(300, 512);
13642 for v in [plus(), minus(), floor_v()] {
13643 let (one, many) = seq_par(|| {
13644 ctx!(c);
13645 Verb::Reduce(b(v.clone())).monad(&y, &mut c, sp()).unwrap()
13646 });
13647 assert_eq!(one.shape, vec![512]);
13648 assert_eq!(floats(&one), floats(&many), "{}", v.name());
13649 }
13650 }
13651
13652 #[test]
13653 fn a_wide_item_integer_reduce_is_exact() {
13654 let n = 300;
13655 let m = 512;
13656 let y = Array::new(
13657 vec![n, m],
13658 Data::I64((0..(n * m) as i64).map(|i| i % 977 - 400).collect::<Vec<i64>>().into()),
13659 );
13660 let (one, many) = seq_par(|| {
13661 ctx!(c);
13662 Verb::Reduce(b(minus())).monad(&y, &mut c, sp()).unwrap()
13663 });
13664 assert_eq!(ints(&one), ints(&many));
13665 }
13666
13667 #[test]
13668 fn a_narrow_item_reduce_chunks_the_items() {
13669 let y = f64_mat(300_000, 8);
13673 let (one, many) = seq_par(|| {
13674 ctx!(c);
13675 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
13676 });
13677 assert_eq!(one.shape, vec![8]);
13678 for (p, q) in floats(&one).iter().zip(floats(&many)) {
13679 assert!((p - q).abs() <= 1e-12 * p.abs().max(1.0), "{p} vs {q}");
13680 }
13681 let ints_y = Array::new(
13682 vec![300_000, 8],
13683 Data::I64((0..300_000 * 8).map(|i| (i % 101) as i64 - 50).collect::<Vec<i64>>().into()),
13684 );
13685 let (one, many) = seq_par(|| {
13686 ctx!(c);
13687 Verb::Reduce(b(plus())).monad(&ints_y, &mut c, sp()).unwrap()
13688 });
13689 assert_eq!(ints(&one), ints(&many));
13690 }
13691
13692 #[test]
13693 fn a_vector_reduce_folds_the_flat_buffer() {
13694 let y = Array::from_f64(noise(BIG * 4));
13695 let (one, many) = seq_par(|| {
13696 ctx!(c);
13697 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
13698 });
13699 let (p, q) = (floats(&one)[0], floats(&many)[0]);
13700 assert!((p - q).abs() <= 1e-12 * p.abs().max(1.0), "{p} vs {q}");
13701
13702 let ints_y = Array::from_i64((0..BIG as i64 * 4).map(|i| i % 1009 - 500).collect());
13705 for v in [plus(), minus(), ceil_v()] {
13706 let (one, many) = seq_par(|| {
13707 ctx!(c);
13708 Verb::Reduce(b(v.clone())).monad(&ints_y, &mut c, sp()).unwrap()
13709 });
13710 assert_eq!(ints(&one), ints(&many), "{}", v.name());
13711 }
13712 }
13713
13714 #[test]
13715 fn a_reduce_that_overflows_falls_back_to_the_sequential_widening() {
13716 let mut v: Vec<i64> = vec![1; BIG];
13717 v[7] = i64::MAX;
13718 let y = Array::from_i64(v);
13719 let (one, many) = seq_par(|| {
13720 ctx!(c);
13721 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
13722 });
13723 assert_eq!(one.dtype(), DType::F64);
13724 assert_eq!(floats(&one), floats(&many));
13725 }
13726
13727 #[test]
13728 fn a_boolean_reduce_matches_the_sequential_promotion() {
13729 let n = BIG;
13730 let y = Array::new(
13731 vec![n],
13732 Data::Bool((0..n).map(|i| (i % 3 == 0) as u8).collect::<Vec<u8>>().into()),
13733 );
13734 let (one, many) = seq_par(|| {
13735 ctx!(c);
13736 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
13737 });
13738 assert_eq!(one.dtype(), DType::I64);
13739 assert_eq!(ints(&one), ints(&many));
13740 assert_eq!(ints(&one)[0], n.div_ceil(3) as i64);
13741 }
13742}