chore(fuzz): Remove unreachable functions in the AST fuzzer

tooling/ast_fuzzer/src/program/mod.rs

@@ -51,7 +51,7 @@
     ctx.set_function_decl(FuncId(1), decl_inner.clone());
     ctx.gen_function(u, FuncId(1))?;
 
     // Parameterless main declaration wrapping the inner "main"
     // function call
     let decl_main = FunctionDeclaration {
         name: "main".into(),
@@ -329,9 +329,11 @@
         Ok(())
     }
 
-    /// As a post-processing step, identify recursive functions and add a call depth parameter to them.
+    /// Post-processing steps that change functions.
     fn rewrite_functions(&mut self, u: &mut Unstructured) -> arbitrary::Result<()> {
-        rewrite::add_recursion_limit(self, u)
+        rewrite::remove_unreachable_functions(self);
+        rewrite::add_recursion_limit(self, u)?;
+        Ok(())
     }
 
     /// Return the generated [Program].

tooling/ast_fuzzer/src/program/rewrite/limit.rs

@@ -42,7 +42,7 @@ pub(crate) fn add_recursion_limit(
 
     // Create proxies for unconstrained functions called from ACIR.
     let mut proxy_functions = HashMap::new();
-    let mut next_func_id = FuncId(ctx.functions.len() as u32);
+    let mut next_func_id = ctx.functions.len() as u32;
 
     for (func_id, func) in &ctx.functions {
         if !func.unconstrained
@@ -52,11 +52,11 @@ pub(crate) fn add_recursion_limit(
             continue;
         }
         let mut proxy = func.clone();
-        proxy.id = next_func_id;
+        proxy.id = FuncId(next_func_id);
         proxy.name = format!("{}_proxy", proxy.name);
         // We will replace the body, update the params, and append the function later.
         proxy_functions.insert(*func_id, proxy);
-        next_func_id = FuncId(next_func_id.0 + 1);
+        next_func_id += 1;
     }
 
     // Rewrite functions.

tooling/ast_fuzzer/src/program/rewrite/mod.rs

@@ -3,8 +3,10 @@ use noirc_frontend::monomorphization::ast::{Expression, Function, Program};
 use super::visitor::visit_expr;
 
 mod limit;
+mod unreachable;
 
 pub(crate) use limit::add_recursion_limit;
+pub(crate) use unreachable::remove_unreachable_functions;
 
 /// Find the next local ID and ident IDs (in that order) that we can use to add
 /// variables to a [Function] during mutations.

tooling/ast_fuzzer/src/program/rewrite/unreachable.rs

@@ -0,0 +1,83 @@
+use std::collections::{BTreeSet, VecDeque};
+
+use im::HashMap;
+use noirc_frontend::monomorphization::ast::{Definition, Expression, FuncId, Ident, Program};
+
+use crate::{
+    program::Context,
+    visitor::{visit_expr, visit_expr_mut},
+};
+
+/// Remove functions that are unreachable from main.
+pub(crate) fn remove_unreachable_functions(ctx: &mut Context) {
+    if ctx.functions.is_empty() {
+        return;
+    }
+    let reachable = find_reachable_functions(ctx);
+
+    // We have to re-assign function IDs, because in the `Program`
+    // the ID of the function is expected to match its position.
+    let remap = reachable
+        .into_iter()
+        .enumerate()
+        .map(|(i, id)| (id, FuncId(i as u32)))
+        .collect::<HashMap<_, _>>();
+
+    let functions = std::mem::take(&mut ctx.functions);
+    let function_declarations = std::mem::take(&mut ctx.function_declarations);
+
+    // Keep only the reachable ones.
+    ctx.functions = functions
+        .into_iter()
+        .filter_map(|(id, mut func)| {
+            remap.get(&id).map(|new_id| {
+                func.id = *new_id;
+                (*new_id, func)
+            })
+        })
+        .collect();
+
+    ctx.function_declarations = function_declarations
+        .into_iter()
+        .filter_map(|(id, func)| remap.get(&id).map(|new_id| (*new_id, func)))
+        .collect();
+
+    // Remap the old IDs to the new ones wherever the functions are referenced.
+    for func in ctx.functions.values_mut() {
+        visit_expr_mut(&mut func.body, &mut |expr| {
+            if let Expression::Ident(Ident { definition: Definition::Function(id), .. }) = expr {
+                let new_id = remap[id];
+                *id = new_id;
+            }
+            true
+        });
+    }
+}
+
+/// Find functions reachable from main.
+fn find_reachable_functions(ctx: &Context) -> BTreeSet<FuncId> {
+    let mut reachable = BTreeSet::new();
+    let mut queue = VecDeque::new();
+
+    // Start from main.
+    queue.push_back(Program::main_id());
+
+    // Find all global functions referred to by their identifier.
+    while let Some(id) = queue.pop_front() {
+        if !reachable.insert(id) {
+            continue;
+        }
+        let func = &ctx.functions[&id];
+
+        visit_expr(&func.body, &mut |expr| {
+            // Regardless of whether it's in a `Call` or stored in a reference,
+            // it will appear in an identifier at some point.
+            if let Expression::Ident(Ident { definition: Definition::Function(id), .. }) = expr {
+                queue.push_back(*id);
+            }
+            true
+        });
+    }
+
+    reachable
+}

tooling/ast_fuzzer/tests/smoke.rs

@@ -47,7 +47,9 @@ fn arb_program_can_be_executed() {
 
         // Print the AST if something goes wrong, then panic.
         let print_ast_and_panic = |msg: &str| -> ! {
-            eprintln!("{}", DisplayAstAsNoir(&program));
+            if maybe_seed.is_none() {
+                eprintln!("{}", DisplayAstAsNoir(&program));
+            }
             panic!("{msg}")
         };