fix(frontend)!: Restrict capturing mutable variable in lambdas

compiler/noirc_frontend/src/elaborator/expressions.rs

@@ -360,26 +360,50 @@ impl Elaborator<'_> {
         (expr_id, typ)
     }
 
-    fn check_can_mutate(&mut self, expr_id: ExprId, location: Location) {
+    pub(super) fn check_can_mutate(&mut self, expr_id: ExprId, location: Location) {
         let expr = self.interner.expression(&expr_id);
         match expr {
             HirExpression::Ident(hir_ident, _) => {
                 if let Some(definition) = self.interner.try_definition(hir_ident.id) {
+                    let name = definition.name.clone();
                     if !definition.mutable {
                         self.push_err(TypeCheckError::CannotMutateImmutableVariable {
-                            name: definition.name.clone(),
+                            name,
                             location,
                         });
+                    } else {
+                        self.check_can_mutate_lambda_capture(hir_ident.id, name, location);
                     }
                 }
             }
+            HirExpression::Index(_) => {
+                self.push_err(TypeCheckError::MutableReferenceToArrayElement { location });
+            }
             HirExpression::MemberAccess(member_access) => {
                 self.check_can_mutate(member_access.lhs, location);
             }
             _ => (),
         }
     }
 
+    // We must check whether the mutable variable we are attempting to mutate
+    // comes from a lambda capture. All captures are immutable so we want to error
+    // if the user attempts to mutate a captured variable inside of a lambda without mutable references.
+    pub(super) fn check_can_mutate_lambda_capture(
+        &mut self,
+        id: DefinitionId,
+        name: String,
+        location: Location,
+    ) {
+        if let Some(lambda_context) = self.lambda_stack.last() {
+            let typ = self.interner.definition_type(id);
+            if !typ.is_mutable_ref() && lambda_context.captures.iter().any(|var| var.ident.id == id)
+            {
+                self.push_err(TypeCheckError::MutableCaptureWithoutRef { name, location });
+            }
+        }
+    }
+
     fn elaborate_index(&mut self, index_expr: IndexExpression) -> (HirExpression, Type) {
         let location = index_expr.index.location;
 
@@ -1126,6 +1150,7 @@ impl Elaborator<'_> {
         lambda: Lambda,
         parameters_type_hints: Option<&Vec<Type>>,
     ) -> (HirExpression, Type) {
+        // dbg!(lambda.clone());
         self.push_scope();
         let scope_index = self.scopes.current_scope_index();
 

compiler/noirc_frontend/src/elaborator/statements.rs

@@ -157,22 +157,9 @@ impl Elaborator<'_> {
         if !mutable {
             let (_, name, location) = self.get_lvalue_error_info(&lvalue);
             self.push_err(TypeCheckError::VariableMustBeMutable { name, location });
-        } else if let Some(lambda_context) = self.lambda_stack.last() {
-            // We must check whether the mutable variable we are attempting to assign
-            // comes from a lambda capture. All captures are immutable so we want to error
-            // if the user attempts to mutate a captured variable inside of a lambda without mutable references.
-            let capture_ids =
-                lambda_context.captures.iter().map(|var| var.ident.id).collect::<Vec<_>>();
+        } else {
             let (id, name, location) = self.get_lvalue_error_info(&lvalue);
-            let typ = self.interner.definition_type(id);
-            for capture_id in capture_ids {
-                if capture_id == id && !typ.is_mutable_ref() {
-                    self.push_err(TypeCheckError::MutableCaptureWithoutRef {
-                        name: name.clone(),
-                        location,
-                    });
-                }
-            }
+            self.check_can_mutate_lambda_capture(id, name, location);
         }
 
         self.unify_with_coercions(&expr_type, &lvalue_type, expression, expr_location, || {

compiler/noirc_frontend/src/elaborator/types.rs

@@ -31,8 +31,8 @@ use crate::{
         traits::{NamedType, ResolvedTraitBound, Trait, TraitConstraint},
     },
     node_interner::{
-        DependencyId, ExprId, FuncId, GlobalValue, ImplSearchErrorKind, NodeInterner, TraitId,
-        TraitImplKind, TraitMethodId,
+        DependencyId, ExprId, FuncId, GlobalValue, ImplSearchErrorKind, TraitId, TraitImplKind,
+        TraitMethodId,
     },
     signed_field::SignedField,
     token::SecondaryAttribute,
@@ -1839,9 +1839,8 @@ impl Elaborator<'_> {
 
             if matches!(expected_object_type.follow_bindings(), Type::MutableReference(_)) {
                 if !matches!(actual_type, Type::MutableReference(_)) {
-                    if let Err(error) = verify_mutable_reference(self.interner, *object) {
-                        self.push_err(TypeCheckError::ResolverError(error));
-                    }
+                    let location = self.interner.id_location(*object);
+                    self.check_can_mutate(*object, location);
 
                     let new_type = Type::MutableReference(Box::new(actual_type));
                     *object_type = new_type.clone();
@@ -1852,8 +1851,6 @@ impl Elaborator<'_> {
 
                     // If that didn't work, then wrap the whole expression in an `&mut`
                     *object = new_object.unwrap_or_else(|| {
-                        let location = self.interner.id_location(*object);
-
                         let new_object =
                             self.interner.push_expr(HirExpression::Prefix(HirPrefixExpression {
                                 operator: UnaryOp::MutableReference,
@@ -2136,30 +2133,3 @@ fn bind_generic(param: &ResolvedGeneric, arg: &Type, bindings: &mut TypeBindings
         bindings.insert(param.type_var.id(), (param.type_var.clone(), param.kind(), arg.clone()));
     }
 }
-
-/// Gives an error if a user tries to create a mutable reference
-/// to an immutable variable.
-fn verify_mutable_reference(interner: &NodeInterner, rhs: ExprId) -> Result<(), ResolverError> {
-    match interner.expression(&rhs) {
-        HirExpression::MemberAccess(member_access) => {
-            verify_mutable_reference(interner, member_access.lhs)
-        }
-        HirExpression::Index(_) => {
-            let location = interner.expr_location(&rhs);
-            Err(ResolverError::MutableReferenceToArrayElement { location })
-        }
-        HirExpression::Ident(ident, _) => {
-            if let Some(definition) = interner.try_definition(ident.id) {
-                if !definition.mutable {
-                    let location = interner.expr_location(&rhs);
-                    let variable = definition.name.clone();
-                    let err =
-                        ResolverError::MutableReferenceToImmutableVariable { location, variable };
-                    return Err(err);
-                }
-            }
-            Ok(())
-        }
-        _ => Ok(()),
-    }
-}

compiler/noirc_frontend/src/hir/resolution/errors.rs

@@ -76,10 +76,6 @@ pub enum ResolverError {
     GenericsOnAssociatedType { location: Location },
     #[error("{0}")]
     ParserError(Box<ParserError>),
-    #[error("Cannot create a mutable reference to {variable}, it was declared to be immutable")]
-    MutableReferenceToImmutableVariable { variable: String, location: Location },
-    #[error("Mutable references to array indices are unsupported")]
-    MutableReferenceToArrayElement { location: Location },
     #[error("Closure environment must be a tuple or unit type")]
     InvalidClosureEnvironment { typ: Type, location: Location },
     #[error("Nested slices, i.e. slices within an array or slice, are not supported")]
@@ -224,8 +220,6 @@ impl ResolverError {
             | ResolverError::NonStructWithGenerics { location }
             | ResolverError::GenericsOnSelfType { location }
             | ResolverError::GenericsOnAssociatedType { location }
-            | ResolverError::MutableReferenceToImmutableVariable { location, .. }
-            | ResolverError::MutableReferenceToArrayElement { location }
             | ResolverError::InvalidClosureEnvironment { location, .. }
             | ResolverError::NestedSlices { location }
             | ResolverError::AbiAttributeOutsideContract { location }
@@ -490,12 +484,6 @@ impl<'a> From<&'a ResolverError> for Diagnostic {
                 *location,
             ),
             ResolverError::ParserError(error) => error.as_ref().into(),
-            ResolverError::MutableReferenceToImmutableVariable { variable, location } => {
-                Diagnostic::simple_error(format!("Cannot mutably reference the immutable variable {variable}"), format!("{variable} is immutable"), *location)
-            },
-            ResolverError::MutableReferenceToArrayElement { location } => {
-                Diagnostic::simple_error("Mutable references to array elements are currently unsupported".into(), "Try storing the element in a fresh variable first".into(), *location)
-            },
             ResolverError::InvalidClosureEnvironment { location, typ } => Diagnostic::simple_error(
                 format!("{typ} is not a valid closure environment type"),
                 "Closure environment must be a tuple or unit type".to_string(), *location),

compiler/noirc_frontend/src/hir/type_check/errors.rs

@@ -102,6 +102,8 @@ pub enum TypeCheckError {
     CannotMutateImmutableVariable { name: String, location: Location },
     #[error("Variable {name} captured in lambda must be a mutable reference")]
     MutableCaptureWithoutRef { name: String, location: Location },
+    #[error("Mutable references to array indices are unsupported")]
+    MutableReferenceToArrayElement { location: Location },
     #[error("No method named '{method_name}' found for type '{object_type}'")]
     UnresolvedMethodCall { method_name: String, object_type: Type, location: Location },
     #[error("Cannot invoke function field '{method_name}' on type '{object_type}' as a method")]
@@ -281,6 +283,8 @@ impl TypeCheckError {
             | TypeCheckError::TupleIndexOutOfBounds { location, .. }
             | TypeCheckError::VariableMustBeMutable { location, .. }
             | TypeCheckError::CannotMutateImmutableVariable { location, .. }
+            | TypeCheckError::MutableCaptureWithoutRef { location, .. }
+            | TypeCheckError::MutableReferenceToArrayElement { location }
             | TypeCheckError::UnresolvedMethodCall { location, .. }
             | TypeCheckError::CannotInvokeStructFieldFunctionType { location, .. }
             | TypeCheckError::IntegerSignedness { location, .. }
@@ -323,8 +327,7 @@ impl TypeCheckError {
             | TypeCheckError::CyclicType { location, .. }
             | TypeCheckError::TypeAnnotationsNeededForIndex { location }
             | TypeCheckError::UnnecessaryUnsafeBlock { location }
-            | TypeCheckError::NestedUnsafeBlock { location }
-            | TypeCheckError::MutableCaptureWithoutRef { location, .. } => *location,
+            | TypeCheckError::NestedUnsafeBlock { location } => *location,
 
             TypeCheckError::DuplicateNamedTypeArg { name: ident, .. }
             | TypeCheckError::NoSuchNamedTypeArg { name: ident, .. } => ident.location(),
@@ -512,6 +515,9 @@ impl<'a> From<&'a TypeCheckError> for Diagnostic {
                 "Use '&mut' instead of 'mut' to capture a mutable variable.".to_string(),
                 *location,
             ),
+            TypeCheckError::MutableReferenceToArrayElement { location } => {
+                Diagnostic::simple_error("Mutable references to array elements are currently unsupported".into(), "Try storing the element in a fresh variable first".into(), *location)
+            },
             TypeCheckError::PublicReturnType { typ, location } => Diagnostic::simple_error(
                 "Functions cannot declare a public return type".to_string(),
                 format!("return type is {typ}"),

compiler/noirc_frontend/src/tests.rs

@@ -3992,6 +3992,26 @@ fn deny_capturing_mut_variable_without_reference_in_lambda() {
     check_errors(src);
 }
 
+#[test]
+fn deny_capturing_mut_variable_without_reference_in_nested_lambda() {
+    let src = r#"
+    fn main() {
+        let mut x = 3;
+        let f = || {
+            let inner = || {
+                x += 2;
+                ^ Mutable variable x captured in lambda must be a mutable reference
+                ~ Use '&mut' instead of 'mut' to capture a mutable variable.
+            };
+            inner();
+        };
+        f();
+        assert(x == 5);
+    }
+    "#;
+    check_errors(src);
+}
+
 #[test]
 fn allow_capturing_mut_variable_only_used_immutably() {
     let src = r#"
@@ -4004,3 +4024,93 @@ fn allow_capturing_mut_variable_only_used_immutably() {
     "#;
     assert_no_errors(src);
 }
+
+#[test]
+fn deny_capturing_mut_var_as_param_to_function() {
+    let src = r#"
+    fn main() {
+        let mut x = 3;
+        let f = || mutate(&mut x);
+                               ^ Mutable variable x captured in lambda must be a mutable reference
+                               ~ Use '&mut' instead of 'mut' to capture a mutable variable.
+        f();
+        assert(x == 3);
+    }
+
+    fn mutate(x: &mut Field) {
+        *x = 5;
+    }
+    "#;
+    check_errors(src);
+}
+
+#[test]
+fn deny_capturing_mut_var_as_param_to_function_in_nested_lambda() {
+    let src = r#"
+    fn main() {
+        let mut x = 3;
+        let f = || { 
+            let inner = || mutate(&mut x); 
+                                       ^ Mutable variable x captured in lambda must be a mutable reference
+                                       ~ Use '&mut' instead of 'mut' to capture a mutable variable.
+            inner();
+        };
+        f();
+        assert(x == 3);
+    }
+
+    fn mutate(x: &mut Field) {
+        *x = 5;
+    }
+    "#;
+    check_errors(src);
+}
+
+#[test]
+fn deny_capturing_mut_var_as_param_to_impl_method() {
+    let src = r#"
+    struct Foo {
+        value: Field,
+    }
+
+    impl Foo {
+        fn mutate(&mut self) {
+            self.value = 2;
+        }
+    }
+
+    fn main() {
+        let mut foo = Foo { value: 1 };
+        let f = || foo.mutate();
+                   ^^^ Mutable variable foo captured in lambda must be a mutable reference
+                   ~~~ Use '&mut' instead of 'mut' to capture a mutable variable.
+        f();
+        assert(foo.value == 2);
+    }
+    "#;
+    check_errors(src);
+}
+
+#[test]
+fn deny_attaching_mut_ref_to_immutable_object() {
+    let src = r#"
+    struct Foo {
+        value: Field,
+    }
+
+    impl Foo {
+        fn mutate(&mut self) {
+            self.value = 2;
+        }
+    }
+
+    fn main() {
+        let foo = Foo { value: 1 };
+        let f = || foo.mutate();
+                   ^^^ Cannot mutate immutable variable `foo`
+        f();
+        assert(foo.value == 2);
+    }
+    "#;
+    check_errors(src);
+}