diff --git a/tooling/ast_fuzzer/src/program/mod.rs b/tooling/ast_fuzzer/src/program/mod.rs index a3e788454ca..83853f0e870 100644 --- a/tooling/ast_fuzzer/src/program/mod.rs +++ b/tooling/ast_fuzzer/src/program/mod.rs @@ -329,9 +329,11 @@ impl Context { 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]. diff --git a/tooling/ast_fuzzer/src/program/rewrite/limit.rs b/tooling/ast_fuzzer/src/program/rewrite/limit.rs index b1ea3492999..c2e4d001163 100644 --- a/tooling/ast_fuzzer/src/program/rewrite/limit.rs +++ b/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. diff --git a/tooling/ast_fuzzer/src/program/rewrite/mod.rs b/tooling/ast_fuzzer/src/program/rewrite/mod.rs index 94db55f9b8d..fccca8007ef 100644 --- a/tooling/ast_fuzzer/src/program/rewrite/mod.rs +++ b/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. diff --git a/tooling/ast_fuzzer/src/program/rewrite/unreachable.rs b/tooling/ast_fuzzer/src/program/rewrite/unreachable.rs new file mode 100644 index 00000000000..d344aae62e3 --- /dev/null +++ b/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::>(); + + 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 { + 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 +} diff --git a/tooling/ast_fuzzer/tests/smoke.rs b/tooling/ast_fuzzer/tests/smoke.rs index 4b708ce3d6e..fdb5b293777 100644 --- a/tooling/ast_fuzzer/tests/smoke.rs +++ b/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}") };