Fix nested macro expansion panic

Author: wtholliday

src/compiler.rs

@@ -554,7 +554,13 @@ impl Compiler {
         // Expand macros in all function bodies.
         for decl in &mut decls {
             if let Decl::Func(ref mut fdecl) = decl {
-                fdecl.expand_macros(&macros);
+                if let Err((loc, msg)) = fdecl.expand_macros(&macros) {
+                    if !self.quiet {
+                        print_error_with_context(loc, &msg);
+                    }
+                    self.last_errors.push(format_error(loc, &msg));
+                    return false;
+                }
             }
         }
 

src/decl.rs

@@ -97,12 +97,25 @@ impl FuncDecl {
     }
 
     /// Expand all macro invocations in this function's arena.
-    pub fn expand_macros(&mut self, macros: &std::collections::HashMap<Name, FuncDecl>) {
-        let n = self.arena.exprs.len();
-        for i in 0..n {
+    pub fn expand_macros(
+        &mut self,
+        macros: &std::collections::HashMap<Name, FuncDecl>,
+    ) -> Result<(), (Loc, String)> {
+        let mut i = 0;
+        let mut expansion_count = 0;
+        while i < self.arena.exprs.len() {
             if let Expr::Macro(name, ref args) = self.arena.exprs[i] {
+                let loc = self.arena.locs[i];
                 let args = args.clone();
                 if let Some(mac) = macros.get(&name) {
+                    expansion_count += 1;
+                    if expansion_count > 100_000 {
+                        return Err((
+                            loc,
+                            format!("macro expansion limit exceeded while expanding {}", name),
+                        ));
+                    }
+
                     let subst: Vec<(Name, ExprID)> = mac
                         .params
                         .iter()
@@ -115,9 +128,12 @@ impl FuncDecl {
 
                     self.arena.exprs[i] = self.arena.exprs[new_body].clone();
                     self.arena.locs[i] = self.arena.locs[new_body];
+                    continue;
                 }
             }
+            i += 1;
         }
+        Ok(())
     }
 }
 

tests/cases/nested_macro.lyte

@@ -0,0 +1,18 @@
+// expected stdout:
+// compilation successful
+// assert(true)
+
+macro inc(x) {
+    x = x + 1
+}
+
+macro inc_twice(x) {
+    @inc(x)
+    @inc(x)
+}
+
+main {
+    var x = 0
+    @inc_twice(x)
+    assert(x == 2)
+}