Skip to content

Commit 6c674c5

Browse files
kaikaliikaikalii
committed
improve performance of deep dips
1 parent 0bd8cbf commit 6c674c5

File tree

7 files changed

+128
-67
lines changed

7 files changed

+128
-67
lines changed

‎src/check.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,11 @@ impl VirtualEnv {
459459
self.sig_node(sn)?;
460460
self.handle_args_outputs(0, n);
461461
}
462+
&DipN(n) => {
463+
let [mut sig] = get_args(args)?;
464+
sig.update_args_outputs(|a, o| (a + n, o + n));
465+
self.handle_sig(sig);
466+
}
462467
ReduceContent | ReduceDepth(_) => {
463468
let [sig] = get_args(args)?;
464469
let args = sig.args().saturating_sub(sig.outputs());

‎src/compile/invert/un.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ pub static UN_PATTERNS: &[&dyn InvertPattern] = &[
208208
&ArrayPat,
209209
&UnpackPat,
210210
&DipPat,
211+
&DipNPat,
211212
&BothPat,
212213
&ImplBothPat,
213214
&BracketPat,
@@ -432,6 +433,14 @@ inverse!(DipPat, input, asm, Dip, span, [f], {
432433
Ok((input, Mod(Dip, eco_vec![inv], span)))
433434
});
434435

436+
inverse!(DipNPat, input, asm, ref, ImplMod(DipN(n), args, span), {
437+
let [f] = args.as_slice() else {
438+
return generic();
439+
};
440+
let inv = f.un_inverse(asm)?;
441+
Ok((input, ImplMod(DipN(*n), eco_vec![inv], *span)))
442+
});
443+
435444
inverse!(BothPat, input, asm, Both, span, [f], {
436445
let inv = f.un_inverse(asm)?;
437446
Ok((

‎src/compile/invert/under.rs

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -343,14 +343,29 @@ macro_rules! under {
343343
};
344344
}
345345

346-
under!(DipPat, input, g_sig, inverse, asm, Dip, span, [f], {
346+
under!(DipPat, input, g_sig, inverse, asm, {
347+
let (input, n, args, span) = match input {
348+
[Mod(Dip, args, span), input @ ..] => (input, 1, args, *span),
349+
[ImplMod(DipN(n), args, span), input @ ..] => (input, *n, args, *span),
350+
_ => return generic(),
351+
};
352+
let m = |sn: SigNode| {
353+
if n == 1 {
354+
Mod(Dip, eco_vec![sn], span)
355+
} else {
356+
ImplMod(DipN(n), eco_vec![sn], span)
357+
}
358+
};
359+
let [f] = args.as_slice() else {
360+
return generic();
361+
};
347362
if f.sig.args() == 0 {
348363
return generic();
349364
}
350365
// F inverse
351366
let inner_g_sig = Signature::new(
352-
g_sig.args().saturating_sub(1),
353-
g_sig.outputs().saturating_sub(1),
367+
g_sig.args().saturating_sub(n),
368+
g_sig.outputs().saturating_sub(n),
354369
);
355370
let (f_before, f_after) = f.under_inverse(inner_g_sig, inverse, asm)?;
356371
// Rest inverse
@@ -365,15 +380,15 @@ under!(DipPat, input, g_sig, inverse, asm, Dip, span, [f], {
365380

366381
// Make before
367382
let mut before = if !inverse || balanced {
368-
Mod(Dip, eco_vec![f_before], span)
383+
m(f_before)
369384
} else {
370385
f_before.node
371386
};
372387
before.push(rest_before);
373388
// Make after
374389
let mut after = rest_after;
375390
let after_inner = if inverse || balanced {
376-
Mod(Dip, eco_vec![f_after], span)
391+
m(f_after)
377392
} else {
378393
f_after.node
379394
};

‎src/compile/modifier.rs

Lines changed: 78 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,12 @@ use pre_eval::PreEvalMode;
1010

1111
const MAX_COMPTIME_DEPTH: usize = if cfg!(debug_assertions) { 5 } else { 20 };
1212

13+
macro_rules! func {
14+
($expr:expr) => {
15+
(|| -> UiuaResult<_> { Ok($expr) })()?
16+
};
17+
}
18+
1319
impl Compiler {
1420
fn desugar_function_pack(
1521
&mut self,
@@ -535,7 +541,7 @@ impl Compiler {
535541
Node::Mod(By, eco_vec![sn], span)
536542
}
537543
}
538-
Reach => {
544+
Reach => func!({
539545
let (sn, _) = self.monadic_modifier_op(modified)?;
540546
let span = self.add_span(modified.modifier.span.clone());
541547
let mut node = sn.node;
@@ -560,7 +566,18 @@ impl Compiler {
560566
));
561567
}
562568
node
563-
}
569+
}),
570+
Dip => func!({
571+
let (f, span) = self.monadic_modifier_op(modified)?;
572+
let span = self.add_span(span);
573+
match f.node {
574+
Node::Mod(Dip, args, _) => Node::ImplMod(ImplPrimitive::DipN(2), args, span),
575+
Node::ImplMod(ImplPrimitive::DipN(n), args, _) => {
576+
Node::ImplMod(ImplPrimitive::DipN(n + 1), args, span)
577+
}
578+
_ => Node::Mod(Dip, eco_vec![f], span),
579+
}
580+
}),
564581
Fork => {
565582
let (f, g, f_span, _) = self.dyadic_modifier_ops(modified)?;
566583
if !modified.pack_expansion && f.node.as_primitive() == Some(Primitive::Identity) {
@@ -581,31 +598,33 @@ impl Compiler {
581598
let Some(sub) = subscript else {
582599
return Ok(None);
583600
};
584-
let sub = self.validate_subscript(sub);
585-
let sub_span = sub.span;
586-
let mut sub = sub
587-
.value
588-
.map_num(|n| self.positive_subscript(n, Both, &sub_span) as u32);
589-
let span = self.add_span(modified.modifier.span.clone());
590-
let op = self.monadic_modifier_op(modified)?.0;
591-
if let Some(side) = &mut sub.side {
592-
if let Some(side_n) = side.n {
593-
if side_n > op.sig.args() {
594-
self.add_error(
595-
modified.modifier.span.clone().merge(sub_span),
596-
format!(
597-
"Sided {}'s quantifier cannot be greater than its \
598-
function's arguments, but {} > {}",
599-
Primitive::Both.format(),
600-
side_n,
601-
op.sig.args()
602-
),
603-
);
604-
side.n = Some(op.sig.args());
601+
func!({
602+
let sub = self.validate_subscript(sub);
603+
let sub_span = sub.span;
604+
let mut sub = sub
605+
.value
606+
.map_num(|n| self.positive_subscript(n, Both, &sub_span) as u32);
607+
let span = self.add_span(modified.modifier.span.clone());
608+
let op = self.monadic_modifier_op(modified)?.0;
609+
if let Some(side) = &mut sub.side {
610+
if let Some(side_n) = side.n {
611+
if side_n > op.sig.args() {
612+
self.add_error(
613+
modified.modifier.span.clone().merge(sub_span),
614+
format!(
615+
"Sided {}'s quantifier cannot be greater than its \
616+
function's arguments, but {} > {}",
617+
Primitive::Both.format(),
618+
side_n,
619+
op.sig.args()
620+
),
621+
);
622+
side.n = Some(op.sig.args());
623+
}
605624
}
606625
}
607-
}
608-
Node::ImplMod(ImplPrimitive::BothImpl(sub), eco_vec![op], span)
626+
Node::ImplMod(ImplPrimitive::BothImpl(sub), eco_vec![op], span)
627+
})
609628
}
610629
Bracket => {
611630
let Some(sub) = subscript else {
@@ -614,39 +633,41 @@ impl Compiler {
614633
let Some(side) = self.subscript_side_only(&sub, Bracket.format()) else {
615634
return Ok(None);
616635
};
617-
let span = self.add_span(modified.modifier.span.clone());
618-
let (a, b, _, _) = self.dyadic_modifier_ops(modified)?;
619-
if a.sig.args() != 2 || b.sig.args() != 2 {
620-
self.add_error(
621-
modified.modifier.span.clone().merge(sub.span),
622-
format!(
623-
"Sided {}'s functions must both have 2 arguments, \
636+
func!({
637+
let span = self.add_span(modified.modifier.span.clone());
638+
let (a, b, _, _) = self.dyadic_modifier_ops(modified)?;
639+
if a.sig.args() != 2 || b.sig.args() != 2 {
640+
self.add_error(
641+
modified.modifier.span.clone().merge(sub.span),
642+
format!(
643+
"Sided {}'s functions must both have 2 arguments, \
624644
but their signatures are {} and {}.",
625-
Primitive::Bracket.format(),
626-
a.sig,
627-
b.sig
628-
),
629-
);
630-
Node::Mod(Bracket, eco_vec![a, b], span)
631-
} else {
632-
let sub_span = self.add_span(sub.span);
633-
let mut node = match side {
634-
SubSide::Left => Node::Mod(
635-
On,
636-
eco_vec![Node::Prim(Flip, sub_span).sig_node().unwrap()],
637-
sub_span,
638-
),
639-
SubSide::Right => Node::Mod(
640-
Dip,
641-
eco_vec![Node::ImplPrim(ImplPrimitive::Over, sub_span)
642-
.sig_node()
643-
.unwrap()],
644-
sub_span,
645-
),
646-
};
647-
node.push(Node::Mod(Bracket, eco_vec![a, b], span));
648-
node
649-
}
645+
Primitive::Bracket.format(),
646+
a.sig,
647+
b.sig
648+
),
649+
);
650+
Node::Mod(Bracket, eco_vec![a, b], span)
651+
} else {
652+
let sub_span = self.add_span(sub.span);
653+
let mut node = match side {
654+
SubSide::Left => Node::Mod(
655+
On,
656+
eco_vec![Node::Prim(Flip, sub_span).sig_node().unwrap()],
657+
sub_span,
658+
),
659+
SubSide::Right => Node::Mod(
660+
Dip,
661+
eco_vec![Node::ImplPrim(ImplPrimitive::Over, sub_span)
662+
.sig_node()
663+
.unwrap()],
664+
sub_span,
665+
),
666+
};
667+
node.push(Node::Mod(Bracket, eco_vec![a, b], span));
668+
node
669+
}
670+
})
650671
}
651672
prim @ (With | Off) => {
652673
let (mut sn, _) = self.monadic_modifier_op(modified)?;

‎src/impl_prim.rs

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ macro_rules! impl_primitive {
5252
impl ImplPrimitive {
5353
pub fn args(&self) -> Option<usize> {
5454
Some(match self {
55-
$($(ImplPrimitive::$variant {..} => $args,)?)*
55+
$($(ImplPrimitive::$variant { .. } => $args,)?)*
5656
ImplPrimitive::UndoDeshape(_) => 2,
5757
ImplPrimitive::UndoTransposeN(n, _) => *n,
5858
ImplPrimitive::UndoReverse { n, .. } => *n,
@@ -67,7 +67,7 @@ macro_rules! impl_primitive {
6767
}
6868
pub fn outputs(&self) -> Option<usize> {
6969
Some(match self {
70-
$($(ImplPrimitive::$variant {..} => $outputs,)?)*
70+
$($(ImplPrimitive::$variant { .. } => $outputs,)?)*
7171
ImplPrimitive::UndoTransposeN(n, _) => *n,
7272
ImplPrimitive::UndoReverse { n, .. } => *n,
7373
ImplPrimitive::UndoRotate(n) => *n,
@@ -81,7 +81,7 @@ macro_rules! impl_primitive {
8181
}
8282
pub fn modifier_args(&self) -> Option<usize> {
8383
match self {
84-
$($(ImplPrimitive::$variant => Some($margs),)?)*
84+
$($(ImplPrimitive::$variant { .. } => Some($margs),)?)*
8585
ImplPrimitive::ReduceDepth(_)
8686
| ImplPrimitive::EachSub(_)
8787
| ImplPrimitive::RowsSub(..) => Some(1),
@@ -234,6 +234,7 @@ impl_primitive!(
234234
(2, RangeStart),
235235
// Implementation details
236236
(2(3), Over),
237+
([1], DipN(usize)),
237238
(1, NBits(usize)),
238239
(1, DeshapeSub(i32)),
239240
(1, ParseSub(usize)),
@@ -479,6 +480,12 @@ impl fmt::Display for ImplPrimitive {
479480
write!(f, "{Keep}")?;
480481
fmt_subscript(f, n as i32)
481482
}
483+
&DipN(n) => {
484+
for _ in 0..n {
485+
write!(f, "{Dip}")?;
486+
}
487+
Ok(())
488+
}
482489
}
483490
}
484491
}

‎src/run_prim.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1250,6 +1250,12 @@ impl ImplPrimitive {
12501250
env.insert_stack(outputs, kept)?;
12511251
env.push_all(temp);
12521252
}
1253+
&ImplPrimitive::DipN(n) => {
1254+
let [f] = get_ops(ops, env)?;
1255+
let dipped = env.pop_n(n)?;
1256+
env.exec(f)?;
1257+
env.push_all(dipped);
1258+
}
12531259
ImplPrimitive::Astar => {
12541260
let [neighbors, heuristic, is_goal] = get_ops(ops, env)?;
12551261
path::path(neighbors, is_goal, Some(heuristic), env)?;

‎todo.md

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@
55
- Fix rename
66
- Fix goto references
77
- `Variants` field for module
8-
- Fix caching with labels
9-
- Fix performance of deep `dip` nesting
108
- Don't serialize sortedness flags
119
- Sided `fork`
1210
- Sided `join`

0 commit comments

Comments
 (0)