feat: add while generator to AST fuzzer

tooling/ast_fuzzer/src/lib.rs

@@ -50,10 +50,13 @@ impl Default for Config {
         ]);
         let stmt_freqs = Freqs::new(&[
             ("drop", 5),
+            ("break", 5),
+            ("continue", 5),
             ("assign", 30),
             ("if", 10),
             ("for", 10),
             ("loop", 10),
+            ("while", 10),
             ("let", 20),
             ("call", 5),
         ]);

tooling/ast_fuzzer/src/program/func.rs

@@ -11,7 +11,8 @@
     hir_def::{self, expr::HirIdent, stmt::HirPattern},
     monomorphization::ast::{
         ArrayLiteral, Assign, BinaryOp, Call, Definition, Expression, For, FuncId, Function,
-        GlobalId, Ident, Index, InlineType, LValue, Literal, LocalId, Parameters, Program, Type,
+        GlobalId, Ident, Index, InlineType, LValue, Let, Literal, LocalId, Parameters, Program,
+        Type, While,
     },
     node_interner::DefinitionId,
     shared::{Signedness, Visibility},
@@ -118,6 +119,9 @@
     /// initially consisting of the function parameters, then extended
     /// by locally defined variables. Block scopes add and remove layers.
     locals: ScopeStack<LocalId>,
+    /// Indicator of being in a loop (and hence able to generate
+    /// break and continue statements)
+    in_loop: bool,
     /// All the functions callable from this one, with the types we can
     /// produce from their return value.
     call_targets: BTreeMap<FuncId, HashSet<Type>>,
@@ -161,7 +165,7 @@
             })
             .collect();
 
-        Self { ctx, id, next_local_id, budget, globals, locals, call_targets }
+        Self { ctx, id, next_local_id, budget, globals, locals, in_loop: false, call_targets }
     }
 
     /// Generate the function body.
@@ -572,7 +576,6 @@
     /// for example loops, variable declarations, etc.
     fn gen_stmt(&mut self, u: &mut Unstructured) -> arbitrary::Result<Expression> {
         let mut freq = Freq::new(u, &self.ctx.config.stmt_freqs)?;
-        // TODO(#7928): While
         // TODO(#7926): Match
         // TODO(#7931): print
         // TODO(#7932): Constrain
@@ -598,6 +601,22 @@
             return self.gen_for(u);
         }
 
+        if freq.enabled_when("loop", self.budget > 1 && self.unconstrained()) {
+            return self.gen_loop(u);
+        }
+
+        if freq.enabled_when("while", self.budget > 1 && self.unconstrained()) {
+            return self.gen_while(u);
+        }
+
+        if freq.enabled_when("break", self.in_loop && self.unconstrained()) {
+            return Ok(Expression::Break);
+        }
+
+        if freq.enabled_when("continue", self.in_loop && self.unconstrained()) {
+            return Ok(Expression::Continue);
+        }
+
         if freq.enabled_when("call", self.budget > 0) {
             if let Some(e) = self.gen_call(u, &Type::Unit, self.max_depth())? {
                 return Ok(e);
@@ -610,16 +629,25 @@
             }
         }
 
-        self.gen_let(u)
+        self.gen_let(u, None)
     }
 
