[sc-489] Use the AST iterators in more places, fixing some bugs

src/term/check/ctrs_arities.rs

@@ -49,20 +49,15 @@ impl Pattern {
     let mut to_check = vec![self];
 
     while let Some(pat) = to_check.pop() {
-      match pat {
-        Pattern::Ctr(name, args) => {
-          let expected = arities.get(name).unwrap();
-          let found = args.len();
-          if *expected != found {
-            return Err(MatchErr::CtrArityMismatch(name.clone(), found, *expected));
-          }
+      if let Pattern::Ctr(name, args) = pat {
+        let expected = arities.get(name).unwrap();
+        let found = args.len();
+        if *expected != found {
+          return Err(MatchErr::CtrArityMismatch(name.clone(), found, *expected));
         }
-        Pattern::Lst(els) | Pattern::Tup(els) => {
-          for el in els {
-            to_check.push(el);
-          }
-        }
-        Pattern::Var(..) | Pattern::Num(..) | Pattern::Str(_) => {}
+      }
+      for child in pat.children() {
+        to_check.push(child);
       }
     }
     Ok(())

src/term/check/unbound_pats.rs

@@ -46,16 +46,12 @@ impl Pattern {
     let mut unbounds = HashSet::new();
     let mut check = vec![self];
     while let Some(pat) = check.pop() {
-      match pat {
-        Pattern::Ctr(nam, args) => {
-          if !is_ctr(nam) {
-            unbounds.insert(nam.clone());
-          }
-          check.extend(args.iter());
+      if let Pattern::Ctr(nam, _) = pat {
+        if !is_ctr(nam) {
+          unbounds.insert(nam.clone());
         }
-        Pattern::Tup(args) | Pattern::Lst(args) => args.iter().for_each(|arg| check.push(arg)),
-        Pattern::Var(_) | Pattern::Num(_) | Pattern::Str(_) => {}
       }
+      check.extend(pat.children());
     }
     unbounds
   }

src/term/mod.rs

@@ -846,6 +846,15 @@ impl Pattern {
     }
   }
 
+  pub fn children_mut(&mut self) -> ChildrenIter<&mut Pattern> {
+    match self {
+      Pattern::Ctr(_, els) | Pattern::Tup(els) | Pattern::Lst(els) => {
+        ChildrenIter::Many(Box::new(els.iter_mut()))
+      }
+      Pattern::Var(_) | Pattern::Num(_) | Pattern::Str(_) => ChildrenIter::zero(),
+    }
+  }
+
   /// Returns an iterator over each subpattern in depth-first, left to right order.
   // TODO: Not lazy.
   pub fn iter(&self) -> impl DoubleEndedIterator<Item = &Pattern> + Clone {

src/term/transform/definition_pruning.rs

@@ -97,20 +97,17 @@ impl Book {
           Some(name) => self.insert_ctrs_used(name, uses, adt_encoding),
           None => self.insert_used(def_name, used, uses, adt_encoding),
         },
-        Term::Lst { els } => {
+        Term::Lst { .. } => {
           self.insert_ctrs_used(&Name::from(LIST), uses, adt_encoding);
-          for term in els {
-            to_find.push(term);
-          }
         }
         Term::Str { .. } => {
           self.insert_ctrs_used(&Name::from(STRING), uses, adt_encoding);
         }
-        _ => {
-          for child in term.children() {
-            to_find.push(child);
-          }
-        }
+        _ => {}
+      }
+
+      for child in term.children() {
+        to_find.push(child);
       }
     }
   }

src/term/transform/desugar_implicit_match_binds.rs

@@ -12,77 +12,62 @@ impl Book {
 
 impl Term {
   pub fn desugar_implicit_match_binds(&mut self, ctrs: &Constructors, adts: &Adts) {
-    let mut to_desugar = vec![self];
-
-    while let Some(term) = to_desugar.pop() {
-      match term {
-        Term::Mat { args, rules } => {
-          // Make all the matched terms variables
-          let mut match_args = vec![];
-          for arg in args.iter_mut() {
-            if let Term::Var { nam } = arg {
-              match_args.push((nam.clone(), None))
-            } else {
-              let nam = Name::new(format!("%matched_{}", match_args.len()));
-              let arg = std::mem::replace(arg, Term::Var { nam: nam.clone() });
-              match_args.push((nam, Some(arg)));
-            }
+    Term::recursive_call(move || {
+      for child in self.children_mut() {
+        child.desugar_implicit_match_binds(ctrs, adts);
+      }
+      if let Term::Mat { args, rules } = self {
+        // Make all the matched terms variables
+        let mut match_args = vec![];
+        for arg in args.iter_mut() {
+          if let Term::Var { nam } = arg {
+            match_args.push((nam.clone(), None))
+          } else {
+            let nam = Name::new(format!("%matched_{}", match_args.len()));
+            let arg = std::mem::replace(arg, Term::Var { nam: nam.clone() });
+            match_args.push((nam, Some(arg)));
           }
+        }
 
-          // Make implicit match binds explicit
-          for rule in rules.iter_mut() {
-            for ((nam, _), pat) in match_args.iter().zip(rule.pats.iter_mut()) {
-              match pat {
-                Pattern::Var(_) => (),
-                Pattern::Ctr(ctr_nam, pat_args) => {
-                  let adt = &adts[ctrs.get(ctr_nam).unwrap()];
-                  let ctr_args = adt.ctrs.get(ctr_nam).unwrap();
-                  if pat_args.is_empty() && !ctr_args.is_empty() {
-                    // Implicit ctr args
-                    *pat_args = ctr_args
-                      .iter()
-                      .map(|field| Pattern::Var(Some(Name::new(format!("{nam}.{field}")))))
-                      .collect();
-                  }
-                }
-                Pattern::Num(NumCtr::Num(_)) => (),
-                Pattern::Num(NumCtr::Succ(_, Some(_))) => (),
-                Pattern::Num(NumCtr::Succ(n, p @ None)) => {
-                  // Implicit num arg
-                  *p = Some(Some(Name::new(format!("{nam}-{n}"))));
+        // Make implicit match binds explicit
+        for rule in rules.iter_mut() {
+          for ((nam, _), pat) in match_args.iter().zip(rule.pats.iter_mut()) {
+            match pat {
+              Pattern::Var(_) => (),
+              Pattern::Ctr(ctr_nam, pat_args) => {
+                let adt = &adts[ctrs.get(ctr_nam).unwrap()];
+                let ctr_args = adt.ctrs.get(ctr_nam).unwrap();
+                if pat_args.is_empty() && !ctr_args.is_empty() {
+                  // Implicit ctr args
+                  *pat_args = ctr_args
+                    .iter()
+                    .map(|field| Pattern::Var(Some(Name::new(format!("{nam}.{field}")))))
+                    .collect();
                 }
-                Pattern::Tup(..) => (),
-                Pattern::Lst(..) => (),
-                Pattern::Str(..) => (),
               }
+              Pattern::Num(NumCtr::Num(_)) => (),
+              Pattern::Num(NumCtr::Succ(_, Some(_))) => (),
+              Pattern::Num(NumCtr::Succ(n, p @ None)) => {
+                // Implicit num arg
+                *p = Some(Some(Name::new(format!("{nam}-{n}"))));
+              }
+              Pattern::Tup(..) => (),
+              Pattern::Lst(..) => (),
+              Pattern::Str(..) => (),
             }
           }
-
-          // Add the binds to the extracted term vars.
-          *term = match_args.into_iter().rev().fold(std::mem::take(term), |nxt, (nam, val)| {
-            if let Some(val) = val {
-              // Non-Var term that was extracted.
-              Term::Let { pat: Pattern::Var(Some(nam)), val: Box::new(val), nxt: Box::new(nxt) }
-            } else {
-              nxt
-            }
-          });
-
-          // Add the next values to check
-          let mut term = term;
-          while let Term::Let { nxt, .. } = term {
-            term = nxt;
-          }
-          let Term::Mat { args: _, rules } = term else { unreachable!() };
-          to_desugar.extend(rules.iter_mut().map(|r| &mut r.body));
         }
 
-        _ => {
-          for child in term.children_mut() {
-            to_desugar.push(child);
+        // Add the binds to the extracted term vars.
+        *self = match_args.into_iter().rev().fold(std::mem::take(self), |nxt, (nam, val)| {
+          if let Some(val) = val {
+            // Non-Var term that was extracted.
+            Term::Let { pat: Pattern::Var(Some(nam)), val: Box::new(val), nxt: Box::new(nxt) }
+          } else {
+            nxt
           }
-        }
+        });
       }
-    }
+    })
   }
 }

src/term/transform/desugar_let_destructors.rs

@@ -13,20 +13,17 @@ impl Book {
 
 impl Term {
   pub fn desugar_let_destructors(&mut self) {
-    Term::recursive_call(move || match self {
-      // Only transform `let`s that are not on variables.
-      Term::Let { pat: Pattern::Var(_), .. } => {
-        for child in self.children_mut() {
-          child.desugar_let_destructors();
-        }
+    Term::recursive_call(move || {
+      for child in self.children_mut() {
+        child.desugar_let_destructors();
       }
-      Term::Let { pat, val, nxt } => {
-        let pat = std::mem::replace(pat, Pattern::Var(None));
-        let mut val = std::mem::take(val);
-        let mut nxt = std::mem::take(nxt);
 
-        val.desugar_let_destructors();
-        nxt.desugar_let_destructors();
+      if let Term::Let { pat, val, nxt } = self
+        && !pat.is_wildcard()
+      {
+        let pat = std::mem::replace(pat, Pattern::Var(None));
+        let val = std::mem::take(val);
+        let nxt = std::mem::take(nxt);
 
         let rules = vec![Rule { pats: vec![pat], body: *nxt }];
 
@@ -39,12 +36,6 @@ impl Term {
           Term::Let { pat, val, nxt: Box::new(Term::Mat { args, rules }) }
         };
       }
-
-      _ => {
-        for child in self.children_mut() {
-          child.desugar_let_destructors();
-        }
-      }
     })
   }
 }

src/term/transform/encode_pattern_matching.rs

@@ -23,23 +23,18 @@ impl Book {
 
 impl Term {
   pub fn encode_simple_matches(&mut self, ctrs: &Constructors, adts: &Adts, adt_encoding: AdtEncoding) {
-    Term::recursive_call(move || match self {
-      Term::Mat { .. } => {
+    Term::recursive_call(move || {
+      for child in self.children_mut() {
+        child.encode_simple_matches(ctrs, adts, adt_encoding)
+      }
+
+      if let Term::Mat { .. } = self {
         debug_assert!(self.is_simple_match(ctrs, adts), "{self}");
         let Term::Mat { args, rules } = self else { unreachable!() };
-        for rule in rules.iter_mut() {
-          rule.body.encode_simple_matches(ctrs, adts, adt_encoding);
-        }
         let arg = std::mem::take(&mut args[0]);
         let rules = std::mem::take(rules);
         *self = encode_match(arg, rules, ctrs, adt_encoding);
       }
-
-      _ => {
-        for child in self.children_mut() {
-          child.encode_simple_matches(ctrs, adts, adt_encoding)
-        }
-      }
     })
   }
 }

src/term/transform/eta_reduction.rs

@@ -14,34 +14,27 @@ impl Term {
   /// Eta-reduces a term and any subterms.
   /// Expects variables to be linear.
   pub fn eta_reduction(&mut self) {
-    Term::recursive_call(move || match self {
-      Term::Lam { tag: lam_tag, nam: Some(lam_var), bod } => {
-        bod.eta_reduction();
-        match bod.as_mut() {
-          Term::App { tag: arg_tag, fun, arg: box Term::Var { nam: var_nam } }
-            if lam_var == var_nam && lam_tag == arg_tag =>
-          {
-            *self = std::mem::take(fun.as_mut());
-          }
-          _ => {}
-        }
+    Term::recursive_call(move || {
+      for child in self.children_mut() {
+        child.eta_reduction()
       }
-      Term::Chn { tag: chn_tag, nam: chn_var, bod } => {
-        bod.eta_reduction();
-        match bod.as_mut() {
-          Term::App { tag: arg_tag, fun, arg: box Term::Lnk { nam: var_nam } }
-            if chn_var == var_nam && chn_tag == arg_tag =>
-          {
-            *self = std::mem::take(fun.as_mut());
-          }
-          _ => {}
+      match self {
+        Term::Lam {
+          tag: lam_tag,
+          nam: lam_var,
+          bod: box Term::App { tag: arg_tag, fun, arg: box Term::Var { nam: var_nam } },
+        } if lam_var == var_nam && lam_tag == arg_tag => {
+          *self = std::mem::take(fun.as_mut());
         }
-      }
-
-      _ => {
-        for child in self.children_mut() {
-          child.eta_reduction()
+        Term::Chn {
+          tag: chn_tag,
+          nam: chn_var,
+          bod: box Term::App { tag: arg_tag, fun, arg: box Term::Lnk { nam: var_nam } },
+        } if chn_var == var_nam && chn_tag == arg_tag => {
+          *self = std::mem::take(fun.as_mut());
         }
+
+        _ => {}
       }
     })
   }

src/term/transform/inline.rs

@@ -25,18 +25,14 @@ impl Term {
     let mut to_inline = vec![self];
 
     while let Some(term) = to_inline.pop() {
-      match term {
-        Term::Ref { nam: def_name } => {
-          if inlineables.contains(def_name) {
-            *term = defs.get(def_name).unwrap().rule().body.clone();
-          }
+      if let Term::Ref { nam: def_name } = term {
+        if inlineables.contains(def_name) {
+          *term = defs.get(def_name).unwrap().rule().body.clone();
         }
+      }
 
-        _ => {
-          for child in term.children_mut() {
-            to_inline.push(child);
-          }
-        }
+      for child in term.children_mut() {
+        to_inline.push(child);
       }
     }
   }

src/term/transform/match_defs_to_term.rs

@@ -10,10 +10,28 @@ impl Book {
 }
 
 impl Definition {
-  /// Converts a pattern matching function with multiple rules and args, into a single rule without pattern matching.
+  /// Converts a pattern matching function with multiple rules and args, into a
+  /// single rule without pattern matching.
+  ///
   /// Moves the pattern matching of the rules into a complex match expression.
   ///
-  /// Preconditions: Rule arities must be correct
+  /// Example:
+  ///
+  /// ```hvm
+  /// if True then else = then
+  /// if False then else = else
+  /// ```
+  ///
+  /// becomes
+  ///
+  /// ```hvm
+  /// if = @%x0 @%x1 @%x2 match %x0, %x1, %x2 {
+  ///   True then else: then
+  ///   False then else: else
+  /// }
+  /// ```
+  ///
+  /// Preconditions: Rule arities must be correct.
   pub fn convert_match_def_to_term(&mut self) {
     let rule = def_rules_to_match(std::mem::take(&mut self.rules));
     self.rules = vec![rule];