1use std::collections::{HashMap, HashSet};
8use std::sync::Arc;
9
10use crate::array::{Array, Data};
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)]
148pub struct EvalCfg {
149 pub agreement: Agreement,
150 pub fmt: FmtOpts,
151 pub tol: Tol,
154 pub rules: Rules,
157}
158
159impl EvalCfg {
160 pub(crate) fn pure<R>(self, f: impl FnOnce(&mut Ctx<'_>) -> R) -> R {
165 let mut sink = |_: &str| debug_assert!(false, "a pure verb wrote to the output sink");
166 let mut env = Env::new(Vec::new());
167 f(&mut Ctx { cfg: self, out: &mut sink, env: &mut env, device: None })
168 }
169}
170
171pub const RECURSION_LIMIT: usize = 64;
180
181pub struct Env {
189 globals: HashMap<String, Array>,
190 frames: Vec<HashMap<String, Array>>,
191 running: Vec<std::sync::Arc<crate::ir::ExplicitDef>>,
194 verbs: HashMap<String, Verb>,
195 args: Vec<Array>,
196}
197
198impl Env {
199 pub fn new(args: Vec<Array>) -> Env {
200 Env {
201 globals: HashMap::new(),
202 frames: Vec::new(),
203 running: Vec::new(),
204 verbs: HashMap::new(),
205 args,
206 }
207 }
208
209 pub fn get(&self, name: &str) -> Option<Array> {
210 if let Some(frame) = self.frames.last() {
211 if let Some(v) = frame.get(name) {
212 return Some(v.clone());
213 }
214 }
215 self.globals.get(name).cloned()
216 }
217
218 pub fn assign(&mut self, name: String, value: Array, scope: crate::ir::Scope) {
219 if scope == crate::ir::Scope::LocalDefault && self.get(&name).is_some() {
220 return;
221 }
222 let target = match (scope, self.frames.last_mut()) {
223 (crate::ir::Scope::Local | crate::ir::Scope::LocalDefault, Some(frame)) => frame,
224 _ => &mut self.globals,
225 };
226 target.insert(name, value);
227 }
228
229 pub fn define(&mut self, name: String, verb: Verb) {
230 self.verbs.insert(name, verb);
231 }
232
233 pub fn undefine(&mut self, name: &str) {
234 self.verbs.remove(name);
235 }
236
237 pub fn verb(&self, name: &str) -> Option<&Verb> {
238 self.verbs.get(name)
239 }
240
241 pub fn arg(&self, i: usize) -> Result<Array> {
242 self.args
243 .get(i)
244 .cloned()
245 .ok_or_else(|| Error::internal("a parameter was read where none is bound"))
246 }
247
248 pub fn enter(
250 &mut self,
251 frame: HashMap<String, Array>,
252 def: std::sync::Arc<crate::ir::ExplicitDef>,
253 span: Span,
254 ) -> Result<()> {
255 if self.frames.len() >= RECURSION_LIMIT {
256 return Err(Error::new(
257 ErrorKind::Domain,
258 format!("explicit definitions called each other more than {RECURSION_LIMIT} deep"),
259 Some(span),
260 )
261 .note("a definition that recurses needs a case that stops"));
262 }
263 self.frames.push(frame);
264 self.running.push(def);
265 Ok(())
266 }
267
268 pub fn leave(&mut self) -> HashMap<String, Array> {
270 self.running.pop();
271 self.frames.pop().unwrap_or_default()
272 }
273
274 pub fn current_def(&self) -> Option<std::sync::Arc<crate::ir::ExplicitDef>> {
276 self.running.last().cloned()
277 }
278}
279
280pub struct Ctx<'a> {
282 pub cfg: EvalCfg,
283 pub out: &'a mut dyn FnMut(&str),
286 pub env: &'a mut Env,
288 pub device: Option<&'a crate::device::Device>,
291}
292
293const MAX_NESTING: usize = 400;
302
303thread_local! {
304 static NESTING: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
305}
306
307pub(crate) fn check_nesting(depth: usize, span: Span) -> Result<()> {
309 if depth > MAX_NESTING {
310 return Err(Error::new(
311 ErrorKind::Limit,
312 format!("this program nests more than {MAX_NESTING} applications deep"),
313 Some(span),
314 ));
315 }
316 Ok(())
317}
318
319pub(crate) struct Nesting;
321
322impl Nesting {
323 pub(crate) fn enter(span: Span) -> Result<Nesting> {
325 let depth = NESTING.with(|c| {
326 let d = c.get() + 1;
327 c.set(d);
328 d
329 });
330 if depth > MAX_NESTING {
331 NESTING.with(|c| c.set(c.get() - 1));
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(Nesting)
339 }
340}
341
342impl Drop for Nesting {
343 fn drop(&mut self) {
344 NESTING.with(|c| c.set(c.get().saturating_sub(1)));
345 }
346}
347
348impl Ctx<'_> {
349 fn with_tol<R>(&mut self, tol: Tol, f: impl FnOnce(&mut Ctx<'_>) -> R) -> R {
351 let cfg = EvalCfg { tol, ..self.cfg };
352 f(&mut Ctx { cfg, out: &mut *self.out, env: &mut *self.env, device: self.device })
353 }
354}
355
356#[derive(Clone, Copy, Debug, PartialEq, Eq)]
358pub enum ScalarMonad {
359 Conj,
361 Neg,
362 Signum,
363 Recip,
364 Sqrt,
365 Exp,
366 Abs,
367 Floor,
368 Ceil,
369 Not,
371 OneMinus,
373 Inc,
375 Dec,
377 Double,
379 Halve,
381 Square,
383 Ln,
385 Pi,
387 Factorial,
390 Imaginary,
392 Polar,
394}
395
396#[derive(Clone, Copy, Debug, PartialEq, Eq)]
398pub enum ScalarDyad {
399 Add,
400 Sub,
401 Mul,
402 DivJ,
404 DivApl,
406 Min,
407 Max,
408 Pow,
409 Residue,
411 Eq,
412 Ne,
413 Lt,
414 Le,
415 Gt,
416 Ge,
417 Lcm,
419 Gcd,
421 Log,
423 Root,
425 Circle,
429 Binomial,
432 MakeComplex,
434 PolarBy,
436}
437
438#[derive(Clone, Copy, Debug, PartialEq, Eq)]
440pub enum Enclose {
441 Always,
443 ExceptSimpleScalar,
445}
446
447#[derive(Clone, Copy, Debug, PartialEq, Eq)]
449pub enum MonadOp {
450 Scalar(ScalarMonad),
451 ShapeOf,
453 Tally,
455 Ravel,
457 TransposeAxes,
459 Head,
461 Behead,
463 Tail,
465 Curtail,
467 Reverse,
469 Nub,
471 GradeUp { origin: i64 },
473 GradeDown { origin: i64 },
475 IotaJ,
477 IotaApl { origin: i64 },
479 Echo,
481 Same,
483 Format,
487 DecodeBits,
489 EncodeBits,
493 Itemize,
495 TableOf,
498 Enclose(Enclose),
500 Open,
504 Raze,
506 First,
509 Enlist,
511 Depth,
514 Indices { origin: i64, boxed_coords: bool },
518 Steps,
520 ToExact,
523 NthPrime,
525 PrimeFactors,
527 MatrixInverse,
529 Roll { origin: i64, fixed: bool, float_at_zero: bool },
534 ComplexParts { polar: bool },
538 SelfClassify,
541 NubSieve,
544 Unicode { pass_chars: bool },
548 Words,
550 Nest,
553 LevelOf,
556 MapPaths,
559 PolyRoots,
563 PolyDeriv,
566 AnagramIndex,
568 CycleForm,
571 Split,
574 Execute { apl: bool },
579 NotYet(&'static str),
581 None,
583}
584
585#[derive(Clone, Copy, Debug, PartialEq, Eq)]
587pub enum DyadOp {
588 Scalar(ScalarDyad),
589 Reshape,
592 Take,
594 Drop,
596 Right,
598 Left,
600 Rotate,
602 AppendLeading,
604 AppendLast,
606 IndexOf { origin: i64 },
609 MemberJ,
611 MemberApl,
613 From,
616 Match,
618 NotMatch,
620 GradeSelect { down: bool },
622 Copy,
625 Decode,
628 Encode,
632 Laminate,
634 Link,
637 Strand,
639 IntervalIndex { offset: i64 },
643 IndexOfLast { origin: i64 },
645 MatrixDivide,
647 PartitionEnclose,
650 Squad { origin: i64 },
652 SelectAxis { axis: usize, rank: usize, origin: i64 },
656 Fetch,
658 PolyEval,
661 PolyIntegral,
664 FormatSpec,
667 TruthTable(u8),
671 ExactForm,
675 Deal { origin: i64, fixed: bool },
678 Boolean(BoolDyad),
681 Less,
683 Union,
685 Intersect,
687 AnagramFrom,
690 Permute,
693 FindSeq,
696 UnicodeForm,
699 PrimeMeta,
703 PrimeExponents,
706 Pick { origin: i64 },
708 Expand,
711 NotYet(&'static str),
712 None,
713}
714
715#[derive(Clone, Copy, Debug, PartialEq, Eq)]
717pub enum BoolDyad {
718 Nor,
720 Nand,
722}
723
724#[derive(Clone, Copy, Debug, PartialEq, Eq)]
727pub struct Prim {
728 pub name: &'static str,
729 pub monad: MonadOp,
730 pub dyad: DyadOp,
731 pub ranks: [i64; 3],
732}
733
734#[derive(Clone, Copy, Debug, PartialEq, Eq)]
738pub enum WindowKind {
739 Prefix,
742 Suffix,
745 Scan,
749}
750
751#[derive(Clone, Copy, Debug, PartialEq, Eq)]
753pub enum Power {
754 Times(u64),
756 Converge,
758}
759
760const CONVERGE_LIMIT: usize = 1 << 20;
762
763pub type MemoCache = Arc<std::sync::Mutex<HashMap<Vec<u64>, Array>>>;
767
768#[derive(Clone, Debug)]
772pub enum Verb {
773 Prim(Prim),
774 Rank(Box<Verb>, [i64; 3]),
776 Reduce(Box<Verb>),
778 Windowed(Box<Verb>, WindowKind),
781 Commute(Box<Verb>),
783 PowerN(Box<Verb>, Power),
785 Fork(Box<Verb>, Box<Verb>, Box<Verb>),
787 NounFork(Array, Box<Verb>, Box<Verb>),
789 Hook(Box<Verb>, Box<Verb>),
791 Atop(Box<Verb>, Box<Verb>),
793 Compose(Box<Verb>, Box<Verb>),
796 BondLeft(Array, Box<Verb>),
799 BondRight(Box<Verb>, Array),
801 Each(Box<Verb>, Enclose),
804 Fit(Box<Verb>, f64),
806 Amend(Array),
808 AmendVerb(Box<Verb>),
811 ShiftFill(Array),
814 Memo(Box<Verb>, MemoCache),
818 Level { u: Box<Verb>, level: i64, spread: bool },
822 Characteristics(Box<Verb>),
825 Before(Box<Verb>, Box<Verb>),
828 UserDerived { def: Box<Verb>, alpha: Box<Verb>, omega: Option<Box<Verb>> },
832 KeyPairs(Box<Verb>),
837 Key(Box<Verb>),
840 Cut(Box<Verb>, i64),
842 PowerV(Box<Verb>, Box<Verb>),
844 PowerUntil(Box<Verb>, Box<Verb>),
846 AlongAxis(Box<Verb>, usize),
850 Explicit(Arc<crate::ir::ExplicitDef>),
854 SelfRef,
857 Named(String),
860 WithObverse(Box<Verb>, Box<Verb>),
863 Agenda(Vec<Verb>, Box<Verb>),
866 Adverse(Box<Verb>, Box<Verb>),
870 Beside(Box<Verb>, Box<Verb>),
874}
875
876impl Verb {
877 pub fn ranks(&self) -> [i64; 3] {
879 match self {
880 Verb::Prim(p) => p.ranks,
881 Verb::Rank(_, r) => *r,
882 Verb::Windowed(_, WindowKind::Prefix) => [RANK_INF, 0, RANK_INF],
885 Verb::Each(..) => [0, 0, 0],
886 Verb::Fit(v, _) => v.ranks(),
887 Verb::Amend(_)
890 | Verb::AmendVerb(_)
891 | Verb::ShiftFill(_)
892 | Verb::Level { .. }
893 | Verb::Characteristics(_)
894 | Verb::UserDerived { .. }
895 | Verb::KeyPairs(_)
896 | Verb::Key(_)
897 | Verb::Cut(..)
898 | Verb::PowerV(..)
899 | Verb::PowerUntil(..)
900 | Verb::AlongAxis(..) => [RANK_INF, RANK_INF, RANK_INF],
901 Verb::Memo(v, _) => v.ranks(),
902 Verb::WithObverse(v, _) | Verb::Adverse(v, _) => v.ranks(),
903 Verb::Beside(..) => [RANK_INF, RANK_INF, RANK_INF],
904 _ => [RANK_INF, RANK_INF, RANK_INF],
905 }
906 }
907
908 pub fn name(&self) -> String {
910 match self {
911 Verb::Prim(p) => p.name.to_string(),
912 Verb::Rank(v, r) => format!("{}\"{}", v.name(), rank_str(*r)),
913 Verb::Reduce(v) => format!("{}/", v.name()),
914 Verb::Windowed(v, WindowKind::Suffix) => format!("{}\\.", v.name()),
915 Verb::Windowed(v, _) => format!("{}\\", v.name()),
916 Verb::Commute(v) => format!("{}~", v.name()),
917 Verb::PowerN(v, Power::Converge) => format!("{}^:_", v.name()),
918 Verb::PowerN(v, Power::Times(n)) => format!("{}^:{n}", v.name()),
919 Verb::Fork(f, g, h) => format!("({} {} {})", f.name(), g.name(), h.name()),
920 Verb::NounFork(_, g, h) => format!("(n {} {})", g.name(), h.name()),
921 Verb::Hook(f, g) => format!("({} {})", f.name(), g.name()),
922 Verb::Atop(f, g) => format!("({}@:{})", f.name(), g.name()),
923 Verb::Compose(f, g) => format!("({}&:{})", f.name(), g.name()),
924 Verb::BondLeft(_, v) => format!("(n&{})", v.name()),
925 Verb::BondRight(v, _) => format!("({}&n)", v.name()),
926 Verb::Each(v, Enclose::Always) => format!("({}&.>)", v.name()),
927 Verb::Each(v, _) => format!("({}¨)", v.name()),
928 Verb::Fit(v, n) => format!("{}!.{n}", v.name()),
929 Verb::Amend(_) => "(m})".to_string(),
930 Verb::AmendVerb(v) => format!("({}}})", v.name()),
931 Verb::ShiftFill(_) => "|.!.n".to_string(),
932 Verb::Characteristics(v) => format!("{} b.", v.name()),
933 Verb::Before(f, g) => format!("({}⍛{})", f.name(), g.name()),
934 Verb::KeyPairs(v) => format!("{}⌸", v.name()),
935 Verb::UserDerived { def, alpha, omega } => match omega {
936 Some(g) => format!("({} {} {})", alpha.name(), def.name(), g.name()),
937 None => format!("({} {})", alpha.name(), def.name()),
938 },
939 Verb::Memo(v, _) => format!("{} M.", v.name()),
940 Verb::Level { u, level, spread } => {
941 format!("{} {} {level}", u.name(), if *spread { "S:" } else { "L:" })
942 }
943 Verb::Key(v) => format!("{}/.", v.name()),
944 Verb::Cut(v, n) => format!("{};.{n}", v.name()),
945 Verb::PowerV(v, w) => format!("{}^:{}", v.name(), w.name()),
946 Verb::PowerUntil(v, w) => format!("{}⍣{}", v.name(), w.name()),
947 Verb::AlongAxis(v, k) => format!("{}[{k}]", v.name()),
948 Verb::Explicit(d) => d.name.clone(),
949 Verb::SelfRef => "$:".to_string(),
950 Verb::Named(n) => n.clone(),
951 Verb::WithObverse(v, w) => format!("({}:.{})", v.name(), w.name()),
952 Verb::Adverse(v, w) => format!("({}::{})", v.name(), w.name()),
953 Verb::Beside(f, g) => format!("({}∘{})", f.name(), g.name()),
954 Verb::Agenda(vs, w) => {
955 let names: Vec<String> = vs.iter().map(Verb::name).collect();
956 format!("({}@.{})", names.join("`"), w.name())
957 }
958 }
959 }
960
961 pub fn uses_tolerance(&self) -> bool {
966 match self {
967 Verb::Prim(p) => {
968 matches!(
969 p.monad,
970 MonadOp::Scalar(ScalarMonad::Floor)
971 | MonadOp::Scalar(ScalarMonad::Ceil)
972 | MonadOp::Nub
973 ) || matches!(
974 p.dyad,
975 DyadOp::Scalar(
976 ScalarDyad::Eq
977 | ScalarDyad::Ne
978 | ScalarDyad::Lt
979 | ScalarDyad::Le
980 | ScalarDyad::Gt
981 | ScalarDyad::Ge
982 ) | DyadOp::Match
983 | DyadOp::NotMatch
984 | DyadOp::MemberJ
985 | DyadOp::MemberApl
986 | DyadOp::IndexOf { .. }
987 | DyadOp::IndexOfLast { .. }
988 )
989 }
990 Verb::Rank(v, _)
991 | Verb::Reduce(v)
992 | Verb::Windowed(v, _)
993 | Verb::Commute(v)
994 | Verb::PowerN(v, _)
995 | Verb::BondLeft(_, v)
996 | Verb::BondRight(v, _)
997 | Verb::Each(v, _)
998 | Verb::Fit(v, _)
999 | Verb::Key(v)
1000 | Verb::Cut(v, _)
1001 | Verb::AlongAxis(v, _) => v.uses_tolerance(),
1002 Verb::PowerV(v, w) | Verb::PowerUntil(v, w) => {
1003 v.uses_tolerance() || w.uses_tolerance()
1004 }
1005 Verb::Amend(_)
1008 | Verb::AmendVerb(_)
1009 | Verb::ShiftFill(_)
1010 | Verb::Characteristics(_)
1011 | Verb::Explicit(_)
1012 | Verb::SelfRef
1013 | Verb::Named(_) => false,
1014 Verb::Memo(v, _) | Verb::Level { u: v, .. } => v.uses_tolerance(),
1015 Verb::WithObverse(v, _) => v.uses_tolerance(),
1016 Verb::Adverse(v, w) | Verb::Beside(v, w) | Verb::Before(v, w) => {
1017 v.uses_tolerance() || w.uses_tolerance()
1018 }
1019 Verb::KeyPairs(v) => v.uses_tolerance(),
1020 Verb::UserDerived { def, alpha, omega } => {
1021 def.uses_tolerance()
1022 || alpha.uses_tolerance()
1023 || omega.as_ref().is_some_and(|g| g.uses_tolerance())
1024 }
1025 Verb::Agenda(vs, w) => {
1026 w.uses_tolerance() || vs.iter().any(Verb::uses_tolerance)
1027 }
1028 Verb::Fork(f, g, h) => {
1029 f.uses_tolerance() || g.uses_tolerance() || h.uses_tolerance()
1030 }
1031 Verb::NounFork(_, g, h)
1032 | Verb::Hook(g, h)
1033 | Verb::Atop(g, h)
1034 | Verb::Compose(g, h) => g.uses_tolerance() || h.uses_tolerance(),
1035 }
1036 }
1037
1038 pub fn is_pure(&self) -> bool {
1043 match self {
1044 Verb::Prim(p) => {
1047 !matches!(p.monad, MonadOp::Echo | MonadOp::Roll { .. })
1048 && !matches!(p.dyad, DyadOp::Deal { .. })
1049 }
1050 Verb::Rank(v, _)
1051 | Verb::Reduce(v)
1052 | Verb::Windowed(v, _)
1053 | Verb::Commute(v)
1054 | Verb::PowerN(v, _) => v.is_pure(),
1055 Verb::Fork(f, g, h) => f.is_pure() && g.is_pure() && h.is_pure(),
1056 Verb::NounFork(_, g, h)
1057 | Verb::Hook(g, h)
1058 | Verb::Atop(g, h)
1059 | Verb::Compose(g, h) => g.is_pure() && h.is_pure(),
1060 Verb::BondLeft(_, v) | Verb::BondRight(v, _) | Verb::Each(v, _) | Verb::Fit(v, _) => {
1061 v.is_pure()
1062 }
1063 Verb::Key(v) | Verb::Cut(v, _) | Verb::AlongAxis(v, _) => v.is_pure(),
1064 Verb::PowerV(v, w) | Verb::PowerUntil(v, w) => v.is_pure() && w.is_pure(),
1065 Verb::WithObverse(v, _) => v.is_pure(),
1066 Verb::Adverse(v, w) | Verb::Beside(v, w) | Verb::Before(v, w) => {
1067 v.is_pure() && w.is_pure()
1068 }
1069 Verb::KeyPairs(v) => v.is_pure(),
1070 Verb::UserDerived { .. } => false,
1073 Verb::Agenda(vs, w) => w.is_pure() && vs.iter().all(Verb::is_pure),
1074 Verb::Amend(_) | Verb::ShiftFill(_) | Verb::Characteristics(_) => true,
1075 Verb::AmendVerb(v) | Verb::Level { u: v, .. } => v.is_pure(),
1076 Verb::Memo(v, _) => v.is_pure(),
1080 Verb::Explicit(_) | Verb::SelfRef | Verb::Named(_) => false,
1085 }
1086 }
1087
1088 pub fn monad(&self, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1090 let _depth = Nesting::enter(span)?;
1091 match self {
1092 Verb::Prim(p) => {
1093 if let MonadOp::Scalar(op) = p.monad {
1096 return scalar_monad(op, y, ctx.cfg.tol, span);
1097 }
1098 if p.monad == MonadOp::Open && is_mixed_simple(y) {
1102 return Ok(y.clone());
1103 }
1104 let frame_rank = y.rank() - effective_rank(p.ranks[0], y.rank());
1105 if frame_rank == 0 {
1106 return monad_op(p, y, ctx, span);
1107 }
1108 let frame = y.shape[..frame_rank].to_vec();
1109 let n: usize = frame.iter().product();
1110 let cells = each_cell(n, y.count(), self.is_pure(), ctx, |i, c| {
1111 monad_op(p, &y.cell_at(frame_rank, i), c, span)
1112 })?;
1113 assemble(&frame, cells, span)
1114 }
1115 Verb::Rank(v, r) => {
1116 let frame_rank = y.rank() - effective_rank(r[0], y.rank());
1117 if frame_rank == 0 {
1118 return v.monad(y, ctx, span);
1121 }
1122 if let Some(a) = reduce_vector_cells(v, y, frame_rank) {
1125 return Ok(a);
1126 }
1127 let frame = y.shape[..frame_rank].to_vec();
1128 let n: usize = frame.iter().product();
1129 let cells = each_cell(n, y.count(), self.is_pure(), ctx, |i, c| {
1130 v.monad(&y.cell_at(frame_rank, i), c, span)
1131 })?;
1132 assemble(&frame, cells, span)
1133 }
1134 Verb::Reduce(v) => reduce(v, y, ctx, span),
1135 Verb::Windowed(v, kind) => {
1136 runs(v, y, *kind == WindowKind::Suffix, ctx, span)
1137 }
1138 Verb::Commute(v) => v.dyad(y, y, ctx, span),
1139 Verb::PowerN(v, p) => power(v, *p, None, y, ctx, span),
1140 Verb::Fork(f, g, h) => {
1141 let l = f.monad(y, ctx, span)?;
1142 let r = h.monad(y, ctx, span)?;
1143 g.dyad(&l, &r, ctx, span)
1144 }
1145 Verb::NounFork(n, g, h) => {
1146 let r = h.monad(y, ctx, span)?;
1147 g.dyad(n, &r, ctx, span)
1148 }
1149 Verb::Hook(f, g) => {
1150 let r = g.monad(y, ctx, span)?;
1151 f.dyad(y, &r, ctx, span)
1152 }
1153 Verb::Atop(f, g) | Verb::Compose(f, g) => {
1154 let r = g.monad(y, ctx, span)?;
1155 f.monad(&r, ctx, span)
1156 }
1157 Verb::BondLeft(m, v) => v.dyad(m, y, ctx, span),
1158 Verb::BondRight(v, n) => v.dyad(y, n, ctx, span),
1159 Verb::Each(u, rule) => {
1160 let n = y.count();
1161 let cells = each_cell(n, n, self.is_pure(), ctx, |i, c| {
1162 let opened = open_cell(&atom(y, i));
1163 Ok(enclose(&u.monad(&opened, c, span)?, *rule))
1164 })?;
1165 assemble(&y.shape, cells, span)
1166 }
1167 Verb::Fit(v, n) => {
1168 let tol = Tol { ct: *n, ..ctx.cfg.tol };
1169 ctx.with_tol(tol, |c| v.monad(y, c, span))
1170 }
1171 Verb::Amend(m) => {
1173 if m.rank() != 0 || y.rank() > 1 {
1174 return Err(Error::new(
1175 ErrorKind::Rank,
1176 "selecting with m} takes one index into a list",
1177 Some(span),
1178 ));
1179 }
1180 from_index(m, y, span)
1181 }
1182 Verb::AmendVerb(u) => {
1184 let m = u.monad(y, ctx, span)?;
1185 Verb::Amend(m).monad(y, ctx, span)
1186 }
1187 Verb::ShiftFill(fill) => shift_fill(&Array::scalar_i64(-1), y, fill, span),
1190 Verb::Memo(u, cache) => memoised(u, cache, None, y, ctx, span),
1191 Verb::Characteristics(u) => characteristics(u, y, span),
1192 Verb::Before(f, g) => {
1193 let l = f.monad(y, ctx, span)?;
1194 g.dyad(&l, y, ctx, span)
1195 }
1196 Verb::KeyPairs(u) => key_pairs(u, y, None, ctx, span),
1197 Verb::UserDerived { def, alpha, omega } => {
1198 with_operands(alpha, omega.as_deref(), ctx, |c| def.monad(y, c, span))
1199 }
1200 Verb::Level { u, level, spread } => {
1201 at_level(u, *level, *spread, y, ctx, span)
1202 }
1203 Verb::Key(u) => oblique(u, y, ctx, span),
1204 Verb::Cut(u, n) => cut(u, None, y, *n, ctx, span),
1205 Verb::PowerV(u, v) => power_v(u, v, None, y, ctx, span),
1206 Verb::PowerUntil(u, v) => power_until(u, v, y, ctx, span),
1207 Verb::AlongAxis(u, k) => along_axis(u, None, y, *k, ctx, span),
1208 Verb::Explicit(d) => crate::ir::call_explicit(d, None, y, ctx, span),
1209 Verb::SelfRef => {
1210 let d = self_ref(ctx, span)?;
1211 crate::ir::call_explicit(&d, None, y, ctx, span)
1212 }
1213 Verb::Named(n) => named_verb(ctx, n, span)?.monad(y, ctx, span),
1214 Verb::WithObverse(v, _) => v.monad(y, ctx, span),
1215 Verb::Adverse(v, w) => match v.monad(y, ctx, span) {
1216 Err(e) if e.kind != ErrorKind::NotYet => w.monad(y, ctx, span),
1217 other => other,
1218 },
1219 Verb::Beside(f, g) => {
1220 let r = g.monad(y, ctx, span)?;
1221 f.monad(&r, ctx, span)
1222 }
1223 Verb::Agenda(vs, w) => {
1224 agenda_pick(vs, w, None, y, ctx, span)?.monad(y, ctx, span)
1225 }
1226 }
1227 }
1228
1229 pub fn dyad(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1231 let _depth = Nesting::enter(span)?;
1232 match self {
1233 Verb::Prim(_) | Verb::Rank(_, _) | Verb::Each(..) => {
1234 self.dyad_ranked(x, y, ctx, span)
1235 }
1236 Verb::Windowed(_, WindowKind::Prefix) => self.dyad_ranked(x, y, ctx, span),
1238 Verb::Windowed(u, WindowKind::Suffix) => outfix(u, x, y, ctx, span),
1241 Verb::Windowed(_, WindowKind::Scan) => {
1242 Err(Error::not_yet("dyadic scan (x f\\ y)", span))
1243 }
1244 Verb::Commute(v) => v.dyad(y, x, ctx, span),
1245 Verb::PowerN(v, p) => power(v, *p, Some(x), y, ctx, span),
1246 Verb::Reduce(v) => table(v, x, y, ctx, span),
1248 Verb::Fork(f, g, h) => {
1249 let l = f.dyad(x, y, ctx, span)?;
1250 let r = h.dyad(x, y, ctx, span)?;
1251 g.dyad(&l, &r, ctx, span)
1252 }
1253 Verb::NounFork(n, g, h) => {
1254 let r = h.dyad(x, y, ctx, span)?;
1255 g.dyad(n, &r, ctx, span)
1256 }
1257 Verb::Hook(f, g) => {
1258 let r = g.monad(y, ctx, span)?;
1259 f.dyad(x, &r, ctx, span)
1260 }
1261 Verb::Atop(f, g) => {
1262 let r = g.dyad(x, y, ctx, span)?;
1263 f.monad(&r, ctx, span)
1264 }
1265 Verb::Compose(f, g) => {
1266 let l = g.monad(x, ctx, span)?;
1267 let r = g.monad(y, ctx, span)?;
1268 f.dyad(&l, &r, ctx, span)
1269 }
1270 Verb::Fit(v, n) => {
1271 let tol = Tol { ct: *n, ..ctx.cfg.tol };
1272 ctx.with_tol(tol, |c| v.dyad(x, y, c, span))
1273 }
1274 Verb::Amend(m) => amend(m, x, y, span),
1275 Verb::AmendVerb(u) => {
1277 let m = u.dyad(x, y, ctx, span)?;
1278 amend(&m, x, y, span)
1279 }
1280 Verb::ShiftFill(fill) => shift_fill(x, y, fill, span),
1281 Verb::Memo(u, cache) => memoised(u, cache, Some(x), y, ctx, span),
1282 Verb::Characteristics(_) => {
1283 Err(Error::domain("u b. has no dyadic meaning", span))
1284 }
1285 Verb::Before(f, g) => {
1286 let l = f.monad(x, ctx, span)?;
1287 g.dyad(&l, y, ctx, span)
1288 }
1289 Verb::KeyPairs(u) => key_pairs(u, x, Some(y), ctx, span),
1290 Verb::UserDerived { def, alpha, omega } => {
1291 with_operands(alpha, omega.as_deref(), ctx, |c| def.dyad(x, y, c, span))
1292 }
1293 Verb::Level { .. } => {
1294 Err(Error::not_yet("a dyadic level or spread (x u L: n y)", span))
1295 }
1296 Verb::Key(u) => key(u, x, y, ctx, span),
1297 Verb::Cut(u, n) => cut(u, Some(x), y, *n, ctx, span),
1298 Verb::PowerV(u, v) => power_v(u, v, Some(x), y, ctx, span),
1299 Verb::PowerUntil(..) => {
1300 Err(Error::not_yet("dyadic power with a function operand (x f⍣g y)", span))
1301 }
1302 Verb::AlongAxis(u, k) => along_axis(u, Some(x), y, *k, ctx, span),
1303 Verb::Explicit(d) => crate::ir::call_explicit(d, Some(x), y, ctx, span),
1304 Verb::SelfRef => {
1305 let d = self_ref(ctx, span)?;
1306 crate::ir::call_explicit(&d, Some(x), y, ctx, span)
1307 }
1308 Verb::Named(n) => named_verb(ctx, n, span)?.dyad(x, y, ctx, span),
1309 Verb::WithObverse(v, _) => v.dyad(x, y, ctx, span),
1310 Verb::Adverse(v, w) => match v.dyad(x, y, ctx, span) {
1311 Err(e) if e.kind != ErrorKind::NotYet => w.dyad(x, y, ctx, span),
1312 other => other,
1313 },
1314 Verb::Beside(f, g) => {
1315 let r = g.monad(y, ctx, span)?;
1316 f.dyad(x, &r, ctx, span)
1317 }
1318 Verb::Agenda(vs, w) => {
1319 agenda_pick(vs, w, Some(x), y, ctx, span)?.dyad(x, y, ctx, span)
1320 }
1321 Verb::BondLeft(..) | Verb::BondRight(..) => {
1323 Err(Error::domain(format!("{} has no dyadic meaning", self.name()), span))
1324 }
1325 }
1326 }
1327
1328 fn dyad_ranked(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1330 let ranks = self.ranks();
1331 let er_l = effective_rank(ranks[1], x.rank());
1332 let er_r = effective_rank(ranks[2], y.rank());
1333 if er_l == 0 && er_r == 0 {
1334 if let Some(op) = self.scalar_dyad_op() {
1337 return scalar_dyad(op, x, y, ctx.cfg, span);
1338 }
1339 }
1340 let fxl = x.rank() - er_l;
1341 let fyl = y.rank() - er_r;
1342 let p = agree(&x.shape[..fxl], &y.shape[..fyl], &x.shape, &y.shape, ctx.cfg.agreement, span)?;
1343 if p.frame.is_empty() {
1344 return self.dyad_cell(x, y, ctx, span);
1345 }
1346 let work = x.count().max(y.count());
1347 let cells = each_cell(p.n, work, self.is_pure(), ctx, |i, c| {
1348 let xc = x.cell_at(fxl, i / p.x_div);
1349 let yc = y.cell_at(fyl, i / p.y_div);
1350 self.dyad_cell(&xc, &yc, c, span)
1351 })?;
1352 assemble(&p.frame, cells, span)
1353 }
1354
1355 fn dyad_cell(&self, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
1357 match self {
1358 Verb::Prim(p) => dyad_op(p, x, y, ctx.cfg, span),
1359 Verb::Rank(v, _) => v.dyad(x, y, ctx, span),
1360 Verb::Windowed(v, _) => infix(v, x, y, ctx, span),
1361 Verb::Each(u, rule) => {
1362 let r = u.dyad(&open_cell(x), &open_cell(y), ctx, span)?;
1363 Ok(enclose(&r, *rule))
1364 }
1365 _ => Err(Error::internal("dyad_cell on a verb without cell ranks")),
1366 }
1367 }
1368
1369 fn scalar_dyad_op(&self) -> Option<ScalarDyad> {
1372 match self {
1373 Verb::Prim(p) => match p.dyad {
1374 DyadOp::Scalar(op) => Some(op),
1375 _ => None,
1376 },
1377 Verb::Rank(v, _) => v.scalar_dyad_op(),
1378 _ => None,
1379 }
1380 }
1381}
1382
1383pub fn effective_rank(r: i64, arg_rank: usize) -> usize {
1386 if r >= 0 {
1387 (r as usize).min(arg_rank)
1388 } else {
1389 arg_rank.saturating_sub(r.unsigned_abs() as usize)
1390 }
1391}
1392
1393fn self_ref(ctx: &Ctx<'_>, span: Span) -> Result<Arc<crate::ir::ExplicitDef>> {
1403 ctx.env.current_def().ok_or_else(|| {
1404 Error::new(
1405 ErrorKind::Value,
1406 "self-reference outside an explicit definition",
1407 Some(span),
1408 )
1409 })
1410}
1411
1412fn named_verb(ctx: &Ctx<'_>, name: &str, span: Span) -> Result<Verb> {
1414 ctx.env.verb(name).cloned().ok_or_else(|| {
1415 Error::new(ErrorKind::Value, format!("undefined verb: {name}"), Some(span))
1416 })
1417}
1418
1419fn each_cell<F>(
1420 n: usize,
1421 work: usize,
1422 pure: bool,
1423 ctx: &mut Ctx<'_>,
1424 f: F,
1425) -> Result<Vec<Array>>
1426where
1427 F: Fn(usize, &mut Ctx<'_>) -> Result<Array> + Sync + Send,
1428{
1429 if pure && n > 1 && par::worth_it(work) {
1430 let cfg = ctx.cfg;
1431 return par::map_indexed(n, |i| cfg.pure(|c| f(i, c))).into_iter().collect();
1432 }
1433 (0..n).map(|i| f(i, ctx)).collect()
1434}
1435
1436fn one_rank(r: i64) -> String {
1439 if r == RANK_INF { "_".to_string() } else { r.to_string() }
1440}
1441
1442fn rank_str(r: [i64; 3]) -> String {
1445 if r[0] == r[1] && r[1] == r[2] {
1446 one_rank(r[0])
1447 } else {
1448 format!("{} {} {}", one_rank(r[0]), one_rank(r[1]), one_rank(r[2]))
1449 }
1450}
1451
1452fn show_shape(shape: &[usize]) -> String {
1454 if shape.is_empty() {
1455 return "(scalar)".to_string();
1456 }
1457 shape.iter().map(|n| n.to_string()).collect::<Vec<_>>().join(" ")
1458}
1459
1460fn strides(shape: &[usize]) -> Vec<usize> {
1464 let mut s = vec![1usize; shape.len()];
1465 for k in (0..shape.len().saturating_sub(1)).rev() {
1466 s[k] = s[k + 1] * shape[k + 1];
1467 }
1468 s
1469}
1470
1471fn odometer(coord: &mut [usize], shape: &[usize]) {
1473 for k in (0..coord.len()).rev() {
1474 coord[k] += 1;
1475 if coord[k] < shape[k] {
1476 return;
1477 }
1478 coord[k] = 0;
1479 }
1480}
1481
1482fn push_elem(dst: &mut Data, src: &Data, i: usize) {
1484 match (dst, src) {
1485 (Data::Bool(a), Data::Bool(b)) => a.push(b[i]),
1486 (Data::I64(a), Data::I64(b)) => a.push(b[i]),
1487 (Data::Ext(a), Data::Ext(b)) => a.push(b[i].clone()),
1488 (Data::Rat(a), Data::Rat(b)) => a.push(b[i].clone()),
1489 (Data::F64(a), Data::F64(b)) => a.push(b[i]),
1490 (Data::Complex(a), Data::Complex(b)) => a.push(b[i]),
1491 (Data::Char(a), Data::Char(b)) => a.push(b[i]),
1492 (Data::Box(a), Data::Box(b)) => a.push(b[i].clone()),
1493 _ => debug_assert!(false, "push_elem across dtypes"),
1494 }
1495}
1496
1497fn fill_data(dtype: DType, n: usize) -> Data {
1499 let mut d = Data::empty(dtype);
1500 for _ in 0..n {
1501 d.push_fill();
1502 }
1503 d
1504}
1505
1506struct Pairing {
1511 frame: Vec<usize>,
1512 n: usize,
1513 x_div: usize,
1514 y_div: usize,
1515}
1516
1517fn frame_mismatch(
1518 xs: &[usize],
1519 ys: &[usize],
1520 fx: &[usize],
1521 fy: &[usize],
1522 axis: usize,
1523 span: Span,
1524) -> Error {
1525 let kind = if fx.len() == 1 && fy.len() == 1 { ErrorKind::Length } else { ErrorKind::Shape };
1528 let note = if axis < fx.len() && axis < fy.len() {
1529 format!("frames first differ at axis {axis}: {} vs {}", fx[axis], fy[axis])
1530 } else {
1531 format!(
1532 "frames have different numbers of axes: {} vs {}, diverging at axis {axis}",
1533 fx.len(),
1534 fy.len()
1535 )
1536 };
1537 Error::new(
1538 kind,
1539 format!(
1540 "arguments do not agree: left shape {}, right shape {}",
1541 show_shape(xs),
1542 show_shape(ys)
1543 ),
1544 Some(span),
1545 )
1546 .note(note)
1547}
1548
1549fn agree(
1552 fx: &[usize],
1553 fy: &[usize],
1554 xs: &[usize],
1555 ys: &[usize],
1556 mode: Agreement,
1557 span: Span,
1558) -> Result<Pairing> {
1559 let common = fx.len().min(fy.len());
1560 match mode {
1561 Agreement::LeadingPrefix => {
1562 for i in 0..common {
1563 if fx[i] != fy[i] {
1564 return Err(frame_mismatch(xs, ys, fx, fy, i, span));
1565 }
1566 }
1567 let (long, short) = if fx.len() >= fy.len() { (fx, fy) } else { (fy, fx) };
1568 let n: usize = long.iter().product();
1569 let surplus: usize = long[short.len()..].iter().product();
1570 let (x_div, y_div) =
1571 if fx.len() >= fy.len() { (1, surplus.max(1)) } else { (surplus.max(1), 1) };
1572 Ok(Pairing { frame: long.to_vec(), n, x_div, y_div })
1573 }
1574 Agreement::ExactOrScalar => {
1575 if fx == fy {
1576 let n: usize = fx.iter().product();
1577 return Ok(Pairing { frame: fx.to_vec(), n, x_div: 1, y_div: 1 });
1578 }
1579 if fx.iter().product::<usize>() == 1 {
1582 let n: usize = fy.iter().product();
1583 return Ok(Pairing { frame: fy.to_vec(), n, x_div: n.max(1), y_div: 1 });
1584 }
1585 if fy.iter().product::<usize>() == 1 {
1586 let n: usize = fx.iter().product();
1587 return Ok(Pairing { frame: fx.to_vec(), n, x_div: 1, y_div: n.max(1) });
1588 }
1589 let axis = (0..common).find(|&i| fx[i] != fy[i]).unwrap_or(common);
1590 Err(frame_mismatch(xs, ys, fx, fy, axis, span))
1591 }
1592 }
1593}
1594
1595fn assemble(frame: &[usize], cells: Vec<Array>, span: Span) -> Result<Array> {
1599 if cells.is_empty() {
1600 return Ok(Array { shape: frame.to_vec(), data: Data::empty(DType::I64) });
1603 }
1604 let mut dt = cells[0].dtype();
1605 for c in &cells[1..] {
1606 dt = DType::promote(dt, c.dtype()).ok_or_else(|| {
1607 let boxed = dt == DType::Box || c.dtype() == DType::Box;
1608 let what = if boxed {
1609 "cannot frame boxed and unboxed results into one array"
1610 } else {
1611 "cannot frame character and numeric results into one array"
1612 };
1613 Error::new(ErrorKind::Type, what, Some(span))
1614 })?;
1615 }
1616 let widen = |c: &Array| -> Result<Data> {
1617 c.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening while framing"))
1618 };
1619
1620 if cells[1..].iter().all(|c| c.shape == cells[0].shape) {
1621 let mut data = Data::empty(dt);
1622 for c in &cells {
1623 if c.dtype() == dt {
1624 data.extend_from(&c.data);
1625 } else {
1626 data.extend_from(&widen(c)?);
1627 }
1628 }
1629 let mut shape = frame.to_vec();
1630 shape.extend_from_slice(&cells[0].shape);
1631 return Ok(Array::new(shape, data));
1632 }
1633
1634 let crank = cells.iter().map(|c| c.rank()).max().unwrap_or(0);
1637 let padded: Vec<Vec<usize>> = cells
1638 .iter()
1639 .map(|c| {
1640 let mut s = vec![1usize; crank - c.rank()];
1641 s.extend_from_slice(&c.shape);
1642 s
1643 })
1644 .collect();
1645 let mut common = vec![0usize; crank];
1646 for s in &padded {
1647 for k in 0..crank {
1648 common[k] = common[k].max(s[k]);
1649 }
1650 }
1651 let cell_n: usize = common.iter().product();
1652 let mut data = Data::empty(dt);
1653 for (c, ps) in cells.iter().zip(&padded) {
1654 let cd = if c.dtype() == dt { c.data.clone() } else { widen(c)? };
1655 let st = strides(ps);
1656 let mut coord = vec![0usize; crank];
1657 for _ in 0..cell_n {
1658 let mut idx = 0usize;
1659 let mut inside = true;
1660 for k in 0..crank {
1661 if coord[k] >= ps[k] {
1662 inside = false;
1663 break;
1664 }
1665 idx += coord[k] * st[k];
1666 }
1667 if inside {
1668 push_elem(&mut data, &cd, idx);
1669 } else {
1670 data.push_fill();
1671 }
1672 odometer(&mut coord, &common);
1673 }
1674 }
1675 let mut shape = frame.to_vec();
1676 shape.extend_from_slice(&common);
1677 Ok(Array::new(shape, data))
1678}
1679
1680fn atom(a: &Array, i: usize) -> Array {
1685 Array { shape: Vec::new(), data: a.data.slice(i, i + 1) }
1686}
1687
1688fn enclose(y: &Array, rule: Enclose) -> Array {
1690 if rule == Enclose::ExceptSimpleScalar && y.rank() == 0 && y.dtype() != DType::Box {
1691 return y.clone();
1692 }
1693 Array::boxed(y.clone())
1694}
1695
1696fn open_cell(y: &Array) -> Array {
1699 match &y.data {
1700 Data::Box(v) if !v.is_empty() => v[0].clone(),
1701 _ => y.clone(),
1702 }
1703}
1704
1705fn first(y: &Array) -> Array {
1708 if y.count() == 0 {
1709 let mut d = Data::empty(y.dtype());
1710 d.push_fill();
1711 return open_cell(&Array { shape: Vec::new(), data: d });
1712 }
1713 open_cell(&atom(y, 0))
1714}
1715
1716fn depth(y: &Array) -> i64 {
1718 match &y.data {
1719 Data::Box(v) => 1 + v.iter().map(depth).max().unwrap_or(0),
1720 _ => i64::from(y.rank() > 0),
1721 }
1722}
1723
1724fn leaves(a: &Array, out: &mut Vec<Array>) {
1726 match &a.data {
1727 Data::Box(v) => {
1728 for b in v.iter() {
1729 leaves(b, out);
1730 }
1731 }
1732 _ => out.push(a.clone()),
1733 }
1734}
1735
1736fn enlist(y: &Array, span: Span) -> Result<Array> {
1738 let mut parts = Vec::new();
1739 leaves(y, &mut parts);
1740 let mut dt = None;
1743 for p in parts.iter().filter(|p| p.count() > 0) {
1744 dt = Some(match dt {
1745 None => p.dtype(),
1746 Some(t) => DType::promote(t, p.dtype()).ok_or_else(|| {
1747 Error::new(
1748 ErrorKind::Type,
1749 "cannot enlist character and numeric data into one vector",
1750 Some(span),
1751 )
1752 })?,
1753 });
1754 }
1755 let dt = dt.unwrap_or(DType::I64);
1756 let mut data = Data::empty(dt);
1757 for p in &parts {
1758 let cast = p.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening in enlist"))?;
1759 data.extend_from(&cast);
1760 }
1761 Ok(Array::new(vec![data.len()], data))
1762}
1763
1764fn spread(a: &Array, shape: &[usize]) -> Array {
1766 let n: usize = shape.iter().product();
1767 let mut data = Data::empty(a.dtype());
1768 for _ in 0..n {
1769 push_elem(&mut data, &a.data, 0);
1770 }
1771 Array::new(shape.to_vec(), data)
1772}
1773
1774fn wider_shape(a: &[usize], b: &[usize]) -> Vec<usize> {
1777 let r = a.len().max(b.len());
1778 let pad = |s: &[usize]| {
1779 let mut v = vec![1usize; r - s.len()];
1780 v.extend_from_slice(s);
1781 v
1782 };
1783 let (pa, pb) = (pad(a), pad(b));
1784 (0..r).map(|k| pa[k].max(pb[k])).collect()
1785}
1786
1787fn raze(y: &Array, span: Span) -> Result<Array> {
1792 let opened: Vec<Array> = (0..y.count()).map(|i| open_cell(&atom(y, i))).collect();
1793 let mut common: Option<Vec<usize>> = None;
1794 for a in opened.iter().filter(|a| a.rank() > 0) {
1795 common = Some(match common {
1796 None => a.shape[1..].to_vec(),
1797 Some(c) => wider_shape(&c, &a.shape[1..]),
1798 });
1799 }
1800 let common = common.unwrap_or_default();
1801 let mut cells: Vec<Array> = Vec::new();
1802 for a in &opened {
1803 if a.rank() == 0 {
1804 cells.push(spread(a, &common));
1805 continue;
1806 }
1807 for i in 0..a.items() {
1808 cells.push(a.item(i));
1809 }
1810 }
1811 if cells.is_empty() {
1812 return Ok(Array::new(vec![0], Data::empty(DType::I64)));
1813 }
1814 let n = cells.len();
1815 assemble(&[n], cells, span)
1816}
1817
1818fn link(x: &Array, y: &Array, span: Span) -> Result<Array> {
1821 let head = Array::boxed(x.clone());
1822 let tail = if y.dtype() == DType::Box { y.clone() } else { Array::boxed(y.clone()) };
1823 catenate(&head, &tail, true, false, span)
1824}
1825
1826fn nest_like(a: &Array, other: &Array) -> Array {
1829 if a.dtype() == DType::Box || other.dtype() != DType::Box {
1830 return a.clone();
1831 }
1832 let cells: Vec<Array> = (0..a.count()).map(|i| atom(a, i)).collect();
1833 Array::new(a.shape.clone(), Data::Box(cells.into()))
1834}
1835
1836fn box_items(y: &Array) -> Array {
1838 if y.dtype() == DType::Box {
1839 return y.clone();
1840 }
1841 let n = y.items();
1842 let boxes: Vec<Array> = (0..n).map(|i| item_or_self(y, i)).collect();
1843 Array::new(vec![n], Data::Box(boxes.into()))
1844}
1845
1846fn strand(x: &Array, y: &Array, span: Span) -> Result<Array> {
1850 let item = enclose(x, Enclose::ExceptSimpleScalar);
1851 let one = |a: &Array| Array::new(vec![1], a.data.clone());
1852 if item.dtype() != DType::Box
1857 && y.dtype() != DType::Box
1858 && DType::promote(item.dtype(), y.dtype()).is_some()
1859 {
1860 return catenate(&one(&item), y, true, false, span);
1861 }
1862 let head = if item.dtype() == DType::Box { item } else { Array::boxed(item) };
1863 catenate(&one(&head), &box_items(y), true, false, span)
1864}
1865
1866fn char_arith(span: Span) -> Error {
1869 Error::new(ErrorKind::Type, "cannot do arithmetic on characters", Some(span))
1870}
1871
1872fn box_arith(span: Span) -> Error {
1873 Error::new(
1874 ErrorKind::Type,
1875 "cannot do arithmetic on boxed values; open them first (J `>`, APL `⊃`)",
1876 Some(span),
1877 )
1878}
1879
1880fn wrong_type(d: DType, span: Span) -> Error {
1883 match d {
1884 DType::Box => box_arith(span),
1885 _ => char_arith(span),
1886 }
1887}
1888
1889fn borrow_i64<'a>(d: &'a Data, tmp: &'a mut Vec<i64>) -> &'a [i64] {
1895 match d {
1896 Data::I64(v) => v,
1897 Data::Bool(v) => {
1898 *tmp = par::map(v, |&b| b as i64);
1899 &tmp[..]
1900 }
1901 _ => &[],
1903 }
1904}
1905
1906fn borrow_f64<'a>(d: &'a Data, tmp: &'a mut Vec<f64>) -> &'a [f64] {
1908 match d {
1909 Data::F64(v) => v,
1910 Data::I64(v) => {
1911 *tmp = par::map(v, |&x| x as f64);
1912 &tmp[..]
1913 }
1914 Data::Bool(v) => {
1915 *tmp = par::map(v, |&x| x as f64);
1916 &tmp[..]
1917 }
1918 Data::Ext(v) => {
1919 *tmp = par::map(v, exact::ext_to_f64);
1920 &tmp[..]
1921 }
1922 Data::Rat(v) => {
1923 *tmp = par::map(v, Rat::to_f64);
1924 &tmp[..]
1925 }
1926 _ => &[],
1927 }
1928}
1929
1930fn borrow_cx<'a>(d: &'a Data, tmp: &'a mut Vec<Cx>) -> &'a [Cx] {
1932 match d {
1933 Data::Complex(v) => v,
1934 Data::Ext(v) => {
1935 *tmp = par::map(v, |x| [exact::ext_to_f64(x), 0.0]);
1936 &tmp[..]
1937 }
1938 Data::Rat(v) => {
1939 *tmp = par::map(v, |x| [x.to_f64(), 0.0]);
1940 &tmp[..]
1941 }
1942 Data::F64(v) => {
1943 *tmp = par::map(v, |&x| [x, 0.0]);
1944 &tmp[..]
1945 }
1946 Data::I64(v) => {
1947 *tmp = par::map(v, |&x| [x as f64, 0.0]);
1948 &tmp[..]
1949 }
1950 Data::Bool(v) => {
1951 *tmp = v.iter().map(|&x| [x as f64, 0.0]).collect();
1952 &tmp[..]
1953 }
1954 _ => &[],
1955 }
1956}
1957
1958fn as_f64<'a>(d: &'a Data, tmp: &'a mut Vec<f64>, span: Span) -> Result<&'a [f64]> {
1960 if !d.dtype().is_numeric() {
1961 return Err(wrong_type(d.dtype(), span));
1962 }
1963 Ok(borrow_f64(d, tmp))
1964}
1965
1966fn arith_type(a: DType, b: DType, span: Span) -> Result<DType> {
1968 if a == DType::Box || b == DType::Box {
1969 return Err(box_arith(span));
1970 }
1971 match DType::promote(a, b) {
1972 Some(DType::Char) => Err(char_arith(span)),
1973 None => Err(Error::new(
1974 ErrorKind::Type,
1975 "cannot mix character and numeric data",
1976 Some(span),
1977 )),
1978 Some(DType::Bool) => Ok(DType::I64),
1979 Some(t) => Ok(t),
1980 }
1981}
1982
1983#[allow(clippy::too_many_arguments)]
1993#[inline]
1994fn zip_chunk<T, U, F>(
1995 xs: &[T],
1996 xoff: usize,
1997 xdiv: usize,
1998 ys: &[T],
1999 yoff: usize,
2000 ydiv: usize,
2001 start: usize,
2002 out: &mut [U],
2003 mut f: F,
2004) -> bool
2005where
2006 T: Copy,
2007 F: FnMut(T, T, &mut U) -> bool,
2008{
2009 let len = out.len();
2010 if len == 0 {
2011 return true;
2012 }
2013 let last = start + len - 1;
2014 let one_x = xdiv > 1 && start / xdiv == last / xdiv;
2015 let one_y = ydiv > 1 && start / ydiv == last / ydiv;
2016 if xdiv == 1 && ydiv == 1 {
2017 let xc = &xs[xoff + start..xoff + start + len];
2018 let yc = &ys[yoff + start..yoff + start + len];
2019 for ((slot, &a), &b) in out.iter_mut().zip(xc).zip(yc) {
2020 if !f(a, b, slot) {
2021 return false;
2022 }
2023 }
2024 } else if xdiv == 1 && one_y {
2025 let b = ys[yoff + start / ydiv];
2026 let xc = &xs[xoff + start..xoff + start + len];
2027 for (slot, &a) in out.iter_mut().zip(xc) {
2028 if !f(a, b, slot) {
2029 return false;
2030 }
2031 }
2032 } else if one_x && ydiv == 1 {
2033 let a = xs[xoff + start / xdiv];
2034 let yc = &ys[yoff + start..yoff + start + len];
2035 for (slot, &b) in out.iter_mut().zip(yc) {
2036 if !f(a, b, slot) {
2037 return false;
2038 }
2039 }
2040 } else {
2041 for (k, slot) in out.iter_mut().enumerate() {
2042 let i = start + k;
2043 if !f(xs[xoff + i / xdiv], ys[yoff + i / ydiv], slot) {
2044 return false;
2045 }
2046 }
2047 }
2048 true
2049}
2050
2051const LANCZOS: [f64; 9] = [
2055 0.999_999_999_999_809_9,
2056 676.520_368_121_885_1,
2057 -1_259.139_216_722_402_8,
2058 771.323_428_777_653_1,
2059 -176.615_029_162_140_6,
2060 12.507_343_278_686_905,
2061 -0.138_571_095_265_720_12,
2062 9.984_369_578_019_572e-6,
2063 1.505_632_735_149_311_6e-7,
2064];
2065
2066fn gamma(x: f64) -> f64 {
2070 use std::f64::consts::PI;
2071 if x < 0.5 {
2072 return PI / ((PI * x).sin() * gamma(1.0 - x));
2074 }
2075 let z = x - 1.0;
2076 let mut a = LANCZOS[0];
2077 for (i, &c) in LANCZOS.iter().enumerate().skip(1) {
2078 a += c / (z + i as f64);
2079 }
2080 let t = z + 7.5;
2081 (2.0 * PI).sqrt() * t.powf(z + 0.5) * (-t).exp() * a
2082}
2083
2084fn factorial(y: f64) -> f64 {
2087 if y.fract() == 0.0 && y.abs() < 1e17 {
2088 let n = y as i64;
2089 if n < 0 {
2090 return if n % 2 == -1 { f64::INFINITY } else { f64::NEG_INFINITY };
2092 }
2093 if n > 170 {
2094 return f64::INFINITY;
2095 }
2096 let mut c = 1.0f64;
2097 for i in 2..=n {
2098 c *= i as f64;
2099 }
2100 return c;
2101 }
2102 gamma(y + 1.0)
2103}
2104
2105const BINOMIAL_PRODUCT_LIMIT: i64 = 4096;
2108
2109fn binomial_product(x: i64, y: f64) -> f64 {
2113 let mut c = 1.0f64;
2114 for i in 1..=x {
2115 c = c * (y - i as f64 + 1.0) / i as f64;
2116 if c == 0.0 {
2117 break;
2118 }
2119 }
2120 c
2121}
2122
2123fn binomial_i64(x: i64, y: i64) -> Option<i64> {
2127 if x < 0 {
2128 if y >= 0 || y < x {
2131 return Some(0);
2132 }
2133 let v = binomial_exact(-y - 1, -x - 1)?;
2134 return if (y - x) % 2 == 0 { Some(v) } else { v.checked_neg() };
2135 }
2136 binomial_exact(x, y)
2137}
2138
2139fn binomial_exact(x: i64, y: i64) -> Option<i64> {
2142 if x > BINOMIAL_PRODUCT_LIMIT {
2143 return None;
2144 }
2145 let mut c: i128 = 1;
2146 for i in 1..=x as i128 {
2147 c = c.checked_mul(y as i128 - i + 1)? / i;
2148 if c == 0 {
2149 break;
2150 }
2151 }
2152 i64::try_from(c).ok()
2153}
2154
2155fn binomial(x: f64, y: f64) -> f64 {
2157 if x.fract() == 0.0 && x.abs() < 1e17 {
2158 let xi = x as i64;
2159 if xi < 0 {
2160 if y.fract() == 0.0 && y < 0.0 && y >= x {
2161 let sign = if (y as i64 - xi) % 2 == 0 { 1.0 } else { -1.0 };
2162 return sign * binomial_product(-y as i64 - 1, -x - 1.0);
2163 }
2164 return 0.0;
2165 }
2166 if xi <= BINOMIAL_PRODUCT_LIMIT {
2167 return binomial_product(xi, y);
2168 }
2169 }
2170 gamma(y + 1.0) / (gamma(x + 1.0) * gamma(y - x + 1.0))
2171}
2172
2173#[inline]
2176fn i64_op(op: ScalarDyad, a: i64, b: i64) -> Option<i64> {
2177 use ScalarDyad::*;
2178 Some(match op {
2179 Add => a.checked_add(b)?,
2180 Sub => a.checked_sub(b)?,
2181 Mul => a.checked_mul(b)?,
2182 Min => a.min(b),
2183 Max => a.max(b),
2184 Residue => {
2185 if a == 0 {
2186 b
2187 } else {
2188 let mut r = b.wrapping_rem(a);
2190 if r != 0 && (r < 0) != (a < 0) {
2191 r += a;
2192 }
2193 r
2194 }
2195 }
2196 Pow => {
2197 if b < 0 {
2198 return None;
2199 }
2200 a.checked_pow(u32::try_from(b).ok()?)?
2201 }
2202 Binomial => binomial_i64(a, b)?,
2203 _ => return None,
2204 })
2205}
2206
2207#[inline]
2209fn f64_op(op: ScalarDyad, a: f64, b: f64, span: Span) -> Result<f64> {
2210 use ScalarDyad::*;
2211 Ok(match op {
2212 Add => a + b,
2213 Sub => a - b,
2214 Mul => a * b,
2215 Min => a.min(b),
2216 Max => a.max(b),
2217 DivJ => {
2218 if b == 0.0 {
2219 if a == 0.0 { 0.0 } else { f64::INFINITY.copysign(a) }
2220 } else {
2221 a / b
2222 }
2223 }
2224 DivApl => {
2225 if b == 0.0 {
2226 if a == 0.0 {
2227 1.0
2228 } else {
2229 return Err(Error::domain("division by zero", span));
2230 }
2231 } else {
2232 a / b
2233 }
2234 }
2235 Pow => {
2236 if a == 0.0 && b == 0.0 {
2237 1.0
2238 } else {
2239 a.powf(b)
2240 }
2241 }
2242 Residue => {
2243 if a.is_infinite() {
2248 if b == 0.0 || (b > 0.0) == (a > 0.0) { b } else { a }
2249 } else if a == 0.0 {
2250 b
2251 } else {
2252 b - a * (b / a).floor()
2253 }
2254 }
2255 Log => {
2256 if a < 0.0 || b < 0.0 {
2257 return Err(Error::not_yet("complex numbers", span));
2258 }
2259 b.ln() / a.ln()
2260 }
2261 Root => {
2262 if b < 0.0 {
2263 return Err(Error::not_yet("complex numbers", span));
2264 }
2265 b.powf(1.0 / a)
2266 }
2267 Circle => return circle(a, b, span),
2268 Binomial => binomial(a, b),
2269 _ => return Err(Error::internal("non-arithmetic op in the float path")),
2270 })
2271}
2272
2273#[inline]
2277fn escapes_reals(op: ScalarDyad, a: f64, b: f64) -> bool {
2278 use ScalarDyad::*;
2279 match op {
2280 Pow => a < 0.0 && b.fract() != 0.0,
2282 Log => a < 0.0 || b < 0.0,
2283 Root => b < 0.0,
2284 Circle => circle_escapes(a, b),
2285 _ => false,
2286 }
2287}
2288
2289#[inline]
2292fn circle_escapes(k: f64, y: f64) -> bool {
2293 if k.fract() != 0.0 {
2294 return false;
2295 }
2296 match k as i64 {
2297 0 | -1 | -2 | -7 => y.abs() > 1.0,
2298 -4 => y.abs() < 1.0,
2299 -6 => y < 1.0,
2300 8 | -8 | -11 | -12 => true,
2303 _ => false,
2304 }
2305}
2306
2307#[inline]
2317fn circle(k: f64, y: f64, span: Span) -> Result<f64> {
2318 if k.fract() != 0.0 {
2319 return Err(Error::domain("the circle function needs an integer left argument", span));
2320 }
2321 let complex = || Error::internal("a circle function left the reals on the real path");
2322 Ok(match k as i64 {
2323 0 => {
2324 if y.abs() > 1.0 {
2325 return Err(complex());
2326 }
2327 (1.0 - y * y).max(0.0).sqrt()
2328 }
2329 1 => y.sin(),
2330 2 => y.cos(),
2331 3 => y.tan(),
2332 4 => (1.0 + y * y).sqrt(),
2333 5 => y.sinh(),
2334 6 => y.cosh(),
2335 7 => y.tanh(),
2336 -1 => {
2337 if y.abs() > 1.0 {
2338 return Err(complex());
2339 }
2340 y.asin()
2341 }
2342 -2 => {
2343 if y.abs() > 1.0 {
2344 return Err(complex());
2345 }
2346 y.acos()
2347 }
2348 -3 => y.atan(),
2349 -4 => {
2350 if y.abs() < 1.0 {
2351 return Err(complex());
2352 }
2353 y.signum() * (y * y - 1.0).max(0.0).sqrt()
2355 }
2356 -5 => y.asinh(),
2357 -6 => {
2358 if y < 1.0 {
2359 return Err(complex());
2360 }
2361 y.acosh()
2362 }
2363 -7 => {
2364 if y.abs() > 1.0 {
2365 return Err(complex());
2366 }
2367 y.atanh()
2368 }
2369 9 | -9 | -10 => y,
2371 10 => y.abs(),
2372 11 => 0.0,
2373 12 => {
2374 if y < 0.0 {
2375 std::f64::consts::PI
2376 } else {
2377 0.0
2378 }
2379 }
2380 8 | -8 | -11 | -12 => return Err(complex()),
2381 _ => {
2382 return Err(Error::domain(
2383 "the circle functions run from _12 to 12",
2384 span,
2385 ));
2386 }
2387 })
2388}
2389
2390#[inline]
2392fn cx_op(op: ScalarDyad, a: Cx, b: Cx, span: Span) -> Result<Cx> {
2393 use ScalarDyad::*;
2394 Ok(match op {
2395 Add => cx::add(a, b),
2396 Sub => cx::sub(a, b),
2397 Mul => cx::mul(a, b),
2398 DivJ => cx::div(a, b),
2399 DivApl => {
2400 if b == cx::ZERO {
2401 if a == cx::ZERO {
2402 cx::ONE
2403 } else {
2404 return Err(Error::domain("division by zero", span));
2405 }
2406 } else {
2407 cx::div(a, b)
2408 }
2409 }
2410 Pow => cx::pow(a, b),
2411 Log => cx::log(a, b),
2412 Root => cx::root(a, b),
2413 Residue => cx::residue(a, b),
2414 Lcm => cx::lcm(a, b),
2415 Gcd => cx::gcd(a, b),
2416 MakeComplex => cx::add(a, cx::mul(cx::I, b)),
2417 PolarBy => cx::mul(a, cx::exp(cx::mul(cx::I, b))),
2418 Circle => {
2419 if a[1] != 0.0 || a[0].fract() != 0.0 {
2420 return Err(Error::domain(
2421 "the circle function needs an integer left argument",
2422 span,
2423 ));
2424 }
2425 cx::circle(a[0] as i64, b).ok_or_else(|| {
2426 Error::domain("the circle functions run from _12 to 12", span)
2427 })?
2428 }
2429 Min | Max => return Err(no_complex_order(span)),
2430 Binomial => {
2431 return Err(Error::not_yet("the binomial function on complex numbers", span));
2432 }
2433 Eq | Ne | Lt | Le | Gt | Ge => {
2434 return Err(Error::internal("a comparison in the complex arithmetic path"));
2435 }
2436 })
2437}
2438
2439fn no_complex_order(span: Span) -> Error {
2442 Error::new(
2443 ErrorKind::Domain,
2444 "complex numbers have no order; only equality (=, ~:) applies to them",
2445 Some(span),
2446 )
2447}
2448
2449#[allow(clippy::too_many_arguments)]
2450#[inline(always)]
2451fn dyad_cx_chunk(
2452 op: ScalarDyad,
2453 xs: &[Cx],
2454 xoff: usize,
2455 xdiv: usize,
2456 ys: &[Cx],
2457 yoff: usize,
2458 ydiv: usize,
2459 start: usize,
2460 out: &mut [Cx],
2461 span: Span,
2462) -> Result<()> {
2463 use ScalarDyad::*;
2464 macro_rules! plain {
2467 ($step:expr) => {{
2468 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut Cx| {
2469 *slot = $step(a, b);
2470 true
2471 });
2472 return Ok(());
2473 }};
2474 }
2475 match op {
2476 Add => plain!(cx::add),
2477 Sub => plain!(cx::sub),
2478 Mul => plain!(cx::mul),
2479 DivJ => plain!(cx::div),
2480 _ => {}
2481 }
2482 let mut err = None;
2483 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut Cx| {
2484 match cx_op(op, a, b, span) {
2485 Ok(v) => {
2486 *slot = v;
2487 true
2488 }
2489 Err(e) => {
2490 err = Some(e);
2491 false
2492 }
2493 }
2494 });
2495 match err {
2496 Some(e) => Err(e),
2497 None => Ok(()),
2498 }
2499}
2500
2501#[allow(clippy::too_many_arguments)]
2502fn dyad_cx(
2503 op: ScalarDyad,
2504 xs: &[Cx],
2505 xoff: usize,
2506 xdiv: usize,
2507 ys: &[Cx],
2508 yoff: usize,
2509 ydiv: usize,
2510 n: usize,
2511 span: Span,
2512) -> Result<Vec<Cx>> {
2513 par::try_fill(n, |start, part| {
2514 dyad_cx_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part, span)
2515 })
2516}
2517
2518#[allow(clippy::too_many_arguments)]
2520fn complex_dyad_data(
2521 op: ScalarDyad,
2522 x: &Data,
2523 xoff: usize,
2524 xdiv: usize,
2525 y: &Data,
2526 yoff: usize,
2527 ydiv: usize,
2528 n: usize,
2529 span: Span,
2530) -> Result<Data> {
2531 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2532 let xs = borrow_cx(x, &mut tx);
2533 let ys = borrow_cx(y, &mut ty);
2534 Ok(Data::Complex(dyad_cx(op, xs, xoff, xdiv, ys, yoff, ydiv, n, span)?.into()))
2535}
2536
2537fn circle_reads_a_part(x: &Data, xoff: usize, xdiv: usize, n: usize) -> bool {
2542 if x.dtype() == DType::Complex {
2543 return false;
2545 }
2546 let mut tmp = Vec::new();
2547 let xs = borrow_f64(x, &mut tmp);
2548 (0..n).all(|i| {
2549 let k = xs[xoff + i / xdiv];
2550 k.fract() == 0.0 && (9.0..=12.0).contains(&k)
2551 })
2552}
2553
2554#[allow(clippy::too_many_arguments)]
2557fn pass_leaves_reals(
2558 op: ScalarDyad,
2559 x: &Data,
2560 xoff: usize,
2561 xdiv: usize,
2562 y: &Data,
2563 yoff: usize,
2564 ydiv: usize,
2565 n: usize,
2566) -> bool {
2567 use ScalarDyad::*;
2568 if !matches!(op, Pow | Log | Root | Circle) {
2569 return false;
2570 }
2571 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2572 let xs = borrow_f64(x, &mut tx);
2573 let ys = borrow_f64(y, &mut ty);
2574 (0..n).any(|i| escapes_reals(op, xs[xoff + i / xdiv], ys[yoff + i / ydiv]))
2575}
2576
2577#[allow(clippy::too_many_arguments)]
2578#[inline(always)]
2579fn dyad_i64_chunk_body(
2580 op: ScalarDyad,
2581 xs: &[i64],
2582 xoff: usize,
2583 xdiv: usize,
2584 ys: &[i64],
2585 yoff: usize,
2586 ydiv: usize,
2587 start: usize,
2588 out: &mut [i64],
2589) -> bool {
2590 use ScalarDyad::*;
2591 macro_rules! overflowing {
2595 ($m:ident) => {{
2596 let mut over = false;
2597 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
2598 let (v, o) = i64::$m(a, b);
2599 *slot = v;
2600 over |= o;
2601 true
2602 });
2603 !over
2604 }};
2605 }
2606 macro_rules! plain {
2607 ($step:expr) => {{
2608 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
2609 *slot = $step(a, b);
2610 true
2611 })
2612 }};
2613 }
2614 match op {
2615 Add => overflowing!(overflowing_add),
2616 Sub => overflowing!(overflowing_sub),
2617 Mul => overflowing!(overflowing_mul),
2618 Min => plain!(i64::min),
2619 Max => plain!(i64::max),
2620 _ => zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut i64| {
2621 match i64_op(op, a, b) {
2622 Some(v) => {
2623 *slot = v;
2624 true
2625 }
2626 None => false,
2627 }
2628 }),
2629 }
2630}
2631
2632multiversioned! {
2633 #[allow(clippy::too_many_arguments)]
2640 fn dyad_i64_chunk(
2641 op: ScalarDyad,
2642 xs: &[i64],
2643 xoff: usize,
2644 xdiv: usize,
2645 ys: &[i64],
2646 yoff: usize,
2647 ydiv: usize,
2648 start: usize,
2649 out: &mut [i64],
2650 ) -> bool = dyad_i64_chunk_body;
2651}
2652
2653#[allow(clippy::too_many_arguments)]
2655fn dyad_i64(
2656 op: ScalarDyad,
2657 xs: &[i64],
2658 xoff: usize,
2659 xdiv: usize,
2660 ys: &[i64],
2661 yoff: usize,
2662 ydiv: usize,
2663 n: usize,
2664) -> Option<Vec<i64>> {
2665 let (out, ok) = par::fill(n, |start, part| {
2666 dyad_i64_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part)
2667 });
2668 ok.then_some(out)
2669}
2670
2671#[allow(clippy::too_many_arguments)]
2672#[inline(always)]
2673fn dyad_f64_chunk_body(
2674 op: ScalarDyad,
2675 xs: &[f64],
2676 xoff: usize,
2677 xdiv: usize,
2678 ys: &[f64],
2679 yoff: usize,
2680 ydiv: usize,
2681 start: usize,
2682 out: &mut [f64],
2683 span: Span,
2684) -> Result<()> {
2685 use ScalarDyad::*;
2686 macro_rules! plain {
2689 ($step:expr) => {{
2690 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut f64| {
2691 *slot = $step(a, b);
2692 true
2693 });
2694 return Ok(());
2695 }};
2696 }
2697 match op {
2698 Add => plain!(|a: f64, b: f64| a + b),
2699 Sub => plain!(|a: f64, b: f64| a - b),
2700 Mul => plain!(|a: f64, b: f64| a * b),
2701 Min => plain!(f64::min),
2702 Max => plain!(f64::max),
2703 _ => {}
2704 }
2705 let mut err = None;
2706 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, out, |a, b, slot: &mut f64| {
2707 match f64_op(op, a, b, span) {
2708 Ok(v) => {
2709 *slot = v;
2710 true
2711 }
2712 Err(e) => {
2713 err = Some(e);
2714 false
2715 }
2716 }
2717 });
2718 match err {
2719 Some(e) => Err(e),
2720 None => Ok(()),
2721 }
2722}
2723
2724multiversioned! {
2725 #[allow(clippy::too_many_arguments)]
2727 fn dyad_f64_chunk(
2728 op: ScalarDyad,
2729 xs: &[f64],
2730 xoff: usize,
2731 xdiv: usize,
2732 ys: &[f64],
2733 yoff: usize,
2734 ydiv: usize,
2735 start: usize,
2736 out: &mut [f64],
2737 span: Span,
2738 ) -> Result<()> = dyad_f64_chunk_body;
2739}
2740
2741#[allow(clippy::too_many_arguments)]
2742fn dyad_f64(
2743 op: ScalarDyad,
2744 xs: &[f64],
2745 xoff: usize,
2746 xdiv: usize,
2747 ys: &[f64],
2748 yoff: usize,
2749 ydiv: usize,
2750 n: usize,
2751 span: Span,
2752) -> Result<Vec<f64>> {
2753 par::try_fill(n, |start, part| {
2754 dyad_f64_chunk(op, xs, xoff, xdiv, ys, yoff, ydiv, start, part, span)
2755 })
2756}
2757
2758fn crossed_types(a: DType, b: DType) -> bool {
2762 let class = |d: DType| match d {
2763 DType::Box => 2,
2764 DType::Char => 1,
2765 _ => 0,
2766 };
2767 class(a) != class(b)
2768}
2769
2770#[allow(clippy::too_many_arguments)]
2771fn compare_data(
2772 op: ScalarDyad,
2773 x: &Data,
2774 xoff: usize,
2775 xdiv: usize,
2776 y: &Data,
2777 yoff: usize,
2778 ydiv: usize,
2779 n: usize,
2780 tol: Tol,
2781 span: Span,
2782) -> Result<Data> {
2783 use ScalarDyad::*;
2784 let (dx, dy) = (x.dtype(), y.dtype());
2785 let equality = matches!(op, Eq | Ne);
2786 let boxed = dx == DType::Box || dy == DType::Box;
2792 if equality && crossed_types(dx, dy) && (!boxed || tol.is_j()) {
2793 let unequal = op == Ne;
2794 return Ok(Data::Bool(vec![u8::from(unequal); n].into()));
2795 }
2796 if boxed {
2797 if !equality {
2800 return Err(box_arith(span));
2801 }
2802 let (Data::Box(a), Data::Box(b)) = (x, y) else {
2803 return Err(Error::not_yet("a scalar function inside a nested array", span));
2806 };
2807 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
2808 for (k, slot) in part.iter_mut().enumerate() {
2809 let i = start + k;
2810 let e = arrays_match(&a[xoff + i / xdiv], &b[yoff + i / ydiv], tol);
2811 *slot = u8::from(if op == Eq { e } else { !e });
2812 }
2813 true
2814 });
2815 return Ok(Data::Bool(out.into()));
2816 }
2817 if dx == DType::Char || dy == DType::Char {
2818 if dx != dy {
2819 return Err(Error::new(
2820 ErrorKind::Type,
2821 "cannot compare character and numeric data",
2822 Some(span),
2823 ));
2824 }
2825 if !equality {
2826 return Err(Error::new(
2827 ErrorKind::Type,
2828 "cannot order character data; only equality applies",
2829 Some(span),
2830 ));
2831 }
2832 let (Data::Char(a), Data::Char(b)) = (x, y) else {
2833 return Err(Error::internal("character comparison on non-character data"));
2834 };
2835 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
2836 zip_chunk(a, xoff, xdiv, b, yoff, ydiv, start, part, |p, q, slot| {
2837 let e = p == q;
2838 *slot = if op == Eq { e as u8 } else { !e as u8 };
2839 true
2840 })
2841 });
2842 return Ok(Data::Bool(out.into()));
2843 }
2844 if DType::promote(dx, dy).is_some_and(DType::is_exact) {
2845 if let Some(d) = exact_compare_data(op, x, xoff, xdiv, y, yoff, ydiv, n) {
2846 return Ok(d);
2847 }
2848 }
2849 if dx == DType::Complex || dy == DType::Complex {
2850 if !equality {
2851 return Err(no_complex_order(span));
2852 }
2853 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2854 let xs = borrow_cx(x, &mut tx);
2855 let ys = borrow_cx(y, &mut ty);
2856 let (out, _) = par::fill(n, |start, part: &mut [u8]| {
2857 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
2858 let e = tol.eq_cx(a, b);
2859 *slot = if op == Eq { e as u8 } else { !e as u8 };
2860 true
2861 })
2862 });
2863 return Ok(Data::Bool(out.into()));
2864 }
2865 let out = if DType::promote(dx, dy) == Some(DType::F64) {
2868 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2869 let xs = borrow_f64(x, &mut tx);
2870 let ys = borrow_f64(y, &mut ty);
2871 par::fill(n, |start, part: &mut [u8]| {
2872 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
2873 *slot = tol_cmp(op, a, b, tol) as u8;
2874 true
2875 })
2876 })
2877 .0
2878 } else {
2879 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2880 let xs = borrow_i64(x, &mut tx);
2881 let ys = borrow_i64(y, &mut ty);
2882 par::fill(n, |start, part: &mut [u8]| {
2883 zip_chunk(xs, xoff, xdiv, ys, yoff, ydiv, start, part, |a, b, slot| {
2884 *slot = cmp_result(op, Some(i64::cmp(&a, &b))) as u8;
2885 true
2886 })
2887 })
2888 .0
2889 };
2890 Ok(Data::Bool(out.into()))
2891}
2892
2893#[inline(always)]
2895pub(crate) fn tol_cmp(op: ScalarDyad, a: f64, b: f64, tol: Tol) -> bool {
2896 use ScalarDyad::*;
2897 match op {
2898 Eq => tol.eq(a, b),
2899 Ne => !tol.eq(a, b),
2900 Lt => tol.lt(a, b),
2901 Le => tol.le(a, b),
2902 Gt => tol.lt(b, a),
2903 Ge => tol.le(b, a),
2904 _ => false,
2905 }
2906}
2907
2908fn cmp_result(op: ScalarDyad, ord: Option<std::cmp::Ordering>) -> bool {
2910 use std::cmp::Ordering::*;
2911 use ScalarDyad::*;
2912 match ord {
2913 None => matches!(op, Ne),
2914 Some(o) => match op {
2915 Eq => o == Equal,
2916 Ne => o != Equal,
2917 Lt => o == Less,
2918 Le => o != Greater,
2919 Gt => o == Greater,
2920 Ge => o != Less,
2921 _ => false,
2922 },
2923 }
2924}
2925
2926fn gcd_i128(a: i128, b: i128) -> i128 {
2928 let (mut a, mut b) = (a.abs(), b.abs());
2929 while b != 0 {
2930 let t = a % b;
2931 a = b;
2932 b = t;
2933 }
2934 a
2935}
2936
2937#[allow(clippy::too_many_arguments)]
2941fn lcm_gcd_data(
2942 op: ScalarDyad,
2943 x: &Data,
2944 xoff: usize,
2945 xdiv: usize,
2946 y: &Data,
2947 yoff: usize,
2948 ydiv: usize,
2949 n: usize,
2950 span: Span,
2951) -> Result<Data> {
2952 let t = arith_type(x.dtype(), y.dtype(), span)?;
2953 if t == DType::Complex {
2954 return complex_dyad_data(op, x, xoff, xdiv, y, yoff, ydiv, n, span);
2956 }
2957 if t.is_exact() {
2958 if let Some(d) = exact_dyad_data(op, t, x, xoff, xdiv, y, yoff, ydiv, n, span)? {
2959 return Ok(d);
2960 }
2961 }
2962 let both_bool = x.dtype() == DType::Bool && y.dtype() == DType::Bool;
2963 let float = t == DType::F64;
2964 let (xs, ys) = if float {
2965 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2966 let xf = borrow_f64(x, &mut tx);
2967 let yf = borrow_f64(y, &mut ty);
2968 let integral = |v: &[f64]| v.iter().all(|&a| a.fract() == 0.0 && fits_i64(a));
2969 if !integral(xf) || !integral(yf) {
2970 return Err(Error::not_yet("LCM/GCD on floats", span));
2971 }
2972 (
2973 xf.iter().map(|&a| a as i64).collect::<Vec<_>>(),
2974 yf.iter().map(|&a| a as i64).collect::<Vec<_>>(),
2975 )
2976 } else {
2977 let (mut tx, mut ty) = (Vec::new(), Vec::new());
2978 (borrow_i64(x, &mut tx).to_vec(), borrow_i64(y, &mut ty).to_vec())
2979 };
2980 let (out, fits) = par::fill(n, |start, part: &mut [i128]| {
2983 let mut fits = true;
2984 zip_chunk(&xs, xoff, xdiv, &ys, yoff, ydiv, start, part, |a, b, slot| {
2985 let (a, b) = (a as i128, b as i128);
2986 let g = gcd_i128(a, b);
2987 let v = if op == ScalarDyad::Gcd {
2988 g
2989 } else if g == 0 {
2990 0
2991 } else {
2992 a / g * b
2993 };
2994 fits &= i64::try_from(v).is_ok();
2995 *slot = v;
2996 true
2997 });
2998 fits
2999 });
3000 if !fits || float {
3001 return Ok(Data::F64(par::map(&out, |&v| v as f64).into()));
3002 }
3003 if both_bool {
3004 return Ok(Data::Bool(par::map(&out, |&v| v as u8).into()));
3005 }
3006 Ok(Data::I64(par::map(&out, |&v| v as i64).into()))
3007}
3008
3009fn to_rat_vec(d: &Data) -> Option<Vec<Rat>> {
3014 Some(match d {
3015 Data::Bool(v) => v.iter().map(|&b| Rat::from_int(Ext::from(b))).collect(),
3016 Data::I64(v) => v.iter().map(|&x| Rat::from_int(Ext::from(x))).collect(),
3017 Data::Ext(v) => v.iter().map(|x| Rat::from_int(x.clone())).collect(),
3018 Data::Rat(v) => v.to_vec(),
3019 Data::F64(_) | Data::Complex(_) | Data::Char(_) | Data::Box(_) => return None,
3020 })
3021}
3022
3023fn rat_window(d: &Data, off: usize, div: usize, n: usize) -> Option<Vec<Rat>> {
3031 if n == 0 {
3032 return Some(Vec::new());
3033 }
3034 let end = off + (n - 1) / div + 1;
3035 if off == 0 && end == d.len() {
3036 return to_rat_vec(d);
3037 }
3038 to_rat_vec(&d.slice(off, end))
3039}
3040
3041fn exact_data(t: DType, out: Vec<Rat>) -> Data {
3049 if t == DType::Ext && out.iter().all(Rat::is_integer) {
3050 return Data::Ext(out.iter().map(|r| r.to_int().expect("whole")).collect());
3051 }
3052 Data::Rat(out.into())
3053}
3054
3055fn too_large(span: Span) -> Error {
3057 Error::domain(
3058 format!(
3059 "the exact result needs more than {} bits; use floats for a value this large",
3060 exact::MAX_BITS
3061 ),
3062 span,
3063 )
3064}
3065
3066fn exact_pow(a: &Rat, b: &Rat, span: Span) -> Result<Option<Rat>> {
3069 let Some(e) = b.to_int().as_ref().and_then(exact::ext_to_i64) else {
3070 return Ok(None);
3071 };
3072 if let Some(v) = a.pow(e) {
3073 return Ok(Some(v));
3074 }
3075 if a.is_zero() && e < 0 { Ok(None) } else { Err(too_large(span)) }
3077}
3078
3079#[allow(clippy::too_many_arguments)]
3083fn exact_dyad_data(
3084 op: ScalarDyad,
3085 t: DType,
3086 x: &Data,
3087 xoff: usize,
3088 xdiv: usize,
3089 y: &Data,
3090 yoff: usize,
3091 ydiv: usize,
3092 n: usize,
3093 span: Span,
3094) -> Result<Option<Data>> {
3095 use ScalarDyad::*;
3096 let (Some(xs), Some(ys)) = (rat_window(x, xoff, xdiv, n), rat_window(y, yoff, ydiv, n))
3097 else {
3098 return Ok(None);
3099 };
3100 let mut out = Vec::with_capacity(n);
3101 for i in 0..n {
3102 let a = &xs[i / xdiv];
3103 let b = &ys[i / ydiv];
3104 let v = match op {
3105 Add => a.add(b),
3106 Sub => a.sub(b),
3107 Mul => a.mul(b),
3108 DivJ | DivApl => match a.div(b) {
3110 Some(v) => v,
3111 None => return Ok(None),
3112 },
3113 Min => a.min(b).clone(),
3114 Max => a.max(b).clone(),
3115 Residue => exact::rat_residue(a, b),
3116 Gcd => exact::rat_gcd(a, b),
3117 Lcm => exact::rat_lcm(a, b),
3118 Pow => match exact_pow(a, b, span)? {
3119 Some(v) => v,
3120 None => return Ok(None),
3121 },
3122 Binomial => match (a.to_int(), b.to_int()) {
3123 (Some(k), Some(m)) => match exact::ext_binomial(&k, &m) {
3124 Some(v) => Rat::from_int(v),
3125 None => return Ok(None),
3126 },
3127 _ => return Ok(None),
3128 },
3129 Root if t == DType::Ext => {
3132 let (Some(k), Some(m)) = (a.to_int(), b.to_int()) else {
3133 return Ok(None);
3134 };
3135 let Some(k) = exact::ext_to_i64(&k).and_then(|k| u32::try_from(k).ok()) else {
3136 return Ok(None);
3137 };
3138 match exact::exact_root(k, &m) {
3139 Some(v) => Rat::from_int(v),
3140 None => return Ok(None),
3141 }
3142 }
3143 Root | Log | Circle | MakeComplex | PolarBy => return Ok(None),
3144 Eq | Ne | Lt | Le | Gt | Ge => return Ok(None),
3146 };
3147 out.push(v);
3148 }
3149 Ok(Some(exact_data(t, out)))
3150}
3151
3152fn exact_monad(op: ScalarMonad, y: &Array) -> Option<Array> {
3155 use ScalarMonad::*;
3156 let v = to_rat_vec(&y.data)?;
3157 let shape = y.shape.clone();
3158 if matches!(op, Floor | Ceil | Signum) {
3161 let out: Vec<Ext> = v
3162 .iter()
3163 .map(|r| match op {
3164 Floor => r.floor(),
3165 Ceil => r.ceil(),
3166 _ => r.signum(),
3167 })
3168 .collect();
3169 return Some(Array { shape, data: Data::Ext(out.into()) });
3170 }
3171 let two = Rat::from_int(Ext::from(2));
3172 let mut out = Vec::with_capacity(v.len());
3173 for r in &v {
3174 let value = match op {
3175 Conj => r.clone(),
3176 Neg => r.neg(),
3177 Abs => r.abs(),
3178 Recip => r.recip()?,
3179 Inc => r.add(&Rat::one()),
3180 Dec => r.sub(&Rat::one()),
3181 OneMinus => Rat::one().sub(r),
3182 Double => r.add(r),
3183 Halve => r.div(&two).expect("two is not zero"),
3184 Square => r.mul(r),
3185 Sqrt => r.sqrt()?,
3186 Factorial => Rat::from_int(r.to_int().as_ref().and_then(exact::ext_factorial)?),
3187 Exp | Ln | Pi | Imaginary | Polar | Not => return None,
3190 Floor | Ceil | Signum => unreachable!("handled above"),
3191 };
3192 out.push(value);
3193 }
3194 Some(Array { shape, data: exact_data(y.dtype(), out) })
3195}
3196
3197fn to_exact(y: &Array, span: Span) -> Result<Array> {
3202 let data = match &y.data {
3203 Data::Ext(_) | Data::Rat(_) => return Ok(y.clone()),
3204 Data::Bool(v) => Data::Ext(v.iter().map(|&b| Ext::from(b)).collect()),
3205 Data::I64(v) => Data::Ext(v.iter().map(|&x| Ext::from(x)).collect()),
3206 Data::F64(v) => {
3207 let mut out = Vec::with_capacity(v.len());
3208 for &x in v.iter() {
3209 out.push(exact::f64_to_rat(x).ok_or_else(|| {
3210 Error::domain("an infinity has no exact value", span)
3211 })?);
3212 }
3213 exact_data(DType::Ext, out)
3214 }
3215 Data::Complex(_) | Data::Char(_) | Data::Box(_) => {
3216 return Err(Error::domain(
3217 format!("x: needs real numbers, not {} data", y.dtype().name()),
3218 span,
3219 ));
3220 }
3221 };
3222 Ok(Array { shape: y.shape.clone(), data })
3223}
3224
3225fn from_exact(y: &Array) -> Array {
3228 let shape = y.shape.clone();
3229 match &y.data {
3230 Data::Ext(v) => match v.iter().map(exact::ext_to_i64).collect::<Option<Vec<i64>>>() {
3231 Some(out) => Array { shape, data: Data::I64(out.into()) },
3232 None => Array { shape, data: Data::F64(v.iter().map(exact::ext_to_f64).collect()) },
3233 },
3234 Data::Rat(v) => Array { shape, data: Data::F64(v.iter().map(Rat::to_f64).collect()) },
3235 _ => y.clone(),
3236 }
3237}
3238
3239fn exact_form(x: &Array, y: &Array, span: Span) -> Result<Array> {
3241 match one_whole(x, "the form x: converts to", span)? {
3242 1 => {
3243 let e = to_exact(y, span)?;
3244 e.cast(DType::Rat).ok_or_else(|| Error::internal("an exact value has no rational form"))
3245 }
3246 2 => {
3247 let e = to_exact(y, span)?;
3248 let v = to_rat_vec(&e.data).ok_or_else(|| Error::internal("x: gave an inexact value"))?;
3249 let mut out = Vec::with_capacity(2 * v.len());
3250 for r in &v {
3251 out.push(r.numer().clone());
3252 out.push(r.denom().clone());
3253 }
3254 let mut shape = y.shape.clone();
3255 shape.push(2);
3256 Ok(Array::new(shape, Data::Ext(out.into())))
3257 }
3258 -1 => Ok(from_exact(y)),
3259 -2 => {
3261 if !y.dtype().is_numeric() {
3262 return Err(Error::domain(
3263 format!("x: needs real numbers, not {} data", y.dtype().name()),
3264 span,
3265 ));
3266 }
3267 Ok(y.clone())
3268 }
3269 n => Err(Error::domain(
3270 format!("x: converts to form 1, 2, _1 or _2, not {n}"),
3271 span,
3272 )),
3273 }
3274}
3275
3276#[allow(clippy::too_many_arguments)]
3280fn exact_compare_data(
3281 op: ScalarDyad,
3282 x: &Data,
3283 xoff: usize,
3284 xdiv: usize,
3285 y: &Data,
3286 yoff: usize,
3287 ydiv: usize,
3288 n: usize,
3289) -> Option<Data> {
3290 let (xs, ys) = (rat_window(x, xoff, xdiv, n)?, rat_window(y, yoff, ydiv, n)?);
3291 let out: Vec<u8> = (0..n)
3292 .map(|i| {
3293 let ord = xs[i / xdiv].cmp(&ys[i / ydiv]);
3294 cmp_result(op, Some(ord)) as u8
3295 })
3296 .collect();
3297 Some(Data::Bool(out.into()))
3298}
3299
3300#[allow(clippy::too_many_arguments)]
3304fn scalar_dyad_data(
3305 op: ScalarDyad,
3306 x: &Data,
3307 xoff: usize,
3308 xdiv: usize,
3309 y: &Data,
3310 yoff: usize,
3311 ydiv: usize,
3312 n: usize,
3313 tol: Tol,
3314 span: Span,
3315) -> Result<Data> {
3316 use ScalarDyad::*;
3317 if matches!(op, Eq | Ne | Lt | Le | Gt | Ge) {
3318 return compare_data(op, x, xoff, xdiv, y, yoff, ydiv, n, tol, span);
3319 }
3320 if matches!(op, Lcm | Gcd) {
3321 return lcm_gcd_data(op, x, xoff, xdiv, y, yoff, ydiv, n, span);
3322 }
3323 let t = arith_type(x.dtype(), y.dtype(), span)?;
3324 if t.is_exact() {
3325 if let Some(d) = exact_dyad_data(op, t, x, xoff, xdiv, y, yoff, ydiv, n, span)? {
3326 return Ok(d);
3327 }
3328 }
3330 if t == DType::I64 && !matches!(op, DivJ | DivApl | Log | Root | Circle) {
3331 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3334 let xs = borrow_i64(x, &mut tx);
3335 let ys = borrow_i64(y, &mut ty);
3336 if let Some(v) = dyad_i64(op, xs, xoff, xdiv, ys, yoff, ydiv, n) {
3337 return Ok(Data::I64(v.into()));
3338 }
3339 }
3341 if t == DType::Complex
3342 || matches!(op, MakeComplex | PolarBy)
3343 || pass_leaves_reals(op, x, xoff, xdiv, y, yoff, ydiv, n)
3344 {
3345 let data = complex_dyad_data(op, x, xoff, xdiv, y, yoff, ydiv, n, span)?;
3346 if op == Circle && circle_reads_a_part(x, xoff, xdiv, n) {
3347 if let Data::Complex(v) = &data {
3348 return Ok(Data::F64(v.iter().map(|z| z[0]).collect()));
3349 }
3350 }
3351 return Ok(data);
3352 }
3353 let (mut tx, mut ty) = (Vec::new(), Vec::new());
3354 let xs = borrow_f64(x, &mut tx);
3355 let ys = borrow_f64(y, &mut ty);
3356 Ok(Data::F64(dyad_f64(op, xs, xoff, xdiv, ys, yoff, ydiv, n, span)?.into()))
3357}
3358
3359fn scalar_dyad(
3361 op: ScalarDyad,
3362 x: &Array,
3363 y: &Array,
3364 cfg: EvalCfg,
3365 span: Span,
3366) -> Result<Array> {
3367 let p = agree(&x.shape, &y.shape, &x.shape, &y.shape, cfg.agreement, span)?;
3368 if p.n == 0 {
3372 return Ok(Array::new(p.frame, Data::empty(empty_result_type(x, y))));
3373 }
3374 let data =
3375 scalar_dyad_data(op, &x.data, 0, p.x_div, &y.data, 0, p.y_div, p.n, cfg.tol, span)?;
3376 Ok(Array::new(p.frame, data))
3377}
3378
3379fn empty_result_type(x: &Array, y: &Array) -> DType {
3382 for a in [x, y] {
3383 if a.dtype().is_numeric() {
3384 return a.dtype();
3385 }
3386 }
3387 DType::I64
3388}
3389
3390fn fits_i64(v: f64) -> bool {
3392 v.is_finite() && v >= i64::MIN as f64 && v < i64::MAX as f64
3393}
3394
3395fn monad_leaves_reals(op: ScalarMonad, d: &Data) -> bool {
3397 use ScalarMonad::*;
3398 match op {
3399 Imaginary | Polar => d.dtype().is_numeric(),
3401 Sqrt | Ln => match d {
3402 Data::I64(v) => par::any(v, |&x| x < 0),
3403 Data::F64(v) => par::any(v, |&x| x < 0.0),
3404 Data::Ext(v) => v.iter().any(|x| x.sign() == num_bigint::Sign::Minus),
3405 Data::Rat(v) => v.iter().any(|x| x < &Rat::zero()),
3406 _ => false,
3407 },
3408 _ => false,
3409 }
3410}
3411
3412fn complex_monad(op: ScalarMonad, y: &Array, span: Span) -> Result<Array> {
3414 use ScalarMonad::*;
3415 let mut tmp = Vec::new();
3416 let v = borrow_cx(&y.data, &mut tmp);
3417 if y.count() > 0 && v.is_empty() {
3418 return Err(wrong_type(y.dtype(), span));
3419 }
3420 let data = match op {
3421 Abs => Data::F64(par::map(v, |&z| cx::abs(z)).into()),
3423 Not => return Err(Error::domain("logical negation needs values of 0 or 1", span)),
3424 Factorial => {
3425 return Err(Error::not_yet("the factorial of a complex number", span));
3426 }
3427 _ => {
3428 let step: fn(Cx) -> Cx = match op {
3429 Conj => cx::conj,
3430 Neg => cx::neg,
3431 Signum => cx::signum,
3432 Recip => cx::recip,
3433 Sqrt => cx::sqrt,
3434 Exp => cx::exp,
3435 Ln => cx::ln,
3436 Floor => cx::floor,
3437 Ceil => cx::ceil,
3438 OneMinus => |z| cx::sub(cx::ONE, z),
3439 Inc => |z| cx::add(z, cx::ONE),
3440 Dec => |z| cx::sub(z, cx::ONE),
3441 Double => |z| cx::add(z, z),
3442 Halve => |z| [z[0] / 2.0, z[1] / 2.0],
3443 Square => |z| cx::mul(z, z),
3444 Pi => |z| [std::f64::consts::PI * z[0], std::f64::consts::PI * z[1]],
3445 Imaginary => |z| cx::mul(cx::I, z),
3446 Polar => |z| cx::exp(cx::mul(cx::I, z)),
3447 Abs | Not | Factorial => unreachable!("handled above"),
3448 };
3449 Data::Complex(par::map(v, |&z| step(z)).into())
3450 }
3451 };
3452 Ok(Array { shape: y.shape.clone(), data })
3453}
3454
3455fn scalar_monad(op: ScalarMonad, y: &Array, tol: Tol, span: Span) -> Result<Array> {
3457 use ScalarMonad::*;
3458 let d = &y.data;
3459 if y.count() == 0 && !d.dtype().is_numeric() {
3462 return Ok(Array::new(y.shape.clone(), Data::empty(DType::I64)));
3463 }
3464 if d.dtype() == DType::Complex || monad_leaves_reals(op, d) {
3465 return complex_monad(op, y, span);
3466 }
3467 if d.dtype().is_exact() {
3468 if let Some(a) = exact_monad(op, y) {
3469 return Ok(a);
3470 }
3471 }
3473 let mut tmp = Vec::new();
3476 let data = match op {
3477 Conj if d.dtype().is_numeric() => d.clone(),
3479 Conj => return Err(wrong_type(d.dtype(), span)),
3480 Imaginary | Polar => return Err(Error::internal("a complex monad on the real path")),
3483 Neg => match d {
3484 Data::Bool(v) => Data::I64(par::map(v, |&b| -(b as i64)).into()),
3485 Data::I64(v) => match par::try_map(v, i64::checked_neg) {
3486 Some(out) => Data::I64(out.into()),
3487 None => Data::F64(par::map(v, |&x| -(x as f64)).into()),
3488 },
3489 Data::F64(v) => Data::F64(par::map(v, |&x| -x).into()),
3490 _ => return Err(wrong_type(d.dtype(), span)),
3491 },
3492 Signum => match d {
3493 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64).into()),
3494 Data::I64(v) => Data::I64(par::map(v, |&x| x.signum()).into()),
3495 Data::F64(v) => Data::F64(
3498 par::map(v, |&x| {
3499 if tol.is_zero(x) {
3500 0.0
3501 } else if x > 0.0 {
3502 1.0
3503 } else if x < 0.0 {
3504 -1.0
3505 } else {
3506 0.0
3507 }
3508 })
3509 .into(),
3510 ),
3511 _ => return Err(wrong_type(d.dtype(), span)),
3512 },
3513 Recip => {
3514 let v = as_f64(d, &mut tmp, span)?;
3518 Data::F64(par::map(v, |&x| if x == 0.0 { f64::INFINITY } else { 1.0 / x }).into())
3519 }
3520 Sqrt => {
3521 let v = as_f64(d, &mut tmp, span)?;
3523 Data::F64(par::map(v, |&x| x.sqrt()).into())
3524 }
3525 Exp => {
3526 let v = as_f64(d, &mut tmp, span)?;
3527 Data::F64(par::map(v, |&x| x.exp()).into())
3528 }
3529 Abs => match d {
3530 Data::Bool(_) => d.clone(),
3531 Data::I64(v) => match par::try_map(v, i64::checked_abs) {
3532 Some(out) => Data::I64(out.into()),
3533 None => Data::F64(par::map(v, |&x| (x as f64).abs()).into()),
3534 },
3535 Data::F64(v) => Data::F64(par::map(v, |&x| x.abs()).into()),
3536 _ => return Err(wrong_type(d.dtype(), span)),
3537 },
3538 Floor | Ceil => match d {
3539 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64).into()),
3540 Data::I64(_) => d.clone(),
3541 Data::F64(v) => {
3542 let round = |x: f64| if op == Floor { tol.floor(x) } else { tol.ceil(x) };
3543 match par::try_map(v, |x| {
3545 let r = round(x);
3546 fits_i64(r).then_some(r as i64)
3547 }) {
3548 Some(out) => Data::I64(out.into()),
3549 None => Data::F64(par::map(v, |&x| round(x)).into()),
3550 }
3551 }
3552 _ => return Err(wrong_type(d.dtype(), span)),
3553 },
3554 Inc | Dec => {
3555 let step = if op == Inc { 1i64 } else { -1 };
3556 match d {
3557 Data::Bool(v) => Data::I64(par::map(v, |&b| b as i64 + step).into()),
3558 Data::I64(v) => match par::try_map(v, |x: i64| x.checked_add(step)) {
3559 Some(out) => Data::I64(out.into()),
3560 None => Data::F64(par::map(v, |&x| x as f64 + step as f64).into()),
3561 },
3562 Data::F64(v) => Data::F64(par::map(v, |&x| x + step as f64).into()),
3563 _ => return Err(wrong_type(d.dtype(), span)),
3564 }
3565 }
3566 Double | Square => match d {
3567 Data::Bool(v) => {
3568 Data::I64(par::map(v, |&b| if op == Double { 2 * b as i64 } else { b as i64 }).into())
3569 }
3570 Data::I64(v) => {
3571 let f = |x: i64| if op == Double { x.checked_mul(2) } else { x.checked_mul(x) };
3572 match par::try_map(v, f) {
3573 Some(out) => Data::I64(out.into()),
3574 None => Data::F64(
3575 par::map(v, |&x| {
3576 let x = x as f64;
3577 if op == Double { x + x } else { x * x }
3578 })
3579 .into(),
3580 ),
3581 }
3582 }
3583 Data::F64(v) => {
3584 Data::F64(par::map(v, |&x| if op == Double { x + x } else { x * x }).into())
3585 }
3586 _ => return Err(wrong_type(d.dtype(), span)),
3587 },
3588 Halve => {
3589 let v = as_f64(d, &mut tmp, span)?;
3590 Data::F64(par::map(v, |&x| x / 2.0).into())
3591 }
3592 Pi => {
3593 let v = as_f64(d, &mut tmp, span)?;
3594 Data::F64(par::map(v, |&x| std::f64::consts::PI * x).into())
3595 }
3596 Factorial => {
3597 let v = as_f64(d, &mut tmp, span)?;
3598 Data::F64(par::map(v, |&x| factorial(x)).into())
3599 }
3600 Ln => {
3601 let v = as_f64(d, &mut tmp, span)?;
3603 Data::F64(par::map(v, |&x| x.ln()).into())
3605 }
3606 OneMinus => match d {
3607 Data::Bool(v) => Data::Bool(par::map(v, |&b| 1 - b).into()),
3608 Data::I64(v) => match par::try_map(v, |x: i64| 1i64.checked_sub(x)) {
3609 Some(out) => Data::I64(out.into()),
3610 None => Data::F64(par::map(v, |&x| 1.0 - x as f64).into()),
3611 },
3612 Data::F64(v) => Data::F64(par::map(v, |&x| 1.0 - x).into()),
3613 _ => return Err(wrong_type(d.dtype(), span)),
3614 },
3615 Not => {
3616 let bad = || Error::domain("logical negation needs values of 0 or 1", span);
3617 match d {
3618 Data::Bool(v) => Data::Bool(par::map(v, |&b| 1 - b).into()),
3619 Data::I64(v) => {
3620 let out = par::try_map(v, |x: i64| match x {
3621 0 => Some(1u8),
3622 1 => Some(0u8),
3623 _ => None,
3624 })
3625 .ok_or_else(bad)?;
3626 Data::Bool(out.into())
3627 }
3628 Data::F64(v) => {
3629 let out = par::try_map(v, |x: f64| {
3630 if x == 0.0 {
3631 Some(1u8)
3632 } else if x == 1.0 {
3633 Some(0u8)
3634 } else {
3635 None
3636 }
3637 })
3638 .ok_or_else(bad)?;
3639 Data::Bool(out.into())
3640 }
3641 _ => return Err(bad()),
3642 }
3643 }
3644 };
3645 Ok(Array { shape: y.shape.clone(), data })
3646}
3647
3648fn transpose_axes(y: &Array) -> Array {
3652 if y.rank() < 2 {
3653 return y.clone();
3654 }
3655 let out_shape: Vec<usize> = y.shape.iter().rev().copied().collect();
3656 let src_strides = strides(&y.shape);
3657 let r = y.rank();
3658 let n = y.count();
3659 let mut data = Data::empty(y.dtype());
3660 let mut coord = vec![0usize; r];
3661 for _ in 0..n {
3662 let idx: usize = (0..r).map(|k| coord[k] * src_strides[r - 1 - k]).sum();
3664 push_elem(&mut data, &y.data, idx);
3665 odometer(&mut coord, &out_shape);
3666 }
3667 Array::new(out_shape, data)
3668}
3669
3670fn iota_j(y: &Array, span: Span) -> Result<Array> {
3673 if y.rank() > 1 {
3674 return Err(Error::new(
3675 ErrorKind::Rank,
3676 "index generator needs a scalar or vector argument",
3677 Some(span),
3678 ));
3679 }
3680 let dims = y
3681 .to_i64_vec()
3682 .ok_or_else(|| Error::domain("index generator needs integer lengths", span))?;
3683 let shape: Vec<usize> = dims.iter().map(|d| d.unsigned_abs() as usize).collect();
3684 let n = crate::limits::elements(&shape, span)?;
3685 let st = strides(&shape);
3686 let mut out = Vec::with_capacity(n);
3687 let mut coord = vec![0usize; shape.len()];
3688 for _ in 0..n {
3689 let mut v = 0usize;
3690 for k in 0..shape.len() {
3691 let c = if dims[k] < 0 { shape[k] - 1 - coord[k] } else { coord[k] };
3692 v += c * st[k];
3693 }
3694 out.push(v as i64);
3695 odometer(&mut coord, &shape);
3696 }
3697 let data = Data::I64(out.into());
3698 let data = if y.dtype() == DType::Ext {
3701 data.cast(DType::Ext).ok_or_else(|| Error::internal("integers have no extended form"))?
3702 } else {
3703 data
3704 };
3705 Ok(Array::new(shape, data))
3706}
3707
3708fn head(y: &Array) -> Array {
3710 if y.rank() == 0 {
3711 return y.clone();
3712 }
3713 if y.items() == 0 {
3714 let cell_shape = y.shape[1..].to_vec();
3715 let n: usize = cell_shape.iter().product();
3716 return Array::new(cell_shape, fill_data(y.dtype(), n));
3717 }
3718 y.item(0)
3719}
3720
3721fn behead(y: &Array, span: Span) -> Result<Array> {
3722 if y.rank() == 0 {
3723 return Err(Error::domain("cannot drop the first item of a scalar", span));
3724 }
3725 if y.items() == 0 {
3726 return Ok(y.clone());
3727 }
3728 let m = y.item_size();
3729 let mut shape = y.shape.clone();
3730 shape[0] -= 1;
3731 Ok(Array::new(shape, y.data.slice(m, y.count())))
3732}
3733
3734fn tail(y: &Array) -> Array {
3736 if y.rank() == 0 {
3737 return y.clone();
3738 }
3739 let n = y.items();
3740 if n == 0 {
3741 let cell_shape = y.shape[1..].to_vec();
3742 let m: usize = cell_shape.iter().product();
3743 return Array::new(cell_shape, fill_data(y.dtype(), m));
3744 }
3745 y.item(n - 1)
3746}
3747
3748fn curtail(y: &Array) -> Array {
3750 if y.rank() == 0 {
3751 return Array::empty(y.dtype());
3752 }
3753 let n = y.items();
3754 if n == 0 {
3755 return y.clone();
3756 }
3757 let m = y.item_size();
3758 let mut shape = y.shape.clone();
3759 shape[0] = n - 1;
3760 Array::new(shape, y.data.slice(0, (n - 1) * m))
3761}
3762
3763fn reverse(y: &Array) -> Array {
3765 if y.rank() == 0 {
3766 return y.clone();
3767 }
3768 let n = y.items();
3769 let m = y.item_size();
3770 let mut data = Data::empty(y.dtype());
3771 for i in (0..n).rev() {
3772 for k in 0..m {
3773 push_elem(&mut data, &y.data, i * m + k);
3774 }
3775 }
3776 Array::new(y.shape.clone(), data)
3777}
3778
3779fn rotate(x: &Array, y: &Array, span: Span) -> Result<Array> {
3782 let counts = axis_counts(x, "rotate", span)?;
3783 if y.rank() == 0 {
3784 return Ok(y.clone());
3785 }
3786 if counts.len() > y.rank() {
3787 return Err(Error::new(
3788 ErrorKind::Length,
3789 format!(
3790 "rotate has {} amounts for an argument of rank {}",
3791 counts.len(),
3792 y.rank()
3793 ),
3794 Some(span),
3795 ));
3796 }
3797 let st = strides(&y.shape);
3798 let n = y.count();
3799 let r = y.rank();
3800 let mut data = Data::empty(y.dtype());
3801 let mut coord = vec![0usize; r];
3802 for _ in 0..n {
3803 let mut idx = 0usize;
3804 for k in 0..r {
3805 let len = y.shape[k] as i64;
3807 let s = counts.get(k).copied().unwrap_or(0);
3808 idx += (coord[k] as i64 + s).rem_euclid(len) as usize * st[k];
3809 }
3810 push_elem(&mut data, &y.data, idx);
3811 odometer(&mut coord, &y.shape);
3812 }
3813 Ok(Array::new(y.shape.clone(), data))
3814}
3815
3816fn elem_key(d: &Data, i: usize) -> u64 {
3819 match d {
3820 Data::Bool(v) => v[i] as u64,
3821 Data::I64(v) => v[i] as u64,
3822 Data::F64(v) => {
3823 let x = v[i];
3824 if x == 0.0 { 0 } else { x.to_bits() }
3825 }
3826 Data::Complex(v) => cx_key(v[i]),
3827 Data::Char(v) => v[i] as u64,
3828 Data::Ext(_) | Data::Rat(_) | Data::Box(_) => 0,
3831 }
3832}
3833
3834fn num_key(d: &Data, i: usize) -> u64 {
3837 match d {
3838 Data::Bool(v) => (v[i] as f64).to_bits(),
3839 Data::I64(v) => (v[i] as f64).to_bits(),
3840 Data::F64(v) => {
3841 let x = v[i];
3842 if x == 0.0 { 0.0f64.to_bits() } else { x.to_bits() }
3843 }
3844 Data::Complex(v) => cx_key(v[i]),
3845 Data::Char(v) => v[i] as u64,
3846 Data::Ext(_) | Data::Rat(_) | Data::Box(_) => 0,
3848 }
3849}
3850
3851fn cx_key(z: Cx) -> u64 {
3853 let bits = |x: f64| if x == 0.0 { 0u64 } else { x.to_bits() };
3854 bits(z[0]) ^ bits(z[1]).rotate_left(32)
3855}
3856
3857fn nub(y: &Array, tol: Tol) -> Array {
3859 if y.rank() == 0 {
3860 return Array::new(vec![1], y.data.clone());
3861 }
3862 let n = y.items();
3863 let m = y.item_size();
3864 let mut keep = Vec::new();
3865 if y.dtype() == DType::Box || y.dtype().is_exact() {
3866 for i in 0..n {
3869 if !keep.iter().any(|&j| arrays_match(&y.item(i), &y.item(j), tol)) {
3870 keep.push(i);
3871 }
3872 }
3873 } else if y.dtype() == DType::F64 && tol.ct != 0.0 {
3874 let mut tv = Vec::new();
3877 let v = borrow_f64(&y.data, &mut tv);
3878 for i in 0..n {
3879 if !keep.iter().any(|&j| (0..m).all(|k| tol.eq(v[i * m + k], v[j * m + k]))) {
3880 keep.push(i);
3881 }
3882 }
3883 } else {
3884 let mut seen: HashSet<Vec<u64>> = HashSet::with_capacity(n);
3885 for i in 0..n {
3886 let key: Vec<u64> = (0..m).map(|k| elem_key(&y.data, i * m + k)).collect();
3887 if seen.insert(key) {
3888 keep.push(i);
3889 }
3890 }
3891 }
3892 let mut data = Data::empty(y.dtype());
3893 for &i in &keep {
3894 for k in 0..m {
3895 push_elem(&mut data, &y.data, i * m + k);
3896 }
3897 }
3898 let mut shape = y.shape.clone();
3899 shape[0] = keep.len();
3900 Array::new(shape, data)
3901}
3902
3903fn cmp_items(d: &Data, i: usize, j: usize, m: usize) -> std::cmp::Ordering {
3907 use std::cmp::Ordering::Equal;
3908 let (a, b) = (i * m, j * m);
3909 let ord = |k: usize| match d {
3910 Data::Bool(v) => v[a + k].cmp(&v[b + k]),
3911 Data::I64(v) => v[a + k].cmp(&v[b + k]),
3912 Data::F64(v) => v[a + k].partial_cmp(&v[b + k]).unwrap_or(Equal),
3913 Data::Complex(v) => v[a + k][0]
3920 .partial_cmp(&v[b + k][0])
3921 .unwrap_or(Equal)
3922 .then_with(|| v[a + k][1].partial_cmp(&v[b + k][1]).unwrap_or(Equal)),
3923 Data::Char(v) => v[a + k].cmp(&v[b + k]),
3924 Data::Ext(v) => v[a + k].cmp(&v[b + k]),
3927 Data::Rat(v) => v[a + k].cmp(&v[b + k]),
3928 Data::Box(_) => Equal,
3930 };
3931 (0..m).map(ord).find(|o| *o != Equal).unwrap_or(Equal)
3932}
3933
3934fn grade_order(y: &Array, down: bool) -> Vec<usize> {
3936 if y.rank() == 0 {
3937 return vec![0];
3938 }
3939 let n = y.items();
3940 let m = y.item_size();
3941 let mut idx: Vec<usize> = (0..n).collect();
3942 if down {
3945 idx.sort_by(|&a, &b| cmp_items(&y.data, b, a, m));
3946 } else {
3947 idx.sort_by(|&a, &b| cmp_items(&y.data, a, b, m));
3948 }
3949 idx
3950}
3951
3952fn select_items(y: &Array, order: &[usize]) -> Array {
3954 let m = y.item_size();
3955 let mut data = Data::empty(y.dtype());
3956 for &i in order {
3957 for k in 0..m {
3958 push_elem(&mut data, &y.data, i * m + k);
3959 }
3960 }
3961 let mut shape = y.shape.clone();
3962 shape[0] = order.len();
3963 Array::new(shape, data)
3964}
3965
3966fn check_gradable(y: &Array, order: ComplexOrder, span: Span) -> Result<()> {
3971 if y.dtype() == DType::Box {
3972 return Err(Error::not_yet("grading boxed arrays (the total array ordering)", span));
3973 }
3974 if y.dtype() == DType::Complex && order != ComplexOrder::RealThenImaginary {
3977 return Err(Error::not_yet("grading complex values by magnitude and angle", span));
3978 }
3979 Ok(())
3980}
3981
3982fn grade_select(
3986 x: &Array,
3987 y: &Array,
3988 down: bool,
3989 order: ComplexOrder,
3990 span: Span,
3991) -> Result<Array> {
3992 check_gradable(y, order, span)?;
3993 let order = grade_order(y, down);
3994 if x.rank() == 0 {
3995 return Ok(x.clone());
3996 }
3997 if let Some(&past) = order.iter().find(|&&i| i >= x.items()) {
3998 return Err(Error::domain(
3999 format!("index {past} is out of range: the argument has {} items", x.items()),
4000 span,
4001 ));
4002 }
4003 Ok(select_items(x, &order))
4004}
4005
4006pub(crate) fn arrays_match(x: &Array, y: &Array, tol: Tol) -> bool {
4009 if x.shape != y.shape {
4010 return false;
4011 }
4012 if x.count() == 0 {
4015 return true;
4016 }
4017 if let (Data::Box(a), Data::Box(b)) = (&x.data, &y.data) {
4018 return a.iter().zip(b.iter()).all(|(p, q)| arrays_match(p, q, tol));
4019 }
4020 let (dx, dy) = (x.dtype(), y.dtype());
4021 match DType::promote(dx, dy) {
4022 None => false,
4023 Some(DType::Char) => match (&x.data, &y.data) {
4024 (Data::Char(a), Data::Char(b)) => a.as_slice() == b.as_slice(),
4025 _ => false,
4026 },
4027 Some(DType::F64) => {
4028 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4029 let a = borrow_f64(&x.data, &mut ta);
4030 let b = borrow_f64(&y.data, &mut tb);
4031 a.iter().zip(b).all(|(p, q)| tol.eq(*p, *q))
4032 }
4033 Some(DType::Complex) => {
4034 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4035 let a = borrow_cx(&x.data, &mut ta);
4036 let b = borrow_cx(&y.data, &mut tb);
4037 a.iter().zip(b).all(|(p, q)| tol.eq_cx(*p, *q))
4038 }
4039 Some(t) if t.is_exact() => match (to_rat_vec(&x.data), to_rat_vec(&y.data)) {
4040 (Some(a), Some(b)) => a == b,
4041 _ => false,
4042 },
4043 Some(_) => {
4044 let (mut ta, mut tb) = (Vec::new(), Vec::new());
4045 let a = borrow_i64(&x.data, &mut ta);
4046 let b = borrow_i64(&y.data, &mut tb);
4047 a.iter().zip(b).all(|(p, q)| p == q)
4048 }
4049 }
4050}
4051
4052fn item_or_self(a: &Array, i: usize) -> Array {
4054 if a.rank() == 0 { a.clone() } else { a.item(i) }
4055}
4056
4057fn member_j(x: &Array, y: &Array, tol: Tol) -> Array {
4060 let cell_rank = y.rank().saturating_sub(1).min(x.rank());
4061 let frame_rank = x.rank() - cell_rank;
4062 let frame: Vec<usize> = x.shape[..frame_rank].to_vec();
4063 let nf: usize = frame.iter().product();
4064 let items = y.items();
4065 let mut out = Vec::with_capacity(nf);
4066 for i in 0..nf {
4067 let cell = x.cell_at(frame_rank, i);
4068 out.push((0..items).any(|j| arrays_match(&cell, &item_or_self(y, j), tol)) as u8);
4069 }
4070 Array::new(frame, Data::Bool(out.into()))
4071}
4072
4073fn member_apl(x: &Array, y: &Array, tol: Tol) -> Array {
4075 let n = x.count();
4076 if x.dtype() == DType::Box
4077 || y.dtype() == DType::Box
4078 || x.dtype().is_exact()
4079 || y.dtype().is_exact()
4080 {
4081 let out: Vec<u8> = (0..n)
4085 .map(|i| {
4086 let e = atom(x, i);
4087 u8::from((0..y.count()).any(|j| arrays_match(&e, &atom(y, j), tol)))
4088 })
4089 .collect();
4090 return Array::new(x.shape.clone(), Data::Bool(out.into()));
4091 }
4092 if (x.dtype() == DType::Char) != (y.dtype() == DType::Char) {
4093 return Array::new(x.shape.clone(), Data::Bool(vec![0u8; n].into()));
4094 }
4095 if tol.ct != 0.0
4096 && (x.dtype() == DType::F64 || y.dtype() == DType::F64)
4097 && x.dtype() != DType::Char
4098 {
4099 let (mut tx, mut ty) = (Vec::new(), Vec::new());
4101 let xs = borrow_f64(&x.data, &mut tx);
4102 let ys = borrow_f64(&y.data, &mut ty);
4103 let out: Vec<u8> =
4104 xs.iter().map(|a| ys.iter().any(|b| tol.eq(*a, *b)) as u8).collect();
4105 return Array::new(x.shape.clone(), Data::Bool(out.into()));
4106 }
4107 let seen: HashSet<u64> = (0..y.count()).map(|i| num_key(&y.data, i)).collect();
4108 let out: Vec<u8> =
4109 (0..n).map(|i| seen.contains(&num_key(&x.data, i)) as u8).collect();
4110 Array::new(x.shape.clone(), Data::Bool(out.into()))
4111}
4112
4113fn index_of(x: &Array, y: &Array, origin: i64, tol: Tol) -> Array {
4115 let cell_rank = x.rank().saturating_sub(1).min(y.rank());
4116 let frame_rank = y.rank() - cell_rank;
4117 let frame: Vec<usize> = y.shape[..frame_rank].to_vec();
4118 let nf: usize = frame.iter().product();
4119 let items = x.items();
4120 let mut out = Vec::with_capacity(nf);
4121 for i in 0..nf {
4122 let cell = y.cell_at(frame_rank, i);
4123 let at = (0..items)
4124 .find(|&j| arrays_match(&cell, &item_or_self(x, j), tol))
4125 .unwrap_or(items);
4126 out.push(origin + at as i64);
4127 }
4128 Array::new(frame, Data::I64(out.into()))
4129}
4130
4131fn from_index(x: &Array, y: &Array, span: Span) -> Result<Array> {
4133 if let Some(spec) = x.as_boxes().and_then(<[Array]>::first) {
4136 let spec = index_spec(spec, y, span)?;
4137 return Ok(select_spec(&spec, y));
4138 }
4139 let idx = x
4140 .to_i64_vec()
4141 .ok_or_else(|| Error::domain("index must be an integer", span))?;
4142 let Some(&i) = idx.first() else {
4143 return Err(Error::internal("from_index with no index"));
4144 };
4145 let n = y.items() as i64;
4146 let k = if i < 0 { i + n } else { i };
4147 if k < 0 || k >= n {
4148 return Err(Error::domain(
4149 format!("index {i} is out of range: the argument has {n} items"),
4150 span,
4151 ));
4152 }
4153 Ok(item_or_self(y, k as usize))
4154}
4155
4156fn cat_promote(a: &Array, other: &Array, rank: usize, axis: usize, span: Span) -> Result<Array> {
4160 if a.rank() == rank {
4161 return Ok(a.clone());
4162 }
4163 if a.rank() == 0 {
4164 let mut shape =
4165 if other.rank() == rank { other.shape.clone() } else { vec![1usize; rank] };
4166 shape[axis] = 1;
4167 let n: usize = shape.iter().product();
4168 let mut data = Data::empty(a.dtype());
4169 for _ in 0..n {
4170 push_elem(&mut data, &a.data, 0);
4171 }
4172 return Ok(Array::new(shape, data));
4173 }
4174 if a.rank() + 1 == rank {
4175 let mut shape = a.shape.clone();
4176 shape.insert(axis, 1);
4177 return Ok(Array::new(shape, a.data.clone()));
4178 }
4179 Err(Error::new(
4180 ErrorKind::Rank,
4181 format!("cannot catenate rank {} with rank {}", a.rank(), other.rank()),
4182 Some(span),
4183 ))
4184}
4185
4186fn catenate(
4188 x: &Array,
4189 y: &Array,
4190 leading: bool,
4191 fill: bool,
4192 span: Span,
4193) -> Result<Array> {
4194 let rank = x.rank().max(y.rank()).max(1);
4195 let axis = if leading { 0 } else { rank - 1 };
4196 let xa = cat_promote(x, y, rank, axis, span)?;
4197 let ya = cat_promote(y, x, rank, axis, span)?;
4198 let mut ragged = false;
4202 let want: Vec<i64> = (0..rank)
4203 .map(|k| {
4204 ragged |= k != axis && xa.shape[k] != ya.shape[k];
4205 xa.shape[k].max(ya.shape[k]) as i64
4206 })
4207 .collect();
4208 if ragged && !fill {
4209 return Err(Error::new(
4210 ErrorKind::Length,
4211 format!(
4212 "cannot catenate: left shape {}, right shape {}",
4213 show_shape(&xa.shape),
4214 show_shape(&ya.shape)
4215 ),
4216 Some(span),
4217 ));
4218 }
4219 let (xa, ya) = if ragged {
4220 let fit = |a: &Array| -> Result<Array> {
4221 let mut to = want.clone();
4222 to[axis] = a.shape[axis] as i64;
4223 take(&Array::from_i64(to), a, false, span)
4224 };
4225 (fit(&xa)?, fit(&ya)?)
4226 } else {
4227 (xa, ya)
4228 };
4229 let (xa, ya) = if !fill && (xa.dtype() == DType::Box) != (ya.dtype() == DType::Box) {
4233 (nest_like(&xa, &ya), nest_like(&ya, &xa))
4234 } else {
4235 (xa, ya)
4236 };
4237 let dt = DType::promote(xa.dtype(), ya.dtype()).ok_or_else(|| {
4238 let boxed = xa.dtype() == DType::Box || ya.dtype() == DType::Box;
4239 let what = if boxed {
4240 "cannot catenate boxed and unboxed data; box the other side first"
4241 } else {
4242 "cannot catenate character and numeric data"
4243 };
4244 Error::new(ErrorKind::Type, what, Some(span))
4245 })?;
4246 let widen = |a: &Array| -> Result<Data> {
4247 if a.dtype() == dt {
4248 Ok(a.data.clone())
4249 } else {
4250 a.data.cast(dt).ok_or_else(|| Error::internal("unsupported widening in catenate"))
4251 }
4252 };
4253 let xd = widen(&xa)?;
4254 let yd = widen(&ya)?;
4255 let outer: usize = xa.shape[..axis].iter().product();
4256 let ix: usize = xa.shape[axis..].iter().product();
4257 let iy: usize = ya.shape[axis..].iter().product();
4258 let mut data = Data::empty(dt);
4259 for o in 0..outer {
4260 for k in 0..ix {
4261 push_elem(&mut data, &xd, o * ix + k);
4262 }
4263 for k in 0..iy {
4264 push_elem(&mut data, &yd, o * iy + k);
4265 }
4266 }
4267 let mut shape = xa.shape.clone();
4268 shape[axis] = xa.shape[axis] + ya.shape[axis];
4269 Ok(Array::new(shape, data))
4270}
4271
4272fn copy_items(x: &Array, y: &Array, apl: bool, span: Span) -> Result<Array> {
4280 let counts = x
4281 .to_i64_vec()
4282 .ok_or_else(|| Error::domain("replication counts must be integers", span))?;
4283 if !apl && counts.iter().any(|&c| c < 0) {
4284 return Err(Error::domain("replication counts must be nonnegative", span));
4285 }
4286 let scalar_y = y.rank() == 0;
4287 let m = y.item_size();
4288 let n = if x.rank() == 0 || !scalar_y { y.items() } else { counts.len() };
4289 let per = if x.rank() == 0 { vec![counts[0]; n] } else { counts };
4290 if per.len() != n {
4291 return Err(Error::new(
4292 ErrorKind::Length,
4293 format!("{} replication count(s) for {n} item(s)", per.len()),
4294 Some(span),
4295 ));
4296 }
4297 let items: u128 = per.iter().map(|&c| c.unsigned_abs() as u128).sum();
4300 let total = crate::limits::count(items * m.max(1) as u128, span)? / m.max(1);
4301 let mut data = Data::empty(y.dtype());
4302 for (i, &c) in per.iter().enumerate() {
4303 let src = if scalar_y { 0 } else { i };
4305 for _ in 0..c.unsigned_abs() {
4306 for k in 0..m {
4307 if c < 0 {
4308 data.push_fill();
4309 } else {
4310 push_elem(&mut data, &y.data, src * m + k);
4311 }
4312 }
4313 }
4314 }
4315 let mut shape = if scalar_y { vec![1] } else { y.shape.clone() };
4317 shape[0] = total;
4318 Ok(Array::new(shape, data))
4319}
4320
4321fn format_chars(y: &Array, opts: &FmtOpts) -> Array {
4330 if y.dtype() == DType::Char {
4331 return y.clone();
4332 }
4333 if y.count() == 0 {
4335 return Array::new(y.shape.clone(), Data::empty(DType::Char));
4336 }
4337 let text = crate::fmt::format_array(y, opts);
4338 if y.dtype() == DType::Box {
4339 let lines = text.lines().filter(|l| !l.is_empty()).count();
4344 let rows: usize =
4345 if y.rank() == 0 { 1 } else { y.shape[..y.rank() - 1].iter().product() };
4346 if lines != rows {
4347 return text_planes(&text, &y.shape[..y.rank().saturating_sub(2)]);
4348 }
4349 }
4350 if y.rank() < 2 {
4351 let chars: Vec<char> = text.chars().collect();
4352 return Array::new(vec![chars.len()], Data::Char(chars.into()));
4353 }
4354 let lines: Vec<&str> = text.lines().filter(|l| !l.is_empty()).collect();
4357 let width = lines.iter().map(|l| l.chars().count()).max().unwrap_or(0);
4358 let mut chars: Vec<char> = Vec::with_capacity(lines.len() * width);
4359 for line in &lines {
4360 chars.extend(line.chars());
4361 chars.resize(chars.len() + width - line.chars().count(), ' ');
4362 }
4363 let mut shape = y.shape[..y.rank() - 1].to_vec();
4366 shape.push(width);
4367 debug_assert_eq!(lines.len(), shape[..shape.len() - 1].iter().product::<usize>());
4368 Array::new(shape, Data::Char(chars.into()))
4369}
4370
4371fn text_planes(text: &str, frame: &[usize]) -> Array {
4374 let lines: Vec<&str> = text.lines().filter(|l| !l.is_empty()).collect();
4375 let width = lines.iter().map(|l| l.chars().count()).max().unwrap_or(0);
4376 let planes: usize = frame.iter().product::<usize>().max(1);
4377 let per = lines.len() / planes;
4378 let mut chars: Vec<char> = Vec::with_capacity(lines.len() * width);
4379 for line in &lines {
4380 chars.extend(line.chars());
4381 chars.resize(chars.len() + width - line.chars().count(), ' ');
4382 }
4383 let mut shape = frame.to_vec();
4384 shape.push(per);
4385 shape.push(width);
4386 Array::new(shape, Data::Char(chars.into()))
4387}
4388
4389fn digits_of(a: &Array, what: &str, span: Span) -> Result<Vec<f64>> {
4391 a.to_f64_vec().ok_or_else(|| Error::domain(format!("{what} needs numeric data"), span))
4392}
4393
4394fn narrow(values: Vec<f64>, integral: bool) -> Data {
4398 if integral && values.iter().all(|&v| v.fract() == 0.0 && fits_i64(v)) {
4399 return Data::I64(values.iter().map(|&v| v as i64).collect::<Vec<_>>().into());
4400 }
4401 Data::F64(values.into())
4402}
4403
4404fn is_integral(a: &Array) -> bool {
4406 !matches!(a.dtype(), DType::F64 | DType::Rat | DType::Char)
4407}
4408
4409fn decode(x: Option<&Array>, y: &Array, span: Span) -> Result<Array> {
4412 let digits = digits_of(y, "decode", span)?;
4413 let radix: Vec<f64> = match x {
4414 None => vec![2.0; digits.len()],
4415 Some(x) => {
4416 let r = digits_of(x, "decode", span)?;
4417 match r.len() {
4418 1 => vec![r[0]; digits.len()],
4419 n if n == digits.len() => r,
4420 n => {
4421 return Err(Error::new(
4422 ErrorKind::Length,
4423 format!("{n} radices for {} digits", digits.len()),
4424 Some(span),
4425 ));
4426 }
4427 }
4428 }
4429 };
4430 let mut acc = 0.0f64;
4431 for (d, b) in digits.iter().zip(&radix) {
4432 acc = acc * b + d;
4433 }
4434 let integral = is_integral(y) && x.is_none_or(is_integral);
4435 Ok(Array::new(vec![], narrow(vec![acc], integral)))
4436}
4437
4438fn bit_width(values: &[f64], span: Span) -> Result<usize> {
4441 if values.is_empty() {
4443 return Ok(0);
4444 }
4445 let mut m = 0.0f64;
4446 for &v in values {
4447 if !v.is_finite() {
4448 return Err(Error::domain("cannot encode an infinite value", span));
4449 }
4450 m = m.max(v.abs());
4451 }
4452 let whole = m.floor();
4453 if whole >= 1e15 {
4454 return Err(Error::domain("the value is too large to encode in binary", span));
4455 }
4456 let mut w = 1usize;
4457 let mut n = whole as i64;
4458 while n > 1 {
4459 n /= 2;
4460 w += 1;
4461 }
4462 Ok(w)
4463}
4464
4465fn encode_one(radix: &[f64], v: f64, out: &mut [f64]) {
4469 let mut rem = v;
4470 for i in (0..radix.len()).rev() {
4471 let b = radix[i];
4472 if b == 0.0 {
4473 out[i] = rem;
4474 rem = 0.0;
4475 } else {
4476 let r = rem - b * (rem / b).floor();
4477 out[i] = r;
4478 rem = (rem - r) / b;
4479 }
4480 }
4481}
4482
4483fn encode(x: &Array, y: &Array, span: Span) -> Result<Array> {
4487 let radix = digits_of(x, "encode", span)?;
4488 let values = digits_of(y, "encode", span)?;
4489 let k = radix.len();
4490 let n = values.len();
4491 let mut out = vec![0.0f64; k * n];
4492 let mut cell = vec![0.0f64; k];
4493 for (j, &v) in values.iter().enumerate() {
4494 encode_one(&radix, v, &mut cell);
4495 for i in 0..k {
4496 out[i * n + j] = cell[i];
4497 }
4498 }
4499 let mut shape = if x.rank() == 0 { Vec::new() } else { vec![k] };
4502 shape.extend_from_slice(&y.shape);
4503 Ok(Array::new(shape, narrow(out, is_integral(x) && is_integral(y))))
4504}
4505
4506fn encode_bits(y: &Array, span: Span) -> Result<Array> {
4508 let values = digits_of(y, "encode", span)?;
4509 let k = bit_width(&values, span)?;
4510 let radix = vec![2.0; k];
4511 let mut out = vec![0.0f64; values.len() * k];
4512 for (j, &v) in values.iter().enumerate() {
4513 encode_one(&radix, v, &mut out[j * k..(j + 1) * k]);
4514 }
4515 let mut shape = y.shape.clone();
4516 shape.push(k);
4517 Ok(Array::new(shape, narrow(out, is_integral(y))))
4518}
4519
4520fn laminate(x: &Array, y: &Array, span: Span) -> Result<Array> {
4525 let spread = |a: &Array, other: &Array| -> Array {
4526 if a.rank() != 0 {
4527 return a.clone();
4528 }
4529 let shape = if other.rank() == 0 { vec![1] } else { other.shape.clone() };
4530 let n: usize = shape.iter().product();
4531 let mut data = Data::empty(a.dtype());
4532 for _ in 0..n {
4533 push_elem(&mut data, &a.data, 0);
4534 }
4535 Array::new(shape, data)
4536 };
4537 assemble(&[2], vec![spread(x, y), spread(y, x)], span)
4538}
4539
4540fn table_of(y: &Array) -> Array {
4542 let shape = match y.rank() {
4543 0 => vec![1, 1],
4544 _ => vec![y.items(), y.item_size()],
4545 };
4546 Array::new(shape, y.data.clone())
4547}
4548
4549fn table(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
4555 let ranks = u.ranks();
4556 let fxl = x.rank() - effective_rank(ranks[1], x.rank());
4557 let fyl = y.rank() - effective_rank(ranks[2], y.rank());
4558 let mut frame = x.shape[..fxl].to_vec();
4559 frame.extend_from_slice(&y.shape[..fyl]);
4560 let nx: usize = x.shape[..fxl].iter().product();
4561 let ny: usize = y.shape[..fyl].iter().product();
4562 let n = nx * ny;
4563 if n == 0 {
4564 return assemble(&frame, Vec::new(), span);
4565 }
4566 if frame.is_empty() {
4567 return u.dyad(x, y, ctx, span);
4568 }
4569 let work = x.count().max(y.count()).max(n);
4570 let cells = each_cell(n, work, u.is_pure(), ctx, |i, c| {
4571 u.dyad(&x.cell_at(fxl, i / ny), &y.cell_at(fyl, i % ny), c, span)
4572 })?;
4573 assemble(&frame, cells, span)
4574}
4575
4576fn monad_op(p: &Prim, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
4578 match p.monad {
4579 MonadOp::Scalar(op) => scalar_monad(op, y, ctx.cfg.tol, span),
4580 MonadOp::ShapeOf => {
4581 Ok(carry_exact(Array::from_i64(y.shape.iter().map(|&n| n as i64).collect()), y))
4582 }
4583 MonadOp::Tally => Ok(carry_exact(Array::scalar_i64(y.items() as i64), y)),
4584 MonadOp::Ravel => Ok(Array::new(vec![y.count()], y.data.clone())),
4585 MonadOp::TransposeAxes => Ok(transpose_axes(y)),
4586 MonadOp::Head => Ok(head(y)),
4587 MonadOp::Behead => behead(y, span),
4588 MonadOp::Tail => Ok(tail(y)),
4589 MonadOp::Curtail => Ok(curtail(y)),
4590 MonadOp::Reverse => Ok(reverse(y)),
4591 MonadOp::Nub => Ok(nub(y, ctx.cfg.tol)),
4592 MonadOp::GradeUp { origin } | MonadOp::GradeDown { origin } => {
4593 check_gradable(y, ctx.cfg.rules.complex_order, span)?;
4594 let down = matches!(p.monad, MonadOp::GradeDown { .. });
4595 let order = grade_order(y, down);
4596 Ok(Array::from_i64(order.iter().map(|&i| origin + i as i64).collect()))
4597 }
4598 MonadOp::IotaJ => iota_j(y, span),
4599 MonadOp::IotaApl { origin } => iota_apl(y, origin, span),
4600 MonadOp::Echo => {
4601 (ctx.out)(&format!("{}\n", crate::fmt::format_array(y, &ctx.cfg.fmt)));
4602 Ok(Array::empty(DType::I64))
4603 }
4604 MonadOp::Same => Ok(y.clone()),
4605 MonadOp::Format => Ok(format_chars(y, &ctx.cfg.fmt)),
4606 MonadOp::DecodeBits => decode(None, y, span).map(|r| carry_exact(r, y)),
4607 MonadOp::EncodeBits => encode_bits(y, span).map(|r| carry_exact(r, y)),
4608 MonadOp::Itemize => {
4609 let mut shape = vec![1usize];
4610 shape.extend_from_slice(&y.shape);
4611 Ok(Array::new(shape, y.data.clone()))
4612 }
4613 MonadOp::TableOf => Ok(table_of(y)),
4614 MonadOp::Enclose(rule) => Ok(enclose(y, rule)),
4615 MonadOp::Open => Ok(open_cell(y)),
4616 MonadOp::Raze => raze(y, span),
4617 MonadOp::First => Ok(first(y)),
4618 MonadOp::Enlist => enlist(y, span),
4619 MonadOp::Depth => Ok(Array::scalar_i64(depth(y))),
4620 MonadOp::Indices { origin, boxed_coords } => {
4621 where_indices(y, origin, boxed_coords, span)
4622 }
4623 MonadOp::Steps => steps(y, span),
4624 MonadOp::ToExact => to_exact(y, span),
4625 MonadOp::NthPrime => {
4626 let n = y
4627 .to_i64_vec()
4628 .ok_or_else(|| Error::domain("the prime index must be an integer", span))?;
4629 let v = n.first().copied().unwrap_or(0);
4630 Ok(carry_exact(Array::scalar_i64(nth_prime(v, span)?), y))
4631 }
4632 MonadOp::PrimeFactors => {
4633 let n = y
4634 .to_i64_vec()
4635 .ok_or_else(|| Error::domain("prime factors need an integer", span))?;
4636 let v = n.first().copied().unwrap_or(0);
4637 Ok(carry_exact(Array::from_i64(prime_factors(v, span)?), y))
4638 }
4639 MonadOp::MatrixInverse => matrix_inverse(y, span),
4640 MonadOp::Roll { origin, fixed, float_at_zero } => {
4641 roll(y, origin, fixed, float_at_zero, span)
4642 }
4643 MonadOp::ComplexParts { polar } => complex_parts(y, polar, span),
4644 MonadOp::SelfClassify => Ok(self_classify(y, ctx.cfg.tol)),
4645 MonadOp::NubSieve => Ok(nub_sieve(y, ctx.cfg.tol)),
4646 MonadOp::Unicode { pass_chars } => unicode(y, pass_chars, span),
4647 MonadOp::Words => words(y, span),
4648 MonadOp::LevelOf => Ok(Array::scalar_i64(boxing_level(y))),
4649 MonadOp::MapPaths => Ok(map_paths(y)),
4650 MonadOp::Nest => Ok(nest(y)),
4651 MonadOp::PolyRoots => poly_roots(y, span),
4652 MonadOp::PolyDeriv => poly_deriv(y, span),
4653 MonadOp::AnagramIndex => anagram_index(y, ctx.cfg.rules.complex_order, span),
4654 MonadOp::CycleForm => cycle_form(y, span),
4655 MonadOp::Split => Ok(split_items(y)),
4656 MonadOp::Execute { apl } => execute(y, apl, ctx, span),
4657 MonadOp::NotYet(what) => Err(Error::not_yet(what, span)),
4658 MonadOp::None => {
4659 Err(Error::domain(format!("{} has no monadic meaning", p.name), span))
4660 }
4661 }
4662}
4663
4664fn complex_parts(y: &Array, polar: bool, span: Span) -> Result<Array> {
4668 let Some(v) = y.to_complex_vec() else {
4669 return Err(wrong_type(y.dtype(), span));
4670 };
4671 let z = v.first().copied().unwrap_or(cx::ZERO);
4672 let pair = if polar { vec![cx::abs(z), cx::arg(z)] } else { vec![z[0], z[1]] };
4673 Ok(Array::from_f64(pair))
4674}
4675
4676fn axis_counts(x: &Array, what: &str, span: Span) -> Result<Vec<i64>> {
4677 if x.rank() > 1 {
4678 return Err(Error::new(
4679 ErrorKind::Rank,
4680 format!("{what} needs a scalar or vector left argument"),
4681 Some(span),
4682 ));
4683 }
4684 if x.count() == 0 {
4687 return Ok(Vec::new());
4688 }
4689 x.to_i64_vec()
4690 .ok_or_else(|| Error::domain(format!("{what} needs integer lengths"), span))
4691}
4692
4693fn reshape(x: &Array, y: &Array, by_items: bool, span: Span) -> Result<Array> {
4704 let dims = axis_counts(x, "reshape", span)?;
4705 if dims.iter().any(|&d| d < 0) {
4706 return Err(Error::domain("reshape lengths must be nonnegative", span));
4707 }
4708 let mut shape: Vec<usize> = dims.iter().map(|&d| d as usize).collect();
4709 let (unit, src) = if by_items {
4711 let item_shape = if y.rank() == 0 { &[][..] } else { &y.shape[1..] };
4712 shape.extend_from_slice(item_shape);
4713 (item_shape.iter().product::<usize>(), y.items().max(usize::from(y.rank() == 0)))
4714 } else {
4715 (1, y.count())
4716 };
4717 let n = crate::limits::elements(&shape, span)?;
4718 let mut data = Data::empty(y.dtype());
4719 if n > 0 && src == 0 {
4720 if by_items {
4721 return Err(Error::new(ErrorKind::Length, "reshape of an empty array", Some(span)));
4722 }
4723 return Ok(Array::new(shape, fill_data(y.dtype(), n)));
4724 }
4725 for i in 0..n {
4726 push_elem(&mut data, &y.data, (i / unit) % src * unit + i % unit);
4729 }
4730 Ok(Array::new(shape, data))
4731}
4732
4733fn leading_run(y: &Array, counts: &[i64], drop: bool) -> Option<Array> {
4737 if y.rank() == 0 || counts.is_empty() || counts[1..].iter().any(|&c| c != 0) {
4738 return None;
4739 }
4740 let n = y.items();
4741 let k = counts[0];
4742 let a = k.unsigned_abs() as usize;
4743 let (lo, keep) = if drop {
4744 let a = a.min(n);
4745 if k >= 0 { (a, n - a) } else { (0, n - a) }
4746 } else {
4747 if a > n {
4749 return None;
4750 }
4751 if k >= 0 { (0, a) } else { (n - a, a) }
4752 };
4753 Some(section(y, lo, lo + keep))
4754}
4755
4756fn take(x: &Array, y: &Array, prototype_fill: bool, span: Span) -> Result<Array> {
4757 let counts = axis_counts(x, "take", span)?;
4758 let fill = if prototype_fill { prototype_of(y) } else { None };
4762 let promoted;
4763 let base = if y.rank() == 0 {
4765 promoted = Array::new(vec![1], y.data.clone());
4766 &promoted
4767 } else {
4768 y
4769 };
4770 if counts.len() > base.rank() {
4771 return Err(Error::not_yet("take with more axes than the rank", span));
4772 }
4773 if let Some(run) = leading_run(base, &counts, false) {
4774 return Ok(run);
4775 }
4776 let mut out_shape = base.shape.clone();
4777 for (a, &k) in counts.iter().enumerate() {
4778 out_shape[a] = k.unsigned_abs() as usize;
4779 }
4780 let n = crate::limits::elements(&out_shape, span)?;
4781 let st = strides(&base.shape);
4782 let mut data = Data::empty(base.dtype());
4783 let mut coord = vec![0usize; out_shape.len()];
4784 for _ in 0..n {
4785 let mut idx = 0usize;
4786 let mut inside = true;
4787 for a in 0..out_shape.len() {
4788 let len = base.shape[a] as i64;
4789 let c = coord[a] as i64;
4790 let s = match counts.get(a) {
4793 Some(&k) if k < 0 => c + len - k.unsigned_abs() as i64,
4794 _ => c,
4795 };
4796 if s < 0 || s >= len {
4797 inside = false;
4798 break;
4799 }
4800 idx += s as usize * st[a];
4801 }
4802 if inside {
4803 push_elem(&mut data, &base.data, idx);
4804 } else if let (Data::Box(v), Some(p)) = (&mut data, &fill) {
4805 v.push(p.clone());
4806 } else {
4807 data.push_fill();
4808 }
4809 odometer(&mut coord, &out_shape);
4810 }
4811 Ok(Array::new(out_shape, data))
4812}
4813
4814fn prototype_of(y: &Array) -> Option<Array> {
4818 fn zeroed(a: &Array) -> Array {
4819 if let Some(items) = a.as_boxes() {
4820 let inner: Vec<Array> = items.iter().map(zeroed).collect();
4821 return Array::new(a.shape.clone(), Data::Box(inner.into()));
4822 }
4823 let dtype = if a.dtype() == DType::Char { DType::Char } else { DType::I64 };
4824 Array::new(a.shape.clone(), fill_data(dtype, a.count()))
4825 }
4826 let first = y.as_boxes()?.first()?;
4827 Some(zeroed(first))
4828}
4829
4830fn drop_(x: &Array, y: &Array, span: Span) -> Result<Array> {
4831 let counts = axis_counts(x, "drop", span)?;
4832 let promoted;
4833 let base = if y.rank() == 0 {
4834 promoted = Array::new(vec![1], y.data.clone());
4835 &promoted
4836 } else {
4837 y
4838 };
4839 if counts.len() > base.rank() {
4840 return Err(Error::not_yet("drop with more axes than the rank", span));
4841 }
4842 if let Some(run) = leading_run(base, &counts, true) {
4843 return Ok(run);
4844 }
4845 let mut out_shape = base.shape.clone();
4846 let mut offset = vec![0usize; base.rank()];
4847 for (a, &k) in counts.iter().enumerate() {
4848 let len = base.shape[a];
4849 let d = (k.unsigned_abs() as usize).min(len);
4850 out_shape[a] = len - d;
4851 if k > 0 {
4852 offset[a] = d;
4853 }
4854 }
4855 let n: usize = out_shape.iter().product();
4856 let st = strides(&base.shape);
4857 let mut data = Data::empty(base.dtype());
4858 let mut coord = vec![0usize; out_shape.len()];
4859 for _ in 0..n {
4860 let idx: usize = (0..out_shape.len()).map(|a| (coord[a] + offset[a]) * st[a]).sum();
4861 push_elem(&mut data, &base.data, idx);
4862 odometer(&mut coord, &out_shape);
4863 }
4864 Ok(Array::new(out_shape, data))
4865}
4866
4867fn dyad_op(p: &Prim, x: &Array, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
4869 let tol = cfg.tol;
4870 match p.dyad {
4871 DyadOp::Scalar(op) => scalar_dyad(op, x, y, cfg, span),
4874 DyadOp::Reshape => reshape(x, y, cfg.agreement == Agreement::LeadingPrefix, span),
4875 DyadOp::Take => take(x, y, cfg.agreement == Agreement::ExactOrScalar, span),
4876 DyadOp::Drop => drop_(x, y, span),
4877 DyadOp::Right => Ok(y.clone()),
4878 DyadOp::Left => Ok(x.clone()),
4879 DyadOp::Rotate => rotate(x, y, span),
4880 DyadOp::AppendLeading => {
4883 catenate(x, y, true, cfg.agreement == Agreement::LeadingPrefix, span)
4884 }
4885 DyadOp::AppendLast => {
4886 catenate(x, y, false, cfg.agreement == Agreement::LeadingPrefix, span)
4887 }
4888 DyadOp::IndexOf { origin } => Ok(index_of(x, y, origin, tol)),
4889 DyadOp::MemberJ => Ok(member_j(x, y, tol)),
4890 DyadOp::MemberApl => Ok(member_apl(x, y, tol)),
4891 DyadOp::From => from_index(x, y, span),
4892 DyadOp::Match => Ok(Array::scalar_bool(arrays_match(x, y, tol))),
4893 DyadOp::NotMatch => Ok(Array::scalar_bool(!arrays_match(x, y, tol))),
4894 DyadOp::GradeSelect { down } => grade_select(x, y, down, cfg.rules.complex_order, span),
4895 DyadOp::Copy => copy_items(x, y, cfg.agreement == Agreement::ExactOrScalar, span),
4896 DyadOp::Decode => decode(Some(x), y, span).map(|r| carry_exact2(r, x, y)),
4897 DyadOp::Encode => encode(x, y, span).map(|r| carry_exact2(r, x, y)),
4898 DyadOp::Laminate => laminate(x, y, span),
4899 DyadOp::Link => link(x, y, span),
4900 DyadOp::Strand => strand(x, y, span),
4901 DyadOp::IntervalIndex { offset } => interval_index(x, y, offset, tol, span),
4902 DyadOp::IndexOfLast { origin } => Ok(index_of_last(x, y, origin, tol)),
4903 DyadOp::MatrixDivide => matrix_divide(x, y, span),
4904 DyadOp::PartitionEnclose => partition_enclose(x, y, span),
4905 DyadOp::Squad { origin } => squad(x, y, origin, span),
4906 DyadOp::SelectAxis { axis, rank, origin } => {
4907 select_axis(x, y, axis, rank, origin, span)
4908 }
4909 DyadOp::Fetch => fetch(x, y, span),
4910 DyadOp::PolyEval => poly_eval(x, y, span),
4911 DyadOp::PolyIntegral => poly_integral(x, y, span),
4912 DyadOp::TruthTable(m) => truth_table(m, x, y, span),
4913 DyadOp::FormatSpec => format_spec(x, y, &cfg.fmt, span),
4914 DyadOp::Deal { origin, fixed } => deal(x, y, origin, fixed, span),
4915 DyadOp::ExactForm => exact_form(x, y, span),
4916 DyadOp::Boolean(op) => bool_dyad(op, x, y, cfg, span),
4917 DyadOp::Less => Ok(set_less(x, y, tol)),
4918 DyadOp::Union => union_items(x, y, tol, span),
4919 DyadOp::Intersect => Ok(intersect_items(x, y, tol)),
4920 DyadOp::AnagramFrom => anagram_from(x, y, span),
4921 DyadOp::Permute => permute(x, y, span),
4922 DyadOp::FindSeq => Ok(find_seq(x, y, tol)),
4923 DyadOp::UnicodeForm => unicode_form(x, y, span),
4924 DyadOp::PrimeMeta => prime_meta(x, y, span).map(|r| carry_exact2(r, x, y)),
4925 DyadOp::PrimeExponents => prime_exponents(x, y, span).map(|r| carry_exact2(r, x, y)),
4926 DyadOp::Pick { origin } => pick(x, y, origin, span),
4927 DyadOp::Expand => expand(x, y, span),
4928 DyadOp::NotYet(what) => Err(Error::not_yet(what, span)),
4929 DyadOp::None => Err(Error::domain(format!("{} has no dyadic meaning", p.name), span)),
4930 }
4931}
4932
4933fn reduce_identity(v: &Verb, n: usize) -> Option<Data> {
4945 let Verb::Prim(p) = v else { return None };
4946 let DyadOp::Scalar(op) = p.dyad else { return None };
4947 let ints = |k: i64| Data::I64(vec![k; n].into());
4948 let bits = |k: u8| Data::Bool(vec![k; n].into());
4949 Some(match op {
4950 ScalarDyad::Add | ScalarDyad::Sub | ScalarDyad::Gcd | ScalarDyad::Residue => ints(0),
4951 ScalarDyad::Mul
4952 | ScalarDyad::DivJ
4953 | ScalarDyad::DivApl
4954 | ScalarDyad::Pow
4955 | ScalarDyad::Lcm
4956 | ScalarDyad::Root
4957 | ScalarDyad::Binomial => ints(1),
4958 ScalarDyad::Min => Data::F64(vec![f64::INFINITY; n].into()),
4959 ScalarDyad::Max => Data::F64(vec![f64::NEG_INFINITY; n].into()),
4960 ScalarDyad::Eq | ScalarDyad::Le | ScalarDyad::Ge => bits(1),
4961 ScalarDyad::Ne | ScalarDyad::Lt | ScalarDyad::Gt => bits(0),
4962 ScalarDyad::MakeComplex | ScalarDyad::PolarBy => return None,
4965 ScalarDyad::Log | ScalarDyad::Circle => return None,
4968 })
4969}
4970
4971fn is_associative(op: ScalarDyad) -> bool {
4977 use ScalarDyad::*;
4978 matches!(op, Add | Mul | Min | Max)
4979}
4980
4981#[inline(always)]
4982fn fold_range_body<T, F>(
4983 v: &[T],
4984 m: usize,
4985 lo: usize,
4986 hi: usize,
4987 j0: usize,
4988 acc: &mut [T],
4989 step: &F,
4990) -> bool
4991where
4992 T: Copy,
4993 F: Fn(T, T) -> (T, bool),
4994{
4995 let w = acc.len();
4996 let base = (hi - 1) * m + j0;
4997 acc.copy_from_slice(&v[base..base + w]);
4998 let mut over = false;
5001 for i in (lo..hi - 1).rev() {
5002 let row = &v[i * m + j0..i * m + j0 + w];
5003 for (slot, &x) in acc.iter_mut().zip(row) {
5004 let (r, o) = step(x, *slot);
5005 *slot = r;
5006 over |= o;
5007 }
5008 }
5009 !over
5010}
5011
5012multiversioned! {
5013 #[allow(clippy::too_many_arguments)]
5014 fn fold_range_vectorised[T: Copy, F: Fn(T, T) -> (T, bool)](
5015 v: &[T],
5016 m: usize,
5017 lo: usize,
5018 hi: usize,
5019 j0: usize,
5020 acc: &mut [T],
5021 step: &F,
5022 ) -> bool = fold_range_body;
5023}
5024
5025const VECTOR_COLUMNS: usize = 16;
5033
5034#[allow(clippy::too_many_arguments)]
5043#[inline]
5044fn fold_range<T, F>(
5045 v: &[T],
5046 m: usize,
5047 lo: usize,
5048 hi: usize,
5049 j0: usize,
5050 acc: &mut [T],
5051 step: &F,
5052) -> bool
5053where
5054 T: Copy,
5055 F: Fn(T, T) -> (T, bool),
5056{
5057 if acc.len() < VECTOR_COLUMNS {
5058 fold_range_body(v, m, lo, hi, j0, acc, step)
5059 } else {
5060 fold_range_vectorised(v, m, lo, hi, j0, acc, step)
5061 }
5062}
5063
5064const FOLD_LANES: usize = 8;
5075
5076const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
5081
5082#[inline(always)]
5087fn fold_lanes_body<T, F>(v: &[T], step: &F) -> Option<T>
5088where
5089 T: Copy,
5090 F: Fn(T, T) -> (T, bool),
5091{
5092 let n = v.len();
5093 let mut over = false;
5094 if n < MIN_LANE_WORK {
5095 let mut acc = v[n - 1];
5096 for &x in v[..n - 1].iter().rev() {
5097 let (r, o) = step(x, acc);
5098 acc = r;
5099 over |= o;
5100 }
5101 return (!over).then_some(acc);
5102 }
5103 let rows = n / FOLD_LANES;
5106 let head = n - rows * FOLD_LANES;
5107 let last = head + (rows - 1) * FOLD_LANES;
5108 let mut acc = [v[last]; FOLD_LANES];
5109 acc.copy_from_slice(&v[last..last + FOLD_LANES]);
5110 for r in (0..rows - 1).rev() {
5111 let row = &v[head + r * FOLD_LANES..head + (r + 1) * FOLD_LANES];
5112 for (slot, &x) in acc.iter_mut().zip(row) {
5113 let (r, o) = step(x, *slot);
5114 *slot = r;
5115 over |= o;
5116 }
5117 }
5118 let mut a = acc[FOLD_LANES - 1];
5119 for &x in acc[..FOLD_LANES - 1].iter().rev() {
5120 let (r, o) = step(x, a);
5121 a = r;
5122 over |= o;
5123 }
5124 for &x in v[..head].iter().rev() {
5125 let (r, o) = step(x, a);
5126 a = r;
5127 over |= o;
5128 }
5129 (!over).then_some(a)
5130}
5131
5132multiversioned! {
5133 fn fold_lanes_vectorised[T: Copy, F: Fn(T, T) -> (T, bool)](
5134 v: &[T],
5135 step: &F,
5136 ) -> Option<T> = fold_lanes_body;
5137}
5138
5139#[inline]
5142fn fold_lanes<T, F>(v: &[T], step: &F) -> Option<T>
5143where
5144 T: Copy,
5145 F: Fn(T, T) -> (T, bool),
5146{
5147 if v.len() < MIN_LANE_WORK {
5148 fold_lanes_body(v, step)
5149 } else {
5150 fold_lanes_vectorised(v, step)
5151 }
5152}
5153
5154fn fold_flat<T, F>(v: &[T], n: usize, assoc: bool, step: &F) -> Option<T>
5157where
5158 T: Copy + Send + Sync,
5159 F: Fn(T, T) -> (T, bool) + Sync + Send,
5160{
5161 if assoc {
5162 return par::try_fold_chunks(
5163 &v[..n],
5164 |part| fold_lanes(part, step),
5165 |a, b| {
5166 let (r, o) = step(a, b);
5167 (!o).then_some(r)
5168 },
5169 );
5170 }
5171 let mut acc = v[n - 1];
5172 let mut over = false;
5173 for &x in v[..n - 1].iter().rev() {
5174 let (r, o) = step(x, acc);
5175 acc = r;
5176 over |= o;
5177 }
5178 (!over).then_some(acc)
5179}
5180
5181fn fold_items<T, F>(v: &[T], n: usize, m: usize, assoc: bool, step: F) -> Option<Vec<T>>
5192where
5193 T: Copy + Default + Send + Sync,
5194 F: Fn(T, T) -> (T, bool) + Sync + Send,
5195{
5196 if m >= par::WIDE_ITEM {
5197 let (out, ok) = par::fill_wide(m, n * m, |j0, acc: &mut [T]| {
5198 fold_range(v, m, 0, n, j0, acc, &step)
5199 });
5200 return ok.then_some(out);
5201 }
5202 if m == 1 {
5203 return fold_flat(v, n, assoc, &step).map(|x| vec![x]);
5204 }
5205 let chunks = if assoc { par::chunks(n, n * m) } else { 1 };
5206 if chunks < 2 {
5207 let mut acc = vec![T::default(); m];
5208 return fold_range(v, m, 0, n, 0, &mut acc, &step).then_some(acc);
5209 }
5210 let per = n.div_ceil(chunks);
5211 let parts = par::map_indexed(n.div_ceil(per), |c| {
5212 let mut acc = vec![T::default(); m];
5213 let ok = fold_range(v, m, c * per, ((c + 1) * per).min(n), 0, &mut acc, &step);
5214 ok.then_some(acc)
5215 });
5216 let mut it = parts.into_iter().rev();
5219 let mut acc = it.next()??;
5220 for part in it {
5221 let part = part?;
5222 let mut over = false;
5223 for (slot, &x) in acc.iter_mut().zip(&part) {
5224 let (r, o) = step(x, *slot);
5225 *slot = r;
5226 over |= o;
5227 }
5228 if over {
5229 return None;
5230 }
5231 }
5232 Some(acc)
5233}
5234
5235fn fold_i64(op: ScalarDyad, v: &[i64], n: usize, m: usize) -> Option<Vec<i64>> {
5236 use ScalarDyad::*;
5237 let assoc = is_associative(op);
5238 match op {
5239 Add => fold_items(v, n, m, assoc, i64::overflowing_add),
5240 Sub => fold_items(v, n, m, assoc, i64::overflowing_sub),
5241 Mul => fold_items(v, n, m, assoc, i64::overflowing_mul),
5242 Min => fold_items(v, n, m, assoc, |a: i64, b: i64| (a.min(b), false)),
5243 Max => fold_items(v, n, m, assoc, |a: i64, b: i64| (a.max(b), false)),
5244 _ => None,
5245 }
5246}
5247
5248fn fold_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize) -> Option<Vec<Cx>> {
5249 use ScalarDyad::*;
5250 let assoc = is_associative(op);
5251 match op {
5252 Add => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::add(a, b), false)),
5253 Sub => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::sub(a, b), false)),
5254 Mul => fold_items(v, n, m, assoc, |a: Cx, b: Cx| (cx::mul(a, b), false)),
5255 _ => None,
5257 }
5258}
5259
5260fn fold_f64(op: ScalarDyad, v: &[f64], n: usize, m: usize) -> Option<Vec<f64>> {
5261 use ScalarDyad::*;
5262 let assoc = is_associative(op);
5263 match op {
5264 Add => fold_items(v, n, m, assoc, |a: f64, b: f64| (a + b, false)),
5265 Sub => fold_items(v, n, m, assoc, |a: f64, b: f64| (a - b, false)),
5266 Mul => fold_items(v, n, m, assoc, |a: f64, b: f64| (a * b, false)),
5267 Min => fold_items(v, n, m, assoc, |a: f64, b: f64| (a.min(b), false)),
5268 Max => fold_items(v, n, m, assoc, |a: f64, b: f64| (a.max(b), false)),
5269 _ => None,
5270 }
5271}
5272
5273fn reduce_typed(op: ScalarDyad, d: &Data, n: usize, m: usize) -> Option<Data> {
5277 use ScalarDyad::*;
5278 if !matches!(op, Add | Sub | Mul | Min | Max) {
5281 return None;
5282 }
5283 match d {
5284 Data::F64(v) => Some(Data::F64(fold_f64(op, v, n, m)?.into())),
5285 Data::Complex(v) => Some(Data::Complex(fold_cx(op, v, n, m)?.into())),
5286 Data::I64(v) => Some(Data::I64(fold_i64(op, v, n, m)?.into())),
5287 Data::Bool(v) => {
5290 let widened = par::map(v, |&b| b as i64);
5291 Some(Data::I64(fold_i64(op, &widened, n, m)?.into()))
5292 }
5293 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
5296 }
5297}
5298
5299#[inline(always)]
5305fn fold_runs_body<T, F>(v: &[T], start: usize, m: usize, out: &mut [T], step: &F) -> bool
5306where
5307 T: Copy,
5308 F: Fn(T, T) -> (T, bool),
5309{
5310 let mut over = false;
5311 for (k, slot) in out.iter_mut().enumerate() {
5312 let run = &v[(start + k) * m..(start + k + 1) * m];
5313 let mut acc = run[m - 1];
5314 for &x in run[..m - 1].iter().rev() {
5315 let (r, o) = step(x, acc);
5316 acc = r;
5317 over |= o;
5318 }
5319 *slot = acc;
5320 }
5321 !over
5322}
5323
5324multiversioned! {
5325 fn fold_runs_vectorised[T: Copy, F: Fn(T, T) -> (T, bool)](
5326 v: &[T],
5327 start: usize,
5328 m: usize,
5329 out: &mut [T],
5330 step: &F,
5331 ) -> bool = fold_runs_body;
5332}
5333
5334fn fold_runs<T, F>(v: &[T], n: usize, m: usize, step: F) -> Option<Vec<T>>
5337where
5338 T: Copy + Default + Send + Sync,
5339 F: Fn(T, T) -> (T, bool) + Sync + Send,
5340{
5341 let wide = m >= VECTOR_COLUMNS;
5345 let (out, ok) = par::fill_wide(n, n * m, |start, part: &mut [T]| {
5346 if wide {
5347 fold_runs_vectorised(v, start, m, part, &step)
5348 } else {
5349 fold_runs_body(v, start, m, part, &step)
5350 }
5351 });
5352 ok.then_some(out)
5353}
5354
5355fn fold_runs_data(op: ScalarDyad, d: &Data, n: usize, m: usize) -> Option<Data> {
5356 use ScalarDyad::*;
5357 match d {
5358 Data::F64(v) => Some(Data::F64(
5359 match op {
5360 Add => fold_runs(v, n, m, |a: f64, b: f64| (a + b, false)),
5361 Sub => fold_runs(v, n, m, |a: f64, b: f64| (a - b, false)),
5362 Mul => fold_runs(v, n, m, |a: f64, b: f64| (a * b, false)),
5363 Min => fold_runs(v, n, m, |a: f64, b: f64| (a.min(b), false)),
5364 Max => fold_runs(v, n, m, |a: f64, b: f64| (a.max(b), false)),
5365 _ => None,
5366 }?
5367 .into(),
5368 )),
5369 Data::I64(v) => Some(Data::I64(
5370 match op {
5371 Add => fold_runs(v, n, m, i64::overflowing_add),
5372 Sub => fold_runs(v, n, m, i64::overflowing_sub),
5373 Mul => fold_runs(v, n, m, i64::overflowing_mul),
5374 Min => fold_runs(v, n, m, |a: i64, b: i64| (a.min(b), false)),
5375 Max => fold_runs(v, n, m, |a: i64, b: i64| (a.max(b), false)),
5376 _ => None,
5377 }?
5378 .into(),
5379 )),
5380 Data::Complex(v) => Some(Data::Complex(
5382 match op {
5383 Add => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::add(a, b), false)),
5384 Sub => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::sub(a, b), false)),
5385 Mul => fold_runs(v, n, m, |a: Cx, b: Cx| (cx::mul(a, b), false)),
5386 _ => None,
5387 }?
5388 .into(),
5389 )),
5390 Data::Bool(v) => {
5393 let widened = par::map(v, |&b| b as i64);
5394 fold_runs_data(op, &Data::I64(widened.into()), n, m)
5395 }
5396 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
5397 }
5398}
5399
5400fn reduce_vector_cells(u: &Verb, y: &Array, frame_rank: usize) -> Option<Array> {
5409 let Verb::Reduce(inner) = u else { return None };
5410 let Verb::Prim(p) = &**inner else { return None };
5411 let DyadOp::Scalar(op) = p.dyad else { return None };
5412 if y.rank() != frame_rank + 1 || !y.dtype().is_numeric() {
5415 return None;
5416 }
5417 let m = y.shape[frame_rank];
5418 if m == 0 {
5421 return None;
5422 }
5423 use ScalarDyad::{Add, Max, Min, Mul, Sub};
5424 if !matches!(op, Add | Sub | Mul | Min | Max) {
5425 return None;
5426 }
5427 let frame = y.shape[..frame_rank].to_vec();
5428 if m == 1 {
5429 return Some(Array::new(frame, y.data.clone()));
5432 }
5433 let n: usize = frame.iter().product();
5434 let data = fold_runs_data(op, &y.data, n, m)?;
5435 Some(Array::new(frame, data))
5436}
5437
5438fn reduce(v: &Verb, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
5440 if y.rank() == 0 {
5441 return Ok(y.clone());
5442 }
5443 let n = y.items();
5444 if n == 1 {
5445 return Ok(y.item(0));
5446 }
5447 let cell_shape = y.shape[1..].to_vec();
5448 let m: usize = cell_shape.iter().product();
5449 if n == 0 {
5450 if matches!(v, Verb::Prim(p) if matches!(p.dyad, DyadOp::AppendLeading | DyadOp::AppendLast))
5453 {
5454 return Ok(Array::new(vec![0], Data::empty(y.dtype())));
5455 }
5456 return match reduce_identity(v, m) {
5457 Some(d) => Ok(Array::new(cell_shape, d)),
5458 None => Err(Error::domain(
5459 format!("empty reduction has no identity for {}", v.name()),
5460 span,
5461 )),
5462 };
5463 }
5464 if y.dtype().is_numeric() {
5465 if let Verb::Prim(p) = v {
5466 if let DyadOp::Scalar(op) = p.dyad {
5467 if let Some(d) = reduce_typed(op, &y.data, n, m) {
5472 return Ok(Array::new(cell_shape, d));
5473 }
5474 let mut acc = y.data.slice((n - 1) * m, n * m);
5477 for i in (0..n - 1).rev() {
5478 acc =
5479 scalar_dyad_data(op, &y.data, i * m, 1, &acc, 0, 1, m, ctx.cfg.tol, span)?;
5480 }
5481 return Ok(Array::new(cell_shape, acc));
5482 }
5483 }
5484 }
5485 let mut acc = y.item(n - 1);
5486 for i in (0..n - 1).rev() {
5487 acc = v.dyad(&y.item(i), &acc, ctx, span)?;
5488 }
5489 Ok(acc)
5490}
5491
5492fn folded_op(u: &Verb) -> Option<ScalarDyad> {
5499 let Verb::Reduce(inner) = u else { return None };
5500 let Verb::Prim(p) = &**inner else { return None };
5501 match p.dyad {
5502 DyadOp::Scalar(op) => Some(op),
5503 _ => None,
5504 }
5505}
5506
5507fn section(y: &Array, lo: usize, hi: usize) -> Array {
5509 let m = y.item_size();
5510 let mut shape = y.shape.clone();
5511 shape[0] = hi - lo;
5512 Array::new(shape, y.data.slice(lo * m, hi * m))
5513}
5514
5515fn as_items(y: &Array) -> Option<Array> {
5518 (y.rank() == 0).then(|| Array::new(vec![1], y.data.clone()))
5519}
5520
5521#[inline(always)]
5522fn scan_flat_body<T, F>(v: &[T], n: usize, m: usize, back: bool, step: F) -> Option<Vec<T>>
5523where
5524 T: Copy + Default,
5525 F: Fn(T, T) -> (T, bool),
5526{
5527 if m == 1 {
5528 let mut out = vec![T::default(); n];
5531 let mut over = false;
5532 if back {
5533 let mut acc = v[n - 1];
5534 out[n - 1] = acc;
5535 for (slot, &x) in out[..n - 1].iter_mut().zip(&v[..n - 1]).rev() {
5536 let (r, o) = step(x, acc);
5537 acc = r;
5538 over |= o;
5539 *slot = acc;
5540 }
5541 } else {
5542 let mut acc = v[0];
5543 out[0] = acc;
5544 for (slot, &x) in out[1..n].iter_mut().zip(&v[1..n]) {
5545 let (r, o) = step(acc, x);
5546 acc = r;
5547 over |= o;
5548 *slot = acc;
5549 }
5550 }
5551 return (!over).then_some(out);
5552 }
5553 let mut out = vec![T::default(); n * m];
5554 let mut acc = vec![T::default(); m];
5555 let mut over = false;
5556 if back {
5557 acc.copy_from_slice(&v[(n - 1) * m..n * m]);
5558 out[(n - 1) * m..n * m].copy_from_slice(&acc);
5559 for i in (0..n - 1).rev() {
5560 for (j, slot) in acc.iter_mut().enumerate() {
5561 let (r, o) = step(v[i * m + j], *slot);
5562 *slot = r;
5563 over |= o;
5564 }
5565 out[i * m..i * m + m].copy_from_slice(&acc);
5566 }
5567 } else {
5568 acc.copy_from_slice(&v[..m]);
5569 out[..m].copy_from_slice(&acc);
5570 for i in 1..n {
5571 for (j, slot) in acc.iter_mut().enumerate() {
5572 let (r, o) = step(*slot, v[i * m + j]);
5573 *slot = r;
5574 over |= o;
5575 }
5576 out[i * m..i * m + m].copy_from_slice(&acc);
5577 }
5578 }
5579 (!over).then_some(out)
5580}
5581
5582multiversioned! {
5583 fn scan_flat_vectorised[T: Copy + Default, F: Fn(T, T) -> (T, bool)](
5584 v: &[T],
5585 n: usize,
5586 m: usize,
5587 back: bool,
5588 step: F,
5589 ) -> Option<Vec<T>> = scan_flat_body;
5590}
5591
5592fn scan_flat<T, F>(v: &[T], n: usize, m: usize, back: bool, step: F) -> Option<Vec<T>>
5604where
5605 T: Copy + Default,
5606 F: Fn(T, T) -> (T, bool),
5607{
5608 if m < VECTOR_COLUMNS {
5609 scan_flat_body(v, n, m, back, step)
5610 } else {
5611 scan_flat_vectorised(v, n, m, back, step)
5612 }
5613}
5614
5615fn scan_i64(op: ScalarDyad, v: &[i64], n: usize, m: usize, back: bool) -> Option<Vec<i64>> {
5616 use ScalarDyad::*;
5617 match op {
5618 Add => scan_flat(v, n, m, back, i64::overflowing_add),
5619 Sub => scan_flat(v, n, m, back, i64::overflowing_sub),
5620 Mul => scan_flat(v, n, m, back, i64::overflowing_mul),
5621 Min => scan_flat(v, n, m, back, |a: i64, b: i64| (a.min(b), false)),
5622 Max => scan_flat(v, n, m, back, |a: i64, b: i64| (a.max(b), false)),
5623 _ => None,
5624 }
5625}
5626
5627fn scan_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize, back: bool) -> Option<Vec<Cx>> {
5628 use ScalarDyad::*;
5629 match op {
5630 Add => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::add(a, b), false)),
5631 Sub => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::sub(a, b), false)),
5632 Mul => scan_flat(v, n, m, back, |a: Cx, b: Cx| (cx::mul(a, b), false)),
5633 _ => None,
5634 }
5635}
5636
5637fn scan_f64(op: ScalarDyad, v: &[f64], n: usize, m: usize, back: bool) -> Option<Vec<f64>> {
5638 use ScalarDyad::*;
5639 match op {
5640 Add => scan_flat(v, n, m, back, |a: f64, b: f64| (a + b, false)),
5641 Sub => scan_flat(v, n, m, back, |a: f64, b: f64| (a - b, false)),
5642 Mul => scan_flat(v, n, m, back, |a: f64, b: f64| (a * b, false)),
5643 Min => scan_flat(v, n, m, back, |a: f64, b: f64| (a.min(b), false)),
5644 Max => scan_flat(v, n, m, back, |a: f64, b: f64| (a.max(b), false)),
5645 _ => None,
5646 }
5647}
5648
5649fn scan_typed(op: ScalarDyad, d: &Data, n: usize, m: usize, back: bool) -> Option<Data> {
5653 use ScalarDyad::*;
5654 if !matches!(op, Add | Sub | Mul | Min | Max) {
5655 return None;
5656 }
5657 let widened = |v: &[i64]| {
5658 let f: Vec<f64> = v.iter().map(|&x| x as f64).collect();
5659 Data::F64(scan_f64(op, &f, n, m, back).expect("the float scan cannot overflow").into())
5660 };
5661 let ints = |v: &[i64]| match scan_i64(op, v, n, m, back) {
5662 Some(out) => Data::I64(out.into()),
5663 None => widened(v),
5664 };
5665 match d {
5666 Data::F64(v) => Some(Data::F64(scan_f64(op, v, n, m, back)?.into())),
5667 Data::Complex(v) => Some(Data::Complex(scan_cx(op, v, n, m, back)?.into())),
5668 Data::I64(v) => Some(ints(v)),
5669 Data::Bool(v) => Some(ints(&v.iter().map(|&b| b as i64).collect::<Vec<_>>())),
5670 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
5673 }
5674}
5675
5676fn window_fold<T, F>(v: &[T], n: usize, m: usize, w: usize, step: F) -> Option<Vec<T>>
5691where
5692 T: Copy + Default + Send + Sync,
5693 F: Fn(T, T) -> (T, bool) + Sync + Send,
5694{
5695 debug_assert!(w >= 1 && n >= w);
5696 if m == 1 {
5697 return window_fold_flat(v, n, w, step);
5698 }
5699 let count = n - w + 1;
5700 let mut out = vec![T::default(); count * m];
5701 let mut pre = vec![T::default(); w * m];
5704 let mut suf = vec![T::default(); w * m];
5705 let mut prev = vec![T::default(); w * m];
5706 let mut over = false;
5707 for b in 0..n.div_ceil(w) {
5708 let bs = b * w;
5709 let be = ((b + 1) * w).min(n);
5710 pre[..m].copy_from_slice(&v[bs * m..bs * m + m]);
5711 for i in 1..be - bs {
5712 let (o, p) = (i * m, (i - 1) * m);
5713 for j in 0..m {
5714 let (r, f) = step(pre[p + j], v[(bs + i) * m + j]);
5715 pre[o + j] = r;
5716 over |= f;
5717 }
5718 }
5719 for e in bs.max(w - 1)..be {
5722 let i = e + 1 - w;
5723 let (oo, po) = (i * m, (e - bs) * m);
5724 if i == bs {
5725 out[oo..oo + m].copy_from_slice(&pre[po..po + m]);
5726 } else {
5727 let so = (i + w - bs) * m;
5728 for j in 0..m {
5729 let (r, f) = step(prev[so + j], pre[po + j]);
5730 out[oo + j] = r;
5731 over |= f;
5732 }
5733 }
5734 }
5735 let last = be - 1 - bs;
5736 suf[last * m..last * m + m].copy_from_slice(&v[(be - 1) * m..be * m]);
5737 for i in (0..last).rev() {
5738 let (o, p) = (i * m, (i + 1) * m);
5739 for j in 0..m {
5740 let (r, f) = step(v[(bs + i) * m + j], suf[p + j]);
5741 suf[o + j] = r;
5742 over |= f;
5743 }
5744 }
5745 std::mem::swap(&mut prev, &mut suf);
5746 }
5747 (!over).then_some(out)
5748}
5749
5750fn window_fold_flat<T, F>(v: &[T], n: usize, w: usize, step: F) -> Option<Vec<T>>
5760where
5761 T: Copy + Default + Send + Sync,
5762 F: Fn(T, T) -> (T, bool) + Sync + Send,
5763{
5764 let (out, ok) = par::fill(n - w + 1, |lo, part: &mut [T]| {
5765 window_fold_range(v, n, w, lo, part, &step)
5766 });
5767 ok.then_some(out)
5768}
5769
5770#[inline(always)]
5771fn window_fold_range_body<T, F>(
5772 v: &[T],
5773 n: usize,
5774 w: usize,
5775 lo: usize,
5776 out: &mut [T],
5777 step: &F,
5778) -> bool
5779where
5780 T: Copy + Default,
5781 F: Fn(T, T) -> (T, bool),
5782{
5783 if out.is_empty() {
5784 return true;
5785 }
5786 let hi = lo + out.len();
5787 let mut pre = vec![T::default(); w];
5788 let mut suf = vec![T::default(); w];
5789 let mut prev = vec![T::default(); w];
5790 let mut over = false;
5791 let mut bs = lo / w * w;
5792 while bs < n && bs <= hi + w - 2 {
5794 let block = &v[bs..(bs + w).min(n)];
5795 let lb = block.len();
5796 let mut acc = block[0];
5797 pre[0] = acc;
5798 for (slot, &x) in pre[1..lb].iter_mut().zip(&block[1..]) {
5799 let (r, o) = step(acc, x);
5800 acc = r;
5801 over |= o;
5802 *slot = acc;
5803 }
5804 for e in bs.max(lo + w - 1)..(bs + lb).min(hi + w - 1) {
5809 let i = e + 1 - w;
5810 out[i - lo] = if i == bs {
5811 pre[e - bs]
5812 } else {
5813 let (r, o) = step(prev[i + w - bs], pre[e - bs]);
5814 over |= o;
5815 r
5816 };
5817 }
5818 let mut acc = block[lb - 1];
5819 suf[lb - 1] = acc;
5820 for (slot, &x) in suf[..lb - 1].iter_mut().zip(&block[..lb - 1]).rev() {
5821 let (r, o) = step(x, acc);
5822 acc = r;
5823 over |= o;
5824 *slot = acc;
5825 }
5826 std::mem::swap(&mut prev, &mut suf);
5827 bs += w;
5828 }
5829 !over
5830}
5831
5832multiversioned! {
5833 fn window_fold_range[T: Copy + Default, F: Fn(T, T) -> (T, bool)](
5838 v: &[T],
5839 n: usize,
5840 w: usize,
5841 lo: usize,
5842 out: &mut [T],
5843 step: &F,
5844 ) -> bool = window_fold_range_body;
5845}
5846
5847fn window_i64(op: ScalarDyad, v: &[i64], n: usize, m: usize, w: usize) -> Option<Vec<i64>> {
5848 use ScalarDyad::*;
5849 match op {
5850 Add => window_fold(v, n, m, w, i64::overflowing_add),
5851 Mul => window_fold(v, n, m, w, i64::overflowing_mul),
5852 Min => window_fold(v, n, m, w, |a: i64, b: i64| (a.min(b), false)),
5853 Max => window_fold(v, n, m, w, |a: i64, b: i64| (a.max(b), false)),
5854 _ => None,
5855 }
5856}
5857
5858fn window_cx(op: ScalarDyad, v: &[Cx], n: usize, m: usize, w: usize) -> Option<Vec<Cx>> {
5859 use ScalarDyad::*;
5860 match op {
5861 Add => window_fold(v, n, m, w, |a: Cx, b: Cx| (cx::add(a, b), false)),
5862 Mul => window_fold(v, n, m, w, |a: Cx, b: Cx| (cx::mul(a, b), false)),
5863 _ => None,
5864 }
5865}
5866
5867fn window_f64(op: ScalarDyad, v: &[f64], n: usize, m: usize, w: usize) -> Option<Vec<f64>> {
5868 use ScalarDyad::*;
5869 match op {
5870 Add => window_fold(v, n, m, w, |a: f64, b: f64| (a + b, false)),
5871 Mul => window_fold(v, n, m, w, |a: f64, b: f64| (a * b, false)),
5872 Min => window_fold(v, n, m, w, |a: f64, b: f64| (a.min(b), false)),
5873 Max => window_fold(v, n, m, w, |a: f64, b: f64| (a.max(b), false)),
5874 _ => None,
5875 }
5876}
5877
5878fn window_typed(op: ScalarDyad, d: &Data, n: usize, m: usize, w: usize) -> Option<Data> {
5882 use ScalarDyad::*;
5883 if !matches!(op, Add | Mul | Min | Max) {
5884 return None;
5885 }
5886 let widened = |v: &[i64]| {
5887 let f: Vec<f64> = v.iter().map(|&x| x as f64).collect();
5888 Data::F64(window_f64(op, &f, n, m, w).expect("the float fold cannot overflow").into())
5889 };
5890 let ints = |v: &[i64]| match window_i64(op, v, n, m, w) {
5891 Some(out) => Data::I64(out.into()),
5892 None => widened(v),
5893 };
5894 match d {
5895 Data::F64(v) => Some(Data::F64(window_f64(op, v, n, m, w)?.into())),
5896 Data::Complex(v) => Some(Data::Complex(window_cx(op, v, n, m, w)?.into())),
5897 Data::I64(v) => Some(ints(v)),
5898 Data::Bool(v) => Some(ints(&v.iter().map(|&b| b as i64).collect::<Vec<_>>())),
5899 Data::Ext(_) | Data::Rat(_) | Data::Char(_) | Data::Box(_) => None,
5902 }
5903}
5904
5905fn runs(u: &Verb, y: &Array, back: bool, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
5907 let promoted = as_items(y);
5908 let base = promoted.as_ref().unwrap_or(y);
5909 let n = base.items();
5910 let m = base.item_size();
5911 if n > 0 && base.dtype().is_numeric() {
5912 if let Some(op) = folded_op(u) {
5913 if back || is_associative(op) {
5916 if let Some(d) = scan_typed(op, &base.data, n, m, back) {
5917 return Ok(Array::new(base.shape.clone(), d));
5918 }
5919 }
5920 }
5921 }
5922 let cells = each_cell(n, n * m, u.is_pure(), ctx, |i, c| {
5923 let part = if back { section(base, i, n) } else { section(base, 0, i + 1) };
5924 u.monad(&part, c, span)
5925 })?;
5926 assemble(&[n], cells, span)
5927}
5928
5929fn empty_windows(u: &Verb, y: &Array, w: usize, ctx: &mut Ctx<'_>, span: Span) -> Array {
5934 let m = y.item_size();
5935 if u.is_pure() {
5936 if let Some(cells) = w.checked_mul(m).filter(|&s| s <= 1 << 20) {
5937 let mut shape = y.shape.clone();
5938 shape[0] = w;
5939 let probe = Array::new(shape, fill_data(y.dtype(), cells));
5940 if let Ok(cell) = u.monad(&probe, ctx, span) {
5941 let mut shape = vec![0usize];
5942 shape.extend_from_slice(&cell.shape);
5943 return Array::new(shape, Data::empty(cell.dtype()));
5944 }
5945 }
5946 }
5947 Array::new(vec![0], Data::empty(DType::I64))
5948}
5949
5950fn window_size(x: &Array, span: Span) -> Result<i64> {
5952 let v = x
5953 .to_i64_vec()
5954 .ok_or_else(|| Error::domain("the window size must be an integer", span))?;
5955 match v.as_slice() {
5956 [k] => Ok(*k),
5957 _ => Err(Error::new(
5958 ErrorKind::Length,
5959 "the window size must be a single number",
5960 Some(span),
5961 )),
5962 }
5963}
5964
5965fn infix(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
5973 let k = window_size(x, span)?;
5974 let promoted = as_items(y);
5975 let base = promoted.as_ref().unwrap_or(y);
5976 let n = base.items();
5977 let m = base.item_size();
5978 if k < 0 {
5979 let w = k.unsigned_abs() as usize;
5980 let count = n.div_ceil(w);
5981 let cells = each_cell(count, n * m, u.is_pure(), ctx, |i, c| {
5982 u.monad(§ion(base, i * w, ((i + 1) * w).min(n)), c, span)
5983 })?;
5984 return assemble(&[count], cells, span);
5985 }
5986 let w = k as usize;
5987 if n < w {
5988 return Ok(empty_windows(u, base, w, ctx, span));
5989 }
5990 let count = n - w + 1;
5991 if w > 0 && base.dtype().is_numeric() {
5992 if let Some(op) = folded_op(u) {
5993 if let Some(d) = window_typed(op, &base.data, n, m, w) {
5994 let mut shape = base.shape.clone();
5995 shape[0] = count;
5996 return Ok(Array::new(shape, d));
5997 }
5998 }
5999 }
6000 let work = count.saturating_mul(w).saturating_mul(m);
6001 let cells = each_cell(count, work, u.is_pure(), ctx, |i, c| {
6002 u.monad(§ion(base, i, i + w), c, span)
6003 })?;
6004 assemble(&[count], cells, span)
6005}
6006
6007fn power(
6010 u: &Verb,
6011 p: Power,
6012 x: Option<&Array>,
6013 y: &Array,
6014 ctx: &mut Ctx<'_>,
6015 span: Span,
6016) -> Result<Array> {
6017 let step = |acc: &Array, c: &mut Ctx<'_>| match x {
6018 Some(x) => u.dyad(x, acc, c, span),
6019 None => u.monad(acc, c, span),
6020 };
6021 match p {
6022 Power::Times(n) => {
6023 let mut acc = y.clone();
6024 for _ in 0..n {
6025 acc = step(&acc, ctx)?;
6026 }
6027 Ok(acc)
6028 }
6029 Power::Converge => {
6030 let mut acc = y.clone();
6031 for _ in 0..CONVERGE_LIMIT {
6032 let next = step(&acc, ctx)?;
6033 if arrays_match(&next, &acc, ctx.cfg.tol) {
6034 return Ok(next);
6035 }
6036 acc = next;
6037 }
6038 Err(Error::domain("the iteration did not converge", span))
6039 }
6040 }
6041}
6042
6043fn power_v(
6046 u: &Verb,
6047 v: &Verb,
6048 x: Option<&Array>,
6049 y: &Array,
6050 ctx: &mut Ctx<'_>,
6051 span: Span,
6052) -> Result<Array> {
6053 let count = match x {
6054 Some(x) => v.dyad(x, y, ctx, span)?,
6055 None => v.monad(y, ctx, span)?,
6056 };
6057 let n = count
6058 .to_i64_vec()
6059 .ok_or_else(|| Error::domain("the power count must be an integer", span))?;
6060 if n.len() != 1 {
6061 return Err(Error::not_yet("a list of power counts (u^:v with several)", span));
6062 }
6063 let n = n[0];
6064 if n < 0 {
6065 return Err(Error::not_yet("a negative power (the verb's inverse)", span));
6066 }
6067 power(u, Power::Times(n as u64), x, y, ctx, span)
6068}
6069
6070fn power_until(
6072 u: &Verb,
6073 test: &Verb,
6074 y: &Array,
6075 ctx: &mut Ctx<'_>,
6076 span: Span,
6077) -> Result<Array> {
6078 let mut acc = y.clone();
6079 for _ in 0..CONVERGE_LIMIT {
6080 let next = u.monad(&acc, ctx, span)?;
6081 let done = test.dyad(&next, &acc, ctx, span)?;
6082 let stop = done
6083 .to_f64_vec()
6084 .ok_or_else(|| Error::domain("the ⍣ test must answer with numbers", span))?;
6085 if !stop.is_empty() && stop.iter().all(|&v| v != 0.0) {
6086 return Ok(next);
6087 }
6088 acc = next;
6089 }
6090 Err(Error::domain("the iteration did not converge", span))
6091}
6092
6093fn along_axis(
6100 u: &Verb,
6101 x: Option<&Array>,
6102 y: &Array,
6103 k: usize,
6104 ctx: &mut Ctx<'_>,
6105 span: Span,
6106) -> Result<Array> {
6107 if k >= y.rank().max(1) {
6108 return Err(Error::new(
6109 ErrorKind::Rank,
6110 format!("axis {k} does not exist on an argument of rank {}", y.rank()),
6111 Some(span),
6112 ));
6113 }
6114 let moved = axis_to_front(y, k);
6115 let r = moved.rank();
6116 let out = match x {
6117 Some(x) => u.dyad(x, &moved, ctx, span)?,
6118 None => u.monad(&moved, ctx, span)?,
6119 };
6120 if out.rank() == r {
6121 return Ok(front_to_axis(&out, k));
6122 }
6123 Ok(out)
6124}
6125
6126fn where_indices(y: &Array, origin: i64, boxed: bool, span: Span) -> Result<Array> {
6134 let counts = y
6135 .to_i64_vec()
6136 .ok_or_else(|| Error::domain("indices needs non-negative integers", span))?;
6137 if counts.iter().any(|&c| c < 0) {
6138 return Err(Error::domain("indices needs non-negative integers", span));
6139 }
6140 if !boxed || y.rank() < 2 {
6141 let mut out = Vec::new();
6142 for (i, &c) in counts.iter().enumerate() {
6143 for _ in 0..c {
6144 out.push(origin + i as i64);
6145 }
6146 }
6147 return Ok(Array::from_i64(out));
6148 }
6149 let r = y.rank();
6150 let mut coord = vec![0usize; r];
6151 let mut out: Vec<Array> = Vec::new();
6152 for &c in &counts {
6153 if c > 0 {
6154 let point =
6155 Array::from_i64(coord.iter().map(|&k| origin + k as i64).collect::<Vec<_>>());
6156 for _ in 0..c {
6157 out.push(point.clone());
6158 }
6159 }
6160 odometer(&mut coord, &y.shape);
6161 }
6162 Ok(Array::new(vec![out.len()], Data::Box(out.into())))
6163}
6164
6165fn interval_index(x: &Array, y: &Array, offset: i64, tol: Tol, span: Span) -> Result<Array> {
6171 let bounds = x
6172 .to_f64_vec()
6173 .ok_or_else(|| Error::domain("interval index needs numeric bounds", span))?;
6174 let vals = y
6175 .to_f64_vec()
6176 .ok_or_else(|| Error::domain("interval index needs numeric values", span))?;
6177 let out: Vec<i64> = vals
6178 .iter()
6179 .map(|&v| offset + bounds.iter().filter(|&&b| tol.lt(b, v)).count() as i64)
6180 .collect();
6181 Ok(Array::new(y.shape.clone(), Data::I64(out.into())))
6182}
6183
6184fn steps(y: &Array, span: Span) -> Result<Array> {
6187 let vals = y.to_f64_vec().ok_or_else(|| Error::domain("steps needs a number", span))?;
6188 let v = match vals.first() {
6189 Some(&v) if v.is_finite() => v,
6190 _ => return Err(Error::domain("steps needs a finite number", span)),
6191 };
6192 let n = (2.0 * v.abs()).floor();
6193 if n > 1e7 {
6194 return Err(Error::domain("steps would produce too many items", span));
6195 }
6196 let n = n as i64 + 1;
6197 let step = if v < 0.0 { -1.0 } else { 1.0 };
6198 let start = -v;
6199 if v.fract() == 0.0 {
6200 let start = start as i64;
6201 let step = step as i64;
6202 return Ok(Array::from_i64((0..n).map(|k| start + k * step).collect()));
6203 }
6204 Ok(Array::from_f64((0..n).map(|k| start + k as f64 * step).collect()))
6205}
6206
6207fn index_of_last(x: &Array, y: &Array, origin: i64, tol: Tol) -> Array {
6209 let cell_rank = x.rank().saturating_sub(1).min(y.rank());
6210 let frame_rank = y.rank() - cell_rank;
6211 let frame: Vec<usize> = y.shape[..frame_rank].to_vec();
6212 let nf: usize = frame.iter().product();
6213 let items = x.items();
6214 let mut out = Vec::with_capacity(nf);
6215 for i in 0..nf {
6216 let cell = y.cell_at(frame_rank, i);
6217 let at = (0..items)
6218 .rev()
6219 .find(|&j| arrays_match(&cell, &item_or_self(x, j), tol))
6220 .unwrap_or(items);
6221 out.push(origin + at as i64);
6222 }
6223 Array::new(frame, Data::I64(out.into()))
6224}
6225
6226fn roll(
6233 y: &Array,
6234 origin: i64,
6235 fixed: bool,
6236 float_at_zero: bool,
6237 span: Span,
6238) -> Result<Array> {
6239 let bounds = y
6240 .to_i64_vec()
6241 .ok_or_else(|| Error::domain("roll needs whole numbers", span))?;
6242 if bounds.iter().any(|&b| b < 0) {
6243 return Err(Error::domain("roll needs non-negative numbers", span));
6244 }
6245 if !float_at_zero && bounds.iter().any(|&b| b == 0) {
6246 return Err(Error::domain("? 0 has no value: the range is empty", span));
6247 }
6248 let any_zero = bounds.contains(&0);
6250 crate::rng::with(fixed, |g| {
6251 if any_zero {
6252 let out: Vec<f64> = bounds
6253 .iter()
6254 .map(|&b| {
6255 if b == 0 {
6256 g.unit()
6257 } else {
6258 (origin + g.below(b as u64) as i64) as f64
6259 }
6260 })
6261 .collect();
6262 return Ok(Array::new(y.shape.clone(), Data::F64(out.into())));
6263 }
6264 let out: Vec<i64> =
6265 bounds.iter().map(|&b| origin + g.below(b as u64) as i64).collect();
6266 Ok(Array::new(y.shape.clone(), Data::I64(out.into())))
6267 })
6268}
6269
6270fn deal(x: &Array, y: &Array, origin: i64, fixed: bool, span: Span) -> Result<Array> {
6272 let want = one_whole(x, "the count dealt", span)?;
6273 let from = one_whole(y, "the range dealt from", span)?;
6274 if want < 0 || from < 0 {
6275 return Err(Error::domain("deal needs non-negative numbers", span));
6276 }
6277 if want > from {
6278 return Err(Error::domain(
6279 format!("cannot deal {want} distinct value(s) from {from}"),
6280 span,
6281 ));
6282 }
6283 if want == 0 {
6284 return Ok(Array::from_i64(Vec::new()));
6285 }
6286 let drawn = crate::rng::with(fixed, |g| g.deal(want as usize, from as u64));
6287 Ok(Array::from_i64(drawn.into_iter().map(|v| v + origin).collect()))
6288}
6289
6290fn one_whole(a: &Array, what: &str, span: Span) -> Result<i64> {
6292 let v = a
6293 .to_i64_vec()
6294 .ok_or_else(|| Error::domain(format!("{what} must be a whole number"), span))?;
6295 match v[..] {
6296 [n] => Ok(n),
6297 _ => Err(Error::new(
6298 ErrorKind::Rank,
6299 format!("{what} must be one number"),
6300 Some(span),
6301 )),
6302 }
6303}
6304
6305fn nth_prime(n: i64, span: Span) -> Result<i64> {
6309 if n < 0 {
6310 return Err(Error::domain("the prime index must not be negative", span));
6311 }
6312 const LIMIT: i64 = 5_000_000;
6313 if n >= LIMIT {
6314 return Err(Error::domain(
6315 format!("prime index {n} is beyond the {LIMIT}th prime"),
6316 span,
6317 ));
6318 }
6319 let k = (n + 1) as f64;
6322 let bound = if n < 6 { 15.0 } else { k * (k.ln() + k.ln().ln()) };
6323 let bound = bound.ceil() as usize + 1;
6324 let mut sieve = vec![true; bound + 1];
6325 sieve[0] = false;
6326 if bound >= 1 {
6327 sieve[1] = false;
6328 }
6329 let mut p = 2usize;
6330 while p * p <= bound {
6331 if sieve[p] {
6332 let mut q = p * p;
6333 while q <= bound {
6334 sieve[q] = false;
6335 q += p;
6336 }
6337 }
6338 p += 1;
6339 }
6340 let mut seen = 0i64;
6341 for (v, &is_p) in sieve.iter().enumerate() {
6342 if is_p {
6343 if seen == n {
6344 return Ok(v as i64);
6345 }
6346 seen += 1;
6347 }
6348 }
6349 Err(Error::internal("the prime sieve was too small"))
6350}
6351
6352fn prime_factors(n: i64, span: Span) -> Result<Vec<i64>> {
6354 if n < 1 {
6355 return Err(Error::domain("prime factors need a positive integer", span));
6356 }
6357 let mut out = Vec::new();
6358 let mut m = n;
6359 let mut d = 2i64;
6360 while d.saturating_mul(d) <= m {
6361 while m % d == 0 {
6362 out.push(d);
6363 m /= d;
6364 }
6365 d += if d == 2 { 1 } else { 2 };
6366 }
6367 if m > 1 {
6368 out.push(m);
6369 }
6370 Ok(out)
6371}
6372
6373fn lstsq(a: &[f64], m: usize, n: usize, b: &[f64], k: usize) -> Option<Vec<f64>> {
6381 let mut r = a.to_vec();
6383 let mut c = b.to_vec();
6384 let at = |i: usize, j: usize, w: usize| i * w + j;
6385 let scale = a.iter().fold(0.0f64, |acc, v| acc.max(v.abs()));
6386 if scale == 0.0 {
6387 return None;
6388 }
6389 for j in 0..n {
6390 let norm = (j..m).map(|i| r[at(i, j, n)] * r[at(i, j, n)]).sum::<f64>().sqrt();
6392 if norm <= 1e-13 * scale {
6393 return None;
6394 }
6395 let alpha = if r[at(j, j, n)] > 0.0 { -norm } else { norm };
6396 let mut v = vec![0.0f64; m];
6397 for i in j..m {
6398 v[i] = r[at(i, j, n)];
6399 }
6400 v[j] -= alpha;
6401 let vnorm2: f64 = (j..m).map(|i| v[i] * v[i]).sum();
6402 if vnorm2 > 0.0 {
6403 for col in j..n {
6404 let dot: f64 = (j..m).map(|i| v[i] * r[at(i, col, n)]).sum();
6405 let f = 2.0 * dot / vnorm2;
6406 for i in j..m {
6407 r[at(i, col, n)] -= f * v[i];
6408 }
6409 }
6410 for col in 0..k {
6411 let dot: f64 = (j..m).map(|i| v[i] * c[at(i, col, k)]).sum();
6412 let f = 2.0 * dot / vnorm2;
6413 for i in j..m {
6414 c[at(i, col, k)] -= f * v[i];
6415 }
6416 }
6417 }
6418 }
6419 let mut x = vec![0.0f64; n * k];
6421 for col in 0..k {
6422 for i in (0..n).rev() {
6423 let mut acc = c[at(i, col, k)];
6424 for j in i + 1..n {
6425 acc -= r[at(i, j, n)] * x[at(j, col, k)];
6426 }
6427 let d = r[at(i, i, n)];
6428 if d.abs() <= 1e-13 * scale {
6429 return None;
6430 }
6431 x[at(i, col, k)] = acc / d;
6432 }
6433 }
6434 Some(x)
6435}
6436
6437fn as_matrix(a: &Array, span: Span) -> Result<(Vec<f64>, usize, usize)> {
6440 let v = a
6441 .to_f64_vec()
6442 .ok_or_else(|| Error::domain("matrix division needs numeric data", span))?;
6443 match a.rank() {
6444 0 => Ok((v, 1, 1)),
6445 1 => {
6446 let m = a.shape[0];
6447 Ok((v, m, 1))
6448 }
6449 2 => Ok((v, a.shape[0], a.shape[1])),
6450 _ => Err(Error::new(
6451 ErrorKind::Rank,
6452 "matrix division needs an argument of rank 2 or less",
6453 Some(span),
6454 )),
6455 }
6456}
6457
6458fn matrix_inverse(y: &Array, span: Span) -> Result<Array> {
6462 let (a, m, n) = as_matrix(y, span)?;
6463 if m < n {
6464 return Err(Error::new(
6465 ErrorKind::Length,
6466 format!("cannot invert a {m} by {n} matrix: it has more columns than rows"),
6467 Some(span),
6468 ));
6469 }
6470 let mut eye = vec![0.0f64; m * m];
6471 for i in 0..m {
6472 eye[i * m + i] = 1.0;
6473 }
6474 let x = lstsq(&a, m, n, &eye, m)
6475 .ok_or_else(|| Error::domain("the matrix is singular", span))?;
6476 let shape = if y.rank() == 2 { vec![n, m] } else { y.shape.clone() };
6479 Ok(Array::new(shape, Data::F64(x.into())))
6480}
6481
6482fn matrix_divide(x: &Array, y: &Array, span: Span) -> Result<Array> {
6484 let (a, m, n) = as_matrix(y, span)?;
6485 let (b, bm, k) = as_matrix(x, span)?;
6486 if bm != m {
6487 return Err(Error::new(
6488 ErrorKind::Length,
6489 format!("the system has {m} rows but the right-hand side has {bm}"),
6490 Some(span),
6491 ));
6492 }
6493 if m < n {
6494 return Err(Error::new(
6495 ErrorKind::Length,
6496 format!("the {m} by {n} system is underdetermined"),
6497 Some(span),
6498 ));
6499 }
6500 let sol = lstsq(&a, m, n, &b, k)
6501 .ok_or_else(|| Error::domain("the system is singular", span))?;
6502 let shape = if x.rank() == 2 { vec![n, k] } else { vec![n] };
6505 Ok(Array::new(shape, Data::F64(sol.into())))
6506}
6507
6508fn squad(x: &Array, y: &Array, origin: i64, span: Span) -> Result<Array> {
6512 if x.rank() > 1 {
6513 return Err(Error::new(
6514 ErrorKind::Rank,
6515 "the index of ⌷ must be a scalar or a vector",
6516 Some(span),
6517 ));
6518 }
6519 let idx = x
6520 .to_i64_vec()
6521 .ok_or_else(|| Error::domain("index must be an integer", span))?;
6522 if idx.len() != y.rank() {
6523 return Err(Error::new(
6524 ErrorKind::Rank,
6525 format!("{} index(es) for an argument of rank {}", idx.len(), y.rank()),
6526 Some(span),
6527 ));
6528 }
6529 let st = strides(&y.shape);
6530 let mut at = 0usize;
6531 for (k, &i) in idx.iter().enumerate() {
6532 let j = i - origin;
6533 if j < 0 || j as usize >= y.shape[k] {
6534 return Err(Error::domain(
6535 format!("index {i} is out of range on axis {k}"),
6536 span,
6537 ));
6538 }
6539 at += j as usize * st[k];
6540 }
6541 Ok(atom(y, at))
6542}
6543
6544fn select_axis(
6551 x: &Array,
6552 y: &Array,
6553 axis: usize,
6554 rank: usize,
6555 origin: i64,
6556 span: Span,
6557) -> Result<Array> {
6558 if rank != 0 && y.rank() != rank {
6559 return Err(Error::new(
6560 ErrorKind::Rank,
6561 format!("{rank} index slot(s) for an argument of rank {}", y.rank()),
6562 Some(span),
6563 ));
6564 }
6565 if axis >= y.rank() {
6566 return Err(Error::new(
6567 ErrorKind::Rank,
6568 format!("axis {axis} does not exist on an argument of rank {}", y.rank()),
6569 Some(span),
6570 ));
6571 }
6572 let idx = x
6573 .to_i64_vec()
6574 .ok_or_else(|| Error::domain("index must be an integer", span))?;
6575 let len = y.shape[axis];
6576 let mut picks = Vec::with_capacity(idx.len());
6577 for &i in &idx {
6578 let j = i - origin;
6579 if j < 0 || j as usize >= len {
6580 return Err(Error::domain(
6581 format!("index {i} is out of range: axis {axis} has {len} items"),
6582 span,
6583 ));
6584 }
6585 picks.push(j as usize);
6586 }
6587 let mut shape = Vec::with_capacity(y.rank() + x.rank());
6588 shape.extend_from_slice(&y.shape[..axis]);
6589 shape.extend_from_slice(&x.shape);
6590 shape.extend_from_slice(&y.shape[axis + 1..]);
6591 let outer: usize = y.shape[..axis].iter().product();
6592 let inner: usize = y.shape[axis + 1..].iter().product();
6593 let mut data = Data::empty(y.dtype());
6594 for o in 0..outer {
6595 for &p in &picks {
6596 let base = (o * len + p) * inner;
6597 for e in 0..inner {
6598 push_elem(&mut data, &y.data, base + e);
6599 }
6600 }
6601 }
6602 Ok(Array::new(shape, data))
6603}
6604
6605fn amend(m: &Array, x: &Array, y: &Array, span: Span) -> Result<Array> {
6609 if y.rank() == 0 {
6610 return Err(Error::new(ErrorKind::Rank, "cannot amend a scalar", Some(span)));
6611 }
6612 if let Some(spec) = m.as_boxes().and_then(<[Array]>::first) {
6614 let spec = index_spec(spec, y, span)?;
6615 return amend_spec(&spec, x, y, span);
6616 }
6617 let idx = m
6618 .to_i64_vec()
6619 .ok_or_else(|| Error::domain("amend indices must be integers", span))?;
6620 let items = y.items() as i64;
6621 let mut at = Vec::with_capacity(idx.len());
6622 for &i in &idx {
6623 let k = if i < 0 { i + items } else { i };
6624 if k < 0 || k >= items {
6625 return Err(Error::domain(
6626 format!("index {i} is out of range: the argument has {items} items"),
6627 span,
6628 ));
6629 }
6630 at.push(k as usize);
6631 }
6632 let cell = y.item_size();
6633 let per_index = if x.count() == cell {
6634 false
6635 } else if x.count() == cell * at.len() {
6636 true
6637 } else {
6638 return Err(Error::new(
6639 ErrorKind::Length,
6640 format!(
6641 "cannot amend {} item(s) of {} element(s) each with {} element(s)",
6642 at.len(),
6643 cell,
6644 x.count()
6645 ),
6646 Some(span),
6647 ));
6648 };
6649 let Some(t) = DType::promote(x.dtype(), y.dtype()) else {
6652 return Err(Error::new(
6653 ErrorKind::Type,
6654 "the replacement and the argument hold different kinds of value",
6655 Some(span),
6656 ));
6657 };
6658 let (Some(src), Some(base)) = (x.data.cast(t), y.data.cast(t)) else {
6659 return Err(Error::new(
6660 ErrorKind::Type,
6661 "the replacement and the argument hold different kinds of value",
6662 Some(span),
6663 ));
6664 };
6665 let mut data = Data::empty(t);
6667 let mut plan: Vec<Option<usize>> = vec![None; y.items()];
6668 for (n, &k) in at.iter().enumerate() {
6669 plan[k] = Some(if per_index { n } else { 0 });
6670 }
6671 for (i, slot) in plan.iter().enumerate() {
6672 match slot {
6673 Some(n) => {
6674 for e in 0..cell {
6675 push_elem(&mut data, &src, n * cell + e);
6676 }
6677 }
6678 None => {
6679 for e in 0..cell {
6680 push_elem(&mut data, &base, i * cell + e);
6681 }
6682 }
6683 }
6684 }
6685 Ok(Array::new(y.shape.clone(), data))
6686}
6687
6688fn fetch(x: &Array, y: &Array, span: Span) -> Result<Array> {
6693 let steps: Vec<Array> = match x.as_boxes() {
6694 Some(bs) => bs.to_vec(),
6695 None => vec![x.clone()],
6696 };
6697 let mut cur = y.clone();
6698 for step in steps {
6699 let idx = if step.count() == 0 {
6702 Vec::new()
6703 } else {
6704 step.to_i64_vec()
6705 .ok_or_else(|| Error::domain("a fetch path holds integers", span))?
6706 };
6707 let base =
6709 if cur.rank() == 0 { Array::new(vec![1], cur.data.clone()) } else { cur.clone() };
6710 if idx.len() > base.rank() {
6711 return Err(Error::new(
6712 ErrorKind::Length,
6713 format!(
6714 "a path step of {} index(es) into a value of rank {}",
6715 idx.len(),
6716 cur.rank()
6717 ),
6718 Some(span),
6719 ));
6720 }
6721 let at = cell_index(&base, &idx, span)?;
6722 cur = open_cell(&base.cell_at(idx.len(), at));
6723 }
6724 Ok(cur)
6725}
6726
6727fn cell_index(y: &Array, idx: &[i64], span: Span) -> Result<usize> {
6729 let mut at = 0usize;
6730 for (k, &i) in idx.iter().enumerate() {
6731 let len = y.shape[k] as i64;
6732 let j = if i < 0 { i + len } else { i };
6733 if j < 0 || j >= len {
6734 return Err(Error::domain(
6735 format!("index {i} is out of range: axis {k} has {len} items"),
6736 span,
6737 ));
6738 }
6739 at = at * y.shape[k] + j as usize;
6740 }
6741 Ok(at)
6742}
6743
6744fn partition_enclose(x: &Array, y: &Array, span: Span) -> Result<Array> {
6753 if y.rank() > 1 {
6756 let last = y.shape[y.rank() - 1];
6757 let rows = y.count() / last.max(1);
6758 let mut cells: Vec<Array> = Vec::new();
6759 let mut width = None;
6760 for r in 0..rows {
6761 let row = Array::new(vec![last], y.data.slice(r * last, (r + 1) * last));
6762 let parts = partition_enclose(x, &row, span)?;
6763 let n = parts.count();
6764 if *width.get_or_insert(n) != n {
6765 return Err(Error::internal("partitions of unequal count"));
6766 }
6767 match parts.data {
6768 Data::Box(v) => cells.extend(v.as_slice().iter().cloned()),
6769 _ => return Err(Error::internal("a partition is boxed")),
6770 }
6771 }
6772 let mut shape = y.shape[..y.rank() - 1].to_vec();
6773 shape.push(width.unwrap_or(0));
6774 return Ok(Array::new(shape, Data::Box(cells.into())));
6775 }
6776 if y.rank() == 0 {
6777 return Err(Error::new(
6778 ErrorKind::Rank,
6779 "partitioned enclose needs an array to partition",
6780 Some(span),
6781 ));
6782 }
6783 let flags = x
6784 .to_i64_vec()
6785 .ok_or_else(|| Error::domain("partition flags must be integers", span))?;
6786 if flags.iter().any(|&f| f < 0) {
6787 return Err(Error::domain("partition flags must not be negative", span));
6788 }
6789 if flags.len() != y.shape[0] {
6790 return Err(Error::new(
6791 ErrorKind::Length,
6792 format!("{} flag(s) for {} item(s)", flags.len(), y.shape[0]),
6793 Some(span),
6794 ));
6795 }
6796 let mut parts: Vec<Array> = Vec::new();
6797 let mut cur: Option<Data> = None;
6798 let mut prev = 0i64;
6799 for (i, &f) in flags.iter().enumerate() {
6800 if f > prev {
6801 if let Some(d) = cur.take() {
6802 parts.push(Array::new(vec![d.len()], d));
6803 }
6804 cur = Some(Data::empty(y.dtype()));
6805 }
6806 prev = f;
6807 if f == 0 {
6808 continue;
6809 }
6810 if let Some(d) = cur.as_mut() {
6811 push_elem(d, &y.data, i);
6812 }
6813 }
6814 if let Some(d) = cur.take() {
6815 parts.push(Array::new(vec![d.len()], d));
6816 }
6817 Ok(Array::new(vec![parts.len()], Data::Box(parts.into())))
6818}
6819
6820fn key(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
6823 let keys = if x.rank() == 0 { Array::new(vec![1], x.data.clone()) } else { x.clone() };
6824 let n = keys.items();
6825 if n != y.items() && !(y.rank() == 0 && n == 1) {
6826 return Err(Error::new(
6827 ErrorKind::Length,
6828 format!("{n} key(s) for {} item(s)", y.items()),
6829 Some(span),
6830 ));
6831 }
6832 let tol = ctx.cfg.tol;
6833 let mut order: Vec<usize> = Vec::new();
6834 let mut groups: Vec<Vec<usize>> = Vec::new();
6835 for i in 0..n {
6836 let k = keys.item(i);
6837 match order.iter().position(|&j| arrays_match(&k, &keys.item(j), tol)) {
6838 Some(g) => groups[g].push(i),
6839 None => {
6840 order.push(i);
6841 groups.push(vec![i]);
6842 }
6843 }
6844 }
6845 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
6846 let mut cells = Vec::with_capacity(groups.len());
6847 for g in &groups {
6848 cells.push(u.monad(&select_items(&items, g), ctx, span)?);
6849 }
6850 assemble(&[groups.len()], cells, span)
6851}
6852
6853fn oblique(u: &Verb, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
6856 if y.rank() < 2 {
6857 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
6858 let n = items.items();
6859 let mut cells = Vec::with_capacity(n);
6860 for i in 0..n {
6861 cells.push(u.monad(&select_items(&items, &[i]), ctx, span)?);
6862 }
6863 return assemble(&[n], cells, span);
6864 }
6865 if y.rank() > 2 {
6866 return Err(Error::not_yet("oblique (u/.) on a rank-3 or higher argument", span));
6867 }
6868 let (rows, cols) = (y.shape[0], y.shape[1]);
6869 let mut cells = Vec::with_capacity(rows + cols - 1);
6870 for d in 0..rows + cols - 1 {
6871 let mut data = Data::empty(y.dtype());
6872 let mut len = 0usize;
6873 for i in 0..rows {
6874 if d >= i && d - i < cols {
6875 push_elem(&mut data, &y.data, i * cols + (d - i));
6876 len += 1;
6877 }
6878 }
6879 cells.push(u.monad(&Array::new(vec![len], data), ctx, span)?);
6880 }
6881 assemble(&[rows + cols - 1], cells, span)
6882}
6883
6884fn cut_ranges(frets: &[bool], mode: i64) -> Vec<(usize, usize)> {
6892 let n = frets.len();
6893 let mut out = Vec::new();
6894 if mode.abs() == 1 {
6895 let mut start: Option<usize> = None;
6896 for (i, &fret) in frets.iter().enumerate() {
6897 if fret {
6898 if let Some(s) = start {
6899 out.push((s, i));
6900 }
6901 start = Some(i);
6902 }
6903 }
6904 if let Some(s) = start {
6905 out.push((s, n));
6906 }
6907 if mode < 0 {
6908 return out.into_iter().map(|(s, e)| (s + 1, e)).collect();
6909 }
6910 } else {
6911 let mut start = 0usize;
6912 for (i, &fret) in frets.iter().enumerate() {
6913 if fret {
6914 out.push((start, i + 1));
6915 start = i + 1;
6916 }
6917 }
6918 if mode < 0 {
6919 return out.into_iter().map(|(s, e)| (s, e - 1)).collect();
6920 }
6921 }
6922 out
6923}
6924
6925fn cut(
6927 u: &Verb,
6928 x: Option<&Array>,
6929 y: &Array,
6930 mode: i64,
6931 ctx: &mut Ctx<'_>,
6932 span: Span,
6933) -> Result<Array> {
6934 if mode == 0 {
6935 let Some(x) = x else {
6936 return u.monad(&reverse_all_axes(y), ctx, span);
6937 };
6938 let (origin, size) = rectangle(x, span)?;
6939 let origin = origin.unwrap_or_else(|| vec![0; size.len()]);
6940 return u.monad(&subarray(y, &origin, &size, span)?, ctx, span);
6941 }
6942 if mode.abs() == 3 {
6943 let Some(x) = x else {
6944 return Err(Error::not_yet("monadic tessellation (u;.3 y)", span));
6945 };
6946 return tessellate(u, x, y, mode < 0, ctx, span);
6947 }
6948 if !matches!(mode, 1 | -1 | 2 | -2) {
6949 return Err(Error::not_yet(format!("cut (u;.{mode})"), span));
6950 }
6951 let items = if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() };
6952 let n = items.items();
6953 let tol = ctx.cfg.tol;
6954 let frets: Vec<bool> = match x {
6955 Some(x) => {
6956 let flags = x
6957 .to_i64_vec()
6958 .ok_or_else(|| Error::domain("cut frets must be integers", span))?;
6959 if let Some(&bad) = flags.iter().find(|&&f| f != 0 && f != 1) {
6962 return Err(Error::domain(format!("{bad} is not a fret: a fret is 0 or 1"), span));
6963 }
6964 if x.rank() == 0 {
6967 vec![flags[0] != 0; n]
6968 } else {
6969 if flags.len() != n {
6970 return Err(Error::new(
6971 ErrorKind::Length,
6972 format!("{} fret(s) for {n} item(s)", flags.len()),
6973 Some(span),
6974 ));
6975 }
6976 flags.iter().map(|&f| f != 0).collect()
6977 }
6978 }
6979 None => {
6980 if n == 0 {
6982 Vec::new()
6983 } else {
6984 let at = if mode.abs() == 1 { 0 } else { n - 1 };
6985 let mark = items.item(at);
6986 (0..n).map(|i| arrays_match(&items.item(i), &mark, tol)).collect()
6987 }
6988 }
6989 };
6990 let ranges = cut_ranges(&frets, mode);
6991 let mut cells = Vec::with_capacity(ranges.len());
6992 for (s, e) in &ranges {
6993 cells.push(u.monad(§ion(&items, *s, *e), ctx, span)?);
6994 }
6995 assemble(&[ranges.len()], cells, span)
6996}
6997
6998fn rectangle(x: &Array, span: Span) -> Result<(Option<Vec<i64>>, Vec<i64>)> {
7001 let values = x
7002 .to_i64_vec()
7003 .ok_or_else(|| Error::domain("a cut rectangle is whole numbers", span))?;
7004 match x.rank() {
7005 0 | 1 => Ok((None, values)),
7006 2 if x.shape[0] == 2 => {
7007 let n = x.shape[1];
7008 Ok((Some(values[..n].to_vec()), values[n..].to_vec()))
7009 }
7010 _ => Err(Error::new(
7011 ErrorKind::Rank,
7012 "a cut rectangle is a vector of sizes, or two rows of origins and sizes",
7013 Some(span),
7014 )),
7015 }
7016}
7017
7018fn subarray(y: &Array, origin: &[i64], size: &[i64], span: Span) -> Result<Array> {
7022 if origin.len() > y.rank() {
7023 return Err(Error::new(
7024 ErrorKind::Rank,
7025 format!("a cut of {} axis/axes into a rank-{} value", origin.len(), y.rank()),
7026 Some(span),
7027 ));
7028 }
7029 let r = y.rank();
7030 let st = strides(&y.shape);
7031 let mut shape = y.shape.clone();
7032 let mut start = vec![0i64; r];
7033 let mut step = vec![1i64; r];
7034 for k in 0..origin.len() {
7035 let len = size[k].unsigned_abs() as usize;
7036 let from = if origin[k] < 0 { origin[k] + y.shape[k] as i64 } else { origin[k] };
7037 if from < 0 || from + len as i64 > y.shape[k] as i64 {
7038 return Err(Error::domain(
7039 format!("a cut of {len} from {from} leaves axis {k} of {}", y.shape[k]),
7040 span,
7041 ));
7042 }
7043 shape[k] = len;
7044 if size[k] < 0 {
7045 start[k] = from + len as i64 - 1;
7046 step[k] = -1;
7047 } else {
7048 start[k] = from;
7049 }
7050 }
7051 Ok(gather(y, &shape, &start, &step, &st))
7052}
7053
7054fn gather(y: &Array, shape: &[usize], start: &[i64], step: &[i64], st: &[usize]) -> Array {
7056 let n: usize = shape.iter().product();
7057 let mut data = Data::empty(y.dtype());
7058 let mut coord = vec![0usize; shape.len()];
7059 for _ in 0..n {
7060 let idx: usize = (0..shape.len())
7061 .map(|k| (start[k] + step[k] * coord[k] as i64) as usize * st[k])
7062 .sum();
7063 push_elem(&mut data, &y.data, idx);
7064 odometer(&mut coord, shape);
7065 }
7066 Array::new(shape.to_vec(), data)
7067}
7068
7069fn tessellate(
7073 u: &Verb,
7074 x: &Array,
7075 y: &Array,
7076 complete: bool,
7077 ctx: &mut Ctx<'_>,
7078 span: Span,
7079) -> Result<Array> {
7080 let (movement, size) = rectangle(x, span)?;
7082 let movement = movement.unwrap_or_else(|| vec![1; size.len()]);
7083 if size.iter().any(|&s| s < 0) {
7084 return Err(Error::not_yet("a tessellation with a negative size", span));
7085 }
7086 if size.len() > y.rank() {
7087 return Err(Error::new(
7088 ErrorKind::Rank,
7089 format!("a tessellation of {} axis/axes into a rank-{} value", size.len(), y.rank()),
7090 Some(span),
7091 ));
7092 }
7093 let mut frame = Vec::with_capacity(size.len());
7094 for k in 0..size.len() {
7095 let (len, step, block) = (y.shape[k] as i64, movement[k], size[k]);
7096 if step <= 0 {
7097 return Err(Error::domain("a tessellation moves by a positive step", span));
7098 }
7099 let count = if complete {
7100 if len < block { 0 } else { (len - block) / step + 1 }
7101 } else {
7102 (len + step - 1) / step
7103 };
7104 frame.push(count as usize);
7105 }
7106 let total: usize = frame.iter().product();
7107 let mut cells = Vec::with_capacity(total);
7108 let mut coord = vec![0usize; frame.len()];
7109 for _ in 0..total {
7110 let origin: Vec<i64> = (0..frame.len()).map(|k| coord[k] as i64 * movement[k]).collect();
7111 let block: Vec<i64> = (0..frame.len())
7113 .map(|k| size[k].min(y.shape[k] as i64 - origin[k]))
7114 .collect();
7115 cells.push(u.monad(&subarray(y, &origin, &block, span)?, ctx, span)?);
7116 odometer(&mut coord, &frame);
7117 }
7118 assemble(&frame, cells, span)
7119}
7120
7121fn reverse_all_axes(y: &Array) -> Array {
7123 if y.rank() == 0 {
7124 return y.clone();
7125 }
7126 let st = strides(&y.shape);
7127 let n = y.count();
7128 let r = y.rank();
7129 let mut data = Data::empty(y.dtype());
7130 let mut coord = vec![0usize; r];
7131 for _ in 0..n {
7132 let idx: usize = (0..r).map(|k| (y.shape[k] - 1 - coord[k]) * st[k]).sum();
7133 push_elem(&mut data, &y.data, idx);
7134 odometer(&mut coord, &y.shape);
7135 }
7136 Array::new(y.shape.clone(), data)
7137}
7138
7139fn axis_to_front(y: &Array, k: usize) -> Array {
7143 if k == 0 || y.rank() < 2 {
7144 return y.clone();
7145 }
7146 let r = y.rank();
7147 let src: Vec<usize> = std::iter::once(k).chain((0..r).filter(|&a| a != k)).collect();
7148 permute_axes(y, &src)
7149}
7150
7151fn front_to_axis(y: &Array, k: usize) -> Array {
7153 if k == 0 || y.rank() < 2 {
7154 return y.clone();
7155 }
7156 let r = y.rank();
7157 let mut src = Vec::with_capacity(r);
7160 for a in 0..r {
7161 src.push(match a.cmp(&k) {
7162 std::cmp::Ordering::Less => a + 1,
7163 std::cmp::Ordering::Equal => 0,
7164 std::cmp::Ordering::Greater => a,
7165 });
7166 }
7167 permute_axes(y, &src)
7168}
7169
7170fn permute_axes(y: &Array, src: &[usize]) -> Array {
7172 let st = strides(&y.shape);
7173 let out_shape: Vec<usize> = src.iter().map(|&a| y.shape[a]).collect();
7174 let n = y.count();
7175 let mut data = Data::empty(y.dtype());
7176 let mut coord = vec![0usize; src.len()];
7177 for _ in 0..n {
7178 let idx: usize = (0..src.len()).map(|a| coord[a] * st[src[a]]).sum();
7179 push_elem(&mut data, &y.data, idx);
7180 odometer(&mut coord, &out_shape);
7181 }
7182 Array::new(out_shape, data)
7183}
7184
7185struct Spec {
7189 width: usize,
7191 cells: Vec<Vec<usize>>,
7193 shape: Vec<usize>,
7196}
7197
7198fn axis_position(v: i64, len: usize, span: Span) -> Result<usize> {
7201 let p = if v < 0 { v + len as i64 } else { v };
7202 if p < 0 || p >= len as i64 {
7203 return Err(Error::domain(
7204 format!("index {v} is out of range: the axis has {len} element(s)"),
7205 span,
7206 ));
7207 }
7208 Ok(p as usize)
7209}
7210
7211fn index_spec(content: &Array, y: &Array, span: Span) -> Result<Spec> {
7221 let too_deep = |n: usize| {
7222 Error::new(
7223 ErrorKind::Rank,
7224 format!("an index specification of {n} axis/axes into a rank-{} value", y.rank()),
7225 Some(span),
7226 )
7227 };
7228 if let Some(items) = content.as_boxes() {
7229 if items.len() > y.rank() {
7230 return Err(too_deep(items.len()));
7231 }
7232 let mut per_axis: Vec<Vec<usize>> = Vec::with_capacity(items.len());
7233 let mut shape: Vec<usize> = Vec::new();
7234 for (k, c) in items.iter().enumerate() {
7235 let len = y.shape[k];
7236 if c.as_boxes().is_some() {
7237 let inner = open_cell(c);
7238 let excluded = inner.to_i64_vec().ok_or_else(|| {
7239 Error::domain("an index complement holds integers", span)
7240 })?;
7241 let mut dropped = vec![false; len];
7242 for v in excluded {
7243 dropped[axis_position(v, len, span)?] = true;
7244 }
7245 let kept: Vec<usize> = (0..len).filter(|i| !dropped[*i]).collect();
7246 shape.push(kept.len());
7247 per_axis.push(kept);
7248 } else {
7249 let idx = c
7250 .to_i64_vec()
7251 .ok_or_else(|| Error::domain("an index holds integers", span))?;
7252 let mut positions = Vec::with_capacity(idx.len());
7253 for v in idx {
7254 positions.push(axis_position(v, len, span)?);
7255 }
7256 shape.extend_from_slice(&c.shape);
7257 per_axis.push(positions);
7258 }
7259 }
7260 let mut cells: Vec<Vec<usize>> = vec![Vec::new()];
7262 for positions in &per_axis {
7263 let mut next = Vec::with_capacity(cells.len() * positions.len());
7264 for prefix in &cells {
7265 for &p in positions {
7266 let mut cell = prefix.clone();
7267 cell.push(p);
7268 next.push(cell);
7269 }
7270 }
7271 cells = next;
7272 }
7273 return Ok(Spec { width: per_axis.len(), cells, shape });
7274 }
7275 let idx = content
7276 .to_i64_vec()
7277 .ok_or_else(|| Error::domain("an index specification holds integers", span))?;
7278 let rank = content.rank();
7279 let width = if rank == 0 { 1 } else { content.shape[rank - 1] };
7280 if width > y.rank() {
7281 return Err(too_deep(width));
7282 }
7283 let shape: Vec<usize> = if rank == 0 { Vec::new() } else { content.shape[..rank - 1].to_vec() };
7284 let count: usize = shape.iter().product();
7285 let mut cells: Vec<Vec<usize>> = Vec::new();
7286 if width == 0 {
7287 cells.resize(count, Vec::new());
7288 } else {
7289 for chunk in idx.chunks(width) {
7290 let mut cell = Vec::with_capacity(width);
7291 for (k, &v) in chunk.iter().enumerate() {
7292 cell.push(axis_position(v, y.shape[k], span)?);
7293 }
7294 cells.push(cell);
7295 }
7296 }
7297 Ok(Spec { width, cells, shape })
7298}
7299
7300fn spec_offset(st: &[usize], cell: &[usize]) -> usize {
7302 cell.iter().enumerate().map(|(k, &p)| p * st[k]).sum()
7303}
7304
7305fn select_spec(spec: &Spec, y: &Array) -> Array {
7307 let st = strides(&y.shape);
7308 let size: usize = y.shape[spec.width..].iter().product();
7309 let mut data = Data::empty(y.dtype());
7310 for cell in &spec.cells {
7311 let base = spec_offset(&st, cell);
7312 for e in 0..size {
7313 push_elem(&mut data, &y.data, base + e);
7314 }
7315 }
7316 let mut shape = spec.shape.clone();
7317 shape.extend_from_slice(&y.shape[spec.width..]);
7318 Array::new(shape, data)
7319}
7320
7321fn amend_spec(spec: &Spec, x: &Array, y: &Array, span: Span) -> Result<Array> {
7324 let size: usize = y.shape[spec.width..].iter().product();
7325 let per_cell = if x.count() == size {
7326 false
7327 } else if x.count() == size * spec.cells.len() {
7328 true
7329 } else {
7330 return Err(Error::new(
7331 ErrorKind::Length,
7332 format!(
7333 "cannot amend {} cell(s) of {size} element(s) each with {} element(s)",
7334 spec.cells.len(),
7335 x.count()
7336 ),
7337 Some(span),
7338 ));
7339 };
7340 let mismatch = || {
7341 Error::new(
7342 ErrorKind::Type,
7343 "the replacement and the argument hold different kinds of value",
7344 Some(span),
7345 )
7346 };
7347 let t = DType::promote(x.dtype(), y.dtype()).ok_or_else(mismatch)?;
7348 let (Some(src), Some(base)) = (x.data.cast(t), y.data.cast(t)) else {
7349 return Err(mismatch());
7350 };
7351 let st = strides(&y.shape);
7352 let mut plan: Vec<Option<usize>> = vec![None; y.count()];
7353 for (n, cell) in spec.cells.iter().enumerate() {
7354 let at = spec_offset(&st, cell);
7355 for e in 0..size {
7356 plan[at + e] = Some(if per_cell { n * size + e } else { e });
7357 }
7358 }
7359 let mut data = Data::empty(t);
7360 for (i, slot) in plan.iter().enumerate() {
7361 match slot {
7362 Some(n) => push_elem(&mut data, &src, *n),
7363 None => push_elem(&mut data, &base, i),
7364 }
7365 }
7366 Ok(Array::new(y.shape.clone(), data))
7367}
7368
7369fn map_paths(y: &Array) -> Array {
7378 fn coord_of(shape: &[usize], mut i: usize) -> Array {
7379 let mut out = vec![0i64; shape.len()];
7380 for k in (0..shape.len()).rev() {
7381 out[k] = (i % shape[k]) as i64;
7382 i /= shape[k];
7383 }
7384 Array::from_i64(out)
7385 }
7386 fn go(y: &Array, prefix: &[Array]) -> Array {
7387 let Some(boxes) = y.as_boxes() else {
7388 if prefix.is_empty() {
7389 return Array::new(vec![0], Data::I64(Vec::new().into()));
7390 }
7391 return Array::new(vec![prefix.len()], Data::Box(prefix.to_vec().into()));
7392 };
7393 let cells: Vec<Array> = boxes
7394 .iter()
7395 .enumerate()
7396 .map(|(i, b)| {
7397 let mut path = prefix.to_vec();
7398 path.push(coord_of(&y.shape, i));
7399 go(b, &path)
7400 })
7401 .collect();
7402 Array::new(y.shape.clone(), Data::Box(cells.into()))
7403 }
7404 go(y, &[])
7405}
7406
7407fn shift_fill(x: &Array, y: &Array, fill: &Array, span: Span) -> Result<Array> {
7412 let counts = axis_counts(x, "shift", span)?;
7413 if y.rank() == 0 {
7414 return Ok(y.clone());
7415 }
7416 if counts.len() > y.rank() {
7417 return Err(Error::new(
7418 ErrorKind::Length,
7419 format!("shift has {} amounts for an argument of rank {}", counts.len(), y.rank()),
7420 Some(span),
7421 ));
7422 }
7423 if fill.count() != 1 {
7424 return Err(Error::new(ErrorKind::Length, "a fill is one atom", Some(span)));
7425 }
7426 let mismatch = || {
7427 Error::new(ErrorKind::Type, "the fill and the argument differ in kind", Some(span))
7428 };
7429 let t = DType::promote(y.dtype(), fill.dtype()).ok_or_else(mismatch)?;
7430 let (Some(base), Some(f)) = (y.data.cast(t), fill.data.cast(t)) else {
7431 return Err(mismatch());
7432 };
7433 let st = strides(&y.shape);
7434 let r = y.rank();
7435 let mut data = Data::empty(t);
7436 let mut coord = vec![0usize; r];
7437 for _ in 0..y.count() {
7438 let mut idx = 0usize;
7439 let mut vacated = false;
7440 for k in 0..r {
7441 let from = coord[k] as i64 + counts.get(k).copied().unwrap_or(0);
7442 if from < 0 || from >= y.shape[k] as i64 {
7443 vacated = true;
7444 break;
7445 }
7446 idx += from as usize * st[k];
7447 }
7448 if vacated {
7449 push_elem(&mut data, &f, 0);
7450 } else {
7451 push_elem(&mut data, &base, idx);
7452 }
7453 odometer(&mut coord, &y.shape);
7454 }
7455 Ok(Array::new(y.shape.clone(), data))
7456}
7457
7458fn memo_key(a: &Array, out: &mut Vec<u64>) -> bool {
7463 out.push(a.rank() as u64);
7464 out.extend(a.shape.iter().map(|&n| n as u64));
7465 out.push(a.dtype() as u64);
7466 match &a.data {
7467 Data::Ext(_) | Data::Rat(_) => false,
7468 Data::Box(items) => items.iter().all(|item| memo_key(item, out)),
7469 d => {
7470 for i in 0..d.len() {
7471 out.push(elem_key(d, i));
7472 }
7473 true
7474 }
7475 }
7476}
7477
7478fn memoised(
7480 u: &Verb,
7481 cache: &MemoCache,
7482 x: Option<&Array>,
7483 y: &Array,
7484 ctx: &mut Ctx<'_>,
7485 span: Span,
7486) -> Result<Array> {
7487 let apply = |ctx: &mut Ctx<'_>| match x {
7488 Some(x) => u.dyad(x, y, ctx, span),
7489 None => u.monad(y, ctx, span),
7490 };
7491 let mut key = vec![u64::from(x.is_some())];
7492 let keyed = x.is_none_or(|x| memo_key(x, &mut key)) && memo_key(y, &mut key);
7493 if !keyed {
7494 return apply(ctx);
7495 }
7496 if let Ok(map) = cache.lock() {
7497 if let Some(hit) = map.get(&key) {
7498 return Ok(hit.clone());
7499 }
7500 }
7501 let out = apply(ctx)?;
7502 if let Ok(mut map) = cache.lock() {
7503 map.insert(key, out.clone());
7504 }
7505 Ok(out)
7506}
7507
7508fn at_level(
7514 u: &Verb,
7515 level: i64,
7516 spread: bool,
7517 y: &Array,
7518 ctx: &mut Ctx<'_>,
7519 span: Span,
7520) -> Result<Array> {
7521 let n = if level < 0 { (boxing_level(y) + level).max(0) } else { level };
7523 if !spread {
7524 return map_level(u, n, y, ctx, span);
7525 }
7526 let mut cells = Vec::new();
7527 collect_level(u, n, y, ctx, span, &mut cells)?;
7528 let count = cells.len();
7529 assemble(&[count], cells, span)
7530}
7531
7532fn map_level(u: &Verb, n: i64, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
7533 let Some(boxes) = y.as_boxes().filter(|_| boxing_level(y) > n) else {
7534 return u.monad(y, ctx, span);
7535 };
7536 let boxes = boxes.to_vec();
7537 let mut cells = Vec::with_capacity(boxes.len());
7538 for b in &boxes {
7539 cells.push(map_level(u, n, b, ctx, span)?);
7540 }
7541 Ok(Array::new(y.shape.clone(), Data::Box(cells.into())))
7542}
7543
7544fn collect_level(
7545 u: &Verb,
7546 n: i64,
7547 y: &Array,
7548 ctx: &mut Ctx<'_>,
7549 span: Span,
7550 out: &mut Vec<Array>,
7551) -> Result<()> {
7552 let Some(boxes) = y.as_boxes().filter(|_| boxing_level(y) > n) else {
7553 out.push(u.monad(y, ctx, span)?);
7554 return Ok(());
7555 };
7556 let boxes = boxes.to_vec();
7557 for b in &boxes {
7558 collect_level(u, n, b, ctx, span, out)?;
7559 }
7560 Ok(())
7561}
7562
7563fn poly_coeffs(y: &Array, span: Span) -> Result<Vec<Cx>> {
7567 let c = y
7568 .data
7569 .cast(DType::Complex)
7570 .ok_or_else(|| Error::domain("a polynomial's coefficients are numbers", span))?;
7571 match c {
7572 Data::Complex(v) => Ok(v.as_slice().to_vec()),
7573 _ => Err(Error::internal("coefficients did not cast to complex")),
7574 }
7575}
7576
7577fn complex_or_real(values: Vec<Cx>) -> Array {
7579 if values.iter().all(|z| z[1] == 0.0) {
7580 return Array::from_f64(values.iter().map(|z| z[0]).collect());
7581 }
7582 Array::new(vec![values.len()], Data::Complex(values.into()))
7583}
7584
7585fn poly_eval(x: &Array, y: &Array, span: Span) -> Result<Array> {
7588 let at = poly_coeffs(y, span)?;
7589 let at = at.first().copied().unwrap_or(cx::ZERO);
7590 let value = match x.as_boxes() {
7591 Some(parts) => {
7592 if parts.len() != 2 {
7593 return Err(Error::domain(
7594 "the root form of a polynomial is `multiplier ; roots`",
7595 span,
7596 ));
7597 }
7598 let multiplier = poly_coeffs(&parts[0], span)?;
7599 let mut v = multiplier.first().copied().unwrap_or(cx::ONE);
7600 for r in poly_coeffs(&parts[1], span)? {
7601 v = cx::mul(v, cx::sub(at, r));
7602 }
7603 v
7604 }
7605 None => {
7606 let c = poly_coeffs(x, span)?;
7607 let mut v = cx::ZERO;
7608 for &k in c.iter().rev() {
7609 v = cx::add(cx::mul(v, at), k);
7610 }
7611 v
7612 }
7613 };
7614 Ok(scalar_complex_or_real(value))
7615}
7616
7617fn scalar_complex_or_real(z: Cx) -> Array {
7618 if z[1] == 0.0 {
7619 return Array::scalar_f64(z[0]);
7620 }
7621 Array::new(vec![], Data::Complex(vec![z].into()))
7622}
7623
7624fn poly_roots(y: &Array, span: Span) -> Result<Array> {
7628 if let Some(parts) = y.as_boxes() {
7629 if parts.len() != 2 {
7630 return Err(Error::domain(
7631 "the root form of a polynomial is `multiplier ; roots`",
7632 span,
7633 ));
7634 }
7635 let multiplier = poly_coeffs(&parts[0], span)?;
7636 let multiplier = multiplier.first().copied().unwrap_or(cx::ONE);
7637 let mut coeffs = vec![multiplier];
7639 for r in poly_coeffs(&parts[1], span)? {
7640 let mut next = vec![cx::ZERO; coeffs.len() + 1];
7641 for (k, &c) in coeffs.iter().enumerate() {
7642 next[k + 1] = cx::add(next[k + 1], c);
7643 next[k] = cx::sub(next[k], cx::mul(c, r));
7644 }
7645 coeffs = next;
7646 }
7647 return Ok(complex_or_real(coeffs));
7648 }
7649 let mut c = poly_coeffs(y, span)?;
7650 while c.len() > 1 && c[c.len() - 1] == cx::ZERO {
7651 c.pop();
7652 }
7653 if c.len() < 2 {
7654 return Err(Error::domain("a polynomial's roots need a coefficient of x", span));
7655 }
7656 let lead = c[c.len() - 1];
7657 let monic: Vec<Cx> = c.iter().map(|&k| cx::div(k, lead)).collect();
7658 let roots = durand_kerner(&monic);
7659 let pair = vec![scalar_complex_or_real(lead), complex_or_real(roots)];
7660 Ok(Array::new(vec![2], Data::Box(pair.into())))
7661}
7662
7663fn durand_kerner(monic: &[Cx]) -> Vec<Cx> {
7670 let d = monic.len() - 1;
7671 let seed = [0.4, 0.9];
7672 let mut z: Vec<Cx> = Vec::with_capacity(d);
7673 let mut p = cx::ONE;
7674 for _ in 0..d {
7675 z.push(p);
7676 p = cx::mul(p, seed);
7677 }
7678 let value = |monic: &[Cx], at: Cx| {
7679 let mut v = cx::ZERO;
7680 for &k in monic.iter().rev() {
7681 v = cx::add(cx::mul(v, at), k);
7682 }
7683 v
7684 };
7685 for _ in 0..500 {
7686 let mut moved: f64 = 0.0;
7687 for i in 0..d {
7688 let mut denom = cx::ONE;
7689 for j in 0..d {
7690 if i != j {
7691 denom = cx::mul(denom, cx::sub(z[i], z[j]));
7692 }
7693 }
7694 if denom == cx::ZERO {
7695 continue;
7696 }
7697 let step = cx::div(value(monic, z[i]), denom);
7698 z[i] = cx::sub(z[i], step);
7699 moved = moved.max(step[0].hypot(step[1]));
7700 }
7701 if moved < 1e-15 {
7702 break;
7703 }
7704 }
7705 for r in &mut z {
7707 if r[1].abs() < 1e-9 {
7708 r[1] = 0.0;
7709 }
7710 if r[0].abs() < 1e-12 {
7711 r[0] = 0.0;
7712 }
7713 }
7714 z.sort_by(|a, b| {
7718 let close = (a[0] - b[0]).abs() <= 1e-9 * (a[0].abs().max(b[0].abs()) + 1.0);
7719 let by_re = if close {
7720 std::cmp::Ordering::Equal
7721 } else {
7722 b[0].partial_cmp(&a[0]).unwrap_or(std::cmp::Ordering::Equal)
7723 };
7724 by_re.then(b[1].partial_cmp(&a[1]).unwrap_or(std::cmp::Ordering::Equal))
7725 });
7726 z
7727}
7728
7729fn poly_deriv(y: &Array, span: Span) -> Result<Array> {
7732 let c = poly_coeffs(y, span)?;
7733 if c.len() < 2 {
7734 return Ok(Array::from_i64(vec![0]));
7735 }
7736 let out: Vec<Cx> =
7737 c.iter().enumerate().skip(1).map(|(k, &v)| cx::mul(v, cx::from_real(k as f64))).collect();
7738 Ok(narrow_numbers(complex_or_real(out)))
7739}
7740
7741fn poly_integral(x: &Array, y: &Array, span: Span) -> Result<Array> {
7743 let c = poly_coeffs(y, span)?;
7744 let k = poly_coeffs(x, span)?;
7745 let mut out = vec![k.first().copied().unwrap_or(cx::ZERO)];
7746 for (i, &v) in c.iter().enumerate() {
7747 out.push(cx::div(v, cx::from_real((i + 1) as f64)));
7748 }
7749 Ok(narrow_numbers(complex_or_real(out)))
7750}
7751
7752fn narrow_numbers(a: Array) -> Array {
7756 let Data::F64(v) = &a.data else { return a };
7757 if v.iter().any(|x| !x.is_finite() || x.fract() != 0.0 || x.abs() > 9e15) {
7758 return a;
7759 }
7760 let values: Vec<i64> = v.iter().map(|&x| x as i64).collect();
7761 Array::new(a.shape, Data::I64(values.into()))
7762}
7763
7764fn characteristics(u: &Verb, y: &Array, span: Span) -> Result<Array> {
7768 let which = y.to_i64_vec().and_then(|v| v.first().copied());
7769 if which != Some(0) {
7770 return Err(Error::not_yet("a verb characteristic other than `u b. 0`", span));
7771 }
7772 let ranks = u.ranks();
7773 Ok(Array::from_f64(
7774 ranks
7775 .iter()
7776 .map(|&r| if r == RANK_INF { f64::INFINITY } else { r as f64 })
7777 .collect(),
7778 ))
7779}
7780
7781fn with_operands<R>(
7784 alpha: &Verb,
7785 omega: Option<&Verb>,
7786 ctx: &mut Ctx<'_>,
7787 f: impl FnOnce(&mut Ctx<'_>) -> Result<R>,
7788) -> Result<R> {
7789 let saved = (ctx.env.verb("⍺⍺").cloned(), ctx.env.verb("⍵⍵").cloned());
7790 ctx.env.define("⍺⍺".to_string(), alpha.clone());
7791 if let Some(g) = omega {
7792 ctx.env.define("⍵⍵".to_string(), g.clone());
7793 }
7794 let out = f(ctx);
7795 match saved.0 {
7796 Some(v) => ctx.env.define("⍺⍺".to_string(), v),
7797 None => ctx.env.undefine("⍺⍺"),
7798 }
7799 match saved.1 {
7800 Some(v) => ctx.env.define("⍵⍵".to_string(), v),
7801 None => ctx.env.undefine("⍵⍵"),
7802 }
7803 out
7804}
7805
7806fn is_mixed_simple(a: &Array) -> bool {
7810 let Some(items) = a.as_boxes() else { return false };
7811 if items.is_empty() || items.iter().any(|b| b.rank() != 0 || b.dtype() == DType::Box) {
7812 return false;
7813 }
7814 let mut common = Some(items[0].dtype());
7815 for b in &items[1..] {
7816 common = common.and_then(|t| DType::promote(t, b.dtype()));
7817 }
7818 common.is_none()
7819}
7820
7821fn nest(y: &Array) -> Array {
7824 if y.dtype() == DType::Box || y.rank() == 0 {
7825 return y.clone();
7826 }
7827 Array::boxed(y.clone())
7828}
7829
7830fn key_pairs(
7834 u: &Verb,
7835 keys: &Array,
7836 values: Option<&Array>,
7837 ctx: &mut Ctx<'_>,
7838 span: Span,
7839) -> Result<Array> {
7840 let base = if keys.rank() == 0 { Array::new(vec![1], keys.data.clone()) } else { keys.clone() };
7841 let n = base.items();
7842 if let Some(v) = values {
7843 if v.items() != n {
7844 return Err(Error::new(
7845 ErrorKind::Length,
7846 format!("{n} key(s) for {} item(s)", v.items()),
7847 Some(span),
7848 ));
7849 }
7850 }
7851 let groups = group_positions(&base, ctx.cfg.tol);
7852 let origin = ctx.cfg.rules.origin;
7853 let mut cells = Vec::with_capacity(groups.len());
7854 for (first, at) in &groups {
7855 let key = item_or_self(&base, *first);
7856 let group = match values {
7857 Some(v) => select_items(v, at),
7858 None => Array::from_i64(at.iter().map(|&i| origin + i as i64).collect()),
7859 };
7860 let monadic = matches!(u, Verb::Explicit(d) if d.left.is_none());
7863 cells.push(if monadic {
7864 u.monad(&group, ctx, span)?
7865 } else {
7866 u.dyad(&key, &group, ctx, span)?
7867 });
7868 }
7869 let count = cells.len();
7870 assemble(&[count], cells, span)
7871}
7872
7873fn group_positions(y: &Array, tol: Tol) -> Vec<(usize, Vec<usize>)> {
7876 let n = y.items();
7877 let mut keys: Vec<Array> = Vec::new();
7878 let mut groups: Vec<(usize, Vec<usize>)> = Vec::new();
7879 for i in 0..n {
7880 let item = y.item(i);
7881 match keys.iter().position(|k| arrays_match(k, &item, tol)) {
7882 Some(at) => groups[at].1.push(i),
7883 None => {
7884 keys.push(item);
7885 groups.push((i, vec![i]));
7886 }
7887 }
7888 }
7889 groups
7890}
7891
7892fn format_spec(x: &Array, y: &Array, fmt: &FmtOpts, span: Span) -> Result<Array> {
7898 let spec = x
7899 .to_i64_vec()
7900 .ok_or_else(|| Error::domain("a format specification is whole numbers", span))?;
7901 if y.dtype() == DType::Box {
7902 return Err(Error::not_yet("format by specification of a nested array", span));
7903 }
7904 let cols = if y.rank() == 0 { 1 } else { y.shape[y.rank() - 1] };
7905 let rows = y.count() / cols.max(1);
7906 let pairs: Vec<(Option<i64>, i64)> = match spec.len() {
7908 1 => vec![(None, spec[0]); cols],
7909 2 => vec![(Some(spec[0]), spec[1]); cols],
7910 n if n == 2 * cols => spec.chunks(2).map(|c| (Some(c[0]), c[1])).collect(),
7911 n => {
7912 return Err(Error::new(
7913 ErrorKind::Length,
7914 format!("{n} specification value(s) for {cols} column(s)"),
7915 Some(span),
7916 ));
7917 }
7918 };
7919 if pairs.iter().any(|&(w, p)| w.is_some_and(|w| w < 0) || p < 0) {
7920 return Err(Error::domain("a format width and precision are nonnegative", span));
7921 }
7922 let numbers = y.to_f64_vec();
7923 let text = |i: usize, p: i64| -> String {
7924 match (&y.data, &numbers) {
7925 (Data::Char(v), _) => v[i].to_string(),
7926 (_, Some(v)) => {
7927 let s = format!("{:.*}", p as usize, v[i]);
7928 if v[i] < 0.0 { format!("{}{}", fmt.neg, &s[1..]) } else { s }
7929 }
7930 _ => String::new(),
7931 }
7932 };
7933 if y.dtype() != DType::Char && numbers.is_none() {
7934 return Err(Error::domain("format by specification takes numbers or characters", span));
7935 }
7936 let widths: Vec<usize> = pairs
7938 .iter()
7939 .enumerate()
7940 .map(|(c, &(w, p))| match w {
7941 Some(w) => w as usize,
7942 None => {
7943 (0..rows).map(|r| text(r * cols + c, p).chars().count()).max().unwrap_or(0) + 1
7944 }
7945 })
7946 .collect();
7947 let line: usize = widths.iter().sum();
7948 let mut out: Vec<char> = Vec::with_capacity(rows * line);
7949 for r in 0..rows {
7950 for c in 0..cols {
7951 let s = text(r * cols + c, pairs[c].1);
7952 let len = s.chars().count();
7953 if len > widths[c] {
7954 return Err(Error::domain(
7955 format!("{s} does not fit a field {} wide", widths[c]),
7956 span,
7957 ));
7958 }
7959 out.extend(std::iter::repeat_n(' ', widths[c] - len));
7960 out.extend(s.chars());
7961 }
7962 }
7963 let mut shape = if y.rank() == 0 { Vec::new() } else { y.shape[..y.rank() - 1].to_vec() };
7964 shape.push(line);
7965 Ok(Array::new(shape, Data::Char(out.into())))
7966}
7967
7968fn iota_apl(y: &Array, origin: i64, span: Span) -> Result<Array> {
7972 if y.rank() > 1 {
7973 return Err(Error::new(
7974 ErrorKind::Rank,
7975 "the index generator takes a shape, which is a scalar or a vector",
7976 Some(span),
7977 ));
7978 }
7979 let dims = y
7980 .to_i64_vec()
7981 .ok_or_else(|| Error::domain("index generator needs an integer argument", span))?;
7982 if dims.iter().any(|&n| n < 0) {
7983 return Err(Error::domain("index generator needs nonnegative lengths", span));
7984 }
7985 if dims.len() <= 1 {
7986 let n = dims.first().copied().unwrap_or(0);
7987 crate::limits::count(n as u128, span)?;
7988 return Ok(Array::from_i64((0..n).map(|i| origin + i).collect()));
7989 }
7990 let shape: Vec<usize> = dims.iter().map(|&n| n as usize).collect();
7991 let total = crate::limits::elements(&shape, span)?;
7992 let mut cells = Vec::with_capacity(total);
7993 let mut coord = vec![0usize; shape.len()];
7994 for _ in 0..total {
7995 cells.push(Array::from_i64(coord.iter().map(|&c| origin + c as i64).collect()));
7996 odometer(&mut coord, &shape);
7997 }
7998 Ok(Array::new(shape, Data::Box(cells.into())))
7999}
8000
8001fn carry_exact(result: Array, y: &Array) -> Array {
8007 if !matches!(y.dtype(), DType::Ext | DType::Rat) {
8008 return result;
8009 }
8010 match result.data.cast(DType::Ext) {
8011 Some(data) => Array::new(result.shape, data),
8012 None => result,
8013 }
8014}
8015
8016fn carry_exact2(result: Array, x: &Array, y: &Array) -> Array {
8017 let widened = carry_exact(result, x);
8018 carry_exact(widened, y)
8019}
8020
8021fn truth_table(m: u8, x: &Array, y: &Array, span: Span) -> Result<Array> {
8024 let table = m & 15;
8025 let bit = |a: i64, b: i64| ((table >> (3 - (2 * a + b))) & 1) as i64;
8026 let xs = x
8027 .to_i64_vec()
8028 .ok_or_else(|| Error::domain("a boolean function takes integers", span))?;
8029 let ys = y
8030 .to_i64_vec()
8031 .ok_or_else(|| Error::domain("a boolean function takes integers", span))?;
8032 let (a, b) = (xs.first().copied().unwrap_or(0), ys.first().copied().unwrap_or(0));
8033 if m < 16 {
8034 if !(0..=1).contains(&a) || !(0..=1).contains(&b) {
8035 return Err(Error::domain(
8036 format!("{m} b. takes 0 and 1; {m} b. + 16 is the same function on every bit"),
8037 span,
8038 ));
8039 }
8040 return Ok(Array::scalar_bool(bit(a, b) != 0));
8041 }
8042 let mut out = 0i64;
8043 for k in 0..64 {
8044 if bit((a >> k) & 1, (b >> k) & 1) != 0 {
8045 out |= 1i64 << k;
8046 }
8047 }
8048 Ok(Array::scalar_i64(out))
8049}
8050
8051pub fn amend_at(
8056 base: &Array,
8057 slots: &[Option<Array>],
8058 value: &Array,
8059 origin: i64,
8060 span: Span,
8061) -> Result<Array> {
8062 if slots.len() != base.rank() {
8063 return Err(Error::new(
8064 ErrorKind::Rank,
8065 format!(
8066 "indexed assignment needs one index per axis: {} slot(s) for a rank-{} value",
8067 slots.len(),
8068 base.rank()
8069 ),
8070 Some(span),
8071 ));
8072 }
8073 let mut axes: Vec<Vec<usize>> = Vec::with_capacity(slots.len());
8075 let mut selected: Vec<usize> = Vec::new();
8076 for (k, slot) in slots.iter().enumerate() {
8077 let len = base.shape[k];
8078 let Some(idx) = slot else {
8079 axes.push((0..len).collect());
8080 selected.push(len);
8081 continue;
8082 };
8083 let Some(values) = idx.to_i64_vec() else {
8084 return Err(Error::new(
8085 ErrorKind::Type,
8086 "an index must be numeric",
8087 Some(span),
8088 ));
8089 };
8090 let mut positions = Vec::with_capacity(values.len());
8091 for v in values {
8092 let p = v - origin;
8093 if p < 0 || p as usize >= len {
8094 return Err(Error::new(
8095 ErrorKind::Domain,
8096 format!("index {v} is outside axis {k}, which has {len} element(s)"),
8097 Some(span),
8098 ));
8099 }
8100 positions.push(p as usize);
8101 }
8102 if idx.rank() > 0 {
8104 selected.push(positions.len());
8105 }
8106 axes.push(positions);
8107 }
8108 let count: usize = axes.iter().map(Vec::len).product();
8109 if value.rank() != 0 && (value.shape != selected || value.count() != count) {
8110 return Err(Error::new(
8111 ErrorKind::Shape,
8112 format!(
8113 "indexed assignment needs a scalar or a {} value, not a {} one",
8114 show_shape(&selected),
8115 show_shape(&value.shape)
8116 ),
8117 Some(span),
8118 ));
8119 }
8120 let dtype = DType::promote(base.dtype(), value.dtype()).ok_or_else(|| {
8123 Error::new(
8124 ErrorKind::Type,
8125 format!(
8126 "cannot put a {} value into a {} array",
8127 value.dtype().name(),
8128 base.dtype().name()
8129 ),
8130 Some(span),
8131 )
8132 })?;
8133 let mut out = base.cast(dtype).ok_or_else(|| Error::internal("promotion failed"))?;
8134 let src = value.cast(dtype).ok_or_else(|| Error::internal("promotion failed"))?;
8135 let strides = row_major_strides(&base.shape);
8136 let mut coords = vec![0usize; axes.len()];
8137 for n in 0..count {
8138 let mut rest = n;
8139 for k in (0..axes.len()).rev() {
8140 let len = axes[k].len();
8141 coords[k] = axes[k][rest % len];
8142 rest /= len;
8143 }
8144 let at: usize = coords.iter().zip(&strides).map(|(c, s)| c * s).sum();
8145 let from = if src.rank() == 0 { 0 } else { n };
8146 put_element(&mut out.data, at, &src.data, from);
8147 }
8148 Ok(out)
8149}
8150
8151fn row_major_strides(shape: &[usize]) -> Vec<usize> {
8152 let mut strides = vec![1usize; shape.len()];
8153 for k in (0..shape.len().saturating_sub(1)).rev() {
8154 strides[k] = strides[k + 1] * shape[k + 1];
8155 }
8156 strides
8157}
8158
8159fn put_element(dst: &mut Data, at: usize, src: &Data, from: usize) {
8161 match (dst, src) {
8162 (Data::Bool(d), Data::Bool(s)) => d.to_mut()[at] = s.as_slice()[from],
8163 (Data::I64(d), Data::I64(s)) => d.to_mut()[at] = s.as_slice()[from],
8164 (Data::Ext(d), Data::Ext(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
8165 (Data::Rat(d), Data::Rat(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
8166 (Data::F64(d), Data::F64(s)) => d.to_mut()[at] = s.as_slice()[from],
8167 (Data::Char(d), Data::Char(s)) => d.to_mut()[at] = s.as_slice()[from],
8168 (Data::Box(d), Data::Box(s)) => d.to_mut()[at] = s.as_slice()[from].clone(),
8169 _ => debug_assert!(false, "amend across types"),
8171 }
8172}
8173
8174fn agenda_pick(
8177 vs: &[Verb],
8178 w: &Verb,
8179 x: Option<&Array>,
8180 y: &Array,
8181 ctx: &mut Ctx<'_>,
8182 span: Span,
8183) -> Result<Verb> {
8184 let chosen = match x {
8185 None => w.monad(y, ctx, span)?,
8186 Some(x) => w.dyad(x, y, ctx, span)?,
8187 };
8188 let at = chosen
8189 .to_i64_vec()
8190 .and_then(|v| v.first().copied())
8191 .ok_or_else(|| Error::domain("an agenda index must be an integer", span))?;
8192 pick_gerund(vs, at, span)
8193}
8194
8195pub(crate) fn pick_gerund(vs: &[Verb], at: i64, span: Span) -> Result<Verb> {
8198 usize::try_from(at)
8199 .ok()
8200 .and_then(|k| vs.get(k))
8201 .cloned()
8202 .ok_or_else(|| {
8203 Error::domain(
8204 format!("agenda {at} is out of range: the gerund has {} verbs", vs.len()),
8205 span,
8206 )
8207 })
8208}
8209
8210fn outfix(u: &Verb, x: &Array, y: &Array, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
8214 let k = one_int(x, "an outfix width", span)?;
8215 let n = y.items() as i64;
8216 let list = as_list(y);
8217 let starts: Vec<i64> = if k < 0 {
8222 let step = -k;
8223 (0..(n + step - 1) / step).map(|i| i * step).collect()
8224 } else {
8225 (0..=(n - k)).collect()
8226 };
8227 let width = k.unsigned_abs() as usize;
8228 let mut cells = Vec::with_capacity(starts.len());
8229 for start in starts {
8230 let start = start as usize;
8231 let keep: Vec<usize> =
8232 (0..n as usize).filter(|&i| i < start || i >= start + width).collect();
8233 cells.push(u.monad(&select_items(&list, &keep), ctx, span)?);
8234 }
8235 assemble(&[cells.len()], cells, span)
8236}
8237
8238pub(crate) fn obverse(v: &Verb) -> Option<Verb> {
8250 let swap = |name: &'static str| -> Option<Verb> {
8251 crate::frontend::j::verb_named(name)
8252 };
8253 Some(match v {
8254 Verb::Prim(p) => {
8255 use ScalarMonad as SM;
8256 let by_monad: Option<&'static str> = match p.monad {
8257 MonadOp::Scalar(SM::Conj | SM::Neg | SM::Recip | SM::OneMinus)
8259 | MonadOp::Reverse
8260 | MonadOp::TransposeAxes => Some(p.name),
8261 MonadOp::Scalar(SM::Exp) => Some("^."),
8262 MonadOp::Scalar(SM::Ln) => Some("^"),
8263 MonadOp::Scalar(SM::Sqrt) => Some("*:"),
8264 MonadOp::Scalar(SM::Square) => Some("%:"),
8265 MonadOp::Scalar(SM::Double) => Some("-:"),
8266 MonadOp::Scalar(SM::Halve) => Some("+:"),
8267 MonadOp::Scalar(SM::Inc) => Some("<:"),
8268 MonadOp::Scalar(SM::Dec) => Some(">:"),
8269 MonadOp::Enclose(_) => Some(">"),
8270 MonadOp::Open => Some("<"),
8271 MonadOp::DecodeBits => Some("#:"),
8272 MonadOp::EncodeBits => Some("#."),
8273 _ => None,
8274 };
8275 swap(by_monad?)?
8276 }
8277 Verb::WithObverse(_, w) => (**w).clone(),
8279 Verb::Atop(f, g) => {
8281 Verb::Atop(Box::new(obverse(g)?), Box::new(obverse(f)?))
8282 }
8283 Verb::Compose(f, g) | Verb::Beside(f, g) => {
8284 Verb::Atop(Box::new(obverse(g)?), Box::new(obverse(f)?))
8285 }
8286 Verb::Rank(f, r) => Verb::Rank(Box::new(obverse(f)?), *r),
8287 Verb::Fit(f, n) => Verb::Fit(Box::new(obverse(f)?), *n),
8288 Verb::PowerN(f, Power::Times(n)) => {
8290 Verb::PowerN(Box::new(obverse(f)?), Power::Times(*n))
8291 }
8292 Verb::BondLeft(m, f) => bond_obverse(m, f, true)?,
8293 Verb::BondRight(f, n) => bond_obverse(n, f, false)?,
8294 _ => return None,
8295 })
8296}
8297
8298fn bond_obverse(n: &Array, f: &Verb, left: bool) -> Option<Verb> {
8302 let Verb::Prim(p) = f else { return None };
8303 let named = |name: &'static str| crate::frontend::j::verb_named(name);
8304 let bond = |name: &'static str, arg: &Array| -> Option<Verb> {
8305 let g = named(name)?;
8306 Some(if left {
8307 Verb::BondLeft(arg.clone(), Box::new(g))
8308 } else {
8309 Verb::BondRight(Box::new(g), arg.clone())
8310 })
8311 };
8312 use ScalarDyad as SD;
8313 let DyadOp::Scalar(op) = p.dyad else { return None };
8314 match (op, left) {
8315 (SD::Sub | SD::DivJ | SD::DivApl, true) => bond(p.name, n),
8317 (SD::Add, _) => bond("-", n),
8318 (SD::Sub, false) => bond("+", n),
8319 (SD::Mul, _) => bond("%", n),
8320 (SD::DivJ | SD::DivApl, false) => bond("*", n),
8321 (SD::Pow, false) => Some(Verb::BondLeft(n.clone(), Box::new(named("%:")?))),
8323 (SD::Pow, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^.")?))),
8324 (SD::Log, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^")?))),
8325 (SD::Root, true) => Some(Verb::BondLeft(n.clone(), Box::new(named("^")?))),
8326 _ => None,
8327 }
8328}
8329
8330fn self_classify(y: &Array, tol: Tol) -> Array {
8335 let items = if y.rank() == 0 { 1 } else { y.items() };
8336 let keys = nub(&as_list(y), tol);
8337 let rows = keys.items();
8338 let mut out = Vec::with_capacity(rows * items);
8339 for i in 0..rows {
8340 let key = item_or_self(&keys, i);
8341 for j in 0..items {
8342 out.push(arrays_match(&key, &item_or_self(y, j), tol) as u8);
8343 }
8344 }
8345 Array::new(vec![rows, items], Data::Bool(out.into()))
8346}
8347
8348fn nub_sieve(y: &Array, tol: Tol) -> Array {
8350 let items = if y.rank() == 0 { 1 } else { y.items() };
8351 let mut seen: Vec<Array> = Vec::new();
8352 let mut out = Vec::with_capacity(items);
8353 for i in 0..items {
8354 let cell = item_or_self(y, i);
8355 let fresh = !seen.iter().any(|s| arrays_match(s, &cell, tol));
8356 if fresh {
8357 seen.push(cell);
8358 }
8359 out.push(fresh as u8);
8360 }
8361 Array::new(vec![items], Data::Bool(out.into()))
8362}
8363
8364fn as_list(y: &Array) -> Array {
8366 if y.rank() == 0 { Array::new(vec![1], y.data.clone()) } else { y.clone() }
8367}
8368
8369fn conforming_cells(y: &Array, item_rank: usize) -> Vec<Array> {
8374 let frame_rank = y.rank().saturating_sub(item_rank);
8375 let nf: usize = y.shape[..frame_rank].iter().product();
8376 (0..nf).map(|i| y.cell_at(frame_rank, i)).collect()
8377}
8378
8379fn item_marks(y: &Array, x: &Array, tol: Tol) -> Vec<bool> {
8381 let n = if y.rank() == 0 { 1 } else { y.items() };
8382 let item_rank = y.rank().saturating_sub(1);
8383 let against = conforming_cells(x, item_rank);
8384 (0..n)
8385 .map(|i| {
8386 let cell = item_or_self(y, i);
8387 against.iter().any(|c| arrays_match(&cell, c, tol))
8388 })
8389 .collect()
8390}
8391
8392fn set_less(x: &Array, y: &Array, tol: Tol) -> Array {
8394 let xs = as_list(x);
8395 let marks = item_marks(&xs, y, tol);
8396 let keep: Vec<usize> = (0..marks.len()).filter(|&i| !marks[i]).collect();
8397 select_items(&xs, &keep)
8398}
8399
8400fn intersect_items(x: &Array, y: &Array, tol: Tol) -> Array {
8402 let xs = as_list(x);
8403 let marks = item_marks(&xs, y, tol);
8404 let keep: Vec<usize> = (0..marks.len()).filter(|&i| marks[i]).collect();
8405 select_items(&xs, &keep)
8406}
8407
8408fn union_items(x: &Array, y: &Array, tol: Tol, span: Span) -> Result<Array> {
8411 let xs = as_list(x);
8412 let ys = as_list(y);
8413 let marks = item_marks(&ys, &xs, tol);
8414 let mut extra: Vec<usize> = Vec::new();
8415 for (i, &seen) in marks.iter().enumerate() {
8416 if seen {
8417 continue;
8418 }
8419 let cell = item_or_self(&ys, i);
8420 if !extra.iter().any(|&j| arrays_match(&item_or_self(&ys, j), &cell, tol)) {
8421 extra.push(i);
8422 }
8423 }
8424 catenate(&xs, &select_items(&ys, &extra), true, false, span)
8425}
8426
8427fn find_seq(x: &Array, y: &Array, tol: Tol) -> Array {
8431 let xs = as_list(x);
8432 let ys = as_list(y);
8433 let (k, n) = (xs.items(), ys.items());
8434 let mut out = vec![0u8; n];
8435 if k > 0 && k <= n {
8436 for (start, slot) in out.iter_mut().enumerate().take(n - k + 1) {
8437 let hit = (0..k).all(|d| {
8438 arrays_match(&item_or_self(&xs, d), &item_or_self(&ys, start + d), tol)
8439 });
8440 *slot = hit as u8;
8441 }
8442 }
8443 Array::new(vec![n], Data::Bool(out.into()))
8444}
8445
8446fn bool_dyad(op: BoolDyad, x: &Array, y: &Array, cfg: EvalCfg, span: Span) -> Result<Array> {
8450 let bit = |a: &Array| -> Result<u8> {
8451 match a.to_i64_vec().as_deref() {
8452 Some([0]) => Ok(0),
8453 Some([1]) => Ok(1),
8454 _ => Err(Error::domain("this verb reads values of 0 or 1", span)),
8455 }
8456 };
8457 let _ = cfg;
8458 let (a, b) = (bit(x)?, bit(y)?);
8459 let v = match op {
8460 BoolDyad::Nor => u8::from(a == 0 && b == 0),
8461 BoolDyad::Nand => u8::from(a == 0 || b == 0),
8462 };
8463 Ok(Array::new(vec![], Data::Bool(vec![v].into())))
8464}
8465
8466fn item_ranks(y: &Array, order: ComplexOrder, span: Span) -> Result<Vec<usize>> {
8472 check_gradable(y, order, span)?;
8473 if !y.dtype().is_numeric() {
8474 return Err(Error::domain("an anagram index needs numbers", span));
8475 }
8476 let order = grade_order(&as_list(y), false);
8477 let mut ranks = vec![0usize; order.len()];
8478 for (place, &i) in order.iter().enumerate() {
8479 ranks[i] = place;
8480 }
8481 Ok(ranks)
8482}
8483
8484fn anagram_index(y: &Array, order: ComplexOrder, span: Span) -> Result<Array> {
8487 let ranks = item_ranks(y, order, span)?;
8488 let n = ranks.len();
8489 let mut index: i128 = 0;
8490 for i in 0..n {
8491 let smaller = ranks[i + 1..].iter().filter(|&&r| r < ranks[i]).count() as i128;
8492 index = index
8493 .checked_mul((n - i) as i128)
8494 .and_then(|v| v.checked_add(smaller))
8495 .ok_or_else(|| Error::not_yet("an anagram index too large for an integer", span))?;
8496 }
8497 i64::try_from(index)
8498 .map(Array::scalar_i64)
8499 .map_err(|_| Error::not_yet("an anagram index too large for an integer", span))
8500}
8501
8502fn anagram_from(x: &Array, y: &Array, span: Span) -> Result<Array> {
8505 let idx = x
8506 .to_i64_vec()
8507 .ok_or_else(|| Error::domain("an anagram index must be an integer", span))?;
8508 let Some(&want) = idx.first() else {
8509 return Err(Error::internal("anagram with no index"));
8510 };
8511 let ys = as_list(y);
8512 let n = ys.items();
8513 let mut total: i128 = 1;
8514 for k in 1..=n as i128 {
8515 total = total
8516 .checked_mul(k)
8517 .ok_or_else(|| Error::not_yet("permuting more items than an integer counts", span))?;
8518 }
8519 let mut at = want as i128;
8520 if at < 0 {
8521 at += total;
8522 }
8523 if at < 0 || at >= total {
8524 return Err(Error::domain(
8525 format!("permutation {want} is out of range: {n} items have {total} of them"),
8526 span,
8527 ));
8528 }
8529 let mut pool: Vec<usize> = (0..n).collect();
8532 let mut order = Vec::with_capacity(n);
8533 let mut fact = total;
8534 for i in 0..n {
8535 fact /= (n - i) as i128;
8536 let d = (at / fact) as usize;
8537 at %= fact;
8538 order.push(pool.remove(d));
8539 }
8540 Ok(select_items(&ys, &order))
8541}
8542
8543fn cycle_form(y: &Array, span: Span) -> Result<Array> {
8547 if y.dtype() == DType::Box {
8548 let perm = cycles_to_direct(y, span)?;
8549 return Ok(Array::from_i64(perm.iter().map(|&i| i as i64).collect()));
8550 }
8551 let perm = direct_permutation(y, span)?;
8552 let mut boxes: Vec<Array> = Vec::new();
8553 let mut done = vec![false; perm.len()];
8554 for start in 0..perm.len() {
8555 if done[start] {
8556 continue;
8557 }
8558 let mut cycle = Vec::new();
8559 let mut at = start;
8560 while !done[at] {
8561 done[at] = true;
8562 cycle.push(at);
8563 at = perm[at];
8564 }
8565 let top = cycle.iter().position(|&v| v == *cycle.iter().max().unwrap()).unwrap();
8568 cycle.rotate_left(top);
8569 boxes.push(Array::boxed(Array::from_i64(
8570 cycle.iter().map(|&i| i as i64).collect(),
8571 )));
8572 }
8573 boxes.sort_by_key(|b| b.as_boxes().map(|s| s[0].to_i64_vec().unwrap()[0]).unwrap_or(0));
8574 let n = boxes.len();
8575 let inner: Vec<Array> =
8576 boxes.into_iter().map(|b| b.as_boxes().unwrap()[0].clone()).collect();
8577 Ok(Array::new(vec![n], Data::Box(inner.into())))
8578}
8579
8580fn direct_permutation(y: &Array, span: Span) -> Result<Vec<usize>> {
8582 let v = y
8583 .to_i64_vec()
8584 .ok_or_else(|| Error::domain("a permutation is a list of integers", span))?;
8585 let n = v.len();
8586 let mut seen = vec![false; n];
8587 let mut out = Vec::with_capacity(n);
8588 for &i in &v {
8589 let k = usize::try_from(i).ok().filter(|&k| k < n && !seen[k]).ok_or_else(|| {
8590 Error::domain(format!("{i} does not belong to a permutation of {n} items"), span)
8591 })?;
8592 seen[k] = true;
8593 out.push(k);
8594 }
8595 Ok(out)
8596}
8597
8598fn cycles_to_direct(y: &Array, span: Span) -> Result<Vec<usize>> {
8602 let boxes = y.as_boxes().ok_or_else(|| Error::internal("cycles from a simple array"))?;
8603 let mut cycles: Vec<Vec<usize>> = Vec::new();
8604 let mut top = 0usize;
8605 for b in boxes {
8606 let v = b
8607 .to_i64_vec()
8608 .ok_or_else(|| Error::domain("a cycle is a list of integers", span))?;
8609 let mut cycle = Vec::with_capacity(v.len());
8610 for &i in &v {
8611 let k = usize::try_from(i)
8612 .map_err(|_| Error::domain(format!("{i} is not an index"), span))?;
8613 top = top.max(k + 1);
8614 cycle.push(k);
8615 }
8616 cycles.push(cycle);
8617 }
8618 let mut perm: Vec<usize> = (0..top).collect();
8619 for cycle in &cycles {
8620 for w in 0..cycle.len() {
8621 perm[cycle[w]] = cycle[(w + 1) % cycle.len()];
8623 }
8624 }
8625 Ok(perm)
8626}
8627
8628fn permute(x: &Array, y: &Array, span: Span) -> Result<Array> {
8633 let ys = as_list(y);
8634 let n = ys.items();
8635 let cyclic = x.dtype() == DType::Box;
8636 if !cyclic && x.rank() == 0 {
8637 return Err(Error::not_yet("permuting by a single atom (x C. y)", span));
8638 }
8639 let mut perm =
8640 if cyclic { cycles_to_direct(x, span)? } else { direct_permutation(&as_list(x), span)? };
8641 if perm.len() > n {
8642 return Err(Error::new(
8643 ErrorKind::Length,
8644 format!("a permutation of {} items applied to {n}", perm.len()),
8645 Some(span),
8646 ));
8647 }
8648 if perm.len() < n {
8649 if cyclic {
8650 perm.extend(perm.len()..n);
8652 } else {
8653 let head: Vec<usize> = (perm.len()..n).collect();
8656 perm.splice(0..0, head);
8657 }
8658 }
8659 Ok(select_items(&ys, &perm))
8660}
8661
8662fn unicode(y: &Array, pass_chars: bool, span: Span) -> Result<Array> {
8668 if y.dtype() == DType::Char {
8669 if pass_chars {
8670 return Ok(y.clone());
8671 }
8672 return Ok(chars_to_codes(y));
8673 }
8674 codes_to_chars(y, span)
8675}
8676
8677fn chars_to_codes(y: &Array) -> Array {
8678 let Data::Char(v) = &y.data else { return y.clone() };
8679 Array::new(y.shape.clone(), Data::I64(v.iter().map(|&c| c as i64).collect()))
8680}
8681
8682fn codes_to_chars(y: &Array, span: Span) -> Result<Array> {
8683 let v = y
8684 .to_i64_vec()
8685 .ok_or_else(|| Error::domain("a codepoint must be an integer", span))?;
8686 let mut out = Vec::with_capacity(v.len());
8687 for &c in &v {
8688 let ch = u32::try_from(c).ok().and_then(char::from_u32).ok_or_else(|| {
8689 Error::domain(format!("{c} is not a Unicode codepoint"), span)
8690 })?;
8691 out.push(ch);
8692 }
8693 Ok(Array::new(y.shape.clone(), Data::Char(out.into())))
8694}
8695
8696fn unicode_form(x: &Array, y: &Array, span: Span) -> Result<Array> {
8699 let form = x
8700 .to_i64_vec()
8701 .ok_or_else(|| Error::domain("a conversion form is an integer", span))?
8702 .first()
8703 .copied()
8704 .unwrap_or(0);
8705 match form {
8706 3 if y.dtype() == DType::Char => Ok(chars_to_codes(y)),
8707 3 => Err(Error::domain("form 3 converts characters to codepoints", span)),
8708 10 => codes_to_chars(y, span),
8709 n => Err(Error::not_yet(format!("the byte-oriented unicode form ({n} u:)"), span)),
8710 }
8711}
8712
8713fn boxing_level(y: &Array) -> i64 {
8715 match y.as_boxes() {
8716 None => 0,
8717 Some(bs) => 1 + bs.iter().map(boxing_level).max().unwrap_or(0),
8718 }
8719}
8720
8721fn split_items(y: &Array) -> Array {
8725 if y.rank() == 0 {
8726 return Array::boxed(y.clone());
8727 }
8728 let last = y.shape[y.rank() - 1];
8729 let outer: Vec<usize> = y.shape[..y.rank() - 1].to_vec();
8730 let n: usize = outer.iter().product();
8731 let mut boxes = Vec::with_capacity(n);
8732 for i in 0..n {
8733 let mut data = Data::empty(y.dtype());
8734 for k in 0..last {
8735 push_elem(&mut data, &y.data, i * last + k);
8736 }
8737 boxes.push(Array::new(vec![last], data));
8738 }
8739 Array::new(outer, Data::Box(boxes.into()))
8740}
8741
8742fn pick(x: &Array, y: &Array, origin: i64, span: Span) -> Result<Array> {
8745 let xs = as_list(x);
8746 let mut cur = y.clone();
8747 for i in 0..xs.items() {
8748 let step = open_cell(&item_or_self(&xs, i));
8749 let idx = step
8750 .to_i64_vec()
8751 .ok_or_else(|| Error::domain("a pick path holds integers", span))?;
8752 let base =
8753 if cur.rank() == 0 { Array::new(vec![1], cur.data.clone()) } else { cur.clone() };
8754 if idx.len() > base.rank() {
8755 return Err(Error::new(
8756 ErrorKind::Length,
8757 format!(
8758 "a path step of {} index(es) into a value of rank {}",
8759 idx.len(),
8760 cur.rank()
8761 ),
8762 Some(span),
8763 ));
8764 }
8765 let zeroed: Vec<i64> = idx.iter().map(|&v| v - origin).collect();
8766 let at = cell_index(&base, &zeroed, span)?;
8767 cur = open_cell(&base.cell_at(idx.len(), at));
8768 }
8769 Ok(cur)
8770}
8771
8772fn prime_meta(x: &Array, y: &Array, span: Span) -> Result<Array> {
8777 let form = one_int(x, "a prime query", span)?;
8778 let n = one_int(y, "a prime query", span)?;
8779 match form {
8780 -1 => Ok(Array::scalar_i64(primes_below(n, span)?)),
8782 0 => Ok(Array::scalar_bool(!is_prime(n))),
8784 1 => Ok(Array::scalar_bool(is_prime(n))),
8785 2 | 3 => {
8787 let (ps, es) = factor_table(n, span)?;
8788 let k = ps.len();
8789 if form == 3 {
8790 return Ok(Array::from_i64(ps));
8791 }
8792 let mut all = ps;
8793 all.extend(es);
8794 Ok(Array::new(vec![2, k], Data::I64(all.into())))
8795 }
8796 4 => Ok(Array::scalar_i64(next_prime(n, span)?)),
8798 -4 => Ok(Array::scalar_i64(previous_prime(n, span)?)),
8799 other => Err(Error::domain(format!("{other} is not a prime query"), span)),
8800 }
8801}
8802
8803fn prime_exponents(x: &Array, y: &Array, span: Span) -> Result<Array> {
8806 let n = one_int(y, "prime exponents", span)?;
8807 let count = x.to_f64_vec().and_then(|v| v.first().copied()).unwrap_or(0.0);
8808 let (ps, es) = factor_table(n, span)?;
8809 if count == f64::NEG_INFINITY {
8810 let k = ps.len();
8811 let mut all = ps;
8812 all.extend(es);
8813 return Ok(Array::new(vec![2, k], Data::I64(all.into())));
8814 }
8815 let want = one_int(x, "prime exponents", span)?;
8816 if want < 0 {
8817 return Err(Error::not_yet(format!("the prime exponent form ({want} q:)"), span));
8818 }
8819 let mut out = Vec::with_capacity(want as usize);
8820 for i in 0..want {
8821 let p = nth_prime(i, span)?;
8822 out.push(ps.iter().position(|&q| q == p).map_or(0, |at| es[at]));
8823 }
8824 Ok(Array::from_i64(out))
8825}
8826
8827fn factor_table(n: i64, span: Span) -> Result<(Vec<i64>, Vec<i64>)> {
8829 let factors = prime_factors(n, span)?;
8830 let mut ps: Vec<i64> = Vec::new();
8831 let mut es: Vec<i64> = Vec::new();
8832 for f in factors {
8833 if ps.last() == Some(&f) {
8834 *es.last_mut().unwrap() += 1;
8835 } else {
8836 ps.push(f);
8837 es.push(1);
8838 }
8839 }
8840 Ok((ps, es))
8841}
8842
8843fn is_prime(n: i64) -> bool {
8844 if n < 2 {
8845 return false;
8846 }
8847 let mut d = 2i64;
8848 while d.saturating_mul(d) <= n {
8849 if n % d == 0 {
8850 return false;
8851 }
8852 d += 1;
8853 }
8854 true
8855}
8856
8857fn primes_below(n: i64, span: Span) -> Result<i64> {
8858 if n < 0 {
8859 return Err(Error::domain("counting the primes below a negative number", span));
8860 }
8861 Ok((2..n).filter(|&k| is_prime(k)).count() as i64)
8862}
8863
8864fn next_prime(n: i64, span: Span) -> Result<i64> {
8865 let mut k = n.checked_add(1).ok_or_else(|| Error::domain("no next prime", span))?;
8866 while !is_prime(k) {
8867 k = k.checked_add(1).ok_or_else(|| Error::domain("no next prime", span))?;
8868 }
8869 Ok(k)
8870}
8871
8872fn previous_prime(n: i64, span: Span) -> Result<i64> {
8873 let mut k = n - 1;
8874 while k >= 2 {
8875 if is_prime(k) {
8876 return Ok(k);
8877 }
8878 k -= 1;
8879 }
8880 Err(Error::domain(format!("there is no prime below {n}"), span))
8881}
8882
8883fn one_int(a: &Array, what: &str, span: Span) -> Result<i64> {
8885 a.to_i64_vec()
8886 .and_then(|v| v.first().copied())
8887 .ok_or_else(|| Error::domain(format!("{what} needs an integer"), span))
8888}
8889
8890fn expand(x: &Array, y: &Array, span: Span) -> Result<Array> {
8893 let mask = x
8894 .to_i64_vec()
8895 .ok_or_else(|| Error::domain("an expansion mask holds 0s and 1s", span))?;
8896 if mask.iter().any(|&b| b != 0 && b != 1) {
8897 return Err(Error::domain("an expansion mask holds 0s and 1s", span));
8898 }
8899 let ys = as_list(y);
8900 let taken = mask.iter().filter(|&&b| b == 1).count();
8901 let n = ys.items();
8902 let spread = n == 1 && taken != 1;
8904 if !spread && taken != n {
8905 return Err(Error::new(
8906 ErrorKind::Length,
8907 format!("an expansion mask taking {taken} item(s) over {n}"),
8908 Some(span),
8909 ));
8910 }
8911 let m = ys.item_size();
8912 let mut data = Data::empty(ys.dtype());
8913 let mut at = 0usize;
8914 for &b in &mask {
8915 if b == 1 {
8916 let from = if spread { 0 } else { at };
8917 for k in 0..m {
8918 push_elem(&mut data, &ys.data, from * m + k);
8919 }
8920 at += 1;
8921 } else {
8922 for _ in 0..m {
8923 data.push_fill();
8924 }
8925 }
8926 }
8927 let mut shape = ys.shape.clone();
8928 if shape.is_empty() {
8929 shape.push(mask.len());
8930 } else {
8931 shape[0] = mask.len();
8932 }
8933 Ok(Array::new(shape, data))
8934}
8935
8936fn execute(y: &Array, apl: bool, ctx: &mut Ctx<'_>, span: Span) -> Result<Array> {
8944 let Data::Char(v) = &y.data else {
8945 return Err(Error::domain("execute reads a character list", span));
8946 };
8947 let src: String = v.iter().collect();
8948 let lang = if apl { crate::Lang::Apl } else { crate::Lang::J };
8949 let dialect = ctx.cfg.rules.dialect();
8952 let nested = crate::compile(lang, &src, &dialect).map_err(|e| nested_error(e, &src, span))?;
8953 if !nested.params.is_empty() {
8954 return Err(Error::domain(
8955 "an executed string cannot take host data: `{name}` has nothing to bind to",
8956 span,
8957 ));
8958 }
8959 let mut rec = None;
8960 let (value, _) = crate::ir::run_block(&nested.stmts, None, ctx, &mut rec)
8961 .map_err(|e| nested_error(e, &src, span))?;
8962 value.ok_or_else(|| Error::domain("the executed string yielded no value", span))
8963}
8964
8965fn nested_error(e: Error, src: &str, span: Span) -> Error {
8969 let inner = e.render(src);
8970 let mut out = Error::new(e.kind, format!("in the executed string: {}", e.msg), Some(span));
8971 out.notes.push(inner.trim_end().to_string());
8972 out
8973}
8974
8975fn words(y: &Array, span: Span) -> Result<Array> {
8981 let Data::Char(v) = &y.data else {
8982 return Err(Error::domain("words reads a character list", span));
8983 };
8984 let src: Vec<char> = v.as_slice().to_vec();
8985 let n = src.len();
8986 let mut out: Vec<Array> = Vec::new();
8987 let mut i = 0usize;
8988 let numeric_start = |k: usize| -> bool {
8989 k < n && (src[k].is_ascii_digit() || src[k] == '_')
8990 };
8991 while i < n {
8992 let c = src[i];
8993 if c == ' ' || c == '\t' {
8994 i += 1;
8995 continue;
8996 }
8997 let start = i;
8998 if c == '\'' {
8999 i += 1;
9000 loop {
9001 if i >= n {
9002 return Err(Error::parse("a word list ends inside a string", span));
9003 }
9004 if src[i] == '\'' {
9005 i += 1;
9006 if i < n && src[i] == '\'' {
9007 i += 1;
9008 continue;
9009 }
9010 break;
9011 }
9012 i += 1;
9013 }
9014 } else if c.is_ascii_alphabetic() {
9015 while i < n && (src[i].is_ascii_alphanumeric() || src[i] == '_') {
9016 i += 1;
9017 }
9018 if i < n && (src[i] == '.' || src[i] == ':') {
9019 i += 1;
9020 }
9021 if src[start..i].iter().collect::<String>() == "NB." {
9023 while i < n && src[i] != '\n' {
9024 i += 1;
9025 }
9026 }
9027 } else if numeric_start(i) {
9028 loop {
9029 while i < n && (src[i].is_ascii_alphanumeric() || src[i] == '.' || src[i] == '_')
9030 {
9031 i += 1;
9032 }
9033 let mut j = i;
9035 while j < n && src[j] == ' ' {
9036 j += 1;
9037 }
9038 if j > i && numeric_start(j) {
9039 i = j;
9040 continue;
9041 }
9042 break;
9043 }
9044 } else {
9045 i += 1;
9046 while i < n && (src[i] == '.' || src[i] == ':') {
9047 i += 1;
9048 }
9049 }
9050 out.push(Array::from_chars(src[start..i].to_vec()));
9051 }
9052 let k = out.len();
9053 Ok(Array::new(vec![k], Data::Box(out.into())))
9054}
9055
9056#[cfg(test)]
9057mod tests {
9058 use super::*;
9059
9060 macro_rules! ctx {
9062 ($name:ident, $agreement:expr) => {
9063 let mut sink = |_: &str| {};
9064 let mut env = Env::new(Vec::new());
9065 #[allow(unused_mut)]
9066 let mut $name = Ctx {
9067 cfg: EvalCfg {
9068 agreement: $agreement,
9069 fmt: FmtOpts::J,
9070 tol: Tol::J,
9071 rules: crate::frontend::Dialect::default()
9074 .rules(if $agreement == Agreement::ExactOrScalar {
9075 crate::Lang::Apl
9076 } else {
9077 crate::Lang::J
9078 })
9079 .expect("the shipped dialect is implemented"),
9080 },
9081 out: &mut sink,
9082 env: &mut env,
9083 device: None,
9084 };
9085 };
9086 ($name:ident) => {
9087 ctx!($name, Agreement::LeadingPrefix);
9088 };
9089 }
9090
9091 fn scalar_prim(name: &'static str, monad: MonadOp, dyad: DyadOp) -> Verb {
9092 Verb::Prim(Prim { name, monad, dyad, ranks: [0, 0, 0] })
9093 }
9094
9095 fn inf_prim(name: &'static str, monad: MonadOp, dyad: DyadOp) -> Verb {
9096 Verb::Prim(Prim { name, monad, dyad, ranks: [RANK_INF, RANK_INF, RANK_INF] })
9097 }
9098
9099 fn plus() -> Verb {
9100 scalar_prim("+", MonadOp::Scalar(ScalarMonad::Conj), DyadOp::Scalar(ScalarDyad::Add))
9101 }
9102 fn minus() -> Verb {
9103 scalar_prim("-", MonadOp::Scalar(ScalarMonad::Neg), DyadOp::Scalar(ScalarDyad::Sub))
9104 }
9105 fn times() -> Verb {
9106 scalar_prim("*", MonadOp::Scalar(ScalarMonad::Signum), DyadOp::Scalar(ScalarDyad::Mul))
9107 }
9108 fn pct() -> Verb {
9109 scalar_prim("%", MonadOp::Scalar(ScalarMonad::Recip), DyadOp::Scalar(ScalarDyad::DivJ))
9110 }
9111 fn div_apl() -> Verb {
9112 scalar_prim("÷", MonadOp::Scalar(ScalarMonad::Recip), DyadOp::Scalar(ScalarDyad::DivApl))
9113 }
9114 fn floor_v() -> Verb {
9115 scalar_prim("<.", MonadOp::Scalar(ScalarMonad::Floor), DyadOp::Scalar(ScalarDyad::Min))
9116 }
9117 fn ceil_v() -> Verb {
9118 scalar_prim(">.", MonadOp::Scalar(ScalarMonad::Ceil), DyadOp::Scalar(ScalarDyad::Max))
9119 }
9120 fn pow_v() -> Verb {
9121 scalar_prim("^", MonadOp::Scalar(ScalarMonad::Exp), DyadOp::Scalar(ScalarDyad::Pow))
9122 }
9123 fn residue_v() -> Verb {
9124 scalar_prim("|", MonadOp::Scalar(ScalarMonad::Abs), DyadOp::Scalar(ScalarDyad::Residue))
9125 }
9126 fn eq_v() -> Verb {
9127 scalar_prim("=", MonadOp::None, DyadOp::Scalar(ScalarDyad::Eq))
9128 }
9129 fn lt_v() -> Verb {
9130 scalar_prim("<", MonadOp::None, DyadOp::Scalar(ScalarDyad::Lt))
9131 }
9132 fn not_v() -> Verb {
9133 scalar_prim("-.", MonadOp::Scalar(ScalarMonad::Not), DyadOp::None)
9134 }
9135 fn sqrt_v() -> Verb {
9136 scalar_prim("%:", MonadOp::Scalar(ScalarMonad::Sqrt), DyadOp::NotYet("dyadic root"))
9137 }
9138 fn dollar() -> Verb {
9139 inf_prim("$", MonadOp::ShapeOf, DyadOp::Reshape)
9140 }
9141 fn pound() -> Verb {
9142 inf_prim("#", MonadOp::Tally, DyadOp::NotYet("copy"))
9143 }
9144 fn comma() -> Verb {
9145 inf_prim(",", MonadOp::Ravel, DyadOp::NotYet("append"))
9146 }
9147 fn transpose_v() -> Verb {
9148 inf_prim("|:", MonadOp::TransposeAxes, DyadOp::NotYet("dyadic transpose"))
9149 }
9150 fn head_v() -> Verb {
9151 inf_prim("{.", MonadOp::Head, DyadOp::Take)
9152 }
9153 fn behead_v() -> Verb {
9154 inf_prim("}.", MonadOp::Behead, DyadOp::Drop)
9155 }
9156 fn iota() -> Verb {
9157 inf_prim("i.", MonadOp::IotaJ, DyadOp::NotYet("index of"))
9158 }
9159 fn iota_apl(origin: i64) -> Verb {
9160 inf_prim("⍳", MonadOp::IotaApl { origin }, DyadOp::NotYet("index of"))
9161 }
9162 fn right_v() -> Verb {
9163 inf_prim("]", MonadOp::Same, DyadOp::Right)
9164 }
9165 fn echo_v() -> Verb {
9166 inf_prim("echo", MonadOp::Echo, DyadOp::None)
9167 }
9168
9169 fn b(v: Verb) -> Box<Verb> {
9170 Box::new(v)
9171 }
9172
9173 fn mat(rows: usize, cols: usize, v: Vec<i64>) -> Array {
9174 Array::new(vec![rows, cols], Data::I64(v.into()))
9175 }
9176
9177 fn ints(a: &Array) -> Vec<i64> {
9178 a.as_i64_slice().expect("integer result").to_vec()
9179 }
9180
9181 fn floats(a: &Array) -> Vec<f64> {
9182 a.as_f64_slice().expect("float result").to_vec()
9183 }
9184
9185 fn bools(a: &Array) -> Vec<u8> {
9186 match &a.data {
9187 Data::Bool(v) => v.to_vec(),
9188 other => panic!("expected boolean result, got {other:?}"),
9189 }
9190 }
9191
9192 fn sp() -> Span {
9193 Span::new(0, 1)
9194 }
9195
9196 fn close(a: f64, b: f64) -> bool {
9197 (a - b).abs() < 1e-9 || (a.is_infinite() && b.is_infinite() && a.signum() == b.signum())
9198 }
9199
9200 #[test]
9203 fn names_of_primitives_and_derived_verbs() {
9204 assert_eq!(plus().name(), "+");
9205 assert_eq!(Verb::Rank(b(plus()), [1, 1, 1]).name(), "+\"1");
9206 assert_eq!(Verb::Rank(b(plus()), [0, 1, RANK_INF]).name(), "+\"0 1 _");
9207 assert_eq!(Verb::Rank(b(plus()), [RANK_INF; 3]).name(), "+\"_");
9208 assert_eq!(Verb::Reduce(b(plus())).name(), "+/");
9209 assert_eq!(Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]).name(), "+/\"1");
9210 assert_eq!(Verb::Fork(b(plus()), b(minus()), b(times())).name(), "(+ - *)");
9211 assert_eq!(
9212 Verb::NounFork(Array::scalar_i64(1), b(plus()), b(minus())).name(),
9213 "(n + -)"
9214 );
9215 assert_eq!(Verb::Hook(b(plus()), b(minus())).name(), "(+ -)");
9216 assert_eq!(Verb::Atop(b(plus()), b(minus())).name(), "(+@:-)");
9217 assert_eq!(Verb::Compose(b(plus()), b(minus())).name(), "(+&:-)");
9218 assert_eq!(Verb::BondLeft(Array::scalar_i64(1), b(plus())).name(), "(n&+)");
9219 assert_eq!(Verb::BondRight(b(plus()), Array::scalar_i64(1)).name(), "(+&n)");
9220 }
9221
9222 #[test]
9223 fn composition_applies_the_right_verb_to_both_arguments() {
9224 ctx!(c);
9225 let v = Verb::Compose(b(plus()), b(times()));
9226 let r = v.monad(&Array::from_i64(vec![-2, 0, 3]), &mut c, sp()).unwrap();
9228 assert_eq!(ints(&r), vec![-1, 0, 1]);
9229 let r = v
9230 .dyad(&Array::scalar_i64(-5), &Array::scalar_i64(7), &mut c, sp())
9231 .unwrap();
9232 assert_eq!(ints(&r), vec![0]);
9233 let bond = Verb::BondLeft(Array::scalar_i64(10), b(minus()));
9235 let r = bond.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
9236 assert_eq!(ints(&r), vec![9, 8]);
9237 let e = bond
9238 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(2), &mut c, sp())
9239 .unwrap_err();
9240 assert_eq!(e.kind, ErrorKind::Domain);
9241 let bond = Verb::BondRight(b(minus()), Array::scalar_i64(10));
9242 let r = bond.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
9243 assert_eq!(ints(&r), vec![-9, -8]);
9244 }
9245
9246 #[test]
9249 fn scalar_monad_covers_the_whole_buffer() {
9250 ctx!(c);
9251 let r = minus().monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
9252 assert_eq!(r.shape, vec![2, 3]);
9253 assert_eq!(ints(&r), vec![-1, -2, -3, -4, -5, -6]);
9254 }
9255
9256 #[test]
9257 fn leading_prefix_agreement_broadcasts_per_row() {
9258 ctx!(c);
9259 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
9260 let y = Array::from_i64(vec![10, 20]);
9261 let r = plus().dyad(&x, &y, &mut c, sp()).unwrap();
9262 assert_eq!(r.shape, vec![2, 3]);
9263 assert_eq!(ints(&r), vec![11, 12, 13, 24, 25, 26]);
9264 let r = plus().dyad(&y, &x, &mut c, sp()).unwrap();
9266 assert_eq!(ints(&r), vec![11, 12, 13, 24, 25, 26]);
9267 }
9268
9269 #[test]
9270 fn exact_or_scalar_rejects_a_prefix_frame() {
9271 ctx!(c, Agreement::ExactOrScalar);
9272 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
9273 let y = Array::from_i64(vec![10, 20]);
9274 let e = plus().dyad(&x, &y, &mut c, sp()).unwrap_err();
9275 assert_eq!(e.kind, ErrorKind::Shape);
9276 assert!(e.msg.contains("2 3"), "{}", e.msg);
9277 assert!(e.msg.contains("right shape 2"), "{}", e.msg);
9278 }
9279
9280 #[test]
9281 fn exact_or_scalar_accepts_equal_frames_and_scalars() {
9282 ctx!(c, Agreement::ExactOrScalar);
9283 let x = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
9284 let r = plus().dyad(&x, &x, &mut c, sp()).unwrap();
9285 assert_eq!(ints(&r), vec![2, 4, 6, 8, 10, 12]);
9286 let r = plus().dyad(&Array::scalar_i64(10), &x, &mut c, sp()).unwrap();
9287 assert_eq!(r.shape, vec![2, 3]);
9288 assert_eq!(ints(&r), vec![11, 12, 13, 14, 15, 16]);
9289 let r = plus().dyad(&x, &Array::scalar_i64(10), &mut c, sp()).unwrap();
9290 assert_eq!(ints(&r), vec![11, 12, 13, 14, 15, 16]);
9291 }
9292
9293 #[test]
9294 fn vector_length_mismatch_is_a_length_error() {
9295 ctx!(c);
9296 let e = plus()
9297 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::from_i64(vec![1, 2, 3, 4, 5]), &mut c, sp())
9298 .unwrap_err();
9299 assert_eq!(e.kind, ErrorKind::Length);
9300 assert!(e.msg.contains("left shape 3"), "{}", e.msg);
9301 assert!(e.msg.contains("right shape 5"), "{}", e.msg);
9302 assert!(e.notes[0].contains("axis 0"), "{:?}", e.notes);
9303 }
9304
9305 #[test]
9306 fn diverging_matrix_frames_name_the_axis() {
9307 ctx!(c);
9308 let e = plus()
9309 .dyad(&mat(2, 3, vec![0; 6]), &mat(2, 4, vec![0; 8]), &mut c, sp())
9310 .unwrap_err();
9311 assert_eq!(e.kind, ErrorKind::Shape);
9312 assert!(e.notes[0].contains("axis 1"), "{:?}", e.notes);
9313 }
9314
9315 #[test]
9316 fn dyadic_rank_pairs_rows_with_the_whole_right_argument() {
9317 ctx!(c);
9318 let v = Verb::Rank(b(plus()), [0, 1, 1]);
9320 let r = v
9321 .dyad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &Array::from_i64(vec![10, 20, 30]), &mut c, sp())
9322 .unwrap();
9323 assert_eq!(r.shape, vec![2, 3]);
9324 assert_eq!(ints(&r), vec![11, 22, 33, 14, 25, 36]);
9325 }
9326
9327 #[test]
9328 fn surplus_frame_axes_repeat_the_shorter_frames_cells() {
9329 ctx!(c);
9330 let v = Verb::Rank(b(head_v()), [0, 0, 1]);
9333 let x = mat(2, 2, vec![1, 1, 2, 2]);
9334 let y = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
9335 let r = v.dyad(&x, &y, &mut c, sp()).unwrap();
9336 assert_eq!(r.shape, vec![2, 2, 2]);
9337 assert_eq!(ints(&r), vec![1, 0, 1, 0, 4, 5, 4, 5]);
9338 }
9339
9340 #[test]
9341 fn an_empty_frame_pairs_its_single_cell_with_every_other_cell() {
9342 ctx!(c, Agreement::ExactOrScalar);
9343 let v = Verb::Rank(b(head_v()), [0, 0, 1]);
9345 let x = Array::from_i64(vec![1, 2]);
9346 let y = Array::from_i64(vec![7, 8, 9]);
9347 let r = v.dyad(&x, &y, &mut c, sp()).unwrap();
9348 assert_eq!(r.shape, vec![2, 2]);
9349 assert_eq!(ints(&r), vec![7, 0, 7, 8]);
9350 }
9351
9352 #[test]
9353 fn negative_rank_leaves_frame_axes() {
9354 ctx!(c);
9355 let v = Verb::Rank(b(dollar()), [-1, -1, -1]);
9357 let r = v.monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
9358 assert_eq!(r.shape, vec![2, 1]);
9359 assert_eq!(ints(&r), vec![3, 3]);
9360 }
9361
9362 #[test]
9363 fn effective_rank_clamps_and_counts_back() {
9364 assert_eq!(effective_rank(0, 3), 0);
9365 assert_eq!(effective_rank(2, 1), 1);
9366 assert_eq!(effective_rank(RANK_INF, 4), 4);
9367 assert_eq!(effective_rank(-1, 3), 2);
9368 assert_eq!(effective_rank(-5, 3), 0);
9369 }
9370
9371 #[test]
9374 fn reduction_folds_right_to_left() {
9375 ctx!(c);
9376 let r = Verb::Reduce(b(minus()))
9378 .monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp())
9379 .unwrap();
9380 assert!(r.shape.is_empty());
9381 assert_eq!(ints(&r), vec![2]);
9382 }
9383
9384 #[test]
9385 fn reduction_of_one_item_and_of_a_scalar() {
9386 ctx!(c);
9387 let r = Verb::Reduce(b(plus()))
9388 .monad(&Array::from_i64(vec![7]), &mut c, sp())
9389 .unwrap();
9390 assert!(r.shape.is_empty());
9391 assert_eq!(ints(&r), vec![7]);
9392 let r = Verb::Reduce(b(plus())).monad(&Array::scalar_i64(7), &mut c, sp()).unwrap();
9393 assert_eq!(ints(&r), vec![7]);
9394 }
9395
9396 #[test]
9397 fn reduction_runs_along_the_leading_axis() {
9398 ctx!(c);
9399 let r = Verb::Reduce(b(plus()))
9400 .monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp())
9401 .unwrap();
9402 assert_eq!(r.shape, vec![3]);
9403 assert_eq!(ints(&r), vec![5, 7, 9]);
9404 }
9405
9406 #[test]
9407 fn rank_wrapped_reduction_sums_the_last_axis() {
9408 ctx!(c);
9409 let v = Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]);
9410 let r = v.monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
9411 assert_eq!(r.shape, vec![2]);
9412 assert_eq!(ints(&r), vec![6, 15]);
9413 }
9414
9415 #[test]
9416 fn empty_reduction_uses_the_identity_cell() {
9417 ctx!(c);
9418 let empty = Array::new(vec![0, 2], Data::I64(vec![].into()));
9419 let r = Verb::Reduce(b(plus())).monad(&empty, &mut c, sp()).unwrap();
9420 assert_eq!(r.shape, vec![2]);
9421 assert_eq!(ints(&r), vec![0, 0]);
9422 let r = Verb::Reduce(b(times())).monad(&empty, &mut c, sp()).unwrap();
9423 assert_eq!(ints(&r), vec![1, 1]);
9424 let r = Verb::Reduce(b(floor_v())).monad(&empty, &mut c, sp()).unwrap();
9425 assert!(floats(&r).iter().all(|&x| x == f64::INFINITY));
9426 let r = Verb::Reduce(b(ceil_v())).monad(&empty, &mut c, sp()).unwrap();
9427 assert!(floats(&r).iter().all(|&x| x == f64::NEG_INFINITY));
9428 let r = Verb::Reduce(b(minus())).monad(&empty, &mut c, sp()).unwrap();
9431 assert_eq!(ints(&r), vec![0, 0]);
9432 let r = Verb::Reduce(b(pct())).monad(&empty, &mut c, sp()).unwrap();
9433 assert_eq!(ints(&r), vec![1, 1]);
9434 let r = Verb::Reduce(b(eq_v())).monad(&empty, &mut c, sp()).unwrap();
9435 assert_eq!(bools(&r), vec![1, 1]);
9436 let r = Verb::Reduce(b(plus()))
9438 .monad(&Array::empty(DType::I64), &mut c, sp())
9439 .unwrap();
9440 assert!(r.shape.is_empty());
9441 assert_eq!(ints(&r), vec![0]);
9442 }
9443
9444 #[test]
9445 fn empty_reduction_without_an_identity_is_a_domain_error() {
9446 ctx!(c);
9447 let v = Verb::Hook(b(plus()), b(minus()));
9451 let e = Verb::Reduce(b(v)).monad(&Array::empty(DType::I64), &mut c, sp()).unwrap_err();
9452 assert_eq!(e.kind, ErrorKind::Domain);
9453 assert!(e.msg.contains("identity"), "{}", e.msg);
9454 }
9455
9456 #[test]
9457 fn reduction_with_a_non_primitive_verb_uses_the_general_fold() {
9458 ctx!(c);
9459 let v = Verb::Reduce(b(Verb::Hook(b(plus()), b(minus()))));
9461 let r = v.monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp()).unwrap();
9462 assert_eq!(ints(&r), vec![2]);
9463 }
9464
9465 #[test]
9466 fn dyadic_reduction_is_the_table() {
9467 ctx!(c);
9468 let v = Verb::Reduce(b(plus()));
9471 let r = v
9472 .dyad(&Array::scalar_i64(2), &Array::from_i64(vec![1, 2, 3]), &mut c, sp())
9473 .unwrap();
9474 assert_eq!(r.shape, vec![3]);
9475 assert_eq!(ints(&r), vec![3, 4, 5]);
9476 let r = v
9479 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::from_i64(vec![10, 20]), &mut c, sp())
9480 .unwrap();
9481 assert_eq!(r.shape, vec![3, 2]);
9482 assert_eq!(ints(&r), vec![11, 21, 12, 22, 13, 23]);
9483 let cat = Verb::Reduce(b(inf_prim(",", MonadOp::Ravel, DyadOp::AppendLeading)));
9485 let r = cat
9486 .dyad(&Array::from_i64(vec![1, 2]), &Array::from_i64(vec![3, 4]), &mut c, sp())
9487 .unwrap();
9488 assert_eq!(r.shape, vec![4]);
9489 assert_eq!(ints(&r), vec![1, 2, 3, 4]);
9490 }
9491
9492 #[test]
9495 fn integer_overflow_promotes_the_whole_result_to_float() {
9496 ctx!(c);
9497 let r = plus()
9498 .dyad(&Array::from_i64(vec![1, i64::MAX]), &Array::scalar_i64(1), &mut c, sp())
9499 .unwrap();
9500 assert_eq!(r.dtype(), DType::F64);
9501 let v = floats(&r);
9502 assert!(close(v[0], 2.0));
9503 assert!(close(v[1], i64::MAX as f64 + 1.0));
9504 let r = plus()
9506 .dyad(&Array::from_i64(vec![1, 2]), &Array::scalar_i64(1), &mut c, sp())
9507 .unwrap();
9508 assert_eq!(r.dtype(), DType::I64);
9509 }
9510
9511 #[test]
9512 fn reduction_overflow_promotes_too() {
9513 ctx!(c);
9514 let r = Verb::Reduce(b(plus()))
9515 .monad(&Array::from_i64(vec![i64::MAX, i64::MAX]), &mut c, sp())
9516 .unwrap();
9517 assert_eq!(r.dtype(), DType::F64);
9518 assert!(close(floats(&r)[0], 2.0 * i64::MAX as f64));
9519 }
9520
9521 #[test]
9522 fn booleans_widen_to_integers_in_arithmetic() {
9523 ctx!(c);
9524 let bits = Array { shape: vec![3], data: Data::Bool(vec![1, 0, 1].into()) };
9525 let r = plus().dyad(&bits, &bits, &mut c, sp()).unwrap();
9526 assert_eq!(r.dtype(), DType::I64);
9527 assert_eq!(ints(&r), vec![2, 0, 2]);
9528 }
9529
9530 #[test]
9531 fn j_division_is_float_and_survives_zero() {
9532 ctx!(c);
9533 let r = pct()
9534 .dyad(&Array::from_i64(vec![1, -1, 0, 6]), &Array::from_i64(vec![0, 0, 0, 4]), &mut c, sp())
9535 .unwrap();
9536 let v = floats(&r);
9537 assert_eq!(v[0], f64::INFINITY);
9538 assert_eq!(v[1], f64::NEG_INFINITY);
9539 assert_eq!(v[2], 0.0);
9540 assert!(close(v[3], 1.5));
9541 }
9542
9543 #[test]
9544 fn apl_division_by_zero_is_a_domain_error_except_zero_by_zero() {
9545 ctx!(c, Agreement::ExactOrScalar);
9546 let r = div_apl()
9547 .dyad(&Array::scalar_i64(0), &Array::scalar_i64(0), &mut c, sp())
9548 .unwrap();
9549 assert!(close(floats(&r)[0], 1.0));
9550 let e = div_apl()
9551 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(0), &mut c, sp())
9552 .unwrap_err();
9553 assert_eq!(e.kind, ErrorKind::Domain);
9554 assert!(e.msg.contains("division by zero"), "{}", e.msg);
9555 let r = div_apl()
9556 .dyad(&Array::scalar_i64(6), &Array::scalar_i64(4), &mut c, sp())
9557 .unwrap();
9558 assert!(close(floats(&r)[0], 1.5));
9559 }
9560
9561 #[test]
9562 fn reciprocal_of_zero_is_infinite() {
9563 ctx!(c);
9564 let r = pct().monad(&Array::from_i64(vec![0, 2]), &mut c, sp()).unwrap();
9565 let v = floats(&r);
9566 assert_eq!(v[0], f64::INFINITY);
9567 assert!(close(v[1], 0.5));
9568 }
9569
9570 #[test]
9571 fn residue_takes_the_sign_of_the_left_argument() {
9572 ctx!(c);
9573 let x = Array::from_i64(vec![3, 3, -3, -3, 0]);
9574 let y = Array::from_i64(vec![5, -5, 5, -5, 5]);
9575 let r = residue_v().dyad(&x, &y, &mut c, sp()).unwrap();
9576 assert_eq!(ints(&r), vec![2, 1, -1, -2, 5]);
9577 let r = residue_v()
9579 .dyad(&Array::from_f64(vec![2.5]), &Array::from_f64(vec![7.0]), &mut c, sp())
9580 .unwrap();
9581 assert!(close(floats(&r)[0], 2.0));
9582 }
9583
9584 #[test]
9585 fn power_stays_integral_when_it_can() {
9586 ctx!(c);
9587 let r = pow_v()
9588 .dyad(&Array::from_i64(vec![2, 0, 5]), &Array::from_i64(vec![10, 0, 1]), &mut c, sp())
9589 .unwrap();
9590 assert_eq!(r.dtype(), DType::I64);
9591 assert_eq!(ints(&r), vec![1024, 1, 5]);
9592 let r = pow_v()
9594 .dyad(&Array::from_i64(vec![2, 4]), &Array::from_i64(vec![-1, 2]), &mut c, sp())
9595 .unwrap();
9596 assert_eq!(r.dtype(), DType::F64);
9597 assert!(close(floats(&r)[0], 0.5));
9598 assert!(close(floats(&r)[1], 16.0));
9599 let r = pow_v()
9601 .dyad(&Array::scalar_i64(10), &Array::scalar_i64(30), &mut c, sp())
9602 .unwrap();
9603 assert_eq!(r.dtype(), DType::F64);
9604 }
9605
9606 #[test]
9607 fn comparisons_yield_booleans() {
9608 ctx!(c);
9609 let r = lt_v()
9610 .dyad(&Array::from_i64(vec![1, 2, 3]), &Array::scalar_i64(2), &mut c, sp())
9611 .unwrap();
9612 assert_eq!(bools(&r), vec![1, 0, 0]);
9613 let r = eq_v()
9614 .dyad(&Array::from_f64(vec![1.0, 2.0]), &Array::from_i64(vec![1, 3]), &mut c, sp())
9615 .unwrap();
9616 assert_eq!(bools(&r), vec![1, 0]);
9617 }
9618
9619 #[test]
9620 fn characters_compare_but_do_not_add() {
9621 ctx!(c);
9622 let a = Array::from_chars(vec!['a', 'b']);
9623 let bb = Array::from_chars(vec!['a', 'c']);
9624 assert_eq!(bools(&eq_v().dyad(&a, &bb, &mut c, sp()).unwrap()), vec![1, 0]);
9625 let e = plus().dyad(&a, &bb, &mut c, sp()).unwrap_err();
9626 assert_eq!(e.kind, ErrorKind::Type);
9627 assert!(e.msg.contains("characters"), "{}", e.msg);
9628 let e = lt_v().dyad(&a, &bb, &mut c, sp()).unwrap_err();
9629 assert_eq!(e.kind, ErrorKind::Type);
9630 let e = plus().dyad(&a, &Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap_err();
9631 assert_eq!(e.kind, ErrorKind::Type);
9632 assert!(e.msg.contains("character"), "{}", e.msg);
9633 let e = plus().monad(&a, &mut c, sp()).unwrap_err();
9634 assert_eq!(e.kind, ErrorKind::Type);
9635 }
9636
9637 #[test]
9638 fn floor_and_ceiling_return_integers_when_they_fit() {
9639 ctx!(c);
9640 let r = floor_v().monad(&Array::from_f64(vec![1.5, -1.5]), &mut c, sp()).unwrap();
9641 assert_eq!(r.dtype(), DType::I64);
9642 assert_eq!(ints(&r), vec![1, -2]);
9643 let r = ceil_v().monad(&Array::from_f64(vec![1.5, -1.5]), &mut c, sp()).unwrap();
9644 assert_eq!(ints(&r), vec![2, -1]);
9645 let r = floor_v().monad(&Array::from_f64(vec![1e30]), &mut c, sp()).unwrap();
9647 assert_eq!(r.dtype(), DType::F64);
9648 let r = floor_v().monad(&Array::from_i64(vec![3]), &mut c, sp()).unwrap();
9650 assert_eq!(ints(&r), vec![3]);
9651 }
9652
9653 #[test]
9654 fn logical_negation_needs_zero_or_one() {
9655 ctx!(c);
9656 let r = not_v().monad(&Array::from_i64(vec![0, 1]), &mut c, sp()).unwrap();
9657 assert_eq!(bools(&r), vec![1, 0]);
9658 let e = not_v().monad(&Array::from_i64(vec![2]), &mut c, sp()).unwrap_err();
9659 assert_eq!(e.kind, ErrorKind::Domain);
9660 }
9661
9662 #[test]
9663 fn signum_abs_and_negation_pick_their_types() {
9664 ctx!(c);
9665 let r = times().monad(&Array::from_i64(vec![-3, 0, 9]), &mut c, sp()).unwrap();
9666 assert_eq!(ints(&r), vec![-1, 0, 1]);
9667 let r = times().monad(&Array::from_f64(vec![-3.0, 0.0, 9.0]), &mut c, sp()).unwrap();
9668 assert_eq!(floats(&r), vec![-1.0, 0.0, 1.0]);
9669 let r = residue_v().monad(&Array::from_i64(vec![-3, 3]), &mut c, sp()).unwrap();
9670 assert_eq!(ints(&r), vec![3, 3]);
9671 let bits = Array { shape: vec![2], data: Data::Bool(vec![0, 1].into()) };
9672 let r = minus().monad(&bits, &mut c, sp()).unwrap();
9673 assert_eq!(r.dtype(), DType::I64);
9674 assert_eq!(ints(&r), vec![0, -1]);
9675 }
9676
9677 #[test]
9678 fn square_root_of_a_negative_number_is_complex() {
9679 ctx!(c);
9680 let r = sqrt_v().monad(&Array::from_i64(vec![9]), &mut c, sp()).unwrap();
9681 assert!(close(floats(&r)[0], 3.0));
9682 let r = sqrt_v().monad(&Array::from_i64(vec![-4]), &mut c, sp()).unwrap();
9683 assert_eq!(r.dtype(), DType::Complex);
9684 assert_eq!(r.as_complex_slice().expect("complex data"), &[[0.0, 2.0]]);
9685 }
9686
9687 #[test]
9690 fn shape_tally_and_ravel() {
9691 ctx!(c);
9692 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
9693 let r = dollar().monad(&m, &mut c, sp()).unwrap();
9694 assert_eq!(r.shape, vec![2]);
9695 assert_eq!(ints(&r), vec![2, 3]);
9696 let r = pound().monad(&m, &mut c, sp()).unwrap();
9697 assert!(r.shape.is_empty());
9698 assert_eq!(ints(&r), vec![2]);
9699 let r = pound().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap();
9701 assert_eq!(ints(&r), vec![1]);
9702 let r = comma().monad(&m, &mut c, sp()).unwrap();
9703 assert_eq!(r.shape, vec![6]);
9704 assert_eq!(ints(&r), vec![1, 2, 3, 4, 5, 6]);
9705 }
9706
9707 #[test]
9708 fn transpose_reverses_the_axes() {
9709 ctx!(c);
9710 let r = transpose_v().monad(&mat(2, 3, vec![1, 2, 3, 4, 5, 6]), &mut c, sp()).unwrap();
9711 assert_eq!(r.shape, vec![3, 2]);
9712 assert_eq!(ints(&r), vec![1, 4, 2, 5, 3, 6]);
9713 let a = Array::new(vec![2, 1, 3], Data::I64(vec![1, 2, 3, 4, 5, 6].into()));
9715 let r = transpose_v().monad(&a, &mut c, sp()).unwrap();
9716 assert_eq!(r.shape, vec![3, 1, 2]);
9717 assert_eq!(ints(&r), vec![1, 4, 2, 5, 3, 6]);
9718 let v = Array::from_i64(vec![1, 2]);
9720 assert_eq!(transpose_v().monad(&v, &mut c, sp()).unwrap(), v);
9721 }
9722
9723 #[test]
9724 fn head_and_behead() {
9725 ctx!(c);
9726 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
9727 let r = head_v().monad(&m, &mut c, sp()).unwrap();
9728 assert_eq!(r.shape, vec![3]);
9729 assert_eq!(ints(&r), vec![1, 2, 3]);
9730 let r = behead_v().monad(&m, &mut c, sp()).unwrap();
9731 assert_eq!(r.shape, vec![1, 3]);
9732 assert_eq!(ints(&r), vec![4, 5, 6]);
9733 let e = Array::new(vec![0, 2], Data::I64(vec![].into()));
9735 let r = head_v().monad(&e, &mut c, sp()).unwrap();
9736 assert_eq!(r.shape, vec![2]);
9737 assert_eq!(ints(&r), vec![0, 0]);
9738 assert_eq!(behead_v().monad(&e, &mut c, sp()).unwrap(), e);
9739 assert_eq!(head_v().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap().shape, Vec::<usize>::new());
9740 let err = behead_v().monad(&Array::scalar_i64(5), &mut c, sp()).unwrap_err();
9741 assert_eq!(err.kind, ErrorKind::Domain);
9742 }
9743
9744 #[test]
9745 fn iota_fills_a_shape_and_reverses_negative_axes() {
9746 ctx!(c);
9747 let r = iota().monad(&Array::from_i64(vec![2, 3]), &mut c, sp()).unwrap();
9748 assert_eq!(r.shape, vec![2, 3]);
9749 assert_eq!(ints(&r), vec![0, 1, 2, 3, 4, 5]);
9750 let r = iota().monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
9752 assert_eq!(r.shape, vec![3]);
9753 assert_eq!(ints(&r), vec![0, 1, 2]);
9754 let r = iota().monad(&Array::scalar_i64(-3), &mut c, sp()).unwrap();
9756 assert_eq!(ints(&r), vec![2, 1, 0]);
9757 let r = iota().monad(&Array::from_i64(vec![2, -3]), &mut c, sp()).unwrap();
9758 assert_eq!(r.shape, vec![2, 3]);
9759 assert_eq!(ints(&r), vec![2, 1, 0, 5, 4, 3]);
9760 let r = iota().monad(&Array::from_i64(vec![-2, 3]), &mut c, sp()).unwrap();
9761 assert_eq!(ints(&r), vec![3, 4, 5, 0, 1, 2]);
9762 let r = iota().monad(&Array::scalar_i64(0), &mut c, sp()).unwrap();
9764 assert_eq!(r.shape, vec![0]);
9765 assert!(ints(&r).is_empty());
9766 let e = iota().monad(&Array::from_f64(vec![1.5]), &mut c, sp()).unwrap_err();
9768 assert_eq!(e.kind, ErrorKind::Domain);
9769 let e = iota().monad(&mat(1, 1, vec![1]), &mut c, sp()).unwrap_err();
9770 assert_eq!(e.kind, ErrorKind::Rank);
9771 }
9772
9773 #[test]
9774 fn apl_iota_starts_at_the_index_origin() {
9775 ctx!(c, Agreement::ExactOrScalar);
9776 let r = iota_apl(1).monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
9777 assert_eq!(ints(&r), vec![1, 2, 3]);
9778 let r = iota_apl(0).monad(&Array::scalar_i64(3), &mut c, sp()).unwrap();
9779 assert_eq!(ints(&r), vec![0, 1, 2]);
9780 let e = iota_apl(1).monad(&Array::scalar_i64(-1), &mut c, sp()).unwrap_err();
9781 assert_eq!(e.kind, ErrorKind::Domain);
9782 let r = iota_apl(1).monad(&Array::from_i64(vec![2, 3]), &mut c, sp()).unwrap();
9785 assert_eq!(r.shape, vec![2, 3]);
9786 assert_eq!(ints(&r.as_boxes().expect("boxed")[4]), vec![2, 2]);
9787 }
9788
9789 #[test]
9790 fn reshape_cycles_the_ravel() {
9791 ctx!(c);
9792 let r = dollar()
9793 .dyad(&Array::from_i64(vec![2, 3]), &Array::from_i64(vec![1, 2]), &mut c, sp())
9794 .unwrap();
9795 assert_eq!(r.shape, vec![2, 3]);
9796 assert_eq!(ints(&r), vec![1, 2, 1, 2, 1, 2]);
9797 let r = dollar()
9799 .dyad(&Array::scalar_i64(3), &Array::from_i64(vec![7]), &mut c, sp())
9800 .unwrap();
9801 assert_eq!(r.shape, vec![3]);
9802 assert_eq!(ints(&r), vec![7, 7, 7]);
9803 let r = dollar()
9805 .dyad(&Array::scalar_i64(2), &Array::from_chars(vec!['a', 'b', 'c']), &mut c, sp())
9806 .unwrap();
9807 assert_eq!(r.dtype(), DType::Char);
9808 let e = dollar()
9810 .dyad(&Array::scalar_i64(2), &Array::empty(DType::I64), &mut c, sp())
9811 .unwrap_err();
9812 assert_eq!(e.kind, ErrorKind::Length);
9813 assert!(e.msg.contains("empty"), "{}", e.msg);
9814 let r = dollar()
9816 .dyad(&Array::scalar_i64(0), &Array::empty(DType::I64), &mut c, sp())
9817 .unwrap();
9818 assert_eq!(r.shape, vec![0]);
9819 let e = dollar()
9820 .dyad(&Array::scalar_i64(-1), &Array::from_i64(vec![1]), &mut c, sp())
9821 .unwrap_err();
9822 assert_eq!(e.kind, ErrorKind::Domain);
9823 }
9824
9825 #[test]
9826 fn take_from_both_ends_and_beyond() {
9827 ctx!(c);
9828 let v = Array::from_i64(vec![1, 2, 3, 4]);
9829 let take = |x: Array, y: &Array, c: &mut Ctx<'_>| head_v().dyad(&x, y, c, sp()).unwrap();
9830 assert_eq!(ints(&take(Array::scalar_i64(2), &v, &mut c)), vec![1, 2]);
9831 assert_eq!(ints(&take(Array::scalar_i64(-2), &v, &mut c)), vec![3, 4]);
9832 let short = Array::from_i64(vec![1, 2, 3]);
9834 assert_eq!(ints(&take(Array::scalar_i64(6), &short, &mut c)), vec![1, 2, 3, 0, 0, 0]);
9835 assert_eq!(ints(&take(Array::scalar_i64(-6), &short, &mut c)), vec![0, 0, 0, 1, 2, 3]);
9837 let r = take(Array::scalar_i64(2), &Array::scalar_i64(5), &mut c);
9839 assert_eq!(r.shape, vec![2]);
9840 assert_eq!(ints(&r), vec![5, 0]);
9841 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
9843 let r = take(Array::scalar_i64(1), &m, &mut c);
9844 assert_eq!(r.shape, vec![1, 3]);
9845 assert_eq!(ints(&r), vec![1, 2, 3]);
9846 let r = take(Array::scalar_i64(-1), &m, &mut c);
9847 assert_eq!(ints(&r), vec![4, 5, 6]);
9848 let r = take(Array::from_i64(vec![2, 2]), &m, &mut c);
9849 assert_eq!(r.shape, vec![2, 2]);
9850 assert_eq!(ints(&r), vec![1, 2, 4, 5]);
9851 let r = take(Array::from_i64(vec![3, -2]), &m, &mut c);
9852 assert_eq!(r.shape, vec![3, 2]);
9853 assert_eq!(ints(&r), vec![2, 3, 5, 6, 0, 0]);
9854 let r = head_v()
9856 .dyad(&Array::scalar_i64(3), &Array::from_chars(vec!['a']), &mut c, sp())
9857 .unwrap();
9858 assert_eq!(r.data, Data::Char(vec!['a', ' ', ' '].into()));
9859 let e = head_v()
9860 .dyad(&Array::from_i64(vec![1, 1]), &Array::from_i64(vec![1, 2]), &mut c, sp())
9861 .unwrap_err();
9862 assert_eq!(e.kind, ErrorKind::NotYet);
9863 }
9864
9865 #[test]
9866 fn drop_from_both_ends_and_beyond() {
9867 ctx!(c);
9868 let v = Array::from_i64(vec![1, 2, 3]);
9869 let drop = |x: Array, y: &Array, c: &mut Ctx<'_>| behead_v().dyad(&x, y, c, sp()).unwrap();
9870 assert_eq!(ints(&drop(Array::scalar_i64(1), &v, &mut c)), vec![2, 3]);
9871 assert_eq!(ints(&drop(Array::scalar_i64(-1), &v, &mut c)), vec![1, 2]);
9872 let r = drop(Array::scalar_i64(5), &v, &mut c);
9874 assert_eq!(r.shape, vec![0]);
9875 assert!(ints(&r).is_empty());
9876 let m = mat(2, 3, vec![1, 2, 3, 4, 5, 6]);
9877 let r = drop(Array::scalar_i64(1), &m, &mut c);
9878 assert_eq!(r.shape, vec![1, 3]);
9879 assert_eq!(ints(&r), vec![4, 5, 6]);
9880 let r = drop(Array::from_i64(vec![0, -1]), &m, &mut c);
9881 assert_eq!(r.shape, vec![2, 2]);
9882 assert_eq!(ints(&r), vec![1, 2, 4, 5]);
9883 }
9884
9885 #[test]
9888 fn cells_of_unequal_shapes_are_padded_with_fills() {
9889 ctx!(c);
9890 let v = Verb::Rank(b(iota()), [0, 0, 0]);
9892 let r = v.monad(&Array::from_i64(vec![1, 2, 3]), &mut c, sp()).unwrap();
9893 assert_eq!(r.shape, vec![3, 3]);
9894 assert_eq!(ints(&r), vec![0, 0, 0, 0, 1, 0, 0, 1, 2]);
9895 }
9896
9897 #[test]
9898 fn framing_aligns_lower_rank_cells_at_the_trailing_axes() {
9899 let cells = vec![Array::from_i64(vec![1, 2]), mat(2, 2, vec![1, 2, 3, 4])];
9900 let r = assemble(&[2], cells, sp()).unwrap();
9901 assert_eq!(r.shape, vec![2, 2, 2]);
9902 assert_eq!(ints(&r), vec![1, 2, 0, 0, 1, 2, 3, 4]);
9903 }
9904
9905 #[test]
9906 fn framing_promotes_cell_types() {
9907 let cells = vec![Array::from_i64(vec![1]), Array::from_f64(vec![2.5])];
9908 let r = assemble(&[2], cells, sp()).unwrap();
9909 assert_eq!(r.dtype(), DType::F64);
9910 assert_eq!(floats(&r), vec![1.0, 2.5]);
9911 let cells = vec![Array::from_i64(vec![1]), Array::from_chars(vec!['a'])];
9913 let e = assemble(&[2], cells, sp()).unwrap_err();
9914 assert_eq!(e.kind, ErrorKind::Type);
9915 }
9916
9917 #[test]
9918 fn framing_over_an_empty_frame_yields_an_empty_result() {
9919 let r = assemble(&[0], Vec::new(), sp()).unwrap();
9920 assert_eq!(r.shape, vec![0]);
9921 assert_eq!(r.count(), 0);
9922 }
9923
9924 #[test]
9927 fn fork_applies_both_tines() {
9928 ctx!(c);
9929 let v = Verb::Fork(b(Verb::Reduce(b(plus()))), b(pct()), b(pound()));
9931 let r = v.monad(&Array::from_i64(vec![1, 2, 3, 4]), &mut c, sp()).unwrap();
9932 assert!(close(floats(&r)[0], 2.5));
9933 let v = Verb::Fork(b(minus()), b(plus()), b(plus()));
9935 let r = v
9936 .dyad(&Array::from_i64(vec![5]), &Array::from_i64(vec![3]), &mut c, sp())
9937 .unwrap();
9938 assert_eq!(ints(&r), vec![10]);
9939 }
9940
9941 #[test]
9942 fn noun_fork_supplies_a_constant_left_argument() {
9943 ctx!(c);
9944 let v = Verb::NounFork(Array::scalar_i64(10), b(minus()), b(right_v()));
9945 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
9946 assert_eq!(ints(&r), vec![9, 8]);
9947 let r = v
9948 .dyad(&Array::scalar_i64(0), &Array::from_i64(vec![1, 2]), &mut c, sp())
9949 .unwrap();
9950 assert_eq!(ints(&r), vec![9, 8]);
9951 }
9952
9953 #[test]
9954 fn hook_reuses_its_right_argument() {
9955 ctx!(c);
9956 let v = Verb::Hook(b(plus()), b(minus()));
9958 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
9959 assert_eq!(ints(&r), vec![0, 0]);
9960 let r = v
9962 .dyad(&Array::from_i64(vec![10]), &Array::from_i64(vec![3]), &mut c, sp())
9963 .unwrap();
9964 assert_eq!(ints(&r), vec![7]);
9965 }
9966
9967 #[test]
9968 fn atop_composes() {
9969 ctx!(c);
9970 let v = Verb::Atop(b(minus()), b(plus()));
9971 let r = v.monad(&Array::from_i64(vec![1, 2]), &mut c, sp()).unwrap();
9972 assert_eq!(ints(&r), vec![-1, -2]);
9973 let r = v
9974 .dyad(&Array::from_i64(vec![1]), &Array::from_i64(vec![2]), &mut c, sp())
9975 .unwrap();
9976 assert_eq!(ints(&r), vec![-3]);
9977 }
9978
9979 #[test]
9980 fn trains_apply_to_the_whole_argument() {
9981 assert_eq!(Verb::Hook(b(plus()), b(minus())).ranks(), [RANK_INF; 3]);
9983 assert_eq!(Verb::Reduce(b(plus())).ranks(), [RANK_INF; 3]);
9984 }
9985
9986 #[test]
9989 fn absent_and_unwritten_meanings_are_reported_differently() {
9990 ctx!(c);
9991 let e = eq_v().monad(&Array::scalar_i64(1), &mut c, sp()).unwrap_err();
9992 assert_eq!(e.kind, ErrorKind::Domain);
9993 assert!(e.msg.contains("no monadic meaning"), "{}", e.msg);
9994 let e = not_v()
9995 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(1), &mut c, sp())
9996 .unwrap_err();
9997 assert_eq!(e.kind, ErrorKind::Domain);
9998 assert!(e.msg.contains("no dyadic meaning"), "{}", e.msg);
9999 let e = pound()
10000 .dyad(&Array::scalar_i64(1), &Array::scalar_i64(1), &mut c, sp())
10001 .unwrap_err();
10002 assert_eq!(e.kind, ErrorKind::NotYet);
10003 assert!(e.msg.contains("copy"), "{}", e.msg);
10004 let _ = echo_v();
10006 }
10007
10008 fn seq_par<T: Send>(f: impl Fn() -> T + Sync + Send) -> (T, T) {
10017 (par::with_threads(1, &f), par::with_threads(4, &f))
10018 }
10019
10020 fn noise(n: usize) -> Vec<f64> {
10022 let mut x = 0x2545_f491_4f6c_dd1du64;
10023 (0..n)
10024 .map(|_| {
10025 x ^= x << 13;
10026 x ^= x >> 7;
10027 x ^= x << 17;
10028 (x >> 11) as f64 / (1u64 << 53) as f64 - 0.5
10029 })
10030 .collect()
10031 }
10032
10033 fn f64_mat(rows: usize, cols: usize) -> Array {
10034 Array::new(vec![rows, cols], Data::F64(noise(rows * cols).into()))
10035 }
10036
10037 const BIG: usize = 200_000;
10039
10040 #[test]
10041 fn an_elementwise_dyad_splits_into_the_same_result() {
10042 let x = Array::from_f64(noise(BIG));
10043 let y = Array::from_f64(noise(BIG).iter().map(|v| v + 0.25).collect());
10044 let (one, many) = seq_par(|| {
10045 ctx!(c);
10046 times().dyad(&x, &y, &mut c, sp()).unwrap()
10047 });
10048 assert_eq!(floats(&one), floats(&many));
10049 let (one, many) = seq_par(|| {
10051 ctx!(c);
10052 plus().dyad(&Array::scalar_f64(0.5), &y, &mut c, sp()).unwrap()
10053 });
10054 assert_eq!(floats(&one), floats(&many));
10055 }
10056
10057 #[test]
10058 fn an_elementwise_dyad_that_overflows_widens_the_same_way() {
10059 let mut v = vec![1i64; BIG];
10062 v[BIG - 3] = i64::MAX;
10063 let x = Array::from_i64(v);
10064 let (one, many) = seq_par(|| {
10065 ctx!(c);
10066 plus().dyad(&x, &x, &mut c, sp()).unwrap()
10067 });
10068 assert_eq!(one.dtype(), DType::F64);
10069 assert_eq!(floats(&one), floats(&many));
10070 }
10071
10072 #[test]
10073 fn an_elementwise_monad_splits_into_the_same_result() {
10074 let y = Array::from_f64(noise(BIG));
10075 for v in [minus(), sqrt_v(), floor_v(), pct()] {
10076 let (one, many) = seq_par(|| {
10077 ctx!(c);
10078 v.monad(&Array::from_f64(y.as_f64_slice().unwrap().iter().map(|x| x.abs()).collect()), &mut c, sp())
10079 .unwrap()
10080 });
10081 assert_eq!(one.data, many.data, "{}", v.name());
10082 }
10083 }
10084
10085 #[test]
10086 fn monadic_cells_run_in_parallel_and_frame_in_order() {
10087 let y = f64_mat(400, 512);
10090 let v = Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]);
10091 let (one, many) = seq_par(|| {
10092 ctx!(c);
10093 v.monad(&y, &mut c, sp()).unwrap()
10094 });
10095 assert_eq!(one.shape, vec![400]);
10096 assert_eq!(floats(&one), floats(&many));
10097 }
10098
10099 #[test]
10100 fn dyadic_cells_run_in_parallel_and_frame_in_order() {
10101 let x = f64_mat(400, 512);
10102 let y = f64_mat(400, 512);
10103 let v = Verb::Rank(b(plus()), [1, 1, 1]);
10105 let (one, many) = seq_par(|| {
10106 ctx!(c);
10107 v.dyad(&x, &y, &mut c, sp()).unwrap()
10108 });
10109 assert_eq!(one.shape, vec![400, 512]);
10110 assert_eq!(floats(&one), floats(&many));
10111 }
10112
10113 #[test]
10114 fn a_verb_that_writes_output_is_not_pure() {
10115 assert!(plus().is_pure());
10116 assert!(Verb::Rank(b(Verb::Reduce(b(plus()))), [1, 1, 1]).is_pure());
10117 assert!(!echo_v().is_pure());
10118 assert!(!Verb::Rank(b(Verb::Atop(b(echo_v()), b(plus()))), [1, 1, 1]).is_pure());
10119 }
10120
10121 #[test]
10122 fn an_impure_verb_keeps_its_cells_in_order() {
10123 let y = Array::new(vec![16, 8192], Data::I64((0..16 * 8192).collect::<Vec<i64>>().into()));
10126 let v = Verb::Rank(b(Verb::Atop(b(echo_v()), b(head_v()))), [1, 1, 1]);
10127 let mut seen: Vec<i64> = Vec::new();
10128 let mut sink = |s: &str| {
10129 if let Some(first) = s.split_whitespace().next() {
10130 if let Ok(n) = first.parse::<i64>() {
10131 seen.push(n);
10132 }
10133 }
10134 };
10135 let mut env = Env::new(Vec::new());
10136 let mut c = Ctx {
10137 cfg: EvalCfg {
10138 agreement: Agreement::LeadingPrefix,
10139 fmt: FmtOpts::J,
10140 tol: Tol::J,
10141 rules: Rules::default(),
10142 },
10143 out: &mut sink,
10144 env: &mut env,
10145 device: None,
10146 };
10147 v.monad(&y, &mut c, sp()).unwrap();
10148 assert_eq!(seen, (0..16).map(|i| i * 8192).collect::<Vec<i64>>());
10149 }
10150
10151 #[test]
10152 fn a_wide_item_reduce_folds_every_column_in_order() {
10153 let y = f64_mat(300, 512);
10156 for v in [plus(), minus(), floor_v()] {
10157 let (one, many) = seq_par(|| {
10158 ctx!(c);
10159 Verb::Reduce(b(v.clone())).monad(&y, &mut c, sp()).unwrap()
10160 });
10161 assert_eq!(one.shape, vec![512]);
10162 assert_eq!(floats(&one), floats(&many), "{}", v.name());
10163 }
10164 }
10165
10166 #[test]
10167 fn a_wide_item_integer_reduce_is_exact() {
10168 let n = 300;
10169 let m = 512;
10170 let y = Array::new(
10171 vec![n, m],
10172 Data::I64((0..(n * m) as i64).map(|i| i % 977 - 400).collect::<Vec<i64>>().into()),
10173 );
10174 let (one, many) = seq_par(|| {
10175 ctx!(c);
10176 Verb::Reduce(b(minus())).monad(&y, &mut c, sp()).unwrap()
10177 });
10178 assert_eq!(ints(&one), ints(&many));
10179 }
10180
10181 #[test]
10182 fn a_narrow_item_reduce_chunks_the_items() {
10183 let y = f64_mat(300_000, 8);
10187 let (one, many) = seq_par(|| {
10188 ctx!(c);
10189 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
10190 });
10191 assert_eq!(one.shape, vec![8]);
10192 for (p, q) in floats(&one).iter().zip(floats(&many)) {
10193 assert!((p - q).abs() <= 1e-12 * p.abs().max(1.0), "{p} vs {q}");
10194 }
10195 let ints_y = Array::new(
10196 vec![300_000, 8],
10197 Data::I64((0..300_000 * 8).map(|i| (i % 101) as i64 - 50).collect::<Vec<i64>>().into()),
10198 );
10199 let (one, many) = seq_par(|| {
10200 ctx!(c);
10201 Verb::Reduce(b(plus())).monad(&ints_y, &mut c, sp()).unwrap()
10202 });
10203 assert_eq!(ints(&one), ints(&many));
10204 }
10205
10206 #[test]
10207 fn a_vector_reduce_folds_the_flat_buffer() {
10208 let y = Array::from_f64(noise(BIG * 4));
10209 let (one, many) = seq_par(|| {
10210 ctx!(c);
10211 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
10212 });
10213 let (p, q) = (floats(&one)[0], floats(&many)[0]);
10214 assert!((p - q).abs() <= 1e-12 * p.abs().max(1.0), "{p} vs {q}");
10215
10216 let ints_y = Array::from_i64((0..BIG as i64 * 4).map(|i| i % 1009 - 500).collect());
10219 for v in [plus(), minus(), ceil_v()] {
10220 let (one, many) = seq_par(|| {
10221 ctx!(c);
10222 Verb::Reduce(b(v.clone())).monad(&ints_y, &mut c, sp()).unwrap()
10223 });
10224 assert_eq!(ints(&one), ints(&many), "{}", v.name());
10225 }
10226 }
10227
10228 #[test]
10229 fn a_reduce_that_overflows_falls_back_to_the_sequential_widening() {
10230 let mut v: Vec<i64> = vec![1; BIG];
10231 v[7] = i64::MAX;
10232 let y = Array::from_i64(v);
10233 let (one, many) = seq_par(|| {
10234 ctx!(c);
10235 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
10236 });
10237 assert_eq!(one.dtype(), DType::F64);
10238 assert_eq!(floats(&one), floats(&many));
10239 }
10240
10241 #[test]
10242 fn a_boolean_reduce_matches_the_sequential_promotion() {
10243 let n = BIG;
10244 let y = Array::new(
10245 vec![n],
10246 Data::Bool((0..n).map(|i| (i % 3 == 0) as u8).collect::<Vec<u8>>().into()),
10247 );
10248 let (one, many) = seq_par(|| {
10249 ctx!(c);
10250 Verb::Reduce(b(plus())).monad(&y, &mut c, sp()).unwrap()
10251 });
10252 assert_eq!(one.dtype(), DType::I64);
10253 assert_eq!(ints(&one), ints(&many));
10254 assert_eq!(ints(&one)[0], n.div_ceil(3) as i64);
10255 }
10256}