From fd4901ceba99fc09a41202ce71a35f5559828f8f Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Tue, 10 Jun 2025 11:03:35 +0100
Subject: [PATCH 01/10] Commutative metamorph

---
 .../fuzz/src/targets/orig_vs_morph.rs         | 61 ++++++++++++++++---
 tooling/ast_fuzzer/src/compare/compiled.rs    |  2 +
 tooling/ast_fuzzer/src/program/mod.rs         |  3 +-
 3 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index c7078d4a2f1..c4424988052 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -1,5 +1,5 @@
-//! Perform random equivalence mutations on the AST and check that the
-//! execution result does not change, a.k.a. metamorphic testing.
+//! Perform random metamorphic mutations on the AST and check that the
+//! execution result does not change.
 
 use crate::{compare_results_compiled, create_ssa_or_die, default_ssa_options};
 use arbitrary::{Arbitrary, Unstructured};
@@ -56,7 +56,7 @@ fn rewrite_function(
     rules: &[rules::Rule],
     max_rewrites: usize,
 ) {
-    // We can call `rewrite::next_local_and_ident_id`) and pass the results to the rewrite rules,
+    // We can call `rewrite::next_local_and_ident_id` and pass the results to the rewrite rules,
     // if they want to add new variables with new local IDs.
     let ctx = rules::Context { unconstrained: func.unconstrained, ..Default::default() };
 
@@ -200,7 +200,7 @@ mod rules {
     use noir_ast_fuzzer::expr;
     use noirc_frontend::{
         ast::BinaryOpKind,
-        monomorphization::ast::{Expression, Type},
+        monomorphization::ast::{Binary, Expression, Type},
     };
 
     #[derive(Clone, Debug, Default)]
@@ -251,7 +251,13 @@ mod rules {
 
     /// Construct all rules that we can apply on a program.
     pub fn all() -> Vec<Rule> {
-        vec![num_plus_minus_zero(), bool_or_self(), bool_xor_self(), bool_xor_rand()]
+        vec![
+            num_plus_minus_zero(),
+            bool_or_self(),
+            bool_xor_self(),
+            bool_xor_rand(),
+            commutative_arithmetic(),
+        ]
     }
 
     /// Transform any numeric value `x` into `x +/- 0`.
@@ -292,6 +298,18 @@ mod rules {
         )
     }
 
+    /// Check if an expression can have a side effect, in which case duplicating or reordering it could
+    /// change the behavior of the program.
+    fn has_side_effect(expr: &Expression) -> bool {
+        expr::exists(expr, |expr| {
+            matches!(
+                expr,
+                Expression::Call(_) // Functions can have side effects, maybe mutating some reference, printing
+                | Expression::Assign(_) // Assignment to a mutable variable could double up effects
+            )
+        })
+    }
+
     /// Common match condition for boolean rules.
     fn bool_rule_matches(ctx: &Context, expr: &Expression) -> bool {
         // If we rewrite `&mut x` into `&mut (x | x)` we will alter the semantics.
@@ -307,13 +325,12 @@ mod rules {
         // unless the expression can have a side effect, which we don't want to duplicate.
         if let Some(typ) = expr.return_type() {
             matches!(typ.as_ref(), Type::Bool)
+                && !has_side_effect(expr)
                 && !expr::exists(expr, |expr| {
                     matches!(
                         expr,
-                        Expression::Call(_) // Functions can have side effects, maybe mutating some reference
-                            | Expression::Assign(_) // Assignment to a mutable variable could double up effects
-                            | Expression::Let(_) // Creating a variable needs a new ID
-                            | Expression::Block(_) // Applying logical operations on blocks would look odd
+                        Expression::Let(_) // Creating a variable needs a new ID
+                    | Expression::Block(_) // Applying logical operations on blocks would look odd
                     )
                 })
         } else {
@@ -352,6 +369,32 @@ mod rules {
             Ok(())
         })
     }
+
+    /// Transform commutative arithmetic operations:
+    /// * `a + b` into `b + a`
+    /// * `a * b` into `b * a`
+    pub fn commutative_arithmetic() -> Rule {
+        Rule::new(
+            |_ctx, expr| {
+                matches!(
+                    expr,
+                    Expression::Binary(Binary {
+                        operator: BinaryOpKind::Add | BinaryOpKind::Multiply,
+                        ..
+                    })
+                ) && !has_side_effect(expr)
+            },
+            |_u, expr| {
+                let Expression::Binary(binary) = expr else {
+                    unreachable!("the rule only matches Binary expressions");
+                };
+
+                std::mem::swap(&mut binary.lhs, &mut binary.rhs);
+
+                Ok(())
+            },
+        )
+    }
 }
 
 #[cfg(test)]
diff --git a/tooling/ast_fuzzer/src/compare/compiled.rs b/tooling/ast_fuzzer/src/compare/compiled.rs
index 9b629123e88..e9043d74ebe 100644
--- a/tooling/ast_fuzzer/src/compare/compiled.rs
+++ b/tooling/ast_fuzzer/src/compare/compiled.rs
@@ -199,7 +199,9 @@ impl CompareMorph {
         let (program2, options) = f(u, program1.clone())?;
         let abi = program_abi(&program1);
 
+        println!("compiling the original");
         let ssa1 = g(program1.clone(), &options);
+        println!("compiling the morph");
         let ssa2 = g(program2.clone(), &options);
 
         let input_program = &ssa1.program;
diff --git a/tooling/ast_fuzzer/src/program/mod.rs b/tooling/ast_fuzzer/src/program/mod.rs
index 92194bca0c0..08938c9ddde 100644
--- a/tooling/ast_fuzzer/src/program/mod.rs
+++ b/tooling/ast_fuzzer/src/program/mod.rs
@@ -533,7 +533,8 @@ impl std::fmt::Display for DisplayAstAsNoir<'_> {
         printer.show_type_in_let = true;
         // Most of the time it doesn't affect testing, except the comptime tests where
         // we parse back the code. For that we use `DisplayAstAsNoirComptime`.
-        // It is quite noisy, so keep it off by default.
+        // Using it also inserts an extra `cast` in the SSA,
+        // which at the moment for example defeats loop unrolling.
         printer.show_type_of_int_literal = false;
         printer.print_program(self.0, f)
     }

From ec10580c78870616d99ba062456dd503c983b3ba Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Tue, 10 Jun 2025 16:18:40 +0100
Subject: [PATCH 02/10] WIP

---
 .../fuzz/src/targets/orig_vs_morph.rs         | 96 +++++++++++--------
 tooling/ast_fuzzer/src/program/expr.rs        | 65 +++++++++++++
 2 files changed, 121 insertions(+), 40 deletions(-)

diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index c4424988052..643cbd4f7b4 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -257,6 +257,7 @@ mod rules {
             bool_xor_self(),
             bool_xor_rand(),
             commutative_arithmetic(),
+            inevitable(),
         ]
     }
 
@@ -298,46 +299,6 @@ mod rules {
         )
     }
 
-    /// Check if an expression can have a side effect, in which case duplicating or reordering it could
-    /// change the behavior of the program.
-    fn has_side_effect(expr: &Expression) -> bool {
-        expr::exists(expr, |expr| {
-            matches!(
-                expr,
-                Expression::Call(_) // Functions can have side effects, maybe mutating some reference, printing
-                | Expression::Assign(_) // Assignment to a mutable variable could double up effects
-            )
-        })
-    }
-
-    /// Common match condition for boolean rules.
-    fn bool_rule_matches(ctx: &Context, expr: &Expression) -> bool {
-        // If we rewrite `&mut x` into `&mut (x | x)` we will alter the semantics.
-        if ctx.is_in_ref_mut {
-            return false;
-        }
-        // We don't want to mess with the arguments of a `println`, because the printer assumes they are bool literals.
-        // Similarly a `constrain` call is expected to have a single boolean expression.
-        if ctx.is_in_special_call {
-            return false;
-        }
-        // We can apply boolean rule on anything that returns a bool,
-        // unless the expression can have a side effect, which we don't want to duplicate.
-        if let Some(typ) = expr.return_type() {
-            matches!(typ.as_ref(), Type::Bool)
-                && !has_side_effect(expr)
-                && !expr::exists(expr, |expr| {
-                    matches!(
-                        expr,
-                        Expression::Let(_) // Creating a variable needs a new ID
-                    | Expression::Block(_) // Applying logical operations on blocks would look odd
-                    )
-                })
-        } else {
-            false
-        }
-    }
-
     /// Transform boolean value `x` into `x | x`.
     pub fn bool_or_self() -> Rule {
         Rule::new(bool_rule_matches, |_u, expr| {
@@ -395,6 +356,61 @@ mod rules {
             },
         )
     }
+
+    /// Transform an expression into an if-then-else with the expression
+    /// repeated in the _then_ and _else_ branch:
+    /// * `x` into `if c { x } else { x }`
+    pub fn inevitable() -> Rule {
+        Rule::new(
+            |ctx, _expr| !ctx.is_in_special_call && !ctx.is_in_ref_mut,
+            |u, expr| {
+                let typ = expr.return_type().map(|typ| typ.into_owned()).unwrap_or(Type::Unit);
+                // *expr = expr::if_else(gen_condition(u)?, expr.clone(), expr.clone(), typ);
+
+                Ok(())
+            },
+        )
+    }
+
+    /// Check if an expression can have a side effect, in which case duplicating or reordering it could
+    /// change the behavior of the program.
+    fn has_side_effect(expr: &Expression) -> bool {
+        expr::exists(expr, |expr| {
+            matches!(
+                expr,
+                Expression::Call(_) // Functions can have side effects, maybe mutating some reference, printing
+                | Expression::Assign(_) // Assignment to a mutable variable could double up effects
+            )
+        })
+    }
+
+    /// Common match condition for boolean rules.
+    fn bool_rule_matches(ctx: &Context, expr: &Expression) -> bool {
+        // If we rewrite `&mut x` into `&mut (x | x)` we will alter the semantics.
+        if ctx.is_in_ref_mut {
+            return false;
+        }
+        // We don't want to mess with the arguments of a `println`, because the printer assumes they are bool literals.
+        // Similarly a `constrain` call is expected to have a single boolean expression.
+        if ctx.is_in_special_call {
+            return false;
+        }
+        // We can apply boolean rule on anything that returns a bool,
+        // unless the expression can have a side effect, which we don't want to duplicate.
+        if let Some(typ) = expr.return_type() {
+            matches!(typ.as_ref(), Type::Bool)
+                && !has_side_effect(expr)
+                && !expr::exists(expr, |expr| {
+                    matches!(
+                        expr,
+                        Expression::Let(_) // Creating a variable needs a new ID
+                    | Expression::Block(_) // Applying logical operations on blocks would look odd
+                    )
+                })
+        } else {
+            false
+        }
+    }
 }
 
 #[cfg(test)]
diff --git a/tooling/ast_fuzzer/src/program/expr.rs b/tooling/ast_fuzzer/src/program/expr.rs
index 0bd57ebc8ca..319d4c50a4f 100644
--- a/tooling/ast_fuzzer/src/program/expr.rs
+++ b/tooling/ast_fuzzer/src/program/expr.rs
@@ -12,6 +12,7 @@ use noirc_frontend::{
     },
     signed_field::SignedField,
 };
+use strum::IntoEnumIterator;
 
 use super::{Name, VariableId, types, visitor::visit_expr};
 
@@ -99,6 +100,70 @@ pub fn gen_literal(u: &mut Unstructured, typ: &Type) -> arbitrary::Result<Expres
     Ok(expr)
 }
 
+// /// Generate an arbitrary pure (free of side effects) expression, returning a specific type.
+// pub fn gen_expr(
+//     u: &mut Unstructured,
+//     typ: &Type,
+//     max_depth: usize,
+// ) -> arbitrary::Result<Expression> {
+//     if max_depth > 0 {
+//         let idx = u.choose_index(3)?;
+//         if idx == 0 {
+//             if let Some(expr) = gen_unary(u, typ, max_depth)? {
+//                 return Ok(expr);
+//             }
+//         }
+//         if idx == 1 {
+//             if let Some(expr) = gen_binary(u, typ, max_depth)? {
+//                 return Ok(expr);
+//             }
+//         }
+//     }
+//     gen_literal(u, typ)
+// }
+
+// /// Generate an arbitrary unary expression, returning a specific type.
+// pub fn gen_unary(
+//     u: &mut Unstructured,
+//     typ: &Type,
+//     max_depth: usize,
+// ) -> arbitrary::Result<Option<Expression>> {
+//     if !types::can_unary_return(typ) {
+//         return Ok(None);
+//     }
+//     let mut make_unary = |op| {
+//         let expr = gen_expr(u, typ, max_depth - 1)?;
+//         Ok(Some(unary(op, expr, typ.clone())))
+//     };
+//     if types::is_numeric(typ) {
+//         make_unary(UnaryOp::Minus)
+//     } else if types::is_bool(typ) {
+//         make_unary(UnaryOp::Not)
+//     } else {
+//         Ok(None)
+//     }
+// }
+
+// /// Generate an arbitrary binary expression, returning a specific type.
+// pub fn gen_binary(
+//     u: &mut Unstructured,
+//     typ: &Type,
+//     max_depth: usize,
+// ) -> arbitrary::Result<Option<Expression>> {
+//     // Collect the operations can return the expected type.
+//     let ops =
+//         BinaryOp::iter().filter(|op| types::can_binary_op_return(op, typ)).collect::<Vec<_>>();
+
+//     // Ideally we checked that the target type can be returned, but just in case.
+//     if ops.is_empty() {
+//         return Ok(None);
+//     }
+
+//     // Choose a random operation.
+//     let op = u.choose_iter(ops)?;
+
+// }
+
 /// Generate a literals for loop ranges with signed/unsigned integers with bits 8, 16, 32 or 64 bits,
 /// in a way that start is not higher than the end, and the maximum difference between them is limited,
 /// so that we don't get huge unrolled loops.

From f5d4681515b4921b6a864d8de99ee498d9039965 Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Wed, 11 Jun 2025 14:49:57 +0100
Subject: [PATCH 03/10] Generate arbitrary literal conditions

