fix: handle large arrays in the monomorphiser and limit their length

compiler/noirc_evaluator/src/acir/mod.rs

@@ -31,10 +31,11 @@ mod types;
 use crate::brillig::Brillig;
 use crate::brillig::brillig_gen::gen_brillig_for;
 use crate::errors::{InternalError, RuntimeError};
+
 use crate::ssa::{
     function_builder::data_bus::DataBus,
     ir::{
-        dfg::DataFlowGraph,
+        dfg::{DataFlowGraph, MAX_ELEMENTS},
         function::{Function, RuntimeType},
         instruction::{
             Binary, BinaryOp, ConstrainError, Instruction, InstructionId, TerminatorInstruction,
@@ -188,8 +189,7 @@ impl<'a> Context<'a> {
         self.acir_context.acir_ir.input_witnesses =
             self.convert_ssa_block_params(entry_block.parameters(), dfg)?;
 
-        let num_return_witnesses =
-            self.get_num_return_witnesses(entry_block.unwrap_terminator(), dfg);
+        let num_return_witnesses = dfg.get_num_return_witnesses(main_func)?;
 
         // Create a witness for each return witness we have to guarantee that the return witnesses match the standard
         // layout for serializing those types as if they were being passed as inputs.
@@ -284,6 +284,25 @@ impl<'a> Context<'a> {
         self.acir_context.acir_ir.input_witnesses = self.acir_context.extract_witnesses(&inputs);
         let returns = main_func.returns().unwrap_or_default();
 
+        // Check the flattened size of return values to avoid OOM during Brillig entry point generation
+        let num_return_values: usize = returns
+            .iter()
+            .map(|result_id| dfg.type_of_value(*result_id).flattened_size().to_usize())
+            .sum();
+        if num_return_values > MAX_ELEMENTS {
+            let entry_block = &dfg[main_func.entry_block()];
+            let call_stack_id = match entry_block.unwrap_terminator() {
+                TerminatorInstruction::Return { call_stack, .. } => *call_stack,
+                _ => unreachable!("ICE: expected return terminator"),
+            };
+            let call_stack = dfg.call_stack_data.get_call_stack(call_stack_id);
+            return Err(RuntimeError::ReturnLimitExceeded {
+                num_witnesses: num_return_values,
+                max_witnesses: MAX_ELEMENTS,
+                call_stack,
+            });
+        }
+
         let outputs: Vec<AcirType> =
             vecmap(returns, |result_id| dfg.type_of_value(*result_id).into());
 
@@ -581,26 +600,6 @@ impl<'a> Context<'a> {
         self.define_result(dfg, instruction, AcirValue::Var(result, typ));
     }
 
-    /// Converts an SSA terminator's return values into their ACIR representations
-    fn get_num_return_witnesses(
-        &self,
-        terminator: &TerminatorInstruction,
-        dfg: &DataFlowGraph,
-    ) -> usize {
-        let return_values = match terminator {
-            TerminatorInstruction::Return { return_values, .. } => return_values,
-            TerminatorInstruction::Unreachable { .. } => return 0,
-            // TODO(https://github.com/noir-lang/noir/issues/4616): Enable recursion on foldable/non-inlined ACIR functions
-            TerminatorInstruction::JmpIf { .. } | TerminatorInstruction::Jmp { .. } => {
-                unreachable!("ICE: Program must have a singular return")
-            }
-        };
-
-        return_values
-            .iter()
-            .fold(0, |acc, value_id| acc + dfg.type_of_value(*value_id).flattened_size().to_usize())
-    }
-
     /// Converts an SSA terminator's return values into their ACIR representations
     fn convert_ssa_return(
         &mut self,

compiler/noirc_evaluator/src/brillig/brillig_gen.rs

@@ -21,7 +21,10 @@ use super::{
         artifact::{BrilligParameter, GeneratedBrillig},
     },
 };
-use crate::{errors::InternalError, ssa::ir::function::Function};
+use crate::{
+    errors::{InternalError, RuntimeError},
+    ssa::ir::function::Function,
+};
 
 /// Generates a complete Brillig entry point artifact for a given SSA-level [Function], linking all dependencies.
 ///
@@ -37,7 +40,7 @@ use crate::{errors::InternalError, ssa::ir::function::Function};
 ///
 /// # Returns
 /// - Ok([GeneratedBrillig]): Fully linked artifact for the entry point that can be executed as a Brillig program.
-/// - Err([InternalError]): If linking fails to find a dependency
+/// - Err([RuntimeError]): If the return value exceeds the witness limit or linking fails
 ///
 /// # Panics
 /// - If the global memory size for the function has not been precomputed.
@@ -46,7 +49,14 @@ pub(crate) fn gen_brillig_for(
     arguments: Vec<BrilligParameter>,
     brillig: &Brillig,
     options: &BrilligOptions,
-) -> Result<GeneratedBrillig<FieldElement>, InternalError> {
+) -> Result<GeneratedBrillig<FieldElement>, RuntimeError> {
+    let return_parameters = FunctionContext::return_values(func);
+
+    // Check if the return value size exceeds the limit before generating the entry point.
+    // This is done early to avoid the expensive entry point codegen which iterates over
+    // each element in the return arrays.
+    func.dfg.get_num_return_witnesses(func)?;
+
     // Create the entry point artifact
     let globals_memory_size = brillig
         .globals_memory_size
@@ -58,7 +68,7 @@ pub(crate) fn gen_brillig_for(
 
     let (mut entry_point, stack_start) = BrilligContext::new_entry_point_artifact(
         arguments,
-        FunctionContext::return_values(func),
+        return_parameters,
         func.id(),
         true,
         globals_memory_size,
@@ -75,7 +85,8 @@ pub(crate) fn gen_brillig_for(
                 return Err(InternalError::General {
                     message: format!("Cannot find linked fn {unresolved_fn_label}"),
                     call_stack: CallStack::new(),
-                });
+                }
+                .into());
             }
         };
         entry_point.link_with(artifact);

compiler/noirc_evaluator/src/brillig/brillig_ir.rs

@@ -17,7 +17,7 @@ pub(crate) mod registers;
 
 mod codegen_binary;
 mod codegen_calls;
-mod codegen_control_flow;
+pub(crate) mod codegen_control_flow;
 mod codegen_intrinsic;
 mod codegen_memory;
 mod codegen_stack;

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

@@ -4,11 +4,11 @@ use acvm::acir::circuit::ErrorSelector;
 use noirc_errors::call_stack::CallStackId;
 use std::collections::{BTreeMap, HashMap};
 
+use super::procedures::ProcedureId;
 use crate::ErrorType;
+use crate::brillig::assert_usize;
 use crate::ssa::ir::{basic_block::BasicBlockId, function::FunctionId};
 
-use super::procedures::ProcedureId;
-
 /// Represents a parameter or a return value of an entry point function.
 #[derive(Debug, Clone, Eq, PartialEq, Hash, PartialOrd, Ord)]
 pub(crate) enum BrilligParameter {
@@ -21,6 +21,20 @@ pub(crate) enum BrilligParameter {
     Vector(Vec<BrilligParameter>, SemanticLength),
 }
 
+impl BrilligParameter {
+    /// Computes the size of a parameter if it was flattened
+    pub(crate) fn flattened_size(&self) -> usize {
+        match self {
+            BrilligParameter::SingleAddr(_) => 1,
+            BrilligParameter::Array(item_types, item_count)
+            | BrilligParameter::Vector(item_types, item_count) => {
+                let item_size: usize = item_types.iter().map(|param| param.flattened_size()).sum();
+                assert_usize(item_count.0) * item_size
+            }
+        }
+    }
+}
+
 /// The result of compiling and linking brillig artifacts.
 /// This is ready to run bytecode with attached metadata.
 #[derive(Debug, Default, Clone)]

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

@@ -265,7 +265,7 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
             for (error_variable, error_param) in
                 error_data_items.into_iter().zip(error_data_types.into_iter())
             {
-                let flattened_size = Self::flattened_size(&error_param);
+                let flattened_size = error_param.flattened_size();
                 match error_param {
                     BrilligParameter::SingleAddr(_) => {
                         ctx.store_instruction(
@@ -323,21 +323,9 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
         );
     }
 
-    /// Computes the size of a parameter if it was flattened
-    pub(super) fn flattened_size(param: &BrilligParameter) -> usize {
-        match param {
-            BrilligParameter::SingleAddr(_) => 1,
-            BrilligParameter::Array(item_types, item_count)
-            | BrilligParameter::Vector(item_types, item_count) => {
-                let item_size: usize = item_types.iter().map(Self::flattened_size).sum();
-                assert_usize(item_count.0) * item_size
-            }
-        }
-    }
-
     /// Computes the size of a parameter if it was flattened
     pub(super) fn flattened_tuple_size(tuple: &[BrilligParameter]) -> usize {
-        tuple.iter().map(Self::flattened_size).sum()
+        tuple.iter().map(|param| param.flattened_size()).sum()
     }
 
     /// Computes the size of a parameter if it was flattened
@@ -423,7 +411,7 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
                         BrilligParameter::Vector(..) => unreachable!("ICE: Cannot flatten vectors"),
                     }
 
-                    target_offset += Self::flattened_size(subitem);
+                    target_offset += subitem.flattened_size();
                 }
             }
         } else {

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

@@ -517,7 +517,11 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
     /// Initializes an array, allocating memory on the heap to store its representation and initializing the reference counter to 1.
     pub(crate) fn codegen_initialize_array(&mut self, array: BrilligArray) {
         // Allocate memory for the ref counter and `size` items.
-        let size = array.size.0 + offsets::ARRAY_META_COUNT;
+        let size = array
+            .size
+            .0
+            .checked_add(offsets::ARRAY_META_COUNT)
+            .expect("Array size overflow: array is too large to be allocated in Brillig");
         self.codegen_allocate_immediate_mem(array.pointer, assert_usize(size));
         self.codegen_initialize_rc(array.pointer, 1);
     }

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

@@ -153,7 +153,7 @@ impl<F: AcirField + DebugToString> BrilligContext<F, Stack> {
                 _ => unreachable!("ICE: cannot match variables against arguments"),
             }
 
-            current_calldata_pointer += Self::flattened_size(argument);
+            current_calldata_pointer += argument.flattened_size();
         }
 
         stack_start
@@ -307,7 +307,7 @@ impl<F: AcirField + DebugToString> BrilligContext<F, Stack> {
                         }
                     }
 
-                    source_offset += Self::flattened_size(subitem);
+                    source_offset += subitem.flattened_size();
                 }
             }
         } else {
@@ -382,7 +382,7 @@ impl<F: AcirField + DebugToString> BrilligContext<F, Stack> {
                 }
             }
 
