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 Words,
613 Nest,
616 LevelOf,
619 MapPaths,
622 PolyRoots,
626 PolyDeriv,
629 AnagramIndex,
631 CycleForm,
634 Split,
637 Execute { apl: bool },
642 NotYet(&'static str),
644 None,
646}
647
648#[derive(Clone, Copy, Debug, PartialEq, Eq)]
650pub enum DyadOp {
651 Scalar(ScalarDyad),
652 Reshape,
655 Take,
657 Drop,
659 Right,
661 Left,
663 Rotate,
665 AppendLeading,
667 AppendLast,
669 IndexOf { origin: i64 },
672 MemberJ,
674 MemberApl,
676 From,
679 Match,
681 NotMatch,
683 GradeSelect { down: bool },
685 Copy,
688 Decode,
691 Encode,
695 CollateGrade { down: bool, origin: i64 },
698 TransposeJ,
701 TransposeApl,
704 DecodeApl,
707 EncodeApl,
710 Laminate,
712 Link,
715 Strand,
717 IntervalIndex { offset: i64, closed: bool },
721 IndexOfLast { origin: i64 },
723 MatrixDivide,
725 PartitionEnclose,
728 Squad { origin: i64 },
730 SelectAxis { axis: usize, rank: usize, origin: i64 },
734 Fetch,
736 PolyEval,
739 PolyIntegral,
742 FormatSpec,
745 TruthTable(u8),
749 ExactForm,
753 Deal { origin: i64, fixed: bool },
756 Boolean(BoolDyad),
759 Less,
761 Union,
763 Intersect,
765 AnagramFrom,
768 Permute,
771 FindSeq,
774 UnicodeForm,
777 PrimeMeta,
781 PrimeExponents,
784 Pick { origin: i64 },
786 Expand,
789 WriteStream,
793 NotYet(&'static str),
794 None,
795}
796
797#[derive(Clone, Copy, Debug, PartialEq, Eq)]
799pub enum BoolDyad {
800 Nor,
802 Nand,
804}
805
806#[derive(Clone, Copy, Debug, PartialEq, Eq)]
809pub struct Prim {
810 pub name: &'static str,
811 pub monad: MonadOp,
812 pub dyad: DyadOp,
813 pub ranks: [i64; 3],
814}
815
816#[derive(Clone, Copy, Debug, PartialEq, Eq)]
820pub enum WindowKind {
821 Prefix,
824 Suffix,
827 Scan,
831}
832
833#[derive(Clone, Debug, PartialEq, Eq)]
835pub enum Power {
836 Times(u64),
838 Converge,
840 Each(Vec<u64>),
843 ConvergeTrace,
845}
846
847const CONVERGE_LIMIT: usize = 1 << 20;
849
850pub type MemoCache = Arc<std::sync::Mutex<HashMap<Vec<u64>, Array>>>;
854
855#[derive(Clone, Debug)]
859pub enum Verb {
860 Prim(Prim),
861 Rank(Box<Verb>, [i64; 3]),
863 Reduce(Box<Verb>),
865 Windowed(Box<Verb>, WindowKind),
868 Commute(Box<Verb>),
870 PowerN(Box<Verb>, Power),
872 Fork(Box<Verb>, Box<Verb>, Box<Verb>),
874 NounFork(Array, Box<Verb>, Box<Verb>),
876 Hook(Box<Verb>, Box<Verb>),
878 Atop(Box<Verb>, Box<Verb>),
880 Compose(Box<Verb>, Box<Verb>),
883 BondLeft(Array, Box<Verb>),
886 BondRight(Box<Verb>, Array),
888 Each(Box<Verb>, Enclose),
891 Fit(Box<Verb>, f64),
893 Amend(Array),
895 AmendVerb(Box<Verb>),
898 ShiftFill(Array),
901 Memo(Box<Verb>, MemoCache),
905 Level { u: Box<Verb>, level: i64, spread: bool },
909 Characteristics(Box<Verb>),
912 Before(Box<Verb>, Box<Verb>),
915 UserDerived { def: Box<Verb>, alpha: Box<Verb>, omega: Option<Box<Verb>> },
919 KeyPairs(Box<Verb>),
924 Key(Box<Verb>),
927 Cut(Box<Verb>, i64),
929 PowerV(Box<Verb>, Box<Verb>),
931 PowerUntil(Box<Verb>, Box<Verb>),
933 AlongAxis(Box<Verb>, usize),
937 Explicit(Arc<crate::ir::ExplicitDef>),
941 SelfRef,
944 Named(String),
947 WithObverse(Box<Verb>, Box<Verb>),
950 Agenda(Vec<Verb>, Box<Verb>),
953 Adverse(Box<Verb>, Box<Verb>),
957 Hypergeometric { num: Vec<crate::complex::Cx>, den: Vec<crate::complex::Cx> },
960 Beside(Box<Verb>, Box<Verb>),
964 Stencil(Box<Verb>, Vec<i64>),
968 Evoke(Vec<Verb>, i64),
974}
975
976impl Verb {
977 pub fn ranks(&self) -> [i64; 3] {
979 match self {
980 Verb::Prim(p) => p.ranks,
981 Verb::Rank(_, r) => *r,
982 Verb::Windowed(_, WindowKind::Prefix) => [RANK_INF, 0, RANK_INF],
985 Verb::Each(..) => [0, 0, 0],
986 Verb::Fit(v, _) => v.ranks(),
987 Verb::Amend(_)
990 | Verb::AmendVerb(_)
991 | Verb::ShiftFill(_)
992 | Verb::Level { .. }
993 | Verb::Characteristics(_)
994 | Verb::UserDerived { .. }
995 | Verb::KeyPairs(_)
996 | Verb::Key(_)
997 | Verb::Cut(..)
998 | Verb::PowerV(..)
999 | Verb::PowerUntil(..)
1000 | Verb::AlongAxis(..) => [RANK_INF, RANK_INF, RANK_INF],
1001 Verb::Memo(v, _) => v.ranks(),
1002 Verb::WithObverse(v, _) | Verb::Adverse(v, _) => v.ranks(),
1003 Verb::Beside(..) => [RANK_INF, RANK_INF, RANK_INF],
1004 Verb::Hypergeometric { .. } => [0, 0, 0],
1006 _ => [RANK_INF, RANK_INF, RANK_INF],
1007 }
1008 }
1009
1010 pub fn name(&self) -> String {
1012 match self {
1013 Verb::Prim(p) => p.name.to_string(),
1014 Verb::Rank(v, r) => format!("{}\"{}", v.name(), rank_str(*r)),
1015 Verb::Reduce(v) => format!("{}/", v.name()),
1016 Verb::Windowed(v, WindowKind::Suffix) => format!("{}\\.", v.name()),
1017 Verb::Windowed(v, _) => format!("{}\\", v.name()),
1018 Verb::Commute(v) => format!("{}~", v.name()),
1019 Verb::PowerN(v, Power::Converge) => format!("{}^:_", v.name()),
1020 Verb::PowerN(v, Power::Times(n)) => format!("{}^:{n}", v.name()),
1021 Verb::PowerN(v, Power::Each(_)) => format!("{}^:n", v.name()),
1022 Verb::PowerN(v, Power::ConvergeTrace) => format!("{}^:a:", v.name()),
1023 Verb::Fork(f, g, h) => format!("({} {} {})", f.name(), g.name(), h.name()),
1024 Verb::NounFork(_, g, h) => format!("(n {} {})", g.name(), h.name()),
1025 Verb::Hook(f, g) => format!("({} {})", f.name(), g.name()),
1026 Verb::Atop(f, g) => format!("({}@:{})", f.name(), g.name()),
1027 Verb::Compose(f, g) => format!("({}&:{})", f.name(), g.name()),
1028 Verb::BondLeft(_, v) => format!("(n&{})", v.name()),
1029 Verb::BondRight(v, _) => format!("({}&n)", v.name()),
1030 Verb::Each(v, Enclose::Always) => format!("({}&.>)", v.name()),
1031 Verb::Each(v, _) => format!("({}¨)", v.name()),
1032 Verb::Fit(v, n) => format!("{}!.{n}", v.name()),
1033 Verb::Amend(_) => "(m})".to_string(),
1034 Verb::AmendVerb(v) => format!("({}}})", v.name()),
1035 Verb::ShiftFill(_) => "|.!.n".to_string(),
1036 Verb::Characteristics(v) => format!("{} b.", v.name()),
1037 Verb::Before(f, g) => format!("({}⍛{})", f.name(), g.name()),
1038 Verb::KeyPairs(v) => format!("{}⌸", v.name()),
1039 Verb::UserDerived { def, alpha, omega } => match omega {
1040 Some(g) => format!("({} {} {})", alpha.name(), def.name(), g.name()),
1041 None => format!("({} {})", alpha.name(), def.name()),
1042 },
1043 Verb::Memo(v, _) => format!("{} M.", v.name()),
1044 Verb::Level { u, level, spread } => {
1045 format!("{} {} {level}", u.name(), if *spread { "S:" } else { "L:" })
1046 }
1047 Verb::Key(v) => format!("{}/.", v.name()),
1048 Verb::Cut(v, n) => format!("{};.{n}", v.name()),
1049 Verb::PowerV(v, w) => format!("{}^:{}", v.name(), w.name()),
1050 Verb::PowerUntil(v, w) => format!("{}⍣{}", v.name(), w.name()),
1051 Verb::AlongAxis(v, k) => format!("{}[{k}]", v.name()),
1052 Verb::Explicit(d) => d.name.clone(),
1053 Verb::SelfRef => "$:".to_string(),
1054 Verb::Named(n) => n.clone(),
1055 Verb::WithObverse(v, w) => format!("({}:.{})", v.name(), w.name()),
1056 Verb::Adverse(v, w) => format!("({}::{})", v.name(), w.name()),
1057 Verb::Beside(f, g) => format!("({}∘{})", f.name(), g.name()),
1058 Verb::Hypergeometric { num, den } => {
1059 format!("({} H. {})", cx_list(num), cx_list(den))
1060 }
1061 Verb::Agenda(vs, w) => {
1062 let names: Vec<String> = vs.iter().map(Verb::name).collect();
1063 format!("({}@.{})", names.join("`"), w.name())
1064 }
1065 Verb::Evoke(vs, n) => {
1066 let names: Vec<String> = vs.iter().map(Verb::name).collect();
1067 format!("({}`:{n})", names.join("`"))
1068 }
1069 Verb::Stencil(u, w) => {
1070 let sizes: Vec<String> = w.iter().map(i64::to_string).collect();
1071 format!("({}⌺{})", u.name(), sizes.join(" "))
1072 }
1073 }
1074 }
1075
1076 pub fn uses_tolerance(&self) -> bool {
1081 match self {
1082 Verb::Prim(p) => {
1083 matches!(
1084 p.monad,
1085 MonadOp::Scalar(ScalarMonad::Floor)
1086 | MonadOp::Scalar(ScalarMonad::Ceil)
1087 | MonadOp::Nub
1088 ) || matches!(
1089 p.dyad,
1090 DyadOp::Scalar(
1091 ScalarDyad::Eq
1092 | ScalarDyad::Ne
1093 | ScalarDyad::Lt
1094 | ScalarDyad::Le
1095 | ScalarDyad::Gt
1096 | ScalarDyad::Ge
1097 ) | DyadOp::Match
1098 | DyadOp::NotMatch
1099 | DyadOp::MemberJ
1100 | DyadOp::MemberApl
1101 | DyadOp::IndexOf { .. }
1102 | DyadOp::IndexOfLast { .. }
1103 )
1104 }
1105 Verb::Rank(v, _)
1106 | Verb::Reduce(v)
1107 | Verb::Windowed(v, _)
1108 | Verb::Commute(v)
1109 | Verb::PowerN(v, _)
1110 | Verb::BondLeft(_, v)
1111 | Verb::BondRight(v, _)
1112 | Verb::Each(v, _)
1113 | Verb::Fit(v, _)
1114 | Verb::Key(v)
1115 | Verb::Cut(v, _)
1116 | Verb::AlongAxis(v, _) => v.uses_tolerance(),
1117 Verb::PowerV(v, w) | Verb::PowerUntil(v, w) => {
1118 v.uses_tolerance() || w.uses_tolerance()
1119 }
1120 Verb::Amend(_)
1123 | Verb::AmendVerb(_)
1124 | Verb::ShiftFill(_)
1125 | Verb::Characteristics(_)
1126 | Verb::Explicit(_)
1127 | Verb::SelfRef
1128 | Verb::Named(_)
1129 | Verb::Hypergeometric { .. } => false,
1130 Verb::Memo(v, _) | Verb::Level { u: v, .. } => v.uses_tolerance(),
1131 Verb::WithObverse(v, _) => v.uses_tolerance(),
1132 Verb::Adverse(v, w) | Verb::Beside(v, w) | Verb::Before(v, w) => {
1133 v.uses_tolerance() || w.uses_tolerance()
1134 }
1135 Verb::KeyPairs(v) => v.uses_tolerance(),
1136 Verb::UserDerived { def, alpha, omega } => {
1137 def.uses_tolerance()
1138 || alpha.uses_tolerance()
1139 || omega.as_ref().is_some_and(|g| g.uses_tolerance())
1140 }
1141 Verb::Agenda(vs, w) => {
1142 w.uses_tolerance() || vs.iter().any(Verb::uses_tolerance)
1143 }
1144 Verb::Evoke(vs, _) => vs.iter().any(Verb::uses_tolerance),
1145 Verb::Stencil(u, _) => u.uses_tolerance(),
1146 Verb::Fork(f, g, h) => {
1147 f.uses_tolerance() || g.uses_tolerance() || h.uses_tolerance()
1148 }
1149 Verb::NounFork(_, g, h)
1150 | Verb::Hook(g, h)
1151 | Verb::Atop(g, h)
1152 | Verb::Compose(g, h) => g.uses_tolerance() || h.uses_tolerance(),
1153 }
1154 }
1155
1156 pub fn is_pure(&self) -> bool {
1161 match self {
1162 Verb::Prim(p) => {
1165 !matches!(
1166 p.monad,
1167 MonadOp::Echo | MonadOp::Roll { .. } | MonadOp::ReadStream
1168 ) && !matches!(p.dyad, DyadOp::Deal { .. } | DyadOp::WriteStream)
1169 }
1170 Verb::Rank(v, _)
1171 | Verb::Reduce(v)
1172 | Verb::Windowed(v, _)
1173 | Verb::Commute(v)
1174 | Verb::PowerN(v, _) => v.is_pure(),
1175 Verb::Fork(f, g, h) => f.is_pure() && g.is_pure() && h.is_pure(),
1176 Verb::NounFork(_, g, h)
1177 | Verb::Hook(g, h)
1178 | Verb::Atop(g, h)
1179 | Verb::Compose(g, h) => g.is_pure() && h.is_pure(),
1180 Verb::BondLeft(_, v) | Verb::BondRight(v, _) | Verb::Each(v, _) | Verb::Fit(v, _) => {
1181 v.is_pure()
1182 }
1183 Verb::Key(v) | Verb::Cut(v, _) | Verb::AlongAxis(v, _) => v.is_pure(),
1184 Verb::Hypergeometric { .. } => true,
1185 Verb::PowerV(v, w) | Verb::PowerUntil(v, w) => v.is_pure() && w.is_pure(),
1186 Verb::WithObverse(v, _) => v.is_pure(),
1187 Verb::Adverse(v, w) | Verb::Beside(v, w) | Verb::Before(v, w) => {
1188 v.is_pure() && w.is_pure()
1189 }
1190 Verb::KeyPairs(v) => v.is_pure(),
1191 Verb::UserDerived { .. } => false,
1194 Verb::Agenda(vs, w) => w.is_pure() && vs.iter().all(Verb::is_pure),
1195 Verb::Evoke(vs, _) => vs.iter().all(Verb::is_pure),
1196 Verb::Stencil(u, _) => u.is_pure(),
1197 Verb::Amend(_) | Verb::ShiftFill(_) | Verb::Characteristics(_) => true,
1198 Verb::AmendVerb(v) | Verb::Level { u: v, .. } => v.is_pure(),
1199 Verb::Memo(v, _) => v.is_pure(),
1203 Verb::Explicit(_) | Verb::SelfRef | Verb::Named(_) => false,
1208 }
1209 }
1210
1211 pub fn monad(&self, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1217 let _depth = Nesting::enter(span)?;
1218 if y.is_row_major() {
1219 return self.monad_rows(y, ctx, span);
1220 }
1221 match self.monad_columns(y, ctx, span) {
1222 Some(r) => r,
1223 None => self.monad_rows(&y.to_row_major(), ctx, span),
1224 }
1225 }
1226
1227 fn monad_rows(&self, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1230 debug_assert!(y.is_row_major());
1231 match self {
1232 Verb::Prim(p) => {
1233 if let MonadOp::Scalar(op) = p.monad {
1236 return scalar_monad(op, y, ctx.cfg, span);
1237 }
1238 if p.monad == MonadOp::Open && is_mixed_simple(y) {
1242 return Ok(y.clone());
1243 }
1244 let frame_rank = y.rank() - effective_rank(p.ranks[0], y.rank());
1245 if frame_rank == 0 {
1246 return monad_op(p, y, ctx, span);
1247 }
1248 let frame = y.shape[..frame_rank].to_vec();
1249 let n: usize = frame.iter().product();
1250 let cells = each_cell(n, y.count(), self.is_pure(), ctx, |i, c| {
1251 monad_op(p, &y.cell_at(frame_rank, i), c, span)
1252 })?;
1253 assemble(&frame, cells, span)
1254 }
1255 Verb::Rank(v, r) => {
1256 let frame_rank = y.rank() - effective_rank(r[0], y.rank());
1257 if frame_rank == 0 {
1258 return v.monad(y, ctx, span);
1261 }
1262 if let Some(a) = reduce_vector_cells(v, y, frame_rank) {
1265 return Ok(a);
1266 }
1267 let frame = y.shape[..frame_rank].to_vec();
1268 let n: usize = frame.iter().product();
1269 let cells = each_cell(n, y.count(), self.is_pure(), ctx, |i, c| {
1270 v.monad(&y.cell_at(frame_rank, i), c, span)
1271 })?;
1272 assemble(&frame, cells, span)
1273 }
1274 Verb::Reduce(v) => reduce(v, y, ctx, span),
1275 Verb::Windowed(v, kind) => {
1276 runs(v, y, *kind == WindowKind::Suffix, ctx, span)
1277 }
1278 Verb::Commute(v) => v.dyad(y, y, ctx, span),
1279 Verb::PowerN(v, p) => power(v, p.clone(), None, y, ctx, span),
1280 Verb::Fork(f, g, h) => {
1281 let l = f.monad(y, ctx, span)?;
1282 let r = h.monad(y, ctx, span)?;
1283 g.dyad(&l, &r, ctx, span)
1284 }
1285 Verb::NounFork(n, g, h) => {
1286 let r = h.monad(y, ctx, span)?;
1287 g.dyad(n, &r, ctx, span)
1288 }
1289 Verb::Hook(f, g) => {
1290 let r = g.monad(y, ctx, span)?;
1291 f.dyad(y, &r, ctx, span)
1292 }
1293 Verb::Atop(f, g) | Verb::Compose(f, g) => {
1294 let r = g.monad(y, ctx, span)?;
1295 f.monad(&r, ctx, span)
1296 }
1297 Verb::BondLeft(m, v) => v.dyad(m, y, ctx, span),
1298 Verb::BondRight(v, n) => v.dyad(y, n, ctx, span),
1299 Verb::Each(u, rule) => {
1300 let n = y.count();
1301 let cells = each_cell(n, n, self.is_pure(), ctx, |i, c| {
1302 let opened = open_cell(&atom(y, i));
1303 Ok(enclose(&u.monad(&opened, c, span)?, *rule))
1304 })?;
1305 assemble(&y.shape, cells, span)
1306 }
1307 Verb::Fit(v, n) => {
1308 let tol = Tol { ct: *n, ..ctx.cfg.tol };
1309 ctx.with_tol(tol, |c| v.monad(y, c, span))
1310 }
1311 Verb::Amend(m) => {
1313 if m.rank() != 0 || y.rank() > 1 {
1314 return Err(Error::new(
1315 ErrorKind::Rank,
1316 "selecting with m} takes one index into a list",
1317 Some(span),
1318 ));
1319 }
1320 from_index(m, y, span)
1321 }
1322 Verb::AmendVerb(u) => {
1324 let m = u.monad(y, ctx, span)?;
1325 Verb::Amend(m).monad(y, ctx, span)
1326 }
1327 Verb::ShiftFill(fill) => shift_fill(&Array::scalar_i64(-1), y, fill, span),
1330 Verb::Memo(u, cache) => memoised(u, cache, None, y, ctx, span),
1331 Verb::Characteristics(u) => characteristics(u, y, span),
1332 Verb::Before(f, g) => {
1333 let l = f.monad(y, ctx, span)?;
1334 g.dyad(&l, y, ctx, span)
1335 }
1336 Verb::KeyPairs(u) => key_pairs(u, y, None, ctx, span),
1337 Verb::UserDerived { def, alpha, omega } => {
1338 with_operands(alpha, omega.as_deref(), ctx, |c| def.monad(y, c, span))
1339 }
1340 Verb::Level { u, level, spread } => {
1341 at_level(u, *level, *spread, y, ctx, span)
1342 }
1343 Verb::Key(u) => oblique(u, y, ctx, span),
1344 Verb::Cut(u, n) => cut(u, None, y, *n, ctx, span),
1345 Verb::PowerV(u, v) => power_v(u, v, None, y, ctx, span),
1346 Verb::PowerUntil(u, v) => power_until(u, v, y, ctx, span),
1347 Verb::AlongAxis(u, k) => along_axis(u, None, y, *k, ctx, span),
1348 Verb::Explicit(d) => crate::ir::call_explicit(d, None, y, ctx, span),
1349 Verb::SelfRef => {
1350 let d = self_ref(ctx, span)?;
1351 crate::ir::call_explicit(&d, None, y, ctx, span)
1352 }
1353 Verb::Named(n) => named_verb(ctx, n, span)?.monad(y, ctx, span),
1354 Verb::WithObverse(v, _) => v.monad(y, ctx, span),
1355 Verb::Adverse(v, w) => match v.monad(y, ctx, span) {
1356 Err(e) if e.kind != ErrorKind::NotYet => w.monad(y, ctx, span),
1357 other => other,
1358 },
1359 Verb::Beside(f, g) => {
1360 let r = g.monad(y, ctx, span)?;
1361 f.monad(&r, ctx, span)
1362 }
1363 Verb::Hypergeometric { num, den } => hypergeometric(num, den, y, span),
1364 Verb::Agenda(vs, w) => {
1365 agenda_pick(vs, w, None, y, ctx, span)?.monad(y, ctx, span)
1366 }
1367 Verb::Evoke(vs, n) => evoke(vs, *n, None, y, ctx, span),
1368 Verb::Stencil(u, w) => stencil(u, w, y, ctx, span),
1369 }
1370 }
1371
1372 fn monad_columns(&self, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Option<Result<Array>> {
1381 debug_assert!(!y.is_row_major());
1382 match self {
1383 Verb::Prim(p) => match p.monad {
1384 MonadOp::Scalar(op) => Some(scalar_monad(op, y, ctx.cfg, span)),
1387 MonadOp::ShapeOf | MonadOp::Tally => Some(monad_op(p, y, ctx, span)),
1389 MonadOp::TransposeAxes => Some(Ok(transpose_axes(y))),
1393 _ => None,
1394 },
1395 Verb::Reduce(v) => reduce_columns(v, y).map(Ok),
1398 Verb::Rank(v, r) => {
1401 if y.rank() != effective_rank(r[0], y.rank()) + 1 {
1402 return None;
1403 }
1404 reduce_rows_columns(v, y).map(Ok)
1405 }
1406 _ => None,
1407 }
1408 }
1409
1410 pub fn dyad(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1416 let _depth = Nesting::enter(span)?;
1417 if x.is_row_major() && y.is_row_major() {
1418 return self.dyad_rows(x, y, ctx, span);
1419 }
1420 if let Some(layout) = self.elementwise_layout(x, y) {
1421 return Ok(self.dyad_rows(x, y, ctx, span)?.with_layout(layout));
1422 }
1423 self.dyad_rows(&x.to_row_major(), &y.to_row_major(), ctx, span)
1424 }
1425
1426 fn elementwise_layout(&self, x: &Array, y: &Array) -> Option<Layout> {
1431 let Verb::Prim(p) = self else { return None };
1432 if !matches!(p.dyad, DyadOp::Scalar(_)) {
1433 return None;
1434 }
1435 if x.rank() == 0 {
1436 return Some(y.layout());
1437 }
1438 if y.rank() == 0 {
1439 return Some(x.layout());
1440 }
1441 (x.shape == y.shape && x.layout() == y.layout()).then(|| x.layout())
1442 }
1443
1444 fn dyad_rows(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1448 match self {
1449 Verb::Prim(_) | Verb::Rank(_, _) | Verb::Each(..) => {
1450 self.dyad_ranked(x, y, ctx, span)
1451 }
1452 Verb::Windowed(_, WindowKind::Prefix) => self.dyad_ranked(x, y, ctx, span),
1454 Verb::Windowed(u, WindowKind::Suffix) => outfix(u, x, y, ctx, span),
1457 Verb::Windowed(_, WindowKind::Scan) => {
1458 Err(Error::not_yet("dyadic scan (x f\\ y)", span))
1459 }
1460 Verb::Commute(v) => v.dyad(y, x, ctx, span),
1461 Verb::PowerN(v, p) => power(v, p.clone(), Some(x), y, ctx, span),
1462 Verb::Reduce(v) => table(v, x, y, ctx, span),
1464 Verb::Fork(f, g, h) => {
1465 let l = f.dyad(x, y, ctx, span)?;
1466 let r = h.dyad(x, y, ctx, span)?;
1467 g.dyad(&l, &r, ctx, span)
1468 }
1469 Verb::NounFork(n, g, h) => {
1470 let r = h.dyad(x, y, ctx, span)?;
1471 g.dyad(n, &r, ctx, span)
1472 }
1473 Verb::Hook(f, g) => {
1474 let r = g.monad(y, ctx, span)?;
1475 f.dyad(x, &r, ctx, span)
1476 }
1477 Verb::Atop(f, g) => {
1478 let r = g.dyad(x, y, ctx, span)?;
1479 f.monad(&r, ctx, span)
1480 }
1481 Verb::Compose(f, g) => {
1482 let l = g.monad(x, ctx, span)?;
1483 let r = g.monad(y, ctx, span)?;
1484 f.dyad(&l, &r, ctx, span)
1485 }
1486 Verb::Fit(v, n) => {
1487 let tol = Tol { ct: *n, ..ctx.cfg.tol };
1488 ctx.with_tol(tol, |c| v.dyad(x, y, c, span))
1489 }
1490 Verb::Amend(m) => amend(m, x, y, span),
1491 Verb::AmendVerb(u) => {
1493 let m = u.dyad(x, y, ctx, span)?;
1494 amend(&m, x, y, span)
1495 }
1496 Verb::ShiftFill(fill) => shift_fill(x, y, fill, span),
1497 Verb::Memo(u, cache) => memoised(u, cache, Some(x), y, ctx, span),
1498 Verb::Characteristics(_) => {
1499 Err(Error::domain("u b. has no dyadic meaning", span))
1500 }
1501 Verb::Before(f, g) => {
1502 let l = f.monad(x, ctx, span)?;
1503 g.dyad(&l, y, ctx, span)
1504 }
1505 Verb::KeyPairs(u) => key_pairs(u, x, Some(y), ctx, span),
1506 Verb::UserDerived { def, alpha, omega } => {
1507 with_operands(alpha, omega.as_deref(), ctx, |c| def.dyad(x, y, c, span))
1508 }
1509 Verb::Level { u, level, spread } => {
1510 at_level_dyad(u, *level, *spread, x, y, ctx, span)
1511 }
1512 Verb::Key(u) => key(u, x, y, ctx, span),
1513 Verb::Cut(u, n) => cut(u, Some(x), y, *n, ctx, span),
1514 Verb::PowerV(u, v) => power_v(u, v, Some(x), y, ctx, span),
1515 Verb::PowerUntil(..) => {
1516 Err(Error::not_yet("dyadic power with a function operand (x f⍣g y)", span))
1517 }
1518 Verb::AlongAxis(u, k) => along_axis(u, Some(x), y, *k, ctx, span),
1519 Verb::Explicit(d) => crate::ir::call_explicit(d, Some(x), y, ctx, span),
1520 Verb::SelfRef => {
1521 let d = self_ref(ctx, span)?;
1522 crate::ir::call_explicit(&d, Some(x), y, ctx, span)
1523 }
1524 Verb::Named(n) => named_verb(ctx, n, span)?.dyad(x, y, ctx, span),
1525 Verb::WithObverse(v, _) => v.dyad(x, y, ctx, span),
1526 Verb::Adverse(v, w) => match v.dyad(x, y, ctx, span) {
1527 Err(e) if e.kind != ErrorKind::NotYet => w.dyad(x, y, ctx, span),
1528 other => other,
1529 },
1530 Verb::Beside(f, g) => {
1531 let r = g.monad(y, ctx, span)?;
1532 f.dyad(x, &r, ctx, span)
1533 }
1534 Verb::Hypergeometric { .. } => {
1535 Err(Error::domain("m H. n has no dyadic meaning", span))
1536 }
1537 Verb::Agenda(vs, w) => {
1538 agenda_pick(vs, w, Some(x), y, ctx, span)?.dyad(x, y, ctx, span)
1539 }
1540 Verb::Evoke(vs, n) => evoke(vs, *n, Some(x), y, ctx, span),
1541 Verb::Stencil(..) => {
1542 Err(Error::domain("f⌺w has no dyadic meaning", span))
1543 }
1544 Verb::BondLeft(..) | Verb::BondRight(..) => {
1546 Err(Error::domain(format!("{} has no dyadic meaning", self.name()), span))
1547 }
1548 }
1549 }
1550
1551 fn dyad_ranked(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1553 let ranks = self.ranks();
1554 let er_l = effective_rank(ranks[1], x.rank());
1555 let er_r = effective_rank(ranks[2], y.rank());
1556 if er_l == 0 && er_r == 0 {
1557 if let Some(op) = self.scalar_dyad_op() {
1560 return scalar_dyad(op, x, y, ctx.cfg, span);
1561 }
1562 }
1563 let fxl = x.rank() - er_l;
1564 let fyl = y.rank() - er_r;
1565 let p = agree(&x.shape[..fxl], &y.shape[..fyl], &x.shape, &y.shape, ctx.cfg.agreement, span)?;
1566 if p.frame.is_empty() {
1567 return self.dyad_cell(x, y, ctx, span);
1568 }
1569 let work = x.count().max(y.count());
1570 let cells = each_cell(p.n, work, self.is_pure(), ctx, |i, c| {
1571 let xc = x.cell_at(fxl, i / p.x_div);
1572 let yc = y.cell_at(fyl, i / p.y_div);
1573 self.dyad_cell(&xc, &yc, c, span)
1574 })?;
1575 assemble(&p.frame, cells, span)
1576 }
1577
1578 fn dyad_cell(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1580 match self {
1581 Verb::Prim(p) if p.dyad == DyadOp::WriteStream => {
1584 stream_number(y, 2, "1!:2 writes", span)?;
1585 (ctx.out)(&format!("{}\n", crate::fmt::format_array(x, &ctx.cfg.fmt)));
1586 Ok(x.clone())
1587 }
1588 Verb::Prim(p) => dyad_op(p, x, y, ctx.cfg, span),
1589 Verb::Rank(v, _) => v.dyad(x, y, ctx, span),
1590 Verb::Windowed(v, _) => infix(v, x, y, ctx, span),
1591 Verb::Each(u, rule) => {
1592 let r = u.dyad(&open_cell(x), &open_cell(y), ctx, span)?;
1593 Ok(enclose(&r, *rule))
1594 }
1595 _ => Err(Error::internal("dyad_cell on a verb without cell ranks")),
1596 }
1597 }
1598
1599 fn scalar_dyad_op(&self) -> Option<ScalarDyad> {
1602 match self {
1603 Verb::Prim(p) => match p.dyad {
1604 DyadOp::Scalar(op) => Some(op),
1605 _ => None,
1606 },
1607 Verb::Rank(v, _) => v.scalar_dyad_op(),
1608 _ => None,
1609 }
1610 }
1611}
1612
1613pub fn effective_rank(r: i64, arg_rank: usize) -> usize {
1616 if r >= 0 {
1617 (r as usize).min(arg_rank)
1618 } else {
1619 arg_rank.saturating_sub(r.unsigned_abs() as usize)
1620 }
1621}
1622
1623fn self_ref(ctx: &Ctx<'_>, span: Span) -> Result<Arc<crate::ir::ExplicitDef>> {
1633 ctx.env.current_def().ok_or_else(|| {
1634 Error::new(
1635 ErrorKind::Value,
1636 "self-reference outside an explicit definition",
1637 Some(span),
1638 )
1639 })
1640}
1641
1642fn named_verb(ctx: &Ctx<'_>, name: &str, span: Span) -> Result<Verb> {
1644 ctx.env.verb(name).cloned().ok_or_else(|| {
1645 Error::new(ErrorKind::Value, format!("undefined verb: {name}"), Some(span))
1646 })
1647}
1648
1649fn each_cell<F>(
1650 n: usize,
1651 work: usize,
1652 pure: bool,
1653 ctx: &mut Ctx<'_>,
1654 f: F,
1655) -> Result<Vec<Array>>
1656where
1657 F: Fn(usize, &mut Ctx<'_>) -> Result<Array> + Sync + Send,
1658{
1659 if pure && n > 1 && par::worth_it(work) {
1660 let cfg = ctx.cfg;
1661 return par::map_indexed(n, |i| cfg.pure(|c| f(i, c))).into_iter().collect();
1662 }
1663 (0..n).map(|i| f(i, ctx)).collect()
1664}
1665
1666fn one_rank(r: i64) -> String {
1669 if r == RANK_INF { "_".to_string() } else { r.to_string() }
1670}
1671
1672fn rank_str(r: [i64; 3]) -> String {
1675 if r[0] == r[1] && r[1] == r[2] {
1676 one_rank(r[0])
1677 } else {
1678 format!("{} {} {}", one_rank(r[0]), one_rank(r[1]), one_rank(r[2]))
1679 }
1680}
1681
1682fn show_shape(shape: &[usize]) -> String {
1684 if shape.is_empty() {
1685 return "(scalar)".to_string();
1686 }
1687 shape.iter().map(|n| n.to_string()).collect::<Vec<_>>().join(" ")
1688}
1689
1690fn strides(shape: &[usize]) -> Vec<usize> {
1694 let mut s = vec![1usize; shape.len()];
1695 for k in (0..shape.len().saturating_sub(1)).rev() {
1696 s[k] = s[k + 1] * shape[k + 1];
1697 }
1698 s
1699}
1700
1701fn odometer(coord: &mut [usize], shape: &[usize]) {
1703 for k in (0..coord.len()).rev() {
1704 coord[k] += 1;
1705 if coord[k] < shape[k] {
1706 return;
1707 }
1708 coord[k] = 0;
1709 }
1710}
1711
1712fn push_elem(dst: &mut Data, src: &Data, i: usize) {
1714 match (dst, src) {
1715 (Data::Bool(a), Data::Bool(b)) => a.push(b[i]),
1716 (Data::I64(a), Data::I64(b)) => a.push(b[i]),
1717 (Data::Ext(a), Data::Ext(b)) => a.push(b[i].clone()),
1718 (Data::Rat(a), Data::Rat(b)) => a.push(b[i].clone()),
1719 (Data::F64(a), Data::F64(b)) => a.push(b[i]),
1720 (Data::Complex(a), Data::Complex(b)) => a.push(b[i]),
1721 (Data::Char(a), Data::Char(b)) => a.push(b[i]),
1722 (Data::Box(a), Data::Box(b)) => a.push(b[i].clone()),
1723 _ => debug_assert!(false, "push_elem across dtypes"),
1724 }
1725}
1726
1727fn fill_data(dtype: DType, n: usize) -> Data {
1729 let mut d = Data::empty(dtype);
1730 for _ in 0..n {
1731 d.push_fill();
1732 }
1733 d
1734}
1735
1736struct Pairing {
1741 frame: Vec<usize>,
1742 n: usize,
1743 x_div: usize,
1744 y_div: usize,
1745}
1746
1747fn frame_mismatch(
1748 xs: &[usize],
1749 ys: &[usize],
1750 fx: &[usize],
1751 fy: &[usize],
1752 axis: usize,
1753 span: Span,
1754) -> Error {
1755 let kind = if fx.len() == 1 && fy.len() == 1 { ErrorKind::Length } else { ErrorKind::Shape };
1758 let note = if axis < fx.len() && axis < fy.len() {
1759 format!("frames first differ at axis {axis}: {} vs {}", fx[axis], fy[axis])
1760 } else {
1761 format!(
1762 "frames have different numbers of axes: {} vs {}, diverging at axis {axis}",
1763 fx.len(),
1764 fy.len()
1765 )
1766 };
1767 Error::new(
1768 kind,
1769 format!(
1770 "arguments do not agree: left shape {}, right shape {}",
1771 show_shape(xs),
1772 show_shape(ys)
1773 ),
1774 Some(span),
1775 )
1776 .note(note)
1777}
1778
1779fn agree(
1782 fx: &[usize],
1783 fy: &[usize],
1784 xs: &[usize],
1785 ys: &[usize],
1786 mode: Agreement,
1787 span: Span,
1788) -> Result<Pairing> {
1789 let common = fx.len().min(fy.len());
1790 match mode {
1791 Agreement::LeadingPrefix => {
1792 for i in 0..common {
1793 if fx[i] != fy[i] {
1794 return Err(frame_mismatch(xs, ys, fx, fy, i, span));
1795 }
1796 }
1797 let (long, short) = if fx.len() >= fy.len() { (fx, fy) } else { (fy, fx) };
1798 let n: usize = long.iter().product();
1799 let surplus: usize = long[short.len()..].iter().product();
1800 let (x_div, y_div) =
1801 if fx.len() >= fy.len() { (1, surplus.max(1)) } else { (surplus.max(1), 1) };
1802 Ok(Pairing { frame: long.to_vec(), n, x_div, y_div })
1803 }
1804 Agreement::ExactOrScalar => {
1805 if fx == fy {
1806 let n: usize = fx.iter().product();
1807 return Ok(Pairing { frame: fx.to_vec(), n, x_div: 1, y_div: 1 });
1808 }
1809 let one = |f: &[usize]| f.iter().product::<usize>() == 1;
1816 if fx.is_empty() || (one(fx) && !fy.is_empty()) {
1817 let n: usize = fy.iter().product();
1818 return Ok(Pairing { frame: fy.to_vec(), n, x_div: n.max(1), y_div: 1 });
1819 }
1820 if fy.is_empty() || one(fy) {
1821 let n: usize = fx.iter().product();
1822 return Ok(Pairing { frame: fx.to_vec(), n, x_div: 1, y_div: n.max(1) });
1823 }
1824 let axis = (0..common).find(|&i| fx[i] != fy[i]).unwrap_or(common);
1825 Err(frame_mismatch(xs, ys, fx, fy, axis, span))
1826 }
1827 }
1828}
1829
1830fn assemble(frame: &[usize], cells: Vec<Array>, span: Span) -> Result<Array> {
1834 if cells.is_empty() {
1835 return Ok(Array::new(frame.to_vec(), Data::empty(DType::I64)));
1838 }
1839 let mut dt = cells[0].dtype();
1840 for c in &cells[1..] {
1841 dt = DType::promote(dt, c.dtype()).ok_or_else(|| {
1842 let boxed = dt == DType::Box || c.dtype() == DType::Box;
1843 let what = if boxed {
1844 "cannot frame boxed and unboxed results into one array"
1845 } else {
1846 "cannot frame character and numeric results into one array"
1847 };
1848 Error::new(ErrorKind::Type, what, Some(span))
1849 })?;
1850 }
1851 let widen = |c: &Array| -> Result<Data> {
1852 c.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening while framing"))
1853 };
1854
1855 if cells[1..].iter().all(|c| c.shape == cells[0].shape) {
1856 let mut data = Data::empty(dt);
1857 for c in &cells {
1858 if c.dtype() == dt {
1859 data.extend_from(&c.data);
1860 } else {
1861 data.extend_from(&widen(c)?);
1862 }
1863 }
1864 let mut shape = frame.to_vec();
1865 shape.extend_from_slice(&cells[0].shape);
1866 return Ok(Array::new(shape, data));
1867 }
1868
1869 let crank = cells.iter().map(|c| c.rank()).max().unwrap_or(0);
1872 let padded: Vec<Vec<usize>> = cells
1873 .iter()
1874 .map(|c| {
1875 let mut s = vec![1usize; crank - c.rank()];
1876 s.extend_from_slice(&c.shape);
1877 s
1878 })
1879 .collect();
1880 let mut common = vec![0usize; crank];
1881 for s in &padded {
1882 for k in 0..crank {
1883 common[k] = common[k].max(s[k]);
1884 }
1885 }
1886 let cell_n: usize = common.iter().product();
1887 let mut data = Data::empty(dt);
1888 for (c, ps) in cells.iter().zip(&padded) {
1889 let cd = if c.dtype() == dt { c.data.clone() } else { widen(c)? };
1890 let st = strides(ps);
1891 let mut coord = vec![0usize; crank];
1892 for _ in 0..cell_n {
1893 let mut idx = 0usize;
1894 let mut inside = true;
1895 for k in 0..crank {
1896 if coord[k] >= ps[k] {
1897 inside = false;
1898 break;
1899 }
1900 idx += coord[k] * st[k];
1901 }
1902 if inside {
1903 push_elem(&mut data, &cd, idx);
1904 } else {
1905 data.push_fill();
1906 }
1907 odometer(&mut coord, &common);
1908 }
1909 }
1910 let mut shape = frame.to_vec();
1911 shape.extend_from_slice(&common);
1912 Ok(Array::new(shape, data))
1913}
1914
1915fn atom(a: &Array, i: usize) -> Array {
1920 debug_assert!(a.is_row_major(), "an atom out of a column-major buffer");
1921 Array::new(Vec::new(), a.data.slice(i, i + 1))
1922}
1923
1924fn enclose(y: &Array, rule: Enclose) -> Array {
1926 if rule == Enclose::ExceptSimpleScalar && y.rank() == 0 && y.dtype() != DType::Box {
1927 return y.clone();
1928 }
1929 Array::boxed(y.clone())
1930}
1931
1932fn open_cell(y: &Array) -> Array {
1935 match &y.data {
1936 Data::Box(v) if !v.is_empty() => v[0].clone(),
1937 _ => y.clone(),
1938 }
1939}
1940
1941fn first(y: &Array) -> Array {
1944 if y.count() == 0 {
1945 let mut d = Data::empty(y.dtype());
1946 d.push_fill();
1947 return open_cell(&Array::new(Vec::new(), d));
1948 }
1949 open_cell(&atom(y, 0))
1950}
1951
1952fn depth(y: &Array) -> i64 {
1954 match &y.data {
1955 Data::Box(v) => 1 + v.iter().map(depth).max().unwrap_or(0),
1956 _ => i64::from(y.rank() > 0),
1957 }
1958}
1959
1960fn leaves(a: &Array, out: &mut Vec<Array>) {
1962 match &a.data {
1963 Data::Box(v) => {
1964 for b in v.iter() {
1965 leaves(b, out);
1966 }
1967 }
1968 _ => out.push(a.clone()),
1969 }
1970}
1971
1972fn enlist(y: &Array, span: Span) -> Result<Array> {
1974 let mut parts = Vec::new();
1975 leaves(y, &mut parts);
1976 let mut dt = None;
1979 for p in parts.iter().filter(|p| p.count() > 0) {
1980 dt = Some(match dt {
1981 None => p.dtype(),
1982 Some(t) => DType::promote(t, p.dtype()).ok_or_else(|| {
1983 Error::new(
1984 ErrorKind::Type,
1985 "cannot enlist character and numeric data into one vector",
1986 Some(span),
1987 )
1988 })?,
1989 });
1990 }
1991 let dt = dt.unwrap_or(DType::I64);
1992 let mut data = Data::empty(dt);
1993 for p in &parts {
1994 let cast = p.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening in enlist"))?;
1995 data.extend_from(&cast);
1996 }
1997 Ok(Array::new(vec![data.len()], data))
1998}
1999
2000fn spread(a: &Array, shape: &[usize]) -> Array {
2002 let n: usize = shape.iter().product();
2003 let mut data = Data::empty(a.dtype());
2004 for _ in 0..n {
2005 push_elem(&mut data, &a.data, 0);
2006 }
2007 Array::new(shape.to_vec(), data)
2008}
2009
2010fn wider_shape(a: &[usize], b: &[usize]) -> Vec<usize> {
2013 let r = a.len().max(b.len());
2014 let pad = |s: &[usize]| {
2015 let mut v = vec![1usize; r - s.len()];
2016 v.extend_from_slice(s);
2017 v
2018 };
2019 let (pa, pb) = (pad(a), pad(b));
2020 (0..r).map(|k| pa[k].max(pb[k])).collect()
2021}
2022
2023fn raze(y: &Array, span: Span) -> Result<Array> {
2028 let opened: Vec<Array> = (0..y.count()).map(|i| open_cell(&atom(y, i))).collect();
2029 let mut common: Option<Vec<usize>> = None;
2030 for a in opened.iter().filter(|a| a.rank() > 0) {
2031 common = Some(match common {
2032 None => a.shape[1..].to_vec(),
2033 Some(c) => wider_shape(&c, &a.shape[1..]),
2034 });
2035 }
2036 let common = common.unwrap_or_default();
2037 let mut cells: Vec<Array> = Vec::new();
2038 for a in &opened {
2039 if a.rank() == 0 {
2040 cells.push(spread(a, &common));
2041 continue;
2042 }
2043 for i in 0..a.items() {
2044 cells.push(a.item(i));
2045 }
2046 }
2047 if cells.is_empty() {
2048 return Ok(Array::new(vec![0], Data::empty(DType::I64)));
2049 }
2050 let n = cells.len();
2051 assemble(&[n], cells, span)
2052}
2053
2054fn link(x: &Array, y: &Array, span: Span) -> Result<Array> {
2057 let head = Array::boxed(x.clone());
2058 let tail = if y.dtype() == DType::Box { y.clone() } else { Array::boxed(y.clone()) };
2059 catenate(&head, &tail, true, false, span)
2060}
2061
2062fn nest_like(a: &Array, other: &Array) -> Array {
2065 if a.dtype() == DType::Box || other.dtype() != DType::Box {
2066 return a.clone();
2067 }
2068 let cells: Vec<Array> = (0..a.count()).map(|i| atom(a, i)).collect();
2069 Array::new(a.shape.clone(), Data::Box(cells.into()))
2070}
2071
2072fn box_items(y: &Array) -> Array {
2074 if y.dtype() == DType::Box {
2075 return y.clone();
2076 }
2077 let n = y.items();
2078 let boxes: Vec<Array> = (0..n).map(|i| item_or_self(y, i)).collect();
2079 Array::new(vec![n], Data::Box(boxes.into()))
2080}
2081
2082fn strand(x: &Array, y: &Array, span: Span) -> Result<Array> {
2086 let item = enclose(x, Enclose::ExceptSimpleScalar);
2087 let one = |a: &Array| Array::new(vec![1], a.data.clone());
2088 if item.dtype() != DType::Box
2093 && y.dtype() != DType::Box
2094 && DType::promote(item.dtype(), y.dtype()).is_some()
2095 {
2096 return catenate(&one(&item), y, true, false, span);
2097 }
2098 let head = if item.dtype() == DType::Box { item } else { Array::boxed(item) };
2099 catenate(&one(&head), &box_items(y), true, false, span)
2100}
2101
2102fn char_arith(span: Span) -> Error {
2105 Error::new(ErrorKind::Type, "cannot do arithmetic on characters", Some(span))
2106}
2107
2108fn box_arith(span: Span) -> Error {
2109 Error::new(
2110 ErrorKind::Type,
2111 "cannot do arithmetic on boxed values; open them first (J `>`, APL `⊃`)",
2112 Some(span),
2113 )
2114}
2115
2116fn wrong_type(d: DType, span: Span) -> Error {
2119 match d {
2120 DType::Box => box_arith(span),
2121 _ => char_arith(span),
2122 }
2123}
2124
2125fn borrow_i64<'a>(d: &'a Data, tmp: &'a mut Vec<i64>) -> &'a [i64] {
2131 match d {
2132 Data::I64(v) => v,
2133 Data::Bool(v) => {
2134 *tmp = par::map(v, |&b| b as i64);
2135 &tmp[..]
2136 }
2137 _ => &[],
2139 }
2140}
2141
2142fn borrow_f64<'a>(d: &'a Data, tmp: &'a mut Vec<f64>) -> &'a [f64] {
2144 match d {
2145 Data::F64(v) => v,
2146 Data::I64(v) => {
2147 *tmp = par::map(v, |&x| x as f64);
2148 &tmp[..]
2149 }
2150 Data::Bool(v) => {
2151 *tmp = par::map(v, |&x| x as f64);
2152 &tmp[..]
2153 }
2154 Data::Ext(v) => {
2155 *tmp = par::map(v, exact::ext_to_f64);
2156 &tmp[..]
2157 }
2158 Data::Rat(v) => {
2159 *tmp = par::map(v, Rat::to_f64);
2160 &tmp[..]
2161 }
2162 _ => &[],
2163 }
2164}
2165
2166fn borrow_cx<'a>(d: &'a Data, tmp: &'a mut Vec<Cx>) -> &'a [Cx] {
2168 match d {
2169 Data::Complex(v) => v,
2170 Data::Ext(v) => {
2171 *tmp = par::map(v, |x| [exact::ext_to_f64(x), 0.0]);
2172 &tmp[..]
2173 }
2174 Data::Rat(v) => {
2175 *tmp = par::map(v, |x| [x.to_f64(), 0.0]);
2176 &tmp[..]
2177 }
2178 Data::F64(v) => {
2179 *tmp = par::map(v, |&x| [x, 0.0]);
2180 &tmp[..]
2181 }
2182 Data::I64(v) => {
2183 *tmp = par::map(v, |&x| [x as f64, 0.0]);
2184 &tmp[..]
2185 }
2186 Data::Bool(v) => {
2187 *tmp = v.iter().map(|&x| [x as f64, 0.0]).collect();
2188 &tmp[..]
2189 }
2190 _ => &[],
2191 }
2192}
2193
2194fn as_f64<'a>(d: &'a Data, tmp: &'a mut Vec<f64>, span: Span) -> Result<&'a [f64]> {
2196 if !d.dtype().is_numeric() {
2197 return Err(wrong_type(d.dtype(), span));
2198 }
2199 Ok(borrow_f64(d, tmp))
2200}
2201
2202fn arith_type(a: DType, b: DType, span: Span) -> Result<DType> {
2204 if a == DType::Box || b == DType::Box {
2205 return Err(box_arith(span));
2206 }
2207 match DType::promote(a, b) {
2208 Some(DType::Char) => Err(char_arith(span)),
2209 None => Err(Error::new(
2210 ErrorKind::Type,
2211 "cannot mix character and numeric data",
2212 Some(span),
2213 )),
2214 Some(DType::Bool) => Ok(DType::I64),
2215 Some(t) => Ok(t),
2216 }
2217}
2218
2219#[allow(clippy::too_many_arguments)]
2229#[inline]
2230fn zip_chunk<T, U, F>(
2231 xs: &[T],
2232 xoff: usize,
2233 xdiv: usize,
2234 ys: &[T],
2235 yoff: usize,
2236 ydiv: usize,
2237 start: usize,
2238 out: &mut [U],
2239 mut f: F,
2240) -> bool
2241where
2242 T: Copy,
2243 F: FnMut(T, T, &mut U) -> bool,
2244{
2245 let len = out.len();
2246 if len == 0 {
2247 return true;
2248 }
2249 let last = start + len - 1;
2250 let one_x = xdiv > 1 && start / xdiv == last / xdiv;
2251 let one_y = ydiv > 1 && start / ydiv == last / ydiv;
2252 if xdiv == 1 && ydiv == 1 {
2253 let xc = &xs[xoff + start..xoff + start + len];
2254 let yc = &ys[yoff + start..yoff + start + len];
2255 for ((slot, &a), &b) in out.iter_mut().zip(xc).zip(yc) {
2256 if !f(a, b, slot) {
2257 return false;
2258 }
2259 }
2260 } else if xdiv == 1 && one_y {
2261 let b = ys[yoff + start / ydiv];
2262 let xc = &xs[xoff + start..xoff + start + len];
2263 for (slot, &a) in out.iter_mut().zip(xc) {
2264 if !f(a, b, slot) {
2265 return false;
2266 }
2267 }
2268 } else if one_x && ydiv == 1 {
2269 let a = xs[xoff + start / xdiv];
2270 let yc = &ys[yoff + start..yoff + start + len];
2271 for (slot, &b) in out.iter_mut().zip(yc) {
2272 if !f(a, b, slot) {
2273 return false;
2274 }
2275 }
2276 } else {
2277 for (k, slot) in out.iter_mut().enumerate() {
2278 let i = start + k;
2279 if !f(xs[xoff + i / xdiv], ys[yoff + i / ydiv], slot) {
2280 return false;
2281 }
2282 }
2283 }
2284 true
2285}
2286
2287const LANCZOS: [f64; 9] = [
2291 0.999_999_999_999_809_9,
2292 676.520_368_121_885_1,
2293 -1_259.139_216_722_402_8,
2294 771.323_428_777_653_1,
2295 -176.615_029_162_140_6,
2296 12.507_343_278_686_905,
2297 -0.138_571_095_265_720_12,
2298 9.984_369_578_019_572e-6,
2299 1.505_632_735_149_311_6e-7,
2300];
2301
2302fn gamma(x: f64) -> f64 {
2306 use std::f64::consts::PI;
2307 if x < 0.5 {
2308 return PI / ((PI * x).sin() * gamma(1.0 - x));
2310 }
2311 let z = x - 1.0;
2312 let mut a = LANCZOS[0];
2313 for (i, &c) in LANCZOS.iter().enumerate().skip(1) {
2314 a += c / (z + i as f64);
2315 }
2316 let t = z + 7.5;
2317 (2.0 * PI).sqrt() * t.powf(z + 0.5) * (-t).exp() * a
2318}
2319
2320fn factorial(y: f64) -> f64 {
2323 if y.fract() == 0.0 && y.abs() < 1e17 {
2324 let n = y as i64;
2325 if n < 0 {
2326 return if n % 2 == -1 { f64::INFINITY } else { f64::NEG_INFINITY };
2328 }
2329 if n > 170 {
2330 return f64::INFINITY;
2331 }
2332 let mut c = 1.0f64;
2333 for i in 2..=n {
2334 c *= i as f64;
2335 }
2336 return c;
2337 }
2338 gamma(y + 1.0)
2339}
2340
2341const BINOMIAL_PRODUCT_LIMIT: i64 = 4096;
2344
2345fn binomial_product(x: i64, y: f64) -> f64 {
2349 let mut c = 1.0f64;
2350 for i in 1..=x {
2351 c = c * (y - i as f64 + 1.0) / i as f64;
2352 if c == 0.0 {
2353 break;
2354 }
2355 }
2356 c
2357}
2358
2359fn binomial_i64(x: i64, y: i64) -> Option<i64> {
2363 if x < 0 {
2364 if y >= 0 || y < x {
2367 return Some(0);
2368 }
2369 let v = binomial_exact(-y - 1, -x - 1)?;
2370 return if (y - x) % 2 == 0 { Some(v) } else { v.checked_neg() };
2371 }
2372 binomial_exact(x, y)
2373}
2374
2375fn binomial_exact(x: i64, y: i64) -> Option<i64> {
2378 if x > BINOMIAL_PRODUCT_LIMIT {
2379 return None;
2380 }
2381 let mut c: i128 = 1;
2382 for i in 1..=x as i128 {
2383 c = c.checked_mul(y as i128 - i + 1)? / i;
2384 if c == 0 {
2385 break;
2386 }
2387 }
2388 i64::try_from(c).ok()
2389}
2390
2391fn binomial(x: f64, y: f64) -> f64 {
2393 if x.fract() == 0.0 && x.abs() < 1e17 {
2394 let xi = x as i64;
2395 if xi < 0 {
2396 if y.fract() == 0.0 && y < 0.0 && y >= x {
2397 let sign = if (y as i64 - xi) % 2 == 0 { 1.0 } else { -1.0 };
2398 return sign * binomial_product(-y as i64 - 1, -x - 1.0);
2399 }
2400 return 0.0;
2401 }
2402 if xi <= BINOMIAL_PRODUCT_LIMIT {
2403 return binomial_product(xi, y);
2404 }
2405 }
2406 gamma(y + 1.0) / (gamma(x + 1.0) * gamma(y - x + 1.0))
2407}
2408
2409#[inline]
2412fn i64_op(op: ScalarDyad, a: i64, b: i64) -> Option<i64> {
2413 use ScalarDyad::*;
2414 Some(match op {
2415 Add => a.checked_add(b)?,
2416 Sub => a.checked_sub(b)?,
2417 Mul => a.checked_mul(b)?,
2418 Min => a.min(b),
2419 Max => a.max(b),
2420 Residue => {
2421 if a == 0 {
2422 b
2423 } else {
2424 let mut r = b.wrapping_rem(a);
2426 if r != 0 && (r < 0) != (a < 0) {
2427 r += a;
2428 }
2429 r
2430 }
2431 }
2432 Pow => {
2433 if b < 0 {
2434 return None;
2435 }
2436 a.checked_pow(u32::try_from(b).ok()?)?
2437 }
2438 Binomial => binomial_i64(a, b)?,
2439 _ => return None,
2440 })
2441}
2442
2443#[inline]
2445fn f64_op(op: ScalarDyad, a: f64, b: f64, span: Span) -> Result<f64> {
2446 use ScalarDyad::*;
2447 Ok(match op {
2448 Add => a + b,
2449 Sub => a - b,
2450 Mul => a * b,
2451 Min => a.min(b),
2452 Max => a.max(b),
2453 DivJ => {
2454 if b == 0.0 {
2455 if a == 0.0 { 0.0 } else { f64::INFINITY.copysign(a) }
2456 } else {
2457 a / b
2458 }
2459 }
2460 DivApl => {
2461 if b == 0.0 {
2462 if a == 0.0 {
2463 1.0
2464 } else {
2465 return Err(Error::domain("division by zero", span));
2466 }
2467 } else {
2468 a / b
2469 }
2470 }
2471 Pow => {
2472 if a == 0.0 && b == 0.0 {
2473 1.0
2474 } else {
2475 a.powf(b)
2476 }
2477 }
2478 Residue => {
2479 if a.is_infinite() {
2484 if b == 0.0 || (b > 0.0) == (a > 0.0) { b } else { a }
2485 } else if a == 0.0 {
2486 b
2487 } else {
2488 b - a * (b / a).floor()
2489 }
2490 }
2491 Log => {
2492 if a < 0.0 || b < 0.0 {
2493 return Err(Error::not_yet("complex numbers", span));
2494 }
2495 b.ln() / a.ln()
2496 }
2497 Root => {
2498 if b < 0.0 {
2499 return Err(Error::not_yet("complex numbers", span));
2500 }
2501 b.powf(1.0 / a)
2502 }
2503 Circle => return circle(a, b, span),
2504 Binomial => binomial(a, b),
2505 _ => return Err(Error::internal("non-arithmetic op in the float path")),
2506 })
2507}
2508
2509#[inline]
2513fn escapes_reals(op: ScalarDyad, a: f64, b: f64) -> bool {
2514 use ScalarDyad::*;
2515 match op {
2516 Pow => a < 0.0 && b.fract() != 0.0,
2518 Log => a < 0.0 || b < 0.0,
2519 Root => b < 0.0,
2520 Circle => circle_escapes(a, b),
2521 _ => false,
2522 }
2523}
2524
2525#[inline]
2528fn circle_escapes(k: f64, y: f64) -> bool {
2529 if k.fract() != 0.0 {
2530 return false;
2531 }
2532 match k as i64 {
2533 0 | -1 | -2 | -7 => y.abs() > 1.0,
2534 -4 => y.abs() < 1.0,
2535 -6 => y < 1.0,
2536 8 | -8 | -11 | -12 => true,
2539 _ => false,
2540 }
2541}
2542
2543#[inline]
2553fn circle(k: f64, y: f64, span: Span) -> Result<f64> {
2554 if k.fract() != 0.0 {
2555 return Err(Error::domain("the circle function needs an integer left argument", span));
2556 }
2557 let complex = || Error::internal("a circle function left the reals on the real path");
2558 Ok(match k as i64 {
2559 0 => {
2560 if y.abs() > 1.0 {
2561 return Err(complex());
2562 }
2563 (1.0 - y * y).max(0.0).sqrt()
2564 }
2565 1 => y.sin(),
2566 2 => y.cos(),
2567 3 => y.tan(),
2568 4 => (1.0 + y * y).sqrt(),
2569 5 => y.sinh(),
2570 6 => y.cosh(),
2571 7 => y.tanh(),
2572 -1 => {
2573 if y.abs() > 1.0 {
2574 return Err(complex());
2575 }
2576 y.asin()
2577 }
2578 -2 => {
2579 if y.abs() > 1.0 {
2580 return Err(complex());
2581 }
2582 y.acos()
2583 }
2584 -3 => y.atan(),
2585 -4 => {
2586 if y.abs() < 1.0 {
2587 return Err(complex());
2588 }
2589 y.signum() * (y * y - 1.0).max(0.0).sqrt()
2591 }
2592 -5 => y.asinh(),
2593 -6 => {
2594 if y < 1.0 {
2595 return Err(complex());
2596 }
2597 y.acosh()
2598 }
2599 -7 => {
2600 if y.abs() > 1.0 {
2601 return Err(complex());
2602 }
2603 y.atanh()
2604 }
2605 9 | -9 | -10 => y,
2607 10 => y.abs(),
2608 11 => 0.0,
2609 12 => {
2610 if y < 0.0 {
2611 std::f64::consts::PI
2612 } else {
2613 0.0
2614 }
2615 }
2616 8 | -8 | -11 | -12 => return Err(complex()),
2617 _ => {
2618 return Err(Error::domain(
2619 "the circle functions run from _12 to 12",
2620 span,
2621 ));
2622 }
2623 })
2624}
2625
2626#[inline]
2628fn cx_op(op: ScalarDyad, a: Cx, b: Cx, span: Span) -> Result<Cx> {
2629 use ScalarDyad::*;
2630 Ok(match op {
2631 Add => cx::add(a, b),
2632 Sub => cx::sub(a, b),
2633 Mul => cx::mul(a, b),
2634 DivJ => cx::div(a, b),
2635 DivApl => {
2636 if b == cx::ZERO {
2637 if a == cx::ZERO {
2638 cx::ONE
2639 } else {
2640 return Err(Error::domain("division by zero", span));
2641 }
2642 } else {
2643 cx::div(a, b)
2644 }
2645 }
2646 Pow => cx::pow(a, b),
2647 Log => cx::log(a, b),
2648 Root => cx::root(a, b),
2649 Residue => cx::residue(a, b),
2650 Lcm => cx::lcm(a, b),
2651 Gcd => cx::gcd(a, b),
2652 MakeComplex => cx::add(a, cx::mul(cx::I, b)),
2653 PolarBy => cx::mul(a, cx::exp(cx::mul(cx::I, b))),
2654 Circle => {
2655 if a[1] != 0.0 || a[0].fract() != 0.0 {
2656 return Err(Error::domain(
2657 "the circle function needs an integer left argument",
2658 span,
2659 ));
2660 }
2661 cx::circle(a[0] as i64, b).ok_or_else(|| {
2662 Error::domain("the circle functions run from _12 to 12", span)
2663 })?
2664 }
2665 Min | Max => return Err(no_complex_order(span)),
2666 Binomial => {
2667 return Err(Error::not_yet("the binomial function on complex numbers", span));
2668 }
2669 Eq | Ne | Lt | Le | Gt | Ge => {
2670 return Err(Error::internal("a comparison in the complex arithmetic path"));
2671 }
2672 })
2673}
2674
2675fn no_complex_order(span: Span) -> Error {
2678 Error::new(
2679 ErrorKind::Domain,
2680 "complex numbers have no order; only equality (=, ~:) applies to them",
2681 Some(span),
2682 )
2683}
2684
2685#[allow(clippy::too_many_arguments)]
2686#[inline(always)]
2687fn dyad_cx_chunk(
2688 op: ScalarDyad,
2689 xs: &[Cx],
2690 xoff: usize,
2691 xdiv: usize,
2692 ys: &[Cx],
2693 yoff: usize,
2694 ydiv: usize,
2695 start: usize,
2696 out: &mut [Cx],
2697 span: Span,
2698) -> Result<()> {
2699 use ScalarDyad::*;
2700 macro_rules! plain {
2703 ($step:expr) => {{
2704 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut Cx| {
2705 *slot = $step(a, b);
2706 true
2707 });
2708 return Ok(());
2709 }};
2710 }
2711 match op {
2712 Add => plain!(cx::add),
2713 Sub => plain!(cx::sub),
2714 Mul => plain!(cx::mul),
2715 DivJ => plain!(cx::div),
2716 _ => {}
2717 }
2718 let mut err = None;
2719 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut Cx| {
2720 match cx_op(op, a, b, span) {
2721 Ok(v) => {
2722 *slot = v;
2723 true
2724 }
2725 Err(e) => {
2726 err = Some(e);
2727 false
2728 }
2729 }
2730 });
2731 match err {
2732 Some(e) => Err(e),
2733 None => Ok(()),
2734 }
2735}
2736
2737#[allow(clippy::too_many_arguments)]
2738fn dyad_cx(
2739 op: ScalarDyad,
2740 xs: &[Cx],
2741 xoff: usize,
2742 xdiv: usize,
2743 ys: &[Cx],
2744 yoff: usize,
2745 ydiv: usize,
2746 n: usize,
2747 span: Span,
2748) -> Result<Vec<Cx>> {
2749 par::try_fill(n, |start, part| {
2750 dyad_cx_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part, span)
2751 })
2752}
2753
2754#[allow(clippy::too_many_arguments)]
2756fn complex_dyad_data(
2757 op: ScalarDyad,
2758 x: &Data,
2759 xoff: usize,
2760 xdiv: usize,
2761 y: &Data,
2762 yoff: usize,
2763 ydiv: usize,
2764 n: usize,
2765 span: Span,
2766) -> Result<Data> {
2767 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2768 let xs = borrow_cx(x, &mut tx);
2769 let ys = borrow_cx(y, &mut ty);
2770 Ok(Data::Complex(dyad_cx(op, xs, xoff, xdiv, ys, yoff, ydiv, n, span)?.into()))
2771}
2772
2773fn circle_reads_a_part(x: &Data, xoff: usize, xdiv: usize, n: usize) -> bool {
2778 if x.dtype() == DType::Complex {
2779 return false;
2781 }
2782 let mut tmp = Vec::new();
2783 let xs = borrow_f64(x, &mut tmp);
2784 (0..n).all(|i| {
2785 let k = xs[xoff + i / xdiv];
2786 k.fract() == 0.0 && (9.0..=12.0).contains(&k)
2787 })
2788}
2789
2790#[allow(clippy::too_many_arguments)]
2793fn pass_leaves_reals(
2794 op: ScalarDyad,
2795 x: &Data,
2796 xoff: usize,
2797 xdiv: usize,
2798 y: &Data,
2799 yoff: usize,
2800 ydiv: usize,
2801 n: usize,
2802) -> bool {
2803 use ScalarDyad::*;
2804 if !matches!(op, Pow | Log | Root | Circle) {
2805 return false;
2806 }
2807 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2808 let xs = borrow_f64(x, &mut tx);
2809 let ys = borrow_f64(y, &mut ty);
2810 (0..n).any(|i| escapes_reals(op, xs[xoff + i / xdiv], ys[yoff + i / ydiv]))
2811}
2812
2813#[allow(clippy::too_many_arguments)]
2814#[inline(always)]
2815fn dyad_i64_chunk_body(
2816 op: ScalarDyad,
2817 xs: &[i64],
2818 xoff: usize,
2819 xdiv: usize,
2820 ys: &[i64],
2821 yoff: usize,
2822 ydiv: usize,
2823 start: usize,
2824 out: &mut [i64],
2825) -> bool {
2826 use ScalarDyad::*;
2827 macro_rules! overflowing {
2831 ($m:ident) => {{
2832 let mut over = false;
2833 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
2834 let (v, o) = i64::$m(a, b);
2835 *slot = v;
2836 over |= o;
2837 true
2838 });
2839 !over
2840 }};
2841 }
2842 macro_rules! plain {
2843 ($step:expr) => {{
2844 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
2845 *slot = $step(a, b);
2846 true
2847 })
2848 }};
2849 }
2850 match op {
2851 Add => overflowing!(overflowing_add),
2852 Sub => overflowing!(overflowing_sub),
2853 Mul => overflowing!(overflowing_mul),
2854 Min => plain!(i64::min),
2855 Max => plain!(i64::max),
2856 _ => zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
2857 match i64_op(op, a, b) {
2858 Some(v) => {
2859 *slot = v;
2860 true
2861 }
2862 None => false,
2863 }
2864 }),
2865 }
2866}
2867
2868multiversioned! {
2869 #[allow(clippy::too_many_arguments)]
2876 fn dyad_i64_chunk(
2877 op: ScalarDyad,
2878 xs: &[i64],
2879 xoff: usize,
2880 xdiv: usize,
2881 ys: &[i64],
2882 yoff: usize,
2883 ydiv: usize,
2884 start: usize,
2885 out: &mut [i64],
2886 ) -> bool = dyad_i64_chunk_body;
2887}
2888
2889#[allow(clippy::too_many_arguments)]
2891fn dyad_i64(
2892 op: ScalarDyad,
2893 xs: &[i64],
2894 xoff: usize,
2895 xdiv: usize,
2896 ys: &[i64],
2897 yoff: usize,
2898 ydiv: usize,
2899 n: usize,
2900) -> Option<Vec<i64>> {
2901 let (out, ok) = par::fill(n, |start, part| {
2902 dyad_i64_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part)
2903 });
2904 ok.then_some(out)
2905}
2906
2907#[allow(clippy::too_many_arguments)]
2908#[inline(always)]
2909fn dyad_f64_chunk_body(
2910 op: ScalarDyad,
2911 xs: &[f64],
2912 xoff: usize,
2913 xdiv: usize,
2914 ys: &[f64],
2915 yoff: usize,
2916 ydiv: usize,
2917 start: usize,
2918 out: &mut [f64],
2919 span: Span,
2920) -> Result<()> {
2921 use ScalarDyad::*;
2922 macro_rules! plain {
2925 ($step:expr) => {{
2926 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut f64| {
2927 *slot = $step(a, b);
2928 true
2929 });
2930 return Ok(());
2931 }};
2932 }
2933 match op {
2934 Add => plain!(|a: f64, b: f64| a + b),
2935 Sub => plain!(|a: f64, b: f64| a - b),
2936 Mul => plain!(|a: f64, b: f64| a * b),
2937 Min => plain!(f64::min),
2938 Max => plain!(f64::max),
2939 _ => {}
2940 }
2941 let mut err = None;
2942 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut f64| {
2943 match f64_op(op, a, b, span) {
2944 Ok(v) => {
2945 *slot = v;
2946 true
2947 }
2948 Err(e) => {
2949 err = Some(e);
2950 false
2951 }
2952 }
2953 });
2954 match err {
2955 Some(e) => Err(e),
2956 None => Ok(()),
2957 }
2958}
2959
2960multiversioned! {
2961 #[allow(clippy::too_many_arguments)]
2963 fn dyad_f64_chunk(
2964 op: ScalarDyad,
2965 xs: &[f64],
2966 xoff: usize,
2967 xdiv: usize,
2968 ys: &[f64],
2969 yoff: usize,
2970 ydiv: usize,
2971 start: usize,
2972 out: &mut [f64],
2973 span: Span,
2974 ) -> Result<()> = dyad_f64_chunk_body;
2975}
2976
2977#[allow(clippy::too_many_arguments)]
2978fn dyad_f64(
2979 op: ScalarDyad,
2980 xs: &[f64],
2981 xoff: usize,
2982 xdiv: usize,
2983 ys: &[f64],
2984 yoff: usize,
2985 ydiv: usize,
2986 n: usize,
2987 span: Span,
2988) -> Result<Vec<f64>> {
2989 par::try_fill(n, |start, part| {
2990 dyad_f64_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part, span)
2991 })
2992}
2993
2994fn crossed_types(a: DType, b: DType) -> bool {
2998 let class = |d: DType| match d {
2999 DType::Box => 2,
3000 DType::Char => 1,
3001 _ => 0,
3002 };
3003 class(a) != class(b)
3004}
3005
3006#[allow(clippy::too_many_arguments)]
3007fn compare_data(
3008 op: ScalarDyad,
3009 x: &Data,
3010 xoff: usize,
3011 xdiv: usize,
3012 y: &Data,
3013 yoff: usize,
3014 ydiv: usize,
3015 n: usize,
3016 tol: Tol,
3017 span: Span,
3018) -> Result<Data> {
3019 use ScalarDyad::*;
3020 let (dx, dy) = (x.dtype(), y.dtype());
3021 let equality = matches!(op, Eq | Ne);
3022 let boxed = dx == DType::Box || dy == DType::Box;
3028 if equality && crossed_types(dx, dy) && (!boxed || tol.is_j()) {
3029 let unequal = op == Ne;
3030 return Ok(Data::Bool(vec![u8::from(unequal); n].into()));
3031 }
3032 if boxed {
3033 if !equality {
3036 return Err(box_arith(span));
3037 }
3038 let (Data::Box(a), Data::Box(b)) = (x, y) else {
3039 return Err(Error::not_yet("a scalar function inside a nested array", span));
3042 };
3043 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
3044 for (k, slot) in part.iter_mut().enumerate() {
3045 let i = start + k;
3046 let e = arrays_match(&a[xoff + i / xdiv], &b[yoff + i / ydiv], tol);
3047 *slot = u8::from(if op == Eq { e } else { !e });
3048 }
3049 true
3050 });
3051 return Ok(Data::Bool(out.into()));
3052 }
3053 if dx == DType::Char || dy == DType::Char {
3054 if dx != dy {
3055 return Err(Error::new(
3056 ErrorKind::Type,
3057 "cannot compare character and numeric data",
3058 Some(span),
3059 ));
3060 }
3061 if !equality {
3062 return Err(Error::new(
3063 ErrorKind::Type,
3064 "cannot order character data; only equality applies",
3065 Some(span),
3066 ));
3067 }
3068 let (Data::Char(a), Data::Char(b)) = (x, y) else {
3069 return Err(Error::internal("character comparison on non-character data"));
3070 };
3071 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
3072 zip_chunk(a, xoff, xdiv, b, yoff, ydiv, start, part, |p, q, slot| {
3073 let e = p == q;
3074 *slot = if op == Eq { e as u8 } else { !e as u8 };
3075 true
3076 })
3077 });
3078 return Ok(Data::Bool(out.into()));
3079 }
3080 if DType::promote(dx, dy).is_some_and(DType::is_exact)
3081 && let Some(d) = exact_compare_data(op, x, xoff, xdiv, y, yoff, ydiv, n)
3082 {
3083 return Ok(d);
3084 }
3085 if dx == DType::Complex || dy == DType::Complex {
3086 if !equality {
3087 return Err(no_complex_order(span));
3088 }
3089 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3090 let xs = borrow_cx(x, &mut tx);
3091 let ys = borrow_cx(y, &mut ty);
3092 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
3093 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
3094 let e = tol.eq_cx(a, b);
3095 *slot = if op == Eq { e as u8 } else { !e as u8 };
3096 true
3097 })
3098 });
3099 return Ok(Data::Bool(out.into()));
3100 }
3101 let out = if DType::promote(dx, dy) == Some(DType::F64) {
3104 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3105 let xs = borrow_f64(x, &mut tx);
3106 let ys = borrow_f64(y, &mut ty);
3107 par::fill(n, |start, part: &mut [u8]| {
3108 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
3109 *slot = tol_cmp(op, a, b, tol) as u8;
3110 true
3111 })
3112 })
3113 .0
3114 } else {
3115 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3116 let xs = borrow_i64(x, &mut tx);
3117 let ys = borrow_i64(y, &mut ty);
3118 par::fill(n, |start, part: &mut [u8]| {
3119 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
3120 *slot = cmp_result(op, Some(i64::cmp(&a, &b))) as u8;
3121 true
3122 })
3123 })
3124 .0
3125 };
3126 Ok(Data::Bool(out.into()))
3127}
3128
3129#[inline(always)]
3131pub(crate) fn tol_cmp(op: ScalarDyad, a: f64, b: f64, tol: Tol) -> bool {
3132 use ScalarDyad::*;
3133 match op {
3134 Eq => tol.eq(a, b),
3135 Ne => !tol.eq(a, b),
3136 Lt => tol.lt(a, b),
3137 Le => tol.le(a, b),
3138 Gt => tol.lt(b, a),
3139 Ge => tol.le(b, a),
3140 _ => false,
3141 }
3142}
3143
3144fn cmp_result(op: ScalarDyad, ord: Option<std::cmp::Ordering>) -> bool {
3146 use std::cmp::Ordering::*;
3147 use ScalarDyad::*;
3148 match ord {
3149 None => matches!(op, Ne),
3150 Some(o) => match op {
3151 Eq => o == Equal,
3152 Ne => o != Equal,
3153 Lt => o == Less,
3154 Le => o != Greater,
3155 Gt => o == Greater,
3156 Ge => o != Less,
3157 _ => false,
3158 },
3159 }
3160}
3161
3162fn gcd_i128(a: i128, b: i128) -> i128 {
3164 let (mut a, mut b) = (a.abs(), b.abs());
3165 while b != 0 {
3166 let t = a % b;
3167 a = b;
3168 b = t;
3169 }
3170 a
3171}
3172
3173fn decimal_parts(v: f64) -> Option<(i128, u32)> {
3177 if !v.is_finite() {
3178 return None;
3179 }
3180 let text = format!("{v:e}");
3181 let (mantissa, exponent) = text.split_once('e')?;
3182 let exponent: i32 = exponent.parse().ok()?;
3183 let (whole, fraction) = mantissa.split_once('.').unwrap_or((mantissa, ""));
3184 let mut digits: i128 = format!("{whole}{fraction}").parse().ok()?;
3185 let mut scale = fraction.len() as i32 - exponent;
3186 while scale < 0 {
3189 digits = digits.checked_mul(10)?;
3190 scale += 1;
3191 }
3192 (scale <= 34).then_some((digits, scale as u32))
3195}
3196
3197fn gcd_decimal(a: f64, b: f64) -> Option<f64> {
3202 let (pa, sa) = decimal_parts(a)?;
3203 let (pb, sb) = decimal_parts(b)?;
3204 let scale = sa.max(sb);
3205 let lift = |p: i128, s: u32| 10i128.checked_pow(scale - s).and_then(|k| p.checked_mul(k));
3206 let g = gcd_i128(lift(pa, sa)?, lift(pb, sb)?);
3207 format!("{g}e-{scale}").parse().ok()
3210}
3211
3212fn gcd_f64(a: f64, b: f64, tol: Tol) -> Option<f64> {
3218 let (mut a, mut b) = (a.abs(), b.abs());
3219 if !a.is_finite() || !b.is_finite() {
3220 return None;
3221 }
3222 for _ in 0..1000 {
3225 if b == 0.0 {
3226 return Some(a);
3227 }
3228 if a == 0.0 {
3229 return Some(b);
3230 }
3231 let q = a / b;
3236 let mut k = q.floor();
3237 if tol.eq(q, k + 1.0) {
3238 k += 1.0;
3239 }
3240 let mut r = a - b * k;
3241 if r <= 0.0 || tol.eq(r, b) {
3242 r = 0.0;
3243 }
3244 a = b;
3245 b = r;
3246 }
3247 Some(a)
3248}
3249
3250#[allow(clippy::too_many_arguments)]
3254fn real_lcm_gcd(
3255 op: ScalarDyad,
3256 xs: &[f64],
3257 xoff: usize,
3258 xdiv: usize,
3259 ys: &[f64],
3260 yoff: usize,
3261 ydiv: usize,
3262 n: usize,
3263 tol: Tol,
3264 span: Span,
3265) -> Result<Data> {
3266 let mut out = vec![0.0f64; n];
3267 let mut ok = true;
3268 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, 0, &mut out, |a, b, slot| {
3269 let Some(g) = gcd_decimal(a, b).or_else(|| gcd_f64(a, b, tol)) else {
3270 ok = false;
3271 return false;
3272 };
3273 *slot = if op == ScalarDyad::Gcd {
3274 g
3275 } else if g == 0.0 {
3276 0.0
3277 } else {
3278 a / g * b
3279 };
3280 true
3281 });
3282 if !ok {
3283 return Err(Error::domain("LCM/GCD needs finite values", span));
3284 }
3285 Ok(Data::F64(out.into()))
3286}
3287
3288#[allow(clippy::too_many_arguments)]
3292fn lcm_gcd_data(
3293 op: ScalarDyad,
3294 x: &Data,
3295 xoff: usize,
3296 xdiv: usize,
3297 y: &Data,
3298 yoff: usize,
3299 ydiv: usize,
3300 n: usize,
3301 tol: Tol,
3302 span: Span,
3303) -> Result<Data> {
3304 let t = arith_type(x.dtype(), y.dtype(), span)?;
3305 if t == DType::Complex {
3306 return complex_dyad_data(op, x, xoff, xdiv, y, yoff, ydiv, n, span);
3308 }
3309 if t.is_exact()
3310 && let Some(d) = exact_dyad_data(op, t, x, xoff, xdiv, y, yoff, ydiv, n, span)?
3311 {
3312 return Ok(d);
3313 }
3314 let both_bool = x.dtype() == DType::Bool && y.dtype() == DType::Bool;
3315 let float = t == DType::F64;
3316 let (xs, ys) = if float {
3317 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3318 let xf = borrow_f64(x, &mut tx);
3319 let yf = borrow_f64(y, &mut ty);
3320 let integral = |v: &[f64]| v.iter().all(|&a| a.fract() == 0.0 && fits_i64(a));
3321 if !integral(xf) || !integral(yf) {
3322 return real_lcm_gcd(op, xf, xoff, xdiv, yf, yoff, ydiv, n, tol, span);
3323 }
3324 (
3325 xf.iter().map(|&a| a as i64).collect::<Vec<_>>(),
3326 yf.iter().map(|&a| a as i64).collect::<Vec<_>>(),
3327 )
3328 } else {
3329 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3330 (borrow_i64(x, &mut tx).to_vec(), borrow_i64(y, &mut ty).to_vec())
3331 };
3332 let (out, fits) = par::fill(n, |start, part: &mut [i128]| {
3335 let mut fits = true;
3336 zip_chunk(&xs, xoff, xdiv, &ys, yoff, ydiv, start, part, |a, b, slot| {
3337 let (a, b) = (a as i128, b as i128);
3338 let g = gcd_i128(a, b);
3339 let v = if op == ScalarDyad::Gcd {
3340 g
3341 } else if g == 0 {
3342 0
3343 } else {
3344 a / g * b
3345 };
3346 fits &= i64::try_from(v).is_ok();
3347 *slot = v;
3348 true
3349 });
3350 fits
3351 });
3352 if !fits || float {
3353 return Ok(Data::F64(par::map(&out, |&v| v as f64).into()));
3354 }
3355 if both_bool {
3356 return Ok(Data::Bool(par::map(&out, |&v| v as u8).into()));
3357 }
3358 Ok(Data::I64(par::map(&out, |&v| v as i64).into()))
3359}
3360
3361fn to_rat_vec(d: &Data) -> Option<Vec<Rat>> {
3366 Some(match d {
3367 Data::Bool(v) => v.iter().map(|&b| Rat::from_int(Ext::from(b))).collect(),
3368 Data::I64(v) => v.iter().map(|&x| Rat::from_int(Ext::from(x))).collect(),
3369 Data::Ext(v) => v.iter().map(|x| Rat::from_int(x.clone())).collect(),
3370 Data::Rat(v) => v.to_vec(),
3371 Data::F64(_) | Data::Complex(_) | Data::Char(_) | Data::Box(_) => return None,
3372 })
3373}
3374
3375fn rat_window(d: &Data, off: usize, div: usize, n: usize) -> Option<Vec<Rat>> {
3383 if n == 0 {
3384 return Some(Vec::new());
3385 }
3386 let end = off + (n - 1) / div + 1;
3387 if off == 0 && end == d.len() {
3388 return to_rat_vec(d);
3389 }
3390 to_rat_vec(&d.slice(off, end))
3391}
3392
3393fn exact_data(t: DType, out: Vec<Rat>) -> Data {
3401 if t == DType::Ext && out.iter().all(Rat::is_integer) {
3402 return Data::Ext(out.iter().map(|r| r.to_int().expect("whole")).collect());
3403 }
3404 Data::Rat(out.into())
3405}
3406
3407fn too_large(span: Span) -> Error {
3409 Error::domain(
3410 format!(
3411 "the exact result needs more than {} bits; use floats for a value this large",
3412 exact::MAX_BITS
3413 ),
3414 span,
3415 )
3416}
3417
3418fn exact_pow(a: &Rat, b: &Rat, span: Span) -> Result<Option<Rat>> {
3421 let Some(e) = b.to_int().as_ref().and_then(exact::ext_to_i64) else {
3422 return Ok(None);
3423 };
3424 if let Some(v) = a.pow(e) {
3425 return Ok(Some(v));
3426 }
3427 if a.is_zero() && e < 0 { Ok(None) } else { Err(too_large(span)) }
3429}
3430
3431#[allow(clippy::too_many_arguments)]
3435fn exact_dyad_data(
3436 op: ScalarDyad,
3437 t: DType,
3438 x: &Data,
3439 xoff: usize,
3440 xdiv: usize,
3441 y: &Data,
3442 yoff: usize,
3443 ydiv: usize,
3444 n: usize,
3445 span: Span,
3446) -> Result<Option<Data>> {
3447 use ScalarDyad::*;
3448 let (Some(xs), Some(ys)) = (rat_window(x, xoff, xdiv, n), rat_window(y, yoff, ydiv, n))
3449 else {
3450 return Ok(None);
3451 };
3452 let mut out = Vec::with_capacity(n);
3453 for i in 0..n {
3454 let a = &xs[i / xdiv];
3455 let b = &ys[i / ydiv];
3456 let v = match op {
3457 Add => a.add(b),
3458 Sub => a.sub(b),
3459 Mul => a.mul(b),
3460 DivJ | DivApl => match a.div(b) {
3462 Some(v) => v,
3463 None => return Ok(None),
3464 },
3465 Min => a.min(b).clone(),
3466 Max => a.max(b).clone(),
3467 Residue => exact::rat_residue(a, b),
3468 Gcd => exact::rat_gcd(a, b),
3469 Lcm => exact::rat_lcm(a, b),
3470 Pow => match exact_pow(a, b, span)? {
3471 Some(v) => v,
3472 None => return Ok(None),
3473 },
3474 Binomial => match (a.to_int(), b.to_int()) {
3475 (Some(k), Some(m)) => match exact::ext_binomial(&k, &m) {
3476 Some(v) => Rat::from_int(v),
3477 None => return Ok(None),
3478 },
3479 _ => return Ok(None),
3480 },
3481 Root if t == DType::Ext => {
3484 let (Some(k), Some(m)) = (a.to_int(), b.to_int()) else {
3485 return Ok(None);
3486 };
3487 let Some(k) = exact::ext_to_i64(&k).and_then(|k| u32::try_from(k).ok()) else {
3488 return Ok(None);
3489 };
3490 match exact::exact_root(k, &m) {
3491 Some(v) => Rat::from_int(v),
3492 None => return Ok(None),
3493 }
3494 }
3495 Root | Log | Circle | MakeComplex | PolarBy => return Ok(None),
3496 Eq | Ne | Lt | Le | Gt | Ge => return Ok(None),
3498 };
3499 out.push(v);
3500 }
3501 Ok(Some(exact_data(t, out)))
3502}
3503
3504fn exact_monad(op: ScalarMonad, y: &Array) -> Option<Array> {
3507 use ScalarMonad::*;
3508 let v = to_rat_vec(&y.data)?;
3509 let shape = y.shape.clone();
3510 if matches!(op, Floor | Ceil | Signum) {
3513 let out: Vec<Ext> = v
3514 .iter()
3515 .map(|r| match op {
3516 Floor => r.floor(),
3517 Ceil => r.ceil(),
3518 _ => r.signum(),
3519 })
3520 .collect();
3521 return Some(Array::new(shape, Data::Ext(out.into())).with_layout(y.layout()));
3522 }
3523 let two = Rat::from_int(Ext::from(2));
3524 let mut out = Vec::with_capacity(v.len());
3525 for r in &v {
3526 let value = match op {
3527 Conj => r.clone(),
3528 Neg => r.neg(),
3529 Abs => r.abs(),
3530 Recip => r.recip()?,
3531 Inc => r.add(&Rat::one()),
3532 Dec => r.sub(&Rat::one()),
3533 OneMinus => Rat::one().sub(r),
3534 Double => r.add(r),
3535 Halve => r.div(&two).expect("two is not zero"),
3536 Square => r.mul(r),
3537 Sqrt => r.sqrt()?,
3538 Factorial => Rat::from_int(r.to_int().as_ref().and_then(exact::ext_factorial)?),
3539 Exp | Ln | Pi | Imaginary | Polar | Not => return None,
3542 Floor | Ceil | Signum => unreachable!("handled above"),
3543 };
3544 out.push(value);
3545 }
3546 Some(Array::new(shape, exact_data(y.dtype(), out)).with_layout(y.layout()))
3547}
3548
3549fn to_exact(y: &Array, span: Span) -> Result<Array> {
3554 let data = match &y.data {
3555 Data::Ext(_) | Data::Rat(_) => return Ok(y.clone()),
3556 Data::Bool(v) => Data::Ext(v.iter().map(|&b| Ext::from(b)).collect()),
3557 Data::I64(v) => Data::Ext(v.iter().map(|&x| Ext::from(x)).collect()),
3558 Data::F64(v) => {
3559 let mut out = Vec::with_capacity(v.len());
3560 for &x in v.iter() {
3561 out.push(exact::f64_to_rat(x).ok_or_else(|| {
3562 Error::domain("an infinity has no exact value", span)
3563 })?);
3564 }
3565 exact_data(DType::Ext, out)
3566 }
3567 Data::Complex(_) | Data::Char(_) | Data::Box(_) => {
3568 return Err(Error::domain(
3569 format!("x: needs real numbers, not {} data", y.dtype().name()),
3570 span,
3571 ));
3572 }
3573 };
3574 Ok(Array::new(y.shape.clone(), data).with_layout(y.layout()))
3575}
3576
3577fn from_exact(y: &Array) -> Array {
3580 let shape = y.shape.clone();
3581 match &y.data {
3582 Data::Ext(v) => match v.iter().map(exact::ext_to_i64).collect::<Option<Vec<i64>>>() {
3583 Some(out) => Array::new(shape, Data::I64(out.into())).with_layout(y.layout()),
3584 None => Array::new(shape, Data::F64(v.iter().map(exact::ext_to_f64).collect()))
3585 .with_layout(y.layout()),
3586 },
3587 Data::Rat(v) => Array::new(shape, Data::F64(v.iter().map(Rat::to_f64).collect()))
3588 .with_layout(y.layout()),
3589 _ => y.clone(),
3590 }
3591}
3592
3593fn exact_form(x: &Array, y: &Array, span: Span) -> Result<Array> {
3595 match one_whole(x, "the form x: converts to", span)? {
3596 1 => {
3597 let e = to_exact(y, span)?;
3598 e.cast(DType::Rat).ok_or_else(|| Error::internal("an exact value has no rational form"))
3599 }
3600 2 => {
3601 let e = to_exact(y, span)?;
3602 let v = to_rat_vec(&e.data).ok_or_else(|| Error::internal("x: gave an inexact value"))?;
3603 let mut out = Vec::with_capacity(2 * v.len());
3604 for r in &v {
3605 out.push(r.numer().clone());
3606 out.push(r.denom().clone());
3607 }
3608 let mut shape = y.shape.clone();
3609 shape.push(2);
3610 Ok(Array::new(shape, Data::Ext(out.into())))
3611 }
3612 -1 => Ok(from_exact(y)),
3613 -2 => {
3615 if !y.dtype().is_numeric() {
3616 return Err(Error::domain(
3617 format!("x: needs real numbers, not {} data", y.dtype().name()),
3618 span,
3619 ));
3620 }
3621 Ok(y.clone())
3622 }
3623 n => Err(Error::domain(
3624 format!("x: converts to form 1, 2, _1 or _2, not {n}"),
3625 span,
3626 )),
3627 }
3628}
3629
3630#[allow(clippy::too_many_arguments)]
3634fn exact_compare_data(
3635 op: ScalarDyad,
3636 x: &Data,
3637 xoff: usize,
3638 xdiv: usize,
3639 y: &Data,
3640 yoff: usize,
3641 ydiv: usize,
3642 n: usize,
3643) -> Option<Data> {
3644 let (xs, ys) = (rat_window(x, xoff, xdiv, n)?, rat_window(y, yoff, ydiv, n)?);
3645 let out: Vec<u8> = (0..n)
3646 .map(|i| {
3647 let ord = xs[i / xdiv].cmp(&ys[i / ydiv]);
3648 cmp_result(op, Some(ord)) as u8
3649 })
3650 .collect();
3651 Some(Data::Bool(out.into()))
3652}
3653
3654#[allow(clippy::too_many_arguments)]
3658fn scalar_dyad_data(
3659 op: ScalarDyad,
3660 x: &Data,
3661 xoff: usize,
3662 xdiv: usize,
3663 y: &Data,
3664 yoff: usize,
3665 ydiv: usize,
3666 n: usize,
3667 tol: Tol,
3668 span: Span,
3669) -> Result<Data> {
3670 use ScalarDyad::*;
3671 if matches!(op, Eq | Ne | Lt | Le | Gt | Ge) {
3672 return compare_data(op, x, xoff, xdiv, y, yoff, ydiv, n, tol, span);
3673 }
3674 if matches!(op, Lcm | Gcd) {
3675 return lcm_gcd_data(op, x, xoff, xdiv, y, yoff, ydiv, n, tol, span);
3676 }
3677 let t = arith_type(x.dtype(), y.dtype(), span)?;
3678 if t.is_exact()
3679 && let Some(d) = exact_dyad_data(op, t, x, xoff, xdiv, y, yoff, ydiv, n, span)?
3680 {
3681 return Ok(d);
3682 }
3683 if t == DType::I64 && !matches!(op, DivJ | DivApl | Log | Root | Circle) {
3685 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3688 let xs = borrow_i64(x, &mut tx);
3689 let ys = borrow_i64(y, &mut ty);
3690 if let Some(v) = dyad_i64(op, xs, xoff, xdiv, ys, yoff, ydiv, n) {
3691 return Ok(Data::I64(v.into()));
3692 }
3693 }
3695 if t == DType::Complex
3696 || matches!(op, MakeComplex | PolarBy)
3697 || pass_leaves_reals(op, x, xoff, xdiv, y, yoff, ydiv, n)
3698 {
3699 let data = complex_dyad_data(op, x, xoff, xdiv, y, yoff, ydiv, n, span)?;
3700 if op == Circle && circle_reads_a_part(x, xoff, xdiv, n) && let Data::Complex(v) = &data {
3701 return Ok(Data::F64(v.iter().map(|z| z[0]).collect()));
3702 }
3703 return Ok(data);
3704 }
3705 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3706 let xs = borrow_f64(x, &mut tx);
3707 let ys = borrow_f64(y, &mut ty);
3708 Ok(Data::F64(dyad_f64(op, xs, xoff, xdiv, ys, yoff, ydiv, n, span)?.into()))
3709}
3710
3711fn frame_pervaded(frame: Vec<usize>, cells: Vec<Array>, span: Span) -> Result<Array> {
3716 if cells.iter().all(|c| c.rank() == 0 && c.dtype() != DType::Box) {
3717 return assemble(&frame, cells, span);
3718 }
3719 let boxes: Vec<Array> = cells.into_iter().collect();
3720 Ok(Array::new(frame, Data::Box(boxes.into())))
3721}
3722
3723fn pervade_dyad(
3729 op: ScalarDyad,
3730 x: &Array,
3731 y: &Array,
3732 cfg: EvalCfg,
3733 span: Span,
3734) -> Result<Array> {
3735 let p = agree(&x.shape, &y.shape, &x.shape, &y.shape, cfg.agreement, span)?;
3736 if p.n == 0 {
3737 return Ok(Array::new(p.frame, Data::empty(DType::Box)));
3738 }
3739 let (xr, yr) = (x.to_row_major(), y.to_row_major());
3740 let mut cells = Vec::with_capacity(p.n);
3741 for i in 0..p.n {
3742 let a = open_cell(&atom(&xr, i / p.x_div));
3743 let b = open_cell(&atom(&yr, i / p.y_div));
3744 cells.push(scalar_dyad(op, &a, &b, cfg, span)?);
3745 }
3746 frame_pervaded(p.frame, cells, span)
3747}
3748
3749fn pervade_monad(op: ScalarMonad, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
3751 if y.count() == 0 {
3752 return Ok(Array::new(y.shape.clone(), Data::empty(DType::Box)));
3753 }
3754 let yr = y.to_row_major();
3755 let mut cells = Vec::with_capacity(y.count());
3756 for i in 0..y.count() {
3757 let a = open_cell(&atom(&yr, i));
3758 cells.push(scalar_monad(op, &a, cfg, span)?);
3759 }
3760 frame_pervaded(y.shape.clone(), cells, span)
3761}
3762
3763fn scalar_dyad(
3764 op: ScalarDyad,
3765 x: &Array,
3766 y: &Array,
3767 cfg: EvalCfg,
3768 span: Span,
3769) -> Result<Array> {
3770 if cfg.rules.lang == crate::Lang::Apl
3771 && (x.dtype() == DType::Box || y.dtype() == DType::Box)
3772 {
3773 return pervade_dyad(op, x, y, cfg, span);
3774 }
3775 let p = agree(&x.shape, &y.shape, &x.shape, &y.shape, cfg.agreement, span)?;
3776 if p.n == 0 {
3780 return Ok(Array::new(p.frame, Data::empty(empty_result_type(x, y))));
3781 }
3782 let data =
3783 scalar_dyad_data(op, &x.data, 0, p.x_div, &y.data, 0, p.y_div, p.n, cfg.tol, span)?;
3784 Ok(Array::new(p.frame, data))
3785}
3786
3787fn empty_result_type(x: &Array, y: &Array) -> DType {
3790 for a in [x, y] {
3791 if a.dtype().is_numeric() {
3792 return a.dtype();
3793 }
3794 }
3795 DType::I64
3796}
3797
3798fn fits_i64(v: f64) -> bool {
3800 v.is_finite() && v >= i64::MIN as f64 && v < i64::MAX as f64
3801}
3802
3803fn monad_leaves_reals(op: ScalarMonad, d: &Data) -> bool {
3805 use ScalarMonad::*;
3806 match op {
3807 Imaginary | Polar => d.dtype().is_numeric(),
3809 Sqrt | Ln => match d {
3810 Data::I64(v) => par::any(v, |&x| x < 0),
3811 Data::F64(v) => par::any(v, |&x| x < 0.0),
3812 Data::Ext(v) => v.iter().any(|x| x.sign() == num_bigint::Sign::Minus),
3813 Data::Rat(v) => v.iter().any(|x| x < &Rat::zero()),
3814 _ => false,
3815 },
3816 _ => false,
3817 }
3818}
3819
3820fn complex_monad(op: ScalarMonad, y: &Array, span: Span) -> Result<Array> {
3822 use ScalarMonad::*;
3823 let mut tmp = Vec::new();
3824 let v = borrow_cx(&y.data, &mut tmp);
3825 if y.count() > 0 && v.is_empty() {
3826 return Err(wrong_type(y.dtype(), span));
3827 }
3828 let data = match op {
3829 Abs => Data::F64(par::map(v, |&z| cx::abs(z)).into()),
3831 Not => return Err(Error::domain("logical negation needs values of 0 or 1", span)),
3832 Factorial => {
3833 return Err(Error::not_yet("the factorial of a complex number", span));
3834 }
3835 _ => {
3836 let step: fn(Cx) -> Cx = match op {
3837 Conj => cx::conj,
3838 Neg => cx::neg,
3839 Signum => cx::signum,
3840 Recip => cx::recip,
3841 Sqrt => cx::sqrt,
3842 Exp => cx::exp,
3843 Ln => cx::ln,
3844 Floor => cx::floor,
3845 Ceil => cx::ceil,
3846 OneMinus => |z| cx::sub(cx::ONE, z),
3847 Inc => |z| cx::add(z, cx::ONE),
3848 Dec => |z| cx::sub(z, cx::ONE),
3849 Double => |z| cx::add(z, z),
3850 Halve => |z| [z[0] / 2.0, z[1] / 2.0],
3851 Square => |z| cx::mul(z, z),
3852 Pi => |z| [std::f64::consts::PI * z[0], std::f64::consts::PI * z[1]],
3853 Imaginary => |z| cx::mul(cx::I, z),
3854 Polar => |z| cx::exp(cx::mul(cx::I, z)),
3855 Abs | Not | Factorial => unreachable!("handled above"),
3856 };
3857 Data::Complex(par::map(v, |&z| step(z)).into())
3858 }
3859 };
3860 Ok(Array::new(y.shape.clone(), data).with_layout(y.layout()))
3861}
3862
3863fn scalar_monad(op: ScalarMonad, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
3865 use ScalarMonad::*;
3866 if cfg.rules.lang == crate::Lang::Apl && y.dtype() == DType::Box {
3867 return pervade_monad(op, y, cfg, span);
3868 }
3869 let tol = cfg.tol;
3870 let d = &y.data;
3871 if y.count() == 0 && !d.dtype().is_numeric() {
3874 return Ok(Array::new(y.shape.clone(), Data::empty(DType::I64)));
3875 }
3876 if d.dtype() == DType::Complex || monad_leaves_reals(op, d) {
3877 return complex_monad(op, y, span);
3878 }
3879 if d.dtype().is_exact() && let Some(a) = exact_monad(op, y) {
3880 return Ok(a);
3881 }
3882 let mut tmp = Vec::new();
3886 let data = match op {
3887 Conj if d.dtype().is_numeric() => d.clone(),
3889 Conj => return Err(wrong_type(d.dtype(), span)),
3890 Imaginary | Polar => return Err(Error::internal("a complex monad on the real path")),
3893 Neg => match d {
3894 Data::Bool(v) => Data::I64(par::map(v, |&b| -(b as i64)).into()),
3895 Data::I64(v) => match par::try_map(v, i64::checked_neg) {
3896 Some(out) => Data::I64(out.into()),
3897 None => Data::F64(par::map(v, |&x| -(x as f64)).into()),
3898 },
3899 Data::F64(v) => Data::F64(par::map(v, |&x| -x).into()),
3900 _ => return Err(wrong_type(d.dtype(), span)),
3901 },
3902 Signum => match d {
3903 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64).into()),
3904 Data::I64(v) => Data::I64(par::map(v, |&x| x.signum()).into()),
3905 Data::F64(v) => Data::F64(
3908 par::map(v, |&x| {
3909 if tol.is_zero(x) {
3910 0.0
3911 } else if x > 0.0 {
3912 1.0
3913 } else if x < 0.0 {
3914 -1.0
3915 } else {
3916 0.0
3917 }
3918 })
3919 .into(),
3920 ),
3921 _ => return Err(wrong_type(d.dtype(), span)),
3922 },
3923 Recip => {
3924 let v = as_f64(d, &mut tmp, span)?;
3928 Data::F64(par::map(v, |&x| if x == 0.0 { f64::INFINITY } else { 1.0 / x }).into())
3929 }
3930 Sqrt => {
3931 let v = as_f64(d, &mut tmp, span)?;
3933 Data::F64(par::map(v, |&x| x.sqrt()).into())
3934 }
3935 Exp => {
3936 let v = as_f64(d, &mut tmp, span)?;
3937 Data::F64(par::map(v, |&x| x.exp()).into())
3938 }
3939 Abs => match d {
3940 Data::Bool(_) => d.clone(),
3941 Data::I64(v) => match par::try_map(v, i64::checked_abs) {
3942 Some(out) => Data::I64(out.into()),
3943 None => Data::F64(par::map(v, |&x| (x as f64).abs()).into()),
3944 },
3945 Data::F64(v) => Data::F64(par::map(v, |&x| x.abs()).into()),
3946 _ => return Err(wrong_type(d.dtype(), span)),
3947 },
3948 Floor | Ceil => match d {
3949 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64).into()),
3950 Data::I64(_) => d.clone(),
3951 Data::F64(v) => {
3952 let round = |x: f64| if op == Floor { tol.floor(x) } else { tol.ceil(x) };
3953 match par::try_map(v, |x| {
3955 let r = round(x);
3956 fits_i64(r).then_some(r as i64)
3957 }) {
3958 Some(out) => Data::I64(out.into()),
3959 None => Data::F64(par::map(v, |&x| round(x)).into()),
3960 }
3961 }
3962 _ => return Err(wrong_type(d.dtype(), span)),
3963 },
3964 Inc | Dec => {
3965 let step = if op == Inc { 1i64 } else { -1 };
3966 match d {
3967 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64 + step).into()),
3968 Data::I64(v) => match par::try_map(v, |x: i64| x.checked_add(step)) {
3969 Some(out) => Data::I64(out.into()),
3970 None => Data::F64(par::map(v, |&x| x as f64 + step as f64).into()),
3971 },
3972 Data::F64(v) => Data::F64(par::map(v, |&x| x + step as f64).into()),
3973 _ => return Err(wrong_type(d.dtype(), span)),
3974 }
3975 }
3976 Double | Square => match d {
3977 Data::Bool(v) => {
3978 Data::I64(par::map(v, |&b| if op == Double { 2 * b as i64 } else { b as i64 }).into())
3979 }
3980 Data::I64(v) => {
3981 let f = |x: i64| if op == Double { x.checked_mul(2) } else { x.checked_mul(x) };
3982 match par::try_map(v, f) {
3983 Some(out) => Data::I64(out.into()),
3984 None => Data::F64(
3985 par::map(v, |&x| {
3986 let x = x as f64;
3987 if op == Double { x + x } else { x * x }
3988 })
3989 .into(),
3990 ),
3991 }
3992 }
3993 Data::F64(v) => {
3994 Data::F64(par::map(v, |&x| if op == Double { x + x } else { x * x }).into())
3995 }
3996 _ => return Err(wrong_type(d.dtype(), span)),
3997 },
3998 Halve => {
3999 let v = as_f64(d, &mut tmp, span)?;
4000 Data::F64(par::map(v, |&x| x / 2.0).into())
4001 }
4002 Pi => {
4003 let v = as_f64(d, &mut tmp, span)?;
4004 Data::F64(par::map(v, |&x| std::f64::consts::PI * x).into())
4005 }
4006 Factorial => {
4007 let v = as_f64(d, &mut tmp, span)?;
4008 Data::F64(par::map(v, |&x| factorial(x)).into())
4009 }
4010 Ln => {
4011 let v = as_f64(d, &mut tmp, span)?;
4013 Data::F64(par::map(v, |&x| x.ln()).into())
4015 }
4016 OneMinus => match d {
4017 Data::Bool(v) => Data::Bool(par::map(v, |&b| 1 - b).into()),
4018 Data::I64(v) => match par::try_map(v, |x: i64| 1i64.checked_sub(x)) {
4019 Some(out) => Data::I64(out.into()),
4020 None => Data::F64(par::map(v, |&x| 1.0 - x as f64).into()),
4021 },
4022 Data::F64(v) => Data::F64(par::map(v, |&x| 1.0 - x).into()),
4023 _ => return Err(wrong_type(d.dtype(), span)),
4024 },
4025 Not => {
4026 let bad = || Error::domain("logical negation needs values of 0 or 1", span);
4027 match d {
4028 Data::Bool(v) => Data::Bool(par::map(v, |&b| 1 - b).into()),
4029 Data::I64(v) => {
4030 let out = par::try_map(v, |x: i64| match x {
4031 0 => Some(1u8),
4032 1 => Some(0u8),
4033 _ => None,
4034 })
4035 .ok_or_else(bad)?;
4036 Data::Bool(out.into())
4037 }
4038 Data::F64(v) => {
4039 let out = par::try_map(v, |x: f64| {
4040 if x == 0.0 {
4041 Some(1u8)
4042 } else if x == 1.0 {
4043 Some(0u8)
4044 } else {
4045 None
4046 }
4047 })
4048 .ok_or_else(bad)?;
4049 Data::Bool(out.into())
4050 }
4051 _ => return Err(bad()),
4052 }
4053 }
4054 };
4055 Ok(Array::new(y.shape.clone(), data).with_layout(y.layout()))
4056}
4057
4058fn transpose_axes(y: &Array) -> Array {
4068 if y.rank() < 2 {
4069 return y.clone();
4070 }
4071 let out_shape: Vec<usize> = y.shape.iter().rev().copied().collect();
4072 let flipped = match y.layout() {
4073 Layout::RowMajor => Layout::ColMajor,
4074 Layout::ColMajor => Layout::RowMajor,
4075 };
4076 Array::new(out_shape, y.data.clone()).with_layout(flipped)
4077}
4078
4079fn iota_j(y: &Array, span: Span) -> Result<Array> {
4082 if y.rank() > 1 {
4083 return Err(Error::new(
4084 ErrorKind::Rank,
4085 "index generator needs a scalar or vector argument",
4086 Some(span),
4087 ));
4088 }
4089 let dims = y
4090 .to_i64_vec()
4091 .ok_or_else(|| Error::domain("index generator needs integer lengths", span))?;
4092 let shape: Vec<usize> = dims.iter().map(|d| d.unsigned_abs() as usize).collect();
4093 let n = crate::limits::elements(&shape, span)?;
4094 let st = strides(&shape);
4095 let mut out = Vec::with_capacity(n);
4096 let mut coord = vec![0usize; shape.len()];
4097 for _ in 0..n {
4098 let mut v = 0usize;
4099 for k in 0..shape.len() {
4100 let c = if dims[k] < 0 { shape[k] - 1 - coord[k] } else { coord[k] };
4101 v += c * st[k];
4102 }
4103 out.push(v as i64);
4104 odometer(&mut coord, &shape);
4105 }
4106 let data = Data::I64(out.into());
4107 let data = if y.dtype() == DType::Ext {
4110 data.cast(DType::Ext).ok_or_else(|| Error::internal("integers have no extended form"))?
4111 } else {
4112 data
4113 };
4114 Ok(Array::new(shape, data))
4115}
4116
4117fn head(y: &Array) -> Array {
4119 if y.rank() == 0 {
4120 return y.clone();
4121 }
4122 if y.items() == 0 {
4123 let cell_shape = y.shape[1..].to_vec();
4124 let n: usize = cell_shape.iter().product();
4125 return Array::new(cell_shape, fill_data(y.dtype(), n));
4126 }
4127 y.item(0)
4128}
4129
4130fn behead(y: &Array, span: Span) -> Result<Array> {
4131 if y.rank() == 0 {
4132 return Err(Error::domain("cannot drop the first item of a scalar", span));
4133 }
4134 if y.items() == 0 {
4135 return Ok(y.clone());
4136 }
4137 let m = y.item_size();
4138 let mut shape = y.shape.clone();
4139 shape[0] -= 1;
4140 Ok(Array::new(shape, y.data.slice(m, y.count())))
4141}
4142
4143fn tail(y: &Array) -> Array {
4145 if y.rank() == 0 {
4146 return y.clone();
4147 }
4148 let n = y.items();
4149 if n == 0 {
4150 let cell_shape = y.shape[1..].to_vec();
4151 let m: usize = cell_shape.iter().product();
4152 return Array::new(cell_shape, fill_data(y.dtype(), m));
4153 }
4154 y.item(n - 1)
4155}
4156
4157fn curtail(y: &Array) -> Array {
4159 if y.rank() == 0 {
4160 return Array::empty(y.dtype());
4161 }
4162 let n = y.items();
4163 if n == 0 {
4164 return y.clone();
4165 }
4166 let m = y.item_size();
4167 let mut shape = y.shape.clone();
4168 shape[0] = n - 1;
4169 Array::new(shape, y.data.slice(0, (n - 1) * m))
4170}
4171
4172fn reverse(y: &Array) -> Array {
4174 if y.rank() == 0 {
4175 return y.clone();
4176 }
4177 let n = y.items();
4178 let m = y.item_size();
4179 let mut data = Data::empty(y.dtype());
4180 for i in (0..n).rev() {
4181 for k in 0..m {
4182 push_elem(&mut data, &y.data, i * m + k);
4183 }
4184 }
4185 Array::new(y.shape.clone(), data)
4186}
4187
4188fn rotate(x: &Array, y: &Array, span: Span) -> Result<Array> {
4191 let counts = axis_counts(x, "rotate", span)?;
4192 if y.rank() == 0 {
4193 return Ok(y.clone());
4194 }
4195 if counts.len() > y.rank() {
4196 return Err(Error::new(
4197 ErrorKind::Length,
4198 format!(
4199 "rotate has {} amounts for an argument of rank {}",
4200 counts.len(),
4201 y.rank()
4202 ),
4203 Some(span),
4204 ));
4205 }
4206 let st = strides(&y.shape);
4207 let n = y.count();
4208 let r = y.rank();
4209 let mut data = Data::empty(y.dtype());
4210 let mut coord = vec![0usize; r];
4211 for _ in 0..n {
4212 let mut idx = 0usize;
4213 for k in 0..r {
4214 let len = y.shape[k] as i64;
4216 let s = counts.get(k).copied().unwrap_or(0);
4217 idx += (coord[k] as i64 + s).rem_euclid(len) as usize * st[k];
4218 }
4219 push_elem(&mut data, &y.data, idx);
4220 odometer(&mut coord, &y.shape);
4221 }
4222 Ok(Array::new(y.shape.clone(), data))
4223}
4224
4225fn elem_key(d: &Data, i: usize) -> u64 {
4228 match d {
4229 Data::Bool(v) => v[i] as u64,
4230 Data::I64(v) => v[i] as u64,
4231 Data::F64(v) => {
4232 let x = v[i];
4233 if x == 0.0 { 0 } else { x.to_bits() }
4234 }
4235 Data::Complex(v) => cx_key(v[i]),
4236 Data::Char(v) => v[i] as u64,
4237 Data::Ext(_) | Data::Rat(_) | Data::Box(_) => 0,
4240 }
4241}
4242
4243fn num_key(d: &Data, i: usize) -> u64 {
4246 match d {
4247 Data::Bool(v) => (v[i] as f64).to_bits(),
4248 Data::I64(v) => (v[i] as f64).to_bits(),
4249 Data::F64(v) => {
4250 let x = v[i];
4251 if x == 0.0 { 0.0f64.to_bits() } else { x.to_bits() }
4252 }
4253 Data::Complex(v) => cx_key(v[i]),
4254 Data::Char(v) => v[i] as u64,
4255 Data::Ext(_) | Data::Rat(_) | Data::Box(_) => 0,
4257 }
4258}
4259
4260fn cx_key(z: Cx) -> u64 {
4262 let bits = |x: f64| if x == 0.0 { 0u64 } else { x.to_bits() };
4263 bits(z[0]) ^ bits(z[1]).rotate_left(32)
4264}
4265
4266fn nub(y: &Array, tol: Tol) -> Array {
4268 if y.rank() == 0 {
4269 return Array::new(vec![1], y.data.clone());
4270 }
4271 let n = y.items();
4272 let m = y.item_size();
4273 let mut keep = Vec::new();
4274 if y.dtype() == DType::Box || y.dtype().is_exact() {
4275 for i in 0..n {
4278 if !keep.iter().any(|&j| arrays_match(&y.item(i), &y.item(j), tol)) {
4279 keep.push(i);
4280 }
4281 }
4282 } else if y.dtype() == DType::F64 && tol.ct != 0.0 {
4283 let mut tv = Vec::new();
4286 let v = borrow_f64(&y.data, &mut tv);
4287 for i in 0..n {
4288 if !keep.iter().any(|&j| (0..m).all(|k| tol.eq(v[i * m + k], v[j * m + k]))) {
4289 keep.push(i);
4290 }
4291 }
4292 } else {
4293 let mut seen: HashSet<Vec<u64>> = HashSet::with_capacity(n);
4294 for i in 0..n {
4295 let key: Vec<u64> = (0..m).map(|k| elem_key(&y.data, i * m + k)).collect();
4296 if seen.insert(key) {
4297 keep.push(i);
4298 }
4299 }
4300 }
4301 let mut data = Data::empty(y.dtype());
4302 for &i in &keep {
4303 for k in 0..m {
4304 push_elem(&mut data, &y.data, i * m + k);
4305 }
4306 }
4307 let mut shape = y.shape.clone();
4308 shape[0] = keep.len();
4309 Array::new(shape, data)
4310}
4311
4312#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4317enum Tao {
4318 J,
4319 Apl2,
4320}
4321
4322impl Tao {
4323 fn of(rules: Rules) -> Tao {
4324 match rules.lang {
4325 crate::Lang::J => Tao::J,
4326 crate::Lang::Apl => Tao::Apl2,
4329 }
4330 }
4331
4332 fn class(self, dt: DType) -> u8 {
4336 match self {
4337 Tao::J => match dt {
4338 DType::Char => 2,
4339 DType::Box => 3,
4340 _ => 0,
4341 },
4342 Tao::Apl2 => match dt {
4343 DType::Char => 0,
4344 DType::Box => 2,
4345 _ => 1,
4346 },
4347 }
4348 }
4349}
4350
4351fn cmp_items_total(x: &Array, y: &Array, tao: Tao) -> std::cmp::Ordering {
4366 use std::cmp::Ordering::Equal;
4367 match tao {
4368 Tao::J => {
4369 let class = |a: &Array| if a.count() == 0 { 0 } else { tao.class(a.dtype()) };
4370 class(x)
4371 .cmp(&class(y))
4372 .then_with(|| x.rank().cmp(&y.rank()))
4373 .then_with(|| x.shape.iter().rev().cmp(y.shape.iter().rev()))
4374 .then_with(|| cmp_atoms(x, y, tao))
4375 }
4376 Tao::Apl2 => x
4377 .rank()
4378 .cmp(&y.rank())
4379 .then_with(|| x.shape.iter().cmp(y.shape.iter()))
4380 .then_with(|| cmp_atoms(x, y, tao))
4381 .then_with(|| {
4382 if x.count() == 0 {
4383 tao.class(x.dtype()).cmp(&tao.class(y.dtype()))
4384 } else {
4385 Equal
4386 }
4387 }),
4388 }
4389}
4390
4391fn cmp_atoms(x: &Array, y: &Array, tao: Tao) -> std::cmp::Ordering {
4394 use std::cmp::Ordering::Equal;
4395 let n = x.count();
4396 if n == 0 {
4397 return Equal;
4398 }
4399 let (xr, yr) = (x.to_row_major(), y.to_row_major());
4400 let (dx, dy) = (xr.row_major_data(), yr.row_major_data());
4401 let opened = |d: &Data, i: usize| -> Array {
4402 match d {
4403 Data::Box(v) => v[i].clone(),
4404 _ => {
4405 let mut one = Data::empty(d.dtype());
4406 push_elem(&mut one, d, i);
4407 Array::new(vec![], one)
4408 }
4409 }
4410 };
4411 if matches!(dx, Data::Box(_)) || matches!(dy, Data::Box(_)) {
4412 return (0..n)
4413 .map(|i| cmp_items_total(&opened(dx, i), &opened(dy, i), tao))
4414 .find(|o| *o != Equal)
4415 .unwrap_or(Equal);
4416 }
4417 let classes = tao.class(dx.dtype()).cmp(&tao.class(dy.dtype()));
4419 if classes != Equal {
4420 return classes;
4421 }
4422 match (dx, dy) {
4423 (Data::Char(a), Data::Char(b)) => a[..n].cmp(&b[..n]),
4424 _ => cmp_numbers(dx, dy, n),
4425 }
4426}
4427
4428fn cmp_numbers(dx: &Data, dy: &Data, n: usize) -> std::cmp::Ordering {
4432 use std::cmp::Ordering::Equal;
4433 let seek = |f: &dyn Fn(usize) -> std::cmp::Ordering| {
4434 (0..n).map(f).find(|o| *o != Equal).unwrap_or(Equal)
4435 };
4436 match DType::promote(dx.dtype(), dy.dtype()) {
4437 Some(DType::Complex) => {
4438 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4439 let (a, b) = (borrow_cx(dx, &mut ta), borrow_cx(dy, &mut tb));
4440 seek(&|k| {
4441 a[k][0]
4442 .partial_cmp(&b[k][0])
4443 .unwrap_or(Equal)
4444 .then_with(|| a[k][1].partial_cmp(&b[k][1]).unwrap_or(Equal))
4445 })
4446 }
4447 Some(DType::F64) => {
4448 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4449 let (a, b) = (borrow_f64(dx, &mut ta), borrow_f64(dy, &mut tb));
4450 seek(&|k| a[k].partial_cmp(&b[k]).unwrap_or(Equal))
4451 }
4452 Some(t) if t.is_exact() => match (to_rat_vec(dx), to_rat_vec(dy)) {
4453 (Some(a), Some(b)) => seek(&|k| a[k].cmp(&b[k])),
4454 _ => Equal,
4455 },
4456 None => Equal,
4458 Some(_) => {
4459 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4460 let (a, b) = (borrow_i64(dx, &mut ta), borrow_i64(dy, &mut tb));
4461 seek(&|k| a[k].cmp(&b[k]))
4462 }
4463 }
4464}
4465
4466fn cmp_items(d: &Data, i: usize, j: usize, m: usize, tao: Tao) -> std::cmp::Ordering {
4470 use std::cmp::Ordering::Equal;
4471 let (a, b) = (i * m, j * m);
4472 let ord = |k: usize| match d {
4473 Data::Bool(v) => v[a + k].cmp(&v[b + k]),
4474 Data::I64(v) => v[a + k].cmp(&v[b + k]),
4475 Data::F64(v) => v[a + k].partial_cmp(&v[b + k]).unwrap_or(Equal),
4476 Data::Complex(v) => v[a + k][0]
4483 .partial_cmp(&v[b + k][0])
4484 .unwrap_or(Equal)
4485 .then_with(|| v[a + k][1].partial_cmp(&v[b + k][1]).unwrap_or(Equal)),
4486 Data::Char(v) => v[a + k].cmp(&v[b + k]),
4487 Data::Ext(v) => v[a + k].cmp(&v[b + k]),
4490 Data::Rat(v) => v[a + k].cmp(&v[b + k]),
4491 Data::Box(v) => cmp_items_total(&v[a + k], &v[b + k], tao),
4494 };
4495 (0..m).map(ord).find(|o| *o != Equal).unwrap_or(Equal)
4496}
4497
4498fn grade_order(y: &Array, down: bool, tao: Tao) -> Vec<usize> {
4500 if y.rank() == 0 {
4501 return vec![0];
4502 }
4503 let n = y.items();
4504 let m = y.item_size();
4505 let mut idx: Vec<usize> = (0..n).collect();
4506 if down {
4509 idx.sort_by(|&a, &b| cmp_items(&y.data, b, a, m, tao));
4510 } else {
4511 idx.sort_by(|&a, &b| cmp_items(&y.data, a, b, m, tao));
4512 }
4513 idx
4514}
4515
4516fn collate_grade(x: &Array, y: &Array, down: bool, origin: i64, span: Span) -> Result<Array> {
4521 let chars_of = |a: &Array| -> Result<Vec<char>> {
4522 match a.row_major_data() {
4523 Data::Char(v) => Ok(v.as_slice().to_vec()),
4524 _ => Err(Error::domain("a collating grade takes characters", span)),
4525 }
4526 };
4527 let (xs, ys) = (chars_of(x)?, chars_of(y)?);
4528 let xshape = if x.rank() == 0 { vec![1] } else { x.shape.clone() };
4529 let width = xshape.len();
4530 let absent: Vec<usize> = xshape.iter().rev().copied().collect();
4534 let mut keys: std::collections::HashMap<char, Vec<usize>> =
4535 std::collections::HashMap::new();
4536 let xst = strides(&xshape);
4537 for (i, &c) in xs.iter().enumerate() {
4538 keys.entry(c).or_insert_with(|| {
4539 (0..width).map(|a| (i / xst[a]) % xshape[a]).rev().collect()
4540 });
4541 }
4542 let key_of = |c: char| keys.get(&c).unwrap_or(&absent).clone();
4543 let n = if y.rank() == 0 { 1 } else { y.items() };
4544 let m = if n == 0 { 0 } else { ys.len() / n };
4545 let item_keys: Vec<Vec<usize>> = (0..n)
4546 .map(|i| ys[i * m..(i + 1) * m].iter().flat_map(|&c| key_of(c)).collect())
4547 .collect();
4548 let mut idx: Vec<usize> = (0..n).collect();
4549 if down {
4550 idx.sort_by(|&a, &b| item_keys[b].cmp(&item_keys[a]));
4551 } else {
4552 idx.sort_by(|&a, &b| item_keys[a].cmp(&item_keys[b]));
4553 }
4554 Ok(Array::from_i64(idx.into_iter().map(|i| origin + i as i64).collect()))
4555}
4556
4557fn atomic_rep(y: &Array, ctx: &Ctx<'_>, span: Span) -> Result<Array> {
4561 let name = match y.as_boxes() {
4562 Some([b]) if y.rank() == 0 => crate::gerund::text_of(b),
4563 _ => None,
4564 };
4565 let Some(name) = name else {
4566 return Err(Error::domain("5!:1 takes a boxed name", span));
4567 };
4568 if let Some(v) = ctx.env.verb(&name) {
4569 let ar = crate::gerund::verb_ar(v).ok_or_else(|| {
4570 Error::not_yet(
4571 format!("the atomic representation of {}", v.name()),
4572 span,
4573 )
4574 })?;
4575 return Ok(Array::boxed(ar.to_array()));
4576 }
4577 match ctx.env.get(&name) {
4578 Some(a) => Ok(Array::boxed(crate::gerund::Ar::Noun(a).to_array())),
4579 None => Err(Error::new(
4580 ErrorKind::Value,
4581 format!("undefined name: {name}"),
4582 Some(span),
4583 )),
4584 }
4585}
4586
4587fn catalogue(y: &Array, span: Span) -> Result<Array> {
4591 let items = if y.rank() == 0 { vec![y.clone()] } else { y.cells(1) };
4592 let opened: Vec<Array> = items
4594 .iter()
4595 .map(|it| match it.as_boxes() {
4596 Some(bs) if it.rank() == 0 => bs[0].clone(),
4597 _ => it.clone(),
4598 })
4599 .collect();
4600 let mut shape: Vec<usize> = Vec::new();
4601 for o in &opened {
4602 shape.extend_from_slice(&o.shape);
4603 }
4604 let total: usize = shape.iter().product();
4605 let mut out = Vec::with_capacity(total);
4606 let mut coord = vec![0usize; shape.len()];
4607 for _ in 0..total {
4608 let mut at = 0usize;
4609 let mut picks = Vec::with_capacity(opened.len());
4610 for o in &opened {
4611 let st = strides(&o.shape);
4612 let idx: usize = (0..o.rank()).map(|a| coord[at + a] * st[a]).sum();
4613 at += o.rank();
4614 let mut data = Data::empty(o.dtype());
4615 push_elem(&mut data, o.row_major_data(), idx);
4616 picks.push(Array::new(vec![], data));
4617 }
4618 out.push(assemble(&[picks.len()], picks, span)?);
4619 odometer(&mut coord, &shape);
4620 }
4621 Ok(Array::new(shape, Data::Box(out.into())))
4622}
4623
4624fn raze_in(y: &Array, tol: Tol, span: Span) -> Result<Array> {
4627 let all = raze(y, span)?;
4628 let n = if all.rank() == 0 { 1 } else { all.items() };
4629 let elements: Vec<Array> = (0..y.count())
4630 .map(|i| {
4631 let mut data = Data::empty(y.dtype());
4632 push_elem(&mut data, y.row_major_data(), i);
4633 let one = Array::new(vec![], data);
4634 match one.as_boxes() {
4635 Some(bs) => bs[0].clone(),
4636 None => one,
4637 }
4638 })
4639 .collect();
4640 let mut out = Vec::with_capacity(elements.len() * n);
4641 for e in &elements {
4642 let row = member_j(&all, e, tol);
4643 out.extend_from_slice(row.to_i64_vec().unwrap_or_default().as_slice());
4644 }
4645 let mut shape = y.shape.clone();
4646 shape.push(n);
4647 Ok(Array::new(shape, Data::Bool(out.into_iter().map(|v| v as u8).collect::<Vec<u8>>().into())))
4648}
4649
4650fn select_items(y: &Array, order: &[usize]) -> Array {
4652 let m = y.item_size();
4653 let mut data = Data::empty(y.dtype());
4654 for &i in order {
4655 for k in 0..m {
4656 push_elem(&mut data, &y.data, i * m + k);
4657 }
4658 }
4659 let mut shape = y.shape.clone();
4660 shape[0] = order.len();
4661 Array::new(shape, data)
4662}
4663
4664fn check_gradable(y: &Array, rules: Rules, span: Span) -> Result<()> {
4669 if y.dtype() == DType::Complex && rules.complex_order != ComplexOrder::RealThenImaginary {
4670 return Err(Error::not_yet("grading complex values by magnitude and angle", span));
4671 }
4672 Ok(())
4673}
4674
4675fn grade_select(x: &Array, y: &Array, down: bool, rules: Rules, span: Span) -> Result<Array> {
4679 check_gradable(y, rules, span)?;
4680 let order = grade_order(y, down, Tao::of(rules));
4681 if x.rank() == 0 {
4682 return Ok(x.clone());
4683 }
4684 if let Some(&past) = order.iter().find(|&&i| i >= x.items()) {
4685 return Err(Error::domain(
4686 format!("index {past} is out of range: the argument has {} items", x.items()),
4687 span,
4688 ));
4689 }
4690 Ok(select_items(x, &order))
4691}
4692
4693pub(crate) fn arrays_match(x: &Array, y: &Array, tol: Tol) -> bool {
4696 if x.shape != y.shape {
4697 return false;
4698 }
4699 if x.layout() != y.layout() {
4702 return arrays_match(&x.to_row_major(), &y.to_row_major(), tol);
4703 }
4704 if x.count() == 0 {
4707 return true;
4708 }
4709 if let (Data::Box(a), Data::Box(b)) = (&x.data, &y.data) {
4710 return a.iter().zip(b.iter()).all(|(p, q)| arrays_match(p, q, tol));
4711 }
4712 let (dx, dy) = (x.dtype(), y.dtype());
4713 match DType::promote(dx, dy) {
4714 None => false,
4715 Some(DType::Char) => match (&x.data, &y.data) {
4716 (Data::Char(a), Data::Char(b)) => a.as_slice() == b.as_slice(),
4717 _ => false,
4718 },
4719 Some(DType::F64) => {
4720 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4721 let a = borrow_f64(&x.data, &mut ta);
4722 let b = borrow_f64(&y.data, &mut tb);
4723 a.iter().zip(b).all(|(p, q)| tol.eq(*p, *q))
4724 }
4725 Some(DType::Complex) => {
4726 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4727 let a = borrow_cx(&x.data, &mut ta);
4728 let b = borrow_cx(&y.data, &mut tb);
4729 a.iter().zip(b).all(|(p, q)| tol.eq_cx(*p, *q))
4730 }
4731 Some(t) if t.is_exact() => match (to_rat_vec(&x.data), to_rat_vec(&y.data)) {
4732 (Some(a), Some(b)) => a == b,
4733 _ => false,
4734 },
4735 Some(_) => {
4736 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4737 let a = borrow_i64(&x.data, &mut ta);
4738 let b = borrow_i64(&y.data, &mut tb);
4739 a.iter().zip(b).all(|(p, q)| p == q)
4740 }
4741 }
4742}
4743
4744fn item_or_self(a: &Array, i: usize) -> Array {
4746 if a.rank() == 0 { a.clone() } else { a.item(i) }
4747}
4748
4749fn member_j(x: &Array, y: &Array, tol: Tol) -> Array {
4752 let cell_rank = y.rank().saturating_sub(1).min(x.rank());
4753 let frame_rank = x.rank() - cell_rank;
4754 let frame: Vec<usize> = x.shape[..frame_rank].to_vec();
4755 let nf: usize = frame.iter().product();
4756 let items = y.items();
4757 let mut out = Vec::with_capacity(nf);
4758 for i in 0..nf {
4759 let cell = x.cell_at(frame_rank, i);
4760 out.push((0..items).any(|j| arrays_match(&cell, &item_or_self(y, j), tol)) as u8);
4761 }
4762 Array::new(frame, Data::Bool(out.into()))
4763}
4764
4765fn member_apl(x: &Array, y: &Array, tol: Tol) -> Array {
4767 let n = x.count();
4768 if x.dtype() == DType::Box
4769 || y.dtype() == DType::Box
4770 || x.dtype().is_exact()
4771 || y.dtype().is_exact()
4772 {
4773 let opened = |a: &Array, i: usize| -> Array {
4779 let e = atom(a, i);
4780 match e.as_boxes() {
4781 Some([b]) if b.rank() == 0 && b.dtype() != DType::Box => b.clone(),
4782 _ => e,
4783 }
4784 };
4785 let out: Vec<u8> = (0..n)
4786 .map(|i| {
4787 let e = opened(x, i);
4788 u8::from((0..y.count()).any(|j| arrays_match(&e, &opened(y, j), tol)))
4789 })
4790 .collect();
4791 return Array::new(x.shape.clone(), Data::Bool(out.into()));
4792 }
4793 if (x.dtype() == DType::Char) != (y.dtype() == DType::Char) {
4794 return Array::new(x.shape.clone(), Data::Bool(vec![0u8; n].into()));
4795 }
4796 if tol.ct != 0.0
4797 && (x.dtype() == DType::F64 || y.dtype() == DType::F64)
4798 && x.dtype() != DType::Char
4799 {
4800 let (mut tx, mut ty) = (Vec::new(), Vec::new());
4802 let xs = borrow_f64(&x.data, &mut tx);
4803 let ys = borrow_f64(&y.data, &mut ty);
4804 let out: Vec<u8> =
4805 xs.iter().map(|a| ys.iter().any(|b| tol.eq(*a, *b)) as u8).collect();
4806 return Array::new(x.shape.clone(), Data::Bool(out.into()));
4807 }
4808 let seen: HashSet<u64> = (0..y.count()).map(|i| num_key(&y.data, i)).collect();
4809 let out: Vec<u8> =
4810 (0..n).map(|i| seen.contains(&num_key(&x.data, i)) as u8).collect();
4811 Array::new(x.shape.clone(), Data::Bool(out.into()))
4812}
4813
4814fn index_of(x: &Array, y: &Array, origin: i64, tol: Tol) -> Array {
4816 let cell_rank = x.rank().saturating_sub(1).min(y.rank());
4817 let frame_rank = y.rank() - cell_rank;
4818 let frame: Vec<usize> = y.shape[..frame_rank].to_vec();
4819 let nf: usize = frame.iter().product();
4820 let items = x.items();
4821 let mut out = Vec::with_capacity(nf);
4822 for i in 0..nf {
4823 let cell = y.cell_at(frame_rank, i);
4824 let at = (0..items)
4825 .find(|&j| arrays_match(&cell, &item_or_self(x, j), tol))
4826 .unwrap_or(items);
4827 out.push(origin + at as i64);
4828 }
4829 Array::new(frame, Data::I64(out.into()))
4830}
4831
4832fn from_index(x: &Array, y: &Array, span: Span) -> Result<Array> {
4834 if let Some(spec) = x.as_boxes().and_then(<[Array]>::first) {
4837 let spec = index_spec(spec, y, span)?;
4838 return Ok(select_spec(&spec, y));
4839 }
4840 let idx = x
4841 .to_i64_vec()
4842 .ok_or_else(|| Error::domain("index must be an integer", span))?;
4843 let Some(&i) = idx.first() else {
4844 return Err(Error::internal("from_index with no index"));
4845 };
4846 let n = y.items() as i64;
4847 let k = if i < 0 { i + n } else { i };
4848 if k < 0 || k >= n {
4849 return Err(Error::domain(
4850 format!("index {i} is out of range: the argument has {n} items"),
4851 span,
4852 ));
4853 }
4854 Ok(item_or_self(y, k as usize))
4855}
4856
4857fn cat_promote(a: &Array, other: &Array, rank: usize, axis: usize, span: Span) -> Result<Array> {
4861 if a.rank() == rank {
4862 return Ok(a.clone());
4863 }
4864 if a.rank() == 0 {
4865 let mut shape =
4866 if other.rank() == rank { other.shape.clone() } else { vec![1usize; rank] };
4867 shape[axis] = 1;
4868 let n: usize = shape.iter().product();
4869 let mut data = Data::empty(a.dtype());
4870 for _ in 0..n {
4871 push_elem(&mut data, &a.data, 0);
4872 }
4873 return Ok(Array::new(shape, data));
4874 }
4875 if a.rank() + 1 == rank {
4876 let mut shape = a.shape.clone();
4877 shape.insert(axis, 1);
4878 return Ok(Array::new(shape, a.data.clone()));
4879 }
4880 Err(Error::new(
4881 ErrorKind::Rank,
4882 format!("cannot catenate rank {} with rank {}", a.rank(), other.rank()),
4883 Some(span),
4884 ))
4885}
4886
4887pub(crate) fn catenate(
4889 x: &Array,
4890 y: &Array,
4891 leading: bool,
4892 fill: bool,
4893 span: Span,
4894) -> Result<Array> {
4895 let rank = x.rank().max(y.rank()).max(1);
4896 let axis = if leading { 0 } else { rank - 1 };
4897 let xa = cat_promote(x, y, rank, axis, span)?;
4898 let ya = cat_promote(y, x, rank, axis, span)?;
4899 let mut ragged = false;
4903 let want: Vec<i64> = (0..rank)
4904 .map(|k| {
4905 ragged |= k != axis && xa.shape[k] != ya.shape[k];
4906 xa.shape[k].max(ya.shape[k]) as i64
4907 })
4908 .collect();
4909 if ragged && !fill {
4910 return Err(Error::new(
4911 ErrorKind::Length,
4912 format!(
4913 "cannot catenate: left shape {}, right shape {}",
4914 show_shape(&xa.shape),
4915 show_shape(&ya.shape)
4916 ),
4917 Some(span),
4918 ));
4919 }
4920 let (xa, ya) = if ragged {
4921 let fit = |a: &Array| -> Result<Array> {
4922 let mut to = want.clone();
4923 to[axis] = a.shape[axis] as i64;
4924 take(&Array::from_i64(to), a, false, false, span)
4925 };
4926 (fit(&xa)?, fit(&ya)?)
4927 } else {
4928 (xa, ya)
4929 };
4930 let (xa, ya) = if !fill && (xa.dtype() == DType::Box) != (ya.dtype() == DType::Box) {
4934 (nest_like(&xa, &ya), nest_like(&ya, &xa))
4935 } else {
4936 (xa, ya)
4937 };
4938 let dt = DType::promote(xa.dtype(), ya.dtype()).ok_or_else(|| {
4939 let boxed = xa.dtype() == DType::Box || ya.dtype() == DType::Box;
4940 let what = if boxed {
4941 "cannot catenate boxed and unboxed data; box the other side first"
4942 } else {
4943 "cannot catenate character and numeric data"
4944 };
4945 Error::new(ErrorKind::Type, what, Some(span))
4946 })?;
4947 let widen = |a: &Array| -> Result<Data> {
4948 if a.dtype() == dt {
4949 Ok(a.data.clone())
4950 } else {
4951 a.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening in catenate"))
4952 }
4953 };
4954 let xd = widen(&xa)?;
4955 let yd = widen(&ya)?;
4956 let outer: usize = xa.shape[..axis].iter().product();
4957 let ix: usize = xa.shape[axis..].iter().product();
4958 let iy: usize = ya.shape[axis..].iter().product();
4959 let mut data = Data::empty(dt);
4960 for o in 0..outer {
4961 for k in 0..ix {
4962 push_elem(&mut data, &xd, o * ix + k);
4963 }
4964 for k in 0..iy {
4965 push_elem(&mut data, &yd, o * iy + k);
4966 }
4967 }
4968 let mut shape = xa.shape.clone();
4969 shape[axis] = xa.shape[axis] + ya.shape[axis];
4970 Ok(Array::new(shape, data))
4971}
4972
4973fn copy_items(x: &Array, y: &Array, apl: bool, span: Span) -> Result<Array> {
4981 let counts = x
4982 .to_i64_vec()
4983 .ok_or_else(|| Error::domain("replication counts must be integers", span))?;
4984 if !apl && counts.iter().any(|&c| c < 0) {
4985 return Err(Error::domain("replication counts must be nonnegative", span));
4986 }
4987 let one_item = apl && x.rank() > 0 && y.rank() > 0 && y.items() == 1 && counts.len() != 1;
4991 let scalar_y = y.rank() == 0 || one_item;
4992 let m = y.item_size();
4993 let n = if x.rank() == 0 || !scalar_y { y.items() } else { counts.len() };
4994 let per = if x.rank() == 0 { vec![counts[0]; n] } else { counts };
4995 if per.len() != n {
4996 return Err(Error::new(
4997 ErrorKind::Length,
4998 format!("{} replication count(s) for {n} item(s)", per.len()),
4999 Some(span),
5000 ));
5001 }
5002 let items: u128 = per.iter().map(|&c| c.unsigned_abs() as u128).sum();
5005 let total = crate::limits::count(items * m.max(1) as u128, span)? / m.max(1);
5006 let mut data = Data::empty(y.dtype());
5007 for (i, &c) in per.iter().enumerate() {
5008 let src = if scalar_y { 0 } else { i };
5010 for _ in 0..c.unsigned_abs() {
5011 for k in 0..m {
5012 if c < 0 {
5013 data.push_fill();
5014 } else {
5015 push_elem(&mut data, &y.data, src * m + k);
5016 }
5017 }
5018 }
5019 }
5020 let mut shape = if y.rank() == 0 { vec![1] } else { y.shape.clone() };
5023 shape[0] = total;
5024 Ok(Array::new(shape, data))
5025}
5026
5027fn format_chars(y: &Array, opts: &FmtOpts) -> Array {
5036 if y.dtype() == DType::Char {
5037 return y.clone();
5038 }
5039 if y.count() == 0 {
5041 return Array::new(y.shape.clone(), Data::empty(DType::Char));
5042 }
5043 let text = crate::fmt::format_array(y, opts);
5044 if y.dtype() == DType::Box {
5045 let lines = text.lines().filter(|l| !l.is_empty()).count();
5050 let rows: usize =
5051 if y.rank() == 0 { 1 } else { y.shape[..y.rank() - 1].iter().product() };
5052 if lines != rows {
5053 return text_planes(&text, &y.shape[..y.rank().saturating_sub(2)]);
5054 }
5055 }
5056 if y.rank() < 2 {
5057 let chars: Vec<char> = text.chars().collect();
5058 return Array::new(vec![chars.len()], Data::Char(chars.into()));
5059 }
5060 let lines: Vec<&str> = text.lines().filter(|l| !l.is_empty()).collect();
5063 let width = lines.iter().map(|l| l.chars().count()).max().unwrap_or(0);
5064 let mut chars: Vec<char> = Vec::with_capacity(lines.len() * width);
5065 for line in &lines {
5066 chars.extend(line.chars());
5067 chars.resize(chars.len() + width - line.chars().count(), ' ');
5068 }
5069 let mut shape = y.shape[..y.rank() - 1].to_vec();
5072 shape.push(width);
5073 debug_assert_eq!(lines.len(), shape[..shape.len() - 1].iter().product::<usize>());
5074 Array::new(shape, Data::Char(chars.into()))
5075}
5076
5077fn text_planes(text: &str, frame: &[usize]) -> Array {
5080 let lines: Vec<&str> = text.lines().filter(|l| !l.is_empty()).collect();
5081 let width = lines.iter().map(|l| l.chars().count()).max().unwrap_or(0);
5082 let planes: usize = frame.iter().product::<usize>().max(1);
5083 let per = lines.len() / planes;
5084 let mut chars: Vec<char> = Vec::with_capacity(lines.len() * width);
5085 for line in &lines {
5086 chars.extend(line.chars());
5087 chars.resize(chars.len() + width - line.chars().count(), ' ');
5088 }
5089 let mut shape = frame.to_vec();
5090 shape.push(per);
5091 shape.push(width);
5092 Array::new(shape, Data::Char(chars.into()))
5093}
5094
5095fn digits_of(a: &Array, what: &str, span: Span) -> Result<Vec<f64>> {
5097 a.to_f64_vec().ok_or_else(|| Error::domain(format!("{what} needs numeric data"), span))
5098}
5099
5100fn narrow(values: Vec<f64>, integral: bool) -> Data {
5104 if integral && values.iter().all(|&v| v.fract() == 0.0 && fits_i64(v)) {
5105 return Data::I64(values.iter().map(|&v| v as i64).collect::<Vec<_>>().into());
5106 }
5107 Data::F64(values.into())
5108}
5109
5110fn is_integral(a: &Array) -> bool {
5112 !matches!(a.dtype(), DType::F64 | DType::Rat | DType::Char)
5113}
5114
5115fn decode_exact(x: Option<&Array>, y: &Array) -> Option<Array> {
5121 let yr = y.to_row_major();
5122 let digits = to_rat_vec(&yr.data)?;
5123 let two = Rat::from_int(Ext::from(2));
5124 let radix: Vec<Rat> = match x {
5125 None => vec![two; digits.len()],
5126 Some(x) => {
5127 let r = to_rat_vec(&x.to_row_major().data)?;
5128 match r.len() {
5129 1 => vec![r[0].clone(); digits.len()],
5130 n if n == digits.len() => r,
5131 _ => return None,
5132 }
5133 }
5134 };
5135 let mut acc = Rat::from_int(Ext::from(0));
5136 for (d, b) in digits.iter().zip(&radix) {
5137 acc = acc.mul(b).add(d);
5138 }
5139 let exact_in = |a: &Array| matches!(a.dtype(), DType::Ext | DType::Rat);
5140 if exact_in(y) || x.is_some_and(exact_in) {
5141 return Some(Array::new(Vec::new(), exact_data(DType::Ext, vec![acc])));
5142 }
5143 let whole = acc.to_int()?;
5146 Some(Array::scalar_i64(exact::ext_to_i64(&whole)?))
5147}
5148
5149fn decode(x: Option<&Array>, y: &Array, span: Span) -> Result<Array> {
5152 if let Some(exact) = decode_exact(x, y) {
5153 return Ok(exact);
5154 }
5155 let digits = digits_of(y, "decode", span)?;
5156 let radix: Vec<f64> = match x {
5157 None => vec![2.0; digits.len()],
5158 Some(x) => {
5159 let r = digits_of(x, "decode", span)?;
5160 match r.len() {
5161 1 => vec![r[0]; digits.len()],
5162 n if n == digits.len() => r,
5163 n => {
5164 return Err(Error::new(
5165 ErrorKind::Length,
5166 format!("{n} radices for {} digits", digits.len()),
5167 Some(span),
5168 ));
5169 }
5170 }
5171 }
5172 };
5173 let mut acc = 0.0f64;
5174 for (d, b) in digits.iter().zip(&radix) {
5175 acc = acc * b + d;
5176 }
5177 let integral = is_integral(y) && x.is_none_or(is_integral);
5178 Ok(Array::new(vec![], narrow(vec![acc], integral)))
5179}
5180
5181fn decode_apl(x: &Array, y: &Array, span: Span) -> Result<Array> {
5185 let digits = digits_of(y, "decode", span)?;
5186 let radices = digits_of(x, "decode", span)?;
5187 let k = if y.rank() == 0 { 1 } else { y.shape[0] };
5192 let n: usize = if y.rank() == 0 { 1 } else { y.shape[1..].iter().product() };
5193 let (rows, width) = match x.rank() {
5194 0 => (1usize, 0usize),
5195 r => (x.shape[..r - 1].iter().product(), x.shape[r - 1]),
5196 };
5197 if width != 0 && width != k {
5198 return Err(Error::new(
5199 ErrorKind::Length,
5200 format!("{width} radices for {k} digits"),
5201 Some(span),
5202 ));
5203 }
5204 let per_row = if x.rank() > 0 && width == 0 { 0 } else { k };
5208 let mut out = vec![0.0f64; rows * n];
5209 for i in 0..rows {
5210 for j in 0..n {
5211 let mut acc = 0.0f64;
5212 for d in 0..per_row {
5213 let b = if width == 0 { radices[0] } else { radices[i * width + d] };
5214 acc = acc * b + digits[d * n + j];
5215 }
5216 out[i * n + j] = acc;
5217 }
5218 }
5219 let mut shape: Vec<usize> = if x.rank() == 0 {
5220 Vec::new()
5221 } else {
5222 x.shape[..x.rank() - 1].to_vec()
5223 };
5224 if y.rank() > 0 {
5225 shape.extend_from_slice(&y.shape[1..]);
5226 }
5227 let integral = is_integral(y) && is_integral(x);
5228 Ok(Array::new(shape, narrow(out, integral)))
5229}
5230
5231fn encode_apl(x: &Array, y: &Array, span: Span) -> Result<Array> {
5234 let radices = digits_of(x, "encode", span)?;
5235 let values = digits_of(y, "encode", span)?;
5236 let k = if x.rank() == 0 { 1 } else { x.shape[0] };
5237 let frames = if k == 0 { 0 } else { radices.len() / k };
5238 let n = values.len();
5239 let mut out = vec![0.0f64; k * frames * n];
5240 let mut radix = vec![0.0f64; k];
5241 let mut cell = vec![0.0f64; k];
5242 for p in 0..frames {
5243 for (i, r) in radix.iter_mut().enumerate() {
5244 *r = radices[i * frames + p];
5245 }
5246 for (j, &v) in values.iter().enumerate() {
5247 encode_one(&radix, v, &mut cell);
5248 for i in 0..k {
5249 out[(i * frames + p) * n + j] = cell[i];
5250 }
5251 }
5252 }
5253 let mut shape = x.shape.clone();
5254 shape.extend_from_slice(&y.shape);
5255 Ok(Array::new(shape, narrow(out, is_integral(x) && is_integral(y))))
5256}
5257
5258fn bit_width(values: &[f64], span: Span) -> Result<usize> {
5261 if values.is_empty() {
5263 return Ok(0);
5264 }
5265 let mut m = 0.0f64;
5266 for &v in values {
5267 if !v.is_finite() {
5268 return Err(Error::domain("cannot encode an infinite value", span));
5269 }
5270 m = m.max(v.abs());
5271 }
5272 let whole = m.floor();
5273 if whole >= 1e15 {
5274 return Err(Error::domain("the value is too large to encode in binary", span));
5275 }
5276 let mut w = 1usize;
5277 let mut n = whole as i64;
5278 while n > 1 {
5279 n /= 2;
5280 w += 1;
5281 }
5282 Ok(w)
5283}
5284
5285fn encode_one(radix: &[f64], v: f64, out: &mut [f64]) {
5289 let mut rem = v;
5290 for i in (0..radix.len()).rev() {
5291 let b = radix[i];
5292 if b == 0.0 {
5293 out[i] = rem;
5294 rem = 0.0;
5295 } else {
5296 let r = rem - b * (rem / b).floor();
5297 out[i] = r;
5298 rem = (rem - r) / b;
5299 }
5300 }
5301}
5302
5303fn encode(x: &Array, y: &Array, span: Span) -> Result<Array> {
5307 let radix = digits_of(x, "encode", span)?;
5308 let values = digits_of(y, "encode", span)?;
5309 let k = radix.len();
5310 let n = values.len();
5311 let mut out = vec![0.0f64; k * n];
5312 let mut cell = vec![0.0f64; k];
5313 for (j, &v) in values.iter().enumerate() {
5314 encode_one(&radix, v, &mut cell);
5315 for i in 0..k {
5316 out[i * n + j] = cell[i];
5317 }
5318 }
5319 let mut shape = if x.rank() == 0 { Vec::new() } else { vec![k] };
5322 shape.extend_from_slice(&y.shape);
5323 Ok(Array::new(shape, narrow(out, is_integral(x) && is_integral(y))))
5324}
5325
5326fn encode_bits(y: &Array, span: Span) -> Result<Array> {
5328 let values = digits_of(y, "encode", span)?;
5329 let k = bit_width(&values, span)?;
5330 let radix = vec![2.0; k];
5331 let mut out = vec![0.0f64; values.len() * k];
5332 for (j, &v) in values.iter().enumerate() {
5333 encode_one(&radix, v, &mut out[j * k..(j + 1) * k]);
5334 }
5335 let mut shape = y.shape.clone();
5336 shape.push(k);
5337 Ok(Array::new(shape, narrow(out, is_integral(y))))
5338}
5339
5340fn laminate(x: &Array, y: &Array, span: Span) -> Result<Array> {
5345 let spread = |a: &Array, other: &Array| -> Array {
5346 if a.rank() != 0 {
5347 return a.clone();
5348 }
5349 let shape = if other.rank() == 0 { vec![1] } else { other.shape.clone() };
5350 let n: usize = shape.iter().product();
5351 let mut data = Data::empty(a.dtype());
5352 for _ in 0..n {
5353 push_elem(&mut data, &a.data, 0);
5354 }
5355 Array::new(shape, data)
5356 };
5357 assemble(&[2], vec![spread(x, y), spread(y, x)], span)
5358}
5359
5360fn table_of(y: &Array) -> Array {
5362 let shape = match y.rank() {
5363 0 => vec![1, 1],
5364 _ => vec![y.items(), y.item_size()],
5365 };
5366 Array::new(shape, y.data.clone())
5367}
5368
5369fn table(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
5375 let ranks = u.ranks();
5376 let fxl = x.rank() - effective_rank(ranks[1], x.rank());
5377 let fyl = y.rank() - effective_rank(ranks[2], y.rank());
5378 let mut frame = x.shape[..fxl].to_vec();
5379 frame.extend_from_slice(&y.shape[..fyl]);
5380 let nx: usize = x.shape[..fxl].iter().product();
5381 let ny: usize = y.shape[..fyl].iter().product();
5382 let n = nx * ny;
5383 if n == 0 {
5384 return assemble(&frame, Vec::new(), span);
5385 }
5386 if frame.is_empty() {
5387 return u.dyad(x, y, ctx, span);
5388 }
5389 let work = x.count().max(y.count()).max(n);
5390 let cells = each_cell(n, work, u.is_pure(), ctx, |i, c| {
5391 u.dyad(&x.cell_at(fxl, i / ny), &y.cell_at(fyl, i % ny), c, span)
5392 })?;
5393 assemble(&frame, cells, span)
5394}
5395
5396fn monad_op(p: &Prim, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
5398 match p.monad {
5399 MonadOp::Scalar(op) => scalar_monad(op, y, ctx.cfg, span),
5400 MonadOp::ShapeOf => {
5401 Ok(carry_exact(Array::from_i64(y.shape.iter().map(|&n| n as i64).collect()), y))
5402 }
5403 MonadOp::Tally => Ok(carry_exact(Array::scalar_i64(y.items() as i64), y)),
5404 MonadOp::Ravel => Ok(Array::new(vec![y.count()], y.data.clone())),
5405 MonadOp::TransposeAxes => Ok(transpose_axes(y)),
5406 MonadOp::Head => Ok(head(y)),
5407 MonadOp::Behead => behead(y, span),
5408 MonadOp::Tail => Ok(tail(y)),
5409 MonadOp::Curtail => Ok(curtail(y)),
5410 MonadOp::Reverse => Ok(reverse(y)),
5411 MonadOp::Nub => Ok(nub(y, ctx.cfg.tol)),
5414 MonadOp::GradeUp { origin } | MonadOp::GradeDown { origin } => {
5415 check_gradable(y, ctx.cfg.rules, span)?;
5416 if ctx.cfg.rules.lang == crate::Lang::Apl && y.rank() == 0 {
5419 return Err(Error::domain("a grade needs an array, not a scalar", span));
5420 }
5421 let down = matches!(p.monad, MonadOp::GradeDown { .. });
5422 let order = grade_order(y, down, Tao::of(ctx.cfg.rules));
5423 Ok(Array::from_i64(order.iter().map(|&i| origin + i as i64).collect()))
5424 }
5425 MonadOp::IotaJ => iota_j(y, span),
5426 MonadOp::IotaApl { origin } => iota_apl(y, origin, span),
5427 MonadOp::Echo => {
5428 (ctx.out)(&format!("{}\n", crate::fmt::format_array(y, &ctx.cfg.fmt)));
5429 Ok(Array::empty(DType::I64))
5430 }
5431 MonadOp::ReadStream => {
5432 stream_number(y, 1, "1!:1 reads", span)?;
5433 let line = ctx.read_line(span)?;
5434 Ok(Array::from_chars(line.chars().collect()))
5435 }
5436 MonadOp::TypeCode => Ok(Array::scalar_i64(type_code(y))),
5437 MonadOp::Same => Ok(y.clone()),
5438 MonadOp::Format => Ok(format_chars(y, &ctx.cfg.fmt)),
5439 MonadOp::DecodeBits => decode(None, y, span).map(|r| carry_exact(r, y)),
5440 MonadOp::EncodeBits => encode_bits(y, span).map(|r| carry_exact(r, y)),
5441 MonadOp::Itemize => {
5442 let mut shape = vec![1usize];
5443 shape.extend_from_slice(&y.shape);
5444 Ok(Array::new(shape, y.data.clone()))
5445 }
5446 MonadOp::TableOf => Ok(table_of(y)),
5447 MonadOp::Enclose(rule) => Ok(enclose(y, rule)),
5448 MonadOp::Open => Ok(open_cell(y)),
5449 MonadOp::Raze => raze(y, span),
5450 MonadOp::Catalogue => catalogue(y, span),
5451 MonadOp::AtomicRep => atomic_rep(y, ctx, span),
5452 MonadOp::RazeIn => raze_in(y, ctx.cfg.tol, span),
5453 MonadOp::First => Ok(first(y)),
5454 MonadOp::Enlist => enlist(y, span),
5455 MonadOp::Depth => Ok(Array::scalar_i64(depth(y))),
5456 MonadOp::Indices { origin, boxed_coords } => {
5457 where_indices(y, origin, boxed_coords, span)
5458 }
5459 MonadOp::Steps => steps(y, span),
5460 MonadOp::ToExact => to_exact(y, span),
5461 MonadOp::NthPrime => {
5462 let n = y
5463 .to_i64_vec()
5464 .ok_or_else(|| Error::domain("the prime index must be an integer", span))?;
5465 let v = n.first().copied().unwrap_or(0);
5466 Ok(carry_exact(Array::scalar_i64(nth_prime(v, span)?), y))
5467 }
5468 MonadOp::PrimeFactors => {
5469 let n = y
5470 .to_i64_vec()
5471 .ok_or_else(|| Error::domain("prime factors need an integer", span))?;
5472 let v = n.first().copied().unwrap_or(0);
5473 Ok(carry_exact(Array::from_i64(prime_factors(v, span)?), y))
5474 }
5475 MonadOp::MatrixInverse => matrix_inverse(y, span),
5476 MonadOp::Roll { origin, fixed, float_at_zero } => {
5477 roll(y, origin, fixed, float_at_zero, span)
5478 }
5479 MonadOp::ComplexParts { polar } => complex_parts(y, polar, span),
5480 MonadOp::SelfClassify => Ok(self_classify(y, ctx.cfg.tol)),
5481 MonadOp::NubSieve => Ok(nub_sieve(y, ctx.cfg.tol)),
5482 MonadOp::Unicode { pass_chars } => unicode(y, pass_chars, span),
5483 MonadOp::Words => words(y, span),
5484 MonadOp::LevelOf => Ok(Array::scalar_i64(boxing_level(y))),
5485 MonadOp::MapPaths => Ok(map_paths(y)),
5486 MonadOp::Nest => Ok(nest(y)),
5487 MonadOp::PolyRoots => poly_roots(y, span),
5488 MonadOp::PolyDeriv => poly_deriv(y, span),
5489 MonadOp::AnagramIndex => anagram_index(y, ctx.cfg.rules, span),
5490 MonadOp::CycleForm => cycle_form(y, span),
5491 MonadOp::Split => Ok(split_items(y)),
5492 MonadOp::Execute { apl } => execute(y, apl, ctx, span),
5493 MonadOp::NotYet(what) => Err(Error::not_yet(what, span)),
5494 MonadOp::None => {
5495 Err(Error::domain(format!("{} has no monadic meaning", p.name), span))
5496 }
5497 }
5498}
5499
5500fn complex_parts(y: &Array, polar: bool, span: Span) -> Result<Array> {
5504 let Some(v) = y.to_complex_vec() else {
5505 return Err(wrong_type(y.dtype(), span));
5506 };
5507 let z = v.first().copied().unwrap_or(cx::ZERO);
5508 let pair = if polar { vec![cx::abs(z), cx::arg(z)] } else { vec![z[0], z[1]] };
5509 Ok(Array::from_f64(pair))
5510}
5511
5512fn axis_counts(x: &Array, what: &str, span: Span) -> Result<Vec<i64>> {
5513 if x.rank() > 1 {
5514 return Err(Error::new(
5515 ErrorKind::Rank,
5516 format!("{what} needs a scalar or vector left argument"),
5517 Some(span),
5518 ));
5519 }
5520 if x.count() == 0 {
5523 return Ok(Vec::new());
5524 }
5525 x.to_i64_vec()
5526 .ok_or_else(|| Error::domain(format!("{what} needs integer lengths"), span))
5527}
5528
5529fn reshape(x: &Array, y: &Array, by_items: bool, span: Span) -> Result<Array> {
5540 let dims = axis_counts(x, "reshape", span)?;
5541 if dims.iter().any(|&d| d < 0) {
5542 return Err(Error::domain("reshape lengths must be nonnegative", span));
5543 }
5544 let mut shape: Vec<usize> = dims.iter().map(|&d| d as usize).collect();
5545 let (unit, src) = if by_items {
5547 let item_shape = if y.rank() == 0 { &[][..] } else { &y.shape[1..] };
5548 shape.extend_from_slice(item_shape);
5549 (item_shape.iter().product::<usize>(), y.items().max(usize::from(y.rank() == 0)))
5550 } else {
5551 (1, y.count())
5552 };
5553 let n = crate::limits::elements(&shape, span)?;
5554 let mut data = Data::empty(y.dtype());
5555 if n > 0 && src == 0 {
5556 if by_items {
5557 return Err(Error::new(ErrorKind::Length, "reshape of an empty array", Some(span)));
5558 }
5559 return Ok(Array::new(shape, fill_data(y.dtype(), n)));
5560 }
5561 for i in 0..n {
5562 push_elem(&mut data, &y.data, (i / unit) % src * unit + i % unit);
5565 }
5566 Ok(Array::new(shape, data))
5567}
5568
5569fn leading_run(y: &Array, counts: &[i64], drop: bool) -> Option<Array> {
5573 if y.rank() == 0 || counts.is_empty() {
5574 return None;
5575 }
5576 let trailing_untouched = counts[1..].iter().enumerate().all(|(a, &c)| {
5580 if drop { c == 0 } else { c.unsigned_abs() as usize == y.shape[a + 1] }
5581 });
5582 if !trailing_untouched {
5583 return None;
5584 }
5585 let n = y.items();
5586 let k = counts[0];
5587 let a = k.unsigned_abs() as usize;
5588 let (lo, keep) = if drop {
5589 let a = a.min(n);
5590 if k >= 0 { (a, n - a) } else { (0, n - a) }
5591 } else {
5592 if a > n {
5594 return None;
5595 }
5596 if k >= 0 { (0, a) } else { (n - a, a) }
5597 };
5598 Some(section(y, lo, lo + keep))
5599}
5600
5601fn count_rank(verb: &str, counts: usize, rank: usize, span: Span) -> Error {
5606 Error::new(
5607 ErrorKind::Length,
5608 format!("{counts} {verb} counts for a rank-{rank} argument"),
5609 Some(span),
5610 )
5611}
5612
5613fn take(x: &Array, y: &Array, prototype_fill: bool, apl: bool, span: Span) -> Result<Array> {
5614 let counts = axis_counts(x, "take", span)?;
5615 let fill = if prototype_fill { prototype_of(y) } else { None };
5619 let promoted;
5620 let base = if y.rank() == 0 {
5623 promoted = Array::new(vec![1; counts.len()], y.data.clone());
5624 &promoted
5625 } else {
5626 y
5627 };
5628 let wrong = if apl {
5630 counts.len() != base.rank()
5631 } else {
5632 counts.len() > base.rank() || (counts.is_empty() && base.rank() > 0)
5633 };
5634 if wrong {
5635 return Err(count_rank("take", counts.len(), base.rank(), span));
5636 }
5637 if let Some(run) = leading_run(base, &counts, false) {
5638 return Ok(run);
5639 }
5640 let mut out_shape = base.shape.clone();
5641 for (a, &k) in counts.iter().enumerate() {
5642 out_shape[a] = k.unsigned_abs() as usize;
5643 }
5644 let n = crate::limits::elements(&out_shape, span)?;
5645 let st = strides(&base.shape);
5646 let mut data = Data::empty(base.dtype());
5647 let mut coord = vec![0usize; out_shape.len()];
5648 for _ in 0..n {
5649 let mut idx = 0usize;
5650 let mut inside = true;
5651 for a in 0..out_shape.len() {
5652 let len = base.shape[a] as i64;
5653 let c = coord[a] as i64;
5654 let s = match counts.get(a) {
5657 Some(&k) if k < 0 => c + len - k.unsigned_abs() as i64,
5658 _ => c,
5659 };
5660 if s < 0 || s >= len {
5661 inside = false;
5662 break;
5663 }
5664 idx += s as usize * st[a];
5665 }
5666 if inside {
5667 push_elem(&mut data, &base.data, idx);
5668 } else if let (Data::Box(v), Some(p)) = (&mut data, &fill) {
5669 v.push(p.clone());
5670 } else {
5671 data.push_fill();
5672 }
5673 odometer(&mut coord, &out_shape);
5674 }
5675 Ok(Array::new(out_shape, data))
5676}
5677
5678fn prototype_of(y: &Array) -> Option<Array> {
5682 fn zeroed(a: &Array) -> Array {
5683 if let Some(items) = a.as_boxes() {
5684 let inner: Vec<Array> = items.iter().map(zeroed).collect();
5685 return Array::new(a.shape.clone(), Data::Box(inner.into()));
5686 }
5687 let dtype = if a.dtype() == DType::Char { DType::Char } else { DType::I64 };
5688 Array::new(a.shape.clone(), fill_data(dtype, a.count()))
5689 }
5690 let first = y.as_boxes()?.first()?;
5691 Some(zeroed(first))
5692}
5693
5694fn drop_(x: &Array, y: &Array, apl: bool, span: Span) -> Result<Array> {
5695 let counts = axis_counts(x, "drop", span)?;
5696 let promoted;
5697 let base = if y.rank() == 0 {
5698 promoted = Array::new(vec![1; counts.len()], y.data.clone());
5699 &promoted
5700 } else {
5701 y
5702 };
5703 let wrong =
5704 if apl { counts.len() != base.rank() } else { counts.len() > base.rank() };
5705 if wrong {
5706 return Err(count_rank("drop", counts.len(), base.rank(), span));
5707 }
5708 if let Some(run) = leading_run(base, &counts, true) {
5709 return Ok(run);
5710 }
5711 let mut out_shape = base.shape.clone();
5712 let mut offset = vec![0usize; base.rank()];
5713 for (a, &k) in counts.iter().enumerate() {
5714 let len = base.shape[a];
5715 let d = (k.unsigned_abs() as usize).min(len);
5716 out_shape[a] = len - d;
5717 if k > 0 {
5718 offset[a] = d;
5719 }
5720 }
5721 let n: usize = out_shape.iter().product();
5722 let st = strides(&base.shape);
5723 let mut data = Data::empty(base.dtype());
5724 let mut coord = vec![0usize; out_shape.len()];
5725 for _ in 0..n {
5726 let idx: usize = (0..out_shape.len()).map(|a| (coord[a] + offset[a]) * st[a]).sum();
5727 push_elem(&mut data, &base.data, idx);
5728 odometer(&mut coord, &out_shape);
5729 }
5730 Ok(Array::new(out_shape, data))
5731}
5732
5733fn dyad_op(p: &Prim, x: &Array, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
5735 let tol = cfg.tol;
5736 match p.dyad {
5737 DyadOp::Scalar(op) => scalar_dyad(op, x, y, cfg, span),
5740 DyadOp::Reshape => reshape(x, y, cfg.agreement == Agreement::LeadingPrefix, span),
5741 DyadOp::Take => {
5742 let apl = cfg.rules.lang == crate::Lang::Apl;
5743 take(x, y, cfg.agreement == Agreement::ExactOrScalar, apl, span)
5744 }
5745 DyadOp::Drop => drop_(x, y, cfg.rules.lang == crate::Lang::Apl, span),
5746 DyadOp::Right => Ok(y.clone()),
5747 DyadOp::Left => Ok(x.clone()),
5748 DyadOp::Rotate => rotate(x, y, span),
5749 DyadOp::AppendLeading => {
5752 catenate(x, y, true, cfg.agreement == Agreement::LeadingPrefix, span)
5753 }
5754 DyadOp::AppendLast => {
5755 catenate(x, y, false, cfg.agreement == Agreement::LeadingPrefix, span)
5756 }
5757 DyadOp::IndexOf { origin } => Ok(index_of(x, y, origin, tol)),
5758 DyadOp::MemberJ => Ok(member_j(x, y, tol)),
5759 DyadOp::MemberApl => Ok(member_apl(x, y, tol)),
5760 DyadOp::From => from_index(x, y, span),
5761 DyadOp::Match => {
5762 let empties_differ = cfg.rules.lang == crate::Lang::Apl
5766 && x.count() == 0
5767 && y.count() == 0
5768 && (x.dtype() == DType::Char) != (y.dtype() == DType::Char);
5769 Ok(Array::scalar_bool(!empties_differ && arrays_match(x, y, tol)))
5770 }
5771 DyadOp::NotMatch => Ok(Array::scalar_bool(!arrays_match(x, y, tol))),
5772 DyadOp::GradeSelect { down } => grade_select(x, y, down, cfg.rules, span),
5773 DyadOp::Copy => copy_items(x, y, cfg.agreement == Agreement::ExactOrScalar, span),
5774 DyadOp::CollateGrade { down, origin } => collate_grade(x, y, down, origin, span),
5775 DyadOp::TransposeJ => transpose_j(x, y, span),
5776 DyadOp::TransposeApl => transpose_apl(x, y, cfg.rules.origin, span),
5777 DyadOp::DecodeApl => decode_apl(x, y, span).map(|r| carry_exact2(r, x, y)),
5778 DyadOp::EncodeApl => encode_apl(x, y, span).map(|r| carry_exact2(r, x, y)),
5779 DyadOp::Decode => decode(Some(x), y, span).map(|r| carry_exact2(r, x, y)),
5780 DyadOp::Encode => encode(x, y, span).map(|r| carry_exact2(r, x, y)),
5781 DyadOp::Laminate => laminate(x, y, span),
5782 DyadOp::Link => link(x, y, span),
5783 DyadOp::Strand => strand(x, y, span),
5784 DyadOp::IntervalIndex { offset, closed } => {
5785 interval_index(x, y, offset, closed, tol, span)
5786 }
5787 DyadOp::IndexOfLast { origin } => Ok(index_of_last(x, y, origin, tol)),
5788 DyadOp::MatrixDivide => matrix_divide(x, y, span),
5789 DyadOp::PartitionEnclose => partition_enclose(x, y, span),
5790 DyadOp::Squad { origin } => squad(x, y, origin, span),
5791 DyadOp::SelectAxis { axis, rank, origin } => {
5792 select_axis(x, y, axis, rank, origin, span)
5793 }
5794 DyadOp::Fetch => fetch(x, y, span),
5795 DyadOp::PolyEval => poly_eval(x, y, span),
5796 DyadOp::PolyIntegral => poly_integral(x, y, span),
5797 DyadOp::TruthTable(m) => truth_table(m, x, y, span),
5798 DyadOp::FormatSpec => format_spec(x, y, &cfg.fmt, span),
5799 DyadOp::Deal { origin, fixed } => deal(x, y, origin, fixed, span),
5800 DyadOp::ExactForm => exact_form(x, y, span),
5801 DyadOp::Boolean(op) => bool_dyad(op, x, y, cfg, span),
5802 DyadOp::Less => {
5803 set_rank(cfg, "without", x, y, span)?;
5804 Ok(set_less(x, y, tol))
5805 }
5806 DyadOp::Union => {
5807 set_rank(cfg, "union", x, y, span)?;
5808 union_items(x, y, tol, span)
5809 }
5810 DyadOp::Intersect => {
5811 set_rank(cfg, "intersection", x, y, span)?;
5812 Ok(intersect_items(x, y, tol))
5813 }
5814 DyadOp::AnagramFrom => anagram_from(x, y, span),
5815 DyadOp::Permute => permute(x, y, span),
5816 DyadOp::FindSeq => {
5817 find_seq(x, y, tol, cfg.rules.lang == crate::Lang::Apl, span)
5818 }
5819 DyadOp::UnicodeForm => unicode_form(x, y, span),
5820 DyadOp::PrimeMeta => prime_meta(x, y, span).map(|r| carry_exact2(r, x, y)),
5821 DyadOp::PrimeExponents => prime_exponents(x, y, span).map(|r| carry_exact2(r, x, y)),
5822 DyadOp::Pick { origin } => pick(x, y, origin, span),
5823 DyadOp::Expand => expand(x, y, span),
5824 DyadOp::WriteStream => Err(Error::internal("1!:2 reached the pure dyad dispatcher")),
5827 DyadOp::NotYet(what) => Err(Error::not_yet(what, span)),
5828 DyadOp::None => Err(Error::domain(format!("{} has no dyadic meaning", p.name), span)),
5829 }
5830}
5831
5832fn reduce_identity(v: &Verb, n: usize) -> Option<Data> {
5844 let Verb::Prim(p) = v else { return None };
5845 let DyadOp::Scalar(op) = p.dyad else { return None };
5846 let ints = |k: i64| Data::I64(vec![k; n].into());
5847 let bits = |k: u8| Data::Bool(vec![k; n].into());
5848 Some(match op {
5849 ScalarDyad::Add | ScalarDyad::Sub | ScalarDyad::Gcd | ScalarDyad::Residue => ints(0),
5850 ScalarDyad::Mul
5851 | ScalarDyad::DivJ
5852 | ScalarDyad::DivApl
5853 | ScalarDyad::Pow
5854 | ScalarDyad::Lcm
5855 | ScalarDyad::Root
5856 | ScalarDyad::Binomial => ints(1),
5857 ScalarDyad::Min => Data::F64(vec![f64::INFINITY; n].into()),
5858 ScalarDyad::Max => Data::F64(vec![f64::NEG_INFINITY; n].into()),
5859 ScalarDyad::Eq | ScalarDyad::Le | ScalarDyad::Ge => bits(1),
5860 ScalarDyad::Ne | ScalarDyad::Lt | ScalarDyad::Gt => bits(0),
5861 ScalarDyad::MakeComplex | ScalarDyad::PolarBy => return None,
5864 ScalarDyad::Log | ScalarDyad::Circle => return None,
5867 })
5868}
5869
5870fn is_associative(op: ScalarDyad) -> bool {
5876 use ScalarDyad::*;
5877 matches!(op, Add | Mul | Min | Max)
5878}
5879
5880#[inline(always)]
5881fn fold_range_body<T, F>(
5882 v: &[T],
5883 m: usize,
5884 lo: usize,
5885 hi: usize,
5886 j0: usize,
5887 acc: &mut [T],
5888 step: &F,
5889) -> bool
5890where
5891 T: Copy,
5892 F: Fn(T, T) -> (T, bool),
5893{
5894 let w = acc.len();
5895 let base = (hi - 1) * m + j0;
5896 acc.copy_from_slice(&v[base..base + w]);
5897 let mut over = false;
5900 for i in (lo..hi - 1).rev() {
5901 let row = &v[i * m + j0..i * m + j0 + w];
5902 for (slot, &x) in acc.iter_mut().zip(row) {
5903 let (r, o) = step(x, *slot);
5904 *slot = r;
5905 over |= o;
5906 }
5907 }
5908 !over
5909}
5910
5911multiversioned! {
5912 #[allow(clippy::too_many_arguments)]
5913 fn fold_range_vectorised[T: Copy, F: Fn(T, T) -> (T, bool)](
5914 v: &[T],
5915 m: usize,
5916 lo: usize,
5917 hi: usize,
5918 j0: usize,
5919 acc: &mut [T],
5920 step: &F,
5921 ) -> bool = fold_range_body;
5922}
5923
5924const VECTOR_COLUMNS: usize = 16;
5932
5933#[allow(clippy::too_many_arguments)]
5942#[inline]
5943fn fold_range<T, F>(
5944 v: &[T],
5945 m: usize,
5946 lo: usize,
5947 hi: usize,
5948 j0: usize,
5949 acc: &mut [T],
5950 step: &F,
5951) -> bool
5952where
5953 T: Copy,
5954 F: Fn(T, T) -> (T, bool),
5955{
5956 if acc.len() < VECTOR_COLUMNS {
5957 fold_range_body(v, m, lo, hi, j0, acc, step)
5958 } else {
5959 fold_range_vectorised(v, m, lo, hi, j0, acc, step)
5960 }
5961}
5962
5963const FOLD_LANES: usize = 8;
5974
5975const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
5980
5981#[inline(always)]
5986fn fold_lanes_body<T, F>(v: &[T], step: &F) -> Option<T>
5987where
5988 T: Copy,
5989 F: Fn(T, T) -> (T, bool),
5990{
5991 let n = v.len();
5992 let mut over = false;
5993 if n < MIN_LANE_WORK {
5994 let mut acc = v[n - 1];
5995 for &x in v[..n - 1].iter().rev() {
5996 let (r, o) = step(x, acc);
5997 acc = r;
5998 over |= o;
5999 }
6000 return (!over).then_some(acc);
6001 }
6002 let rows = n / FOLD_LANES;
6005 let head = n - rows * FOLD_LANES;
6006 let last = head + (rows - 1) * FOLD_LANES;
6007 let mut acc = [v[last]; FOLD_LANES];
6008 acc.copy_from_slice(&v[last..last + FOLD_LANES]);
6009 for r in (0..rows - 1).rev() {
6010 let row = &v[head + r * FOLD_LANES..head + (r + 1) * FOLD_LANES];
6011 for (slot, &x) in acc.iter_mut().zip(row) {
6012 let (r, o) = step(x, *slot);
6013 *slot = r;
6014 over |= o;
6015 }
6016 }
6017 let mut a = acc[FOLD_LANES - 1];
6018 for &x in acc[..FOLD_LANES - 1].iter().rev() {
6019 let (r, o) = step(x, a);
6020 a = r;
6021 over |= o;
6022 }
6023 for &x in v[..head].iter().rev() {
6024 let (r, o) = step(x, a);
6025 a = r;
6026 over |= o;
6027 }
6028 (!over).then_some(a)
6029}
6030
6031multiversioned! {
6032 fn fold_lanes_vectorised[T: Copy, F: Fn(T, T) -> (T, bool)](
6033 v: &[T],
6034 step: &F,
6035 ) -> Option<T> = fold_lanes_body;
6036}
6037
6038#[inline]
6041fn fold_lanes<T, F>(v: &[T], step: &F) -> Option<T>
6042where
6043 T: Copy,
6044 F: Fn(T, T) -> (T, bool),
6045{
6046 if v.len() < MIN_LANE_WORK {
6047 fold_lanes_body(v, step)
6048 } else {
6049 fold_lanes_vectorised(v, step)
6050 }
6051}
6052
6053fn fold_flat<T, F>(v: &[T], n: usize, assoc: bool, step: &F) -> Option<T>
6056where
6057 T: Copy + Send + Sync,
6058 F: Fn(T, T) -> (T, bool) + Sync + Send,
6059{
6060 if assoc {
6061 return par::try_fold_chunks(
6062 &v[..n],
6063 |part| fold_lanes(part, step),
6064 |a, b| {
6065 let (r, o) = step(a, b);
6066 (!o).then_some(r)
6067 },
6068 );
6069 }
6070 let mut acc = v[n - 1];
6071 let mut over = false;
6072 for &x in v[..n - 1].iter().rev() {
6073 let (r, o) = step(x, acc);
6074 acc = r;
6075 over |= o;
6076 }
6077 (!over).then_some(acc)
6078}
6079
6080fn fold_items<T, F>(v: &[T], n: usize, m: usize, assoc: bool, step: F) -> Option<Vec<T>>
6091where
6092 T: Copy + Default + Send + Sync,
6093 F: Fn(T, T) -> (T, bool) + Sync + Send,
6094{
6095 if m >= par::WIDE_ITEM {
6096 let (out, ok) = par::fill_wide(m, n * m, |j0, acc: &mut [T]| {
6097 fold_range(v, m, 0, n, j0, acc, &step)
6098 });
6099 return ok.then_some(out);
6100 }
6101 if m == 1 {
6102 return fold_flat(v, n, assoc, &step).map(|x| vec![x]);
6103 }
6104 let chunks = if assoc { par::chunks(n, n * m) } else { 1 };
6105 if chunks < 2 {
6106 let mut acc = vec![T::default(); m];
6107 return fold_range(v, m, 0, n, 0, &mut acc, &step).then_some(acc);
6108 }
6109 let per = n.div_ceil(chunks);
6110 let parts = par::map_indexed(n.div_ceil(per), |c| {
6111 let mut acc = vec![T::default(); m];
6112 let ok = fold_range(v, m, c * per, ((c + 1) * per).min(n), 0, &mut acc, &step);
6113 ok.then_some(acc)
6114 });
6115 let mut it = parts.into_iter().rev();
6118 let mut acc = it.next()??;
6119 for part in it {
6120 let part = part?;
6121 let mut over = false;
6122 for (slot, &x) in acc.iter_mut().zip(&part) {
6123 let (r, o) = step(x, *slot);
6124 *slot = r;
6125 over |= o;
6126 }
6127 if over {
6128 return None;
6129 }
6130 }
6131 Some(acc)
6132}
6133
6134fn fold_i64(op: ScalarDyad, v: &[i64], n: usize, m: usize) -> Option<Vec<i64>> {
6135 use ScalarDyad::*;
6136 let assoc = is_associative(op);
6137 match op {
6138 Add => fold_items(v, n, m, assoc, i64::overflowing_add),
6139 Sub => fold_items(v, n, m, assoc, i64::overflowing_sub),
6140 Mul => fold_items(v, n, m, assoc, i64::overflowing_mul),
6141 Min => fold_items(v, n, m, assoc, |a: i64, b: i64| (a.min(b), false)),
6142 Max => fold_items(v, n, m, assoc, |a: i64, b: i64| (a.max(b), false)),
6143 _ => None,
6144 }
6145}
6146
6147fn fold_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize) -> Option<Vec<Cx>> {
6148 use ScalarDyad::*;
6149 let assoc = is_associative(op);
6150 match op {
6151 Add => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::add(a, b), false)),
6152 Sub => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::sub(a, b), false)),
6153 Mul => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::mul(a, b), false)),
6154 _ => None,
6156 }
6157}
6158
6159fn fold_f64(op: ScalarDyad, v: &[f64], n: usize, m: usize) -> Option<Vec<f64>> {
6160 use ScalarDyad::*;
6161 let assoc = is_associative(op);
6162 match op {
6163 Add => fold_items(v, n, m, assoc, |a: f64, b: f64| (a + b, false)),
6164 Sub => fold_items(v, n, m, assoc, |a: f64, b: f64| (a - b, false)),
6165 Mul => fold_items(v, n, m, assoc, |a: f64, b: f64| (a * b, false)),
6166 Min => fold_items(v, n, m, assoc, |a: f64, b: f64| (a.min(b), false)),
6167 Max => fold_items(v, n, m, assoc, |a: f64, b: f64| (a.max(b), false)),
6168 _ => None,
6169 }
6170}
6171
6172fn reduce_typed(op: ScalarDyad, d: &Data, n: usize, m: usize) -> Option<Data> {
6176 use ScalarDyad::*;
6177 if !matches!(op, Add | Sub | Mul | Min | Max) {
6180 return None;
6181 }
6182 match d {
6183 Data::F64(v) => Some(Data::F64(fold_f64(op, v, n, m)?.into())),
6184 Data::Complex(v) => Some(Data::Complex(fold_cx(op, v, n, m)?.into())),
6185 Data::I64(v) => Some(Data::I64(fold_i64(op, v, n, m)?.into())),
6186 Data::Bool(v) => {
6189 let widened = par::map(v, |&b| b as i64);
6190 Some(Data::I64(fold_i64(op, &widened, n, m)?.into()))
6191 }
6192 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
6195 }
6196}
6197
6198#[inline(always)]
6204fn fold_runs_body<T, F>(v: &[T], start: usize, m: usize, out: &mut [T], step: &F) -> bool
6205where
6206 T: Copy,
6207 F: Fn(T, T) -> (T, bool),
6208{
6209 let mut over = false;
6210 for (k, slot) in out.iter_mut().enumerate() {
6211 let run = &v[(start + k) * m..(start + k + 1) * m];
6212 let mut acc = run[m - 1];
6213 for &x in run[..m - 1].iter().rev() {
6214 let (r, o) = step(x, acc);
6215 acc = r;
6216 over |= o;
6217 }
6218 *slot = acc;
6219 }
6220 !over
6221}
6222
6223multiversioned! {
6224 fn fold_runs_vectorised[T: Copy, F: Fn(T, T) -> (T, bool)](
6225 v: &[T],
6226 start: usize,
6227 m: usize,
6228 out: &mut [T],
6229 step: &F,
6230 ) -> bool = fold_runs_body;
6231}
6232
6233fn fold_runs<T, F>(v: &[T], n: usize, m: usize, step: F) -> Option<Vec<T>>
6236where
6237 T: Copy + Default + Send + Sync,
6238 F: Fn(T, T) -> (T, bool) + Sync + Send,
6239{
6240 let wide = m >= VECTOR_COLUMNS;
6244 let (out, ok) = par::fill_wide(n, n * m, |start, part: &mut [T]| {
6245 if wide {
6246 fold_runs_vectorised(v, start, m, part, &step)
6247 } else {
6248 fold_runs_body(v, start, m, part, &step)
6249 }
6250 });
6251 ok.then_some(out)
6252}
6253
6254fn fold_runs_data(op: ScalarDyad, d: &Data, n: usize, m: usize) -> Option<Data> {
6255 use ScalarDyad::*;
6256 match d {
6257 Data::F64(v) => Some(Data::F64(
6258 match op {
6259 Add => fold_runs(v, n, m, |a: f64, b: f64| (a + b, false)),
6260 Sub => fold_runs(v, n, m, |a: f64, b: f64| (a - b, false)),
6261 Mul => fold_runs(v, n, m, |a: f64, b: f64| (a * b, false)),
6262 Min => fold_runs(v, n, m, |a: f64, b: f64| (a.min(b), false)),
6263 Max => fold_runs(v, n, m, |a: f64, b: f64| (a.max(b), false)),
6264 _ => None,
6265 }?
6266 .into(),
6267 )),
6268 Data::I64(v) => Some(Data::I64(
6269 match op {
6270 Add => fold_runs(v, n, m, i64::overflowing_add),
6271 Sub => fold_runs(v, n, m, i64::overflowing_sub),
6272 Mul => fold_runs(v, n, m, i64::overflowing_mul),
6273 Min => fold_runs(v, n, m, |a: i64, b: i64| (a.min(b), false)),
6274 Max => fold_runs(v, n, m, |a: i64, b: i64| (a.max(b), false)),
6275 _ => None,
6276 }?
6277 .into(),
6278 )),
6279 Data::Complex(v) => Some(Data::Complex(
6281 match op {
6282 Add => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::add(a, b), false)),
6283 Sub => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::sub(a, b), false)),
6284 Mul => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::mul(a, b), false)),
6285 _ => None,
6286 }?
6287 .into(),
6288 )),
6289 Data::Bool(v) => {
6292 let widened = par::map(v, |&b| b as i64);
6293 fold_runs_data(op, &Data::I64(widened.into()), n, m)
6294 }
6295 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
6296 }
6297}
6298
6299fn run_slices<T: Clone>(b: &Buf<T>, runs: usize, len: usize) -> Vec<&[T]> {
6315 if let Some(parts) = b.parts() && parts.len() == runs && parts.iter().all(|p| p.len() == len) {
6316 return parts.iter().map(Buf::as_slice).collect();
6317 }
6318 let flat = b.as_slice();
6319 (0..runs).map(|c| &flat[c * len..(c + 1) * len]).collect()
6320}
6321
6322fn fold_columns<T, F>(cols: &[&[T]], len: usize, assoc: bool, step: F) -> Option<Vec<T>>
6331where
6332 T: Copy + Default + Send + Sync,
6333 F: Fn(T, T) -> (T, bool) + Sync + Send,
6334{
6335 if par::worth_it(len) {
6340 let mut out = Vec::with_capacity(cols.len());
6341 for c in cols {
6342 out.push(fold_flat(c, len, assoc, &step)?);
6343 }
6344 return Some(out);
6345 }
6346 let (out, ok) = par::fill_wide(cols.len(), cols.len() * len, |start, part: &mut [T]| {
6347 let mut ok = true;
6348 for (k, slot) in part.iter_mut().enumerate() {
6349 match fold_flat(cols[start + k], len, assoc, &step) {
6350 Some(v) => *slot = v,
6351 None => ok = false,
6352 }
6353 }
6354 ok
6355 });
6356 ok.then_some(out)
6357}
6358
6359fn fold_columns_data(op: ScalarDyad, d: &Data, runs: usize, len: usize) -> Option<Data> {
6361 use ScalarDyad::*;
6362 if !matches!(op, Add | Sub | Mul | Min | Max) {
6363 return None;
6364 }
6365 let assoc = is_associative(op);
6366 macro_rules! by {
6367 ($v:expr, $add:expr, $sub:expr, $mul:expr, $min:expr, $max:expr) => {{
6368 let cols = run_slices($v, runs, len);
6369 match op {
6370 Add => fold_columns(&cols, len, assoc, $add),
6371 Sub => fold_columns(&cols, len, assoc, $sub),
6372 Mul => fold_columns(&cols, len, assoc, $mul),
6373 Min => fold_columns(&cols, len, assoc, $min),
6374 Max => fold_columns(&cols, len, assoc, $max),
6375 _ => None,
6376 }?
6377 }};
6378 }
6379 match d {
6380 Data::F64(v) => Some(Data::F64(
6381 by!(
6382 v,
6383 |a: f64, b: f64| (a + b, false),
6384 |a: f64, b: f64| (a - b, false),
6385 |a: f64, b: f64| (a * b, false),
6386 |a: f64, b: f64| (a.min(b), false),
6387 |a: f64, b: f64| (a.max(b), false)
6388 )
6389 .into(),
6390 )),
6391 Data::I64(v) => Some(Data::I64(
6392 by!(
6393 v,
6394 i64::overflowing_add,
6395 i64::overflowing_sub,
6396 i64::overflowing_mul,
6397 |a: i64, b: i64| (a.min(b), false),
6398 |a: i64, b: i64| (a.max(b), false)
6399 )
6400 .into(),
6401 )),
6402 Data::Complex(v) => {
6403 if !matches!(op, Add | Sub | Mul) {
6404 return None;
6405 }
6406 Some(Data::Complex(
6407 by!(
6408 v,
6409 |a: Cx, b: Cx| (cx::add(a, b), false),
6410 |a: Cx, b: Cx| (cx::sub(a, b), false),
6411 |a: Cx, b: Cx| (cx::mul(a, b), false),
6412 |_: Cx, _: Cx| unreachable!("refused above"),
6413 |_: Cx, _: Cx| unreachable!("refused above")
6414 )
6415 .into(),
6416 ))
6417 }
6418 Data::Bool(v) => {
6423 let widened: Vec<Buf<i64>> = run_slices(v, runs, len)
6424 .iter()
6425 .map(|c| Buf::from_vec(par::map(c, |&b| b as i64)))
6426 .collect();
6427 fold_columns_data(op, &Data::I64(Buf::join(widened)), runs, len)
6428 }
6429 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
6430 }
6431}
6432
6433fn reduce_columns(v: &Verb, y: &Array) -> Option<Array> {
6438 let Verb::Prim(p) = v else { return None };
6439 let DyadOp::Scalar(op) = p.dyad else { return None };
6440 if !y.dtype().is_numeric() {
6441 return None;
6442 }
6443 let n = y.shape[0];
6444 let m: usize = y.shape[1..].iter().product();
6445 if n == 0 || m == 0 {
6448 return None;
6449 }
6450 let shape = y.shape[1..].to_vec();
6451 if n == 1 {
6454 return Some(Array::col_major(shape, y.data.clone()));
6455 }
6456 let data = fold_columns_data(op, &y.data, m, n)?;
6457 Some(Array::col_major(shape, data))
6458}
6459
6460fn fold_across<T, F>(cols: &[&[T]], rows: usize, step: F) -> Option<Vec<T>>
6463where
6464 T: Copy + Default + Send + Sync,
6465 F: Fn(T, T) -> (T, bool) + Sync + Send,
6466{
6467 let (last, rest) = cols.split_last()?;
6468 let (out, ok) = par::fill(rows, |start, part: &mut [T]| {
6469 let mut over = false;
6470 for (k, slot) in part.iter_mut().enumerate() {
6471 let i = start + k;
6472 let mut acc = last[i];
6473 for c in rest.iter().rev() {
6474 let (r, o) = step(c[i], acc);
6475 acc = r;
6476 over |= o;
6477 }
6478 *slot = acc;
6479 }
6480 !over
6481 });
6482 ok.then_some(out)
6483}
6484
6485fn fold_across_data(op: ScalarDyad, d: &Data, rows: usize, cols: usize) -> Option<Data> {
6486 use ScalarDyad::*;
6487 if !matches!(op, Add | Sub | Mul | Min | Max) {
6488 return None;
6489 }
6490 macro_rules! by {
6491 ($v:expr, $add:expr, $sub:expr, $mul:expr, $min:expr, $max:expr) => {{
6492 let parts = run_slices($v, cols, rows);
6493 match op {
6494 Add => fold_across(&parts, rows, $add),
6495 Sub => fold_across(&parts, rows, $sub),
6496 Mul => fold_across(&parts, rows, $mul),
6497 Min => fold_across(&parts, rows, $min),
6498 Max => fold_across(&parts, rows, $max),
6499 _ => None,
6500 }?
6501 }};
6502 }
6503 match d {
6504 Data::F64(v) => Some(Data::F64(
6505 by!(
6506 v,
6507 |a: f64, b: f64| (a + b, false),
6508 |a: f64, b: f64| (a - b, false),
6509 |a: f64, b: f64| (a * b, false),
6510 |a: f64, b: f64| (a.min(b), false),
6511 |a: f64, b: f64| (a.max(b), false)
6512 )
6513 .into(),
6514 )),
6515 Data::I64(v) => Some(Data::I64(
6516 by!(
6517 v,
6518 i64::overflowing_add,
6519 i64::overflowing_sub,
6520 i64::overflowing_mul,
6521 |a: i64, b: i64| (a.min(b), false),
6522 |a: i64, b: i64| (a.max(b), false)
6523 )
6524 .into(),
6525 )),
6526 Data::Complex(v) => {
6527 if !matches!(op, Add | Sub | Mul) {
6528 return None;
6529 }
6530 Some(Data::Complex(
6531 by!(
6532 v,
6533 |a: Cx, b: Cx| (cx::add(a, b), false),
6534 |a: Cx, b: Cx| (cx::sub(a, b), false),
6535 |a: Cx, b: Cx| (cx::mul(a, b), false),
6536 |_: Cx, _: Cx| unreachable!("refused above"),
6537 |_: Cx, _: Cx| unreachable!("refused above")
6538 )
6539 .into(),
6540 ))
6541 }
6542 Data::Bool(v) => {
6543 let widened: Vec<Buf<i64>> = run_slices(v, cols, rows)
6544 .iter()
6545 .map(|c| Buf::from_vec(par::map(c, |&b| b as i64)))
6546 .collect();
6547 fold_across_data(op, &Data::I64(Buf::join(widened)), rows, cols)
6548 }
6549 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
6550 }
6551}
6552
6553fn reduce_rows_columns(u: &Verb, y: &Array) -> Option<Array> {
6556 let Verb::Reduce(inner) = u else { return None };
6557 let Verb::Prim(p) = &**inner else { return None };
6558 let DyadOp::Scalar(op) = p.dyad else { return None };
6559 if y.rank() != 2 || !y.dtype().is_numeric() {
6562 return None;
6563 }
6564 let (rows, cols) = (y.shape[0], y.shape[1]);
6565 if rows == 0 || cols == 0 {
6568 return None;
6569 }
6570 if cols == 1 {
6571 return Some(Array::new(vec![rows], y.data.clone()));
6573 }
6574 let data = fold_across_data(op, &y.data, rows, cols)?;
6575 Some(Array::new(vec![rows], data))
6576}
6577
6578fn reduce_vector_cells(u: &Verb, y: &Array, frame_rank: usize) -> Option<Array> {
6587 let Verb::Reduce(inner) = u else { return None };
6588 let Verb::Prim(p) = &**inner else { return None };
6589 let DyadOp::Scalar(op) = p.dyad else { return None };
6590 if y.rank() != frame_rank + 1 || !y.dtype().is_numeric() {
6593 return None;
6594 }
6595 let m = y.shape[frame_rank];
6596 if m == 0 {
6599 return None;
6600 }
6601 use ScalarDyad::{Add, Max, Min, Mul, Sub};
6602 if !matches!(op, Add | Sub | Mul | Min | Max) {
6603 return None;
6604 }
6605 let frame = y.shape[..frame_rank].to_vec();
6606 if m == 1 {
6607 return Some(Array::new(frame, y.data.clone()));
6610 }
6611 let n: usize = frame.iter().product();
6612 let data = fold_runs_data(op, &y.data, n, m)?;
6613 Some(Array::new(frame, data))
6614}
6615
6616fn reduce(v: &Verb, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
6618 if y.rank() == 0 {
6619 return Ok(y.clone());
6620 }
6621 let n = y.items();
6622 if n == 1 {
6623 return Ok(y.item(0));
6624 }
6625 let cell_shape = y.shape[1..].to_vec();
6626 let m: usize = cell_shape.iter().product();
6627 if n == 0 {
6628 if matches!(v, Verb::Prim(p) if matches!(p.dyad, DyadOp::AppendLeading | DyadOp::AppendLast))
6631 {
6632 return Ok(Array::new(vec![0], Data::empty(y.dtype())));
6633 }
6634 return match reduce_identity(v, m) {
6635 Some(d) => Ok(Array::new(cell_shape, d)),
6636 None => Err(Error::domain(
6637 format!("empty reduction has no identity for {}", v.name()),
6638 span,
6639 )),
6640 };
6641 }
6642 if y.dtype().is_numeric() && let Verb::Prim(p) = v && let DyadOp::Scalar(op) = p.dyad {
6643 if let Some(d) = reduce_typed(op, y.row_major_data(), n, m) {
6648 return Ok(Array::new(cell_shape, d));
6649 }
6650 let mut acc = y.data.slice((n - 1) * m, n * m);
6653 for i in (0..n - 1).rev() {
6654 acc =
6655 scalar_dyad_data(op, &y.data, i * m, 1, &acc, 0, 1, m, ctx.cfg.tol, span)?;
6656 }
6657 return Ok(Array::new(cell_shape, acc));
6658 }
6659 let mut acc = y.item(n - 1);
6660 for i in (0..n - 1).rev() {
6661 acc = v.dyad(&y.item(i), &acc, ctx, span)?;
6662 }
6663 Ok(acc)
6664}
6665
6666fn folded_op(u: &Verb) -> Option<ScalarDyad> {
6673 let Verb::Reduce(inner) = u else { return None };
6674 let Verb::Prim(p) = &**inner else { return None };
6675 match p.dyad {
6676 DyadOp::Scalar(op) => Some(op),
6677 _ => None,
6678 }
6679}
6680
6681fn section(y: &Array, lo: usize, hi: usize) -> Array {
6683 let m = y.item_size();
6684 let mut shape = y.shape.clone();
6685 shape[0] = hi - lo;
6686 Array::new(shape, y.data.slice(lo * m, hi * m))
6687}
6688
6689fn as_items(y: &Array) -> Option<Array> {
6692 (y.rank() == 0).then(|| Array::new(vec![1], y.data.clone()))
6693}
6694
6695#[inline(always)]
6696fn scan_flat_body<T, F>(v: &[T], n: usize, m: usize, back: bool, step: F) -> Option<Vec<T>>
6697where
6698 T: Copy + Default,
6699 F: Fn(T, T) -> (T, bool),
6700{
6701 if m == 1 {
6702 let mut out = vec![T::default(); n];
6705 let mut over = false;
6706 if back {
6707 let mut acc = v[n - 1];
6708 out[n - 1] = acc;
6709 for (slot, &x) in out[..n - 1].iter_mut().zip(&v[..n - 1]).rev() {
6710 let (r, o) = step(x, acc);
6711 acc = r;
6712 over |= o;
6713 *slot = acc;
6714 }
6715 } else {
6716 let mut acc = v[0];
6717 out[0] = acc;
6718 for (slot, &x) in out[1..n].iter_mut().zip(&v[1..n]) {
6719 let (r, o) = step(acc, x);
6720 acc = r;
6721 over |= o;
6722 *slot = acc;
6723 }
6724 }
6725 return (!over).then_some(out);
6726 }
6727 let mut out = vec![T::default(); n * m];
6728 let mut acc = vec![T::default(); m];
6729 let mut over = false;
6730 if back {
6731 acc.copy_from_slice(&v[(n - 1) * m..n * m]);
6732 out[(n - 1) * m..n * m].copy_from_slice(&acc);
6733 for i in (0..n - 1).rev() {
6734 for (j, slot) in acc.iter_mut().enumerate() {
6735 let (r, o) = step(v[i * m + j], *slot);
6736 *slot = r;
6737 over |= o;
6738 }
6739 out[i * m..i * m + m].copy_from_slice(&acc);
6740 }
6741 } else {
6742 acc.copy_from_slice(&v[..m]);
6743 out[..m].copy_from_slice(&acc);
6744 for i in 1..n {
6745 for (j, slot) in acc.iter_mut().enumerate() {
6746 let (r, o) = step(*slot, v[i * m + j]);
6747 *slot = r;
6748 over |= o;
6749 }
6750 out[i * m..i * m + m].copy_from_slice(&acc);
6751 }
6752 }
6753 (!over).then_some(out)
6754}
6755
6756multiversioned! {
6757 fn scan_flat_vectorised[T: Copy + Default, F: Fn(T, T) -> (T, bool)](
6758 v: &[T],
6759 n: usize,
6760 m: usize,
6761 back: bool,
6762 step: F,
6763 ) -> Option<Vec<T>> = scan_flat_body;
6764}
6765
6766fn scan_flat<T, F>(v: &[T], n: usize, m: usize, back: bool, step: F) -> Option<Vec<T>>
6778where
6779 T: Copy + Default,
6780 F: Fn(T, T) -> (T, bool),
6781{
6782 if m < VECTOR_COLUMNS {
6783 scan_flat_body(v, n, m, back, step)
6784 } else {
6785 scan_flat_vectorised(v, n, m, back, step)
6786 }
6787}
6788
6789fn scan_i64(op: ScalarDyad, v: &[i64], n: usize, m: usize, back: bool) -> Option<Vec<i64>> {
6790 use ScalarDyad::*;
6791 match op {
6792 Add => scan_flat(v, n, m, back, i64::overflowing_add),
6793 Sub => scan_flat(v, n, m, back, i64::overflowing_sub),
6794 Mul => scan_flat(v, n, m, back, i64::overflowing_mul),
6795 Min => scan_flat(v, n, m, back, |a: i64, b: i64| (a.min(b), false)),
6796 Max => scan_flat(v, n, m, back, |a: i64, b: i64| (a.max(b), false)),
6797 _ => None,
6798 }
6799}
6800
6801fn scan_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize, back: bool) -> Option<Vec<Cx>> {
6802 use ScalarDyad::*;
6803 match op {
6804 Add => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::add(a, b), false)),
6805 Sub => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::sub(a, b), false)),
6806 Mul => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::mul(a, b), false)),
6807 _ => None,
6808 }
6809}
6810
6811fn scan_f64(op: ScalarDyad, v: &[f64], n: usize, m: usize, back: bool) -> Option<Vec<f64>> {
6812 use ScalarDyad::*;
6813 match op {
6814 Add => scan_flat(v, n, m, back, |a: f64, b: f64| (a + b, false)),
6815 Sub => scan_flat(v, n, m, back, |a: f64, b: f64| (a - b, false)),
6816 Mul => scan_flat(v, n, m, back, |a: f64, b: f64| (a * b, false)),
6817 Min => scan_flat(v, n, m, back, |a: f64, b: f64| (a.min(b), false)),
6818 Max => scan_flat(v, n, m, back, |a: f64, b: f64| (a.max(b), false)),
6819 _ => None,
6820 }
6821}
6822
6823fn scan_typed(op: ScalarDyad, d: &Data, n: usize, m: usize, back: bool) -> Option<Data> {
6827 use ScalarDyad::*;
6828 if !matches!(op, Add | Sub | Mul | Min | Max) {
6829 return None;
6830 }
6831 let widened = |v: &[i64]| {
6832 let f: Vec<f64> = v.iter().map(|&x| x as f64).collect();
6833 Data::F64(scan_f64(op, &f, n, m, back).expect("the float scan cannot overflow").into())
6834 };
6835 let ints = |v: &[i64]| match scan_i64(op, v, n, m, back) {
6836 Some(out) => Data::I64(out.into()),
6837 None => widened(v),
6838 };
6839 match d {
6840 Data::F64(v) => Some(Data::F64(scan_f64(op, v, n, m, back)?.into())),
6841 Data::Complex(v) => Some(Data::Complex(scan_cx(op, v, n, m, back)?.into())),
6842 Data::I64(v) => Some(ints(v)),
6843 Data::Bool(v) => Some(ints(&v.iter().map(|&b| b as i64).collect::<Vec<_>>())),
6844 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
6847 }
6848}
6849
6850fn window_fold<T, F>(v: &[T], n: usize, m: usize, w: usize, step: F) -> Option<Vec<T>>
6865where
6866 T: Copy + Default + Send + Sync,
6867 F: Fn(T, T) -> (T, bool) + Sync + Send,
6868{
6869 debug_assert!(w >= 1 && n >= w);
6870 if m == 1 {
6871 return window_fold_flat(v, n, w, step);
6872 }
6873 let count = n - w + 1;
6874 let mut out = vec![T::default(); count * m];
6875 let mut pre = vec![T::default(); w * m];
6878 let mut suf = vec![T::default(); w * m];
6879 let mut prev = vec![T::default(); w * m];
6880 let mut over = false;
6881 for b in 0..n.div_ceil(w) {
6882 let bs = b * w;
6883 let be = ((b + 1) * w).min(n);
6884 pre[..m].copy_from_slice(&v[bs * m..bs * m + m]);
6885 for i in 1..be - bs {
6886 let (o, p) = (i * m, (i - 1) * m);
6887 for j in 0..m {
6888 let (r, f) = step(pre[p + j], v[(bs + i) * m + j]);
6889 pre[o + j] = r;
6890 over |= f;
6891 }
6892 }
6893 for e in bs.max(w - 1)..be {
6896 let i = e + 1 - w;
6897 let (oo, po) = (i * m, (e - bs) * m);
6898 if i == bs {
6899 out[oo..oo + m].copy_from_slice(&pre[po..po + m]);
6900 } else {
6901 let so = (i + w - bs) * m;
6902 for j in 0..m {
6903 let (r, f) = step(prev[so + j], pre[po + j]);
6904 out[oo + j] = r;
6905 over |= f;
6906 }
6907 }
6908 }
6909 let last = be - 1 - bs;
6910 suf[last * m..last * m + m].copy_from_slice(&v[(be - 1) * m..be * m]);
6911 for i in (0..last).rev() {
6912 let (o, p) = (i * m, (i + 1) * m);
6913 for j in 0..m {
6914 let (r, f) = step(v[(bs + i) * m + j], suf[p + j]);
6915 suf[o + j] = r;
6916 over |= f;
6917 }
6918 }
6919 std::mem::swap(&mut prev, &mut suf);
6920 }
6921 (!over).then_some(out)
6922}
6923
6924fn window_fold_flat<T, F>(v: &[T], n: usize, w: usize, step: F) -> Option<Vec<T>>
6934where
6935 T: Copy + Default + Send + Sync,
6936 F: Fn(T, T) -> (T, bool) + Sync + Send,
6937{
6938 let (out, ok) = par::fill(n - w + 1, |lo, part: &mut [T]| {
6939 window_fold_range(v, n, w, lo, part, &step)
6940 });
6941 ok.then_some(out)
6942}
6943
6944#[inline(always)]
6945fn window_fold_range_body<T, F>(
6946 v: &[T],
6947 n: usize,
6948 w: usize,
6949 lo: usize,
6950 out: &mut [T],
6951 step: &F,
6952) -> bool
6953where
6954 T: Copy + Default,
6955 F: Fn(T, T) -> (T, bool),
6956{
6957 if out.is_empty() {
6958 return true;
6959 }
6960 let hi = lo + out.len();
6961 let mut pre = vec![T::default(); w];
6962 let mut suf = vec![T::default(); w];
6963 let mut prev = vec![T::default(); w];
6964 let mut over = false;
6965 let mut bs = lo / w * w;
6966 while bs < n && bs <= hi + w - 2 {
6968 let block = &v[bs..(bs + w).min(n)];
6969 let lb = block.len();
6970 let mut acc = block[0];
6971 pre[0] = acc;
6972 for (slot, &x) in pre[1..lb].iter_mut().zip(&block[1..]) {
6973 let (r, o) = step(acc, x);
6974 acc = r;
6975 over |= o;
6976 *slot = acc;
6977 }
6978 for e in bs.max(lo + w - 1)..(bs + lb).min(hi + w - 1) {
6983 let i = e + 1 - w;
6984 out[i - lo] = if i == bs {
6985 pre[e - bs]
6986 } else {
6987 let (r, o) = step(prev[i + w - bs], pre[e - bs]);
6988 over |= o;
6989 r
6990 };
6991 }
6992 let mut acc = block[lb - 1];
6993 suf[lb - 1] = acc;
6994 for (slot, &x) in suf[..lb - 1].iter_mut().zip(&block[..lb - 1]).rev() {
6995 let (r, o) = step(x, acc);
6996 acc = r;
6997 over |= o;
6998 *slot = acc;
6999 }
7000 std::mem::swap(&mut prev, &mut suf);
7001 bs += w;
7002 }
7003 !over
7004}
7005
7006multiversioned! {
7007 fn window_fold_range[T: Copy + Default, F: Fn(T, T) -> (T, bool)](
7012 v: &[T],
7013 n: usize,
7014 w: usize,
7015 lo: usize,
7016 out: &mut [T],
7017 step: &F,
7018 ) -> bool = window_fold_range_body;
7019}
7020
7021pub(crate) fn windows_into<T, F>(v: &[T], w: usize, lo: usize, out: &mut [T], step: &F) -> bool
7030where
7031 T: Copy + Default,
7032 F: Fn(T, T) -> (T, bool),
7033{
7034 window_fold_range(v, v.len(), w, lo, out, step)
7035}
7036
7037fn window_i64(op: ScalarDyad, v: &[i64], n: usize, m: usize, w: usize) -> Option<Vec<i64>> {
7038 use ScalarDyad::*;
7039 match op {
7040 Add => window_fold(v, n, m, w, i64::overflowing_add),
7041 Mul => window_fold(v, n, m, w, i64::overflowing_mul),
7042 Min => window_fold(v, n, m, w, |a: i64, b: i64| (a.min(b), false)),
7043 Max => window_fold(v, n, m, w, |a: i64, b: i64| (a.max(b), false)),
7044 _ => None,
7045 }
7046}
7047
7048fn window_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize, w: usize) -> Option<Vec<Cx>> {
7049 use ScalarDyad::*;
7050 match op {
7051 Add => window_fold(v, n, m, w, |a: Cx, b: Cx| (cx::add(a, b), false)),
7052 Mul => window_fold(v, n, m, w, |a: Cx, b: Cx| (cx::mul(a, b), false)),
7053 _ => None,
7054 }
7055}
7056
7057fn window_f64(op: ScalarDyad, v: &[f64], n: usize, m: usize, w: usize) -> Option<Vec<f64>> {
7058 use ScalarDyad::*;
7059 match op {
7060 Add => window_fold(v, n, m, w, |a: f64, b: f64| (a + b, false)),
7061 Mul => window_fold(v, n, m, w, |a: f64, b: f64| (a * b, false)),
7062 Min => window_fold(v, n, m, w, |a: f64, b: f64| (a.min(b), false)),
7063 Max => window_fold(v, n, m, w, |a: f64, b: f64| (a.max(b), false)),
7064 _ => None,
7065 }
7066}
7067
7068fn window_typed(op: ScalarDyad, d: &Data, n: usize, m: usize, w: usize) -> Option<Data> {
7072 use ScalarDyad::*;
7073 if !matches!(op, Add | Mul | Min | Max) {
7074 return None;
7075 }
7076 let widened = |v: &[i64]| {
7077 let f: Vec<f64> = v.iter().map(|&x| x as f64).collect();
7078 Data::F64(window_f64(op, &f, n, m, w).expect("the float fold cannot overflow").into())
7079 };
7080 let ints = |v: &[i64]| match window_i64(op, v, n, m, w) {
7081 Some(out) => Data::I64(out.into()),
7082 None => widened(v),
7083 };
7084 match d {
7085 Data::F64(v) => Some(Data::F64(window_f64(op, v, n, m, w)?.into())),
7086 Data::Complex(v) => Some(Data::Complex(window_cx(op, v, n, m, w)?.into())),
7087 Data::I64(v) => Some(ints(v)),
7088 Data::Bool(v) => Some(ints(&v.iter().map(|&b| b as i64).collect::<Vec<_>>())),
7089 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
7092 }
7093}
7094
7095fn runs(u: &Verb, y: &Array, back: bool, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
7097 let promoted = as_items(y);
7098 let base = promoted.as_ref().unwrap_or(y);
7099 let n = base.items();
7100 let m = base.item_size();
7101 if n > 0 && base.dtype().is_numeric() && let Some(op) = folded_op(u) {
7102 if (back || is_associative(op))
7105 && let Some(d) = scan_typed(op, base.row_major_data(), n, m, back)
7106 {
7107 return Ok(Array::new(base.shape.clone(), d));
7108 }
7109 }
7110 let cells = each_cell(n, n * m, u.is_pure(), ctx, |i, c| {
7111 let part = if back { section(base, i, n) } else { section(base, 0, i + 1) };
7112 u.monad(&part, c, span)
7113 })?;
7114 assemble(&[n], cells, span)
7115}
7116
7117fn empty_windows(u: &Verb, y: &Array, w: usize, ctx: &mut Ctx<'_>, span: Span) -> Array {
7122 let m = y.item_size();
7123 if u.is_pure() && let Some(cells) = w.checked_mul(m).filter(|&s| s <= 1 << 20) {
7124 let mut shape = y.shape.clone();
7125 shape[0] = w;
7126 let probe = Array::new(shape, fill_data(y.dtype(), cells));
7127 if let Ok(cell) = u.monad(&probe, ctx, span) {
7128 let mut shape = vec![0usize];
7129 shape.extend_from_slice(&cell.shape);
7130 return Array::new(shape, Data::empty(cell.dtype()));
7131 }
7132 }
7133 Array::new(vec![0], Data::empty(DType::I64))
7134}
7135
7136fn window_size(x: &Array, span: Span) -> Result<i64> {
7138 let v = x
7139 .to_i64_vec()
7140 .ok_or_else(|| Error::domain("the window size must be an integer", span))?;
7141 match v.as_slice() {
7142 [k] => Ok(*k),
7143 _ => Err(Error::new(
7144 ErrorKind::Length,
7145 "the window size must be a single number",
7146 Some(span),
7147 )),
7148 }
7149}
7150
7151fn infix(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
7159 let k = window_size(x, span)?;
7160 let promoted = as_items(y);
7161 let base = promoted.as_ref().unwrap_or(y);
7162 let n = base.items();
7163 let m = base.item_size();
7164 if k < 0 {
7165 let w = k.unsigned_abs() as usize;
7166 let count = n.div_ceil(w);
7167 let cells = each_cell(count, n * m, u.is_pure(), ctx, |i, c| {
7168 u.monad(§ion(base, i * w, ((i + 1) * w).min(n)), c, span)
7169 })?;
7170 return assemble(&[count], cells, span);
7171 }
7172 let w = k as usize;
7173 if n < w {
7174 return Ok(empty_windows(u, base, w, ctx, span));
7175 }
7176 let count = n - w + 1;
7177 if w > 0 && base.dtype().is_numeric()
7178 && let Some(op) = folded_op(u) && let Some(d) = window_typed(op, &base.data, n, m, w)
7179 {
7180 let mut shape = base.shape.clone();
7181 shape[0] = count;
7182 return Ok(Array::new(shape, d));
7183 }
7184 let work = count.saturating_mul(w).saturating_mul(m);
7185 let cells = each_cell(count, work, u.is_pure(), ctx, |i, c| {
7186 u.monad(§ion(base, i, i + w), c, span)
7187 })?;
7188 assemble(&[count], cells, span)
7189}
7190
7191fn power(
7194 u: &Verb,
7195 p: Power,
7196 x: Option<&Array>,
7197 y: &Array,
7198 ctx: &mut Ctx<'_>,
7199 span: Span,
7200) -> Result<Array> {
7201 let step = |acc: &Array, c: &mut Ctx<'_>| match x {
7202 Some(x) => u.dyad(x, acc, c, span),
7203 None => u.monad(acc, c, span),
7204 };
7205 match p {
7206 Power::Times(n) => {
7207 let mut acc = y.clone();
7208 for _ in 0..n {
7209 acc = step(&acc, ctx)?;
7210 }
7211 Ok(acc)
7212 }
7213 Power::Converge => {
7214 let mut acc = y.clone();
7215 for _ in 0..CONVERGE_LIMIT {
7216 let next = step(&acc, ctx)?;
7217 if arrays_match(&next, &acc, ctx.cfg.tol) {
7218 return Ok(next);
7219 }
7220 acc = next;
7221 }
7222 Err(Error::domain("the iteration did not converge", span))
7223 }
7224 Power::Each(ref counts) => {
7228 let mut acc = y.clone();
7229 let mut done = 0u64;
7230 let mut order: Vec<usize> = (0..counts.len()).collect();
7231 order.sort_by_key(|&i| counts[i]);
7232 let mut cells: Vec<Option<Array>> = vec![None; counts.len()];
7233 for i in order {
7234 while done < counts[i] {
7235 acc = step(&acc, ctx)?;
7236 done += 1;
7237 }
7238 cells[i] = Some(acc.clone());
7239 }
7240 let cells: Vec<Array> = cells.into_iter().map(|c| c.expect("every count filled")).collect();
7241 assemble(&[cells.len()], cells, span)
7242 }
7243 Power::ConvergeTrace => {
7244 let mut acc = y.clone();
7245 let mut cells = vec![acc.clone()];
7246 for _ in 0..CONVERGE_LIMIT {
7247 let next = step(&acc, ctx)?;
7248 if arrays_match(&next, &acc, ctx.cfg.tol) {
7249 return assemble(&[cells.len()], cells, span);
7250 }
7251 cells.push(next.clone());
7252 acc = next;
7253 }
7254 Err(Error::domain("the iteration did not converge", span))
7255 }
7256 }
7257}
7258
7259fn power_v(
7262 u: &Verb,
7263 v: &Verb,
7264 x: Option<&Array>,
7265 y: &Array,
7266 ctx: &mut Ctx<'_>,
7267 span: Span,
7268) -> Result<Array> {
7269 let count = match x {
7270 Some(x) => v.dyad(x, y, ctx, span)?,
7271 None => v.monad(y, ctx, span)?,
7272 };
7273 let n = count
7274 .to_i64_vec()
7275 .ok_or_else(|| Error::domain("the power count must be an integer", span))?;
7276 if n.len() != 1 {
7277 return Err(Error::not_yet("a list of power counts (u^:v with several)", span));
7278 }
7279 let n = n[0];
7280 if n < 0 {
7281 return Err(Error::not_yet("a negative power (the verb's inverse)", span));
7282 }
7283 power(u, Power::Times(n as u64), x, y, ctx, span)
7284}
7285
7286fn power_until(
7288 u: &Verb,
7289 test: &Verb,
7290 y: &Array,
7291 ctx: &mut Ctx<'_>,
7292 span: Span,
7293) -> Result<Array> {
7294 let mut acc = y.clone();
7295 for _ in 0..CONVERGE_LIMIT {
7296 let next = u.monad(&acc, ctx, span)?;
7297 let done = test.dyad(&next, &acc, ctx, span)?;
7298 let stop = done
7299 .to_f64_vec()
7300 .ok_or_else(|| Error::domain("the ⍣ test must answer with numbers", span))?;
7301 if !stop.is_empty() && stop.iter().all(|&v| v != 0.0) {
7302 return Ok(next);
7303 }
7304 acc = next;
7305 }
7306 Err(Error::domain("the iteration did not converge", span))
7307}
7308
7309fn along_axis(
7316 u: &Verb,
7317 x: Option<&Array>,
7318 y: &Array,
7319 k: usize,
7320 ctx: &mut Ctx<'_>,
7321 span: Span,
7322) -> Result<Array> {
7323 if k >= y.rank().max(1) {
7324 return Err(Error::new(
7325 ErrorKind::Rank,
7326 format!("axis {k} does not exist on an argument of rank {}", y.rank()),
7327 Some(span),
7328 ));
7329 }
7330 let moved = axis_to_front(y, k);
7331 let r = moved.rank();
7332 let out = match x {
7333 Some(x) => u.dyad(x, &moved, ctx, span)?,
7334 None => u.monad(&moved, ctx, span)?,
7335 };
7336 if out.rank() == r {
7337 return Ok(front_to_axis(&out, k));
7338 }
7339 Ok(out)
7340}
7341
7342fn where_indices(y: &Array, origin: i64, boxed: bool, span: Span) -> Result<Array> {
7350 let counts = y
7351 .to_i64_vec()
7352 .ok_or_else(|| Error::domain("indices needs non-negative integers", span))?;
7353 if counts.iter().any(|&c| c < 0) {
7354 return Err(Error::domain("indices needs non-negative integers", span));
7355 }
7356 if !boxed || y.rank() < 2 {
7357 let mut out = Vec::new();
7358 for (i, &c) in counts.iter().enumerate() {
7359 for _ in 0..c {
7360 out.push(origin + i as i64);
7361 }
7362 }
7363 return Ok(Array::from_i64(out));
7364 }
7365 let r = y.rank();
7366 let mut coord = vec![0usize; r];
7367 let mut out: Vec<Array> = Vec::new();
7368 for &c in &counts {
7369 if c > 0 {
7370 let point =
7371 Array::from_i64(coord.iter().map(|&k| origin + k as i64).collect::<Vec<_>>());
7372 for _ in 0..c {
7373 out.push(point.clone());
7374 }
7375 }
7376 odometer(&mut coord, &y.shape);
7377 }
7378 Ok(Array::new(vec![out.len()], Data::Box(out.into())))
7379}
7380
7381fn interval_index(
7387 x: &Array,
7388 y: &Array,
7389 offset: i64,
7390 closed: bool,
7391 tol: Tol,
7392 span: Span,
7393) -> Result<Array> {
7394 let bounds = x
7395 .to_f64_vec()
7396 .ok_or_else(|| Error::domain("interval index needs numeric bounds", span))?;
7397 let vals = y
7398 .to_f64_vec()
7399 .ok_or_else(|| Error::domain("interval index needs numeric values", span))?;
7400 let out: Vec<i64> = vals
7401 .iter()
7402 .map(|&v| {
7403 let count =
7406 bounds.iter().filter(|&&b| if closed { !tol.lt(v, b) } else { tol.lt(b, v) });
7407 offset + count.count() as i64
7408 })
7409 .collect();
7410 Ok(Array::new(y.shape.clone(), Data::I64(out.into())))
7411}
7412
7413fn steps(y: &Array, span: Span) -> Result<Array> {
7416 let vals = y.to_f64_vec().ok_or_else(|| Error::domain("steps needs a number", span))?;
7417 let v = match vals.first() {
7418 Some(&v) if v.is_finite() => v,
7419 _ => return Err(Error::domain("steps needs a finite number", span)),
7420 };
7421 let n = (2.0 * v.abs()).floor();
7422 if n > 1e7 {
7423 return Err(Error::domain("steps would produce too many items", span));
7424 }
7425 let n = n as i64 + 1;
7426 let step = if v < 0.0 { -1.0 } else { 1.0 };
7427 let start = -v;
7428 if v.fract() == 0.0 {
7429 let start = start as i64;
7430 let step = step as i64;
7431 return Ok(Array::from_i64((0..n).map(|k| start + k * step).collect()));
7432 }
7433 Ok(Array::from_f64((0..n).map(|k| start + k as f64 * step).collect()))
7434}
7435
7436fn index_of_last(x: &Array, y: &Array, origin: i64, tol: Tol) -> Array {
7438 let cell_rank = x.rank().saturating_sub(1).min(y.rank());
7439 let frame_rank = y.rank() - cell_rank;
7440 let frame: Vec<usize> = y.shape[..frame_rank].to_vec();
7441 let nf: usize = frame.iter().product();
7442 let items = x.items();
7443 let mut out = Vec::with_capacity(nf);
7444 for i in 0..nf {
7445 let cell = y.cell_at(frame_rank, i);
7446 let at = (0..items)
7447 .rev()
7448 .find(|&j| arrays_match(&cell, &item_or_self(x, j), tol))
7449 .unwrap_or(items);
7450 out.push(origin + at as i64);
7451 }
7452 Array::new(frame, Data::I64(out.into()))
7453}
7454
7455fn roll(
7462 y: &Array,
7463 origin: i64,
7464 fixed: bool,
7465 float_at_zero: bool,
7466 span: Span,
7467) -> Result<Array> {
7468 let bounds = y
7469 .to_i64_vec()
7470 .ok_or_else(|| Error::domain("roll needs whole numbers", span))?;
7471 if bounds.iter().any(|&b| b < 0) {
7472 return Err(Error::domain("roll needs non-negative numbers", span));
7473 }
7474 if !float_at_zero && bounds.contains(&0) {
7475 return Err(Error::domain("? 0 has no value: the range is empty", span));
7476 }
7477 let any_zero = bounds.contains(&0);
7479 crate::rng::with(fixed, |g| {
7480 if any_zero {
7481 let out: Vec<f64> = bounds
7482 .iter()
7483 .map(|&b| {
7484 if b == 0 {
7485 g.unit()
7486 } else {
7487 (origin + g.below(b as u64) as i64) as f64
7488 }
7489 })
7490 .collect();
7491 return Ok(Array::new(y.shape.clone(), Data::F64(out.into())));
7492 }
7493 let out: Vec<i64> =
7494 bounds.iter().map(|&b| origin + g.below(b as u64) as i64).collect();
7495 Ok(Array::new(y.shape.clone(), Data::I64(out.into())))
7496 })
7497}
7498
7499fn deal(x: &Array, y: &Array, origin: i64, fixed: bool, span: Span) -> Result<Array> {
7501 let want = one_whole(x, "the count dealt", span)?;
7502 let from = one_whole(y, "the range dealt from", span)?;
7503 if want < 0 || from < 0 {
7504 return Err(Error::domain("deal needs non-negative numbers", span));
7505 }
7506 if want > from {
7507 return Err(Error::domain(
7508 format!("cannot deal {want} distinct value(s) from {from}"),
7509 span,
7510 ));
7511 }
7512 if want == 0 {
7513 return Ok(Array::from_i64(Vec::new()));
7514 }
7515 let drawn = crate::rng::with(fixed, |g| g.deal(want as usize, from as u64));
7516 Ok(Array::from_i64(drawn.into_iter().map(|v| v + origin).collect()))
7517}
7518
7519fn one_whole(a: &Array, what: &str, span: Span) -> Result<i64> {
7521 let v = a
7522 .to_i64_vec()
7523 .ok_or_else(|| Error::domain(format!("{what} must be a whole number"), span))?;
7524 match v[..] {
7525 [n] => Ok(n),
7526 _ => Err(Error::new(
7527 ErrorKind::Rank,
7528 format!("{what} must be one number"),
7529 Some(span),
7530 )),
7531 }
7532}
7533
7534fn nth_prime(n: i64, span: Span) -> Result<i64> {
7538 if n < 0 {
7539 return Err(Error::domain("the prime index must not be negative", span));
7540 }
7541 const LIMIT: i64 = 5_000_000;
7542 if n >= LIMIT {
7543 return Err(Error::domain(
7544 format!("prime index {n} is beyond the {LIMIT}th prime"),
7545 span,
7546 ));
7547 }
7548 let k = (n + 1) as f64;
7551 let bound = if n < 6 { 15.0 } else { k * (k.ln() + k.ln().ln()) };
7552 let bound = bound.ceil() as usize + 1;
7553 let mut sieve = vec![true; bound + 1];
7554 sieve[0] = false;
7555 if bound >= 1 {
7556 sieve[1] = false;
7557 }
7558 let mut p = 2usize;
7559 while p * p <= bound {
7560 if sieve[p] {
7561 let mut q = p * p;
7562 while q <= bound {
7563 sieve[q] = false;
7564 q += p;
7565 }
7566 }
7567 p += 1;
7568 }
7569 let mut seen = 0i64;
7570 for (v, &is_p) in sieve.iter().enumerate() {
7571 if is_p {
7572 if seen == n {
7573 return Ok(v as i64);
7574 }
7575 seen += 1;
7576 }
7577 }
7578 Err(Error::internal("the prime sieve was too small"))
7579}
7580
7581fn prime_factors(n: i64, span: Span) -> Result<Vec<i64>> {
7583 if n < 1 {
7584 return Err(Error::domain("prime factors need a positive integer", span));
7585 }
7586 let mut out = Vec::new();
7587 let mut m = n;
7588 let mut d = 2i64;
7589 while d.saturating_mul(d) <= m {
7590 while m % d == 0 {
7591 out.push(d);
7592 m /= d;
7593 }
7594 d += if d == 2 { 1 } else { 2 };
7595 }
7596 if m > 1 {
7597 out.push(m);
7598 }
7599 Ok(out)
7600}
7601
7602fn lstsq(a: &[f64], m: usize, n: usize, b: &[f64], k: usize) -> Option<Vec<f64>> {
7610 let mut r = a.to_vec();
7612 let mut c = b.to_vec();
7613 let at = |i: usize, j: usize, w: usize| i * w + j;
7614 let scale = a.iter().fold(0.0f64, |acc, v| acc.max(v.abs()));
7615 if scale == 0.0 {
7616 return None;
7617 }
7618 for j in 0..n {
7619 let norm = (j..m).map(|i| r[at(i, j, n)] * r[at(i, j, n)]).sum::<f64>().sqrt();
7621 if norm <= 1e-13 * scale {
7622 return None;
7623 }
7624 let alpha = if r[at(j, j, n)] > 0.0 { -norm } else { norm };
7625 let mut v = vec![0.0f64; m];
7626 for i in j..m {
7627 v[i] = r[at(i, j, n)];
7628 }
7629 v[j] -= alpha;
7630 let vnorm2: f64 = (j..m).map(|i| v[i] * v[i]).sum();
7631 if vnorm2 > 0.0 {
7632 for col in j..n {
7633 let dot: f64 = (j..m).map(|i| v[i] * r[at(i, col, n)]).sum();
7634 let f = 2.0 * dot / vnorm2;
7635 for i in j..m {
7636 r[at(i, col, n)] -= f * v[i];
7637 }
7638 }
7639 for col in 0..k {
7640 let dot: f64 = (j..m).map(|i| v[i] * c[at(i, col, k)]).sum();
7641 let f = 2.0 * dot / vnorm2;
7642 for i in j..m {
7643 c[at(i, col, k)] -= f * v[i];
7644 }
7645 }
7646 }
7647 }
7648 let mut x = vec![0.0f64; n * k];
7650 for col in 0..k {
7651 for i in (0..n).rev() {
7652 let mut acc = c[at(i, col, k)];
7653 for j in i + 1..n {
7654 acc -= r[at(i, j, n)] * x[at(j, col, k)];
7655 }
7656 let d = r[at(i, i, n)];
7657 if d.abs() <= 1e-13 * scale {
7658 return None;
7659 }
7660 x[at(i, col, k)] = acc / d;
7661 }
7662 }
7663 Some(x)
7664}
7665
7666fn as_matrix(a: &Array, span: Span) -> Result<(Vec<f64>, usize, usize)> {
7669 let v = a
7670 .to_f64_vec()
7671 .ok_or_else(|| Error::domain("matrix division needs numeric data", span))?;
7672 match a.rank() {
7673 0 => Ok((v, 1, 1)),
7674 1 => {
7675 let m = a.shape[0];
7676 Ok((v, m, 1))
7677 }
7678 2 => Ok((v, a.shape[0], a.shape[1])),
7679 _ => Err(Error::new(
7680 ErrorKind::Rank,
7681 "matrix division needs an argument of rank 2 or less",
7682 Some(span),
7683 )),
7684 }
7685}
7686
7687fn matrix_inverse(y: &Array, span: Span) -> Result<Array> {
7691 let (a, m, n) = as_matrix(y, span)?;
7692 if m < n {
7693 return Err(Error::new(
7694 ErrorKind::Length,
7695 format!("cannot invert a {m} by {n} matrix: it has more columns than rows"),
7696 Some(span),
7697 ));
7698 }
7699 let mut eye = vec![0.0f64; m * m];
7700 for i in 0..m {
7701 eye[i * m + i] = 1.0;
7702 }
7703 let x = lstsq(&a, m, n, &eye, m)
7704 .ok_or_else(|| Error::domain("the matrix is singular", span))?;
7705 let shape = if y.rank() == 2 { vec![n, m] } else { y.shape.clone() };
7708 Ok(Array::new(shape, Data::F64(x.into())))
7709}
7710
7711fn matrix_divide(x: &Array, y: &Array, span: Span) -> Result<Array> {
7713 let (a, m, n) = as_matrix(y, span)?;
7714 let (b, bm, k) = as_matrix(x, span)?;
7715 if bm != m {
7716 return Err(Error::new(
7717 ErrorKind::Length,
7718 format!("the system has {m} rows but the right-hand side has {bm}"),
7719 Some(span),
7720 ));
7721 }
7722 if m < n {
7723 return Err(Error::new(
7724 ErrorKind::Length,
7725 format!("the {m} by {n} system is underdetermined"),
7726 Some(span),
7727 ));
7728 }
7729 let sol = lstsq(&a, m, n, &b, k)
7730 .ok_or_else(|| Error::domain("the system is singular", span))?;
7731 let shape = if x.rank() == 2 { vec![n, k] } else { vec![n] };
7734 Ok(Array::new(shape, Data::F64(sol.into())))
7735}
7736
7737fn squad(x: &Array, y: &Array, origin: i64, span: Span) -> Result<Array> {
7741 if x.rank() > 1 {
7742 return Err(Error::new(
7743 ErrorKind::Rank,
7744 "the index of ⌷ must be a scalar or a vector",
7745 Some(span),
7746 ));
7747 }
7748 let items: Vec<Array> = if x.rank() == 0 { vec![x.clone()] } else { x.cells(1) };
7751 if items.len() != y.rank() {
7752 return Err(Error::new(
7753 ErrorKind::Rank,
7754 format!("{} index(es) for an argument of rank {}", items.len(), y.rank()),
7755 Some(span),
7756 ));
7757 }
7758 let mut specs = Vec::with_capacity(items.len());
7759 let mut shape = Vec::new();
7760 for (k, item) in items.iter().enumerate() {
7761 let spec = match item.as_boxes() {
7762 Some(bs) if item.rank() == 0 => bs[0].clone(),
7763 _ => item.clone(),
7764 };
7765 let idx = spec
7766 .to_i64_vec()
7767 .ok_or_else(|| Error::domain("index must be an integer", span))?;
7768 for &i in &idx {
7769 let j = i - origin;
7770 if j < 0 || j as usize >= y.shape[k] {
7771 return Err(Error::domain(
7772 format!("index {i} is out of range on axis {k}"),
7773 span,
7774 ));
7775 }
7776 }
7777 shape.extend_from_slice(&spec.shape);
7778 specs.push((spec.shape.clone(), idx));
7779 }
7780 let y = y.to_row_major();
7781 let st = strides(&y.shape);
7782 let total: usize = shape.iter().product();
7783 let mut data = Data::empty(y.dtype());
7784 let mut coord = vec![0usize; shape.len()];
7785 for _ in 0..total {
7786 let mut at = 0usize;
7787 let mut used = 0usize;
7788 for (k, (sshape, idx)) in specs.iter().enumerate() {
7789 let sst = strides(sshape);
7790 let pick: usize = (0..sshape.len()).map(|a| coord[used + a] * sst[a]).sum();
7791 used += sshape.len();
7792 at += (idx[pick] - origin) as usize * st[k];
7793 }
7794 push_elem(&mut data, y.row_major_data(), at);
7795 odometer(&mut coord, &shape);
7796 }
7797 Ok(Array::new(shape, data))
7798}
7799
7800fn select_axis(
7807 x: &Array,
7808 y: &Array,
7809 axis: usize,
7810 rank: usize,
7811 origin: i64,
7812 span: Span,
7813) -> Result<Array> {
7814 if rank != 0 && y.rank() != rank {
7815 return Err(Error::new(
7816 ErrorKind::Rank,
7817 format!("{rank} index slot(s) for an argument of rank {}", y.rank()),
7818 Some(span),
7819 ));
7820 }
7821 if axis >= y.rank() {
7822 return Err(Error::new(
7823 ErrorKind::Rank,
7824 format!("axis {axis} does not exist on an argument of rank {}", y.rank()),
7825 Some(span),
7826 ));
7827 }
7828 let idx = x
7829 .to_i64_vec()
7830 .ok_or_else(|| Error::domain("index must be an integer", span))?;
7831 let len = y.shape[axis];
7832 let mut picks = Vec::with_capacity(idx.len());
7833 for &i in &idx {
7834 let j = i - origin;
7835 if j < 0 || j as usize >= len {
7836 return Err(Error::domain(
7837 format!("index {i} is out of range: axis {axis} has {len} items"),
7838 span,
7839 ));
7840 }
7841 picks.push(j as usize);
7842 }
7843 let mut shape = Vec::with_capacity(y.rank() + x.rank());
7844 shape.extend_from_slice(&y.shape[..axis]);
7845 shape.extend_from_slice(&x.shape);
7846 shape.extend_from_slice(&y.shape[axis + 1..]);
7847 let outer: usize = y.shape[..axis].iter().product();
7848 let inner: usize = y.shape[axis + 1..].iter().product();
7849 let mut data = Data::empty(y.dtype());
7850 for o in 0..outer {
7851 for &p in &picks {
7852 let base = (o * len + p) * inner;
7853 for e in 0..inner {
7854 push_elem(&mut data, &y.data, base + e);
7855 }
7856 }
7857 }
7858 Ok(Array::new(shape, data))
7859}
7860
7861fn amend(m: &Array, x: &Array, y: &Array, span: Span) -> Result<Array> {
7865 if y.rank() == 0 {
7866 return Err(Error::new(ErrorKind::Rank, "cannot amend a scalar", Some(span)));
7867 }
7868 if let Some(spec) = m.as_boxes().and_then(<[Array]>::first) {
7870 let spec = index_spec(spec, y, span)?;
7871 return amend_spec(&spec, x, y, span);
7872 }
7873 let idx = m
7874 .to_i64_vec()
7875 .ok_or_else(|| Error::domain("amend indices must be integers", span))?;
7876 let items = y.items() as i64;
7877 let mut at = Vec::with_capacity(idx.len());
7878 for &i in &idx {
7879 let k = if i < 0 { i + items } else { i };
7880 if k < 0 || k >= items {
7881 return Err(Error::domain(
7882 format!("index {i} is out of range: the argument has {items} items"),
7883 span,
7884 ));
7885 }
7886 at.push(k as usize);
7887 }
7888 let cell = y.item_size();
7889 let per_index = if x.count() == cell {
7890 false
7891 } else if x.count() == cell * at.len() {
7892 true
7893 } else {
7894 return Err(Error::new(
7895 ErrorKind::Length,
7896 format!(
7897 "cannot amend {} item(s) of {} element(s) each with {} element(s)",
7898 at.len(),
7899 cell,
7900 x.count()
7901 ),
7902 Some(span),
7903 ));
7904 };
7905 let Some(t) = DType::promote(x.dtype(), y.dtype()) else {
7908 return Err(Error::new(
7909 ErrorKind::Type,
7910 "the replacement and the argument hold different kinds of value",
7911 Some(span),
7912 ));
7913 };
7914 let (Some(src), Some(base)) = (x.data.cast(t), y.data.cast(t)) else {
7915 return Err(Error::new(
7916 ErrorKind::Type,
7917 "the replacement and the argument hold different kinds of value",
7918 Some(span),
7919 ));
7920 };
7921 let mut data = Data::empty(t);
7923 let mut plan: Vec<Option<usize>> = vec![None; y.items()];
7924 for (n, &k) in at.iter().enumerate() {
7925 plan[k] = Some(if per_index { n } else { 0 });
7926 }
7927 for (i, slot) in plan.iter().enumerate() {
7928 match slot {
7929 Some(n) => {
7930 for e in 0..cell {
7931 push_elem(&mut data, &src, n * cell + e);
7932 }
7933 }
7934 None => {
7935 for e in 0..cell {
7936 push_elem(&mut data, &base, i * cell + e);
7937 }
7938 }
7939 }
7940 }
7941 Ok(Array::new(y.shape.clone(), data))
7942}
7943
7944fn fetch(x: &Array, y: &Array, span: Span) -> Result<Array> {
7949 let steps: Vec<Array> = match x.as_boxes() {
7950 Some(bs) => bs.to_vec(),
7951 None => vec![x.clone()],
7952 };
7953 let mut cur = y.clone();
7954 for step in steps {
7955 let idx = if step.count() == 0 {
7958 Vec::new()
7959 } else {
7960 step.to_i64_vec()
7961 .ok_or_else(|| Error::domain("a fetch path holds integers", span))?
7962 };
7963 let base =
7965 if cur.rank() == 0 { Array::new(vec![1], cur.data.clone()) } else { cur.clone() };
7966 if idx.len() > base.rank() {
7967 return Err(Error::new(
7968 ErrorKind::Length,
7969 format!(
7970 "a path step of {} index(es) into a value of rank {}",
7971 idx.len(),
7972 cur.rank()
7973 ),
7974 Some(span),
7975 ));
7976 }
7977 let at = cell_index(&base, &idx, span)?;
7978 cur = open_cell(&base.cell_at(idx.len(), at));
7979 }
7980 Ok(cur)
7981}
7982
7983fn cell_index(y: &Array, idx: &[i64], span: Span) -> Result<usize> {
7985 let mut at = 0usize;
7986 for (k, &i) in idx.iter().enumerate() {
7987 let len = y.shape[k] as i64;
7988 let j = if i < 0 { i + len } else { i };
7989 if j < 0 || j >= len {
7990 return Err(Error::domain(
7991 format!("index {i} is out of range: axis {k} has {len} items"),
7992 span,
7993 ));
7994 }
7995 at = at * y.shape[k] + j as usize;
7996 }
7997 Ok(at)
7998}
7999
8000fn partition_enclose(x: &Array, y: &Array, span: Span) -> Result<Array> {
8009 if y.rank() > 1 {
8012 let last = y.shape[y.rank() - 1];
8013 let rows = y.count() / last.max(1);
8014 let mut cells: Vec<Array> = Vec::new();
8015 let mut width = None;
8016 for r in 0..rows {
8017 let row = Array::new(vec![last], y.data.slice(r * last, (r + 1) * last));
8018 let parts = partition_enclose(x, &row, span)?;
8019 let n = parts.count();
8020 if *width.get_or_insert(n) != n {
8021 return Err(Error::internal("partitions of unequal count"));
8022 }
8023 match parts.data {
8024 Data::Box(v) => cells.extend(v.as_slice().iter().cloned()),
8025 _ => return Err(Error::internal("a partition is boxed")),
8026 }
8027 }
8028 let mut shape = y.shape[..y.rank() - 1].to_vec();
8029 shape.push(width.unwrap_or(0));
8030 return Ok(Array::new(shape, Data::Box(cells.into())));
8031 }
8032 if y.rank() == 0 {
8033 return Err(Error::new(
8034 ErrorKind::Rank,
8035 "partitioned enclose needs an array to partition",
8036 Some(span),
8037 ));
8038 }
8039 let flags = x
8040 .to_i64_vec()
8041 .ok_or_else(|| Error::domain("partition flags must be integers", span))?;
8042 if flags.iter().any(|&f| f < 0) {
8043 return Err(Error::domain("partition flags must not be negative", span));
8044 }
8045 if flags.len() != y.shape[0] {
8046 return Err(Error::new(
8047 ErrorKind::Length,
8048 format!("{} flag(s) for {} item(s)", flags.len(), y.shape[0]),
8049 Some(span),
8050 ));
8051 }
8052 let mut parts: Vec<Array> = Vec::new();
8053 let mut cur: Option<Data> = None;
8054 let mut prev = 0i64;
8055 for (i, &f) in flags.iter().enumerate() {
8056 if f > prev {
8057 if let Some(d) = cur.take() {
8058 parts.push(Array::new(vec![d.len()], d));
8059 }
8060 cur = Some(Data::empty(y.dtype()));
8061 }
8062 prev = f;
8063 if f == 0 {
8064 continue;
8065 }
8066 if let Some(d) = cur.as_mut() {
8067 push_elem(d, &y.data, i);
8068 }
8069 }
8070 if let Some(d) = cur.take() {
8071 parts.push(Array::new(vec![d.len()], d));
8072 }
8073 Ok(Array::new(vec![parts.len()], Data::Box(parts.into())))
8074}
8075
8076fn key(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
8079 let keys = if x.rank() == 0 { Array::new(vec![1], x.data.clone()) } else { x.clone() };
8080 let n = keys.items();
8081 if n != y.items() && !(y.rank() == 0 && n == 1) {
8082 return Err(Error::new(
8083 ErrorKind::Length,
8084 format!("{n} key(s) for {} item(s)", y.items()),
8085 Some(span),
8086 ));
8087 }
8088 let tol = ctx.cfg.tol;
8089 let mut order: Vec<usize> = Vec::new();
8090 let mut groups: Vec<Vec<usize>> = Vec::new();
8091 for i in 0..n {
8092 let k = keys.item(i);
8093 match order.iter().position(|&j| arrays_match(&k, &keys.item(j), tol)) {
8094 Some(g) => groups[g].push(i),
8095 None => {
8096 order.push(i);
8097 groups.push(vec![i]);
8098 }
8099 }
8100 }
8101 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
8102 let mut cells = Vec::with_capacity(groups.len());
8103 for g in &groups {
8104 cells.push(u.monad(&select_items(&items, g), ctx, span)?);
8105 }
8106 assemble(&[groups.len()], cells, span)
8107}
8108
8109fn oblique(u: &Verb, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
8112 if y.rank() < 2 {
8113 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
8114 let n = items.items();
8115 let mut cells = Vec::with_capacity(n);
8116 for i in 0..n {
8117 cells.push(u.monad(&select_items(&items, &[i]), ctx, span)?);
8118 }
8119 return assemble(&[n], cells, span);
8120 }
8121 if y.rank() > 2 {
8122 return Err(Error::not_yet("oblique (u/.) on a rank-3 or higher argument", span));
8123 }
8124 let (rows, cols) = (y.shape[0], y.shape[1]);
8125 let mut cells = Vec::with_capacity(rows + cols - 1);
8126 for d in 0..rows + cols - 1 {
8127 let mut data = Data::empty(y.dtype());
8128 let mut len = 0usize;
8129 for i in 0..rows {
8130 if d >= i && d - i < cols {
8131 push_elem(&mut data, &y.data, i * cols + (d - i));
8132 len += 1;
8133 }
8134 }
8135 cells.push(u.monad(&Array::new(vec![len], data), ctx, span)?);
8136 }
8137 assemble(&[rows + cols - 1], cells, span)
8138}
8139
8140fn cut_ranges(frets: &[bool], mode: i64) -> Vec<(usize, usize)> {
8148 let n = frets.len();
8149 let mut out = Vec::new();
8150 if mode.abs() == 1 {
8151 let mut start: Option<usize> = None;
8152 for (i, &fret) in frets.iter().enumerate() {
8153 if fret {
8154 if let Some(s) = start {
8155 out.push((s, i));
8156 }
8157 start = Some(i);
8158 }
8159 }
8160 if let Some(s) = start {
8161 out.push((s, n));
8162 }
8163 if mode < 0 {
8164 return out.into_iter().map(|(s, e)| (s + 1, e)).collect();
8165 }
8166 } else {
8167 let mut start = 0usize;
8168 for (i, &fret) in frets.iter().enumerate() {
8169 if fret {
8170 out.push((start, i + 1));
8171 start = i + 1;
8172 }
8173 }
8174 if mode < 0 {
8175 return out.into_iter().map(|(s, e)| (s, e - 1)).collect();
8176 }
8177 }
8178 out
8179}
8180
8181fn cut(
8183 u: &Verb,
8184 x: Option<&Array>,
8185 y: &Array,
8186 mode: i64,
8187 ctx: &mut Ctx<'_>,
8188 span: Span,
8189) -> Result<Array> {
8190 if mode == 0 {
8191 let Some(x) = x else {
8192 return u.monad(&reverse_all_axes(y), ctx, span);
8193 };
8194 let (origin, size) = rectangle(x, span)?;
8195 let origin = origin.unwrap_or_else(|| vec![0; size.len()]);
8196 return u.monad(&subarray(y, &origin, &size, span)?, ctx, span);
8197 }
8198 if mode.abs() == 3 {
8199 let Some(x) = x else {
8200 return Err(Error::not_yet("monadic tessellation (u;.3 y)", span));
8201 };
8202 return tessellate(u, x, y, mode < 0, ctx, span);
8203 }
8204 if !matches!(mode, 1 | -1 | 2 | -2) {
8205 return Err(Error::not_yet(format!("cut (u;.{mode})"), span));
8206 }
8207 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
8208 let n = items.items();
8209 let tol = ctx.cfg.tol;
8210 let frets: Vec<bool> = match x {
8211 Some(x) => {
8212 let flags = x
8213 .to_i64_vec()
8214 .ok_or_else(|| Error::domain("cut frets must be integers", span))?;
8215 if let Some(&bad) = flags.iter().find(|&&f| f != 0 && f != 1) {
8218 return Err(Error::domain(format!("{bad} is not a fret: a fret is 0 or 1"), span));
8219 }
8220 if x.rank() == 0 {
8223 vec![flags[0] != 0; n]
8224 } else {
8225 if flags.len() != n {
8226 return Err(Error::new(
8227 ErrorKind::Length,
8228 format!("{} fret(s) for {n} item(s)", flags.len()),
8229 Some(span),
8230 ));
8231 }
8232 flags.iter().map(|&f| f != 0).collect()
8233 }
8234 }
8235 None => {
8236 if n == 0 {
8238 Vec::new()
8239 } else {
8240 let at = if mode.abs() == 1 { 0 } else { n - 1 };
8241 let mark = items.item(at);
8242 (0..n).map(|i| arrays_match(&items.item(i), &mark, tol)).collect()
8243 }
8244 }
8245 };
8246 let ranges = cut_ranges(&frets, mode);
8247 let mut cells = Vec::with_capacity(ranges.len());
8248 for (s, e) in &ranges {
8249 cells.push(u.monad(§ion(&items, *s, *e), ctx, span)?);
8250 }
8251 assemble(&[ranges.len()], cells, span)
8252}
8253
8254fn rectangle(x: &Array, span: Span) -> Result<(Option<Vec<i64>>, Vec<i64>)> {
8257 let values = x
8258 .to_i64_vec()
8259 .ok_or_else(|| Error::domain("a cut rectangle is whole numbers", span))?;
8260 match x.rank() {
8261 0 | 1 => Ok((None, values)),
8262 2 if x.shape[0] == 2 => {
8263 let n = x.shape[1];
8264 Ok((Some(values[..n].to_vec()), values[n..].to_vec()))
8265 }
8266 _ => Err(Error::new(
8267 ErrorKind::Rank,
8268 "a cut rectangle is a vector of sizes, or two rows of origins and sizes",
8269 Some(span),
8270 )),
8271 }
8272}
8273
8274fn subarray(y: &Array, origin: &[i64], size: &[i64], span: Span) -> Result<Array> {
8278 if origin.len() > y.rank() {
8279 return Err(Error::new(
8280 ErrorKind::Rank,
8281 format!("a cut of {} axis/axes into a rank-{} value", origin.len(), y.rank()),
8282 Some(span),
8283 ));
8284 }
8285 let r = y.rank();
8286 let st = strides(&y.shape);
8287 let mut shape = y.shape.clone();
8288 let mut start = vec![0i64; r];
8289 let mut step = vec![1i64; r];
8290 for k in 0..origin.len() {
8291 let len = size[k].unsigned_abs() as usize;
8292 let from = if origin[k] < 0 { origin[k] + y.shape[k] as i64 } else { origin[k] };
8293 if from < 0 || from + len as i64 > y.shape[k] as i64 {
8294 return Err(Error::domain(
8295 format!("a cut of {len} from {from} leaves axis {k} of {}", y.shape[k]),
8296 span,
8297 ));
8298 }
8299 shape[k] = len;
8300 if size[k] < 0 {
8301 start[k] = from + len as i64 - 1;
8302 step[k] = -1;
8303 } else {
8304 start[k] = from;
8305 }
8306 }
8307 Ok(gather(y, &shape, &start, &step, &st))
8308}
8309
8310fn gather(y: &Array, shape: &[usize], start: &[i64], step: &[i64], st: &[usize]) -> Array {
8312 let n: usize = shape.iter().product();
8313 let mut data = Data::empty(y.dtype());
8314 let mut coord = vec![0usize; shape.len()];
8315 for _ in 0..n {
8316 let idx: usize = (0..shape.len())
8317 .map(|k| (start[k] + step[k] * coord[k] as i64) as usize * st[k])
8318 .sum();
8319 push_elem(&mut data, &y.data, idx);
8320 odometer(&mut coord, shape);
8321 }
8322 Array::new(shape.to_vec(), data)
8323}
8324
8325fn tessellate(
8329 u: &Verb,
8330 x: &Array,
8331 y: &Array,
8332 complete: bool,
8333 ctx: &mut Ctx<'_>,
8334 span: Span,
8335) -> Result<Array> {
8336 let (movement, size) = rectangle(x, span)?;
8338 if size.iter().any(|&s| s < 0) && movement.is_none() {
8343 return Err(Error::not_yet(
8344 "a negative block size without a movement row (x u;.3 y)",
8345 span,
8346 ));
8347 }
8348 let movement = movement.unwrap_or_else(|| vec![1; size.len()]);
8349 if size.len() > y.rank() {
8350 return Err(Error::new(
8351 ErrorKind::Rank,
8352 format!("a tessellation of {} axis/axes into a rank-{} value", size.len(), y.rank()),
8353 Some(span),
8354 ));
8355 }
8356 let mut frame = Vec::with_capacity(size.len());
8357 for k in 0..size.len() {
8358 let (len, step, block) = (y.shape[k] as i64, movement[k], size[k].abs());
8359 if step <= 0 {
8360 return Err(Error::domain("a tessellation moves by a positive step", span));
8361 }
8362 let count = if complete {
8363 if len < block { 0 } else { (len - block) / step + 1 }
8364 } else {
8365 (len + step - 1) / step
8366 };
8367 frame.push(count as usize);
8368 }
8369 let total: usize = frame.iter().product();
8370 let mut cells = Vec::with_capacity(total);
8371 let mut coord = vec![0usize; frame.len()];
8372 for _ in 0..total {
8373 let origin: Vec<i64> = (0..frame.len()).map(|k| coord[k] as i64 * movement[k]).collect();
8374 let block: Vec<i64> = (0..frame.len())
8377 .map(|k| {
8378 let len = size[k].abs().min(y.shape[k] as i64 - origin[k]);
8379 if size[k] < 0 { -len } else { len }
8380 })
8381 .collect();
8382 cells.push(u.monad(&subarray(y, &origin, &block, span)?, ctx, span)?);
8383 odometer(&mut coord, &frame);
8384 }
8385 assemble(&frame, cells, span)
8386}
8387
8388fn reverse_all_axes(y: &Array) -> Array {
8390 if y.rank() == 0 {
8391 return y.clone();
8392 }
8393 let st = strides(&y.shape);
8394 let n = y.count();
8395 let r = y.rank();
8396 let mut data = Data::empty(y.dtype());
8397 let mut coord = vec![0usize; r];
8398 for _ in 0..n {
8399 let idx: usize = (0..r).map(|k| (y.shape[k] - 1 - coord[k]) * st[k]).sum();
8400 push_elem(&mut data, &y.data, idx);
8401 odometer(&mut coord, &y.shape);
8402 }
8403 Array::new(y.shape.clone(), data)
8404}
8405
8406fn axis_to_front(y: &Array, k: usize) -> Array {
8410 if k == 0 || y.rank() < 2 {
8411 return y.clone();
8412 }
8413 let r = y.rank();
8414 let src: Vec<usize> = std::iter::once(k).chain((0..r).filter(|&a| a != k)).collect();
8415 permute_axes(y, &src)
8416}
8417
8418fn front_to_axis(y: &Array, k: usize) -> Array {
8420 if k == 0 || y.rank() < 2 {
8421 return y.clone();
8422 }
8423 let r = y.rank();
8424 let mut src = Vec::with_capacity(r);
8427 for a in 0..r {
8428 src.push(match a.cmp(&k) {
8429 std::cmp::Ordering::Less => a + 1,
8430 std::cmp::Ordering::Equal => 0,
8431 std::cmp::Ordering::Greater => a,
8432 });
8433 }
8434 permute_axes(y, &src)
8435}
8436
8437fn transpose_to(y: &Array, dest: &[usize], span: Span) -> Result<Array> {
8442 let rank_out = dest.iter().copied().max().map_or(0, |m| m + 1);
8443 let mut out_shape = vec![usize::MAX; rank_out];
8444 for (a, &d) in dest.iter().enumerate() {
8445 out_shape[d] = out_shape[d].min(y.shape[a]);
8446 }
8447 if out_shape.contains(&usize::MAX) {
8448 return Err(Error::new(
8449 ErrorKind::Domain,
8450 "a transpose must name every axis of the result",
8451 Some(span),
8452 ));
8453 }
8454 let y = y.to_row_major();
8455 let st = strides(&y.shape);
8456 let n: usize = out_shape.iter().product();
8457 let mut data = Data::empty(y.dtype());
8458 let mut coord = vec![0usize; rank_out];
8459 for _ in 0..n {
8460 let idx: usize = dest.iter().enumerate().map(|(a, &d)| coord[d] * st[a]).sum();
8461 push_elem(&mut data, &y.data, idx);
8462 odometer(&mut coord, &out_shape);
8463 }
8464 Ok(Array::new(out_shape, data))
8465}
8466
8467fn transpose_apl(x: &Array, y: &Array, io: i64, span: Span) -> Result<Array> {
8470 let axes = x
8471 .to_i64_vec()
8472 .ok_or_else(|| Error::domain("a transpose is given whole numbers", span))?;
8473 if axes.len() != y.rank() {
8474 return Err(Error::new(
8475 ErrorKind::Length,
8476 format!("{} axes for a rank-{} value", axes.len(), y.rank()),
8477 Some(span),
8478 ));
8479 }
8480 let mut dest = Vec::with_capacity(axes.len());
8481 for a in axes {
8482 let d = a - io;
8483 if d < 0 || d as usize >= y.rank() {
8484 return Err(Error::new(
8485 ErrorKind::Domain,
8486 format!("axis {a} is outside a rank-{} value", y.rank()),
8487 Some(span),
8488 ));
8489 }
8490 dest.push(d as usize);
8491 }
8492 transpose_to(y, &dest, span)
8493}
8494
8495fn transpose_j(x: &Array, y: &Array, span: Span) -> Result<Array> {
8499 let groups: Vec<Vec<i64>> = match x.as_boxes() {
8500 Some(bs) => bs
8501 .iter()
8502 .map(|b| {
8503 b.to_i64_vec().ok_or_else(|| {
8504 Error::domain("a transpose is given whole numbers", span)
8505 })
8506 })
8507 .collect::<Result<Vec<_>>>()?,
8508 None => x
8509 .to_i64_vec()
8510 .ok_or_else(|| Error::domain("a transpose is given whole numbers", span))?
8511 .into_iter()
8512 .map(|a| vec![a])
8513 .collect(),
8514 };
8515 let r = y.rank();
8516 let mut group_of = vec![None; r];
8519 for (g, axes) in groups.iter().enumerate() {
8520 for &a in axes {
8521 let k = if a < 0 { a + r as i64 } else { a };
8522 if k < 0 || k as usize >= r {
8523 return Err(Error::new(
8524 ErrorKind::Domain,
8525 format!("axis {a} is outside a rank-{r} value"),
8526 Some(span),
8527 ));
8528 }
8529 if group_of[k as usize].is_some() {
8530 return Err(Error::new(
8531 ErrorKind::Domain,
8532 format!("axis {a} is named twice in a transpose"),
8533 Some(span),
8534 ));
8535 }
8536 group_of[k as usize] = Some(g);
8537 }
8538 }
8539 let leading = group_of.iter().filter(|g| g.is_none()).count();
8540 let mut dest = vec![0usize; r];
8541 let mut next = 0;
8542 for a in 0..r {
8543 match group_of[a] {
8544 None => {
8545 dest[a] = next;
8546 next += 1;
8547 }
8548 Some(g) => dest[a] = leading + g,
8549 }
8550 }
8551 transpose_to(y, &dest, span)
8552}
8553
8554fn permute_axes(y: &Array, src: &[usize]) -> Array {
8556 let st = strides(&y.shape);
8557 let out_shape: Vec<usize> = src.iter().map(|&a| y.shape[a]).collect();
8558 let n = y.count();
8559 let mut data = Data::empty(y.dtype());
8560 let mut coord = vec![0usize; src.len()];
8561 for _ in 0..n {
8562 let idx: usize = (0..src.len()).map(|a| coord[a] * st[src[a]]).sum();
8563 push_elem(&mut data, &y.data, idx);
8564 odometer(&mut coord, &out_shape);
8565 }
8566 Array::new(out_shape, data)
8567}
8568
8569struct Spec {
8573 width: usize,
8575 cells: Vec<Vec<usize>>,
8577 shape: Vec<usize>,
8580}
8581
8582fn axis_position(v: i64, len: usize, span: Span) -> Result<usize> {
8585 let p = if v < 0 { v + len as i64 } else { v };
8586 if p < 0 || p >= len as i64 {
8587 return Err(Error::domain(
8588 format!("index {v} is out of range: the axis has {len} element(s)"),
8589 span,
8590 ));
8591 }
8592 Ok(p as usize)
8593}
8594
8595fn index_spec(content: &Array, y: &Array, span: Span) -> Result<Spec> {
8605 let too_deep = |n: usize| {
8606 Error::new(
8607 ErrorKind::Rank,
8608 format!("an index specification of {n} axis/axes into a rank-{} value", y.rank()),
8609 Some(span),
8610 )
8611 };
8612 if let Some(items) = content.as_boxes() {
8613 if items.len() > y.rank() {
8614 return Err(too_deep(items.len()));
8615 }
8616 let mut per_axis: Vec<Vec<usize>> = Vec::with_capacity(items.len());
8617 let mut shape: Vec<usize> = Vec::new();
8618 for (k, c) in items.iter().enumerate() {
8619 let len = y.shape[k];
8620 if c.as_boxes().is_some() {
8621 let inner = open_cell(c);
8622 let excluded = inner.to_i64_vec().ok_or_else(|| {
8623 Error::domain("an index complement holds integers", span)
8624 })?;
8625 let mut dropped = vec![false; len];
8626 for v in excluded {
8627 dropped[axis_position(v, len, span)?] = true;
8628 }
8629 let kept: Vec<usize> = (0..len).filter(|i| !dropped[*i]).collect();
8630 shape.push(kept.len());
8631 per_axis.push(kept);
8632 } else {
8633 let idx = c
8634 .to_i64_vec()
8635 .ok_or_else(|| Error::domain("an index holds integers", span))?;
8636 let mut positions = Vec::with_capacity(idx.len());
8637 for v in idx {
8638 positions.push(axis_position(v, len, span)?);
8639 }
8640 shape.extend_from_slice(&c.shape);
8641 per_axis.push(positions);
8642 }
8643 }
8644 let mut cells: Vec<Vec<usize>> = vec![Vec::new()];
8646 for positions in &per_axis {
8647 let mut next = Vec::with_capacity(cells.len() * positions.len());
8648 for prefix in &cells {
8649 for &p in positions {
8650 let mut cell = prefix.clone();
8651 cell.push(p);
8652 next.push(cell);
8653 }
8654 }
8655 cells = next;
8656 }
8657 return Ok(Spec { width: per_axis.len(), cells, shape });
8658 }
8659 let idx = content
8660 .to_i64_vec()
8661 .ok_or_else(|| Error::domain("an index specification holds integers", span))?;
8662 let rank = content.rank();
8663 let width = if rank == 0 { 1 } else { content.shape[rank - 1] };
8664 if width > y.rank() {
8665 return Err(too_deep(width));
8666 }
8667 let shape: Vec<usize> = if rank == 0 { Vec::new() } else { content.shape[..rank - 1].to_vec() };
8668 let count: usize = shape.iter().product();
8669 let mut cells: Vec<Vec<usize>> = Vec::new();
8670 if width == 0 {
8671 cells.resize(count, Vec::new());
8672 } else {
8673 for chunk in idx.chunks(width) {
8674 let mut cell = Vec::with_capacity(width);
8675 for (k, &v) in chunk.iter().enumerate() {
8676 cell.push(axis_position(v, y.shape[k], span)?);
8677 }
8678 cells.push(cell);
8679 }
8680 }
8681 Ok(Spec { width, cells, shape })
8682}
8683
8684fn spec_offset(st: &[usize], cell: &[usize]) -> usize {
8686 cell.iter().enumerate().map(|(k, &p)| p * st[k]).sum()
8687}
8688
8689fn select_spec(spec: &Spec, y: &Array) -> Array {
8691 let st = strides(&y.shape);
8692 let size: usize = y.shape[spec.width..].iter().product();
8693 let mut data = Data::empty(y.dtype());
8694 for cell in &spec.cells {
8695 let base = spec_offset(&st, cell);
8696 for e in 0..size {
8697 push_elem(&mut data, &y.data, base + e);
8698 }
8699 }
8700 let mut shape = spec.shape.clone();
8701 shape.extend_from_slice(&y.shape[spec.width..]);
8702 Array::new(shape, data)
8703}
8704
8705fn amend_spec(spec: &Spec, x: &Array, y: &Array, span: Span) -> Result<Array> {
8708 let size: usize = y.shape[spec.width..].iter().product();
8709 let per_cell = if x.count() == size {
8710 false
8711 } else if x.count() == size * spec.cells.len() {
8712 true
8713 } else {
8714 return Err(Error::new(
8715 ErrorKind::Length,
8716 format!(
8717 "cannot amend {} cell(s) of {size} element(s) each with {} element(s)",
8718 spec.cells.len(),
8719 x.count()
8720 ),
8721 Some(span),
8722 ));
8723 };
8724 let mismatch = || {
8725 Error::new(
8726 ErrorKind::Type,
8727 "the replacement and the argument hold different kinds of value",
8728 Some(span),
8729 )
8730 };
8731 let t = DType::promote(x.dtype(), y.dtype()).ok_or_else(mismatch)?;
8732 let (Some(src), Some(base)) = (x.data.cast(t), y.data.cast(t)) else {
8733 return Err(mismatch());
8734 };
8735 let st = strides(&y.shape);
8736 let mut plan: Vec<Option<usize>> = vec![None; y.count()];
8737 for (n, cell) in spec.cells.iter().enumerate() {
8738 let at = spec_offset(&st, cell);
8739 for e in 0..size {
8740 plan[at + e] = Some(if per_cell { n * size + e } else { e });
8741 }
8742 }
8743 let mut data = Data::empty(t);
8744 for (i, slot) in plan.iter().enumerate() {
8745 match slot {
8746 Some(n) => push_elem(&mut data, &src, *n),
8747 None => push_elem(&mut data, &base, i),
8748 }
8749 }
8750 Ok(Array::new(y.shape.clone(), data))
8751}
8752
8753fn map_paths(y: &Array) -> Array {
8762 fn coord_of(shape: &[usize], mut i: usize) -> Array {
8763 let mut out = vec![0i64; shape.len()];
8764 for k in (0..shape.len()).rev() {
8765 out[k] = (i % shape[k]) as i64;
8766 i /= shape[k];
8767 }
8768 Array::from_i64(out)
8769 }
8770 fn go(y: &Array, prefix: &[Array]) -> Array {
8771 let Some(boxes) = y.as_boxes() else {
8772 if prefix.is_empty() {
8773 return Array::new(vec![0], Data::I64(Vec::new().into()));
8774 }
8775 return Array::new(vec![prefix.len()], Data::Box(prefix.to_vec().into()));
8776 };
8777 let cells: Vec<Array> = boxes
8778 .iter()
8779 .enumerate()
8780 .map(|(i, b)| {
8781 let mut path = prefix.to_vec();
8782 path.push(coord_of(&y.shape, i));
8783 go(b, &path)
8784 })
8785 .collect();
8786 Array::new(y.shape.clone(), Data::Box(cells.into()))
8787 }
8788 go(y, &[])
8789}
8790
8791fn shift_fill(x: &Array, y: &Array, fill: &Array, span: Span) -> Result<Array> {
8796 let counts = axis_counts(x, "shift", span)?;
8797 if y.rank() == 0 {
8798 return Ok(y.clone());
8799 }
8800 if counts.len() > y.rank() {
8801 return Err(Error::new(
8802 ErrorKind::Length,
8803 format!("shift has {} amounts for an argument of rank {}", counts.len(), y.rank()),
8804 Some(span),
8805 ));
8806 }
8807 if fill.count() != 1 {
8808 return Err(Error::new(ErrorKind::Length, "a fill is one atom", Some(span)));
8809 }
8810 let mismatch = || {
8811 Error::new(ErrorKind::Type, "the fill and the argument differ in kind", Some(span))
8812 };
8813 let t = DType::promote(y.dtype(), fill.dtype()).ok_or_else(mismatch)?;
8814 let (Some(base), Some(f)) = (y.data.cast(t), fill.data.cast(t)) else {
8815 return Err(mismatch());
8816 };
8817 let st = strides(&y.shape);
8818 let r = y.rank();
8819 let mut data = Data::empty(t);
8820 let mut coord = vec![0usize; r];
8821 for _ in 0..y.count() {
8822 let mut idx = 0usize;
8823 let mut vacated = false;
8824 for k in 0..r {
8825 let from = coord[k] as i64 + counts.get(k).copied().unwrap_or(0);
8826 if from < 0 || from >= y.shape[k] as i64 {
8827 vacated = true;
8828 break;
8829 }
8830 idx += from as usize * st[k];
8831 }
8832 if vacated {
8833 push_elem(&mut data, &f, 0);
8834 } else {
8835 push_elem(&mut data, &base, idx);
8836 }
8837 odometer(&mut coord, &y.shape);
8838 }
8839 Ok(Array::new(y.shape.clone(), data))
8840}
8841
8842fn memo_key(a: &Array, out: &mut Vec<u64>) -> bool {
8847 out.push(a.rank() as u64);
8848 out.extend(a.shape.iter().map(|&n| n as u64));
8849 out.push(a.dtype() as u64);
8850 match &a.data {
8851 Data::Ext(_) | Data::Rat(_) => false,
8852 Data::Box(items) => items.iter().all(|item| memo_key(item, out)),
8853 d => {
8854 for i in 0..d.len() {
8855 out.push(elem_key(d, i));
8856 }
8857 true
8858 }
8859 }
8860}
8861
8862fn memoised(
8864 u: &Verb,
8865 cache: &MemoCache,
8866 x: Option<&Array>,
8867 y: &Array,
8868 ctx: &mut Ctx<'_>,
8869 span: Span,
8870) -> Result<Array> {
8871 let apply = |ctx: &mut Ctx<'_>| match x {
8872 Some(x) => u.dyad(x, y, ctx, span),
8873 None => u.monad(y, ctx, span),
8874 };
8875 let mut key = vec![u64::from(x.is_some())];
8876 let keyed = x.is_none_or(|x| memo_key(x, &mut key)) && memo_key(y, &mut key);
8877 if !keyed {
8878 return apply(ctx);
8879 }
8880 if let Ok(map) = cache.lock() && let Some(hit) = map.get(&key) {
8881 return Ok(hit.clone());
8882 }
8883 let out = apply(ctx)?;
8884 if let Ok(mut map) = cache.lock() {
8885 map.insert(key, out.clone());
8886 }
8887 Ok(out)
8888}
8889
8890fn at_level(
8896 u: &Verb,
8897 level: i64,
8898 spread: bool,
8899 y: &Array,
8900 ctx: &mut Ctx<'_>,
8901 span: Span,
8902) -> Result<Array> {
8903 let n = if level < 0 { (boxing_level(y) + level).max(0) } else { level };
8905 if !spread {
8906 return map_level(u, n, y, ctx, span);
8907 }
8908 let mut cells = Vec::new();
8909 collect_level(u, n, y, ctx, span, &mut cells)?;
8910 let count = cells.len();
8911 assemble(&[count], cells, span)
8912}
8913
8914fn map_level(u: &Verb, n: i64, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
8915 let Some(boxes) = y.as_boxes().filter(|_| boxing_level(y) > n) else {
8916 return u.monad(y, ctx, span);
8917 };
8918 let boxes = boxes.to_vec();
8919 let mut cells = Vec::with_capacity(boxes.len());
8920 for b in &boxes {
8921 cells.push(map_level(u, n, b, ctx, span)?);
8922 }
8923 Ok(Array::new(y.shape.clone(), Data::Box(cells.into())))
8924}
8925
8926fn at_level_dyad(
8931 u: &Verb,
8932 level: i64,
8933 spread: bool,
8934 x: &Array,
8935 y: &Array,
8936 ctx: &mut Ctx<'_>,
8937 span: Span,
8938) -> Result<Array> {
8939 let depth = |a: &Array| if level < 0 { (boxing_level(a) + level).max(0) } else { level };
8942 let (nx, ny) = (depth(x), depth(y));
8943 if !spread {
8944 return map_level_dyad(u, nx, ny, x, y, ctx, span);
8945 }
8946 let mut cells = Vec::new();
8947 collect_level_dyad(u, nx, ny, x, y, ctx, span, &mut cells)?;
8948 let count = cells.len();
8949 assemble(&[count], cells, span)
8950}
8951
8952struct LevelPairs {
8954 left: Vec<Array>,
8955 right: Vec<Array>,
8956 shape: Vec<usize>,
8957}
8958
8959fn level_pairs(
8962 nx: i64,
8963 ny: i64,
8964 x: &Array,
8965 y: &Array,
8966 span: Span,
8967) -> Result<Option<LevelPairs>> {
8968 let bx = x.as_boxes().filter(|_| boxing_level(x) > nx);
8969 let by = y.as_boxes().filter(|_| boxing_level(y) > ny);
8970 Ok(match (bx, by) {
8971 (None, None) => None,
8972 (Some(bx), None) => {
8973 let n = bx.len();
8974 Some(LevelPairs {
8975 left: bx.to_vec(),
8976 right: vec![y.clone(); n],
8977 shape: x.shape.clone(),
8978 })
8979 }
8980 (None, Some(by)) => {
8981 let n = by.len();
8982 Some(LevelPairs {
8983 left: vec![x.clone(); n],
8984 right: by.to_vec(),
8985 shape: y.shape.clone(),
8986 })
8987 }
8988 (Some(bx), Some(by)) => {
8989 if x.shape != y.shape {
8990 return Err(Error::new(
8991 ErrorKind::Length,
8992 format!(
8993 "the levels do not agree: left shape {}, right shape {}",
8994 show_shape(&x.shape),
8995 show_shape(&y.shape)
8996 ),
8997 Some(span),
8998 ));
8999 }
9000 Some(LevelPairs { left: bx.to_vec(), right: by.to_vec(), shape: x.shape.clone() })
9001 }
9002 })
9003}
9004
9005fn map_level_dyad(
9006 u: &Verb,
9007 nx: i64,
9008 ny: i64,
9009 x: &Array,
9010 y: &Array,
9011 ctx: &mut Ctx<'_>,
9012 span: Span,
9013) -> Result<Array> {
9014 let Some(step) = level_pairs(nx, ny, x, y, span)? else {
9015 return u.dyad(x, y, ctx, span);
9016 };
9017 let mut cells = Vec::with_capacity(step.left.len());
9018 for (a, b) in step.left.iter().zip(step.right.iter()) {
9019 cells.push(map_level_dyad(u, nx, ny, a, b, ctx, span)?);
9020 }
9021 Ok(Array::new(step.shape, Data::Box(cells.into())))
9022}
9023
9024#[allow(clippy::too_many_arguments)]
9025fn collect_level_dyad(
9026 u: &Verb,
9027 nx: i64,
9028 ny: i64,
9029 x: &Array,
9030 y: &Array,
9031 ctx: &mut Ctx<'_>,
9032 span: Span,
9033 out: &mut Vec<Array>,
9034) -> Result<()> {
9035 let Some(step) = level_pairs(nx, ny, x, y, span)? else {
9036 out.push(u.dyad(x, y, ctx, span)?);
9037 return Ok(());
9038 };
9039 for (a, b) in step.left.iter().zip(step.right.iter()) {
9040 collect_level_dyad(u, nx, ny, a, b, ctx, span, out)?;
9041 }
9042 Ok(())
9043}
9044
9045fn collect_level(
9046 u: &Verb,
9047 n: i64,
9048 y: &Array,
9049 ctx: &mut Ctx<'_>,
9050 span: Span,
9051 out: &mut Vec<Array>,
9052) -> Result<()> {
9053 let Some(boxes) = y.as_boxes().filter(|_| boxing_level(y) > n) else {
9054 out.push(u.monad(y, ctx, span)?);
9055 return Ok(());
9056 };
9057 let boxes = boxes.to_vec();
9058 for b in &boxes {
9059 collect_level(u, n, b, ctx, span, out)?;
9060 }
9061 Ok(())
9062}
9063
9064fn poly_coeffs(y: &Array, span: Span) -> Result<Vec<Cx>> {
9068 let c = y
9069 .data
9070 .cast(DType::Complex)
9071 .ok_or_else(|| Error::domain("a polynomial's coefficients are numbers", span))?;
9072 match c {
9073 Data::Complex(v) => Ok(v.as_slice().to_vec()),
9074 _ => Err(Error::internal("coefficients did not cast to complex")),
9075 }
9076}
9077
9078const HYPERGEOMETRIC_TERMS: usize = 1 << 16;
9082
9083fn cx_list(v: &[Cx]) -> String {
9085 v.iter()
9086 .map(|z| if z[1] == 0.0 { format!("{}", z[0]) } else { format!("{}j{}", z[0], z[1]) })
9087 .collect::<Vec<_>>()
9088 .join(" ")
9089}
9090
9091fn hypergeometric(num: &[Cx], den: &[Cx], y: &Array, span: Span) -> Result<Array> {
9099 let (num, den) = cancel_parameters(num, den);
9100 let at = poly_coeffs(y, span)?;
9101 let mut out = Vec::with_capacity(at.len());
9102 for z in &at {
9103 out.push(hypergeometric_at(&num, &den, *z, span)?);
9104 }
9105 let mut a = complex_or_real(out);
9106 a.shape = y.shape.clone();
9107 Ok(a)
9108}
9109
9110fn cancel_parameters(num: &[Cx], den: &[Cx]) -> (Vec<Cx>, Vec<Cx>) {
9113 let mut left: Vec<Cx> = Vec::with_capacity(num.len());
9114 let mut right: Vec<Cx> = den.to_vec();
9115 for a in num {
9116 match right.iter().position(|b| b == a) {
9117 Some(i) => {
9118 right.remove(i);
9119 }
9120 None => left.push(*a),
9121 }
9122 }
9123 (left, right)
9124}
9125
9126fn hypergeometric_at(num: &[Cx], den: &[Cx], z: Cx, span: Span) -> Result<Cx> {
9127 let real = |v: &[Cx]| v.iter().all(|c| c[1] == 0.0);
9131 if z[1] == 0.0 && real(num) && real(den) {
9132 let n: Vec<f64> = num.iter().map(|c| c[0]).collect();
9133 let d: Vec<f64> = den.iter().map(|c| c[0]).collect();
9134 return Ok([hypergeometric_real(&n, &d, z[0], span)?, 0.0]);
9135 }
9136 let mut sum = cx::ONE;
9137 let mut term = cx::ONE;
9138 for k in 0..HYPERGEOMETRIC_TERMS {
9139 let kk = [k as f64, 0.0];
9140 let mut ratio = z;
9141 for a in num {
9142 ratio = cx::mul(ratio, cx::add(*a, kk));
9143 }
9144 for b in den {
9145 ratio = cx::div(ratio, cx::add(*b, kk));
9146 }
9147 term = cx::div(cx::mul(term, ratio), [k as f64 + 1.0, 0.0]);
9148 if !term[0].is_finite() || !term[1].is_finite() {
9149 return Ok(term);
9152 }
9153 let before = sum;
9154 sum = cx::add(sum, term);
9155 if sum == before {
9157 return Ok(sum);
9158 }
9159 }
9160 Err(Error::domain(
9161 format!("the hypergeometric series did not converge within {HYPERGEOMETRIC_TERMS} terms"),
9162 span,
9163 ))
9164}
9165
9166fn hypergeometric_real(num: &[f64], den: &[f64], z: f64, span: Span) -> Result<f64> {
9167 let mut sum = 1.0f64;
9168 let mut term = 1.0f64;
9169 for k in 0..HYPERGEOMETRIC_TERMS {
9170 let kk = k as f64;
9171 let mut ratio = z;
9172 for a in num {
9173 ratio *= a + kk;
9174 }
9175 for b in den {
9176 ratio /= b + kk;
9177 }
9178 term = term * ratio / (kk + 1.0);
9179 if !term.is_finite() {
9180 return Ok(term);
9181 }
9182 let before = sum;
9183 sum += term;
9184 if sum == before {
9185 return Ok(sum);
9186 }
9187 }
9188 Err(Error::domain(
9189 format!("the hypergeometric series did not converge within {HYPERGEOMETRIC_TERMS} terms"),
9190 span,
9191 ))
9192}
9193
9194fn complex_or_real(values: Vec<Cx>) -> Array {
9196 if values.iter().all(|z| z[1] == 0.0) {
9197 return Array::from_f64(values.iter().map(|z| z[0]).collect());
9198 }
9199 Array::new(vec![values.len()], Data::Complex(values.into()))
9200}
9201
9202fn poly_eval(x: &Array, y: &Array, span: Span) -> Result<Array> {
9205 let at = poly_coeffs(y, span)?;
9206 let at = at.first().copied().unwrap_or(cx::ZERO);
9207 let value = match x.as_boxes() {
9208 Some(parts) => {
9209 if parts.len() != 2 {
9210 return Err(Error::domain(
9211 "the root form of a polynomial is `multiplier ; roots`",
9212 span,
9213 ));
9214 }
9215 let multiplier = poly_coeffs(&parts[0], span)?;
9216 let mut v = multiplier.first().copied().unwrap_or(cx::ONE);
9217 for r in poly_coeffs(&parts[1], span)? {
9218 v = cx::mul(v, cx::sub(at, r));
9219 }
9220 v
9221 }
9222 None => {
9223 let c = poly_coeffs(x, span)?;
9224 let mut v = cx::ZERO;
9225 for &k in c.iter().rev() {
9226 v = cx::add(cx::mul(v, at), k);
9227 }
9228 v
9229 }
9230 };
9231 Ok(scalar_complex_or_real(value))
9232}
9233
9234fn scalar_complex_or_real(z: Cx) -> Array {
9235 if z[1] == 0.0 {
9236 return Array::scalar_f64(z[0]);
9237 }
9238 Array::new(vec![], Data::Complex(vec![z].into()))
9239}
9240
9241fn poly_roots(y: &Array, span: Span) -> Result<Array> {
9245 if let Some(parts) = y.as_boxes() {
9246 if parts.len() != 2 {
9247 return Err(Error::domain(
9248 "the root form of a polynomial is `multiplier ; roots`",
9249 span,
9250 ));
9251 }
9252 let multiplier = poly_coeffs(&parts[0], span)?;
9253 let multiplier = multiplier.first().copied().unwrap_or(cx::ONE);
9254 let mut coeffs = vec![multiplier];
9256 for r in poly_coeffs(&parts[1], span)? {
9257 let mut next = vec![cx::ZERO; coeffs.len() + 1];
9258 for (k, &c) in coeffs.iter().enumerate() {
9259 next[k + 1] = cx::add(next[k + 1], c);
9260 next[k] = cx::sub(next[k], cx::mul(c, r));
9261 }
9262 coeffs = next;
9263 }
9264 return Ok(complex_or_real(coeffs));
9265 }
9266 let mut c = poly_coeffs(y, span)?;
9267 while c.len() > 1 && c[c.len() - 1] == cx::ZERO {
9268 c.pop();
9269 }
9270 if c.iter().all(|&k| k == cx::ZERO) {
9274 let pair = vec![Array::scalar_i64(0), Array::new(vec![0], Data::empty(DType::I64))];
9275 return Ok(Array::new(vec![2], Data::Box(pair.into())));
9276 }
9277 if c.len() < 2 {
9278 return Err(Error::domain("a polynomial's roots need a coefficient of x", span));
9279 }
9280 let lead = c[c.len() - 1];
9281 let monic: Vec<Cx> = c.iter().map(|&k| cx::div(k, lead)).collect();
9282 let roots = durand_kerner(&monic);
9283 let pair = vec![scalar_complex_or_real(lead), complex_or_real(roots)];
9284 Ok(Array::new(vec![2], Data::Box(pair.into())))
9285}
9286
9287fn durand_kerner(monic: &[Cx]) -> Vec<Cx> {
9294 let d = monic.len() - 1;
9295 let seed = [0.4, 0.9];
9296 let mut z: Vec<Cx> = Vec::with_capacity(d);
9297 let mut p = cx::ONE;
9298 for _ in 0..d {
9299 z.push(p);
9300 p = cx::mul(p, seed);
9301 }
9302 let value = |monic: &[Cx], at: Cx| {
9303 let mut v = cx::ZERO;
9304 for &k in monic.iter().rev() {
9305 v = cx::add(cx::mul(v, at), k);
9306 }
9307 v
9308 };
9309 for _ in 0..500 {
9310 let mut moved: f64 = 0.0;
9311 for i in 0..d {
9312 let mut denom = cx::ONE;
9313 for j in 0..d {
9314 if i != j {
9315 denom = cx::mul(denom, cx::sub(z[i], z[j]));
9316 }
9317 }
9318 if denom == cx::ZERO {
9319 continue;
9320 }
9321 let step = cx::div(value(monic, z[i]), denom);
9322 z[i] = cx::sub(z[i], step);
9323 moved = moved.max(step[0].hypot(step[1]));
9324 }
9325 if moved < 1e-15 {
9326 break;
9327 }
9328 }
9329 for r in &mut z {
9331 if r[1].abs() < 1e-9 {
9332 r[1] = 0.0;
9333 }
9334 if r[0].abs() < 1e-12 {
9335 r[0] = 0.0;
9336 }
9337 }
9338 z.sort_by(|a, b| {
9342 let close = (a[0] - b[0]).abs() <= 1e-9 * (a[0].abs().max(b[0].abs()) + 1.0);
9343 let by_re = if close {
9344 std::cmp::Ordering::Equal
9345 } else {
9346 b[0].partial_cmp(&a[0]).unwrap_or(std::cmp::Ordering::Equal)
9347 };
9348 by_re.then(b[1].partial_cmp(&a[1]).unwrap_or(std::cmp::Ordering::Equal))
9349 });
9350 z
9351}
9352
9353fn poly_deriv(y: &Array, span: Span) -> Result<Array> {
9356 let c = poly_coeffs(y, span)?;
9357 if c.len() < 2 {
9358 return Ok(Array::from_i64(vec![0]));
9359 }
9360 let out: Vec<Cx> =
9361 c.iter().enumerate().skip(1).map(|(k, &v)| cx::mul(v, cx::from_real(k as f64))).collect();
9362 Ok(narrow_numbers(complex_or_real(out)))
9363}
9364
9365fn poly_integral(x: &Array, y: &Array, span: Span) -> Result<Array> {
9367 let c = poly_coeffs(y, span)?;
9368 let k = poly_coeffs(x, span)?;
9369 let mut out = vec![k.first().copied().unwrap_or(cx::ZERO)];
9370 for (i, &v) in c.iter().enumerate() {
9371 out.push(cx::div(v, cx::from_real((i + 1) as f64)));
9372 }
9373 Ok(narrow_numbers(complex_or_real(out)))
9374}
9375
9376fn narrow_numbers(a: Array) -> Array {
9380 let Data::F64(v) = &a.data else { return a };
9381 if v.iter().any(|x| !x.is_finite() || x.fract() != 0.0 || x.abs() > 9e15) {
9382 return a;
9383 }
9384 let values: Vec<i64> = v.iter().map(|&x| x as i64).collect();
9385 Array::new(a.shape, Data::I64(values.into()))
9386}
9387
9388fn characteristics(u: &Verb, y: &Array, span: Span) -> Result<Array> {
9392 let which = y.to_i64_vec().and_then(|v| v.first().copied());
9393 let chars = |s: String| Ok(Array::from_chars(s.chars().collect()));
9394 match which {
9395 Some(0) => {
9396 let ranks = u.ranks();
9397 Ok(Array::from_f64(
9398 ranks
9399 .iter()
9400 .map(|&r| if r == RANK_INF { f64::INFINITY } else { r as f64 })
9401 .collect(),
9402 ))
9403 }
9404 Some(-1) => match obverse(u) {
9408 Some(v) => chars(v.name()),
9409 None => Err(Error::not_yet(
9410 format!("the obverse of {} (no inverse is known)", u.name()),
9411 span,
9412 )),
9413 },
9414 Some(1) => match reduce_identity(u, 1).as_ref().map(identity_spelling) {
9415 Some(s) => chars(s),
9416 None => Err(Error::not_yet(
9417 format!("the identity function of {} (u b. 1)", u.name()),
9418 span,
9419 )),
9420 },
9421 _ => Err(Error::not_yet("a verb characteristic other than 0, 1 and _1", span)),
9422 }
9423}
9424
9425fn identity_spelling(d: &Data) -> String {
9428 let one = Array::new(Vec::new(), d.slice(0, 1));
9429 let text = crate::fmt::format_array(&one, &crate::fmt::FmtOpts::J);
9430 format!("{} $~ }}.@$", text.trim())
9431}
9432
9433fn with_operands<R>(
9436 alpha: &Verb,
9437 omega: Option<&Verb>,
9438 ctx: &mut Ctx<'_>,
9439 f: impl FnOnce(&mut Ctx<'_>) -> Result<R>,
9440) -> Result<R> {
9441 let saved = (ctx.env.verb("⍺⍺").cloned(), ctx.env.verb("⍵⍵").cloned());
9442 ctx.env.define("⍺⍺".to_string(), alpha.clone());
9443 if let Some(g) = omega {
9444 ctx.env.define("⍵⍵".to_string(), g.clone());
9445 }
9446 let out = f(ctx);
9447 match saved.0 {
9448 Some(v) => ctx.env.define("⍺⍺".to_string(), v),
9449 None => ctx.env.undefine("⍺⍺"),
9450 }
9451 match saved.1 {
9452 Some(v) => ctx.env.define("⍵⍵".to_string(), v),
9453 None => ctx.env.undefine("⍵⍵"),
9454 }
9455 out
9456}
9457
9458fn is_mixed_simple(a: &Array) -> bool {
9462 let Some(items) = a.as_boxes() else { return false };
9463 if items.is_empty() || items.iter().any(|b| b.rank() != 0 || b.dtype() == DType::Box) {
9464 return false;
9465 }
9466 let mut common = Some(items[0].dtype());
9467 for b in &items[1..] {
9468 common = common.and_then(|t| DType::promote(t, b.dtype()));
9469 }
9470 common.is_none()
9471}
9472
9473fn nest(y: &Array) -> Array {
9476 if y.dtype() == DType::Box || y.rank() == 0 {
9477 return y.clone();
9478 }
9479 Array::boxed(y.clone())
9480}
9481
9482fn key_pairs(
9486 u: &Verb,
9487 keys: &Array,
9488 values: Option<&Array>,
9489 ctx: &mut Ctx<'_>,
9490 span: Span,
9491) -> Result<Array> {
9492 let base = if keys.rank() == 0 { Array::new(vec![1], keys.data.clone()) } else { keys.clone() };
9493 let n = base.items();
9494 if let Some(v) = values && v.items() != n {
9495 return Err(Error::new(
9496 ErrorKind::Length,
9497 format!("{n} key(s) for {} item(s)", v.items()),
9498 Some(span),
9499 ));
9500 }
9501 let groups = group_positions(&base, ctx.cfg.tol);
9502 let origin = ctx.cfg.rules.origin;
9503 let mut cells = Vec::with_capacity(groups.len());
9504 for (first, at) in &groups {
9505 let key = item_or_self(&base, *first);
9506 let group = match values {
9507 Some(v) => select_items(v, at),
9508 None => Array::from_i64(at.iter().map(|&i| origin + i as i64).collect()),
9509 };
9510 let monadic = matches!(u, Verb::Explicit(d) if d.left.is_none());
9513 cells.push(if monadic {
9514 u.monad(&group, ctx, span)?
9515 } else {
9516 u.dyad(&key, &group, ctx, span)?
9517 });
9518 }
9519 let count = cells.len();
9520 assemble(&[count], cells, span)
9521}
9522
9523fn group_positions(y: &Array, tol: Tol) -> Vec<(usize, Vec<usize>)> {
9526 let n = y.items();
9527 let mut keys: Vec<Array> = Vec::new();
9528 let mut groups: Vec<(usize, Vec<usize>)> = Vec::new();
9529 for i in 0..n {
9530 let item = y.item(i);
9531 match keys.iter().position(|k| arrays_match(k, &item, tol)) {
9532 Some(at) => groups[at].1.push(i),
9533 None => {
9534 keys.push(item);
9535 groups.push((i, vec![i]));
9536 }
9537 }
9538 }
9539 groups
9540}
9541
9542fn format_spec(x: &Array, y: &Array, fmt: &FmtOpts, span: Span) -> Result<Array> {
9548 let spec = x
9549 .to_i64_vec()
9550 .ok_or_else(|| Error::domain("a format specification is whole numbers", span))?;
9551 if y.dtype() == DType::Box {
9552 return Err(Error::not_yet("format by specification of a nested array", span));
9553 }
9554 let cols = if y.rank() == 0 { 1 } else { y.shape[y.rank() - 1] };
9555 let rows = y.count() / cols.max(1);
9556 let pairs: Vec<(Option<i64>, i64)> = match spec.len() {
9558 1 => vec![(None, spec[0]); cols],
9559 2 => vec![(Some(spec[0]), spec[1]); cols],
9560 n if n == 2 * cols => spec.chunks(2).map(|c| (Some(c[0]), c[1])).collect(),
9561 n => {
9562 return Err(Error::new(
9563 ErrorKind::Length,
9564 format!("{n} specification value(s) for {cols} column(s)"),
9565 Some(span),
9566 ));
9567 }
9568 };
9569 if pairs.iter().any(|&(w, p)| w.is_some_and(|w| w < 0) || p < 0) {
9570 return Err(Error::domain("a format width and precision are nonnegative", span));
9571 }
9572 let numbers = y.to_f64_vec();
9573 let text = |i: usize, p: i64| -> String {
9574 match (&y.data, &numbers) {
9575 (Data::Char(v), _) => v[i].to_string(),
9576 (_, Some(v)) => {
9577 let s = format!("{:.*}", p as usize, v[i]);
9578 if v[i] < 0.0 { format!("{}{}", fmt.neg, &s[1..]) } else { s }
9579 }
9580 _ => String::new(),
9581 }
9582 };
9583 if y.dtype() != DType::Char && numbers.is_none() {
9584 return Err(Error::domain("format by specification takes numbers or characters", span));
9585 }
9586 let widths: Vec<usize> = pairs
9588 .iter()
9589 .enumerate()
9590 .map(|(c, &(w, p))| match w {
9591 Some(w) => w as usize,
9592 None => {
9593 (0..rows).map(|r| text(r * cols + c, p).chars().count()).max().unwrap_or(0) + 1
9594 }
9595 })
9596 .collect();
9597 let line: usize = widths.iter().sum();
9598 let mut out: Vec<char> = Vec::with_capacity(rows * line);
9599 for r in 0..rows {
9600 for c in 0..cols {
9601 let s = text(r * cols + c, pairs[c].1);
9602 let len = s.chars().count();
9603 if len > widths[c] {
9604 return Err(Error::domain(
9605 format!("{s} does not fit a field {} wide", widths[c]),
9606 span,
9607 ));
9608 }
9609 out.extend(std::iter::repeat_n(' ', widths[c] - len));
9610 out.extend(s.chars());
9611 }
9612 }
9613 let mut shape = if y.rank() == 0 { Vec::new() } else { y.shape[..y.rank() - 1].to_vec() };
9614 shape.push(line);
9615 Ok(Array::new(shape, Data::Char(out.into())))
9616}
9617
9618fn iota_apl(y: &Array, origin: i64, span: Span) -> Result<Array> {
9622 if y.rank() > 1 {
9623 return Err(Error::new(
9624 ErrorKind::Rank,
9625 "the index generator takes a shape, which is a scalar or a vector",
9626 Some(span),
9627 ));
9628 }
9629 let dims = y
9630 .to_i64_vec()
9631 .ok_or_else(|| Error::domain("index generator needs an integer argument", span))?;
9632 if dims.iter().any(|&n| n < 0) {
9633 return Err(Error::domain("index generator needs nonnegative lengths", span));
9634 }
9635 if dims.len() <= 1 {
9636 let n = dims.first().copied().unwrap_or(0);
9637 crate::limits::count(n as u128, span)?;
9638 return Ok(Array::from_i64((0..n).map(|i| origin + i).collect()));
9639 }
9640 let shape: Vec<usize> = dims.iter().map(|&n| n as usize).collect();
9641 let total = crate::limits::elements(&shape, span)?;
9642 let mut cells = Vec::with_capacity(total);
9643 let mut coord = vec![0usize; shape.len()];
9644 for _ in 0..total {
9645 cells.push(Array::from_i64(coord.iter().map(|&c| origin + c as i64).collect()));
9646 odometer(&mut coord, &shape);
9647 }
9648 Ok(Array::new(shape, Data::Box(cells.into())))
9649}
9650
9651fn carry_exact(result: Array, y: &Array) -> Array {
9657 if !matches!(y.dtype(), DType::Ext | DType::Rat) {
9658 return result;
9659 }
9660 match result.data.cast(DType::Ext) {
9661 Some(data) => Array::new(result.shape, data),
9662 None => result,
9663 }
9664}
9665
9666fn carry_exact2(result: Array, x: &Array, y: &Array) -> Array {
9667 let widened = carry_exact(result, x);
9668 carry_exact(widened, y)
9669}
9670
9671fn truth_table(m: u8, x: &Array, y: &Array, span: Span) -> Result<Array> {
9674 let table = m & 15;
9675 let bit = |a: i64, b: i64| ((table >> (3 - (2 * a + b))) & 1) as i64;
9676 let xs = x
9677 .to_i64_vec()
9678 .ok_or_else(|| Error::domain("a boolean function takes integers", span))?;
9679 let ys = y
9680 .to_i64_vec()
9681 .ok_or_else(|| Error::domain("a boolean function takes integers", span))?;
9682 let (a, b) = (xs.first().copied().unwrap_or(0), ys.first().copied().unwrap_or(0));
9683 if m < 16 {
9684 if !(0..=1).contains(&a) || !(0..=1).contains(&b) {
9685 return Err(Error::domain(
9686 format!("{m} b. takes 0 and 1; {m} b. + 16 is the same function on every bit"),
9687 span,
9688 ));
9689 }
9690 return Ok(Array::scalar_bool(bit(a, b) != 0));
9691 }
9692 let mut out = 0i64;
9693 for k in 0..64 {
9694 if bit((a >> k) & 1, (b >> k) & 1) != 0 {
9695 out |= 1i64 << k;
9696 }
9697 }
9698 Ok(Array::scalar_i64(out))
9699}
9700
9701pub fn amend_at(
9706 base: &Array,
9707 slots: &[Option<Array>],
9708 value: &Array,
9709 origin: i64,
9710 span: Span,
9711) -> Result<Array> {
9712 if slots.len() != base.rank() {
9713 return Err(Error::new(
9714 ErrorKind::Rank,
9715 format!(
9716 "indexed assignment needs one index per axis: {} slot(s) for a rank-{} value",
9717 slots.len(),
9718 base.rank()
9719 ),
9720 Some(span),
9721 ));
9722 }
9723 if !base.is_row_major() || !value.is_row_major() {
9726 let (b, v) = (base.to_row_major(), value.to_row_major());
9727 return amend_at(&b, slots, &v, origin, span);
9728 }
9729 let mut axes: Vec<Vec<usize>> = Vec::with_capacity(slots.len());
9731 let mut selected: Vec<usize> = Vec::new();
9732 for (k, slot) in slots.iter().enumerate() {
9733 let len = base.shape[k];
9734 let Some(idx) = slot else {
9735 axes.push((0..len).collect());
9736 selected.push(len);
9737 continue;
9738 };
9739 let Some(values) = idx.to_i64_vec() else {
9740 return Err(Error::new(
9741 ErrorKind::Type,
9742 "an index must be numeric",
9743 Some(span),
9744 ));
9745 };
9746 let mut positions = Vec::with_capacity(values.len());
9747 for v in values {
9748 let p = v - origin;
9749 if p < 0 || p as usize >= len {
9750 return Err(Error::new(
9751 ErrorKind::Domain,
9752 format!("index {v} is outside axis {k}, which has {len} element(s)"),
9753 Some(span),
9754 ));
9755 }
9756 positions.push(p as usize);
9757 }
9758 if idx.rank() > 0 {
9760 selected.push(positions.len());
9761 }
9762 axes.push(positions);
9763 }
9764 let count: usize = axes.iter().map(Vec::len).product();
9765 if value.rank() != 0 && (value.shape != selected || value.count() != count) {
9766 return Err(Error::new(
9767 ErrorKind::Shape,
9768 format!(
9769 "indexed assignment needs a scalar or a {} value, not a {} one",
9770 show_shape(&selected),
9771 show_shape(&value.shape)
9772 ),
9773 Some(span),
9774 ));
9775 }
9776 let dtype = DType::promote(base.dtype(), value.dtype()).ok_or_else(|| {
9779 Error::new(
9780 ErrorKind::Type,
9781 format!(
9782 "cannot put a {} value into a {} array",
9783 value.dtype().name(),
9784 base.dtype().name()
9785 ),
9786 Some(span),
9787 )
9788 })?;
9789 let mut out = base.cast(dtype).ok_or_else(|| Error::internal("promotion failed"))?;
9790 let src = value.cast(dtype).ok_or_else(|| Error::internal("promotion failed"))?;
9791 let strides = row_major_strides(&base.shape);
9792 let mut coords = vec![0usize; axes.len()];
9793 for n in 0..count {
9794 let mut rest = n;
9795 for k in (0..axes.len()).rev() {
9796 let len = axes[k].len();
9797 coords[k] = axes[k][rest % len];
9798 rest /= len;
9799 }
9800 let at: usize = coords.iter().zip(&strides).map(|(c, s)| c * s).sum();
9801 let from = if src.rank() == 0 { 0 } else { n };
9802 put_element(&mut out.data, at, &src.data, from);
9803 }
9804 Ok(out)
9805}
9806
9807fn row_major_strides(shape: &[usize]) -> Vec<usize> {
9808 let mut strides = vec![1usize; shape.len()];
9809 for k in (0..shape.len().saturating_sub(1)).rev() {
9810 strides[k] = strides[k + 1] * shape[k + 1];
9811 }
9812 strides
9813}
9814
9815fn put_element(dst: &mut Data, at: usize, src: &Data, from: usize) {
9817 match (dst, src) {
9818 (Data::Bool(d), Data::Bool(s)) => d.to_mut()[at] = s.as_slice()[from],
9819 (Data::I64(d), Data::I64(s)) => d.to_mut()[at] = s.as_slice()[from],
9820 (Data::Ext(d), Data::Ext(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
9821 (Data::Rat(d), Data::Rat(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
9822 (Data::F64(d), Data::F64(s)) => d.to_mut()[at] = s.as_slice()[from],
9823 (Data::Char(d), Data::Char(s)) => d.to_mut()[at] = s.as_slice()[from],
9824 (Data::Box(d), Data::Box(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
9825 _ => debug_assert!(false, "amend across types"),
9827 }
9828}
9829
9830fn agenda_pick(
9833 vs: &[Verb],
9834 w: &Verb,
9835 x: Option<&Array>,
9836 y: &Array,
9837 ctx: &mut Ctx<'_>,
9838 span: Span,
9839) -> Result<Verb> {
9840 let chosen = match x {
9841 None => w.monad(y, ctx, span)?,
9842 Some(x) => w.dyad(x, y, ctx, span)?,
9843 };
9844 let at = chosen
9845 .to_i64_vec()
9846 .and_then(|v| v.first().copied())
9847 .ok_or_else(|| Error::domain("an agenda index must be an integer", span))?;
9848 pick_gerund(vs, at, span)
9849}
9850
9851pub(crate) fn pick_gerund(vs: &[Verb], at: i64, span: Span) -> Result<Verb> {
9854 usize::try_from(at)
9855 .ok()
9856 .and_then(|k| vs.get(k))
9857 .cloned()
9858 .ok_or_else(|| {
9859 Error::domain(
9860 format!("agenda {at} is out of range: the gerund has {} verbs", vs.len()),
9861 span,
9862 )
9863 })
9864}
9865
9866fn evoke(
9871 vs: &[Verb],
9872 form: i64,
9873 x: Option<&Array>,
9874 y: &Array,
9875 ctx: &mut Ctx<'_>,
9876 span: Span,
9877) -> Result<Array> {
9878 if vs.is_empty() {
9879 return Err(Error::domain("an evoked gerund is empty", span));
9880 }
9881 if form == 0 {
9882 let mut cells = Vec::with_capacity(vs.len());
9883 for v in vs {
9884 cells.push(match x {
9885 None => v.monad(y, ctx, span)?,
9886 Some(x) => v.dyad(x, y, ctx, span)?,
9887 });
9888 }
9889 return assemble(&[vs.len()], cells, span);
9890 }
9891 if x.is_some() {
9892 return Err(Error::domain("m`:3 has no dyadic meaning", span));
9893 }
9894 let items = if y.rank() == 0 { vec![y.clone()] } else { y.cells(1) };
9895 let Some((last, rest)) = items.split_last() else {
9896 return Err(Error::domain("m`:3 needs an argument with items", span));
9897 };
9898 let mut acc = last.clone();
9899 for (i, item) in rest.iter().enumerate().rev() {
9900 acc = vs[i % vs.len()].dyad(item, &acc, ctx, span)?;
9901 }
9902 Ok(acc)
9903}
9904
9905fn stencil(u: &Verb, w: &[i64], y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
9910 if w.len() > y.rank() {
9911 return Err(Error::new(
9912 ErrorKind::Rank,
9913 format!("a stencil of {} axis/axes into a rank-{} value", w.len(), y.rank()),
9914 Some(span),
9915 ));
9916 }
9917 if w.iter().any(|&n| n <= 0) {
9918 return Err(Error::domain("a stencil window is a positive size", span));
9919 }
9920 let y = y.to_row_major();
9921 let k = w.len();
9922 let st = strides(&y.shape);
9923 let frame: Vec<usize> = y.shape[..k].to_vec();
9924 let mut wshape: Vec<usize> = w.iter().map(|&n| n as usize).collect();
9926 wshape.extend_from_slice(&y.shape[k..]);
9927 let inner: usize = y.shape[k..].iter().product();
9928 let total: usize = frame.iter().product();
9929 let mut cells = Vec::with_capacity(total);
9930 let mut at = vec![0usize; frame.len()];
9931 let mut coord = vec![0usize; k];
9932 for _ in 0..total {
9933 let mut data = Data::empty(y.dtype());
9934 coord.iter_mut().for_each(|c| *c = 0);
9935 let count: usize = w.iter().map(|&n| n as usize).product();
9936 for _ in 0..count {
9937 let mut base = 0usize;
9938 let mut inside = true;
9939 for a in 0..k {
9940 let off = at[a] as i64 + coord[a] as i64 - (w[a] - 1) / 2;
9941 if off < 0 || off >= y.shape[a] as i64 {
9942 inside = false;
9943 break;
9944 }
9945 base += off as usize * st[a];
9946 }
9947 for j in 0..inner {
9948 if inside {
9949 push_elem(&mut data, &y.data, base + j);
9950 } else {
9951 data.push_fill();
9952 }
9953 }
9954 odometer(&mut coord, &wshape[..k]);
9955 }
9956 cells.push(u.monad(&Array::new(wshape.clone(), data), ctx, span)?);
9957 odometer(&mut at, &frame);
9958 }
9959 assemble(&frame, cells, span)
9960}
9961
9962fn outfix(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
9966 let k = one_int(x, "an outfix width", span)?;
9967 let n = y.items() as i64;
9968 let list = as_list(y);
9969 let starts: Vec<i64> = if k < 0 {
9974 let step = -k;
9975 (0..(n + step - 1) / step).map(|i| i * step).collect()
9976 } else {
9977 (0..=(n - k)).collect()
9978 };
9979 let width = k.unsigned_abs() as usize;
9980 let mut cells = Vec::with_capacity(starts.len());
9981 for start in starts {
9982 let start = start as usize;
9983 let keep: Vec<usize> =
9984 (0..n as usize).filter(|&i| i < start || i >= start + width).collect();
9985 cells.push(u.monad(&select_items(&list, &keep), ctx, span)?);
9986 }
9987 assemble(&[cells.len()], cells, span)
9988}
9989
9990pub(crate) fn obverse(v: &Verb) -> Option<Verb> {
10002 let swap = |name: &'static str| -> Option<Verb> {
10003 crate::frontend::j::verb_named(name)
10004 };
10005 Some(match v {
10006 Verb::Prim(p) => {
10007 use ScalarMonad as SM;
10008 if matches!(
10012 p.monad,
10013 MonadOp::Scalar(SM::Conj | SM::Neg | SM::Recip | SM::OneMinus)
10014 | MonadOp::Reverse
10015 | MonadOp::TransposeAxes
10016 ) {
10017 return Some(v.clone());
10018 }
10019 if matches!(p.monad, MonadOp::Scalar(SM::Imaginary)) {
10022 return Some(Verb::Atop(Box::new(swap("-")?), Box::new(swap("j.")?)));
10023 }
10024 let by_monad: Option<&'static str> = match p.monad {
10025 MonadOp::Scalar(SM::Exp) => Some("^."),
10026 MonadOp::Scalar(SM::Ln) => Some("^"),
10027 MonadOp::Scalar(SM::Sqrt) => Some("*:"),
10028 MonadOp::Scalar(SM::Square) => Some("%:"),
10029 MonadOp::Scalar(SM::Double) => Some("-:"),
10030 MonadOp::Scalar(SM::Halve) => Some("+:"),
10031 MonadOp::Scalar(SM::Inc) => Some("<:"),
10032 MonadOp::Scalar(SM::Dec) => Some(">:"),
10033 MonadOp::Enclose(_) => Some(">"),
10034 MonadOp::Open => Some("<"),
10035 MonadOp::DecodeBits => Some("#:"),
10036 MonadOp::EncodeBits => Some("#."),
10037 _ => None,
10038 };
10039 swap(by_monad?)?
10040 }
10041 Verb::WithObverse(_, w) => (**w).clone(),
10043 Verb::Atop(f, g) => {
10045 Verb::Atop(Box::new(obverse(g)?), Box::new(obverse(f)?))
10046 }
10047 Verb::Compose(f, g) | Verb::Beside(f, g) => {
10048 Verb::Atop(Box::new(obverse(g)?), Box::new(obverse(f)?))
10049 }
10050 Verb::Rank(f, r) => Verb::Rank(Box::new(obverse(f)?), *r),
10051 Verb::Fit(f, n) => Verb::Fit(Box::new(obverse(f)?), *n),
10052 Verb::PowerN(f, Power::Times(n)) => {
10054 Verb::PowerN(Box::new(obverse(f)?), Power::Times(*n))
10055 }
10056 Verb::BondLeft(m, f) => bond_obverse(m, f, true)?,
10057 Verb::BondRight(f, n) => bond_obverse(n, f, false)?,
10058 _ => return None,
10059 })
10060}
10061
10062fn bond_obverse(n: &Array, f: &Verb, left: bool) -> Option<Verb> {
10066 let Verb::Prim(p) = f else { return None };
10067 let named = |name: &'static str| crate::frontend::j::verb_named(name);
10068 let bond = |name: &'static str, arg: &Array| -> Option<Verb> {
10069 let g = named(name)?;
10070 Some(if left {
10071 Verb::BondLeft(arg.clone(), Box::new(g))
10072 } else {
10073 Verb::BondRight(Box::new(g), arg.clone())
10074 })
10075 };
10076 use ScalarDyad as SD;
10077 let DyadOp::Scalar(op) = p.dyad else { return None };
10078 match (op, left) {
10079 (SD::Sub | SD::DivJ | SD::DivApl, true) => bond(p.name, n),
10081 (SD::Add, _) => Some(Verb::BondRight(Box::new(named("-")?), n.clone())),
10085 (SD::Mul, _) => Some(Verb::BondRight(Box::new(named("%")?), n.clone())),
10086 (SD::Sub, false) => bond("+", n),
10087 (SD::DivJ | SD::DivApl, false) => bond("*", n),
10088 (SD::Pow, false) => Some(Verb::BondLeft(n.clone(), Box::new(named("%:")?))),
10090 (SD::Pow, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^.")?))),
10091 (SD::Log, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^")?))),
10092 (SD::Root, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^")?))),
10093 _ => None,
10094 }
10095}
10096
10097fn self_classify(y: &Array, tol: Tol) -> Array {
10102 let items = if y.rank() == 0 { 1 } else { y.items() };
10103 let keys = nub(&as_list(y), tol);
10104 let rows = keys.items();
10105 let mut out = Vec::with_capacity(rows * items);
10106 for i in 0..rows {
10107 let key = item_or_self(&keys, i);
10108 for j in 0..items {
10109 out.push(arrays_match(&key, &item_or_self(y, j), tol) as u8);
10110 }
10111 }
10112 Array::new(vec![rows, items], Data::Bool(out.into()))
10113}
10114
10115fn nub_sieve(y: &Array, tol: Tol) -> Array {
10117 let items = if y.rank() == 0 { 1 } else { y.items() };
10118 let mut seen: Vec<Array> = Vec::new();
10119 let mut out = Vec::with_capacity(items);
10120 for i in 0..items {
10121 let cell = item_or_self(y, i);
10122 let fresh = !seen.iter().any(|s| arrays_match(s, &cell, tol));
10123 if fresh {
10124 seen.push(cell);
10125 }
10126 out.push(fresh as u8);
10127 }
10128 Array::new(vec![items], Data::Bool(out.into()))
10129}
10130
10131fn as_list(y: &Array) -> Array {
10133 if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() }
10134}
10135
10136fn conforming_cells(y: &Array, item_rank: usize) -> Vec<Array> {
10141 let frame_rank = y.rank().saturating_sub(item_rank);
10142 let nf: usize = y.shape[..frame_rank].iter().product();
10143 (0..nf).map(|i| y.cell_at(frame_rank, i)).collect()
10144}
10145
10146fn item_marks(y: &Array, x: &Array, tol: Tol) -> Vec<bool> {
10148 let n = if y.rank() == 0 { 1 } else { y.items() };
10149 let item_rank = y.rank().saturating_sub(1);
10150 let against = conforming_cells(x, item_rank);
10151 (0..n)
10152 .map(|i| {
10153 let cell = item_or_self(y, i);
10154 against.iter().any(|c| arrays_match(&cell, c, tol))
10155 })
10156 .collect()
10157}
10158
10159fn set_less(x: &Array, y: &Array, tol: Tol) -> Array {
10161 let xs = as_list(x);
10162 let marks = item_marks(&xs, y, tol);
10163 let keep: Vec<usize> = (0..marks.len()).filter(|&i| !marks[i]).collect();
10164 select_items(&xs, &keep)
10165}
10166
10167fn set_rank(cfg: EvalCfg, what: &str, x: &Array, y: &Array, span: Span) -> Result<()> {
10171 if cfg.rules.lang == crate::Lang::Apl && (x.rank() > 1 || y.rank() > 1) {
10172 return Err(Error::new(
10173 ErrorKind::Rank,
10174 format!("{what} takes vectors, not rank {} and rank {}", x.rank(), y.rank()),
10175 Some(span),
10176 ));
10177 }
10178 Ok(())
10179}
10180
10181fn intersect_items(x: &Array, y: &Array, tol: Tol) -> Array {
10183 let xs = as_list(x);
10184 let marks = item_marks(&xs, y, tol);
10185 let keep: Vec<usize> = (0..marks.len()).filter(|&i| marks[i]).collect();
10186 select_items(&xs, &keep)
10187}
10188
10189fn union_items(x: &Array, y: &Array, tol: Tol, span: Span) -> Result<Array> {
10192 let xs = as_list(x);
10193 let ys = as_list(y);
10194 let marks = item_marks(&ys, &xs, tol);
10195 let mut extra: Vec<usize> = Vec::new();
10196 for (i, &seen) in marks.iter().enumerate() {
10197 if seen {
10198 continue;
10199 }
10200 let cell = item_or_self(&ys, i);
10201 if !extra.iter().any(|&j| arrays_match(&item_or_self(&ys, j), &cell, tol)) {
10202 extra.push(i);
10203 }
10204 }
10205 catenate(&xs, &select_items(&ys, &extra), true, false, span)
10206}
10207
10208fn find_seq(x: &Array, y: &Array, tol: Tol, apl: bool, span: Span) -> Result<Array> {
10218 let (xr, yr) = (x.rank(), y.rank());
10219 if apl && xr > yr {
10220 return Ok(Array::new(y.shape.clone(), Data::Bool(vec![0u8; y.count()].into())));
10222 }
10223 if !apl && xr.max(1) != yr {
10224 return Err(Error::new(
10225 ErrorKind::Rank,
10226 format!("a rank-{xr} pattern in a rank-{yr} argument"),
10227 Some(span),
10228 ));
10229 }
10230 let mut pattern = vec![1usize; yr];
10231 pattern[yr - xr..].copy_from_slice(&x.shape);
10232 let n = y.count();
10233 let mut out = vec![0u8; n];
10234 let (xrm, yrm) = (x.to_row_major(), y.to_row_major());
10235 let yst = strides(&y.shape);
10236 let cells: usize = pattern.iter().product();
10237 let mut at = vec![0usize; yr];
10238 for slot in out.iter_mut() {
10239 if (0..yr).all(|a| at[a] + pattern[a] <= y.shape[a]) {
10240 let mut off = vec![0usize; yr];
10241 let mut hit = true;
10242 for k in 0..cells {
10243 let i: usize = (0..yr).map(|a| (at[a] + off[a]) * yst[a]).sum();
10244 if !arrays_match(&atom(&xrm, k), &atom(&yrm, i), tol) {
10245 hit = false;
10246 break;
10247 }
10248 odometer(&mut off, &pattern);
10249 }
10250 *slot = hit as u8;
10251 }
10252 odometer(&mut at, &y.shape);
10253 }
10254 Ok(Array::new(y.shape.clone(), Data::Bool(out.into())))
10255}
10256
10257fn bool_dyad(op: BoolDyad, x: &Array, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
10261 let bit = |a: &Array| -> Result<u8> {
10262 match a.to_i64_vec().as_deref() {
10263 Some([0]) => Ok(0),
10264 Some([1]) => Ok(1),
10265 _ => Err(Error::domain("this verb reads values of 0 or 1", span)),
10266 }
10267 };
10268 let _ = cfg;
10269 let (a, b) = (bit(x)?, bit(y)?);
10270 let v = match op {
10271 BoolDyad::Nor => u8::from(a == 0 && b == 0),
10272 BoolDyad::Nand => u8::from(a == 0 || b == 0),
10273 };
10274 Ok(Array::new(vec![], Data::Bool(vec![v].into())))
10275}
10276
10277fn item_ranks(y: &Array, rules: Rules, span: Span) -> Result<Vec<usize>> {
10283 check_gradable(y, rules, span)?;
10284 if !y.dtype().is_numeric() {
10285 return Err(Error::domain("an anagram index needs numbers", span));
10286 }
10287 let order = grade_order(&as_list(y), false, Tao::of(rules));
10288 let mut ranks = vec![0usize; order.len()];
10289 for (place, &i) in order.iter().enumerate() {
10290 ranks[i] = place;
10291 }
10292 Ok(ranks)
10293}
10294
10295fn anagram_index(y: &Array, rules: Rules, span: Span) -> Result<Array> {
10298 let ranks = item_ranks(y, rules, span)?;
10299 let n = ranks.len();
10300 let mut index: i128 = 0;
10301 for i in 0..n {
10302 let smaller = ranks[i + 1..].iter().filter(|&&r| r < ranks[i]).count() as i128;
10303 index = index
10304 .checked_mul((n - i) as i128)
10305 .and_then(|v| v.checked_add(smaller))
10306 .ok_or_else(|| Error::not_yet("an anagram index too large for an integer", span))?;
10307 }
10308 i64::try_from(index)
10309 .map(Array::scalar_i64)
10310 .map_err(|_| Error::not_yet("an anagram index too large for an integer", span))
10311}
10312
10313fn anagram_from(x: &Array, y: &Array, span: Span) -> Result<Array> {
10316 let idx = x
10317 .to_i64_vec()
10318 .ok_or_else(|| Error::domain("an anagram index must be an integer", span))?;
10319 let Some(&want) = idx.first() else {
10320 return Err(Error::internal("anagram with no index"));
10321 };
10322 let ys = as_list(y);
10323 let n = ys.items();
10324 let mut total: i128 = 1;
10325 for k in 1..=n as i128 {
10326 total = total
10327 .checked_mul(k)
10328 .ok_or_else(|| Error::not_yet("permuting more items than an integer counts", span))?;
10329 }
10330 let mut at = want as i128;
10331 if at < 0 {
10332 at += total;
10333 }
10334 if at < 0 || at >= total {
10335 return Err(Error::domain(
10336 format!("permutation {want} is out of range: {n} items have {total} of them"),
10337 span,
10338 ));
10339 }
10340 let mut pool: Vec<usize> = (0..n).collect();
10343 let mut order = Vec::with_capacity(n);
10344 let mut fact = total;
10345 for i in 0..n {
10346 fact /= (n - i) as i128;
10347 let d = (at / fact) as usize;
10348 at %= fact;
10349 order.push(pool.remove(d));
10350 }
10351 Ok(select_items(&ys, &order))
10352}
10353
10354fn cycle_form(y: &Array, span: Span) -> Result<Array> {
10358 if y.dtype() == DType::Box {
10359 let perm = cycles_to_direct(y, span)?;
10360 return Ok(Array::from_i64(perm.iter().map(|&i| i as i64).collect()));
10361 }
10362 let perm = direct_permutation(y, span)?;
10363 let mut boxes: Vec<Array> = Vec::new();
10364 let mut done = vec![false; perm.len()];
10365 for start in 0..perm.len() {
10366 if done[start] {
10367 continue;
10368 }
10369 let mut cycle = Vec::new();
10370 let mut at = start;
10371 while !done[at] {
10372 done[at] = true;
10373 cycle.push(at);
10374 at = perm[at];
10375 }
10376 let top = cycle.iter().position(|&v| v == *cycle.iter().max().unwrap()).unwrap();
10379 cycle.rotate_left(top);
10380 boxes.push(Array::boxed(Array::from_i64(
10381 cycle.iter().map(|&i| i as i64).collect(),
10382 )));
10383 }
10384 boxes.sort_by_key(|b| b.as_boxes().map(|s| s[0].to_i64_vec().unwrap()[0]).unwrap_or(0));
10385 let n = boxes.len();
10386 let inner: Vec<Array> =
10387 boxes.into_iter().map(|b| b.as_boxes().unwrap()[0].clone()).collect();
10388 Ok(Array::new(vec![n], Data::Box(inner.into())))
10389}
10390
10391fn direct_permutation(y: &Array, span: Span) -> Result<Vec<usize>> {
10393 let v = y
10394 .to_i64_vec()
10395 .ok_or_else(|| Error::domain("a permutation is a list of integers", span))?;
10396 let n = v.len();
10397 let mut seen = vec![false; n];
10398 let mut out = Vec::with_capacity(n);
10399 for &i in &v {
10400 let k = usize::try_from(i).ok().filter(|&k| k < n && !seen[k]).ok_or_else(|| {
10401 Error::domain(format!("{i} does not belong to a permutation of {n} items"), span)
10402 })?;
10403 seen[k] = true;
10404 out.push(k);
10405 }
10406 Ok(out)
10407}
10408
10409fn cycles_to_direct(y: &Array, span: Span) -> Result<Vec<usize>> {
10413 let boxes = y.as_boxes().ok_or_else(|| Error::internal("cycles from a simple array"))?;
10414 let mut cycles: Vec<Vec<usize>> = Vec::new();
10415 let mut top = 0usize;
10416 for b in boxes {
10417 let v = b
10418 .to_i64_vec()
10419 .ok_or_else(|| Error::domain("a cycle is a list of integers", span))?;
10420 let mut cycle = Vec::with_capacity(v.len());
10421 for &i in &v {
10422 let k = usize::try_from(i)
10423 .map_err(|_| Error::domain(format!("{i} is not an index"), span))?;
10424 top = top.max(k + 1);
10425 cycle.push(k);
10426 }
10427 cycles.push(cycle);
10428 }
10429 let mut perm: Vec<usize> = (0..top).collect();
10430 for cycle in &cycles {
10431 for w in 0..cycle.len() {
10432 perm[cycle[w]] = cycle[(w + 1) % cycle.len()];
10434 }
10435 }
10436 Ok(perm)
10437}
10438
10439fn permute(x: &Array, y: &Array, span: Span) -> Result<Array> {
10444 let ys = as_list(y);
10445 let n = ys.items();
10446 let cyclic = x.dtype() == DType::Box;
10447 if !cyclic && x.rank() == 0 {
10448 return Err(Error::not_yet("permuting by a single atom (x C. y)", span));
10449 }
10450 let mut perm =
10451 if cyclic { cycles_to_direct(x, span)? } else { direct_permutation(&as_list(x), span)? };
10452 if perm.len() > n {
10453 return Err(Error::new(
10454 ErrorKind::Length,
10455 format!("a permutation of {} items applied to {n}", perm.len()),
10456 Some(span),
10457 ));
10458 }
10459 if perm.len() < n {
10460 if cyclic {
10461 perm.extend(perm.len()..n);
10463 } else {
10464 let head: Vec<usize> = (perm.len()..n).collect();
10467 perm.splice(0..0, head);
10468 }
10469 }
10470 Ok(select_items(&ys, &perm))
10471}
10472
10473fn unicode(y: &Array, pass_chars: bool, span: Span) -> Result<Array> {
10479 if y.dtype() == DType::Char {
10480 if pass_chars {
10481 return Ok(y.clone());
10482 }
10483 return Ok(chars_to_codes(y));
10484 }
10485 codes_to_chars(y, span)
10486}
10487
10488fn chars_to_codes(y: &Array) -> Array {
10489 let Data::Char(v) = &y.data else { return y.clone() };
10490 Array::new(y.shape.clone(), Data::I64(v.iter().map(|&c| c as i64).collect()))
10491}
10492
10493fn codes_to_chars(y: &Array, span: Span) -> Result<Array> {
10494 let v = y
10495 .to_i64_vec()
10496 .ok_or_else(|| Error::domain("a codepoint must be an integer", span))?;
10497 let mut out = Vec::with_capacity(v.len());
10498 for &c in &v {
10499 let ch = u32::try_from(c).ok().and_then(char::from_u32).ok_or_else(|| {
10500 Error::domain(format!("{c} is not a Unicode codepoint"), span)
10501 })?;
10502 out.push(ch);
10503 }
10504 Ok(Array::new(y.shape.clone(), Data::Char(out.into())))
10505}
10506
10507fn unicode_form(x: &Array, y: &Array, span: Span) -> Result<Array> {
10510 let form = x
10511 .to_i64_vec()
10512 .ok_or_else(|| Error::domain("a conversion form is an integer", span))?
10513 .first()
10514 .copied()
10515 .unwrap_or(0);
10516 match form {
10517 3 if y.dtype() == DType::Char => Ok(chars_to_codes(y)),
10518 3 => Err(Error::domain("form 3 converts characters to codepoints", span)),
10519 10 => codes_to_chars(y, span),
10520 n => Err(Error::not_yet(format!("the byte-oriented unicode form ({n} u:)"), span)),
10521 }
10522}
10523
10524fn boxing_level(y: &Array) -> i64 {
10526 match y.as_boxes() {
10527 None => 0,
10528 Some(bs) => 1 + bs.iter().map(boxing_level).max().unwrap_or(0),
10529 }
10530}
10531
10532fn split_items(y: &Array) -> Array {
10536 if y.rank() == 0 {
10537 return Array::boxed(y.clone());
10538 }
10539 let last = y.shape[y.rank() - 1];
10540 let outer: Vec<usize> = y.shape[..y.rank() - 1].to_vec();
10541 let n: usize = outer.iter().product();
10542 let mut boxes = Vec::with_capacity(n);
10543 for i in 0..n {
10544 let mut data = Data::empty(y.dtype());
10545 for k in 0..last {
10546 push_elem(&mut data, &y.data, i * last + k);
10547 }
10548 boxes.push(Array::new(vec![last], data));
10549 }
10550 Array::new(outer, Data::Box(boxes.into()))
10551}
10552
10553fn pick(x: &Array, y: &Array, origin: i64, span: Span) -> Result<Array> {
10556 let xs = as_list(x);
10557 let mut cur = y.clone();
10558 for i in 0..xs.items() {
10559 let step = open_cell(&item_or_self(&xs, i));
10560 let idx = step
10561 .to_i64_vec()
10562 .ok_or_else(|| Error::domain("a pick path holds integers", span))?;
10563 let base =
10564 if cur.rank() == 0 { Array::new(vec![1], cur.data.clone()) } else { cur.clone() };
10565 if idx.len() > base.rank() {
10566 return Err(Error::new(
10567 ErrorKind::Length,
10568 format!(
10569 "a path step of {} index(es) into a value of rank {}",
10570 idx.len(),
10571 cur.rank()
10572 ),
10573 Some(span),
10574 ));
10575 }
10576 let zeroed: Vec<i64> = idx.iter().map(|&v| v - origin).collect();
10577 let at = cell_index(&base, &zeroed, span)?;
10578 cur = open_cell(&base.cell_at(idx.len(), at));
10579 }
10580 Ok(cur)
10581}
10582
10583fn prime_meta(x: &Array, y: &Array, span: Span) -> Result<Array> {
10588 let form = one_int(x, "a prime query", span)?;
10589 let n = one_int(y, "a prime query", span)?;
10590 match form {
10591 -1 => Ok(Array::scalar_i64(primes_below(n, span)?)),
10593 0 => Ok(Array::scalar_bool(!is_prime(n))),
10595 1 => Ok(Array::scalar_bool(is_prime(n))),
10596 2 | 3 => {
10598 let (ps, es) = factor_table(n, span)?;
10599 let k = ps.len();
10600 if form == 3 {
10601 return Ok(Array::from_i64(ps));
10602 }
10603 let mut all = ps;
10604 all.extend(es);
10605 Ok(Array::new(vec![2, k], Data::I64(all.into())))
10606 }
10607 4 => Ok(Array::scalar_i64(next_prime(n, span)?)),
10609 -4 => Ok(Array::scalar_i64(previous_prime(n, span)?)),
10610 other => Err(Error::domain(format!("{other} is not a prime query"), span)),
10611 }
10612}
10613
10614fn prime_exponents(x: &Array, y: &Array, span: Span) -> Result<Array> {
10617 let n = one_int(y, "prime exponents", span)?;
10618 let count = x.to_f64_vec().and_then(|v| v.first().copied()).unwrap_or(0.0);
10619 let (ps, es) = factor_table(n, span)?;
10620 if count == f64::NEG_INFINITY {
10621 let k = ps.len();
10622 let mut all = ps;
10623 all.extend(es);
10624 return Ok(Array::new(vec![2, k], Data::I64(all.into())));
10625 }
10626 let want = one_int(x, "prime exponents", span)?;
10627 if want < 0 {
10628 return Err(Error::not_yet(format!("the prime exponent form ({want} q:)"), span));
10629 }
10630 let mut out = Vec::with_capacity(want as usize);
10631 for i in 0..want {
10632 let p = nth_prime(i, span)?;
10633 out.push(ps.iter().position(|&q| q == p).map_or(0, |at| es[at]));
10634 }
10635 Ok(Array::from_i64(out))
10636}
10637
10638fn factor_table(n: i64, span: Span) -> Result<(Vec<i64>, Vec<i64>)> {
10640 let factors = prime_factors(n, span)?;
10641 let mut ps: Vec<i64> = Vec::new();
10642 let mut es: Vec<i64> = Vec::new();
10643 for f in factors {
10644 if ps.last() == Some(&f) {
10645 *es.last_mut().unwrap() += 1;
10646 } else {
10647 ps.push(f);
10648 es.push(1);
10649 }
10650 }
10651 Ok((ps, es))
10652}
10653
10654fn is_prime(n: i64) -> bool {
10655 if n < 2 {
10656 return false;
10657 }
10658 let mut d = 2i64;
10659 while d.saturating_mul(d) <= n {
10660 if n % d == 0 {
10661 return false;
10662 }
10663 d += 1;
10664 }
10665 true
10666}
10667
10668fn primes_below(n: i64, span: Span) -> Result<i64> {
10669 if n < 0 {
10670 return Err(Error::domain("counting the primes below a negative number", span));
10671 }
10672 Ok((2..n).filter(|&k| is_prime(k)).count() as i64)
10673}
10674
10675fn next_prime(n: i64, span: Span) -> Result<i64> {
10676 let mut k = n.checked_add(1).ok_or_else(|| Error::domain("no next prime", span))?;
10677 while !is_prime(k) {
10678 k = k.checked_add(1).ok_or_else(|| Error::domain("no next prime", span))?;
10679 }
10680 Ok(k)
10681}
10682
10683fn previous_prime(n: i64, span: Span) -> Result<i64> {
10684 let mut k = n - 1;
10685 while k >= 2 {
10686 if is_prime(k) {
10687 return Ok(k);
10688 }
10689 k -= 1;
10690 }
10691 Err(Error::domain(format!("there is no prime below {n}"), span))
10692}
10693
10694fn one_int(a: &Array, what: &str, span: Span) -> Result<i64> {
10696 a.to_i64_vec()
10697 .and_then(|v| v.first().copied())
10698 .ok_or_else(|| Error::domain(format!("{what} needs an integer"), span))
10699}
10700
10701fn expand(x: &Array, y: &Array, span: Span) -> Result<Array> {
10704 let mask = x
10705 .to_i64_vec()
10706 .ok_or_else(|| Error::domain("an expansion mask holds 0s and 1s", span))?;
10707 if mask.iter().any(|&b| b != 0 && b != 1) {
10708 return Err(Error::domain("an expansion mask holds 0s and 1s", span));
10709 }
10710 let ys = as_list(y);
10711 let taken = mask.iter().filter(|&&b| b == 1).count();
10712 let n = ys.items();
10713 let spread = n == 1 && taken != 1;
10715 if !spread && taken != n {
10716 return Err(Error::new(
10717 ErrorKind::Length,
10718 format!("an expansion mask taking {taken} item(s) over {n}"),
10719 Some(span),
10720 ));
10721 }
10722 let m = ys.item_size();
10723 let mut data = Data::empty(ys.dtype());
10724 let mut at = 0usize;
10725 for &b in &mask {
10726 if b == 1 {
10727 let from = if spread { 0 } else { at };
10728 for k in 0..m {
10729 push_elem(&mut data, &ys.data, from * m + k);
10730 }
10731 at += 1;
10732 } else {
10733 for _ in 0..m {
10734 data.push_fill();
10735 }
10736 }
10737 }
10738 let mut shape = ys.shape.clone();
10739 if shape.is_empty() {
10740 shape.push(mask.len());
10741 } else {
10742 shape[0] = mask.len();
10743 }
10744 Ok(Array::new(shape, data))
10745}
10746
10747fn execute(y: &Array, apl: bool, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
10755 let Data::Char(v) = &y.data else {
10756 return Err(Error::domain("execute reads a character list", span));
10757 };
10758 let src: String = v.iter().collect();
10759 execute_source(&src, apl, ctx, span)
10760}
10761
10762pub(crate) fn execute_source(
10765 src: &str,
10766 apl: bool,
10767 ctx: &mut Ctx<'_>,
10768 span: Span,
10769) -> Result<Array> {
10770 let lang = if apl { crate::Lang::Apl } else { crate::Lang::J };
10771 let dialect = ctx.cfg.rules.dialect();
10774 let nested = crate::compile(lang, src, &dialect).map_err(|e| nested_error(e, src, span))?;
10775 if !nested.params.is_empty() {
10776 return Err(Error::domain(
10777 "an executed string cannot take host data: `{name}` has nothing to bind to",
10778 span,
10779 ));
10780 }
10781 let mut rec = None;
10782 let (value, _) = crate::ir::run_block(&nested.stmts, None, ctx, &mut rec)
10783 .map_err(|e| nested_error(e, src, span))?;
10784 value.ok_or_else(|| Error::domain("the executed string yielded no value", span))
10785}
10786
10787fn stream_number(y: &Array, open: i64, what: &str, span: Span) -> Result<()> {
10794 let closed = || {
10795 Err(Error::sandbox(
10796 format!("{what} the standard stream {open} only; a file is outside the program"),
10797 span,
10798 ))
10799 };
10800 if matches!(y.data, Data::Box(_)) {
10801 return closed();
10802 }
10803 match y.to_i64_vec().as_deref() {
10804 Some([n]) if *n == open => Ok(()),
10805 Some([_]) => closed(),
10806 _ => Err(Error::domain(format!("{what} one stream number"), span)),
10807 }
10808}
10809
10810fn type_code(y: &Array) -> i64 {
10813 match y.dtype() {
10814 DType::Bool => 1,
10815 DType::Char => 2,
10816 DType::I64 => 4,
10817 DType::F64 => 8,
10818 DType::Complex => 16,
10819 DType::Box => 32,
10820 DType::Ext => 64,
10821 DType::Rat => 128,
10822 }
10823}
10824
10825fn nested_error(e: Error, src: &str, span: Span) -> Error {
10829 let inner = e.render(src);
10830 let mut out = Error::new(e.kind, format!("in the executed string: {}", e.msg), Some(span));
10831 out.notes.push(inner.trim_end().to_string());
10832 out
10833}
10834
10835fn words(y: &Array, span: Span) -> Result<Array> {
10841 let Data::Char(v) = &y.data else {
10842 return Err(Error::domain("words reads a character list", span));
10843 };
10844 let src: Vec<char> = v.as_slice().to_vec();
10845 let n = src.len();
10846 let mut out: Vec<Array> = Vec::new();
10847 let mut i = 0usize;
10848 let numeric_start = |k: usize| -> bool {
10849 k < n && (src[k].is_ascii_digit() || src[k] == '_')
10850 };
10851 while i < n {
10852 let c = src[i];
10853 if c == ' ' || c == '\t' {
10854 i += 1;
10855 continue;
10856 }
10857 let start = i;
10858 if c == '\'' {
10859 i += 1;
10860 loop {
10861 if i >= n {
10862 return Err(Error::parse("a word list ends inside a string", span));
10863 }
10864 if src[i] == '\'' {
10865 i += 1;
10866 if i < n && src[i] == '\'' {
10867 i += 1;
10868 continue;
10869 }
10870 break;
10871 }
10872 i += 1;
10873 }
10874 } else if c.is_ascii_alphabetic() {
10875 while i < n && (src[i].is_ascii_alphanumeric() || src[i] == '_') {
10876 i += 1;
10877 }
10878 if i < n && (src[i] == '.' || src[i] == ':') {
10879 i += 1;
10880 }
10881 if src[start..i].iter().collect::<String>() == "NB." {
10883 while i < n && src[i] != '\n' {
10884 i += 1;
10885 }
10886 }
10887 } else if numeric_start(i) {
10888 loop {
10889 while i < n && (src[i].is_ascii_alphanumeric() || src[i] == '.' || src[i] == '_')
10890 {
10891 i += 1;
10892 }
10893 let mut j = i;
10895 while j < n && src[j] == ' ' {
10896 j += 1;
10897 }
10898 if j > i && numeric_start(j) {
10899 i = j;
10900 continue;
10901 }
10902 break;
10903 }
10904 } else {
10905 i += 1;
10906 while i < n && (src[i] == '.' || src[i] == ':') {
10907 i += 1;
10908 }
10909 }
10910 out.push(Array::from_chars(src[start..i].to_vec()));
10911 }
10912 let k = out.len();
10913 Ok(Array::new(vec![k], Data::Box(out.into())))
10914}
10915
10916#[cfg(test)]
10917mod tests {
10918 use super::*;
10919
10920 macro_rules! ctx {
10922 ($name:ident, $agreement:expr) => {
10923 let mut sink = |_: &str| {};
10924 let mut env = Env::new(Vec::new());
10925 #[allow(unused_mut)]
10926 let mut $name = Ctx {
10927 cfg: EvalCfg {
10928 agreement: $agreement,
10929 fmt: FmtOpts::J,
10930 tol: Tol::J,
10931 rules: crate::frontend::Dialect::default()
10934 .rules(if $agreement == Agreement::ExactOrScalar {
10935 crate::Lang::Apl
10936 } else {
10937 crate::Lang::J
10938 })
10939 .expect("the shipped dialect is implemented"),
10940 },
10941 out: &mut sink,
10942 inp: None,
10943 env: &mut env,
10944 device: None,
10945 };
10946 };
10947 ($name:ident) => {
10948 ctx!($name, Agreement::LeadingPrefix);
10949 };
10950 }
10951
10952 fn scalar_prim(name: &'static str, monad: MonadOp, dyad: DyadOp) -> Verb {
10953 Verb::Prim(Prim { name, monad, dyad, ranks: [0, 0, 0] })
10954 }
10955
10956 fn inf_prim(name: &'static str, monad: MonadOp, dyad: DyadOp) -> Verb {
10957 Verb::Prim(Prim { name, monad, dyad, ranks: [RANK_INF, RANK_INF, RANK_INF] })
10958 }
10959
10960 fn plus() -> Verb {
10961 scalar_prim("+", MonadOp::Scalar(ScalarMonad::Conj), DyadOp::Scalar(ScalarDyad::Add))
10962 }
10963 fn minus() -> Verb {
10964 scalar_prim("-", MonadOp::Scalar(ScalarMonad::Neg), DyadOp::Scalar(ScalarDyad::Sub))
10965 }
10966 fn times() -> Verb {
10967 scalar_prim("*", MonadOp::Scalar(ScalarMonad::Signum), DyadOp::Scalar(ScalarDyad::Mul))
10968 }
10969 fn pct() -> Verb {
10970 scalar_prim("%", MonadOp::Scalar(ScalarMonad::Recip), DyadOp::Scalar(ScalarDyad::DivJ))
10971 }
10972 fn div_apl() -> Verb {
10973 scalar_prim("÷", MonadOp::Scalar(ScalarMonad::Recip), DyadOp::Scalar(ScalarDyad::DivApl))
10974 }
10975 fn floor_v() -> Verb {
10976 scalar_prim("<.", MonadOp::Scalar(ScalarMonad::Floor), DyadOp::Scalar(ScalarDyad::Min))
10977 }
10978 fn ceil_v() -> Verb {
10979 scalar_prim(">.", MonadOp::Scalar(ScalarMonad::Ceil), DyadOp::Scalar(ScalarDyad::Max))
10980 }
10981 fn pow_v() -> Verb {
10982 scalar_prim("^", MonadOp::Scalar(ScalarMonad::Exp), DyadOp::Scalar(ScalarDyad::Pow))
10983 }
10984 fn residue_v() -> Verb {
10985 scalar_prim("|", MonadOp::Scalar(ScalarMonad::Abs), DyadOp::Scalar(ScalarDyad::Residue))
10986 }
10987 fn eq_v() -> Verb {
10988 scalar_prim("=", MonadOp::None, DyadOp::Scalar(ScalarDyad::Eq))
10989 }
10990 fn lt_v() -> Verb {
10991 scalar_prim("<", MonadOp::None, DyadOp::Scalar(ScalarDyad::Lt))
10992 }
10993 fn not_v() -> Verb {
10994 scalar_prim("-.", MonadOp::Scalar(ScalarMonad::Not), DyadOp::None)
10995 }
10996 fn sqrt_v() -> Verb {
10997 scalar_prim("%:", MonadOp::Scalar(ScalarMonad::Sqrt), DyadOp::NotYet("dyadic root"))
10998 }
10999 fn dollar() -> Verb {
11000 inf_prim("$", MonadOp::ShapeOf, DyadOp::Reshape)
11001 }
11002 fn pound() -> Verb {
11003 inf_prim("#", MonadOp::Tally, DyadOp::NotYet("copy"))
11004 }
11005 fn comma() -> Verb {
11006 inf_prim(",", MonadOp::Ravel, DyadOp::NotYet("append"))
11007 }
11008 fn transpose_v() -> Verb {
11009 inf_prim("|:", MonadOp::TransposeAxes, DyadOp::NotYet("dyadic transpose"))
11010 }
11011 fn head_v() -> Verb {
11012 inf_prim("{.", MonadOp::Head, DyadOp::Take)
11013 }
11014 fn behead_v() -> Verb {
11015 inf_prim("}.", MonadOp::Behead, DyadOp::Drop)
11016 }
11017 fn iota() -> Verb {
11018 inf_prim("i.", MonadOp::IotaJ, DyadOp::NotYet("index of"))
11019 }
11020 fn iota_apl(origin: i64) -> Verb {
11021 inf_prim("⍳", MonadOp::IotaApl { origin }, DyadOp::NotYet("index of"))
11022 }
11023 fn right_v() -> Verb {
11024 inf_prim("]", MonadOp::Same, DyadOp::Right)
11025 }
11026 fn echo_v() -> Verb {
11027 inf_prim("echo", MonadOp::Echo, DyadOp::None)
11028 }
11029
11030 fn b(v: Verb) -> Box<Verb> {
11031 Box::new(v)
11032 }
11033
11034 fn mat(rows: usize, cols: usize, v: Vec<i64>) -> Array {
11035 Array::new(vec![rows, cols], Data::I64(v.into()))
11036 }
11037
11038 fn ints(a: &Array) -> Vec<i64> {
11040 a.to_row_major().as_i64_slice().expect("integer result").to_vec()
11041 }
11042
11043 fn floats(a: &Array) -> Vec<f64> {
11044 a.to_row_major().as_f64_slice().expect("float result").to_vec()
11045 }
11046
11047 fn bools(a: &Array) -> Vec<u8> {
11048 match &a.to_row_major().data {
11049 Data::Bool(v) => v.to_vec(),
11050 other => panic!("expected boolean result, got {other:?}"),
11051 }
11052 }
11053
11054 fn sp() -> Span {
11055 Span::new(0, 1)
11056 }
11057
11058 fn close(a: f64, b: f64) -> bool {
11059 (a - b).abs() < 1e-9 || (a.is_infinite() && b.is_infinite() && a.signum() == b.signum())
11060 }
11061
11062 #[test]
11065 fn names_of_primitives_and_derived_verbs() {
11066 assert_eq!(plus().name(), "+");
11067 assert_eq!(Verb::Rank(b(plus()), [1, 1, 1]).name(), "+\"1");
11068 assert_eq!(Verb::Rank(b(plus()), [0, 1, RANK_INF]).name(), "+\"0 1 _");
11069 assert_eq!(Verb::Rank(b(plus()), [RANK_INF; 3]).name(), "+\"_");
11070 assert_eq!(Verb::Reduce(b(plus())).name(), "+/");
11071 assert_eq!(Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]).name(), "+/\"1");
11072 assert_eq!(Verb::Fork(b(plus()), b(minus()), b(times())).name(), "(+ - *)");
11073 assert_eq!(
11074 Verb::NounFork(Array::scalar_i64(1), b(plus()), b(minus())).name(),
11075 "(n + -)"
11076 );
11077 assert_eq!(Verb::Hook(b(plus()), b(minus())).name(), "(+ -)");
11078 assert_eq!(Verb::Atop(b(plus()), b(minus())).name(), "(+@:-)");
11079 assert_eq!(Verb::Compose(b(plus()), b(minus())).name(), "(+&:-)");
11080 assert_eq!(Verb::BondLeft(Array::scalar_i64(1), b(plus())).name(), "(n&+)");
11081 assert_eq!(Verb::BondRight(b(plus()), Array::scalar_i64(1)).name(), "(+&n)");
11082 }
11083
11084 #[test]
11085 fn composition_applies_the_right_verb_to_both_arguments() {
11086 ctx!(c);
11087 let v = Verb::Compose(b(plus()), b(times()));
11088 let r = v.monad(&Array::from_i64(vec![-2, 0, 3]), &mut c, sp()).unwrap();
11090 assert_eq!(ints(&r), vec![-1, 0, 1]);
11091 let r = v
11092 .dyad(&Array::scalar_i64(-5), &Array::scalar_i64(7), &mut c, sp())
11093 .unwrap();
11094 assert_eq!(ints(&r), vec![0]);
11095 let bond = Verb::BondLeft(Array::scalar_i64(10), b(minus()));
11097 let r = bond.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
11098 assert_eq!(ints(&r), vec![9, 8]);
11099 let e = bond
11100 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(2), &mut c, sp())
11101 .unwrap_err();
11102 assert_eq!(e.kind, ErrorKind::Domain);
11103 let bond = Verb::BondRight(b(minus()), Array::scalar_i64(10));
11104 let r = bond.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
11105 assert_eq!(ints(&r), vec![-9, -8]);
11106 }
11107
11108 #[test]
11111 fn scalar_monad_covers_the_whole_buffer() {
11112 ctx!(c);
11113 let r = minus().monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
11114 assert_eq!(r.shape, vec![2, 3]);
11115 assert_eq!(ints(&r), vec![-1, -2, -3, -4, -5, -6]);
11116 }
11117
11118 #[test]
11119 fn leading_prefix_agreement_broadcasts_per_row() {
11120 ctx!(c);
11121 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
11122 let y = Array::from_i64(vec![10, 20]);
11123 let r = plus().dyad(&x, &y, &mut c, sp()).unwrap();
11124 assert_eq!(r.shape, vec![2, 3]);
11125 assert_eq!(ints(&r), vec![11, 12, 13, 24, 25, 26]);
11126 let r = plus().dyad(&y, &x, &mut c, sp()).unwrap();
11128 assert_eq!(ints(&r), vec![11, 12, 13, 24, 25, 26]);
11129 }
11130
11131 #[test]
11132 fn exact_or_scalar_rejects_a_prefix_frame() {
11133 ctx!(c, Agreement::ExactOrScalar);
11134 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
11135 let y = Array::from_i64(vec![10, 20]);
11136 let e = plus().dyad(&x, &y, &mut c, sp()).unwrap_err();
11137 assert_eq!(e.kind, ErrorKind::Shape);
11138 assert!(e.msg.contains("2 3"), "{}", e.msg);
11139 assert!(e.msg.contains("right shape 2"), "{}", e.msg);
11140 }
11141
11142 #[test]
11143 fn exact_or_scalar_accepts_equal_frames_and_scalars() {
11144 ctx!(c, Agreement::ExactOrScalar);
11145 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
11146 let r = plus().dyad(&x, &x, &mut c, sp()).unwrap();
11147 assert_eq!(ints(&r), vec![2, 4, 6, 8, 10, 12]);
11148 let r = plus().dyad(&Array::scalar_i64(10), &x, &mut c, sp()).unwrap();
11149 assert_eq!(r.shape, vec![2, 3]);
11150 assert_eq!(ints(&r), vec![11, 12, 13, 14, 15, 16]);
11151 let r = plus().dyad(&x, &Array::scalar_i64(10), &mut c, sp()).unwrap();
11152 assert_eq!(ints(&r), vec![11, 12, 13, 14, 15, 16]);
11153 }
11154
11155 #[test]
11156 fn vector_length_mismatch_is_a_length_error() {
11157 ctx!(c);
11158 let e = plus()
11159 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::from_i64(vec![1, 2, 3, 4, 5]), &mut c, sp())
11160 .unwrap_err();
11161 assert_eq!(e.kind, ErrorKind::Length);
11162 assert!(e.msg.contains("left shape 3"), "{}", e.msg);
11163 assert!(e.msg.contains("right shape 5"), "{}", e.msg);
11164 assert!(e.notes[0].contains("axis 0"), "{:?}", e.notes);
11165 }
11166
11167 #[test]
11168 fn diverging_matrix_frames_name_the_axis() {
11169 ctx!(c);
11170 let e = plus()
11171 .dyad(&mat(2, 3, vec![0; 6]), &mat(2, 4, vec![0; 8]), &mut c, sp())
11172 .unwrap_err();
11173 assert_eq!(e.kind, ErrorKind::Shape);
11174 assert!(e.notes[0].contains("axis 1"), "{:?}", e.notes);
11175 }
11176
11177 #[test]
11178 fn dyadic_rank_pairs_rows_with_the_whole_right_argument() {
11179 ctx!(c);
11180 let v = Verb::Rank(b(plus()), [0, 1, 1]);
11182 let r = v
11183 .dyad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &Array::from_i64(vec![10, 20, 30]), &mut c, sp())
11184 .unwrap();
11185 assert_eq!(r.shape, vec![2, 3]);
11186 assert_eq!(ints(&r), vec![11, 22, 33, 14, 25, 36]);
11187 }
11188
11189 #[test]
11190 fn surplus_frame_axes_repeat_the_shorter_frames_cells() {
11191 ctx!(c);
11192 let v = Verb::Rank(b(head_v()), [0, 0, 1]);
11195 let x = mat(2, 2, vec![1, 1, 2, 2]);
11196 let y = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
11197 let r = v.dyad(&x, &y, &mut c, sp()).unwrap();
11198 assert_eq!(r.shape, vec![2, 2, 2]);
11199 assert_eq!(ints(&r), vec![1, 0, 1, 0, 4, 5, 4, 5]);
11200 }
11201
11202 #[test]
11203 fn an_empty_frame_pairs_its_single_cell_with_every_other_cell() {
11204 ctx!(c, Agreement::ExactOrScalar);
11205 let v = Verb::Rank(b(head_v()), [0, 0, 1]);
11207 let x = Array::from_i64(vec![1, 2]);
11208 let y = Array::from_i64(vec![7, 8, 9]);
11209 let r = v.dyad(&x, &y, &mut c, sp()).unwrap();
11210 assert_eq!(r.shape, vec![2, 2]);
11211 assert_eq!(ints(&r), vec![7, 0, 7, 8]);
11212 }
11213
11214 #[test]
11215 fn negative_rank_leaves_frame_axes() {
11216 ctx!(c);
11217 let v = Verb::Rank(b(dollar()), [-1, -1, -1]);
11219 let r = v.monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
11220 assert_eq!(r.shape, vec![2, 1]);
11221 assert_eq!(ints(&r), vec![3, 3]);
11222 }
11223
11224 #[test]
11225 fn effective_rank_clamps_and_counts_back() {
11226 assert_eq!(effective_rank(0, 3), 0);
11227 assert_eq!(effective_rank(2, 1), 1);
11228 assert_eq!(effective_rank(RANK_INF, 4), 4);
11229 assert_eq!(effective_rank(-1, 3), 2);
11230 assert_eq!(effective_rank(-5, 3), 0);
11231 }
11232
11233 #[test]
11236 fn reduction_folds_right_to_left() {
11237 ctx!(c);
11238 let r = Verb::Reduce(b(minus()))
11240 .monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp())
11241 .unwrap();
11242 assert!(r.shape.is_empty());
11243 assert_eq!(ints(&r), vec![2]);
11244 }
11245
11246 #[test]
11247 fn reduction_of_one_item_and_of_a_scalar() {
11248 ctx!(c);
11249 let r = Verb::Reduce(b(plus()))
11250 .monad(&Array::from_i64(vec![7]), &mut c, sp())
11251 .unwrap();
11252 assert!(r.shape.is_empty());
11253 assert_eq!(ints(&r), vec![7]);
11254 let r = Verb::Reduce(b(plus())).monad(&Array::scalar_i64(7), &mut c, sp()).unwrap();
11255 assert_eq!(ints(&r), vec![7]);
11256 }
11257
11258 #[test]
11259 fn reduction_runs_along_the_leading_axis() {
11260 ctx!(c);
11261 let r = Verb::Reduce(b(plus()))
11262 .monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp())
11263 .unwrap();
11264 assert_eq!(r.shape, vec![3]);
11265 assert_eq!(ints(&r), vec![5, 7, 9]);
11266 }
11267
11268 #[test]
11269 fn rank_wrapped_reduction_sums_the_last_axis() {
11270 ctx!(c);
11271 let v = Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]);
11272 let r = v.monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
11273 assert_eq!(r.shape, vec![2]);
11274 assert_eq!(ints(&r), vec![6, 15]);
11275 }
11276
11277 #[test]
11278 fn empty_reduction_uses_the_identity_cell() {
11279 ctx!(c);
11280 let empty = Array::new(vec![0, 2], Data::I64(vec![].into()));
11281 let r = Verb::Reduce(b(plus())).monad(&empty, &mut c, sp()).unwrap();
11282 assert_eq!(r.shape, vec![2]);
11283 assert_eq!(ints(&r), vec![0, 0]);
11284 let r = Verb::Reduce(b(times())).monad(&empty, &mut c, sp()).unwrap();
11285 assert_eq!(ints(&r), vec![1, 1]);
11286 let r = Verb::Reduce(b(floor_v())).monad(&empty, &mut c, sp()).unwrap();
11287 assert!(floats(&r).iter().all(|&x| x == f64::INFINITY));
11288 let r = Verb::Reduce(b(ceil_v())).monad(&empty, &mut c, sp()).unwrap();
11289 assert!(floats(&r).iter().all(|&x| x == f64::NEG_INFINITY));
11290 let r = Verb::Reduce(b(minus())).monad(&empty, &mut c, sp()).unwrap();
11293 assert_eq!(ints(&r), vec![0, 0]);
11294 let r = Verb::Reduce(b(pct())).monad(&empty, &mut c, sp()).unwrap();
11295 assert_eq!(ints(&r), vec![1, 1]);
11296 let r = Verb::Reduce(b(eq_v())).monad(&empty, &mut c, sp()).unwrap();
11297 assert_eq!(bools(&r), vec![1, 1]);
11298 let r = Verb::Reduce(b(plus()))
11300 .monad(&Array::empty(DType::I64), &mut c, sp())
11301 .unwrap();
11302 assert!(r.shape.is_empty());
11303 assert_eq!(ints(&r), vec![0]);
11304 }
11305
11306 #[test]
11307 fn empty_reduction_without_an_identity_is_a_domain_error() {
11308 ctx!(c);
11309 let v = Verb::Hook(b(plus()), b(minus()));
11313 let e = Verb::Reduce(b(v)).monad(&Array::empty(DType::I64), &mut c, sp()).unwrap_err();
11314 assert_eq!(e.kind, ErrorKind::Domain);
11315 assert!(e.msg.contains("identity"), "{}", e.msg);
11316 }
11317
11318 #[test]
11319 fn reduction_with_a_non_primitive_verb_uses_the_general_fold() {
11320 ctx!(c);
11321 let v = Verb::Reduce(b(Verb::Hook(b(plus()), b(minus()))));
11323 let r = v.monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp()).unwrap();
11324 assert_eq!(ints(&r), vec![2]);
11325 }
11326
11327 #[test]
11328 fn dyadic_reduction_is_the_table() {
11329 ctx!(c);
11330 let v = Verb::Reduce(b(plus()));
11333 let r = v
11334 .dyad(&Array::scalar_i64(2), &Array::from_i64(vec![1, 2, 3]), &mut c, sp())
11335 .unwrap();
11336 assert_eq!(r.shape, vec![3]);
11337 assert_eq!(ints(&r), vec![3, 4, 5]);
11338 let r = v
11341 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::from_i64(vec![10, 20]), &mut c, sp())
11342 .unwrap();
11343 assert_eq!(r.shape, vec![3, 2]);
11344 assert_eq!(ints(&r), vec![11, 21, 12, 22, 13, 23]);
11345 let cat = Verb::Reduce(b(inf_prim(",", MonadOp::Ravel, DyadOp::AppendLeading)));
11347 let r = cat
11348 .dyad(&Array::from_i64(vec![1, 2]), &Array::from_i64(vec![3, 4]), &mut c, sp())
11349 .unwrap();
11350 assert_eq!(r.shape, vec![4]);
11351 assert_eq!(ints(&r), vec![1, 2, 3, 4]);
11352 }
11353
11354 #[test]
11357 fn integer_overflow_promotes_the_whole_result_to_float() {
11358 ctx!(c);
11359 let r = plus()
11360 .dyad(&Array::from_i64(vec![1, i64::MAX]), &Array::scalar_i64(1), &mut c, sp())
11361 .unwrap();
11362 assert_eq!(r.dtype(), DType::F64);
11363 let v = floats(&r);
11364 assert!(close(v[0], 2.0));
11365 assert!(close(v[1], i64::MAX as f64 + 1.0));
11366 let r = plus()
11368 .dyad(&Array::from_i64(vec![1, 2]), &Array::scalar_i64(1), &mut c, sp())
11369 .unwrap();
11370 assert_eq!(r.dtype(), DType::I64);
11371 }
11372
11373 #[test]
11374 fn reduction_overflow_promotes_too() {
11375 ctx!(c);
11376 let r = Verb::Reduce(b(plus()))
11377 .monad(&Array::from_i64(vec![i64::MAX, i64::MAX]), &mut c, sp())
11378 .unwrap();
11379 assert_eq!(r.dtype(), DType::F64);
11380 assert!(close(floats(&r)[0], 2.0 * i64::MAX as f64));
11381 }
11382
11383 #[test]
11384 fn booleans_widen_to_integers_in_arithmetic() {
11385 ctx!(c);
11386 let bits = Array::new(vec![3], Data::Bool(vec![1, 0, 1].into()));
11387 let r = plus().dyad(&bits, &bits, &mut c, sp()).unwrap();
11388 assert_eq!(r.dtype(), DType::I64);
11389 assert_eq!(ints(&r), vec![2, 0, 2]);
11390 }
11391
11392 #[test]
11393 fn j_division_is_float_and_survives_zero() {
11394 ctx!(c);
11395 let r = pct()
11396 .dyad(&Array::from_i64(vec![1, -1, 0, 6]), &Array::from_i64(vec![0, 0, 0, 4]), &mut c, sp())
11397 .unwrap();
11398 let v = floats(&r);
11399 assert_eq!(v[0], f64::INFINITY);
11400 assert_eq!(v[1], f64::NEG_INFINITY);
11401 assert_eq!(v[2], 0.0);
11402 assert!(close(v[3], 1.5));
11403 }
11404
11405 #[test]
11406 fn apl_division_by_zero_is_a_domain_error_except_zero_by_zero() {
11407 ctx!(c, Agreement::ExactOrScalar);
11408 let r = div_apl()
11409 .dyad(&Array::scalar_i64(0), &Array::scalar_i64(0), &mut c, sp())
11410 .unwrap();
11411 assert!(close(floats(&r)[0], 1.0));
11412 let e = div_apl()
11413 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(0), &mut c, sp())
11414 .unwrap_err();
11415 assert_eq!(e.kind, ErrorKind::Domain);
11416 assert!(e.msg.contains("division by zero"), "{}", e.msg);
11417 let r = div_apl()
11418 .dyad(&Array::scalar_i64(6), &Array::scalar_i64(4), &mut c, sp())
11419 .unwrap();
11420 assert!(close(floats(&r)[0], 1.5));
11421 }
11422
11423 #[test]
11424 fn reciprocal_of_zero_is_infinite() {
11425 ctx!(c);
11426 let r = pct().monad(&Array::from_i64(vec![0, 2]), &mut c, sp()).unwrap();
11427 let v = floats(&r);
11428 assert_eq!(v[0], f64::INFINITY);
11429 assert!(close(v[1], 0.5));
11430 }
11431
11432 #[test]
11433 fn residue_takes_the_sign_of_the_left_argument() {
11434 ctx!(c);
11435 let x = Array::from_i64(vec![3, 3, -3, -3, 0]);
11436 let y = Array::from_i64(vec![5, -5, 5, -5, 5]);
11437 let r = residue_v().dyad(&x, &y, &mut c, sp()).unwrap();
11438 assert_eq!(ints(&r), vec![2, 1, -1, -2, 5]);
11439 let r = residue_v()
11441 .dyad(&Array::from_f64(vec![2.5]), &Array::from_f64(vec![7.0]), &mut c, sp())
11442 .unwrap();
11443 assert!(close(floats(&r)[0], 2.0));
11444 }
11445
11446 #[test]
11447 fn power_stays_integral_when_it_can() {
11448 ctx!(c);
11449 let r = pow_v()
11450 .dyad(&Array::from_i64(vec![2, 0, 5]), &Array::from_i64(vec![10, 0, 1]), &mut c, sp())
11451 .unwrap();
11452 assert_eq!(r.dtype(), DType::I64);
11453 assert_eq!(ints(&r), vec![1024, 1, 5]);
11454 let r = pow_v()
11456 .dyad(&Array::from_i64(vec![2, 4]), &Array::from_i64(vec![-1, 2]), &mut c, sp())
11457 .unwrap();
11458 assert_eq!(r.dtype(), DType::F64);
11459 assert!(close(floats(&r)[0], 0.5));
11460 assert!(close(floats(&r)[1], 16.0));
11461 let r = pow_v()
11463 .dyad(&Array::scalar_i64(10), &Array::scalar_i64(30), &mut c, sp())
11464 .unwrap();
11465 assert_eq!(r.dtype(), DType::F64);
11466 }
11467
11468 #[test]
11469 fn comparisons_yield_booleans() {
11470 ctx!(c);
11471 let r = lt_v()
11472 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::scalar_i64(2), &mut c, sp())
11473 .unwrap();
11474 assert_eq!(bools(&r), vec![1, 0, 0]);
11475 let r = eq_v()
11476 .dyad(&Array::from_f64(vec![1.0, 2.0]), &Array::from_i64(vec![1, 3]), &mut c, sp())
11477 .unwrap();
11478 assert_eq!(bools(&r), vec![1, 0]);
11479 }
11480
11481 #[test]
11482 fn characters_compare_but_do_not_add() {
11483 ctx!(c);
11484 let a = Array::from_chars(vec!['a', 'b']);
11485 let bb = Array::from_chars(vec!['a', 'c']);
11486 assert_eq!(bools(&eq_v().dyad(&a, &bb, &mut c, sp()).unwrap()), vec![1, 0]);
11487 let e = plus().dyad(&a, &bb, &mut c, sp()).unwrap_err();
11488 assert_eq!(e.kind, ErrorKind::Type);
11489 assert!(e.msg.contains("characters"), "{}", e.msg);
11490 let e = lt_v().dyad(&a, &bb, &mut c, sp()).unwrap_err();
11491 assert_eq!(e.kind, ErrorKind::Type);
11492 let e = plus().dyad(&a, &Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap_err();
11493 assert_eq!(e.kind, ErrorKind::Type);
11494 assert!(e.msg.contains("character"), "{}", e.msg);
11495 let e = plus().monad(&a, &mut c, sp()).unwrap_err();
11496 assert_eq!(e.kind, ErrorKind::Type);
11497 }
11498
11499 #[test]
11500 fn floor_and_ceiling_return_integers_when_they_fit() {
11501 ctx!(c);
11502 let r = floor_v().monad(&Array::from_f64(vec![1.5, -1.5]), &mut c, sp()).unwrap();
11503 assert_eq!(r.dtype(), DType::I64);
11504 assert_eq!(ints(&r), vec![1, -2]);
11505 let r = ceil_v().monad(&Array::from_f64(vec![1.5, -1.5]), &mut c, sp()).unwrap();
11506 assert_eq!(ints(&r), vec![2, -1]);
11507 let r = floor_v().monad(&Array::from_f64(vec![1e30]), &mut c, sp()).unwrap();
11509 assert_eq!(r.dtype(), DType::F64);
11510 let r = floor_v().monad(&Array::from_i64(vec![3]), &mut c, sp()).unwrap();
11512 assert_eq!(ints(&r), vec![3]);
11513 }
11514
11515 #[test]
11516 fn logical_negation_needs_zero_or_one() {
11517 ctx!(c);
11518 let r = not_v().monad(&Array::from_i64(vec![0, 1]), &mut c, sp()).unwrap();
11519 assert_eq!(bools(&r), vec![1, 0]);
11520 let e = not_v().monad(&Array::from_i64(vec![2]), &mut c, sp()).unwrap_err();
11521 assert_eq!(e.kind, ErrorKind::Domain);
11522 }
11523
11524 #[test]
11525 fn signum_abs_and_negation_pick_their_types() {
11526 ctx!(c);
11527 let r = times().monad(&Array::from_i64(vec![-3, 0, 9]), &mut c, sp()).unwrap();
11528 assert_eq!(ints(&r), vec![-1, 0, 1]);
11529 let r = times().monad(&Array::from_f64(vec![-3.0, 0.0, 9.0]), &mut c, sp()).unwrap();
11530 assert_eq!(floats(&r), vec![-1.0, 0.0, 1.0]);
11531 let r = residue_v().monad(&Array::from_i64(vec![-3, 3]), &mut c, sp()).unwrap();
11532 assert_eq!(ints(&r), vec![3, 3]);
11533 let bits = Array::new(vec![2], Data::Bool(vec![0, 1].into()));
11534 let r = minus().monad(&bits, &mut c, sp()).unwrap();
11535 assert_eq!(r.dtype(), DType::I64);
11536 assert_eq!(ints(&r), vec![0, -1]);
11537 }
11538
11539 #[test]
11540 fn square_root_of_a_negative_number_is_complex() {
11541 ctx!(c);
11542 let r = sqrt_v().monad(&Array::from_i64(vec![9]), &mut c, sp()).unwrap();
11543 assert!(close(floats(&r)[0], 3.0));
11544 let r = sqrt_v().monad(&Array::from_i64(vec![-4]), &mut c, sp()).unwrap();
11545 assert_eq!(r.dtype(), DType::Complex);
11546 assert_eq!(r.as_complex_slice().expect("complex data"), &[[0.0, 2.0]]);
11547 }
11548
11549 #[test]
11552 fn shape_tally_and_ravel() {
11553 ctx!(c);
11554 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
11555 let r = dollar().monad(&m, &mut c, sp()).unwrap();
11556 assert_eq!(r.shape, vec![2]);
11557 assert_eq!(ints(&r), vec![2, 3]);
11558 let r = pound().monad(&m, &mut c, sp()).unwrap();
11559 assert!(r.shape.is_empty());
11560 assert_eq!(ints(&r), vec![2]);
11561 let r = pound().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap();
11563 assert_eq!(ints(&r), vec![1]);
11564 let r = comma().monad(&m, &mut c, sp()).unwrap();
11565 assert_eq!(r.shape, vec![6]);
11566 assert_eq!(ints(&r), vec![1, 2, 3, 4, 5, 6]);
11567 }
11568
11569 #[test]
11570 fn transpose_reverses_the_axes() {
11571 ctx!(c);
11572 let r = transpose_v().monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
11573 assert_eq!(r.shape, vec![3, 2]);
11574 assert_eq!(ints(&r), vec![1, 4, 2, 5, 3, 6]);
11575 let a = Array::new(vec![2, 1, 3], Data::I64(vec![1, 2, 3, 4, 5, 6].into()));
11577 let r = transpose_v().monad(&a, &mut c, sp()).unwrap();
11578 assert_eq!(r.shape, vec![3, 1, 2]);
11579 assert_eq!(ints(&r), vec![1, 4, 2, 5, 3, 6]);
11580 let v = Array::from_i64(vec![1, 2]);
11582 assert_eq!(transpose_v().monad(&v, &mut c, sp()).unwrap(), v);
11583 }
11584
11585 #[test]
11586 fn head_and_behead() {
11587 ctx!(c);
11588 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
11589 let r = head_v().monad(&m, &mut c, sp()).unwrap();
11590 assert_eq!(r.shape, vec![3]);
11591 assert_eq!(ints(&r), vec![1, 2, 3]);
11592 let r = behead_v().monad(&m, &mut c, sp()).unwrap();
11593 assert_eq!(r.shape, vec![1, 3]);
11594 assert_eq!(ints(&r), vec![4, 5, 6]);
11595 let e = Array::new(vec![0, 2], Data::I64(vec![].into()));
11597 let r = head_v().monad(&e, &mut c, sp()).unwrap();
11598 assert_eq!(r.shape, vec![2]);
11599 assert_eq!(ints(&r), vec![0, 0]);
11600 assert_eq!(behead_v().monad(&e, &mut c, sp()).unwrap(), e);
11601 assert_eq!(head_v().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap().shape, Vec::<usize>::new());
11602 let err = behead_v().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap_err();
11603 assert_eq!(err.kind, ErrorKind::Domain);
11604 }
11605
11606 #[test]
11607 fn iota_fills_a_shape_and_reverses_negative_axes() {
11608 ctx!(c);
11609 let r = iota().monad(&Array::from_i64(vec![2, 3]), &mut c, sp()).unwrap();
11610 assert_eq!(r.shape, vec![2, 3]);
11611 assert_eq!(ints(&r), vec![0, 1, 2, 3, 4, 5]);
11612 let r = iota().monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
11614 assert_eq!(r.shape, vec![3]);
11615 assert_eq!(ints(&r), vec![0, 1, 2]);
11616 let r = iota().monad(&Array::scalar_i64(-3), &mut c, sp()).unwrap();
11618 assert_eq!(ints(&r), vec![2, 1, 0]);
11619 let r = iota().monad(&Array::from_i64(vec![2, -3]), &mut c, sp()).unwrap();
11620 assert_eq!(r.shape, vec![2, 3]);
11621 assert_eq!(ints(&r), vec![2, 1, 0, 5, 4, 3]);
11622 let r = iota().monad(&Array::from_i64(vec![-2, 3]), &mut c, sp()).unwrap();
11623 assert_eq!(ints(&r), vec![3, 4, 5, 0, 1, 2]);
11624 let r = iota().monad(&Array::scalar_i64(0), &mut c, sp()).unwrap();
11626 assert_eq!(r.shape, vec![0]);
11627 assert!(ints(&r).is_empty());
11628 let e = iota().monad(&Array::from_f64(vec![1.5]), &mut c, sp()).unwrap_err();
11630 assert_eq!(e.kind, ErrorKind::Domain);
11631 let e = iota().monad(&mat(1, 1, vec![1]), &mut c, sp()).unwrap_err();
11632 assert_eq!(e.kind, ErrorKind::Rank);
11633 }
11634
11635 #[test]
11636 fn apl_iota_starts_at_the_index_origin() {
11637 ctx!(c, Agreement::ExactOrScalar);
11638 let r = iota_apl(1).monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
11639 assert_eq!(ints(&r), vec![1, 2, 3]);
11640 let r = iota_apl(0).monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
11641 assert_eq!(ints(&r), vec![0, 1, 2]);
11642 let e = iota_apl(1).monad(&Array::scalar_i64(-1), &mut c, sp()).unwrap_err();
11643 assert_eq!(e.kind, ErrorKind::Domain);
11644 let r = iota_apl(1).monad(&Array::from_i64(vec![2, 3]), &mut c, sp()).unwrap();
11647 assert_eq!(r.shape, vec![2, 3]);
11648 assert_eq!(ints(&r.as_boxes().expect("boxed")[4]), vec![2, 2]);
11649 }
11650
11651 #[test]
11652 fn reshape_cycles_the_ravel() {
11653 ctx!(c);
11654 let r = dollar()
11655 .dyad(&Array::from_i64(vec![2, 3]), &Array::from_i64(vec![1, 2]), &mut c, sp())
11656 .unwrap();
11657 assert_eq!(r.shape, vec![2, 3]);
11658 assert_eq!(ints(&r), vec![1, 2, 1, 2, 1, 2]);
11659 let r = dollar()
11661 .dyad(&Array::scalar_i64(3), &Array::from_i64(vec![7]), &mut c, sp())
11662 .unwrap();
11663 assert_eq!(r.shape, vec![3]);
11664 assert_eq!(ints(&r), vec![7, 7, 7]);
11665 let r = dollar()
11667 .dyad(&Array::scalar_i64(2), &Array::from_chars(vec!['a', 'b', 'c']), &mut c, sp())
11668 .unwrap();
11669 assert_eq!(r.dtype(), DType::Char);
11670 let e = dollar()
11672 .dyad(&Array::scalar_i64(2), &Array::empty(DType::I64), &mut c, sp())
11673 .unwrap_err();
11674 assert_eq!(e.kind, ErrorKind::Length);
11675 assert!(e.msg.contains("empty"), "{}", e.msg);
11676 let r = dollar()
11678 .dyad(&Array::scalar_i64(0), &Array::empty(DType::I64), &mut c, sp())
11679 .unwrap();
11680 assert_eq!(r.shape, vec![0]);
11681 let e = dollar()
11682 .dyad(&Array::scalar_i64(-1), &Array::from_i64(vec![1]), &mut c, sp())
11683 .unwrap_err();
11684 assert_eq!(e.kind, ErrorKind::Domain);
11685 }
11686
11687 #[test]
11688 fn take_from_both_ends_and_beyond() {
11689 ctx!(c);
11690 let v = Array::from_i64(vec![1, 2, 3, 4]);
11691 let take = |x: Array, y: &Array, c: &mut Ctx<'_>| head_v().dyad(&x, y, c, sp()).unwrap();
11692 assert_eq!(ints(&take(Array::scalar_i64(2), &v, &mut c)), vec![1, 2]);
11693 assert_eq!(ints(&take(Array::scalar_i64(-2), &v, &mut c)), vec![3, 4]);
11694 let short = Array::from_i64(vec![1, 2, 3]);
11696 assert_eq!(ints(&take(Array::scalar_i64(6), &short, &mut c)), vec![1, 2, 3, 0, 0, 0]);
11697 assert_eq!(ints(&take(Array::scalar_i64(-6), &short, &mut c)), vec![0, 0, 0, 1, 2, 3]);
11699 let r = take(Array::scalar_i64(2), &Array::scalar_i64(5), &mut c);
11701 assert_eq!(r.shape, vec![2]);
11702 assert_eq!(ints(&r), vec![5, 0]);
11703 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
11705 let r = take(Array::scalar_i64(1), &m, &mut c);
11706 assert_eq!(r.shape, vec![1, 3]);
11707 assert_eq!(ints(&r), vec![1, 2, 3]);
11708 let r = take(Array::scalar_i64(-1), &m, &mut c);
11709 assert_eq!(ints(&r), vec![4, 5, 6]);
11710 let r = take(Array::from_i64(vec![2, 2]), &m, &mut c);
11711 assert_eq!(r.shape, vec![2, 2]);
11712 assert_eq!(ints(&r), vec![1, 2, 4, 5]);
11713 let r = take(Array::from_i64(vec![3, -2]), &m, &mut c);
11714 assert_eq!(r.shape, vec![3, 2]);
11715 assert_eq!(ints(&r), vec![2, 3, 5, 6, 0, 0]);
11716 let r = head_v()
11718 .dyad(&Array::scalar_i64(3), &Array::from_chars(vec!['a']), &mut c, sp())
11719 .unwrap();
11720 assert_eq!(r.data, Data::Char(vec!['a', ' ', ' '].into()));
11721 let e = head_v()
11724 .dyad(&Array::from_i64(vec![1, 1]), &Array::from_i64(vec![1, 2]), &mut c, sp())
11725 .unwrap_err();
11726 assert_eq!(e.kind, ErrorKind::Length);
11727 let r = head_v()
11728 .dyad(&Array::from_i64(vec![1, 2]), &Array::scalar_i64(5), &mut c, sp())
11729 .unwrap();
11730 assert_eq!(r.shape, vec![1, 2]);
11731 assert_eq!(ints(&r), vec![5, 0]);
11732 }
11733
11734 #[test]
11735 fn drop_from_both_ends_and_beyond() {
11736 ctx!(c);
11737 let v = Array::from_i64(vec![1, 2, 3]);
11738 let drop = |x: Array, y: &Array, c: &mut Ctx<'_>| behead_v().dyad(&x, y, c, sp()).unwrap();
11739 assert_eq!(ints(&drop(Array::scalar_i64(1), &v, &mut c)), vec![2, 3]);
11740 assert_eq!(ints(&drop(Array::scalar_i64(-1), &v, &mut c)), vec![1, 2]);
11741 let r = drop(Array::scalar_i64(5), &v, &mut c);
11743 assert_eq!(r.shape, vec![0]);
11744 assert!(ints(&r).is_empty());
11745 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
11746 let r = drop(Array::scalar_i64(1), &m, &mut c);
11747 assert_eq!(r.shape, vec![1, 3]);
11748 assert_eq!(ints(&r), vec![4, 5, 6]);
11749 let r = drop(Array::from_i64(vec![0, -1]), &m, &mut c);
11750 assert_eq!(r.shape, vec![2, 2]);
11751 assert_eq!(ints(&r), vec![1, 2, 4, 5]);
11752 }
11753
11754 #[test]
11757 fn cells_of_unequal_shapes_are_padded_with_fills() {
11758 ctx!(c);
11759 let v = Verb::Rank(b(iota()), [0, 0, 0]);
11761 let r = v.monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp()).unwrap();
11762 assert_eq!(r.shape, vec![3, 3]);
11763 assert_eq!(ints(&r), vec![0, 0, 0, 0, 1, 0, 0, 1, 2]);
11764 }
11765
11766 #[test]
11767 fn framing_aligns_lower_rank_cells_at_the_trailing_axes() {
11768 let cells = vec![Array::from_i64(vec![1, 2]), mat(2, 2, vec![1, 2, 3, 4])];
11769 let r = assemble(&[2], cells, sp()).unwrap();
11770 assert_eq!(r.shape, vec![2, 2, 2]);
11771 assert_eq!(ints(&r), vec![1, 2, 0, 0, 1, 2, 3, 4]);
11772 }
11773
11774 #[test]
11775 fn framing_promotes_cell_types() {
11776 let cells = vec![Array::from_i64(vec![1]), Array::from_f64(vec![2.5])];
11777 let r = assemble(&[2], cells, sp()).unwrap();
11778 assert_eq!(r.dtype(), DType::F64);
11779 assert_eq!(floats(&r), vec![1.0, 2.5]);
11780 let cells = vec![Array::from_i64(vec![1]), Array::from_chars(vec!['a'])];
11782 let e = assemble(&[2], cells, sp()).unwrap_err();
11783 assert_eq!(e.kind, ErrorKind::Type);
11784 }
11785
11786 #[test]
11787 fn framing_over_an_empty_frame_yields_an_empty_result() {
11788 let r = assemble(&[0], Vec::new(), sp()).unwrap();
11789 assert_eq!(r.shape, vec![0]);
11790 assert_eq!(r.count(), 0);
11791 }
11792
11793 #[test]
11796 fn fork_applies_both_tines() {
11797 ctx!(c);
11798 let v = Verb::Fork(b(Verb::Reduce(b(plus()))), b(pct()), b(pound()));
11800 let r = v.monad(&Array::from_i64(vec![1, 2, 3, 4]), &mut c, sp()).unwrap();
11801 assert!(close(floats(&r)[0], 2.5));
11802 let v = Verb::Fork(b(minus()), b(plus()), b(plus()));
11804 let r = v
11805 .dyad(&Array::from_i64(vec![5]), &Array::from_i64(vec![3]), &mut c, sp())
11806 .unwrap();
11807 assert_eq!(ints(&r), vec![10]);
11808 }
11809
11810 #[test]
11811 fn noun_fork_supplies_a_constant_left_argument() {
11812 ctx!(c);
11813 let v = Verb::NounFork(Array::scalar_i64(10), b(minus()), b(right_v()));
11814 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
11815 assert_eq!(ints(&r), vec![9, 8]);
11816 let r = v
11817 .dyad(&Array::scalar_i64(0), &Array::from_i64(vec![1, 2]), &mut c, sp())
11818 .unwrap();
11819 assert_eq!(ints(&r), vec![9, 8]);
11820 }
11821
11822 #[test]
11823 fn hook_reuses_its_right_argument() {
11824 ctx!(c);
11825 let v = Verb::Hook(b(plus()), b(minus()));
11827 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
11828 assert_eq!(ints(&r), vec![0, 0]);
11829 let r = v
11831 .dyad(&Array::from_i64(vec![10]), &Array::from_i64(vec![3]), &mut c, sp())
11832 .unwrap();
11833 assert_eq!(ints(&r), vec![7]);
11834 }
11835
11836 #[test]
11837 fn atop_composes() {
11838 ctx!(c);
11839 let v = Verb::Atop(b(minus()), b(plus()));
11840 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
11841 assert_eq!(ints(&r), vec![-1, -2]);
11842 let r = v
11843 .dyad(&Array::from_i64(vec![1]), &Array::from_i64(vec![2]), &mut c, sp())
11844 .unwrap();
11845 assert_eq!(ints(&r), vec![-3]);
11846 }
11847
11848 #[test]
11849 fn trains_apply_to_the_whole_argument() {
11850 assert_eq!(Verb::Hook(b(plus()), b(minus())).ranks(), [RANK_INF; 3]);
11852 assert_eq!(Verb::Reduce(b(plus())).ranks(), [RANK_INF; 3]);
11853 }
11854
11855 #[test]
11858 fn absent_and_unwritten_meanings_are_reported_differently() {
11859 ctx!(c);
11860 let e = eq_v().monad(&Array::scalar_i64(1), &mut c, sp()).unwrap_err();
11861 assert_eq!(e.kind, ErrorKind::Domain);
11862 assert!(e.msg.contains("no monadic meaning"), "{}", e.msg);
11863 let e = not_v()
11864 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(1), &mut c, sp())
11865 .unwrap_err();
11866 assert_eq!(e.kind, ErrorKind::Domain);
11867 assert!(e.msg.contains("no dyadic meaning"), "{}", e.msg);
11868 let e = pound()
11869 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(1), &mut c, sp())
11870 .unwrap_err();
11871 assert_eq!(e.kind, ErrorKind::NotYet);
11872 assert!(e.msg.contains("copy"), "{}", e.msg);
11873 let _ = echo_v();
11875 }
11876
11877 fn seq_par<T: Send>(f: impl Fn() -> T + Sync + Send) -> (T, T) {
11886 (par::with_threads(1, &f), par::with_threads(4, &f))
11887 }
11888
11889 fn noise(n: usize) -> Vec<f64> {
11891 let mut x = 0x2545_f491_4f6c_dd1du64;
11892 (0..n)
11893 .map(|_| {
11894 x ^= x << 13;
11895 x ^= x >> 7;
11896 x ^= x << 17;
11897 (x >> 11) as f64 / (1u64 << 53) as f64 - 0.5
11898 })
11899 .collect()
11900 }
11901
11902 fn f64_mat(rows: usize, cols: usize) -> Array {
11903 Array::new(vec![rows, cols], Data::F64(noise(rows * cols).into()))
11904 }
11905
11906 const BIG: usize = 200_000;
11908
11909 #[test]
11910 fn an_elementwise_dyad_splits_into_the_same_result() {
11911 let x = Array::from_f64(noise(BIG));
11912 let y = Array::from_f64(noise(BIG).iter().map(|v| v + 0.25).collect());
11913 let (one, many) = seq_par(|| {
11914 ctx!(c);
11915 times().dyad(&x, &y, &mut c, sp()).unwrap()
11916 });
11917 assert_eq!(floats(&one), floats(&many));
11918 let (one, many) = seq_par(|| {
11920 ctx!(c);
11921 plus().dyad(&Array::scalar_f64(0.5), &y, &mut c, sp()).unwrap()
11922 });
11923 assert_eq!(floats(&one), floats(&many));
11924 }
11925
11926 #[test]
11927 fn an_elementwise_dyad_that_overflows_widens_the_same_way() {
11928 let mut v = vec![1i64; BIG];
11931 v[BIG - 3] = i64::MAX;
11932 let x = Array::from_i64(v);
11933 let (one, many) = seq_par(|| {
11934 ctx!(c);
11935 plus().dyad(&x, &x, &mut c, sp()).unwrap()
11936 });
11937 assert_eq!(one.dtype(), DType::F64);
11938 assert_eq!(floats(&one), floats(&many));
11939 }
11940
11941 #[test]
11942 fn an_elementwise_monad_splits_into_the_same_result() {
11943 let y = Array::from_f64(noise(BIG));
11944 for v in [minus(), sqrt_v(), floor_v(), pct()] {
11945 let (one, many) = seq_par(|| {
11946 ctx!(c);
11947 v.monad(&Array::from_f64(y.as_f64_slice().unwrap().iter().map(|x| x.abs()).collect()), &mut c, sp())
11948 .unwrap()
11949 });
11950 assert_eq!(one.data, many.data, "{}", v.name());
11951 }
11952 }
11953
11954 #[test]
11955 fn monadic_cells_run_in_parallel_and_frame_in_order() {
11956 let y = f64_mat(400, 512);
11959 let v = Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]);
11960 let (one, many) = seq_par(|| {
11961 ctx!(c);
11962 v.monad(&y, &mut c, sp()).unwrap()
11963 });
11964 assert_eq!(one.shape, vec![400]);
11965 assert_eq!(floats(&one), floats(&many));
11966 }
11967
11968 #[test]
11969 fn dyadic_cells_run_in_parallel_and_frame_in_order() {
11970 let x = f64_mat(400, 512);
11971 let y = f64_mat(400, 512);
11972 let v = Verb::Rank(b(plus()), [1, 1, 1]);
11974 let (one, many) = seq_par(|| {
11975 ctx!(c);
11976 v.dyad(&x, &y, &mut c, sp()).unwrap()
11977 });
11978 assert_eq!(one.shape, vec![400, 512]);
11979 assert_eq!(floats(&one), floats(&many));
11980 }
11981
11982 #[test]
11983 fn a_verb_that_writes_output_is_not_pure() {
11984 assert!(plus().is_pure());
11985 assert!(Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]).is_pure());
11986 assert!(!echo_v().is_pure());
11987 assert!(!Verb::Rank(b(Verb::Atop(b(echo_v()), b(plus()))), [1, 1, 1]).is_pure());
11988 }
11989
11990 #[test]
11991 fn an_impure_verb_keeps_its_cells_in_order() {
11992 let y = Array::new(vec![16, 8192], Data::I64((0..16 * 8192).collect::<Vec<i64>>().into()));
11995 let v = Verb::Rank(b(Verb::Atop(b(echo_v()), b(head_v()))), [1, 1, 1]);
11996 let mut seen: Vec<i64> = Vec::new();
11997 let mut sink = |s: &str| {
11998 if let Some(first) = s.split_whitespace().next() && let Ok(n) = first.parse::<i64>() {
11999 seen.push(n);
12000 }
12001 };
12002 let mut env = Env::new(Vec::new());
12003 let mut c = Ctx {
12004 cfg: EvalCfg {
12005 agreement: Agreement::LeadingPrefix,
12006 fmt: FmtOpts::J,
12007 tol: Tol::J,
12008 rules: Rules::default(),
12009 },
12010 out: &mut sink,
12011 inp: None,
12012 env: &mut env,
12013 device: None,
12014 };
12015 v.monad(&y, &mut c, sp()).unwrap();
12016 assert_eq!(seen, (0..16).map(|i| i * 8192).collect::<Vec<i64>>());
12017 }
12018
12019 #[test]
12020 fn a_wide_item_reduce_folds_every_column_in_order() {
12021 let y = f64_mat(300, 512);
12024 for v in [plus(), minus(), floor_v()] {
12025 let (one, many) = seq_par(|| {
12026 ctx!(c);
12027 Verb::Reduce(b(v.clone())).monad(&y, &mut c, sp()).unwrap()
12028 });
12029 assert_eq!(one.shape, vec![512]);
12030 assert_eq!(floats(&one), floats(&many), "{}", v.name());
12031 }
12032 }
12033
12034 #[test]
12035 fn a_wide_item_integer_reduce_is_exact() {
12036 let n = 300;
12037 let m = 512;
12038 let y = Array::new(
12039 vec![n, m],
12040 Data::I64((0..(n * m) as i64).map(|i| i % 977 - 400).collect::<Vec<i64>>().into()),
12041 );
12042 let (one, many) = seq_par(|| {
12043 ctx!(c);
12044 Verb::Reduce(b(minus())).monad(&y, &mut c, sp()).unwrap()
12045 });
12046 assert_eq!(ints(&one), ints(&many));
12047 }
12048
12049 #[test]
12050 fn a_narrow_item_reduce_chunks_the_items() {
12051 let y = f64_mat(300_000, 8);
12055 let (one, many) = seq_par(|| {
12056 ctx!(c);
12057 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
12058 });
12059 assert_eq!(one.shape, vec![8]);
12060 for (p, q) in floats(&one).iter().zip(floats(&many)) {
12061 assert!((p - q).abs() <= 1e-12 * p.abs().max(1.0), "{p} vs {q}");
12062 }
12063 let ints_y = Array::new(
12064 vec![300_000, 8],
12065 Data::I64((0..300_000 * 8).map(|i| (i % 101) as i64 - 50).collect::<Vec<i64>>().into()),
12066 );
12067 let (one, many) = seq_par(|| {
12068 ctx!(c);
12069 Verb::Reduce(b(plus())).monad(&ints_y, &mut c, sp()).unwrap()
12070 });
12071 assert_eq!(ints(&one), ints(&many));
12072 }
12073
12074 #[test]
12075 fn a_vector_reduce_folds_the_flat_buffer() {
12076 let y = Array::from_f64(noise(BIG * 4));
12077 let (one, many) = seq_par(|| {
12078 ctx!(c);
12079 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
12080 });
12081 let (p, q) = (floats(&one)[0], floats(&many)[0]);
12082 assert!((p - q).abs() <= 1e-12 * p.abs().max(1.0), "{p} vs {q}");
12083
12084 let ints_y = Array::from_i64((0..BIG as i64 * 4).map(|i| i % 1009 - 500).collect());
12087 for v in [plus(), minus(), ceil_v()] {
12088 let (one, many) = seq_par(|| {
12089 ctx!(c);
12090 Verb::Reduce(b(v.clone())).monad(&ints_y, &mut c, sp()).unwrap()
12091 });
12092 assert_eq!(ints(&one), ints(&many), "{}", v.name());
12093 }
12094 }
12095
12096 #[test]
12097 fn a_reduce_that_overflows_falls_back_to_the_sequential_widening() {
12098 let mut v: Vec<i64> = vec![1; BIG];
12099 v[7] = i64::MAX;
12100 let y = Array::from_i64(v);
12101 let (one, many) = seq_par(|| {
12102 ctx!(c);
12103 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
12104 });
12105 assert_eq!(one.dtype(), DType::F64);
12106 assert_eq!(floats(&one), floats(&many));
12107 }
12108
12109 #[test]
12110 fn a_boolean_reduce_matches_the_sequential_promotion() {
12111 let n = BIG;
12112 let y = Array::new(
12113 vec![n],
12114 Data::Bool((0..n).map(|i| (i % 3 == 0) as u8).collect::<Vec<u8>>().into()),
12115 );
12116 let (one, many) = seq_par(|| {
12117 ctx!(c);
12118 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
12119 });
12120 assert_eq!(one.dtype(), DType::I64);
12121 assert_eq!(ints(&one), ints(&many));
12122 assert_eq!(ints(&one)[0], n.div_ceil(3) as i64);
12123 }
12124}