chore: manage call stacks using a tree

compiler/noirc_evaluator/src/acir/acir_variable.rs

@@ -23,7 +23,7 @@ use std::{borrow::Cow, hash::Hash};
 use crate::brillig::brillig_ir::artifact::GeneratedBrillig;
 use crate::errors::{InternalBug, InternalError, RuntimeError, SsaReport};
 use crate::ssa::ir::{
-    dfg::CallStack, instruction::Endian, types::NumericType, types::Type as SsaType,
+    call_stack::CallStack, instruction::Endian, types::NumericType, types::Type as SsaType,
 };
 
 use super::big_int::BigIntContext;

compiler/noirc_evaluator/src/acir/generated_acir.rs

@@ -16,7 +16,7 @@ use super::brillig_directive;
 use crate::{
     brillig::brillig_ir::artifact::GeneratedBrillig,
     errors::{InternalError, RuntimeError, SsaReport},
-    ssa::ir::dfg::CallStack,
+    ssa::ir::call_stack::CallStack,
     ErrorType,
 };
 

compiler/noirc_evaluator/src/acir/mod.rs

@@ -35,7 +35,8 @@ use crate::ssa::ir::instruction::Hint;
 use crate::ssa::{
     function_builder::data_bus::DataBus,
     ir::{
-        dfg::{CallStack, DataFlowGraph},
+        call_stack::CallStack,
+        dfg::DataFlowGraph,
         function::{Function, FunctionId, RuntimeType},
         instruction::{
             Binary, BinaryOp, ConstrainError, Instruction, InstructionId, Intrinsic,
@@ -675,7 +676,7 @@ impl<'a> Context<'a> {
         brillig: &Brillig,
     ) -> Result<Vec<SsaReport>, RuntimeError> {
         let instruction = &dfg[instruction_id];
-        self.acir_context.set_call_stack(dfg.get_call_stack(instruction_id));
+        self.acir_context.set_call_stack(dfg.get_instruction_call_stack(instruction_id));
         let mut warnings = Vec::new();
         match instruction {
             Instruction::Binary(binary) => {
@@ -1822,7 +1823,7 @@ impl<'a> Context<'a> {
     ) -> Result<(Vec<AcirVar>, Vec<SsaReport>), RuntimeError> {
         let (return_values, call_stack) = match terminator {
             TerminatorInstruction::Return { return_values, call_stack } => {
-                (return_values, call_stack.clone())
+                (return_values, *call_stack)
             }
             // TODO(https://github.com/noir-lang/noir/issues/4616): Enable recursion on foldable/non-inlined ACIR functions
             _ => unreachable!("ICE: Program must have a singular return"),
@@ -1841,6 +1842,7 @@ impl<'a> Context<'a> {
             }
         }
 
+        let call_stack = dfg.call_stack_data.get_call_stack(call_stack);
         let warnings = if has_constant_return {
             vec![SsaReport::Warning(InternalWarning::ReturnConstant { call_stack })]
         } else {
@@ -2903,7 +2905,7 @@ mod test {
         ssa::{
             function_builder::FunctionBuilder,
             ir::{
-                dfg::CallStack,
+                call_stack::CallStack,
                 function::FunctionId,
                 instruction::BinaryOp,
                 map::Id,
@@ -2932,7 +2934,9 @@ mod test {
         // Set a call stack for testing whether `brillig_locations` in the `GeneratedAcir` was accurately set.
         let mut stack = CallStack::unit(Location::dummy());
         stack.push_back(Location::dummy());
-        builder.set_call_stack(stack);
+        let call_stack =
+            builder.current_function.dfg.call_stack_data.get_or_insert_locations(stack);
+        builder.set_call_stack(call_stack);
 
         let foo_v0 = builder.add_parameter(Type::field());
         let foo_v1 = builder.add_parameter(Type::field());

compiler/noirc_evaluator/src/brillig/brillig_gen.rs

@@ -18,7 +18,7 @@ use super::{
 };
 use crate::{
     errors::InternalError,
-    ssa::ir::{dfg::CallStack, function::Function},
+    ssa::ir::{call_stack::CallStack, function::Function},
 };
 
 /// Converting an SSA function into Brillig bytecode.

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

@@ -7,7 +7,7 @@ use crate::brillig::brillig_ir::registers::Stack;
 use crate::brillig::brillig_ir::{
     BrilligBinaryOp, BrilligContext, ReservedRegisters, BRILLIG_MEMORY_ADDRESSING_BIT_SIZE,
 };
-use crate::ssa::ir::dfg::CallStack;
+use crate::ssa::ir::call_stack::CallStack;
 use crate::ssa::ir::instruction::{ConstrainError, Hint};
 use crate::ssa::ir::{
     basic_block::BasicBlockId,
@@ -201,7 +201,7 @@ impl<'block> BrilligBlock<'block> {
     /// Converts an SSA instruction into a sequence of Brillig opcodes.
     fn convert_ssa_instruction(&mut self, instruction_id: InstructionId, dfg: &DataFlowGraph) {
         let instruction = &dfg[instruction_id];
-        self.brillig_context.set_call_stack(dfg.get_call_stack(instruction_id));
+        self.brillig_context.set_call_stack(dfg.get_instruction_call_stack(instruction_id));
 
         self.initialize_constants(
             &self.function_context.constant_allocation.allocated_at_location(

compiler/noirc_evaluator/src/brillig/brillig_ir.rs

@@ -30,7 +30,7 @@ pub(crate) use instructions::BrilligBinaryOp;
 use registers::{RegisterAllocator, ScratchSpace};
 
 use self::{artifact::BrilligArtifact, debug_show::DebugToString, registers::Stack};
-use crate::ssa::ir::dfg::CallStack;
+use crate::ssa::ir::call_stack::CallStack;
 use acvm::{
     acir::brillig::{MemoryAddress, Opcode as BrilligOpcode},
     AcirField,

compiler/noirc_evaluator/src/brillig/brillig_ir/artifact.rs

@@ -2,7 +2,7 @@ use acvm::acir::brillig::Opcode as BrilligOpcode;
 use acvm::acir::circuit::ErrorSelector;
 use std::collections::{BTreeMap, HashMap};
 
-use crate::ssa::ir::{basic_block::BasicBlockId, dfg::CallStack, function::FunctionId};
+use crate::ssa::ir::{basic_block::BasicBlockId, call_stack::CallStack, function::FunctionId};
 use crate::ErrorType;
 
 use super::procedures::ProcedureId;

compiler/noirc_evaluator/src/errors.rs

@@ -12,7 +12,7 @@ use iter_extended::vecmap;
 use noirc_errors::{CustomDiagnostic as Diagnostic, FileDiagnostic};
 use thiserror::Error;
 
-use crate::ssa::ir::{dfg::CallStack, types::NumericType};
+use crate::ssa::ir::{call_stack::CallStack, types::NumericType};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, PartialEq, Eq, Clone, Error)]

compiler/noirc_evaluator/src/ssa/checks/check_for_underconstrained_values.rs

@@ -382,7 +382,7 @@ impl DependencyContext {
             .keys()
             .map(|brillig_call| {
                 SsaReport::Bug(InternalBug::UncheckedBrilligCall {
-                    call_stack: function.dfg.get_call_stack(*brillig_call),
+                    call_stack: function.dfg.get_instruction_call_stack(*brillig_call),
                 })
             })
             .collect();
@@ -513,7 +513,7 @@ impl Context {
             // There is a value not in the set, which means that the inputs/outputs of this call have not been properly constrained
             if unused_inputs {
                 warnings.push(SsaReport::Bug(InternalBug::IndependentSubgraph {
-                    call_stack: function.dfg.get_call_stack(
+                    call_stack: function.dfg.get_instruction_call_stack(
                         self.brillig_return_to_instruction_id[&brillig_output_in_set],
                     ),
                 }));

compiler/noirc_evaluator/src/ssa/function_builder/mod.rs

@@ -18,7 +18,8 @@ use crate::ssa::ir::{
 use super::{
     ir::{
         basic_block::BasicBlock,
-        dfg::{CallStack, InsertInstructionResult},
+        call_stack::{CallStack, CallStackId},
+        dfg::InsertInstructionResult,
         function::RuntimeType,
         instruction::{ConstrainError, InstructionId, Intrinsic},
         types::NumericType,
@@ -34,10 +35,10 @@ use super::{
 /// Contrary to the name, this struct has the capacity to build as many
 /// functions as needed, although it is limited to one function at a time.
 pub(crate) struct FunctionBuilder {
-    pub(super) current_function: Function,
+    pub(crate) current_function: Function,
     current_block: BasicBlockId,
     finished_functions: Vec<Function>,
-    call_stack: CallStack,
+    call_stack: CallStackId,
     error_types: BTreeMap<ErrorSelector, HirType>,
 }
 
@@ -53,7 +54,7 @@ impl FunctionBuilder {
             current_block: new_function.entry_block(),
             current_function: new_function,
             finished_functions: Vec::new(),
-            call_stack: CallStack::new(),
+            call_stack: CallStackId::root(),
             error_types: BTreeMap::default(),
         }
     }
@@ -78,11 +79,14 @@ impl FunctionBuilder {
         function_id: FunctionId,
         runtime_type: RuntimeType,
     ) {
+        let call_stack = self.current_function.dfg.get_call_stack(self.call_stack);
         let mut new_function = Function::new(name, function_id);
         new_function.set_runtime(runtime_type);
         self.current_block = new_function.entry_block();
-
         let old_function = std::mem::replace(&mut self.current_function, new_function);
+        // Copy the call stack to the new function
+        self.call_stack =
+            self.current_function.dfg.call_stack_data.get_or_insert_locations(call_stack);
         self.finished_functions.push(old_function);
     }
 
@@ -171,7 +175,7 @@ impl FunctionBuilder {
             instruction,
             block,
             ctrl_typevars,
-            self.call_stack.clone(),
+            self.call_stack,
         )
     }
 
@@ -196,17 +200,17 @@ impl FunctionBuilder {
     }
 
     pub(crate) fn set_location(&mut self, location: Location) -> &mut FunctionBuilder {
-        self.call_stack = CallStack::unit(location);
+        self.call_stack = self.current_function.dfg.call_stack_data.add_location_to_root(location);
         self
     }
 
-    pub(crate) fn set_call_stack(&mut self, call_stack: CallStack) -> &mut FunctionBuilder {
+    pub(crate) fn set_call_stack(&mut self, call_stack: CallStackId) -> &mut FunctionBuilder {
         self.call_stack = call_stack;
         self
     }
 
     pub(crate) fn get_call_stack(&self) -> CallStack {
-        self.call_stack.clone()
+        self.current_function.dfg.get_call_stack(self.call_stack)
     }
 
     /// Insert a Load instruction at the end of the current block, loading from the given address
@@ -378,7 +382,7 @@ impl FunctionBuilder {
         destination: BasicBlockId,
         arguments: Vec<ValueId>,
     ) {
-        let call_stack = self.call_stack.clone();
+        let call_stack = self.call_stack;
         self.terminate_block_with(TerminatorInstruction::Jmp {
             destination,
             arguments,
@@ -394,7 +398,7 @@ impl FunctionBuilder {
         then_destination: BasicBlockId,
         else_destination: BasicBlockId,
     ) {
-        let call_stack = self.call_stack.clone();
+        let call_stack = self.call_stack;
         self.terminate_block_with(TerminatorInstruction::JmpIf {
             condition,
             then_destination,
@@ -405,7 +409,7 @@ impl FunctionBuilder {
 
     /// Terminate the current block with a return instruction
     pub(crate) fn terminate_with_return(&mut self, return_values: Vec<ValueId>) {
-        let call_stack = self.call_stack.clone();
+        let call_stack = self.call_stack;
         self.terminate_block_with(TerminatorInstruction::Return { return_values, call_stack });
     }
 

compiler/noirc_evaluator/src/ssa/ir/basic_block.rs

@@ -1,5 +1,5 @@
 use super::{
-    dfg::CallStack,
+    call_stack::CallStackId,
     instruction::{InstructionId, TerminatorInstruction},
     map::Id,
     value::ValueId,
@@ -123,7 +123,7 @@ impl BasicBlock {
             terminator,
             TerminatorInstruction::Return {
                 return_values: Vec::new(),
-                call_stack: CallStack::new(),
+                call_stack: CallStackId::root(),
             },
         )
     }

compiler/noirc_evaluator/src/ssa/ir/call_stack.rs

@@ -0,0 +1,128 @@
+use fxhash::FxHashMap;
+use serde::{Deserialize, Serialize};
+
+use noirc_errors::Location;
+
+pub(crate) type CallStack = im::Vector<Location>;
+
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub(crate) struct CallStackId(u32);
+
+impl CallStackId {
+    pub(crate) fn root() -> Self {
+        Self::new(0)
+    }
+
+    fn new(id: usize) -> Self {
+        Self(id as u32)
+    }
+
+    pub(crate) fn index(&self) -> usize {
+        self.0 as usize
+    }
+
+    pub(crate) fn is_root(&self) -> bool {
+        self.0 == 0
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub(crate) struct LocationNode {
+    pub(crate) parent: Option<CallStackId>,
+    pub(crate) children: Vec<CallStackId>,
+    pub(crate) children_hash: FxHashMap<u64, CallStackId>,
+    pub(crate) value: Location,
+}
+
+impl LocationNode {
+    pub(crate) fn new(parent: Option<CallStackId>, value: Location) -> Self {
+        LocationNode { parent, children: Vec::new(), children_hash: FxHashMap::default(), value }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub(crate) struct CallStackHelper {
+    locations: Vec<LocationNode>,
+}
+
+impl Default for CallStackHelper {
+    /// Generates a new helper, with an empty location tree
+    fn default() -> Self {
+        let mut result = CallStackHelper { locations: Vec::new() };
+        result.add_location_to_root(Location::dummy());
+        result
+    }
+}
+
+impl CallStackHelper {
+    /// Construct a CallStack from a CallStackId
+    pub(crate) fn get_call_stack(&self, mut call_stack: CallStackId) -> CallStack {
+        let mut result = im::Vector::new();
+        while let Some(parent) = self.locations[call_stack.index()].parent {
+            result.push_back(self.locations[call_stack.index()].value);
+            call_stack = parent;
+        }
+        result
+    }
+
+    /// Returns a new CallStackId which extends the call_stack with the provided call_stack.
+    pub(crate) fn extend_call_stack(
+        &mut self,
+        mut call_stack: CallStackId,
+        locations: &CallStack,
+    ) -> CallStackId {
+        for location in locations {
+            call_stack = self.add_child(call_stack, *location);
+        }
+        call_stack
+    }
+
+    /// Adds a location to the call stack
+    pub(crate) fn add_child(&mut self, call_stack: CallStackId, location: Location) -> CallStackId {
+        let key = fxhash::hash64(&location);
+        if let Some(result) = self.locations[call_stack.index()].children_hash.get(&key) {
+            if self.locations[result.index()].value == location {
+                return *result;
+            }
+        }
+        let new_location = LocationNode::new(Some(call_stack), location);
+        let key = fxhash::hash64(&new_location.value);
+        self.locations.push(new_location);
+        let new_location_id = CallStackId::new(self.locations.len() - 1);
+
+        self.locations[call_stack.index()].children.push(new_location_id);
+        self.locations[call_stack.index()].children_hash.insert(key, new_location_id);
+        new_location_id
+    }
+
+    /// Retrieve the CallStackId corresponding to call_stack with the last 'len' locations removed.
+    pub(crate) fn unwind_call_stack(
+        &self,
+        mut call_stack: CallStackId,
+        mut len: usize,
+    ) -> CallStackId {
+        while len > 0 {
+            if let Some(parent) = self.locations[call_stack.index()].parent {
+                len -= 1;
+                call_stack = parent;
+            } else {
+                break;
+            }
+        }
+        call_stack
+    }
+
+    pub(crate) fn add_location_to_root(&mut self, location: Location) -> CallStackId {
+        if self.locations.is_empty() {
+            self.locations.push(LocationNode::new(None, location));
+            CallStackId::root()
+        } else {
+            self.add_child(CallStackId::root(), location)
+        }
+    }
+
+    /// Get (or create) a CallStackId corresponding to the given locations
+    pub(crate) fn get_or_insert_locations(&mut self, locations: CallStack) -> CallStackId {
+        self.extend_call_stack(CallStackId::root(), &locations)
+    }
+}