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 d095e4efd5..10c33ae7ff 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 cbb3049a90..b4c96de196 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, + block_parameters: HashSet, available_constants: HashMap, } @@ -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, all_block_parameters: HashSet) -> 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 94656ed6f4..8b3ed91b67 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.