fix: properly propagate constant value to assigned var in mast phase

examples/generic_for_loop.no

@@ -11,6 +11,15 @@ fn join(const LEN: Field, arr1: [Field; LLEN], arr2: [Field; RLEN]) -> [Field; L
     return arr;
 }
 
+fn accumulate_mut(const LEN: Field) -> Field {
+    // shouldn't fold mutable variables as constants
+    let mut zz = LEN;
+    for ii in 0..3 {
+        zz = zz + zz;
+    }
+    return zz;
+}
+
 fn main(pub xx: Field) {
     let arr1 = [xx + 1, xx + 2];
     let arr2 = [xx + 3, xx + 4];
@@ -21,4 +30,7 @@ fn main(pub xx: Field) {
     assert_eq(arr[1], arr1[1]);
     assert_eq(arr[2], arr2[0]);
     assert_eq(arr[3], arr2[1]);
+
+    let res = accumulate_mut(1);
+    assert_eq(res, 8);
 }

src/mast/mod.rs

@@ -563,15 +563,33 @@ fn monomorphize_expr<B: Backend>(
                 | Op2::Multiplication
                 | Op2::Division
                 | Op2::BoolAnd
-                | Op2::BoolOr => lhs_mono.typ,
+                | Op2::BoolOr => lhs_mono.clone().typ,
             };
 
-            let cst = match (lhs_mono.constant, rhs_mono.constant) {
-                (Some(lhs), Some(rhs)) => match op {
-                    Op2::Addition => Some(lhs + rhs),
-                    Op2::Subtraction => Some(lhs - rhs),
-                    Op2::Multiplication => Some(lhs * rhs),
-                    Op2::Division => Some(lhs / rhs),
+            let ExprMonoInfo {
+                typ: lhs_typ,
+                constant: lhs_cst,
+                ..
+            } = lhs_mono;
+            let ExprMonoInfo {
+                typ: rhs_typ,
+                constant: rhs_cst,
+                ..
+            } = rhs_mono;
+
+            // check if the constant values can be folded
+            let cst = match (lhs_typ, rhs_typ) {
+                (
+                    Some(TyKind::Field { constant: true }),
+                    Some(TyKind::Field { constant: true }),
+                ) => match (lhs_cst, rhs_cst) {
+                    (Some(lhs), Some(rhs)) => match op {
+                        Op2::Addition => Some(lhs + rhs),
+                        Op2::Subtraction => Some(lhs - rhs),
+                        Op2::Multiplication => Some(lhs * rhs),
+                        Op2::Division => Some(lhs / rhs),
+                        _ => None,
+                    },
                     _ => None,
                 },
                 _ => None,
@@ -895,7 +913,14 @@ pub fn monomorphize_stmt<B: Backend>(
         StmtKind::Assign { mutable, lhs, rhs } => {
             let rhs_mono = monomorphize_expr(ctx, rhs, mono_fn_env)?;
             let typ = rhs_mono.typ.as_ref().expect("expected a type");
-            let type_info = MTypeInfo::new(typ, lhs.span, None);
+
+            let type_info = if *mutable && matches!(typ, TyKind::Field { constant: true }) {
+                // mutable variable shouldn't be constant
+                let new_typ = &TyKind::Field { constant: false };
+                MTypeInfo::new(new_typ, lhs.span, None)
+            } else {
+                MTypeInfo::new(typ, lhs.span, rhs_mono.constant)
+            };
 
             // store the type of lhs in the env
             mono_fn_env.store_type(&lhs.value, &type_info)?;