fix: Make Ssa-gen use existing reference when compiling &mut foo.bar.baz

compiler/noirc_evaluator/src/ssa/ssa_gen/mod.rs

@@ -372,7 +372,12 @@ impl FunctionContext<'_> {
                 ))
             }
             UnaryOp::Reference { mutable: _ } => {
-                Ok(self.codegen_reference(&unary.rhs)?.map(|rhs| {
+                let rhs = self.codegen_reference(&unary.rhs)?;
+                // If skip is set then `rhs` is a member access expression which is already a reference
+                if unary.skip {
+                    return Ok(rhs);
+                }
+                Ok(rhs.map(|rhs| {
                     match rhs {
                         value::Value::Normal(value) => {
                             let rhs_type = self.builder.current_function.dfg.type_of_value(value);

compiler/noirc_frontend/src/elaborator/expressions.rs

@@ -67,7 +67,8 @@ impl Elaborator<'_> {
             ExpressionKind::Constrain(constrain) => self.elaborate_constrain(constrain),
             ExpressionKind::Constructor(constructor) => self.elaborate_constructor(*constructor),
             ExpressionKind::MemberAccess(access) => {
-                return self.elaborate_member_access(*access, expr.location);
+                let (expr_id, typ, _) = self.elaborate_member_access(*access, expr.location, false);
+                return (expr_id, typ);
             }
             ExpressionKind::Cast(cast) => self.elaborate_cast(*cast, expr.location),
             ExpressionKind::Infix(infix) => return self.elaborate_infix(*infix, expr.location),
@@ -118,6 +119,23 @@ impl Elaborator<'_> {
         (id, typ)
     }
 
+    /// Elaborate a member access expression without adding the automatic dereferencing operations
+    /// to it, treating it as an offset instead. This also returns a boolean indicating whether the
+    /// result skipped any required auto-derefs (and thus needs dereferencing to be used as a value
+    /// instead of a reference). This flag is used when `&mut foo.bar.baz` is used to cancel out
+    /// the `&mut`.
+    fn elaborate_reference_expression(&mut self, expr: Expression) -> (ExprId, Type, bool) {
+        match expr.kind {
+            ExpressionKind::MemberAccess(access) => {
+                self.elaborate_member_access(*access, expr.location, true)
+            }
+            _ => {
+                let (expr_id, typ) = self.elaborate_expression(expr);
+                (expr_id, typ, false)
+            }
+        }
+    }
+
     fn elaborate_interned_statement_as_expr(
         &mut self,
         id: InternedStatementKind,
@@ -403,40 +421,60 @@ impl Elaborator<'_> {
 
     fn elaborate_prefix(&mut self, prefix: PrefixExpression, location: Location) -> (ExprId, Type) {
         let rhs_location = prefix.rhs.location;
+        let operator = prefix.operator;
 
-        let (rhs, rhs_type) = self.elaborate_expression(prefix.rhs);
-        let trait_method_id = self.interner.get_prefix_operator_trait_method(&prefix.operator);
+        let (rhs, rhs_type, skip_op) = if matches!(operator, UnaryOp::Reference { .. }) {
+            let (rhs, rhs_type, needs_deref) = self.elaborate_reference_expression(prefix.rhs);
 
-        let operator = prefix.operator;
+            // If the reference expression delayed a needed deref, we can skip the `&mut _`
+            // operation since the expression is already a reference.
+            (rhs, rhs_type, needs_deref)
+        } else {
+            let (rhs, rhs_type) = self.elaborate_expression(prefix.rhs);
+            (rhs, rhs_type, false)
+        };
+
+        let trait_method_id = self.interner.get_prefix_operator_trait_method(&operator);
 
         if let UnaryOp::Reference { mutable } = operator {
             if mutable {
-                self.check_can_mutate(rhs, rhs_location);
+                // If skip_op is set we already know we have a mutable reference
+                if !skip_op {
+                    self.check_can_mutate(rhs, rhs_location);
+                }
             } else {
                 self.use_unstable_feature(UnstableFeature::Ownership, location);
             }
         }
 
-        let expr = HirExpression::Prefix(HirPrefixExpression { operator, rhs, trait_method_id });
+        let expr = HirExpression::Prefix(HirPrefixExpression {
+            operator,
+            rhs,
+            trait_method_id,
+            skip: skip_op,
+        });
         let expr_id = self.interner.push_expr(expr);
         self.interner.push_expr_location(expr_id, location);
 
-        let result = self.prefix_operand_type_rules(&operator, &rhs_type, location);
-        let typ = self.handle_operand_type_rules_result(
-            result,
-            &rhs_type,
-            trait_method_id,
-            expr_id,
-            location,
-        );
+        let typ = if skip_op {
+            rhs_type
+        } else {
+            let result = self.prefix_operand_type_rules(&operator, &rhs_type, location);
+            self.handle_operand_type_rules_result(
+                result,
+                &rhs_type,
+                trait_method_id,
+                expr_id,
+                location,
+            )
+        };
 
         self.interner.push_expr_type(expr_id, typ.clone());
         (expr_id, typ)
     }
 
     pub(super) fn check_can_mutate(&mut self, expr_id: ExprId, location: Location) {
-        let expr = self.interner.expression(&expr_id);
-        match expr {
+        match self.interner.expression(&expr_id) {
             HirExpression::Ident(hir_ident, _) => {
                 if let Some(definition) = self.interner.try_definition(hir_ident.id) {
                     let name = definition.name.clone();
@@ -1014,20 +1052,33 @@ impl Elaborator<'_> {
         ret
     }
 
+    /// This method also returns whether or not its lhs still needs to be dereferenced depending on
+    /// `is_offset`:
+    /// - `is_offset = false`: Auto-dereferencing will occur, and this will always return false
+    /// - `is_offset = true`: Auto-dereferencing is disabled, and this will return true if the lhs
+    ///   is a reference.
     fn elaborate_member_access(
         &mut self,
         access: MemberAccessExpression,
         location: Location,
-    ) -> (ExprId, Type) {
-        let (lhs, lhs_type) = self.elaborate_expression(access.lhs);
+        is_offset: bool,
+    ) -> (ExprId, Type, bool) {
+        // We don't need the boolean 'skipped auto-derefs' from elaborate_reference_expression
+        // since if we have skipped any then `lhs_type` will be a reference and we will need to
+        // skip the deref (if is_offset is true) here anyway to extract the field out of the reference.
+        // This is more reliable than using the boolean return value here since the return value
+        // doesn't account for reference variables which we need to account for.
+        let (lhs, lhs_type, _) = self.elaborate_reference_expression(access.lhs);
+        let is_reference = lhs_type.is_ref();
+
         let rhs = access.rhs;
         let rhs_location = rhs.location();
         // `is_offset` is only used when lhs is a reference and we want to return a reference to rhs
-        let access = HirMemberAccess { lhs, rhs, is_offset: false };
+        let access = HirMemberAccess { lhs, rhs, is_offset };
         let expr_id = self.intern_expr(HirExpression::MemberAccess(access.clone()), location);
         let typ = self.type_check_member_access(access, expr_id, lhs_type, rhs_location);
         self.interner.push_expr_type(expr_id, typ.clone());
-        (expr_id, typ)
+        (expr_id, typ, is_offset && is_reference)
     }
 
     pub fn intern_expr(&mut self, expr: HirExpression, location: Location) -> ExprId {

compiler/noirc_frontend/src/elaborator/types.rs

@@ -999,11 +999,10 @@ impl Elaborator<'_> {
         if let Type::Reference(element, _mut) = typ.follow_bindings() {
             let location = self.interner.id_location(object);
 
-            let object = self.interner.push_expr(HirExpression::Prefix(HirPrefixExpression {
-                operator: UnaryOp::Dereference { implicitly_added: true },
-                rhs: object,
-                trait_method_id: None,
-            }));
+            let object = self.interner.push_expr(HirExpression::Prefix(HirPrefixExpression::new(
+                UnaryOp::Dereference { implicitly_added: true },
+                object,
+            )));
             self.interner.push_expr_type(object, element.as_ref().clone());
             self.interner.push_expr_location(object, location);
 
@@ -1591,11 +1590,10 @@ impl Elaborator<'_> {
 
         let dereference_lhs = |this: &mut Self, lhs_type, element| {
             let old_lhs = *access_lhs;
-            *access_lhs = this.interner.push_expr(HirExpression::Prefix(HirPrefixExpression {
-                operator: crate::ast::UnaryOp::Dereference { implicitly_added: true },
-                rhs: old_lhs,
-                trait_method_id: None,
-            }));
+            *access_lhs = this.interner.push_expr(HirExpression::Prefix(HirPrefixExpression::new(
+                crate::ast::UnaryOp::Dereference { implicitly_added: true },
+                old_lhs,
+            )));
             this.interner.push_expr_type(old_lhs, lhs_type);
             this.interner.push_expr_type(*access_lhs, element);
 
@@ -2069,12 +2067,9 @@ impl Elaborator<'_> {
 
                     // If that didn't work, then wrap the whole expression in an `&mut`
                     *object = new_object.unwrap_or_else(|| {
-                        let new_object =
-                            self.interner.push_expr(HirExpression::Prefix(HirPrefixExpression {
-                                operator: UnaryOp::Reference { mutable },
-                                rhs: *object,
-                                trait_method_id: None,
-                            }));
+                        let new_object = self.interner.push_expr(HirExpression::Prefix(
+                            HirPrefixExpression::new(UnaryOp::Reference { mutable }, *object),
+                        ));
                         self.interner.push_expr_type(new_object, new_type);
                         self.interner.push_expr_location(new_object, location);
                         new_object

compiler/noirc_frontend/src/hir/comptime/interpreter.rs

@@ -43,7 +43,7 @@ use crate::{
 };
 
 use super::errors::{IResult, InterpreterError};
-use super::value::{Closure, Value, unwrap_rc};
+use super::value::{Closure, StructFields, Value, unwrap_rc};
 
 mod builtin;
 mod cast;
@@ -826,6 +826,10 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
             _ => self.evaluate(prefix.rhs)?,
         };
 
+        if prefix.skip {
+            return Ok(rhs);
+        }
+
         if self.elaborator.interner.get_selected_impl_for_expression(id).is_some() {
             self.evaluate_overloaded_prefix(prefix, rhs, id)
         } else {
@@ -962,11 +966,11 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
         Ok(Value::Enum(constructor.variant_index, fields, typ))
     }
 
-    fn evaluate_access(&mut self, access: HirMemberAccess, id: ExprId) -> IResult<Value> {
-        let (fields, struct_type) = match self.evaluate(access.lhs)? {
-            Value::Struct(fields, typ) => (fields, typ),
+    fn get_fields(&mut self, value: Value, id: ExprId) -> IResult<(StructFields, Type)> {
+        match value {
+            Value::Struct(fields, typ) => Ok((fields, typ)),
             Value::Tuple(fields) => {
-                let (fields, field_types): (HashMap<Rc<String>, Shared<Value>>, Vec<Type>) = fields
+                let (fields, field_types) = fields
                     .into_iter()
                     .enumerate()
                     .map(|(i, field)| {
@@ -975,14 +979,21 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
                         (key_val_pair, field_type)
                     })
                     .unzip();
-                (fields, Type::Tuple(field_types))
+                Ok((fields, Type::Tuple(field_types)))
             }
+            Value::Pointer(element, ..) => self.get_fields(element.unwrap_or_clone(), id),
             value => {
                 let location = self.elaborator.interner.expr_location(&id);
                 let typ = value.get_type().into_owned();
-                return Err(InterpreterError::NonTupleOrStructInMemberAccess { typ, location });
+                Err(InterpreterError::NonTupleOrStructInMemberAccess { typ, location })
             }
-        };
+        }
+    }
+
+    fn evaluate_access(&mut self, access: HirMemberAccess, id: ExprId) -> IResult<Value> {
+        let lhs = self.evaluate(access.lhs)?;
+        let is_offset = access.is_offset && lhs.get_type().is_ref();
+        let (fields, struct_type) = self.get_fields(lhs, id)?;
 
         let field = fields.get(access.rhs.as_string()).cloned().ok_or_else(|| {
             let location = self.elaborator.interner.expr_location(&id);
@@ -994,8 +1005,12 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
 
         // Return a reference to the field so that `&mut foo.bar.baz` can use this reference.
         // We set auto_deref to true so that when it is used elsewhere it is dereferenced
-        // automatically.
-        Ok(Value::Pointer(field, true, false))
+        // automatically. In some cases in the frontend the leading `&mut` will cancel out
+        // with a field access which is expected to only offset into the struct and thus return
+        // a reference already. In this case we set auto_deref to false because the outer `&mut`
+        // will also be removed in that case so the pointer should be explicit.
+        let auto_deref = !is_offset;
+        Ok(Value::Pointer(field, auto_deref, false))
     }
 
     fn evaluate_call(&mut self, call: HirCallExpression, id: ExprId) -> IResult<Value> {

compiler/noirc_frontend/src/hir_def/expr.rs

@@ -144,6 +144,20 @@ pub struct HirPrefixExpression {
     /// The trait method id for the operator trait method that corresponds to this operator,
     /// if such a trait exists (for example, there's no trait for the dereference operator).
     pub trait_method_id: Option<TraitItemId>,
+
+    /// If this is true we should skip this operation and directly return `rhs` instead.
+    /// This is used for compiling `&mut foo.bar.baz` where `foo.bar.baz` already returns
+    /// a reference and we do not want to create a new reference. Additionally, this node
+    /// is kept so that `nargo expand` still expands into the full `&mut foo.bar.baz` instead
+    /// of removing the leading `&mut`.
+    pub skip: bool,
+}
+
+impl HirPrefixExpression {
+    /// Creates a basic HirPrefixExpression with `trait_method_id = None` and `skip = false`
+    pub fn new(operator: UnaryOp, rhs: ExprId) -> Self {
+        Self { operator, rhs, trait_method_id: None, skip: false }
+    }
 }
 
 #[derive(Debug, Clone)]

compiler/noirc_frontend/src/hir_def/types.rs

@@ -1299,6 +1299,10 @@ impl Type {
         matches!(self.follow_bindings_shallow().as_ref(), Type::Reference(_, true))
     }
 
+    pub(crate) fn is_ref(&self) -> bool {
+        matches!(self.follow_bindings_shallow().as_ref(), Type::Reference(_, _))
+    }
+
     /// 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.

compiler/noirc_frontend/src/monomorphization/ast.rs

@@ -283,6 +283,11 @@ pub struct Unary {
     pub rhs: Box<Expression>,
     pub result_type: Type,
     pub location: Location,
+
+    /// Carried over from `HirPrefixExpression::skip`. This also flags whether we should directly
+    /// return `rhs` and skip performing this operation. Compared to replacing this node with rhs
+    /// directly, keeping this flag keeps `--show-monomorphized` output the same.
+    pub skip: bool,
 }
 
 pub type BinaryOp = BinaryOpKind;

compiler/noirc_frontend/src/monomorphization/mod.rs

@@ -543,7 +543,13 @@ impl<'interner> Monomorphizer<'interner> {
                     let operator = prefix.operator;
                     let rhs = Box::new(rhs);
                     let result_type = Self::convert_type(&self.interner.id_type(expr), location)?;
-                    ast::Expression::Unary(ast::Unary { operator, rhs, result_type, location })
+                    ast::Expression::Unary(ast::Unary {
+                        operator,
+                        rhs,
+                        result_type,
+                        location,
+                        skip: prefix.skip,
+                    })
                 }
             }
 
@@ -2274,6 +2280,7 @@ impl<'interner> Monomorphizer<'interner> {
                     result_type,
                     operator: Reference { mutable: *mutable },
                     location,
+                    skip: false,
                 })
             }
         }
@@ -2363,6 +2370,7 @@ impl<'interner> Monomorphizer<'interner> {
                     rhs: Box::new(result),
                     result_type: ast::Type::Bool,
                     location,
+                    skip: false,
                 });
             }
             // All the comparison operators require special handling since their `cmp` method

test_programs/execution_success/reference_cancelling/Nargo.toml

@@ -0,0 +1,6 @@
+[package]
+name = "reference_cancelling"
+type = "bin"
+authors = [""]
+
+[dependencies]

test_programs/execution_success/reference_cancelling/src/main.nr

@@ -0,0 +1,23 @@
+fn main() {
+    assert(comptime { store_through_outer_pointer() });
+    assert(store_through_outer_pointer());
+
+    assert(comptime { store_through_middle_pointer() });
+    assert(store_through_middle_pointer());
+}
+
+fn store_through_outer_pointer() -> bool {
+    let a = &mut ((false,),);
+    let b = &mut a.0.0;
+    *a = ((true,),);
+    println(*b);
+    *b
+}
+
+fn store_through_middle_pointer() -> bool {
+    let a = &mut ((false,),);
+    let b = &mut a.0.0;
+    a.0 = (true,);
+    println(*b);
+    *b
+}