---
 Cargo.lock                                    |   1 +
 tooling/ast_fuzzer/fuzz/Cargo.toml            |   1 +
 .../fuzz/src/targets/orig_vs_morph.rs         | 167 ++++++++++++++++--
 tooling/ast_fuzzer/src/compare/compiled.rs    |   5 +-
 tooling/ast_fuzzer/src/lib.rs                 |   2 +-
 tooling/ast_fuzzer/src/program/expr.rs        |  65 -------
 tooling/ast_fuzzer/src/program/mod.rs         |   2 +-
 tooling/ast_fuzzer/src/program/types.rs       |  38 ++--
 8 files changed, 178 insertions(+), 103 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1d1a60cc8d1..2cbe736bfde 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3444,6 +3444,7 @@ dependencies = [
  "noirc_abi",
  "noirc_evaluator",
  "noirc_frontend",
+ "strum",
 ]
 
 [[package]]
diff --git a/tooling/ast_fuzzer/fuzz/Cargo.toml b/tooling/ast_fuzzer/fuzz/Cargo.toml
index 01065a81a4b..f1e90919d33 100644
--- a/tooling/ast_fuzzer/fuzz/Cargo.toml
+++ b/tooling/ast_fuzzer/fuzz/Cargo.toml
@@ -15,6 +15,7 @@ cargo-fuzz = true
 arbitrary.workspace = true
 color-eyre.workspace = true
 libfuzzer-sys.workspace = true
+strum.workspace = true
 
 acir.workspace = true
 noirc_abi.workspace = true
diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index 643cbd4f7b4..2e043af114c 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -196,6 +196,7 @@ fn is_special_call(call: &Call) -> bool {
 
 /// Metamorphic transformation rules.
 mod rules {
+    use super::helpers::{gen_expr, has_side_effect};
     use arbitrary::{Arbitrary, Unstructured};
     use noir_ast_fuzzer::expr;
     use noirc_frontend::{
@@ -362,28 +363,26 @@ mod rules {
     /// * `x` into `if c { x } else { x }`
     pub fn inevitable() -> Rule {
         Rule::new(
-            |ctx, _expr| !ctx.is_in_special_call && !ctx.is_in_ref_mut,
+            |ctx, expr| {
+                !ctx.is_in_special_call
+                    && !ctx.is_in_ref_mut
+                    // If we're in ACIR then don't turn loop ranges into expressions.
+                    && (ctx.unconstrained || !ctx.is_in_range)
+                    // Avoid creating new variables because:
+                    // * `let x = 1;` transformed into `if true { let x = 1; } else { let x = 1; }` would leave `x` undefined.
+                    // * Creating variables in blocks would need to be assigned new IDs, and subsequent expressions would need to be rewritten,
+                    //   for which we currently lack the context. `for` introduces the index variable.
+                    && !expr::exists(expr, |expr| matches!(expr, Expression::Let(_) | Expression::For(_)))
+            },
             |u, expr| {
                 let typ = expr.return_type().map(|typ| typ.into_owned()).unwrap_or(Type::Unit);
-                // *expr = expr::if_else(gen_condition(u)?, expr.clone(), expr.clone(), typ);
-
+                let cond = gen_expr(u, &Type::Bool, 2)?;
+                *expr = expr::if_else(cond, expr.clone(), expr.clone(), typ);
                 Ok(())
             },
         )
     }
 
-    /// Check if an expression can have a side effect, in which case duplicating or reordering it could
-    /// change the behavior of the program.
-    fn has_side_effect(expr: &Expression) -> bool {
-        expr::exists(expr, |expr| {
-            matches!(
-                expr,
-                Expression::Call(_) // Functions can have side effects, maybe mutating some reference, printing
-                | Expression::Assign(_) // Assignment to a mutable variable could double up effects
-            )
-        })
-    }
-
     /// Common match condition for boolean rules.
     fn bool_rule_matches(ctx: &Context, expr: &Expression) -> bool {
         // If we rewrite `&mut x` into `&mut (x | x)` we will alter the semantics.
@@ -413,6 +412,144 @@ mod rules {
     }
 }
 
+mod helpers {
+    use std::sync::OnceLock;
+
+    use arbitrary::Unstructured;
+    use noir_ast_fuzzer::{expr, types};
+    use noirc_frontend::{
+        ast::{IntegerBitSize, UnaryOp},
+        monomorphization::ast::{BinaryOp, Expression, Type},
+        shared::Signedness,
+    };
+    use strum::IntoEnumIterator;
+
+    /// Check if an expression can have a side effect, in which case duplicating or reordering it could
+    /// change the behavior of the program.
+    pub(super) fn has_side_effect(expr: &Expression) -> bool {
+        expr::exists(expr, |expr| {
+            matches!(
+                expr,
+                Expression::Call(_) // Functions can have side effects, maybe mutating some reference, printing
+                | Expression::Assign(_) // Assignment to a mutable variable could double up effects
+            )
+        })
+    }
+
+    /// Generate an arbitrary pure (free of side effects) expression, returning a specific type.
+    pub(super) fn gen_expr(
+        u: &mut Unstructured,
+        typ: &Type,
+        max_depth: usize,
+    ) -> arbitrary::Result<Expression> {
+        if max_depth > 0 {
+            let idx = u.choose_index(3)?;
+            if idx == 0 {
+                if let Some(expr) = gen_unary(u, typ, max_depth)? {
+                    return Ok(expr);
+                }
+            }
+            if idx == 1 {
+                if let Some(expr) = gen_binary(u, typ, max_depth)? {
+                    return Ok(expr);
+                }
+            }
+        }
+        expr::gen_literal(u, typ)
+    }
+
+    /// Generate an arbitrary unary expression, returning a specific type.
+    pub(super) fn gen_unary(
+        u: &mut Unstructured,
+        typ: &Type,
+        max_depth: usize,
+    ) -> arbitrary::Result<Option<Expression>> {
+        if !types::can_unary_return(typ) {
+            return Ok(None);
+        }
+        let mut make_unary = |op| {
+            let expr = gen_expr(u, typ, max_depth.saturating_sub(1))?;
+            Ok(Some(expr::unary(op, expr, typ.clone())))
+        };
+        if types::is_numeric(typ) {
+            make_unary(UnaryOp::Minus)
+        } else if types::is_bool(typ) {
+            make_unary(UnaryOp::Not)
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Generate an arbitrary binary expression, returning a specific type.
+    ///
+    /// The operation should not have any side effects, ie. it must not fail.
+    pub(super) fn gen_binary(
+        u: &mut Unstructured,
+        typ: &Type,
+        max_depth: usize,
+    ) -> arbitrary::Result<Option<Expression>> {
+        // Collect the operations can return the expected type.
+        // Do not introduce new errors in randomly generated code that only exists in the morph.
+        let ops = BinaryOp::iter()
+            .filter(|op| {
+                types::can_binary_op_return(op, typ)
+                    && !types::can_binary_op_overflow(op)
+                    && !types::can_binary_op_err_by_zero(op)
+            })
+            .collect::<Vec<_>>();
+
+        // Ideally we checked that the target type can be returned, but just in case.
+        if ops.is_empty() {
+            return Ok(None);
+        }
+
+        // Choose a random operation.
+        let op = u.choose_iter(ops)?;
+
+        let type_options = TYPES.get_or_init(|| {
+            let mut types = vec![Type::Bool, Type::Field];
+            for sign in [Signedness::Signed, Signedness::Unsigned] {
+                for size in IntegerBitSize::iter() {
+                    if sign.is_signed() && size.bit_size() == 1 {
+                        continue;
+                    }
+                    types.push(Type::Integer(sign, size));
+                }
+            }
+            types
+        });
+
+        // Select input types that can produce the output we want.
+        let type_options = type_options
+            .iter()
+            .filter(|input| types::can_binary_op_return_from_input(&op, input, typ))
+            .collect::<Vec<_>>();
+
+        // Choose a type for the LHS and RHS.
+        let lhs_type = u.choose_iter(type_options)?;
+        let rhs_type = match op {
+            BinaryOp::ShiftLeft | BinaryOp::ShiftRight => &types::U8,
+            _ => lhs_type,
+        };
+
+        // Generate expressions for LHS and RHS.
+        let lhs_expr = gen_expr(u, lhs_type, max_depth.saturating_sub(1))?;
+        let rhs_expr = gen_expr(u, rhs_type, max_depth.saturating_sub(1))?;
+
+        let mut expr = expr::binary(lhs_expr, op, rhs_expr);
+
+        // If we have chosen e.g. u8 and need u32 we need to cast.
+        if !(lhs_type == typ || types::is_bool(typ) && op.is_comparator()) {
+            expr = expr::cast(expr, typ.clone());
+        }
+
+        Ok(Some(expr))
+    }
+
+    /// Types we can consider using in this context.
+    static TYPES: OnceLock<Vec<Type>> = OnceLock::new();
+}
+
 #[cfg(test)]
 mod tests {
     /// ```ignore
diff --git a/tooling/ast_fuzzer/src/compare/compiled.rs b/tooling/ast_fuzzer/src/compare/compiled.rs
index e9043d74ebe..89c1ff00719 100644
--- a/tooling/ast_fuzzer/src/compare/compiled.rs
+++ b/tooling/ast_fuzzer/src/compare/compiled.rs
@@ -79,6 +79,9 @@ impl Comparable for NargoError<FieldElement> {
             return false;
         };
 
+        // let msg1 = e1.user_defined_failure_message(&abi1.error_types);
+        // let msg2 = e2.user_defined_failure_message(&abi2.error_types);
+
         match (ee1, ee2) {
             (
                 AssertionFailed(ResolvedAssertionPayload::String(c), _, _),
@@ -199,9 +202,7 @@ impl CompareMorph {
         let (program2, options) = f(u, program1.clone())?;
         let abi = program_abi(&program1);
 
-        println!("compiling the original");
         let ssa1 = g(program1.clone(), &options);
-        println!("compiling the morph");
         let ssa2 = g(program2.clone(), &options);
 
         let input_program = &ssa1.program;
diff --git a/tooling/ast_fuzzer/src/lib.rs b/tooling/ast_fuzzer/src/lib.rs
index b10c2220c0c..8fc4e07bf48 100644
--- a/tooling/ast_fuzzer/src/lib.rs
+++ b/tooling/ast_fuzzer/src/lib.rs
@@ -13,7 +13,7 @@ pub use compare::input_values_to_ssa;
 pub use input::arb_inputs;
 use program::freq::Freqs;
 pub use program::{DisplayAstAsNoir, DisplayAstAsNoirComptime, arb_program, arb_program_comptime};
-pub use program::{expr, rewrite, visitor};
+pub use program::{expr, rewrite, types, visitor};
 
 /// AST generation configuration.
 #[derive(Debug, Clone)]
diff --git a/tooling/ast_fuzzer/src/program/expr.rs b/tooling/ast_fuzzer/src/program/expr.rs
index 319d4c50a4f..0bd57ebc8ca 100644
--- a/tooling/ast_fuzzer/src/program/expr.rs
+++ b/tooling/ast_fuzzer/src/program/expr.rs
@@ -12,7 +12,6 @@ use noirc_frontend::{
     },
     signed_field::SignedField,
 };
-use strum::IntoEnumIterator;
 
 use super::{Name, VariableId, types, visitor::visit_expr};
 
@@ -100,70 +99,6 @@ pub fn gen_literal(u: &mut Unstructured, typ: &Type) -> arbitrary::Result<Expres
     Ok(expr)
 }
 
-// /// Generate an arbitrary pure (free of side effects) expression, returning a specific type.
-// pub fn gen_expr(
-//     u: &mut Unstructured,
-//     typ: &Type,
-//     max_depth: usize,
-// ) -> arbitrary::Result<Expression> {
-//     if max_depth > 0 {
-//         let idx = u.choose_index(3)?;
-//         if idx == 0 {
-//             if let Some(expr) = gen_unary(u, typ, max_depth)? {
-//                 return Ok(expr);
-//             }
-//         }
-//         if idx == 1 {
-//             if let Some(expr) = gen_binary(u, typ, max_depth)? {
-//                 return Ok(expr);
-//             }
-//         }
-//     }
-//     gen_literal(u, typ)
-// }
-
-// /// Generate an arbitrary unary expression, returning a specific type.
-// pub fn gen_unary(
-//     u: &mut Unstructured,
-//     typ: &Type,
-//     max_depth: usize,
-// ) -> arbitrary::Result<Option<Expression>> {
-//     if !types::can_unary_return(typ) {
-//         return Ok(None);
-//     }
-//     let mut make_unary = |op| {
-//         let expr = gen_expr(u, typ, max_depth - 1)?;
-//         Ok(Some(unary(op, expr, typ.clone())))
-//     };
-//     if types::is_numeric(typ) {
-//         make_unary(UnaryOp::Minus)
-//     } else if types::is_bool(typ) {
-//         make_unary(UnaryOp::Not)
-//     } else {
-//         Ok(None)
-//     }
-// }
-
-// /// Generate an arbitrary binary expression, returning a specific type.
-// pub fn gen_binary(
-//     u: &mut Unstructured,
-//     typ: &Type,
-//     max_depth: usize,
-// ) -> arbitrary::Result<Option<Expression>> {
-//     // Collect the operations can return the expected type.
-//     let ops =
-//         BinaryOp::iter().filter(|op| types::can_binary_op_return(op, typ)).collect::<Vec<_>>();
-
-//     // Ideally we checked that the target type can be returned, but just in case.
-//     if ops.is_empty() {
-//         return Ok(None);
-//     }
-
-//     // Choose a random operation.
-//     let op = u.choose_iter(ops)?;
-
-// }
-
 /// Generate a literals for loop ranges with signed/unsigned integers with bits 8, 16, 32 or 64 bits,
 /// in a way that start is not higher than the end, and the maximum difference between them is limited,
 /// so that we don't get huge unrolled loops.
diff --git a/tooling/ast_fuzzer/src/program/mod.rs b/tooling/ast_fuzzer/src/program/mod.rs
index 08938c9ddde..aa0c2bbe3d1 100644
--- a/tooling/ast_fuzzer/src/program/mod.rs
+++ b/tooling/ast_fuzzer/src/program/mod.rs
@@ -21,7 +21,7 @@ pub(crate) mod freq;
 mod func;
 pub mod rewrite;
 mod scope;
-mod types;
+pub mod types;
 pub mod visitor;
 
 #[cfg(test)]
diff --git a/tooling/ast_fuzzer/src/program/types.rs b/tooling/ast_fuzzer/src/program/types.rs
index 89687df7c98..d2cc34a8077 100644
--- a/tooling/ast_fuzzer/src/program/types.rs
+++ b/tooling/ast_fuzzer/src/program/types.rs
@@ -10,13 +10,13 @@ use noirc_frontend::{
 };
 use strum::IntoEnumIterator as _;
 
-pub(crate) const U8: Type = Type::Integer(Signedness::Unsigned, IntegerBitSize::Eight);
-pub(crate) const U32: Type = Type::Integer(Signedness::Unsigned, IntegerBitSize::ThirtyTwo);
+pub const U8: Type = Type::Integer(Signedness::Unsigned, IntegerBitSize::Eight);
+pub const U32: Type = Type::Integer(Signedness::Unsigned, IntegerBitSize::ThirtyTwo);
 
 /// Calculate the depth of a type.
 ///
 /// Leaf types have a depth of 0.
-pub(crate) fn type_depth(typ: &Type) -> usize {
+pub fn type_depth(typ: &Type) -> usize {
     match typ {
         Type::Field | Type::Bool | Type::String(_) | Type::Unit | Type::Integer(_, _) => 0,
         Type::Array(_, typ) => 1 + type_depth(typ),
@@ -28,7 +28,7 @@ pub(crate) fn type_depth(typ: &Type) -> usize {
 /// We can only use globals that can be evaluated at comptime.
 /// Some types don't compile in Noir, so avoid those as we couldn't
 /// put any related failures into an integration test.
-pub(crate) fn can_be_global(typ: &Type) -> bool {
+pub fn can_be_global(typ: &Type) -> bool {
     !matches!(
         typ,
         Type::Integer(Signedness::Signed, IntegerBitSize::One | IntegerBitSize::HundredTwentyEight)
@@ -41,7 +41,7 @@ pub(crate) fn can_be_global(typ: &Type) -> bool {
 /// We decided we will avoid 0 length arrays in the main inputs and outputs, because we don't generate
 /// witnesses for them anyway, and they are tricky to handle consistently when they can be regular inputs
 /// as well as part of the databus. They are not expected in real programs as they don't do anything useful.
-pub(crate) fn can_be_main(typ: &Type) -> bool {
+pub fn can_be_main(typ: &Type) -> bool {
     match typ {
         Type::String(size) => *size > 0,
         Type::Array(size, typ) => *size > 0 && can_be_main(typ),
@@ -54,7 +54,7 @@ pub(crate) fn can_be_main(typ: &Type) -> bool {
 /// Collect all the sub-types produced by a type.
 ///
 /// It's like a _power set_ of the type.
-pub(crate) fn types_produced(typ: &Type) -> HashSet<Type> {
+pub fn types_produced(typ: &Type) -> HashSet<Type> {
     /// Recursively visit subtypes.
     fn visit(acc: &mut HashSet<Type>, typ: &Type) {
         if acc.contains(typ) {
@@ -131,7 +131,7 @@ pub(crate) fn types_produced(typ: &Type) -> HashSet<Type> {
     acc
 }
 
-pub(crate) fn to_hir_type(typ: &Type) -> hir_def::types::Type {
+pub fn to_hir_type(typ: &Type) -> hir_def::types::Type {
     use hir_def::types::{Kind as HirKind, Type as HirType};
 
     // Meet the expectations of `Type::evaluate_to_u32`.
@@ -169,27 +169,27 @@ pub(crate) fn to_hir_type(typ: &Type) -> hir_def::types::Type {
 }
 
 /// Check if the type is a number.
-pub(crate) fn is_numeric(typ: &Type) -> bool {
+pub fn is_numeric(typ: &Type) -> bool {
     matches!(typ, Type::Field | Type::Integer(_, _))
 }
 
 /// Check if a type is `Unit`.
-pub(crate) fn is_unit(typ: &Type) -> bool {
+pub fn is_unit(typ: &Type) -> bool {
     matches!(typ, Type::Unit)
 }
 
 /// Check if the type works with `UnaryOp::Not`
-pub(crate) fn is_bool(typ: &Type) -> bool {
+pub fn is_bool(typ: &Type) -> bool {
     matches!(typ, Type::Bool)
 }
 
 /// Check if the type is a reference wrapping another.
-pub(crate) fn is_reference(typ: &Type) -> bool {
+pub fn is_reference(typ: &Type) -> bool {
     matches!(typ, Type::Reference(_, _))
 }
 
 /// Check if the type can be used with a `println` statement.
-pub(crate) fn is_printable(typ: &Type) -> bool {
+pub fn is_printable(typ: &Type) -> bool {
     match typ {
         Type::Reference(_, _) => false,
         Type::Field | Type::Integer(_, _) | Type::String(_) | Type::Bool | Type::Unit => true,
@@ -201,7 +201,7 @@ pub(crate) fn is_printable(typ: &Type) -> bool {
 }
 
 /// Can the type be returned by some `UnaryOp`.
-pub(crate) fn can_unary_return(typ: &Type) -> bool {
+pub fn can_unary_return(typ: &Type) -> bool {
     match typ {
         Type::Field => true,
         Type::Bool => true,
@@ -221,12 +221,12 @@ pub(crate) fn can_unary_return(typ: &Type) -> bool {
 }
 
 /// Can the type be returned by some `BinaryOp`.
-pub(crate) fn can_binary_return(typ: &Type) -> bool {
+pub fn can_binary_return(typ: &Type) -> bool {
     BinaryOp::iter().any(|op| can_binary_op_return(&op, typ))
 }
 
 /// Check if a certain binary operation can return a target type as output.
-pub(crate) fn can_binary_op_return(op: &BinaryOp, typ: &Type) -> bool {
+pub fn can_binary_op_return(op: &BinaryOp, typ: &Type) -> bool {
     use BinaryOpKind::*;
     match typ {
         Type::Bool => op.is_comparator(),
@@ -242,7 +242,7 @@ pub(crate) fn can_binary_op_return(op: &BinaryOp, typ: &Type) -> bool {
 
 /// Can the binary operation result in an overflow.
 /// These are operations that commonly fail with random constants.
-pub(crate) fn can_binary_op_overflow(op: &BinaryOp) -> bool {
+pub fn can_binary_op_overflow(op: &BinaryOp) -> bool {
     use BinaryOpKind::*;
     matches!(op, Add | Subtract | Multiply | ShiftLeft)
 }
@@ -251,13 +251,13 @@ pub(crate) fn can_binary_op_overflow(op: &BinaryOp) -> bool {
 /// These operations can fail because of an unfortunate combination of
 /// expressions, such as `a / (a % a)` or `a % (a % a)`, but it's less
 /// like to occur than the overflows.
-pub(crate) fn can_binary_op_err_by_zero(op: &BinaryOp) -> bool {
+pub fn can_binary_op_err_by_zero(op: &BinaryOp) -> bool {
     use BinaryOpKind::*;
     matches!(op, Divide | Modulo)
 }
 
 /// Check if a certain binary operation can take a type as input and produce the target output.
-pub(crate) fn can_binary_op_return_from_input(op: &BinaryOp, input: &Type, output: &Type) -> bool {
+pub fn can_binary_op_return_from_input(op: &BinaryOp, input: &Type, output: &Type) -> bool {
     match (input, output) {
         (Type::Field, Type::Field) => op.is_valid_for_field_type() && !op.is_equality(),
         (Type::Field, Type::Bool) => op.is_equality(),
@@ -285,6 +285,6 @@ pub(crate) fn can_binary_op_return_from_input(op: &BinaryOp, input: &Type, outpu
 }
 
 /// Reference an expression into a target type
-pub(crate) fn ref_mut(typ: Type) -> Type {
+pub fn ref_mut(typ: Type) -> Type {
     Type::Reference(Box::new(typ), true)
 }

From 198d0491efcf658090c1813a7663fc720d0c49fa Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Wed, 11 Jun 2025 16:05:09 +0100
Subject: [PATCH 04/10] Make error types available in comparisons

---
 tooling/ast_fuzzer/src/compare/compiled.rs | 94 +++++++++++++++++++---
 tooling/ast_fuzzer/src/compare/comptime.rs | 10 ++-
 2 files changed, 90 insertions(+), 14 deletions(-)

diff --git a/tooling/ast_fuzzer/src/compare/compiled.rs b/tooling/ast_fuzzer/src/compare/compiled.rs
index 89c1ff00719..1d8020b58d2 100644
--- a/tooling/ast_fuzzer/src/compare/compiled.rs
+++ b/tooling/ast_fuzzer/src/compare/compiled.rs
@@ -1,4 +1,6 @@
 //! Compare an arbitrary AST compiled into bytecode and executed with the VM.
+use std::collections::BTreeMap;
+
 use acir::{FieldElement, native_types::WitnessStack};
 use acvm::pwg::{OpcodeResolutionError, ResolvedAssertionPayload};
 use arbitrary::Unstructured;
@@ -6,7 +8,7 @@ use bn254_blackbox_solver::Bn254BlackBoxSolver;
 use color_eyre::eyre::{self, WrapErr};
 use nargo::{NargoError, errors::ExecutionError, foreign_calls::DefaultForeignCallBuilder};
 use noirc_abi::{Abi, InputMap, input_parser::InputValue};
-use noirc_evaluator::ssa::SsaProgramArtifact;
+use noirc_evaluator::{ErrorType, ssa::SsaProgramArtifact};
 use noirc_frontend::monomorphization::ast::Program;
 
 use crate::{Config, arb_inputs, arb_program, program_abi};
@@ -30,15 +32,56 @@ impl From<(SsaProgramArtifact, CompareOptions)> for CompareArtifact {
     }
 }
 
+/// These are the error types in the `SsaProgramArtifact`, which are not the same as the ones in the ABI,
+/// but they can provide extra information when comparing errors.
+type SsaErrorTypes = BTreeMap<acir::circuit::ErrorSelector, ErrorType>;
+
 /// The execution result is the value returned from the circuit and any output from `println`.
 type ExecResult = (Result<WitnessStack<FieldElement>, NargoError<FieldElement>>, String);
 
+pub struct NargoErrorWithTypes(NargoError<FieldElement>, SsaErrorTypes);
+
+impl NargoErrorWithTypes {
+    /// Copy of `NargoError::user_defined_failure_message` accepting `SsaErrorTypes` instead of ABI errors.
+    fn user_defined_failure_message(&self) -> Option<String> {
+        match &self.0 {
+            NargoError::ExecutionError(error) => match error {
+                ExecutionError::AssertionFailed(payload, _, _) => match payload {
+                    ResolvedAssertionPayload::String(message) => Some(message.to_string()),
+                    ResolvedAssertionPayload::Raw(raw) => {
+                        let ssa_type = self.1.get(&raw.selector)?;
+                        match ssa_type {
+                            ErrorType::String(message) => Some(message.to_string()),
+                            ErrorType::Dynamic(_hir_type) => {
+                                // This would be the case if we have a format string that needs to be filled with the raw payload
+                                // decoded as ABI type. The code generator shouldn't produce this kind. It shouldn't be too difficult
+                                // to map the type, but the mapper in `crate::abi` doesn't handle format strings at the moment.
+                                panic!("didn't expect dynamic error types")
+                            }
+                        }
+                    }
+                },
+                ExecutionError::SolvingError(error, _) => match error {
+                    OpcodeResolutionError::BlackBoxFunctionFailed(_, reason) => {
+                        Some(reason.to_string())
+                    }
+                    _ => None,
+                },
+            },
+            NargoError::ForeignCallError(error) => Some(error.to_string()),
+            _ => None,
+        }
+    }
+}
+
 /// The result of the execution of compiled programs, decoded by their ABI.
-pub type CompareCompiledResult = CompareResult<InputValue, NargoError<FieldElement>>;
+pub type CompareCompiledResult = CompareResult<InputValue, NargoErrorWithTypes>;
 
 impl CompareCompiledResult {
     pub fn new(
         abi: &Abi,
+        ets1: &SsaErrorTypes,
+        ets2: &SsaErrorTypes,
         (res1, print1): ExecResult,
         (res2, print2): ExecResult,
     ) -> eyre::Result<Self> {
@@ -49,14 +92,17 @@ impl CompareCompiledResult {
         };
 
         match (res1, res2) {
-            (Err(e1), Err(e2)) => Ok(CompareResult::BothFailed(e1, e2)),
+            (Err(e1), Err(e2)) => Ok(CompareResult::BothFailed(
+                NargoErrorWithTypes(e1, ets1.clone()),
+                NargoErrorWithTypes(e2, ets2.clone()),
+            )),
             (Err(e1), Ok(ws2)) => Ok(CompareResult::LeftFailed(
-                e1,
+                NargoErrorWithTypes(e1, ets1.clone()),
                 ExecOutput { return_value: decode(ws2)?, print_output: print2 },
             )),
             (Ok(ws1), Err(e2)) => Ok(CompareResult::RightFailed(
                 ExecOutput { return_value: decode(ws1)?, print_output: print1 },
-                e2,
+                NargoErrorWithTypes(e2, ets2.clone()),
             )),
             (Ok(ws1), Ok(ws2)) => {
                 let o1 = ExecOutput { return_value: decode(ws1)?, print_output: print1 };
@@ -67,20 +113,24 @@ impl CompareCompiledResult {
     }
 }
 
-impl Comparable for NargoError<FieldElement> {
+impl Comparable for NargoErrorWithTypes {
     fn equivalent(e1: &Self, e2: &Self) -> bool {
         use ExecutionError::*;
 
         // For now consider non-execution errors as failures we need to investigate.
-        let NargoError::ExecutionError(ee1) = e1 else {
+        let NargoError::ExecutionError(ee1) = &e1.0 else {
             return false;
         };
-        let NargoError::ExecutionError(ee2) = e2 else {
+        let NargoError::ExecutionError(ee2) = &e2.0 else {
             return false;
         };
 
-        // let msg1 = e1.user_defined_failure_message(&abi1.error_types);
-        // let msg2 = e2.user_defined_failure_message(&abi2.error_types);
+        println!("{:?}", e1.1);
+        println!("{:?}", e2.1);
+
+        let msg1 = e1.user_defined_failure_message();
+        let msg2 = e2.user_defined_failure_message();
+        let is_same_msg = msg1.is_some() && msg2.is_some() && msg1 == msg2;
 
         match (ee1, ee2) {
             (
@@ -90,7 +140,7 @@ impl Comparable for NargoError<FieldElement> {
                 // Looks like the workaround we have for comptime failures originating from overflows and similar assertion failures.
                 true
             }
-            (AssertionFailed(p1, _, _), AssertionFailed(p2, _, _)) => p1 == p2,
+            (AssertionFailed(p1, _, _), AssertionFailed(p2, _, _)) => p1 == p2 || is_same_msg,
             (SolvingError(s1, _), SolvingError(s2, _)) => format!("{s1}") == format!("{s2}"),
             (SolvingError(s, _), AssertionFailed(p, _, _))
             | (AssertionFailed(p, _, _), SolvingError(s, _)) => match (s, p) {
@@ -98,7 +148,7 @@ impl Comparable for NargoError<FieldElement> {
                     OpcodeResolutionError::UnsatisfiedConstrain { .. },
                     ResolvedAssertionPayload::String(s),
                 ) => s == "Attempted to divide by zero",
-                _ => false,
+                _ => is_same_msg,
             },
         }
     }
@@ -110,6 +160,18 @@ impl Comparable for InputValue {
     }
 }
 
+impl std::fmt::Display for NargoErrorWithTypes {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        std::fmt::Display::fmt(&self.0, f)
+    }
+}
+
+impl std::fmt::Debug for NargoErrorWithTypes {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        std::fmt::Debug::fmt(&self.0, f)
+    }
+}
+
 /// Compare the execution of equivalent programs, compiled in different ways.
 pub struct CompareCompiled<P> {
     pub program: P,
@@ -147,7 +209,13 @@ impl<P> CompareCompiled<P> {
         let (res1, print1) = do_exec(&self.ssa1.artifact.program);
         let (res2, print2) = do_exec(&self.ssa2.artifact.program);
 
-        CompareCompiledResult::new(&self.abi, (res1, print1), (res2, print2))
+        CompareCompiledResult::new(
+            &self.abi,
+            &self.ssa1.artifact.error_types,
+            &self.ssa2.artifact.error_types,
+            (res1, print1),
+            (res2, print2),
+        )
     }
 }
 
diff --git a/tooling/ast_fuzzer/src/compare/comptime.rs b/tooling/ast_fuzzer/src/compare/comptime.rs
index f54ad7e682a..dd938f7137b 100644
--- a/tooling/ast_fuzzer/src/compare/comptime.rs
+++ b/tooling/ast_fuzzer/src/compare/comptime.rs
@@ -129,6 +129,8 @@ impl CompareComptime {
                     let res1 = Err(NargoError::ExecutionError(err));
                     return CompareCompiledResult::new(
                         &self.abi,
+                        &Default::default(), // We failed to compile the program, so no error types.
+                        &self.ssa.artifact.error_types,
                         (res1, "".to_string()),
                         (res2, print2),
                     );
@@ -143,7 +145,13 @@ impl CompareComptime {
         // Execute the 1st (comptime) program.
         let (res1, print1) = do_exec(&program1.program);
 
-        CompareCompiledResult::new(&self.abi, (res1, comptime_print + &print1), (res2, print2))
+        CompareCompiledResult::new(
+            &self.abi,
+            &Default::default(), // We have a fully compiled program at this point, no access to the SSA error types, just ABI error types.
+            &self.ssa.artifact.error_types,
+            (res1, comptime_print + &print1),
+            (res2, print2),
+        )
     }
 
     /// Generate a random comptime-viable AST, reverse it into

From c9c58d040cef076ca13c360f4cc4329b11a889d6 Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Wed, 11 Jun 2025 16:10:57 +0100
Subject: [PATCH 05/10] Avoid negative/large int literals

---
 tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs | 10 ++++++++++
 tooling/ast_fuzzer/src/compare/compiled.rs           |  3 ---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index 2e043af114c..d0d1dc29c2f 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -508,11 +508,21 @@ mod helpers {
 
         let type_options = TYPES.get_or_init(|| {
             let mut types = vec![Type::Bool, Type::Field];
+
             for sign in [Signedness::Signed, Signedness::Unsigned] {
                 for size in IntegerBitSize::iter() {
                     if sign.is_signed() && size.bit_size() == 1 {
                         continue;
                     }
+                    // Avoid negative literals; the frontend makes them difficult to work with in expressions
+                    // where no type inference information is available.
+                    if sign.is_signed() {
+                        continue;
+                    }
+                    // Avoid large integers; frontend doesn't like them.
+                    if size.bit_size() > 32 {
+                        continue;
+                    }
                     types.push(Type::Integer(sign, size));
                 }
             }
diff --git a/tooling/ast_fuzzer/src/compare/compiled.rs b/tooling/ast_fuzzer/src/compare/compiled.rs
index 1d8020b58d2..79a99e68720 100644
--- a/tooling/ast_fuzzer/src/compare/compiled.rs
+++ b/tooling/ast_fuzzer/src/compare/compiled.rs
@@ -125,9 +125,6 @@ impl Comparable for NargoErrorWithTypes {
             return false;
         };
 
-        println!("{:?}", e1.1);
-        println!("{:?}", e2.1);
-
         let msg1 = e1.user_defined_failure_message();
         let msg2 = e2.user_defined_failure_message();
         let is_same_msg = msg1.is_some() && msg2.is_some() && msg1 == msg2;

From 49a2694e4cd06dd185546da62d6856c898c2b790 Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Thu, 12 Jun 2025 13:42:39 +0100
Subject: [PATCH 06/10] Allow duplicating blocks which define variables

---
 Cargo.lock                                    |   2 +-
 .../src/ssa/ssa_gen/context.rs                |  11 +-
 cspell.json                                   |   1 +
 tooling/ast_fuzzer/fuzz/src/lib.rs            |   3 +-
 .../fuzz/src/targets/orig_vs_morph.rs         | 151 ++++++++++++---
 tooling/ast_fuzzer/src/program/rewrite/mod.rs |   2 +-
 tooling/ast_fuzzer/src/program/visitor.rs     | 181 ++++++++++--------
 7 files changed, 234 insertions(+), 117 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index a222be5e277..ee810cb488e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3414,8 +3414,8 @@ dependencies = [
  "noirc_abi",
  "noirc_evaluator",
  "noirc_frontend",
- "strum",
  "proptest",
+ "strum",
 ]
 
 [[package]]
diff --git a/compiler/noirc_evaluator/src/ssa/ssa_gen/context.rs b/compiler/noirc_evaluator/src/ssa/ssa_gen/context.rs
index 5ff91ccaa67..e8b9cdcfbe3 100644
--- a/compiler/noirc_evaluator/src/ssa/ssa_gen/context.rs
+++ b/compiler/noirc_evaluator/src/ssa/ssa_gen/context.rs
@@ -792,11 +792,18 @@ impl<'a> FunctionContext<'a> {
     /// Looks up the value of a given local variable. Expects the variable to have
     /// been previously defined or panics otherwise.
     pub(super) fn lookup(&self, id: LocalId) -> Values {
-        self.definitions.get(&id).expect("lookup: variable not defined").clone()
+        self.definitions
+            .get(&id)
+            .unwrap_or_else(|| panic!("lookup: variable {id:?} not defined"))
+            .clone()
     }
 
     pub(super) fn lookup_global(&self, id: GlobalId) -> Values {
-        self.shared_context.globals.get(&id).expect("lookup_global: variable not defined").clone()
+        self.shared_context
+            .globals
+            .get(&id)
+            .unwrap_or_else(|| panic!("lookup_global: variable {id:?} not defined"))
+            .clone()
     }
 
     /// Extract the given field of the tuple. Panics if the given Values is not
diff --git a/cspell.json b/cspell.json
index 4167a516ca6..88be7381c96 100644
--- a/cspell.json
+++ b/cspell.json
@@ -74,6 +74,7 @@
         "crypdoku",
         "csat",
         "ctstring",
+        "curr",
         "curvegroup",
         "databus",
         "debouncer",
diff --git a/tooling/ast_fuzzer/fuzz/src/lib.rs b/tooling/ast_fuzzer/fuzz/src/lib.rs
index 2ee78b57421..04c86eac55a 100644
--- a/tooling/ast_fuzzer/fuzz/src/lib.rs
+++ b/tooling/ast_fuzzer/fuzz/src/lib.rs
@@ -70,7 +70,8 @@ where
     // and print the AST, then resume the panic, because
     // `Program` has a `RefCell` in it, which is not unwind safe.
     if show_ast() {
-        eprintln!("---\n{}\n---", DisplayAstAsNoir(&program));
+        //eprintln!("---\n{}\n---", DisplayAstAsNoir(&program));
+        eprintln!("---\n{}\n---", (&program));
     }
 
     ssa::create_program_with_passes(program, options, primary, secondary).unwrap_or_else(|e| {
diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index 82e149728fa..885c2b37203 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -5,11 +5,11 @@ use crate::{compare_results_compiled, create_ssa_or_die, default_ssa_options};
 use arbitrary::{Arbitrary, Unstructured};
 use color_eyre::eyre;
 use noir_ast_fuzzer::compare::{CompareMorph, CompareOptions};
-use noir_ast_fuzzer::visitor;
 use noir_ast_fuzzer::{Config, visitor::visit_expr_mut};
+use noir_ast_fuzzer::{rewrite, visitor};
 use noirc_frontend::ast::UnaryOp;
 use noirc_frontend::monomorphization::ast::{
-    Call, Definition, Expression, Function, Ident, Program, Unary,
+    Call, Definition, Expression, Function, Ident, IdentId, LocalId, Program, Unary,
 };
 
 pub fn fuzz(u: &mut Unstructured) -> eyre::Result<()> {
@@ -61,11 +61,44 @@ fn rewrite_function(
     let ctx = rules::Context { unconstrained: func.unconstrained, ..Default::default() };
 
     let estimate = estimate_applicable_rules(&ctx, &func.body, rules);
-    let mut morph = MorphContext { target: max_rewrites.min(estimate), estimate, count: 0, rules };
+    let mut morph = MorphContext {
+        target: max_rewrites.min(estimate),
+        estimate,
+        count: 0,
+        rules,
+        locals: LocalContext::new(func),
+    };
 
     morph.rewrite_expr(&ctx, u, &mut func.body);
 }
 
+/// Context necessary to generate new local IDs during rewrites.
+///
+/// Potentially a place to reconstruct local variable scopes.
+struct LocalContext {
+    next_local_id: u32,
+    next_ident_id: u32,
+}
+
+impl LocalContext {
+    fn new(func: &Function) -> Self {
+        let (next_local_id, next_ident_id) = rewrite::next_local_and_ident_id(func);
+        Self { next_local_id, next_ident_id }
+    }
+
+    fn next_local_id(&mut self) -> LocalId {
+        let id = self.next_local_id;
+        self.next_local_id += 1;
+        LocalId(id)
+    }
+
+    fn next_ident_id(&mut self) -> IdentId {
+        let id = self.next_ident_id;
+        self.next_ident_id += 1;
+        IdentId(id)
+    }
+}
+
 /// Recursively apply rules while keeping a tally on how many we have done.
 struct MorphContext<'a> {
     /// Number of rewrites we want to achieve.
@@ -76,6 +109,8 @@ struct MorphContext<'a> {
     count: usize,
     /// Rules to apply.
     rules: &'a [rules::Rule],
+    /// Book keeping for local variables.
+    locals: LocalContext,
 }
 
 impl MorphContext<'_> {
@@ -148,14 +183,14 @@ impl MorphContext<'_> {
     /// Check if a rule can be applied on an expression. If it can, apply it based on some arbitrary
     /// criteria, returning a flag showing whether it was applied.
     fn try_apply_rule(
-        &self,
+        &mut self,
         ctx: &rules::Context,
         u: &mut Unstructured,
         expr: &mut Expression,
         rule: &rules::Rule,
     ) -> arbitrary::Result<bool> {
         if rule.matches(ctx, expr) && u.ratio(self.target, self.estimate)? {
-            rule.rewrite(u, expr)?;
+            rule.rewrite(u, &mut self.locals, expr)?;
             Ok(true)
         } else {
             Ok(false)
@@ -196,6 +231,8 @@ fn is_special_call(call: &Call) -> bool {
 
 /// Metamorphic transformation rules.
 mod rules {
+    use crate::targets::orig_vs_morph::{LocalContext, helpers::reassign_ids};
+
     use super::helpers::{gen_expr, has_side_effect};
     use arbitrary::{Arbitrary, Unstructured};
     use noir_ast_fuzzer::expr;
@@ -219,7 +256,8 @@ mod rules {
     /// Check if the rule can be applied on an expression.
     type MatchFn = dyn Fn(&Context, &Expression) -> bool;
     /// Apply the rule on an expression, mutating/replacing it in-place.
-    type RewriteFn = dyn Fn(&mut Unstructured, &mut Expression) -> arbitrary::Result<()>;
+    type RewriteFn =
+        dyn Fn(&mut Unstructured, &mut LocalContext, &mut Expression) -> arbitrary::Result<()>;
 
     /// Metamorphic transformation rule.
     pub struct Rule {
@@ -230,7 +268,12 @@ mod rules {
     impl Rule {
         pub fn new(
             matches: impl Fn(&Context, &Expression) -> bool + 'static,
-            rewrite: impl Fn(&mut Unstructured, &mut Expression) -> arbitrary::Result<()> + 'static,
+            rewrite: impl Fn(
+                &mut Unstructured,
+                &mut LocalContext,
+                &mut Expression,
+            ) -> arbitrary::Result<()>
+            + 'static,
         ) -> Self {
             Self { matches: Box::new(matches), rewrite: Box::new(rewrite) }
         }
@@ -244,9 +287,10 @@ mod rules {
         pub fn rewrite(
             &self,
             u: &mut Unstructured,
+            locals: &mut LocalContext,
             expr: &mut Expression,
         ) -> arbitrary::Result<()> {
-            (self.rewrite)(u, expr)
+            (self.rewrite)(u, locals, expr)
         }
     }
 
@@ -257,11 +301,12 @@ mod rules {
             bool_or_self(),
             bool_xor_self(),
             bool_xor_rand(),
-            commutative_arithmetic(),
-            inevitable(),
+            num_commute(),
+            any_inevitable(),
         ]
     }
 
+    // TODO: any numeric value `x` into `x * 1` or `x / 1`.
     /// Transform any numeric value `x` into `x +/- 0`.
     pub fn num_plus_minus_zero() -> Rule {
         Rule::new(
@@ -285,7 +330,7 @@ mod rules {
                     false
                 }
             },
-            |u, expr| {
+            |u, _locals, expr| {
                 let typ = expr.return_type().expect("only called on matching type").into_owned();
 
                 let op =
@@ -302,7 +347,7 @@ mod rules {
 
     /// Transform boolean value `x` into `x | x`.
     pub fn bool_or_self() -> Rule {
-        Rule::new(bool_rule_matches, |_u, expr| {
+        Rule::new(bool_rule_matches, |_u, _locals, expr| {
             expr::replace(expr, |expr| expr::binary(expr.clone(), BinaryOpKind::Or, expr));
             Ok(())
         })
@@ -310,7 +355,7 @@ mod rules {
 
     /// Transform boolean value `x` into `x ^ x ^ x`.
     pub fn bool_xor_self() -> Rule {
-        Rule::new(bool_rule_matches, |_u, expr| {
+        Rule::new(bool_rule_matches, |_u, _locals, expr| {
             expr::replace(expr, |expr| {
                 let rhs = expr::binary(expr.clone(), BinaryOpKind::Xor, expr.clone());
                 expr::binary(expr, BinaryOpKind::Xor, rhs)
@@ -321,7 +366,7 @@ mod rules {
 
     /// Transform boolean value `x` into `rnd ^ x ^ rnd`.
     pub fn bool_xor_rand() -> Rule {
-        Rule::new(bool_rule_matches, |u, expr| {
+        Rule::new(bool_rule_matches, |u, _locals, expr| {
             // This is where we could access the scope to look for a random bool variable.
             let rnd = expr::gen_literal(u, &Type::Bool)?;
             expr::replace(expr, |expr| {
@@ -335,7 +380,7 @@ mod rules {
     /// Transform commutative arithmetic operations:
     /// * `a + b` into `b + a`
     /// * `a * b` into `b * a`
-    pub fn commutative_arithmetic() -> Rule {
+    pub fn num_commute() -> Rule {
         Rule::new(
             |_ctx, expr| {
                 matches!(
@@ -346,7 +391,7 @@ mod rules {
                     })
                 ) && !has_side_effect(expr)
             },
-            |_u, expr| {
+            |_u, _locals, expr| {
                 let Expression::Binary(binary) = expr else {
                     unreachable!("the rule only matches Binary expressions");
                 };
@@ -358,26 +403,28 @@ mod rules {
         )
     }
 
-    /// Transform an expression into an if-then-else with the expression
+    /// Transform any expression into an if-then-else with the itself
     /// repeated in the _then_ and _else_ branch:
     /// * `x` into `if c { x } else { x }`
-    pub fn inevitable() -> Rule {
+    pub fn any_inevitable() -> Rule {
         Rule::new(
             |ctx, expr| {
                 !ctx.is_in_special_call
                     && !ctx.is_in_ref_mut
-                    // If we're in ACIR then don't turn loop ranges into expressions.
+                    // If we're in ACIR then don't turn loop ranges into non-constant expressions.
                     && (ctx.unconstrained || !ctx.is_in_range)
-                    // Avoid creating new variables because:
-                    // * `let x = 1;` transformed into `if true { let x = 1; } else { let x = 1; }` would leave `x` undefined.
-                    // * Creating variables in blocks would need to be assigned new IDs, and subsequent expressions would need to be rewritten,
-                    //   for which we currently lack the context. `for` introduces the index variable.
-                    && !expr::exists(expr, |expr| matches!(expr, Expression::Let(_) | Expression::For(_)))
+                    // `let x = 1;` transformed into `if true { let x = 1; } else { let x = 1; }` would leave `x` undefined.
+                    && !matches!(expr, Expression::Let(_))
             },
-            |u, expr| {
+            |u, locals, expr| {
                 let typ = expr.return_type().map(|typ| typ.into_owned()).unwrap_or(Type::Unit);
                 let cond = gen_expr(u, &Type::Bool, 2)?;
-                *expr = expr::if_else(cond, expr.clone(), expr.clone(), typ);
+
+                // Duplicate the expression, then assign new IDs to all variables created in it.
+                let mut alt = expr.clone();
+                reassign_ids(locals, &mut alt);
+
+                expr::replace(expr, |expr| expr::if_else(cond, expr, alt, typ));
                 Ok(())
             },
         )
@@ -413,17 +460,19 @@ mod rules {
 }
 
 mod helpers {
-    use std::sync::OnceLock;
+    use std::{cell::RefCell, collections::HashMap, sync::OnceLock};
 
     use arbitrary::Unstructured;
-    use noir_ast_fuzzer::{expr, types};
+    use noir_ast_fuzzer::{expr, types, visitor::visit_expr_be_mut};
     use noirc_frontend::{
         ast::{IntegerBitSize, UnaryOp},
-        monomorphization::ast::{BinaryOp, Expression, Type},
+        monomorphization::ast::{BinaryOp, Definition, Expression, LocalId, Type},
         shared::Signedness,
     };
     use strum::IntoEnumIterator;
 
+    use crate::targets::orig_vs_morph::LocalContext;
+
     /// Check if an expression can have a side effect, in which case duplicating or reordering it could
     /// change the behavior of the program.
     pub(super) fn has_side_effect(expr: &Expression) -> bool {
@@ -558,6 +607,50 @@ mod helpers {
 
     /// Types we can consider using in this context.
     static TYPES: OnceLock<Vec<Type>> = OnceLock::new();
+
+    pub(super) fn reassign_ids(locals: &mut LocalContext, expr: &mut Expression) {
+        fn replace_local_id(
+            locals: &mut LocalContext,
+            replacements: &mut HashMap<LocalId, LocalId>,
+            id: &mut LocalId,
+        ) {
+            let curr = *id;
+            let next = locals.next_local_id();
+            replacements.insert(curr, next);
+            *id = next;
+        }
+
+        let replacements = RefCell::new(HashMap::new());
+
+        visit_expr_be_mut(
+            expr,
+            &mut |expr| {
+                match expr {
+                    Expression::Let(let_) => {
+                        replace_local_id(locals, &mut replacements.borrow_mut(), &mut let_.id)
+                    }
+                    Expression::For(for_) => replace_local_id(
+                        locals,
+                        &mut replacements.borrow_mut(),
+                        &mut for_.index_variable,
+                    ),
+                    Expression::Ident(ident) => {
+                        ident.id = locals.next_ident_id();
+                    }
+                    _ => (),
+                }
+                (true, ())
+            },
+            &mut |_, _| {},
+            &mut |ident| {
+                if let Definition::Local(id) = &mut ident.definition {
+                    if let Some(replacement) = replacements.borrow().get(id) {
+                        *id = *replacement;
+                    }
+                }
+            },
+        );
+    }
 }
 
 #[cfg(test)]
diff --git a/tooling/ast_fuzzer/src/program/rewrite/mod.rs b/tooling/ast_fuzzer/src/program/rewrite/mod.rs
index fccca8007ef..bcd2fdb1bb5 100644
--- a/tooling/ast_fuzzer/src/program/rewrite/mod.rs
+++ b/tooling/ast_fuzzer/src/program/rewrite/mod.rs
@@ -10,7 +10,7 @@ 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.
-fn next_local_and_ident_id(func: &Function) -> (u32, u32) {
+pub fn next_local_and_ident_id(func: &Function) -> (u32, u32) {
     let mut next_local_id = func.parameters.iter().map(|p| p.0.0 + 1).max().unwrap_or_default();
     let mut next_ident_id = 0;
 
diff --git a/tooling/ast_fuzzer/src/program/visitor.rs b/tooling/ast_fuzzer/src/program/visitor.rs
index 26eb8f4181c..cc5895bee0b 100644
--- a/tooling/ast_fuzzer/src/program/visitor.rs
+++ b/tooling/ast_fuzzer/src/program/visitor.rs
@@ -1,4 +1,12 @@
-use noirc_frontend::monomorphization::ast::{Expression, LValue, Literal};
+use noirc_frontend::monomorphization::ast::{Expression, Ident, LValue, Literal};
+
+/// Visit all identifiers under the [Expression].
+pub fn visit_ident_mut<V>(expr: &mut Expression, v: &mut V)
+where
+    V: FnMut(&mut Ident),
+{
+    visit_expr_be_mut(expr, &mut |_| (true, ()), &mut |_, _| {}, v);
+}
 
 /// Visit the contents of an [Expression] representing the AST.
 ///
@@ -10,7 +18,7 @@ pub fn visit_expr_mut<V>(expr: &mut Expression, v: &mut V)
 where
     V: FnMut(&mut Expression) -> bool,
 {
-    visit_expr_be_mut(expr, &mut |e| (v(e), ()), &mut |_, _| {});
+    visit_expr_be_mut(expr, &mut |e| (v(e), ()), &mut |_, _| {}, &mut |_| {});
 }
 
 /// Visit for the contents of an [Expression] representing the AST,
@@ -23,111 +31,115 @@ where
 ///
 /// Compared to [visit_expr_mut], this version allows the caller to maintain
 /// scopes and context, facilitated by a _token_ passed between _begin_ and _end_.
-pub fn visit_expr_be_mut<B, E, T>(expr: &mut Expression, b: &mut B, e: &mut E)
+///
+/// It also takes a function to modify [Ident]s, which can be located in
+/// `Expression::Ident` or `Expression::Assign`.
+pub fn visit_expr_be_mut<B, E, T, I>(expr: &mut Expression, b: &mut B, e: &mut E, i: &mut I)
 where
     B: FnMut(&mut Expression) -> (bool, T),
     E: FnMut(&Expression, T),
+    I: FnMut(&mut Ident),
 {
     let (go, token) = b(expr);
     if !go {
         return e(expr, token);
     }
     match expr {
-        Expression::Ident(_) => {}
+        Expression::Ident(ident) => i(ident),
         Expression::Literal(literal) => match literal {
             Literal::Array(array_literal) => {
                 for expr in array_literal.contents.iter_mut() {
-                    visit_expr_be_mut(expr, b, e);
+                    visit_expr_be_mut(expr, b, e, i);
                 }
             }
             Literal::Slice(array_literal) => {
                 for expr in array_literal.contents.iter_mut() {
-                    visit_expr_be_mut(expr, b, e);
+                    visit_expr_be_mut(expr, b, e, i);
                 }
             }
             Literal::Integer(_, _, _) | Literal::Bool(_) | Literal::Unit | Literal::Str(_) => {}
             Literal::FmtStr(_, _, expr) => {
-                visit_expr_be_mut(expr, b, e);
+                visit_expr_be_mut(expr, b, e, i);
             }
         },
         Expression::Block(exprs) => {
             for expr in exprs.iter_mut() {
-                visit_expr_be_mut(expr, b, e);
+                visit_expr_be_mut(expr, b, e, i);
             }
         }
         Expression::Unary(unary) => {
-            visit_expr_be_mut(&mut unary.rhs, b, e);
+            visit_expr_be_mut(&mut unary.rhs, b, e, i);
         }
         Expression::Binary(binary) => {
-            visit_expr_be_mut(&mut binary.lhs, b, e);
-            visit_expr_be_mut(&mut binary.rhs, b, e);
+            visit_expr_be_mut(&mut binary.lhs, b, e, i);
+            visit_expr_be_mut(&mut binary.rhs, b, e, i);
         }
         Expression::Index(index) => {
-            visit_expr_be_mut(&mut index.collection, b, e);
-            visit_expr_be_mut(&mut index.index, b, e);
+            visit_expr_be_mut(&mut index.collection, b, e, i);
+            visit_expr_be_mut(&mut index.index, b, e, i);
         }
         Expression::Cast(cast) => {
-            visit_expr_be_mut(&mut cast.lhs, b, e);
+            visit_expr_be_mut(&mut cast.lhs, b, e, i);
         }
         Expression::For(for_) => {
-            visit_expr_be_mut(&mut for_.start_range, b, e);
-            visit_expr_be_mut(&mut for_.end_range, b, e);
-            visit_expr_be_mut(&mut for_.block, b, e);
+            visit_expr_be_mut(&mut for_.start_range, b, e, i);
+            visit_expr_be_mut(&mut for_.end_range, b, e, i);
+            visit_expr_be_mut(&mut for_.block, b, e, i);
         }
         Expression::Loop(expr) => {
-            visit_expr_be_mut(expr, b, e);
+            visit_expr_be_mut(expr, b, e, i);
         }
         Expression::While(while_) => {
-            visit_expr_be_mut(&mut while_.condition, b, e);
-            visit_expr_be_mut(&mut while_.body, b, e);
+            visit_expr_be_mut(&mut while_.condition, b, e, i);
+            visit_expr_be_mut(&mut while_.body, b, e, i);
         }
         Expression::If(if_) => {
-            visit_expr_be_mut(&mut if_.condition, b, e);
-            visit_expr_be_mut(&mut if_.consequence, b, e);
+            visit_expr_be_mut(&mut if_.condition, b, e, i);
+            visit_expr_be_mut(&mut if_.consequence, b, e, i);
             if let Some(ref mut alternative) = if_.alternative {
-                visit_expr_be_mut(alternative, b, e);
+                visit_expr_be_mut(alternative, b, e, i);
             }
         }
         Expression::Match(match_) => {
             for case in match_.cases.iter_mut() {
-                visit_expr_be_mut(&mut case.branch, b, e);
+                visit_expr_be_mut(&mut case.branch, b, e, i);
             }
             if let Some(ref mut case) = match_.default_case {
-                visit_expr_be_mut(case, b, e);
+                visit_expr_be_mut(case, b, e, i);
             }
         }
         Expression::Tuple(exprs) => {
             for expr in exprs.iter_mut() {
-                visit_expr_be_mut(expr, b, e);
+                visit_expr_be_mut(expr, b, e, i);
             }
         }
         Expression::ExtractTupleField(expr, _) => {
-            visit_expr_be_mut(expr, b, e);
+            visit_expr_be_mut(expr, b, e, i);
         }
         Expression::Call(call) => {
-            visit_expr_be_mut(&mut call.func, b, e);
+            visit_expr_be_mut(&mut call.func, b, e, i);
             for arg in call.arguments.iter_mut() {
-                visit_expr_be_mut(arg, b, e);
+                visit_expr_be_mut(arg, b, e, i);
             }
         }
         Expression::Let(let_) => {
-            visit_expr_be_mut(&mut let_.expression, b, e);
+            visit_expr_be_mut(&mut let_.expression, b, e, i);
         }
         Expression::Constrain(expr, _, _) => {
-            visit_expr_be_mut(expr, b, e);
+            visit_expr_be_mut(expr, b, e, i);
         }
         Expression::Assign(assign) => {
-            visit_lvalue_mut(&mut assign.lvalue, b, e);
-            visit_expr_be_mut(&mut assign.expression, b, e);
+            visit_lvalue_mut(&mut assign.lvalue, b, e, i);
+            visit_expr_be_mut(&mut assign.expression, b, e, i);
         }
         Expression::Semi(expr) => {
-            visit_expr_be_mut(expr, b, e);
+            visit_expr_be_mut(expr, b, e, i);
         }
         Expression::Clone(expr) => {
-            visit_expr_be_mut(expr, b, e);
+            visit_expr_be_mut(expr, b, e, i);
         }
         Expression::Drop(expr) => {
-            visit_expr_be_mut(expr, b, e);
+            visit_expr_be_mut(expr, b, e, i);
         }
         Expression::Break => {}
         Expression::Continue => {}
@@ -136,22 +148,23 @@ where
     e(expr, token);
 }
 
-fn visit_lvalue_mut<B, E, T>(lvalue: &mut LValue, b: &mut B, e: &mut E)
+fn visit_lvalue_mut<B, E, T, I>(lvalue: &mut LValue, b: &mut B, e: &mut E, i: &mut I)
 where
     B: FnMut(&mut Expression) -> (bool, T),
     E: FnMut(&Expression, T),
+    I: FnMut(&mut Ident),
 {
     match lvalue {
-        LValue::Ident(_) => {}
+        LValue::Ident(ident) => i(ident),
         LValue::Index { array, index, .. } => {
-            visit_lvalue_mut(array.as_mut(), b, e);
-            visit_expr_be_mut(index.as_mut(), b, e);
+            visit_lvalue_mut(array.as_mut(), b, e, i);
+            visit_expr_be_mut(index.as_mut(), b, e, i);
         }
         LValue::MemberAccess { object, .. } => {
-            visit_lvalue_mut(object.as_mut(), b, e);
+            visit_lvalue_mut(object.as_mut(), b, e, i);
         }
         LValue::Dereference { reference, .. } => {
-            visit_lvalue_mut(reference.as_mut(), b, e);
+            visit_lvalue_mut(reference.as_mut(), b, e, i);
         }
     }
 }
@@ -167,7 +180,7 @@ pub fn visit_expr<V>(expr: &Expression, v: &mut V)
 where
     V: FnMut(&Expression) -> bool,
 {
-    visit_expr_be(expr, &mut |e| (v(e), ()), &mut |_, _| {});
+    visit_expr_be(expr, &mut |e| (v(e), ()), &mut |_, _| {}, &mut |_| {});
 }
 
 /// Visit the contents of an [Expression] representing the AST,
@@ -181,10 +194,11 @@ where
 ///
 /// Compared to [visit_expr], this version allows the caller to maintain
 /// scopes and context, facilitated by a _token_ passed between _begin_ and _end_.
-pub fn visit_expr_be<B, E, T>(expr: &Expression, b: &mut B, e: &mut E)
+pub fn visit_expr_be<B, E, T, I>(expr: &Expression, b: &mut B, e: &mut E, i: &mut I)
 where
     B: FnMut(&Expression) -> (bool, T),
     E: FnMut(&Expression, T),
+    I: FnMut(&Ident),
 {
     let (go, token) = b(expr);
 
@@ -193,101 +207,101 @@ where
     }
 
     match expr {
-        Expression::Ident(_) => {}
+        Expression::Ident(ident) => i(ident),
         Expression::Literal(literal) => match literal {
             Literal::Array(array_literal) => {
                 for expr in &array_literal.contents {
-                    visit_expr_be(expr, b, e);
+                    visit_expr_be(expr, b, e, i);
                 }
             }
             Literal::Slice(array_literal) => {
                 for expr in &array_literal.contents {
-                    visit_expr_be(expr, b, e);
+                    visit_expr_be(expr, b, e, i);
                 }
             }
             Literal::Integer(_, _, _) | Literal::Bool(_) | Literal::Unit | Literal::Str(_) => {}
             Literal::FmtStr(_, _, expr) => {
-                visit_expr_be(expr, b, e);
+                visit_expr_be(expr, b, e, i);
             }
         },
         Expression::Block(exprs) => {
             for expr in exprs {
-                visit_expr_be(expr, b, e);
+                visit_expr_be(expr, b, e, i);
             }
         }
         Expression::Unary(unary) => {
-            visit_expr_be(&unary.rhs, b, e);
+            visit_expr_be(&unary.rhs, b, e, i);
         }
         Expression::Binary(binary) => {
-            visit_expr_be(&binary.lhs, b, e);
-            visit_expr_be(&binary.rhs, b, e);
+            visit_expr_be(&binary.lhs, b, e, i);
+            visit_expr_be(&binary.rhs, b, e, i);
         }
         Expression::Index(index) => {
-            visit_expr_be(&index.collection, b, e);
-            visit_expr_be(&index.index, b, e);
+            visit_expr_be(&index.collection, b, e, i);
+            visit_expr_be(&index.index, b, e, i);
         }
         Expression::Cast(cast) => {
-            visit_expr_be(&cast.lhs, b, e);
+            visit_expr_be(&cast.lhs, b, e, i);
         }
         Expression::For(for_) => {
-            visit_expr_be(&for_.start_range, b, e);
-            visit_expr_be(&for_.end_range, b, e);
-            visit_expr_be(&for_.block, b, e);
+            visit_expr_be(&for_.start_range, b, e, i);
+            visit_expr_be(&for_.end_range, b, e, i);
+            visit_expr_be(&for_.block, b, e, i);
         }
         Expression::Loop(expr) => {
-            visit_expr_be(expr, b, e);
+            visit_expr_be(expr, b, e, i);
         }
         Expression::While(while_) => {
-            visit_expr_be(&while_.condition, b, e);
-            visit_expr_be(&while_.body, b, e);
+            visit_expr_be(&while_.condition, b, e, i);
+            visit_expr_be(&while_.body, b, e, i);
         }
         Expression::If(if_) => {
-            visit_expr_be(&if_.condition, b, e);
-            visit_expr_be(&if_.consequence, b, e);
+            visit_expr_be(&if_.condition, b, e, i);
+            visit_expr_be(&if_.consequence, b, e, i);
             if let Some(ref alternative) = if_.alternative {
-                visit_expr_be(alternative, b, e);
+                visit_expr_be(alternative, b, e, i);
             }
         }
         Expression::Match(match_) => {
             for case in &match_.cases {
-                visit_expr_be(&case.branch, b, e);
+                visit_expr_be(&case.branch, b, e, i);
             }
             if let Some(ref case) = match_.default_case {
-                visit_expr_be(case, b, e);
+                visit_expr_be(case, b, e, i);
             }
         }
         Expression::Tuple(exprs) => {
             for expr in exprs {
-                visit_expr_be(expr, b, e);
+                visit_expr_be(expr, b, e, i);
             }
         }
         Expression::ExtractTupleField(expr, _) => {
-            visit_expr_be(expr, b, e);
+            visit_expr_be(expr, b, e, i);
         }
         Expression::Call(call) => {
-            visit_expr_be(&call.func, b, e);
+            visit_expr_be(&call.func, b, e, i);
             for arg in &call.arguments {
-                visit_expr_be(arg, b, e);
+                visit_expr_be(arg, b, e, i);
             }
         }
         Expression::Let(let_) => {
-            visit_expr_be(&let_.expression, b, e);
+            visit_expr_be(&let_.expression, b, e, i);
         }
         Expression::Constrain(expr, _, _) => {
-            visit_expr_be(expr, b, e);
+            visit_expr_be(expr, b, e, i);
         }
         Expression::Assign(assign) => {
-            visit_lvalue(&assign.lvalue, b, e);
-            visit_expr_be(&assign.expression, b, e);
+            visit_lvalue(&assign.lvalue, b, e, i);
+            visit_expr_be(&assign.expression, b, e, i);
         }
         Expression::Semi(expr) => {
-            visit_expr_be(expr, b, e);
+            visit_expr_be(expr, b, e, i);
         }
         Expression::Clone(expr) => {
-            visit_expr_be(expr, b, e);
+            visit_expr_be(expr, b, e, i);
         }
         Expression::Drop(expr) => {
-            visit_expr_be(expr, b, e);
+            visit_expr_be(expr, b, e, i);
         }
         Expression::Break => {}
         Expression::Continue => {}
@@ -296,22 +310,23 @@ where
     e(expr, token);
 }
 
-fn visit_lvalue<B, E, T>(lvalue: &LValue, b: &mut B, e: &mut E)
+fn visit_lvalue<B, E, T, I>(lvalue: &LValue, b: &mut B, e: &mut E, i: &mut I)
 where
     B: FnMut(&Expression) -> (bool, T),
     E: FnMut(&Expression, T),
+    I: FnMut(&Ident),
 {
     match lvalue {
-        LValue::Ident(_) => {}
+        LValue::Ident(ident) => i(ident),
         LValue::Index { array, index, .. } => {
-            visit_lvalue(array.as_ref(), b, e);
-            visit_expr_be(index.as_ref(), b, e);
+            visit_lvalue(array.as_ref(), b, e, i);
+            visit_expr_be(index.as_ref(), b, e, i);
         }
         LValue::MemberAccess { object, .. } => {
-            visit_lvalue(object.as_ref(), b, e);
+            visit_lvalue(object.as_ref(), b, e, i);
         }
         LValue::Dereference { reference, .. } => {
-            visit_lvalue(reference.as_ref(), b, e);
+            visit_lvalue(reference.as_ref(), b, e, i);
         }
     }
 }

From 541e51416471e22113de9beb7418cbbde397296d Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Thu, 12 Jun 2025 14:05:23 +0100
Subject: [PATCH 07/10] num_mul_one and num_div_one

---
 .../fuzz/src/targets/orig_vs_morph.rs         | 87 +++++++++++--------
 1 file changed, 51 insertions(+), 36 deletions(-)

diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index 885c2b37203..327c0f19ab5 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -234,7 +234,7 @@ mod rules {
     use crate::targets::orig_vs_morph::{LocalContext, helpers::reassign_ids};
 
     use super::helpers::{gen_expr, has_side_effect};
-    use arbitrary::{Arbitrary, Unstructured};
+    use arbitrary::Unstructured;
     use noir_ast_fuzzer::expr;
     use noirc_frontend::{
         ast::BinaryOpKind,
@@ -297,7 +297,10 @@ mod rules {
     /// Construct all rules that we can apply on a program.
     pub fn all() -> Vec<Rule> {
         vec![
-            num_plus_minus_zero(),
+            num_add_zero(),
+            num_sub_zero(),
+            num_mul_one(),
+            num_div_one(),
             bool_or_self(),
             bool_xor_self(),
             bool_xor_rand(),
@@ -306,43 +309,33 @@ mod rules {
         ]
     }
 
-    // TODO: any numeric value `x` into `x * 1` or `x / 1`.
-    /// Transform any numeric value `x` into `x +/- 0`.
-    pub fn num_plus_minus_zero() -> Rule {
-        Rule::new(
-            |ctx, expr| {
-                // Because of #8305 we can't reliably use expressions in ranges in ACIR.
-                if ctx.is_in_range && !ctx.unconstrained {
-                    return false;
-                }
-                // If we rewrite `&mut x` into `&mut (x - 0)` we will alter the semantics.
-                if ctx.is_in_ref_mut {
-                    return false;
-                }
-                // Appending 0 to a block would look odd.
-                if matches!(expr, Expression::Block(_)) {
-                    return false;
-                }
-                // We can apply this rule on anything that returns a number.
-                if let Some(typ) = expr.return_type() {
-                    matches!(typ.as_ref(), Type::Field | Type::Integer(_, _))
-                } else {
-                    false
-                }
-            },
-            |u, _locals, expr| {
-                let typ = expr.return_type().expect("only called on matching type").into_owned();
+    /// Transform any numeric value `x` into `x <op> <rhs>`
+    fn num_op(op: BinaryOpKind, rhs: u32) -> Rule {
+        Rule::new(num_rule_matches, move |_u, _locals, expr| {
+            let typ = expr.return_type().expect("only called on matching type").into_owned();
+            expr::replace(expr, |expr| expr::binary(expr, op, expr::int_literal(rhs, false, typ)));
+            Ok(())
+        })
+    }
 
-                let op =
-                    if bool::arbitrary(u)? { BinaryOpKind::Add } else { BinaryOpKind::Subtract };
+    /// Transform any numeric value `x` into `x+0`
+    pub fn num_add_zero() -> Rule {
+        num_op(BinaryOpKind::Add, 0)
+    }
 
-                expr::replace(expr, |expr| {
-                    expr::binary(expr, op, expr::int_literal(0u32, false, typ))
-                });
+    /// Transform any numeric value `x` into `x-0`
+    pub fn num_sub_zero() -> Rule {
+        num_op(BinaryOpKind::Subtract, 0)
+    }
 
-                Ok(())
-            },
-        )
+    /// Transform any numeric value `x` into `x*1`
+    pub fn num_mul_one() -> Rule {
+        num_op(BinaryOpKind::Multiply, 1)
+    }
+
+    /// Transform any numeric value `x` into `x/1`
+    pub fn num_div_one() -> Rule {
+        num_op(BinaryOpKind::Divide, 1)
     }
 
     /// Transform boolean value `x` into `x | x`.
@@ -457,6 +450,28 @@ mod rules {
             false
         }
     }
+
+    /// Common condition for numeric rules
+    fn num_rule_matches(ctx: &Context, expr: &Expression) -> bool {
+        // Because of #8305 we can't reliably use expressions in ranges in ACIR.
+        if ctx.is_in_range && !ctx.unconstrained {
+            return false;
+        }
+        // If we rewrite `&mut x` into `&mut (x - 0)` we will alter the semantics.
+        if ctx.is_in_ref_mut {
+            return false;
+        }
+        // Appending 0 to a block would look odd.
+        if matches!(expr, Expression::Block(_)) {
+            return false;
+        }
+        // We can apply this rule on anything that returns a number.
+        if let Some(typ) = expr.return_type() {
+            matches!(typ.as_ref(), Type::Field | Type::Integer(_, _))
+        } else {
+            false
+        }
+    }
 }
 
 mod helpers {

From 0ed20594eaec0610bb29140940280eb2a9c03176 Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Thu, 12 Jun 2025 14:57:03 +0100
Subject: [PATCH 08/10] Add int_break_up

---
 tooling/ast_fuzzer/fuzz/src/lib.rs            |  3 +-
 .../fuzz/src/targets/orig_vs_morph.rs         | 45 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/tooling/ast_fuzzer/fuzz/src/lib.rs b/tooling/ast_fuzzer/fuzz/src/lib.rs
index 04c86eac55a..2ee78b57421 100644
--- a/tooling/ast_fuzzer/fuzz/src/lib.rs
+++ b/tooling/ast_fuzzer/fuzz/src/lib.rs
@@ -70,8 +70,7 @@ where
     // and print the AST, then resume the panic, because
     // `Program` has a `RefCell` in it, which is not unwind safe.
     if show_ast() {
-        //eprintln!("---\n{}\n---", DisplayAstAsNoir(&program));
-        eprintln!("---\n{}\n---", (&program));
+        eprintln!("---\n{}\n---", DisplayAstAsNoir(&program));
     }
 
     ssa::create_program_with_passes(program, options, primary, secondary).unwrap_or_else(|e| {
diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index 327c0f19ab5..ac5f4a6f20a 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -238,7 +238,8 @@ mod rules {
     use noir_ast_fuzzer::expr;
     use noirc_frontend::{
         ast::BinaryOpKind,
-        monomorphization::ast::{Binary, Expression, Type},
+        monomorphization::ast::{Binary, Expression, Literal, Type},
+        signed_field::SignedField,
     };
 
     #[derive(Clone, Debug, Default)]
@@ -306,6 +307,7 @@ mod rules {
             bool_xor_rand(),
             num_commute(),
             any_inevitable(),
+            int_break_up(),
         ]
     }
 
@@ -338,6 +340,47 @@ mod rules {
         num_op(BinaryOpKind::Divide, 1)
     }
 
+    /// Break an integer literal `a` into `b + c`.
+    pub fn int_break_up() -> Rule {
+        Rule::new(
+            |ctx, expr| {
+                if ctx.is_in_range && !ctx.unconstrained || ctx.is_in_ref_mut {
+                    return false;
+                }
+                matches!(expr, Expression::Literal(Literal::Integer(_, Type::Integer(_, _), _)))
+            },
+            |u, _locals, expr| {
+                let Expression::Literal(Literal::Integer(a, typ, loc)) = expr else {
+                    unreachable!("matched only integer literals, got {expr}");
+                };
+                let mut b_expr = expr::gen_literal(u, typ)?;
+                let Expression::Literal(Literal::Integer(b, _, _)) = &mut b_expr else {
+                    unreachable!("generated a literal of the same type");
+                };
+
+                // Make them have the same sign, so they are on the same side of 0 and a single number
+                // can add up to them without overflow. (e.g. there is no x such that `i32::MIN + x == i32::MAX`)
+                if a.is_negative() && !b.is_negative() {
+                    *b = SignedField::negative(b.absolute_value());
+                } else if !a.is_negative() && b.is_negative() {
+                    *b = SignedField::positive(b.absolute_value());
+                }
+
+                let (op, c) = if *a >= *b {
+                    (BinaryOpKind::Add, (*a - *b))
+                } else {
+                    (BinaryOpKind::Subtract, (*b - *a))
+                };
+
+                let c_expr = Expression::Literal(Literal::Integer(c, typ.clone(), loc.clone()));
+
+                *expr = expr::binary(b_expr, op, c_expr);
+
+                Ok(())
+            },
+        )
+    }
+
     /// Transform boolean value `x` into `x | x`.
     pub fn bool_or_self() -> Rule {
         Rule::new(bool_rule_matches, |_u, _locals, expr| {

From b62f77883150aceeb17047f1348edc5fd2cab5a4 Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Thu, 12 Jun 2025 16:20:15 +0100
Subject: [PATCH 09/10] Build up the scope during travesal and use bool vars

---
 .../fuzz/src/targets/orig_vs_morph.rs         | 250 ++++++++++++------
 tooling/ast_fuzzer/src/lib.rs                 |   2 +-
 tooling/ast_fuzzer/src/program/mod.rs         |   2 +-
 tooling/ast_fuzzer/src/program/scope.rs       |  30 +--
 4 files changed, 187 insertions(+), 97 deletions(-)

diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index ac5f4a6f20a..82c4c367401 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -1,11 +1,14 @@
 //! Perform random metamorphic mutations on the AST and check that the
 //! execution result does not change.
 
+use std::cell::{Cell, RefCell};
+
 use crate::{compare_results_compiled, create_ssa_or_die, default_ssa_options};
 use arbitrary::{Arbitrary, Unstructured};
 use color_eyre::eyre;
 use noir_ast_fuzzer::compare::{CompareMorph, CompareOptions};
-use noir_ast_fuzzer::{Config, visitor::visit_expr_mut};
+use noir_ast_fuzzer::scope::ScopeStack;
+use noir_ast_fuzzer::{Config, visitor::visit_expr_be_mut};
 use noir_ast_fuzzer::{rewrite, visitor};
 use noirc_frontend::ast::UnaryOp;
 use noirc_frontend::monomorphization::ast::{
@@ -61,12 +64,12 @@ fn rewrite_function(
     let ctx = rules::Context { unconstrained: func.unconstrained, ..Default::default() };
 
     let estimate = estimate_applicable_rules(&ctx, &func.body, rules);
-    let mut morph = MorphContext {
+    let morph = MorphContext {
         target: max_rewrites.min(estimate),
         estimate,
-        count: 0,
+        count: Cell::new(0),
         rules,
-        locals: LocalContext::new(func),
+        vars: RefCell::new(VariableContext::new(func)),
     };
 
     morph.rewrite_expr(&ctx, u, &mut func.body);
@@ -75,15 +78,23 @@ fn rewrite_function(
 /// Context necessary to generate new local IDs during rewrites.
 ///
 /// Potentially a place to reconstruct local variable scopes.
-struct LocalContext {
+struct VariableContext {
     next_local_id: u32,
     next_ident_id: u32,
+    locals: ScopeStack<LocalId>,
 }
 
-impl LocalContext {
+impl VariableContext {
     fn new(func: &Function) -> Self {
         let (next_local_id, next_ident_id) = rewrite::next_local_and_ident_id(func);
-        Self { next_local_id, next_ident_id }
+
+        let locals = ScopeStack::new(
+            func.parameters
+                .iter()
+                .map(|(id, mutable, name, typ, _vis)| (*id, *mutable, name.clone(), typ.clone())),
+        );
+
+        Self { next_local_id, next_ident_id, locals }
     }
 
     fn next_local_id(&mut self) -> LocalId {
@@ -106,91 +117,146 @@ struct MorphContext<'a> {
     /// (Over)estimate of the maximum number we could hope to apply.
     estimate: usize,
     /// Number of rewrites applied so far, up to the `target`.
-    count: usize,
+    count: Cell<usize>,
     /// Rules to apply.
     rules: &'a [rules::Rule],
     /// Book keeping for local variables.
-    locals: LocalContext,
+    vars: RefCell<VariableContext>,
 }
 
 impl MorphContext<'_> {
     /// Check if we have reached the target.
     fn limit_reached(&self) -> bool {
-        self.target == 0 || self.estimate == 0 || self.count == self.target
+        self.target == 0 || self.estimate == 0 || self.count.get() == self.target
     }
 
-    fn rewrite_expr(&mut self, ctx: &rules::Context, u: &mut Unstructured, expr: &mut Expression) {
-        visit_expr_mut(expr, &mut |expr: &mut Expression| {
-            if self.limit_reached() {
-                return false;
-            }
-            match expr {
-                Expression::For(for_) => {
-                    // Separate context for just the ranges.
-                    let range_ctx = rules::Context { is_in_range: true, ..*ctx };
-                    self.rewrite_expr(&range_ctx, u, &mut for_.start_range);
-                    self.rewrite_expr(&range_ctx, u, &mut for_.end_range);
-                    // Original context for the body.
-                    self.rewrite_expr(ctx, u, &mut for_.block);
-                    // No need to visit children, we just visited them.
-                    false
+    fn rewrite_expr(&self, ctx: &rules::Context, u: &mut Unstructured, expr: &mut Expression) {
+        visit_expr_be_mut(
+            expr,
+            &mut |expr: &mut Expression| {
+                if self.limit_reached() {
+                    return (false, false);
                 }
-                Expression::Unary(
-                    unary @ Unary { operator: UnaryOp::Reference { mutable: true }, .. },
-                ) => {
-                    let ctx = rules::Context { is_in_ref_mut: true, ..*ctx };
-                    self.rewrite_expr(&ctx, u, &mut unary.rhs);
-                    false
+
+                // Check if we are entering a new scope.
+                let entered = matches!(
+                    expr,
+                    Expression::Block(_)
+                        | Expression::For(_)
+                        | Expression::Loop(_)
+                        | Expression::While(_)
+                        | Expression::If(_)
+                );
+
+                if entered {
+                    self.vars.borrow_mut().locals.enter();
                 }
-                Expression::Call(call) if is_special_call(call) => {
-                    let ctx = rules::Context { is_in_special_call: true, ..*ctx };
-                    for arg in call.arguments.iter_mut() {
-                        self.rewrite_expr(&ctx, u, arg);
-                    }
-                    false
+
+                // We apply the rules on this expression, but its
+                // children will be visited after this call.
+                let cont = self.apply_rules(ctx, u, expr);
+
+                (cont, entered)
+            },
+            &mut |expr, entered| {
+                // A `let` variable becomes visible *after* we have have processed all its children,
+                // so, only its siblings can see it.
+                if let Expression::Let(let_) = expr {
+                    let typ = let_.expression.return_type().expect("let should have a type");
+                    self.vars.borrow_mut().locals.add(
+                        let_.id,
+                        let_.mutable,
+                        let_.name.clone(),
+                        typ.into_owned(),
+                    );
+                }
+                if entered {
+                    self.vars.borrow_mut().locals.exit();
+                }
+            },
+            &mut |_| {},
+        );
+    }
+
+    fn apply_rules(
+        &self,
+        ctx: &rules::Context,
+        u: &mut Unstructured,
+        expr: &mut Expression,
+    ) -> bool {
+        match expr {
+            Expression::For(for_) => {
+                // Separate context for just the ranges.
+                let range_ctx = rules::Context { is_in_range: true, ..*ctx };
+                self.rewrite_expr(&range_ctx, u, &mut for_.start_range);
+                self.rewrite_expr(&range_ctx, u, &mut for_.end_range);
+                // Original context for the body.
+                self.vars.borrow_mut().locals.add(
+                    for_.index_variable,
+                    false,
+                    for_.index_name.clone(),
+                    for_.index_type.clone(),
+                );
+                self.rewrite_expr(ctx, u, &mut for_.block);
+                self.vars.borrow_mut().locals.remove(&for_.index_variable);
+                // No need to visit children, we just visited them.
+                false
+            }
+            Expression::Unary(
+                unary @ Unary { operator: UnaryOp::Reference { mutable: true }, .. },
+            ) => {
+                let ctx = rules::Context { is_in_ref_mut: true, ..*ctx };
+                self.rewrite_expr(&ctx, u, &mut unary.rhs);
+                false
+            }
+            Expression::Call(call) if is_special_call(call) => {
+                let ctx = rules::Context { is_in_special_call: true, ..*ctx };
+                for arg in call.arguments.iter_mut() {
+                    self.rewrite_expr(&ctx, u, arg);
                 }
-                // The rest can just have the rules applied on them, using the same context.
-                _ => {
-                    for rule in self.rules {
-                        match self.try_apply_rule(ctx, u, expr, rule) {
-                            Ok(false) => {
-                                // We couldn't, or decided not to apply this rule; try the next one.
-                                continue;
-                            }
-                            Err(_) => {
-                                // We ran out of randomness; stop visiting the AST.
-                                return false;
-                            }
-
-                            Ok(true) => {
-                                // We applied a rule on this expression.
-                                self.count += 1;
-                                // We could visit the children of this morphed expression, which could result in repeatedly applying
-                                // the same rule over and over again. When we have 100% application rate (e.g. a small function),
-                                // it would be wasting all the potential on the first rule that matched, e.g. `(x - (0 + (0 - 0)))`.
-                                // It would also throw off the estimate if we introduce new items on which we can apply rules.
-                                return false;
-                            }
+                false
+            }
+            // The rest can just have the rules applied on them, using the same context.
+            _ => {
+                for rule in self.rules {
+                    match self.try_apply_rule(ctx, u, expr, rule) {
+                        Ok(false) => {
+                            // We couldn't, or decided not to apply this rule; try the next one.
+                            continue;
+                        }
+                        Err(_) => {
+                            // We ran out of randomness; stop visiting the AST.
+                            return false;
+                        }
+
+                        Ok(true) => {
+                            // We applied a rule on this expression.
+                            self.count.set(self.count.get() + 1);
+                            // We could visit the children of this morphed expression, which could result in repeatedly applying
+                            // the same rule over and over again. When we have 100% application rate (e.g. a small function),
+                            // it would be wasting all the potential on the first rule that matched, e.g. `(x - (0 + (0 - 0)))`.
+                            // It would also throw off the estimate if we introduce new items on which we can apply rules.
+                            return false;
                         }
                     }
-                    // If we made it this far, we did not apply any rule, so look deeper in the AST.
-                    true
                 }
+                // If we made it this far, we did not apply any rule, so look deeper in the AST.
+                true
             }
-        });
+        }
     }
 
     /// Check if a rule can be applied on an expression. If it can, apply it based on some arbitrary
     /// criteria, returning a flag showing whether it was applied.
     fn try_apply_rule(
-        &mut self,
+        &self,
         ctx: &rules::Context,
         u: &mut Unstructured,
         expr: &mut Expression,
         rule: &rules::Rule,
     ) -> arbitrary::Result<bool> {
         if rule.matches(ctx, expr) && u.ratio(self.target, self.estimate)? {
-            rule.rewrite(u, &mut self.locals, expr)?;
+            rule.rewrite(u, &mut self.vars.borrow_mut(), expr)?;
             Ok(true)
         } else {
             Ok(false)
@@ -231,14 +297,14 @@ fn is_special_call(call: &Call) -> bool {
 
 /// Metamorphic transformation rules.
 mod rules {
-    use crate::targets::orig_vs_morph::{LocalContext, helpers::reassign_ids};
+    use crate::targets::orig_vs_morph::{VariableContext, helpers::reassign_ids};
 
     use super::helpers::{gen_expr, has_side_effect};
     use arbitrary::Unstructured;
     use noir_ast_fuzzer::expr;
     use noirc_frontend::{
         ast::BinaryOpKind,
-        monomorphization::ast::{Binary, Expression, Literal, Type},
+        monomorphization::ast::{Binary, Definition, Expression, Ident, Literal, Type},
         signed_field::SignedField,
     };
 
@@ -258,7 +324,7 @@ mod rules {
     type MatchFn = dyn Fn(&Context, &Expression) -> bool;
     /// Apply the rule on an expression, mutating/replacing it in-place.
     type RewriteFn =
-        dyn Fn(&mut Unstructured, &mut LocalContext, &mut Expression) -> arbitrary::Result<()>;
+        dyn Fn(&mut Unstructured, &mut VariableContext, &mut Expression) -> arbitrary::Result<()>;
 
     /// Metamorphic transformation rule.
     pub struct Rule {
@@ -271,7 +337,7 @@ mod rules {
             matches: impl Fn(&Context, &Expression) -> bool + 'static,
             rewrite: impl Fn(
                 &mut Unstructured,
-                &mut LocalContext,
+                &mut VariableContext,
                 &mut Expression,
             ) -> arbitrary::Result<()>
             + 'static,
@@ -288,10 +354,10 @@ mod rules {
         pub fn rewrite(
             &self,
             u: &mut Unstructured,
-            locals: &mut LocalContext,
+            vars: &mut VariableContext,
             expr: &mut Expression,
         ) -> arbitrary::Result<()> {
-            (self.rewrite)(u, locals, expr)
+            (self.rewrite)(u, vars, expr)
         }
     }
 
@@ -452,13 +518,37 @@ mod rules {
                     // `let x = 1;` transformed into `if true { let x = 1; } else { let x = 1; }` would leave `x` undefined.
                     && !matches!(expr, Expression::Let(_))
             },
-            |u, locals, expr| {
+            |u, vars, expr| {
                 let typ = expr.return_type().map(|typ| typ.into_owned()).unwrap_or(Type::Unit);
-                let cond = gen_expr(u, &Type::Bool, 2)?;
+
+                // Find a bool expression we can use. For simplicity just consider actual bool variables,
+                // not things that can produce variables, so we have less logic to repeat for the `FunctionContext`.
+                let bool_vars = vars
+                    .locals
+                    .current()
+                    .variables()
+                    .filter_map(|(id, (_, _, typ))| (*typ == Type::Bool).then_some(id))
+                    .collect::<Vec<_>>();
+
+                // If we don't have a bool variable, generate some random expression.
+                let cond = if bool_vars.is_empty() {
+                    gen_expr(u, &Type::Bool, 2)?
+                } else {
+                    let id = u.choose_iter(bool_vars)?;
+                    let (mutable, name, typ) = vars.locals.current().get_variable(id);
+                    Expression::Ident(Ident {
+                        location: None,
+                        definition: Definition::Local(*id),
+                        mutable: *mutable,
+                        name: name.clone(),
+                        typ: typ.clone(),
+                        id: vars.next_ident_id(),
+                    })
+                };
 
                 // Duplicate the expression, then assign new IDs to all variables created in it.
                 let mut alt = expr.clone();
-                reassign_ids(locals, &mut alt);
+                reassign_ids(vars, &mut alt);
 
                 expr::replace(expr, |expr| expr::if_else(cond, expr, alt, typ));
                 Ok(())
@@ -529,7 +619,7 @@ mod helpers {
     };
     use strum::IntoEnumIterator;
 
-    use crate::targets::orig_vs_morph::LocalContext;
+    use crate::targets::orig_vs_morph::VariableContext;
 
     /// Check if an expression can have a side effect, in which case duplicating or reordering it could
     /// change the behavior of the program.
@@ -666,14 +756,14 @@ mod helpers {
     /// Types we can consider using in this context.
     static TYPES: OnceLock<Vec<Type>> = OnceLock::new();
 
-    pub(super) fn reassign_ids(locals: &mut LocalContext, expr: &mut Expression) {
+    pub(super) fn reassign_ids(vars: &mut VariableContext, expr: &mut Expression) {
         fn replace_local_id(
-            locals: &mut LocalContext,
+            vars: &mut VariableContext,
             replacements: &mut HashMap<LocalId, LocalId>,
             id: &mut LocalId,
         ) {
             let curr = *id;
-            let next = locals.next_local_id();
+            let next = vars.next_local_id();
             replacements.insert(curr, next);
             *id = next;
         }
@@ -685,15 +775,15 @@ mod helpers {
             &mut |expr| {
                 match expr {
                     Expression::Let(let_) => {
-                        replace_local_id(locals, &mut replacements.borrow_mut(), &mut let_.id)
+                        replace_local_id(vars, &mut replacements.borrow_mut(), &mut let_.id)
                     }
                     Expression::For(for_) => replace_local_id(
-                        locals,
+                        vars,
                         &mut replacements.borrow_mut(),
                         &mut for_.index_variable,
                     ),
                     Expression::Ident(ident) => {
-                        ident.id = locals.next_ident_id();
+                        ident.id = vars.next_ident_id();
                     }
                     _ => (),
                 }
diff --git a/tooling/ast_fuzzer/src/lib.rs b/tooling/ast_fuzzer/src/lib.rs
index 20596e57d90..6e08054f88f 100644
--- a/tooling/ast_fuzzer/src/lib.rs
+++ b/tooling/ast_fuzzer/src/lib.rs
@@ -15,7 +15,7 @@ pub use compare::input_values_to_ssa;
 pub use input::arb_inputs;
 use program::freq::Freqs;
 pub use program::{DisplayAstAsNoir, DisplayAstAsNoirComptime, arb_program, arb_program_comptime};
-pub use program::{expr, rewrite, types, visitor};
+pub use program::{expr, rewrite, scope, types, visitor};
 
 /// AST generation configuration.
 #[derive(Debug, Clone)]
diff --git a/tooling/ast_fuzzer/src/program/mod.rs b/tooling/ast_fuzzer/src/program/mod.rs
index aa0c2bbe3d1..d74320331f5 100644
--- a/tooling/ast_fuzzer/src/program/mod.rs
+++ b/tooling/ast_fuzzer/src/program/mod.rs
@@ -20,7 +20,7 @@ pub mod expr;
 pub(crate) mod freq;
 mod func;
 pub mod rewrite;
-mod scope;
+pub mod scope;
 pub mod types;
 pub mod visitor;
 
diff --git a/tooling/ast_fuzzer/src/program/scope.rs b/tooling/ast_fuzzer/src/program/scope.rs
index 30fd137e0d9..11193ab6d8c 100644
--- a/tooling/ast_fuzzer/src/program/scope.rs
+++ b/tooling/ast_fuzzer/src/program/scope.rs
@@ -8,7 +8,7 @@ pub(crate) type Variable = (/*mutable*/ bool, Name, Type);
 
 /// A layer of variables available to choose from in blocks.
 #[derive(Debug, Clone)]
-pub(crate) struct Scope<K: Ord> {
+pub struct Scope<K: Ord> {
     /// ID and type of variables created in all visible scopes,
     /// which includes this scope and its ancestors.
     variables: im::OrdMap<K, Variable>,
@@ -23,7 +23,7 @@ where
     K: Ord + Clone + Copy + Debug,
 {
     /// Create the initial scope from function parameters.
-    pub(super) fn new(vars: impl Iterator<Item = (K, bool, Name, Type)>) -> Self {
+    pub fn new(vars: impl Iterator<Item = (K, bool, Name, Type)>) -> Self {
         let mut scope = Self { variables: im::OrdMap::new(), producers: im::OrdMap::new() };
         for (id, mutable, name, typ) in vars {
             scope.add(id, mutable, name, typ);
@@ -72,7 +72,7 @@ where
     }
 
     /// Choose a random producer of a type, if there is one.
-    pub(super) fn choose_producer(
+    pub fn choose_producer(
         &self,
         u: &mut Unstructured,
         typ: &Type,
@@ -87,7 +87,7 @@ where
     }
 
     /// Get a variable in scope.
-    pub(super) fn get_variable(&self, id: &K) -> &Variable {
+    pub fn get_variable(&self, id: &K) -> &Variable {
         self.variables.get(id).unwrap_or_else(|| panic!("variable doesn't exist: {:?}", id))
     }
 }
@@ -97,62 +97,62 @@ where
     K: Ord,
 {
     /// Check if there are any variables in scope.
-    pub(super) fn is_empty(&self) -> bool {
+    pub fn is_empty(&self) -> bool {
         self.variables.is_empty()
     }
 
     /// Iterate the variables in scope.
-    pub(super) fn variables(&self) -> impl ExactSizeIterator<Item = (&K, &Variable)> {
+    pub fn variables(&self) -> impl ExactSizeIterator<Item = (&K, &Variable)> {
         self.variables.iter()
     }
 
     /// Iterate the IDs of the variables in scope.
-    pub(super) fn variable_ids(&self) -> impl ExactSizeIterator<Item = &K> {
+    pub fn variable_ids(&self) -> impl ExactSizeIterator<Item = &K> {
         self.variables.keys()
     }
 
     /// Iterate the types we can produce from other variables.
-    pub(super) fn types_produced(&self) -> impl ExactSizeIterator<Item = &Type> {
+    pub fn types_produced(&self) -> impl ExactSizeIterator<Item = &Type> {
         self.producers.keys()
     }
 }
 
 /// Scope stack as we exit and enter blocks
-pub(crate) struct ScopeStack<K: Ord>(Vec<Scope<K>>);
+pub struct ScopeStack<K: Ord>(Vec<Scope<K>>);
 
 impl<K> ScopeStack<K>
 where
     K: Ord + Clone + Copy + Debug,
 {
     /// Create a stack from the base variables.
-    pub(super) fn new(vars: impl Iterator<Item = (K, bool, Name, Type)>) -> Self {
+    pub fn new(vars: impl Iterator<Item = (K, bool, Name, Type)>) -> Self {
         Self(vec![Scope::new(vars)])
     }
 
     /// The top scope in the stack.
-    pub(super) fn current(&self) -> &Scope<K> {
+    pub fn current(&self) -> &Scope<K> {
         self.0.last().expect("there is always the base layer")
     }
 
     /// Push a new scope on top of the current one.
-    pub(super) fn enter(&mut self) {
+    pub fn enter(&mut self) {
         // Instead of shallow cloning an immutable map, we could loop through layers when looking up variables.
         self.0.push(self.current().clone());
     }
 
     /// Remove the last layer of block variables.
-    pub(super) fn exit(&mut self) {
+    pub fn exit(&mut self) {
         self.0.pop();
         assert!(!self.0.is_empty(), "never pop the base layer");
     }
 
     /// Add a new variable to the current scope.
-    pub(super) fn add(&mut self, id: K, mutable: bool, name: String, typ: Type) {
+    pub fn add(&mut self, id: K, mutable: bool, name: String, typ: Type) {
         self.0.last_mut().expect("there is always a layer").add(id, mutable, name, typ);
     }
 
     /// Remove a variable from all scopes.
-    pub(super) fn remove(&mut self, id: &K) {
+    pub fn remove(&mut self, id: &K) {
         for scope in self.0.iter_mut() {
             scope.remove(id);
         }

From e880ad1f93f65f465450bf45cac4628eae681eb0 Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Thu, 12 Jun 2025 16:21:24 +0100
Subject: [PATCH 10/10] Fix clippy

---
 tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
index 82c4c367401..86690c66748 100644
--- a/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
+++ b/tooling/ast_fuzzer/fuzz/src/targets/orig_vs_morph.rs
@@ -438,7 +438,7 @@ mod rules {
                     (BinaryOpKind::Subtract, (*b - *a))
                 };
 
-                let c_expr = Expression::Literal(Literal::Integer(c, typ.clone(), loc.clone()));
+                let c_expr = Expression::Literal(Literal::Integer(c, typ.clone(), *loc));
 
                 *expr = expr::binary(b_expr, op, c_expr);