From c56fb2fa6adf6232af1a31ee2abd57ddff3ff3b6 Mon Sep 17 00:00:00 2001
From: sirasistant <sirasistant@gmail.com>
Date: Mon, 12 Feb 2024 13:13:42 +0000
Subject: [PATCH] feat: Deallocate stack items at the instruction level

---
 .../src/brillig/brillig_gen/brillig_block.rs  |  6 ++++-
 .../brillig_gen/brillig_block_variables.rs    | 23 ++++++++++++++++---
 .../noirc_evaluator/src/brillig/brillig_ir.rs |  6 ++++-
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
index d095e4efd5f..10c33ae7ffe 100644
--- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
+++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
@@ -665,7 +665,11 @@ impl<'block> BrilligBlock<'block> {
             .expect("Last uses for instruction should have been computed");
 
         for dead_variable in dead_variables {
-            self.variables.remove_variable(dead_variable);
+            self.variables.remove_variable(
+                dead_variable,
+                self.function_context,
+                self.brillig_context,
+            );
         }
         self.brillig_context.set_call_stack(CallStack::new());
     }
diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs
index cbb3049a904..b4c96de1969 100644
--- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs
+++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs
@@ -19,6 +19,7 @@ use super::brillig_fn::FunctionContext;
 #[derive(Debug, Default)]
 pub(crate) struct BlockVariables {
     available_variables: HashSet<ValueId>,
+    block_parameters: HashSet<ValueId>,
     available_constants: HashMap<ValueId, BrilligVariable>,
 }
 
@@ -26,7 +27,8 @@ impl BlockVariables {
     /// Creates a BlockVariables instance. It uses the variables that are live in to the block and the global available variables (block parameters)
     pub(crate) fn new(live_in: HashSet<ValueId>, all_block_parameters: HashSet<ValueId>) -> Self {
         BlockVariables {
-            available_variables: live_in.into_iter().chain(all_block_parameters).collect(),
+            available_variables: live_in.into_iter().chain(all_block_parameters.clone()).collect(),
+            block_parameters: all_block_parameters,
             ..Default::default()
         }
     }
@@ -81,8 +83,23 @@ impl BlockVariables {
     }
 
     /// Removes a variable so it's not used anymore within this block.
-    pub(crate) fn remove_variable(&mut self, value_id: &ValueId) {
-        self.available_variables.remove(value_id);
+    pub(crate) fn remove_variable(
+        &mut self,
+        value_id: &ValueId,
+        function_context: &mut FunctionContext,
+        brillig_context: &mut BrilligContext,
+    ) {
+        assert!(self.available_variables.remove(value_id), "ICE: Variable is not available");
+        // Block parameters should not be deallocated
+        if !self.block_parameters.contains(value_id) {
+            let variable = function_context
+                .ssa_value_allocations
+                .get(value_id)
+                .expect("ICE: Variable allocation not found");
+            variable.extract_registers().iter().for_each(|register| {
+                brillig_context.deallocate_register(*register);
+            });
+        }
     }
 
     /// For a given SSA value id, return the corresponding cached allocation.
diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir.rs
index 94656ed6f46..8b3ed91b670 100644
--- a/compiler/noirc_evaluator/src/brillig/brillig_ir.rs
+++ b/compiler/noirc_evaluator/src/brillig/brillig_ir.rs
@@ -468,6 +468,9 @@ impl BrilligContext {
             sources.push(*return_register);
             destinations.push(destination_register);
         }
+        destinations
+            .iter()
+            .for_each(|destination| self.registers.ensure_register_is_allocated(*destination));
         self.mov_registers_to_registers_instruction(sources, destinations);
         self.stop_instruction();
     }
@@ -935,11 +938,12 @@ impl BrilligContext {
     ) {
         // Allocate our result registers and write into them
         // We assume the return values of our call are held in 0..num results register indices
-        let (sources, destinations) = result_registers
+        let (sources, destinations): (Vec<_>, Vec<_>) = result_registers
             .iter()
             .enumerate()
             .map(|(i, result_register)| (self.register(i), *result_register))
             .unzip();
+        sources.iter().for_each(|source| self.registers.ensure_register_is_allocated(*source));
         self.mov_registers_to_registers_instruction(sources, destinations);
 
         // Restore all the same registers we have, in exact reverse order.