feat(optimization): Try to optimize ownership of assign statements

compiler/noirc_frontend/src/ownership/last_uses.rs

@@ -1,6 +1,5 @@
 //! This module contains the last use analysis pass which is run on each function before
-//! the ownership pass when the experimental ownership scheme is enabled. This pass does
-//! not run without this experimental flag - and if it did its results would go unused.
+//! the ownership pass.
 //!
 //! The purpose of this pass is to find which instance of a variable is the variable's
 //! last use. Note that a variable may have multiple last uses. This can happen if the
@@ -105,6 +104,8 @@
     ///           println((x, b));
     ///       }
     ///   }
+    ///   - The exception to this is if the variable is assigned to within a loop, we may be able
+    ///   to move the last use of the variable before it is assigned a fresh value.
     ///   ```
     ///   In the snippet above, `x` will have loop index 0 which does not match its last use
     ///   within the for loop (1 loop deep = loop index of 1). However, `b` has loop index 1
@@ -123,17 +124,47 @@
     ///   }
     ///   ```
     ///   `x` above has two last uses, one in each if branch.
-    last_uses: HashMap<LocalId, (/*loop index*/ usize, Branches)>,
+    last_uses: HashMap<LocalId, UseData>,
+
+    /// When an assignment comes we want to cut the last_uses of the variable being assigned to
+    /// so that the last use before the assignment - if there is one, can be moved since we're
+    /// reassigning the value anyway.
+    ///
+    /// E.g. `foo = bar(foo)` should be able to move `foo` into `bar` even if it is used afterward
+    /// since the result of `bar(foo)` will be assigned to `foo` afterward anyway.
+    last_assignment_uses: Vec<(LocalId, Branches)>,
+}
+
+struct UseData {
+    /// The loop index this variable was defined in
+    defined_in: LoopIndex,
+
+    /// The loop index the most recent instance of this variable was used in
+    last_used_in: LoopIndex,
+
+    /// The Branches conditional structure which locates the last use of the variable.
+    branches: Branches,
 }
 
+impl UseData {
+    fn new(defined_in: LoopIndex) -> Self {
+        Self { defined_in, last_used_in: defined_in, branches: Branches::None }
+    }
+}
+
+type LoopIndex = usize;
+
 impl Context {
     /// Traverse the given function and return the last use(s) of each local variable.
     /// A variable may have multiple last uses if it was last used within a conditional expression.
     pub(super) fn find_last_uses_of_variables(
         function: &Function,
     ) -> HashMap<LocalId, Vec<IdentId>> {
-        let mut context =
-            LastUseContext { current_loop_and_branch: Vec::new(), last_uses: HashMap::default() };
+        let mut context = LastUseContext {
+            current_loop_and_branch: Vec::new(),
+            last_uses: HashMap::default(),
+            last_assignment_uses: Vec::new(),
+        };
         context.push_loop_scope();
         for (parameter, ..) in &function.parameters {
             context.declare_variable(*parameter);
@@ -159,13 +190,13 @@
         *path = BranchesPath::IfOrMatch { branch_index, rest };
     }
 
     fn unbranch(&mut self) {
         let path =
             self.current_loop_and_branch.last_mut().expect("We should always have at least 1 path");
         let rest = std::mem::replace(path, BranchesPath::Here);
 
         match rest {
             BranchesPath::Here => panic!("unbranch called without any branches"),
             BranchesPath::IfOrMatch { branch_index: _, rest } => *path = *rest,
         }
     }
@@ -176,7 +207,7 @@
 
     fn declare_variable(&mut self, id: LocalId) {
         let loop_index = self.loop_index();
-        self.last_uses.insert(id, (loop_index, Branches::None));
+        self.last_uses.insert(id, UseData::new(loop_index));
     }
 
     /// Remember a new use of the given variable, possibly overwriting or
@@ -192,12 +223,12 @@
             self.current_loop_and_branch.last().expect("We should always have at least 1 scope");
         let loop_index = self.loop_index();
 
-        if let Some((variable_loop_index, uses)) = self.last_uses.get_mut(&id) {
-            if *variable_loop_index == loop_index {
-                Self::remember_use_of_variable_rec(uses, path, variable);
-            } else {
-                *uses = Branches::None;
+        if let Some(use_data) = self.last_uses.get_mut(&id) {
+            if use_data.defined_in != loop_index {
+                use_data.branches = Branches::None;
             }
+            use_data.last_used_in = loop_index;
+            Self::remember_use_of_variable_rec(&mut use_data.branches, path, variable);
         }
     }
 
@@ -239,10 +270,21 @@
     }
 
     fn get_variables_to_move(self) -> HashMap<LocalId, Vec<IdentId>> {
-        self.last_uses
+        let mut last_uses: HashMap<LocalId, Vec<IdentId>> = self
+            .last_uses
             .into_iter()
-            .map(|(definition, (_, last_uses))| (definition, last_uses.flatten_uses()))
-            .collect()
+            .filter_map(|(definition, use_data)| {
+                // Only consider moves where the last use was not in a nested loop.
+                // All uses of a variable inside a loop it was not defined in need to be clones.
+                (use_data.defined_in == use_data.last_used_in)
+                    .then(|| (definition, use_data.branches.flatten_uses()))
+            })
+            .collect();
+
+        for (local, assignment_uses) in self.last_assignment_uses {
+            last_uses.entry(local).or_default().extend(assignment_uses.flatten_uses());
+        }
+        last_uses
     }
 
     fn track_variables_in_expression(&mut self, expr: &Expression) {
@@ -261,7 +303,7 @@
             Expression::While(while_expr) => self.track_variables_in_while(while_expr),
             Expression::If(if_expr) => self.track_variables_in_if(if_expr),
             Expression::Match(match_expr) => self.track_variables_in_match(match_expr),
             Expression::Tuple(elems) => self.track_variables_in_tuple(elems),
             Expression::ExtractTupleField(tuple, _index) => {
                 self.track_variables_in_expression(tuple);
             }
@@ -340,12 +382,12 @@
 
         self.branch(0);
         self.track_variables_in_expression(&if_expr.consequence);
         self.unbranch();
 
         if let Some(alt) = &if_expr.alternative {
             self.branch(1);
             self.track_variables_in_expression(alt);
             self.unbranch();
         }
     }
 
@@ -357,7 +399,7 @@
 
             self.branch(i);
             self.track_variables_in_expression(&case.branch);
             self.unbranch();
         }
 
         if let Some(default_case) = &match_expr.default_case {
@@ -367,8 +409,8 @@
         }
     }
 
     fn track_variables_in_tuple(&mut self, elems: &[Expression]) {
         for elem in elems {
             self.track_variables_in_expression(elem);
         }
     }
