feat(ownership): Do not clone indexed call results containing arrays

compiler/noirc_frontend/src/ownership/mod.rs

@@ -131,32 +131,44 @@ impl Context {
     ///
     /// Note that this also matches on dereference operations to exempt their LHS from clones,
     /// but their LHS is always exempt from clones so this is unchanged.
-    fn handle_reference_expression(&mut self, expr: &mut Expression) {
+    ///
+    /// # Returns
+    /// A boolean representing whether or not the expression was borrowed by reference or moved.
+    /// This helps the caller determine whether we need to clone the result of a borrowed expression.
+    /// `true` represents an expression that was moved
+    /// `false` represents an expression borrowed by reference
+    fn handle_reference_expression(&mut self, expr: &mut Expression) -> bool {
         match expr {
-            Expression::Ident(_) => (),
+            Expression::Ident(_) => false,
             Expression::Block(exprs) => {
                 let len_minus_one = exprs.len().saturating_sub(1);
                 for expr in exprs.iter_mut().take(len_minus_one) {
                     // In `&{ a; b; ...; z }` we're only taking the reference of `z`.
                     self.handle_expression(expr);
                 }
                 if let Some(expr) = exprs.last_mut() {
-                    self.handle_reference_expression(expr);
+                    self.handle_reference_expression(expr)
+                } else {
+                    true
                 }
             }
             Expression::Unary(Unary { rhs, operator: UnaryOp::Dereference { .. }, .. }) => {
-                self.handle_reference_expression(rhs);
+                self.handle_reference_expression(rhs)
             }
             Expression::ExtractTupleField(tuple, _index) => self.handle_reference_expression(tuple),
 
             Expression::Index(index) => {
-                self.handle_reference_expression(&mut index.collection);
+                let is_moved = self.handle_reference_expression(&mut index.collection);
                 self.handle_expression(&mut index.index);
+                is_moved
             }
 
             // If we have something like `f(arg)` then we want to treat those variables normally
             // rather than avoid cloning them. So we shouldn't recur in `handle_reference_expression`.
-            other => self.handle_expression(other),
+            other => {
+                self.handle_expression(other);
+                true
+            }
         }
     }
 
@@ -277,10 +289,11 @@ impl Context {
         };
 
         // Don't clone the collection, cloning only the resulting element is cheaper.
-        self.handle_reference_expression(&mut index.collection);
+        let is_moved = self.handle_reference_expression(&mut index.collection);
         self.handle_expression(&mut index.index);
 
-        if contains_array_or_str_type(&index.element_type) {
+        // If the index collection is being borrowed we need to clone the result.
+        if !is_moved && contains_array_or_str_type(&index.element_type) {
             clone_expr(index_expr);
         }
     }

compiler/noirc_frontend/src/ownership/tests.rs

@@ -113,3 +113,61 @@ fn can_move_within_loop() {
     }
     ");
 }
+
+#[test]
+fn borrows_on_nested_index() {
+    let src = "
+    unconstrained fn main(x: Field, y: pub Field) {
+        let EXPONENTIATE: [[[Field; 2]; 2]; 2] = [[[1, 1], [0, 0]], [[1, 1], [0, 0]]];
+        let mut acc: Field = 0;
+        for i in 0..2 {
+            for j in 0..2 {
+                acc += EXPONENTIATE[i][j][i];
+            }
+        }
+        assert(acc != 0);
+    }
+    ";
+
+    let program = get_monomorphized_no_emit_test(src).unwrap();
+    // We expect no clones
+    insta::assert_snapshot!(program, @r"
+    unconstrained fn main$f0(x$l0: Field, y$l1: pub Field) -> () {
+        let EXPONENTIATE$l2 = [[[1, 1], [0, 0]], [[1, 1], [0, 0]]];
+        let mut acc$l3 = 0;
+        for i$l4 in 0 .. 2 {
+            for j$l5 in 0 .. 2 {
+                acc$l3 = (acc$l3 + EXPONENTIATE$l2[i$l4][j$l5][i$l4])
+            }
+        };
+        assert((acc$l3 != 0));
+    }
+    ");
+}
+
+#[test]
+fn moves_call_array_result() {
+    let src = "
+    unconstrained fn main(i: u32) -> pub u32 {
+        let _a = foo()[1][0][1];
+        let _s = foo()[1][0][1];
+        i
+    }
+    unconstrained fn foo() -> [[[[u128; 0]; 2]; 1]; 2] {
+        [[[[], []]], [[[], []]]]
+    }
+    ";
+
+    let program = get_monomorphized_no_emit_test(src).unwrap();
+    // We expect no clones
+    insta::assert_snapshot!(program, @r"
+    unconstrained fn main$f0(i$l0: u32) -> pub u32 {
+        let _a$l1 = foo$f1()[1][0][1];
+        let _s$l2 = foo$f1()[1][0][1];
+        i$l0
+    }
+    unconstrained fn foo$f1() -> [[[[u128; 0]; 2]; 1]; 2] {
+        [[[[], []]], [[[], []]]]
+    }
+    ");
+}