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;
 

compiler/noirc_frontend/src/elaborator/statements.rs

@@ -155,8 +155,11 @@ impl Elaborator<'_> {
         let (lvalue, lvalue_type, mutable) = self.elaborate_lvalue(assign.lvalue);
 
         if !mutable {
-            let (name, location) = self.get_lvalue_name_and_location(&lvalue);
+            let (_, name, location) = self.get_lvalue_error_info(&lvalue);
             self.push_err(TypeCheckError::VariableMustBeMutable { name, location });
+        } else {
+            let (id, name, location) = self.get_lvalue_error_info(&lvalue);
+            self.check_can_mutate_lambda_capture(id, name, location);
         }
 
         self.unify_with_coercions(&expr_type, &lvalue_type, expression, expr_location, || {
@@ -334,20 +337,20 @@ impl Elaborator<'_> {
         (expr, self.interner.next_type_variable())
     }
 
-    fn get_lvalue_name_and_location(&self, lvalue: &HirLValue) -> (String, Location) {
+    fn get_lvalue_error_info(&self, lvalue: &HirLValue) -> (DefinitionId, String, Location) {
         match lvalue {
             HirLValue::Ident(name, _) => {
                 let location = name.location;
 
                 if let Some(definition) = self.interner.try_definition(name.id) {
-                    (definition.name.clone(), location)
+                    (name.id, definition.name.clone(), location)
                 } else {
-                    ("(undeclared variable)".into(), location)
+                    (DefinitionId::dummy_id(), "(undeclared variable)".into(), location)
                 }
             }
-            HirLValue::MemberAccess { object, .. } => self.get_lvalue_name_and_location(object),
-            HirLValue::Index { array, .. } => self.get_lvalue_name_and_location(array),
-            HirLValue::Dereference { lvalue, .. } => self.get_lvalue_name_and_location(lvalue),
+            HirLValue::MemberAccess { object, .. } => self.get_lvalue_error_info(object),
+            HirLValue::Index { array, .. } => self.get_lvalue_error_info(array),
+            HirLValue::Dereference { lvalue, .. } => self.get_lvalue_error_info(lvalue),
         }
     }
 
@@ -449,8 +452,8 @@ impl Elaborator<'_> {
                     Type::Slice(elem_type) => *elem_type,
                     Type::Error => Type::Error,
                     Type::String(_) => {
-                        let (_lvalue_name, lvalue_location) =
-                            self.get_lvalue_name_and_location(&lvalue);
+                        let (_id, _lvalue_name, lvalue_location) =
+                            self.get_lvalue_error_info(&lvalue);
                         self.push_err(TypeCheckError::StringIndexAssign {
                             location: lvalue_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

@@ -100,6 +100,10 @@ pub enum TypeCheckError {
     VariableMustBeMutable { name: String, location: Location },
     #[error("Cannot mutate immutable variable `{name}`")]
     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")]
@@ -279,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, .. }
@@ -322,8 +328,10 @@ impl TypeCheckError {
             | TypeCheckError::TypeAnnotationsNeededForIndex { location }
             | TypeCheckError::UnnecessaryUnsafeBlock { location }
             | TypeCheckError::NestedUnsafeBlock { location } => *location,
+
             TypeCheckError::DuplicateNamedTypeArg { name: ident, .. }
             | TypeCheckError::NoSuchNamedTypeArg { name: ident, .. } => ident.location(),
+
             TypeCheckError::NoMatchingImplFound(no_matching_impl_found_error) => {
                 no_matching_impl_found_error.location
             }
@@ -502,6 +510,14 @@ impl<'a> From<&'a TypeCheckError> for Diagnostic {
             | TypeCheckError::InvalidShiftSize { location } => {
                 Diagnostic::simple_error(error.to_string(), String::new(), *location)
             }
+            TypeCheckError::MutableCaptureWithoutRef { name, location } => Diagnostic::simple_error(
+                format!("Mutable variable {name} captured in lambda must be a mutable reference"),
+                "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/hir_def/types.rs

@@ -1242,6 +1242,10 @@ impl Type {
         }
     }
 
+    pub(crate) fn is_mutable_ref(&self) -> bool {
+        matches!(self.follow_bindings_shallow().as_ref(), Type::MutableReference(_))
+    }
+
     /// True if this type can be used as a parameter to `main` or a contract function.
     /// This is only false for unsized types like slices or slices that do not make sense
     /// as a program input such as named generics or mutable references.