make anti compilation smarter

uiua-lang · Sep 22, 2024 · 12ead79 · 12ead79
1 parent 39ea58b
commit 12ead79
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 50 deletions.
diff --git a/site/primitives.json b/site/primitives.json
@@ -965,13 +965,6 @@
     "class": "Stack",
     "description": "Duplicate the second-to-top value to the top of the stack"
   },
-  "pad": {
-    "args": 2,
-    "outputs": 1,
-    "class": "DyadicArray",
-    "description": "Pad an array",
-    "experimental": true
-  },
   "parse": {
     "glyph": "⋕",
     "args": 1,

diff --git a/src/algorithm/invert.rs b/src/algorithm/invert.rs
@@ -340,6 +340,61 @@ pub(crate) fn invert_instrs(instrs: &[Instr], comp: &mut Compiler) -> Option<Eco
     None
 }
 
+/// Invert a sequence of instructions with anti
+pub(crate) fn anti_instrs(instrs: &[Instr], comp: &mut Compiler) -> Option<EcoVec<Instr>> {
+    if instrs.is_empty() {
+        return Some(EcoVec::new());
+    }
+    dbgln!("anti-inverting {:?}", FmtInstrs(instrs, &comp.asm));
+
+    let mut inverted = EcoVec::new();
+    let mut curr_instrs = instrs;
+    'find_pattern: loop {
+        for pattern in ON_INVERT_PATTERNS {
+            if let Some((input, mut inv)) = pattern.invert_extract(curr_instrs, comp) {
+                dbgln!(
+                    "matched pattern {:?} on {:?} to {:?}",
+                    pattern,
+                    FmtInstrs(&curr_instrs[..curr_instrs.len() - input.len()], &comp.asm),
+                    FmtInstrs(&inv, &comp.asm)
+                );
+                inv.extend(inverted);
+                inverted = inv;
+                if input.is_empty() {
+                    dbgln!(
+                        "inverted {:?} to {:?}",
+                        FmtInstrs(instrs, &comp.asm),
+                        FmtInstrs(&inverted, &comp.asm)
+                    );
+                    return resolve_uns(inverted, comp);
+                }
+                curr_instrs = input;
+                continue 'find_pattern;
+            }
+        }
+        break;
+    }
+
+    dbgln!(
+        "anti-inverting {:?} failed with remaining {:?}",
+        FmtInstrs(instrs, &comp.asm),
+        FmtInstrs(curr_instrs, &comp.asm)
+    );
+
+    let sig = instrs_signature(instrs).ok()?;
+    if sig.args >= 2 {
+        let mut instrs = instrs.to_vec();
+        let span = instrs.iter().find_map(|instr| instr.span()).unwrap_or(0);
+        instrs.insert(0, Instr::copy_inline(span));
+        instrs.push(Instr::pop_inline(1, span));
+        let mut inverse = invert_instrs(&instrs, comp)?;
+        inverse.push(Instr::Prim(Primitive::Pop, span));
+        Some(inverse)
+    } else {
+        invert_instrs(instrs, comp)
+    }
+}
+
 type Under = (EcoVec<Instr>, EcoVec<Instr>);
 
 /// Calculate the "before" and "after" instructions for `under`ing a sequence of instructions.