@@ -400,6 +442,41 @@
     fn track_variables_in_assign(&mut self, assign: &ast::Assign) {
         self.track_variables_in_lvalue(&assign.lvalue);
         self.track_variables_in_expression(&assign.expression);
+
+        // Try to make the last use of a variable in the current loop index a move
+        if let Some(id) = Self::try_get_lvalue_local_id(&assign.lvalue) {
+            let loop_index = self.loop_index();
+            if let Some(use_data) = self.last_uses.get_mut(&id) {
+                // Don't apply this optimization if the variable was last used in a loop
+                // which has since ended, e.g:
+                // ```noir
+                // for _ in 0 .. 10 {
+                //     foo(bar);
+                // }
+                // bar = new_value;
+                // ```
+                // without this check we'd make `foo(bar)` a move despite `bar` being used
+                // again on the next iteration of the loop.
+                if use_data.last_used_in <= loop_index {
+                    self.last_assignment_uses.push((id, use_data.branches.clone()));
+                }
+            }
+        }
+    }
+
+    fn try_get_lvalue_local_id(lvalue: &ast::LValue) -> Option<LocalId> {
+        // Handling cases other than the full value being assigned is more complex
+        // so we only handle the case where we assign the entire value so that we don't
+        // accidentally move too much. E.g. if we only assign `foo.bar = baz(foo)` we
+        // don't want to move all of `foo` in that case. So we only handle identifiers
+        // for the lvalue.
+        match lvalue {
+            ast::LValue::Ident(ident) => match &ident.definition {
+                ast::Definition::Local(local_id) => Some(*local_id),
+                _ => None,
+            },
+            _ => None,
+        }
     }
 
     /// A variable in an lvalue position is never moved (otherwise you wouldn't

compiler/noirc_frontend/src/ownership/mod.rs

@@ -320,6 +320,21 @@ impl Context {
     }
 
     fn handle_call(&mut self, call: &mut crate::monomorphization::ast::Call) {
+        // Hack: We don't have immutable references yet so we pretend some array builtins
+        // accepts arrays as a reference to avoid unnecessary clones.
+        if let crate::monomorphization::ast::Expression::Ident(variable) = call.func.as_ref() {
+            if let crate::monomorphization::ast::Definition::Builtin(name) = &variable.definition {
+                // Having {array,slice}_refcount not increment refcounts makes debugging
+                // refcounts much easier
+                if name == "array_len" || name == "slice_refcount" || name == "array_refcount" {
+                    for arg in &mut call.arguments {
+                        self.handle_reference_expression(arg);
+                    }
+                    return;
+                }
+            }
+        }
+
         self.handle_expression(&mut call.func);
         for arg in &mut call.arguments {
             self.handle_expression(arg);

test_programs/execution_success/slice_push_back_copies/Nargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "arithmetic_generics"
+type = "bin"
+authors = [""]
+compiler_version = ">=0.32.0"
+
+[dependencies]

test_programs/execution_success/slice_push_back_copies/Prover.toml

@@ -0,0 +1,34 @@
+arr = [
+    1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    2, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    3, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    4, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    5, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    6, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    7, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    8, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    9, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+   10, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+
+    1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    2, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    3, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    4, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    5, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    6, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    7, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    8, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    9, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+   10, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+
+    1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    2, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    3, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    4, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    5, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    6, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    7, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    8, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+    9, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+   10, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+]

test_programs/execution_success/slice_push_back_copies/src/main.nr

@@ -0,0 +1,10 @@
+pub unconstrained fn main(arr: [Field; 300]) {
+    let mut serialized_args: [Field] = &[];
+    let arr_serialized: [[Field; 1]; 300] = arr.map(|x| [x]);
+
+    for i in 0..arr.len() {
+        serialized_args = serialized_args.append(arr_serialized[i].as_slice());
+    }
+    let calldata = serialized_args.push_front(611740813);
+    std::println(calldata);
+}