-            return_data_index += Self::flattened_size(return_param);
+            return_data_index += return_param.flattened_size();
         }
 
         let return_pointer = self.make_usize_constant_instruction(return_data_offset.into());

compiler/noirc_evaluator/src/errors.rs

@@ -94,6 +94,10 @@ pub enum RuntimeError {
     },
     #[error("SSA validation failed: {message}")]
     SsaValidationError { message: String, call_stack: CallStack },
+    #[error(
+        "The return value has {num_witnesses} elements which exceeds the limit of {max_witnesses}"
+    )]
+    ReturnLimitExceeded { num_witnesses: usize, max_witnesses: usize, call_stack: CallStack },
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Error)]
@@ -145,7 +149,8 @@ impl RuntimeError {
             | RuntimeError::UnknownReference { call_stack }
             | RuntimeError::RecursionLimit { call_stack, .. }
             | RuntimeError::UnconstrainedCallingConstrained { call_stack, .. }
-            | RuntimeError::SsaValidationError { call_stack, .. } => call_stack,
+            | RuntimeError::SsaValidationError { call_stack, .. }
+            | RuntimeError::ReturnLimitExceeded { call_stack, .. } => call_stack,
         }
     }
 }

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

@@ -170,16 +170,19 @@ impl BrilligTaintedIds {
                 // If the result value is an array, create an empty descendant set for
                 // every element to be accessed further on and record the indices
                 // of the resulting sets for future reference
-                Some(length) => {
+                Some(length)
+                    if length.0 <= crate::ssa::ir::dfg::MAX_ELEMENTS.try_into().unwrap() =>
+                {
                     array_elements.insert(*result, vec![]);
                     for _ in 0..length.0 {
                         array_elements[result].push(results_status.len());
                         results_status
                             .push(ResultStatus::Unconstrained { descendants: HashSet::new() });
                     }
                 }
-                // Otherwise initialize a descendant set with the current value
-                None => {
+                // For very large arrays or non-arrays, treat the whole result as a single value
+                // to avoid memory/time issues when tracking individual elements
+                Some(_) | None => {
                     results_status.push(ResultStatus::Unconstrained {
                         descendants: HashSet::from([*result]),
                     });

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

@@ -3,7 +3,9 @@ use std::{borrow::Cow, sync::Arc};
 use crate::{
     brillig::assert_u32,
     ssa::{
+        RuntimeError,
         function_builder::data_bus::DataBus,
+        ir::function::Function,
         ir::instruction::ArrayOffset,
         opt::pure::{FunctionPurities, Purity},
     },
@@ -166,6 +168,13 @@ impl From<GlobalsGraph> for DataFlowGraph {
     }
 }
 
+/// Maximum number of elements allowed for return values, or for some arrays.
+/// This limit prevents hangings or out-of-memory issues when dealing with very large arrays.
+/// 2^24 = 16,777,216 witnesses.
+/// In practice, the number of witnesses is limited by the CRS size, which is usually around 2^20.
+/// So this limit should not interfere with real use cases.
+pub(crate) const MAX_ELEMENTS: usize = 1 << 24;
+
 impl DataFlowGraph {
     /// Runtime type of the function.
     pub(crate) fn runtime(&self) -> RuntimeType {
@@ -853,6 +862,31 @@ impl DataFlowGraph {
         let results = self.instruction_results(instruction_id);
         results.contains(&return_data)
     }
+
+    /// Computes the number of flattened values returned by the SSA terminator
+    /// Error if it exceeds MAX_ELEMENTS
+    pub(crate) fn get_num_return_witnesses(&self, func: &Function) -> Result<usize, RuntimeError> {
+        if let Some(TerminatorInstruction::Return { return_values, call_stack }) =
+            func.return_instruction()
+        {
+            let num_return_values: usize = return_values
+                .iter()
+                .map(|result_id| self.type_of_value(*result_id).flattened_size().to_usize())
+                .sum();
+
+            if num_return_values > MAX_ELEMENTS {
+                let call_stack = func.dfg.call_stack_data.get_call_stack(*call_stack);
+                return Err(RuntimeError::ReturnLimitExceeded {
+                    num_witnesses: num_return_values,
+                    max_witnesses: MAX_ELEMENTS,
+                    call_stack,
+                });
+            }
+            return Ok(num_return_values);
+        }
+
+        Ok(0)
+    }
 }
 
 impl std::ops::Index<InstructionId> for DataFlowGraph {

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

@@ -186,13 +186,19 @@ impl Function {
     /// None might be returned if the function ends up with all of its block
     /// terminators being `jmp`, `jmpif` or `unreachable`.
     pub(crate) fn returns(&self) -> Option<&[ValueId]> {
-        for block in self.reachable_blocks() {
-            let terminator = self.dfg[block].terminator();
-            if let Some(TerminatorInstruction::Return { return_values, .. }) = terminator {
-                return Some(return_values);
-            }
+        match self.return_instruction()? {
+            TerminatorInstruction::Return { return_values, .. } => Some(return_values),
+            _ => None,
         }
-        None
+    }
+
+    /// Retrieve the return instruction of this function, if any.
+    pub(crate) fn return_instruction(&self) -> Option<&TerminatorInstruction> {
+        self.reachable_blocks().into_iter().find_map(|block| {
+            self.dfg[block]
+                .terminator()
+                .filter(|t| matches!(t, TerminatorInstruction::Return { .. }))
+        })
     }
 
     /// Collects all the reachable blocks of this function.