-    /// Generate a `Let` statement.
-    fn gen_let(&mut self, u: &mut Unstructured) -> arbitrary::Result<Expression> {
+    /// Generate a `Let` statement, optionally requesting mutability.
+    fn gen_let(
+        &mut self,
+        u: &mut Unstructured,
+        mutable: Option<bool>,
+    ) -> arbitrary::Result<Expression> {
         // Generate a type or choose an existing one.
         let max_depth = self.max_depth();
         let typ = self.ctx.gen_type(u, max_depth, false)?;
         let id = self.next_local_id();
-        let mutable = bool::arbitrary(u)?;
+
+        let mutable = match mutable {
+            Some(m) => m,
+            None => bool::arbitrary(u)?,
+        };
+
         let name = make_name(id.0 as usize, false);
         let expr = self.gen_expr(u, &typ, max_depth, Flags::TOP)?;
 
@@ -764,7 +792,9 @@
         // Decrease budget so we don't nest for loops endlessly.
         self.decrease_budget(1);
 
+        let was_in_loop = std::mem::replace(&mut self.in_loop, true);
         let block = self.gen_block(u, &Type::Unit)?;
+        self.in_loop = was_in_loop;
 
         let expr = Expression::For(For {
             index_variable: idx_id,
@@ -837,10 +867,6 @@
 
     /// Generate a `loop` loop.
     fn gen_loop(&mut self, u: &mut Unstructured) -> arbitrary::Result<Expression> {
-        // For now, the only way to break out of the loop is the guard
-        // condition implemented below, as Break or Continue currently don't figure in
-        // generated expressions. This should be amended in TODO(#7928): While
-
         // Declare break index variable visible in the loop body. Do not include it
         // in the locals the generator would be able to manipulate, as it could
         // lead to the loop becoming infinite.
@@ -858,7 +884,9 @@
         let let_idx = expr::let_var(idx_id, true, idx_name, expr::u32_literal(0));
 
         // Get the randomized loop body
+        let was_in_loop = std::mem::replace(&mut self.in_loop, true);
         let mut loop_body = self.gen_block(u, &Type::Unit)?;
+        self.in_loop = was_in_loop;
 
         // Increment the index in the beginning of the body.
         expr::prepend(
@@ -883,6 +911,66 @@
 
         Ok(Expression::Block(vec![let_idx, Expression::Loop(Box::new(loop_body))]))
     }
+
+    /// Generate a `while` loop.
+    fn gen_while(&mut self, u: &mut Unstructured) -> arbitrary::Result<Expression> {
+        // Declare break index variable visible in the loop body. Do not include it
+        // in the locals the generator would be able to manipulate, as it could
+        // lead to the loop becoming infinite.
+        let idx_type = types::U32;
+        let idx_id = self.next_local_id();
+        let idx_name = format!("idx_{}", make_name(idx_id.0 as usize, false));
+        let idx_ident =
+            expr::ident_inner(VariableId::Local(idx_id), true, idx_name.clone(), idx_type);
+        let idx_expr = Expression::Ident(idx_ident.clone());
+
+        // Decrease budget so we don't nest endlessly.
+        self.decrease_budget(1);
+
+        // Start building the loop harness, initialize index to 0
+        let mut stmts = vec![Expression::Let(Let {
+            id: idx_id,
+            mutable: true,
+            name: idx_name,
+            expression: Box::new(expr::u32_literal(0)),
+        })];
+
+        // Get the randomized loop body
+        let was_in_loop = std::mem::replace(&mut self.in_loop, true);
+        let mut loop_body = self.gen_block(u, &Type::Unit)?;
+        self.in_loop = was_in_loop;
+
+        // Increment the index in the beginning of the body.
+        expr::prepend(
+            &mut loop_body,
+            expr::assign(
+                idx_ident,
+                expr::binary(idx_expr.clone(), BinaryOp::Add, expr::u32_literal(1)),
+            ),
+        );
+
+        // Put everything into if/else
+        let inner_block = Expression::Block(vec![expr::if_else(
+            expr::binary(
+                idx_expr,
+                BinaryOp::Equal,
+                expr::u32_literal(self.ctx.config.max_loop_size as u32),
+            ),
+            Expression::Break,
+            loop_body,
+            Type::Unit,
+        )]);
+
+        // Generate the `while` condition with depth 1
+        let condition = self.gen_expr(u, &Type::Bool, 1, Flags::CONDITION)?;
+
+        stmts.push(Expression::While(While {
+            condition: Box::new(condition),
+            body: Box::new(inner_block),
+        }));
+
+        Ok(Expression::Block(stmts))
+    }
 }
 
 /// Find the next local ID we can use to add variables to a [Function] during mutations.
@@ -908,9 +996,9 @@
     let mut ctx = Context::default();
     ctx.config.max_loop_size = 10;
     ctx.add_main_decl(&mut u);
     let mut fctx = FunctionContext::new(&mut ctx, FuncId(0));
     fctx.budget = 2;
     let loop_code = format!("{}", fctx.gen_loop(&mut u).unwrap()).replace(" ", "");
 
     println!("{loop_code}");
     assert!(
@@ -926,3 +1014,28 @@
         )
     );
 }
+
+#[test]
+fn test_while() {
+    let mut u = Unstructured::new(&[0u8; 1]);
+    let mut ctx = Context::default();
+    ctx.config.max_loop_size = 10;
+    ctx.add_main_decl(&mut u);
+    let mut fctx = FunctionContext::new(&mut ctx, FuncId(0));
+    fctx.budget = 2;
+    let while_code = format!("{}", fctx.gen_while(&mut u).unwrap()).replace(" ", "");
+
+    println!("{while_code}");
+    assert!(
+        while_code.starts_with(
+            &r#"{
+    let mut idx_a$l0 = 0;
+    while (!false) {
+        if (idx_a$l0 == 10) {
+            break
+        } else {
+            idx_a$l0 = (idx_a$l0 + 1)"#
+                .replace(" ", "")
+        )
+    );
+}