feat: Consider block parameters in variable liveness

compiler/noirc_evaluator/src/brillig/brillig_gen.rs

@@ -14,7 +14,7 @@ use crate::ssa::ir::function::Function;
 pub(crate) fn convert_ssa_function(func: &Function, enable_debug_trace: bool) -> BrilligArtifact {
     let mut brillig_context = BrilligContext::new(enable_debug_trace);
 
-    let mut function_context = FunctionContext::new(func, &mut brillig_context);
+    let mut function_context = FunctionContext::new(func);
 
     brillig_context.enter_context(FunctionContext::function_id_to_function_label(func.id()));
 

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs

@@ -7,7 +7,7 @@ use crate::brillig::brillig_ir::{
 use crate::ssa::ir::dfg::CallStack;
 use crate::ssa::ir::instruction::ConstrainError;
 use crate::ssa::ir::{
-    basic_block::{BasicBlock, BasicBlockId},
+    basic_block::BasicBlockId,
     dfg::DataFlowGraph,
     function::FunctionId,
     instruction::{
@@ -49,8 +49,7 @@ impl<'block> BrilligBlock<'block> {
         dfg: &DataFlowGraph,
     ) {
         let live_in = function_context.liveness.get_live_in(&block_id);
-        let variables =
-            BlockVariables::new(live_in.clone(), function_context.all_block_parameters());
+        let variables = BlockVariables::new(live_in.clone());
 
         brillig_context.set_allocated_registers(
             variables
@@ -72,9 +71,9 @@ impl<'block> BrilligBlock<'block> {
         let block_label = self.create_block_label_for_current_function(self.block_id);
         self.brillig_context.enter_context(block_label);
 
-        // Convert the block parameters
+        self.convert_block_params(dfg);
+
         let block = &dfg[self.block_id];
-        self.convert_block_params(block, dfg);
 
         // Convert all of the instructions into the block
         for instruction_id in block.instructions() {
@@ -134,12 +133,8 @@ impl<'block> BrilligBlock<'block> {
                 let target_block = &dfg[*destination_block];
                 for (src, dest) in arguments.iter().zip(target_block.parameters()) {
                     // Destinations are block parameters so they should have been allocated previously.
-                    let destination = self.variables.get_block_param(
-                        self.function_context,
-                        *destination_block,
-                        *dest,
-                        dfg,
-                    );
+                    let destination =
+                        self.variables.get_allocation(self.function_context, *dest, dfg);
                     let source = self.convert_ssa_value(*src, dfg);
                     self.pass_variable(source, destination);
                 }
@@ -206,10 +201,14 @@ impl<'block> BrilligBlock<'block> {
         }
     }
 
-    /// Converts SSA Block parameters into Brillig Registers.
-    fn convert_block_params(&mut self, block: &BasicBlock, dfg: &DataFlowGraph) {
-        for param_id in block.parameters() {
-            let value = &dfg[*param_id];
+    /// Allocates the block parameters that the given block is defining
+    fn convert_block_params(&mut self, dfg: &DataFlowGraph) {
+        // We don't allocate the block parameters here, we allocate the parameters the block is defining.
+        // Since predecessors to a block have to know where the parameters of the block are allocated to pass data to it,
+        // the block parameters need to be defined/allocated before the given block. Variable liveness provides when the block parameters are defined.
+        // For the entry block, the defined block params will be the params of the function + any extra params of blocks it's the immediate dominator of.
+        for param_id in self.function_context.liveness.defined_block_params(&self.block_id) {
+            let value = &dfg[param_id];
             let param_type = match value {
                 Value::Param { typ, .. } => typ,
                 _ => unreachable!("ICE: Only Param type values should appear in block parameters"),
@@ -220,10 +219,10 @@ impl<'block> BrilligBlock<'block> {
                 // Be a valid pointer to the array.
                 // For slices, two registers are passed, the pointer to the data and a register holding the size of the slice.
                 Type::Numeric(_) | Type::Array(..) | Type::Slice(..) | Type::Reference(_) => {
-                    self.variables.get_block_param(
+                    self.variables.define_variable(
                         self.function_context,
-                        self.block_id,
-                        *param_id,
+                        self.brillig_context,
+                        param_id,
                         dfg,
                     );
                 }

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs

@@ -9,7 +9,6 @@ use crate::{
         BrilligContext,
     },
     ssa::ir::{
-        basic_block::BasicBlockId,
         dfg::DataFlowGraph,
         types::{CompositeType, Type},
         value::ValueId,
@@ -21,18 +20,13 @@ 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>,
 }
 
 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.clone()).collect(),
-            block_parameters: all_block_parameters,
-            ..Default::default()
-        }
+    pub(crate) fn new(live_in: HashSet<ValueId>) -> Self {
+        BlockVariables { available_variables: live_in, ..Default::default() }
     }
 
     /// Returns all non-constant variables that have not been removed at this point.
@@ -92,16 +86,13 @@ impl BlockVariables {
         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);
-            });
-        }
+        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.
@@ -155,27 +146,6 @@ impl BlockVariables {
     pub(crate) fn dump_constants(&mut self) {
         self.available_constants.clear();
     }
-
-    /// For a given block parameter, return the allocation that was done globally to the function.
-    pub(crate) fn get_block_param(
-        &mut self,
-        function_context: &FunctionContext,
-        block_id: BasicBlockId,
-        value_id: ValueId,
-        dfg: &DataFlowGraph,
-    ) -> BrilligVariable {
-        let value_id = dfg.resolve(value_id);
-        assert!(
-            function_context
-                .block_parameters
-                .get(&block_id)
-                .expect("Block not found")
-                .contains(&value_id),
-            "Value is not a block parameter"
-        );
-
-        *function_context.ssa_value_allocations.get(&value_id).expect("Block param not found")
-    }
 }
 
 /// Computes the length of an array. This will match with the indexes that SSA will issue

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs

@@ -4,7 +4,6 @@ use crate::{
     brillig::brillig_ir::{
         artifact::{BrilligParameter, Label},
         brillig_variable::{get_bit_size_from_ssa_type, BrilligVariable},
-        BrilligContext,
     },
     ssa::ir::{
         basic_block::BasicBlockId,
@@ -14,57 +13,37 @@ use crate::{
         value::ValueId,
     },
 };
-use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
+use fxhash::FxHashMap as HashMap;
 
-use super::{brillig_block_variables::allocate_value, variable_liveness::VariableLiveness};
+use super::variable_liveness::VariableLiveness;
 
 pub(crate) struct FunctionContext {
     pub(crate) function_id: FunctionId,
     /// Map from SSA values its allocation. Since values can be only defined once in SSA form, we insert them here on when we allocate them at their definition.
     pub(crate) ssa_value_allocations: HashMap<ValueId, BrilligVariable>,
-    /// Block parameters are pre allocated at the function level.
-    pub(crate) block_parameters: HashMap<BasicBlockId, Vec<ValueId>>,
     /// The block ids of the function in reverse post order.
     pub(crate) blocks: Vec<BasicBlockId>,
     /// Liveness information for each variable in the function.
     pub(crate) liveness: VariableLiveness,
 }
 
 impl FunctionContext {
-    /// Creates a new function context. It will allocate parameters for all blocks and compute the liveness of every variable.
-    pub(crate) fn new(function: &Function, brillig_context: &mut BrilligContext) -> Self {
+    /// Creates a new function context. It will compute the liveness of every variable.
+    pub(crate) fn new(function: &Function) -> Self {
         let id = function.id();
 
         let mut reverse_post_order = Vec::new();
         reverse_post_order.extend_from_slice(PostOrder::with_function(function).as_slice());
         reverse_post_order.reverse();
 
-        let mut block_parameters = HashMap::default();
-        let mut ssa_variable_to_register_or_memory = HashMap::default();
-
-        for &block_id in &reverse_post_order {
-            let block = &function.dfg[block_id];
-            let parameters = block.parameters().to_vec();
-            parameters.iter().for_each(|&value_id| {
-                let variable = allocate_value(value_id, brillig_context, &function.dfg);
-                ssa_variable_to_register_or_memory.insert(value_id, variable);
-            });
-            block_parameters.insert(block_id, parameters);
-        }
-
         Self {
             function_id: id,
-            ssa_value_allocations: ssa_variable_to_register_or_memory,
-            block_parameters,
+            ssa_value_allocations: HashMap::default(),
             blocks: reverse_post_order,
             liveness: VariableLiveness::from_function(function),
         }
     }
 
-    pub(crate) fn all_block_parameters(&self) -> HashSet<ValueId> {
-        self.block_parameters.values().flat_map(|parameters| parameters.iter()).cloned().collect()
-    }
-
     /// Creates a function label from a given SSA function id.
     pub(crate) fn function_id_to_function_label(function_id: FunctionId) -> Label {
         function_id.to_string()

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs

@@ -377,9 +377,9 @@ mod tests {
         builder.set_runtime(RuntimeType::Brillig);
 
         let ssa = builder.finish();
-        let mut brillig_context = create_context();
+        let brillig_context = create_context();
 
-        let function_context = FunctionContext::new(ssa.main(), &mut brillig_context);
+        let function_context = FunctionContext::new(ssa.main());
         (ssa, function_context, brillig_context)
     }