@@ -718,18 +773,9 @@ fn resolve_uns(instrs: EcoVec<Instr>, comp: &mut Compiler) -> Option<EcoVec<Inst
                     if (instrs.peek())
                         .is_some_and(|instr| matches!(instr, Instr::Prim(Primitive::Anti, _))) =>
                 {
-                    let Some(Instr::Prim(Primitive::Anti, span)) = instrs.next() else {
-                        unreachable!()
-                    };
-                    let mut instrs = f.instrs(&comp.asm).to_vec();
-                    if f.signature().args >= 2 {
-                        instrs.insert(0, Instr::copy_inline(span));
-                        instrs.push(Instr::pop_inline(1, span));
-                    }
-                    let mut inverse = invert_instrs(&instrs, comp)?;
-                    if f.signature().args >= 2 {
-                        inverse.push(Instr::Prim(Primitive::Pop, span));
-                    }
+                    instrs.next();
+                    let instrs = f.instrs(&comp.asm).to_vec();
+                    let inverse = anti_instrs(&instrs, comp)?;
                     resolved.extend(inverse);
                 }
                 Instr::PushFunc(f) => {
@@ -842,35 +888,26 @@ fn under_anti_pattern<'a>(
     g_sig: Signature,
     comp: &mut Compiler,
 ) -> Option<(&'a [Instr], Under)> {
-    let [Instr::PushFunc(f), Instr::Prim(Primitive::Anti, span), input @ ..] = input else {
+    let [Instr::PushFunc(f), Instr::Prim(Primitive::Anti, _), input @ ..] = input else {
         return None;
     };
-    let mut instrs = EcoVec::from(f.instrs(&comp.asm));
-    if f.signature().args >= 2 {
-        instrs.insert(0, Instr::copy_inline(*span));
-        instrs.push(Instr::pop_inline(1, *span));
-    }
-    let mut befores = invert_instrs(&instrs, comp)?;
-    if f.signature().args >= 2 {
-        befores.push(Instr::Prim(Primitive::Pop, *span));
-    }
+    let instrs = EcoVec::from(f.instrs(&comp.asm));
+    let befores = anti_instrs(&instrs, comp)?;
     if let [Instr::PushFunc(_), Instr::PushFunc(_), Instr::Prim(Primitive::SetInverse, _)] =
         befores.as_slice()
     {
         if let Some(under) = under_instrs(&befores, g_sig, comp) {
             return Some((input, under));
         }
     }
-    let (befores, mut afters) = if let Some((befores, afters)) = under_instrs(&befores, g_sig, comp)
-    {
-        (befores, afters)
-    } else {
-        (befores, instrs)
-    };
-    if f.signature().args >= 2 {
-        afters.push(Instr::Prim(Primitive::Pop, *span));
-    }
-    Some((input, (befores, afters)))
+    Some((
+        input,
+        if let Some((befores, afters)) = under_instrs(&befores, g_sig, comp) {
+            (befores, afters)
+        } else {
+            (befores, instrs)
+        },
+    ))
 }
 
 fn under_call_pattern<'a>(

diff --git a/src/compile/modifier.rs b/src/compile/modifier.rs
@@ -1,5 +1,7 @@
 //! Compiler code for modifiers
 
+use crate::algorithm::invert::anti_instrs;
+
 use super::*;
 
 impl Compiler {
@@ -937,10 +939,7 @@ impl Compiler {
                 let f_res = self.compile_operand_word(f);
                 self.in_inverse = !self.in_inverse;
                 let (mut new_func, f_sig) = f_res?;
-                let spandex = self.add_span(span.clone());
-                if f_sig.args >= 2 {
-                    on(&mut new_func.instrs, f_sig, spandex);
-                } else {
+                if f_sig.args < 2 {
                     self.emit_diagnostic(
                         format!(
                             "Prefer {} over {} for functions \
@@ -951,15 +950,11 @@ impl Compiler {
                         DiagnosticKind::Style,
                         span.clone(),
                     );
-                };
+                }
 
-                if let Some(inverted) = invert_instrs(&new_func.instrs, self) {
-                    let mut sig = self.sig_of(&inverted, &span)?;
+                if let Some(inverted) = anti_instrs(&new_func.instrs, self) {
+                    let sig = self.sig_of(&inverted, &span)?;
                     new_func.instrs = inverted;
-                    if f_sig.args >= 2 {
-                        new_func.instrs.push(Instr::Prim(Primitive::Pop, spandex));
-                        sig.outputs -= 1;
-                    }
                     finish!(new_func, sig);
                 } else {
                     return Err(self.fatal_error(span, "No inverse found"));