feat(fuzz): Add experimental comptime code generator

compiler/noirc_frontend/src/monomorphization/ast.rs

@@ -458,7 +458,11 @@ impl std::fmt::Display for Program {
 
 impl std::fmt::Display for Function {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        super::printer::AstPrinter::default().print_function(self, None, f)
+        super::printer::AstPrinter::default().print_function(
+            self,
+            f,
+            super::printer::FunctionPrintOptions::default(),
+        )
     }
 }
 

compiler/noirc_frontend/src/monomorphization/printer.rs

@@ -8,6 +8,16 @@ use super::ast::{
 use iter_extended::vecmap;
 use std::fmt::{Display, Formatter};
 
+#[derive(Default)]
+pub struct FunctionPrintOptions {
+    pub return_visibility: Option<Visibility>,
+    /// Wraps function body in a `comptime` block. Used to make
+    /// comptime function callers in fuzzing.
+    pub comptime_wrap_body: bool,
+    /// Marks function as comptime. Used in fuzzing.
+    pub comptime: bool,
+}
+
 #[derive(Debug)]
 pub struct AstPrinter {
     indent_level: u32,
@@ -43,8 +53,10 @@ impl AstPrinter {
             self.print_global(id, global, f)?;
         }
         for function in &program.functions {
-            let vis = (function.id == Program::main_id()).then_some(program.return_visibility);
-            self.print_function(function, vis, f)?;
+            let return_visibility =
+                (function.id == Program::main_id()).then_some(program.return_visibility);
+            let fpo = FunctionPrintOptions { return_visibility, ..Default::default() };
+            self.print_function(function, f, fpo)?;
         }
         Ok(())
     }
@@ -64,15 +76,16 @@ impl AstPrinter {
     pub fn print_function(
         &mut self,
         function: &Function,
-        return_visibility: Option<Visibility>,
         f: &mut Formatter,
+        options: FunctionPrintOptions,
     ) -> std::fmt::Result {
         let params = vecmap(&function.parameters, |(id, mutable, name, typ)| {
             format!("{}{}: {}", if *mutable { "mut " } else { "" }, self.fmt_local(name, *id), typ)
         })
         .join(", ");
 
-        let vis = return_visibility
+        let vis = options
+            .return_visibility
             .map(|vis| match vis {
                 Visibility::Private => "".to_string(),
                 Visibility::Public => "pub ".to_string(),
@@ -82,14 +95,25 @@ impl AstPrinter {
             .unwrap_or_default();
 
         let unconstrained = if function.unconstrained { "unconstrained " } else { "" };
+        let comptime = if options.comptime { "comptime " } else { "" };
         let name = self.fmt_func(&function.name, function.id);
         let return_type = &function.return_type;
 
-        write!(f, "{unconstrained}fn {name}({params}) -> {vis}{return_type} {{",)?;
+        write!(f, "{comptime}{unconstrained}fn {name}({params}) -> {vis}{return_type} {{",)?;
         self.in_unconstrained = function.unconstrained;
+        if options.comptime_wrap_body {
+            self.indent_level += 1;
+            self.next_line(f)?;
+            write!(f, "comptime {{")?;
+        }
         self.indent_level += 1;
         self.print_expr_expect_block(&function.body, f)?;
         self.indent_level -= 1;
+        if options.comptime_wrap_body {
+            self.next_line(f)?;
+            self.indent_level -= 1;
+            write!(f, "}}")?;
+        }
         self.in_unconstrained = false;
         self.next_line(f)?;
         writeln!(f, "}}")?;

tooling/ast_fuzzer/examples/sample_comptime.rs

@@ -0,0 +1,24 @@
+//! Print a random comptime AST
+//!
+//! ```shell
+//! cargo run -p noir_ast_fuzzer --example sample_comptime
+//! ```
+use arbitrary::Unstructured;
+use noir_ast_fuzzer::{Config, DisplayAstAsNoirComptime, arb_program_comptime};
+use rand::RngCore;
+
+fn main() {
+    let data = {
+        let mut rng = rand::thread_rng();
+        let mut data = [0u8; 1024 * 1024];
+        rng.fill_bytes(&mut data);
+        data
+    };
+    let mut u = Unstructured::new(&data);
+
+    let mut config = Config::default();
+    config.max_globals = 0;
+
+    let program = arb_program_comptime(&mut u, config).expect("arb_program");
+    println!("{}", DisplayAstAsNoirComptime(&program));
+}

tooling/ast_fuzzer/src/lib.rs

@@ -7,13 +7,15 @@ pub use abi::program_abi;
 pub use input::arb_inputs;
 use program::freq::Freqs;
 pub use program::visitor::{visit_expr, visit_expr_mut};
-pub use program::{DisplayAstAsNoir, arb_program};
+pub use program::{DisplayAstAsNoir, DisplayAstAsNoirComptime, arb_program, arb_program_comptime};
 
 /// AST generation configuration.
 #[derive(Debug, Clone)]
 pub struct Config {
     /// Maximum number of global definitions.
     pub max_globals: usize,
+    /// Minimum number of functions (other than main) to generate.
+    pub min_functions: usize,
     /// Maximum number of functions (other than main) to generate.
     pub max_functions: usize,
     /// Maximum number of arguments a function can have.
@@ -40,6 +42,8 @@ pub struct Config {
     pub stmt_freqs_acir: Freqs,
     /// Frequency of statements in Brillig functions.
     pub stmt_freqs_brillig: Freqs,
+    /// Whether to force all functions to be unconstrained.
+    pub force_brillig: bool,
 }
 
 impl Default for Config {
@@ -75,6 +79,7 @@ impl Default for Config {
         ]);
         Self {
             max_globals: 3,
+            min_functions: 0,
             max_functions: 5,
             max_function_args: 3,
             max_function_size: 25,
@@ -88,6 +93,7 @@ impl Default for Config {
             expr_freqs,
             stmt_freqs_acir,
             stmt_freqs_brillig,
+            force_brillig: false,
         }
     }
 }

tooling/ast_fuzzer/src/program/func.rs

@@ -184,7 +184,7 @@
     }
 
     /// Generate the function body.
-    pub fn gen_body(mut self, u: &mut Unstructured) -> arbitrary::Result<Expression> {
+    pub fn gen_body(&mut self, u: &mut Unstructured) -> arbitrary::Result<Expression> {
         // If we don't limit the budget according to the available data,
         // it gives us a lot of `false` and 0 and we end up with deep `!(!false)` if expressions.
         self.budget = self.budget.min(u.len());
@@ -196,6 +196,16 @@
         Ok(body)
     }
 
+    /// Generate the function body, wrapping a function call with literal arguments.
+    /// This is used to test comptime functions, which can only take those.
+    pub fn gen_body_with_lit_call(
+        &mut self,
+        u: &mut Unstructured,
+        callee_id: FuncId,
+    ) -> arbitrary::Result<Expression> {
+        self.gen_lit_call(u, callee_id)
+    }
+
     /// Get the function declaration.
     fn decl(&self) -> &FunctionDeclaration {
         self.ctx.function_decl(self.id)
@@ -924,6 +934,43 @@
         self.gen_expr_from_source(u, call_expr, &callee.return_type, typ, self.max_depth())
     }
 
+    /// Generate a call to a specific function, with arbitrary literals
+    /// for arguments (useful for generating comptime wrapper calls)
+    fn gen_lit_call(
+        &mut self,
+        u: &mut Unstructured,
+        callee_id: FuncId,
+    ) -> arbitrary::Result<Expression> {
+        let callee = self.ctx.function_decl(callee_id).clone();
+        let param_types = callee.params.iter().map(|p| p.3.clone()).collect::<Vec<_>>();
+
+        let mut args = Vec::new();
+        for typ in &param_types {
+            args.push(expr::gen_literal(u, typ)?);
+        }
+
+        let call_expr = Expression::Call(Call {
+            func: Box::new(Expression::Ident(Ident {
+                location: None,
+                definition: Definition::Function(callee_id),
+                mutable: false,
+                name: callee.name.clone(),
+                typ: Type::Function(
+                    param_types,
+                    Box::new(callee.return_type.clone()),
+                    Box::new(Type::Unit),
+                    callee.unconstrained,
+                ),
+                id: self.next_ident_id(),
+            })),
+            arguments: args,
+            return_type: callee.return_type,
+            location: Location::dummy(),
+        });
+
+        Ok(call_expr)
+    }
+
     /// Generate a `loop` loop.
     fn gen_loop(&mut self, u: &mut Unstructured) -> arbitrary::Result<Expression> {
         // Declare break index variable visible in the loop body. Do not include it
@@ -1089,9 +1136,9 @@
 #[test]
 fn test_while() {
     let mut u = Unstructured::new(&[0u8; 1]);
     let mut ctx = Context::default();
     ctx.config.max_loop_size = 10;
     ctx.config.vary_loop_size = false;
     ctx.gen_main_decl(&mut u);
     let mut fctx = FunctionContext::new(&mut ctx, FuncId(0));
     fctx.budget = 2;

tooling/ast_fuzzer/src/program/mod.rs

@@ -9,7 +9,7 @@
     ast::IntegerBitSize,
     monomorphization::{
         ast::{Expression, FuncId, Function, GlobalId, InlineType, LocalId, Program, Type},
-        printer::AstPrinter,
+        printer::{AstPrinter, FunctionPrintOptions},
     },
     shared::{Signedness, Visibility},
 };
@@ -38,6 +38,43 @@
     Ok(program)
 }
 
+/// Generate an arbitrary monomorphized AST to be reversed into a valid comptime
+/// Noir, with a single comptime function called from main with literal arguments.
+pub fn arb_program_comptime(u: &mut Unstructured, config: Config) -> arbitrary::Result<Program> {
+    let mut config = config.clone();
+    // Comptime should use Brillig feature set
+    config.force_brillig = true;
+    // Generate only main function
+    config.min_functions = 0;
+    config.max_functions = 0;
+
+    //
+    let mut ctx = Context::new(config);
+
+    let decl_inner = ctx.gen_function_decl(u, 1)?;
+    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(),
+        params: vec![],
+        return_type: decl_inner.return_type.clone(),
+        param_visibilities: vec![],
+        return_visibility: Visibility::Public,
+        inline_type: InlineType::default(),
+        unconstrained: false,
+    };
+
+    ctx.set_function_decl(FuncId(0), decl_main);
+    ctx.gen_function_with_body(u, FuncId(0), |u, fctx| fctx.gen_body_with_lit_call(u, FuncId(1)))?;
+    ctx.rewrite_functions(u)?;
+
+    let program = ctx.finalize();
+    Ok(program)
+}
+
 /// ID of variables in scope.
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 enum VariableId {
@@ -80,6 +117,11 @@
         self.function_declarations.get(&id).expect("function should exist")
     }
 
+    /// Get a function declaration.
+    fn set_function_decl(&mut self, id: FuncId, decl: FunctionDeclaration) {
+        self.function_declarations.insert(id, decl);
+    }
+
     /// Get the main function declaration.
     fn main_decl(&self) -> &FunctionDeclaration {
         self.function_declarations.get(&Program::main_id()).expect("main should exist")
@@ -109,9 +151,11 @@
         Ok((name, typ, val))
     }
 
-    /// Generate random function names and signatures.
+    /// Generate random function names and signatures, optionally specifying their exact number.
     fn gen_function_decls(&mut self, u: &mut Unstructured) -> arbitrary::Result<()> {
-        let num_non_main_fns = u.int_in_range(0..=self.config.max_functions)?;
+        let num_non_main_fns =
+            u.int_in_range(self.config.min_functions..=self.config.max_functions)?;
+
         for i in 0..(1 + num_non_main_fns) {
             let d = self.gen_function_decl(u, i)?;
             self.function_declarations.insert(FuncId(i as u32), d);
@@ -172,7 +216,7 @@
             } else {
                 *u.choose(&[InlineType::Inline, InlineType::InlineAlways])?
             },
-            unconstrained: bool::arbitrary(u)?,
+            unconstrained: if self.config.force_brillig { true } else { bool::arbitrary(u)? },
         };
 
         Ok(decl)
@@ -189,23 +233,40 @@
     fn gen_functions(&mut self, u: &mut Unstructured) -> arbitrary::Result<()> {
         let ids = self.function_declarations.keys().cloned().collect::<Vec<_>>();
         for id in ids {
-            let body = FunctionContext::new(self, id).gen_body(u)?;
-            let decl = self.function_decl(id);
-            let func = Function {
-                id,
-                name: decl.name.clone(),
-                parameters: decl.params.clone(),
-                body,
-                return_type: decl.return_type.clone(),
-                unconstrained: decl.unconstrained,
-                inline_type: decl.inline_type,
-                func_sig: decl.signature(),
-            };
-            self.functions.insert(id, func);
+            self.gen_function(u, id)?;
         }
         Ok(())
     }
 
+    /// Generate random function body.
+    fn gen_function(&mut self, u: &mut Unstructured, id: FuncId) -> arbitrary::Result<()> {
+        self.gen_function_with_body(u, id, |u, fctx| fctx.gen_body(u))
+    }
+
+    /// Generate function with a specified body generator.
+    fn gen_function_with_body(
+        &mut self,
+        u: &mut Unstructured,
+        id: FuncId,
+        f: impl Fn(&mut Unstructured, &mut FunctionContext) -> arbitrary::Result<Expression>,
+    ) -> arbitrary::Result<()> {
+        let mut fctx = FunctionContext::new(self, id);
+        let body = f(u, &mut fctx)?;
+        let decl = self.function_decl(id);
+        let func = Function {
+            id,
+            name: decl.name.clone(),
+            parameters: decl.params.clone(),
+            body,
+            return_type: decl.return_type.clone(),
+            unconstrained: decl.unconstrained,
+            inline_type: decl.inline_type,
+            func_sig: decl.signature(),
+        };
+        self.functions.insert(id, func);
+        Ok(())
+    }
+
     /// As a post-processing step, identify recursive functions and add a call depth parameter to them.
     fn rewrite_functions(&mut self, u: &mut Unstructured) -> arbitrary::Result<()> {
         rewrite::add_recursion_limit(self, u)
@@ -335,3 +396,26 @@
         printer.print_program(self.0, f)
     }
 }
+
+/// Wrapper around `Program` that prints its AST as close to
+/// Noir syntax as we can get, making it `comptime`. The AST must
+/// be specifically prepared to include a main function consisting
+/// of a `comptime` wrapped call to a `comptime` (or `unconstrained`)
+/// marked function.
+pub struct DisplayAstAsNoirComptime<'a>(pub &'a Program);
+
+impl std::fmt::Display for DisplayAstAsNoirComptime<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut printer = AstPrinter::default();
+        printer.show_id = false;
+        for function in &self.0.functions {
+            let mut fpo = FunctionPrintOptions::default();
+            if function.id == Program::main_id() {
+                fpo.comptime_wrap_body = true;
+                fpo.return_visibility = Some(Visibility::Public);
+            } 
+            printer.print_function(function, f, fpo)?;
+        }
+        Ok(())
+    }
+}