From 51ec4e95270d85b1fd9452c15e4e5b4f6b50c1cf Mon Sep 17 00:00:00 2001 From: Akosh Farkash Date: Wed, 14 May 2025 14:42:49 +0100 Subject: [PATCH 1/3] Add ctx_limit to all functions --- tooling/ast_fuzzer/src/program/rewrite.rs | 92 +++++++++++++++-------- 1 file changed, 61 insertions(+), 31 deletions(-) diff --git a/tooling/ast_fuzzer/src/program/rewrite.rs b/tooling/ast_fuzzer/src/program/rewrite.rs index ffa0a757769..d29a6d552fc 100644 --- a/tooling/ast_fuzzer/src/program/rewrite.rs +++ b/tooling/ast_fuzzer/src/program/rewrite.rs @@ -1,7 +1,6 @@ -use std::collections::BTreeMap; +use std::collections::{HashMap, HashSet}; use arbitrary::Unstructured; -use im::HashMap; use nargo::errors::Location; use noirc_frontend::{ ast::BinaryOpKind, @@ -16,44 +15,66 @@ use super::{ visitor::{visit_expr, visit_expr_mut}, }; -/// Find recursive functions and add a `ctx_limit: &mut u32` parameter to them, -/// which we use to limit the number of recursive calls. This is complicated by -/// the fact that we cannot pass mutable references from ACIR to Brillig. To -/// overcome that, we create a proxy function for unconstrained functions that -/// take `mut ctx_limit: u32` instead, and pass it on as a mutable ref. +/// To avoid the potential of infinite recursion at runtime, add a `ctx_limit: &mut u32` +/// parameter to all functions, which we use to limit the number of recursive calls. +/// +/// This is complicated by the fact that we cannot pass mutable references from ACIR to Brillig. +/// To overcome that, we create a proxy function for unconstrained functions that take +/// `mut ctx_limit: u32` instead, and pass it on as a mutable ref. +/// +/// Originally only actually recursive functions (ie. one that called something else) +/// received this extra parameters, but in order to support higher order functions +/// which can be passed a recursive or a non-recursive function as an argument, +/// all functions get the extra parameter. pub(crate) fn add_recursion_limit( ctx: &mut Context, u: &mut Unstructured, ) -> arbitrary::Result<()> { // Collect recursive functions, ie. the ones which call other functions. - // Remember if they are unconstrained; those need proxies as well. let recursive_functions = ctx .functions .iter() - .filter_map(|(id, func)| expr::has_call(&func.body).then_some((*id, func.unconstrained))) - .collect::>(); + .filter_map(|(id, func)| expr::has_call(&func.body).then_some(*id)) + .collect::>(); + + // Collect functions called from ACIR; they will need proxy functions. + let called_from_acir = ctx.functions.values().filter(|func| !func.unconstrained).fold( + HashSet::::new(), + |mut acc, func| { + acc.extend(expr::callees(&func.body)); + acc + }, + ); - // Create proxies for unconstrained recursive functions. - // We could check whether they are called from ACIR, but that would require further traversals. + let unconstrained_functions = ctx + .functions + .iter() + .filter_map(|(id, func)| func.unconstrained.then_some(*id)) + .collect::>(); + + // 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); - for (func_id, unconstrained) in &recursive_functions { - if !*unconstrained || *func_id == Program::main_id() { + for (func_id, func) in &ctx.functions { + if !func.unconstrained + || *func_id == Program::main_id() + || !called_from_acir.contains(func_id) + { continue; } - let mut proxy = ctx.functions[func_id].clone(); + let mut proxy = func.clone(); proxy.id = next_func_id; proxy.name = format!("{}_proxy", proxy.name); - // We will replace the body and update the params later. + // 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); } - // Rewrite recursive functions. - for (func_id, unconstrained) in recursive_functions.iter() { - let func = ctx.functions.get_mut(func_id).unwrap(); + // Rewrite functions. + for (func_id, func) in ctx.functions.iter_mut() { let is_main = *func_id == Program::main_id(); + let is_recursive = recursive_functions.contains(func_id); // We'll need a new ID for variables or parameters. We could speed this up by // 1) caching this value in a "function meta" construct, or @@ -89,7 +110,7 @@ pub(crate) fn add_recursion_limit( expr::u32_literal(ctx.config.max_recursive_calls as u32), ); expr::prepend(&mut func.body, init_limit); - } else { + } else if is_recursive { // In non-main we look at the limit and return a random value if it's zero, // otherwise decrease it by one and continue with the original body. let limit_type = types::ref_mut(types::U32); @@ -132,6 +153,18 @@ pub(crate) fn add_recursion_limit( func.return_type.clone(), ) }); + } else { + // For non-recursive functions just add an unused parameter. + // In non-main we look at the limit and return a random value if it's zero, + // otherwise decrease it by one and continue with the original body. + let limit_type = types::ref_mut(types::U32); + func.parameters.push(( + limit_id, + false, + format!("_{limit_name}"), + limit_type.clone(), + Visibility::Private, + )); } // Add the non-reference version of the parameter to the proxy function. @@ -194,26 +227,23 @@ pub(crate) fn add_recursion_limit( // Update calls to pass along the limit and call the proxy if necessary. visit_expr_mut(&mut func.body, &mut |expr: &mut Expression| { if let Expression::Call(call) = expr { - let Expression::Ident(func) = call.func.as_mut() else { + let Expression::Ident(ident) = call.func.as_mut() else { unreachable!("functions are called by ident"); }; - let Definition::Function(func_id) = func.definition else { + let Definition::Function(callee_id) = ident.definition else { unreachable!("function definition expected"); }; - // If the callee isn't recursive, it won't have the extra parameter. - let Some(callee_unconstrained) = recursive_functions.get(&func_id) else { - return true; - }; - let Type::Function(param_types, _, _, _) = &mut func.typ else { + let Type::Function(param_types, _, _, _) = &mut ident.typ else { unreachable!("function type expected"); }; - if *callee_unconstrained && !unconstrained { + let callee_unconstrained = unconstrained_functions.contains(&callee_id); + if callee_unconstrained && !func.unconstrained { // Calling Brillig from ACIR: call the proxy. - let Some(proxy) = proxy_functions.get(&func_id) else { + let Some(proxy) = proxy_functions.get(&callee_id) else { unreachable!("expected to have a proxy"); }; - func.name = proxy.name.clone(); - func.definition = Definition::Function(proxy.id); + ident.name = proxy.name.clone(); + ident.definition = Definition::Function(proxy.id); // Pass the limit by value. let limit_expr = if is_main { expr::ident( From a60d506e0ef3824606ae6a1aa74ea50a3d2b96c4 Mon Sep 17 00:00:00 2001 From: Akosh Farkash Date: Wed, 14 May 2025 14:43:03 +0100 Subject: [PATCH 2/3] Do not generate ACIR if main is Brillig --- tooling/ast_fuzzer/src/program/mod.rs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tooling/ast_fuzzer/src/program/mod.rs b/tooling/ast_fuzzer/src/program/mod.rs index e9f3453d165..9eb86406fe8 100644 --- a/tooling/ast_fuzzer/src/program/mod.rs +++ b/tooling/ast_fuzzer/src/program/mod.rs @@ -169,6 +169,16 @@ impl Context { let is_main = i == 0; let num_params = u.int_in_range(0..=self.config.max_function_args)?; + // If `main` is unconstrained, it won't call ACIR, so no point generating ACIR functions. + let unconstrained = self.config.force_brillig + || (!is_main + && self + .functions + .get(&Program::main_id()) + .map(|func| func.unconstrained) + .unwrap_or_default()) + || bool::arbitrary(u)?; + let mut params = Vec::new(); for p in 0..num_params { let id = LocalId(p as u32); @@ -227,7 +237,7 @@ impl Context { } else { *u.choose(&[InlineType::Inline, InlineType::InlineAlways])? }, - unconstrained: self.config.force_brillig || bool::arbitrary(u)?, + unconstrained, }; Ok(decl) From c110a1e2045ea7275e20d244f7ce59e61ab847bd Mon Sep 17 00:00:00 2001 From: Akosh Farkash Date: Wed, 14 May 2025 15:55:51 +0100 Subject: [PATCH 3/3] Update rewrite test --- tooling/ast_fuzzer/src/program/tests.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tooling/ast_fuzzer/src/program/tests.rs b/tooling/ast_fuzzer/src/program/tests.rs index c9b54d80bff..e3227d7c0c0 100644 --- a/tooling/ast_fuzzer/src/program/tests.rs +++ b/tooling/ast_fuzzer/src/program/tests.rs @@ -196,7 +196,7 @@ fn test_recursion_limit_rewrite() { // - main passes the limit to foo by ref // - foo passes the limit to bar_proxy by value // - bar_proxy passes the limit to baz by ref - // - bar does not passes the limit to qux + // - bar passes the limit to qux, even though it's unused // - baz passes the limit to itself by ref let code = format!("{}", DisplayAstAsNoir(&program)); @@ -220,7 +220,7 @@ fn test_recursion_limit_rewrite() { } else { *ctx_limit = ((*ctx_limit) - 1); baz(ctx_limit); - qux() + qux(ctx_limit) } } unconstrained fn baz(ctx_limit: &mut u32) -> () { @@ -231,13 +231,10 @@ fn test_recursion_limit_rewrite() { baz(ctx_limit) } } - unconstrained fn qux() -> () { + unconstrained fn qux(_ctx_limit: &mut u32) -> () { } unconstrained fn bar_proxy(mut ctx_limit: u32) -> () { bar((&mut ctx_limit)) } - unconstrained fn baz_proxy(mut ctx_limit: u32) -> () { - baz((&mut ctx_limit)) - } "); }