chore: allow injecting custom brillig stdlib implementations in ACIRgen

compiler/noirc_evaluator/src/acir/acir_context/brillig_call.rs

@@ -13,10 +13,32 @@ use iter_extended::{try_vecmap, vecmap};
 use crate::brillig::brillig_ir::artifact::GeneratedBrillig;
 use crate::errors::{InternalError, RuntimeError};
 
-use super::generated_acir::BrilligStdlibFunc;
+use super::generated_acir::{BrilligStdlibFunc, PLACEHOLDER_BRILLIG_INDEX};
 use super::{AcirContext, AcirDynamicArray, AcirType, AcirValue, AcirVar};
 
 impl<F: AcirField, B: BlackBoxFunctionSolver<F>> AcirContext<F, B> {
+    /// Generates a brillig call to a handwritten section of brillig bytecode.
+    pub(crate) fn stdlib_brillig_call(
+        &mut self,
+        predicate: AcirVar,
+        brillig_stdlib_func: BrilligStdlibFunc,
+        stdlib_func_bytecode: &GeneratedBrillig<F>,
+        inputs: Vec<AcirValue>,
+        outputs: Vec<AcirType>,
+        attempt_execution: bool,
+    ) -> Result<Vec<AcirValue>, RuntimeError> {
+        self.brillig_call(
+            predicate,
+            stdlib_func_bytecode,
+            inputs,
+            outputs,
+            attempt_execution,
+            false,
+            PLACEHOLDER_BRILLIG_INDEX,
+            Some(brillig_stdlib_func),
+        )
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn brillig_call(
         &mut self,

compiler/noirc_evaluator/src/acir/acir_context/generated_acir/brillig_directive.rs

@@ -4,10 +4,51 @@ use acvm::acir::{
         BinaryFieldOp, BinaryIntOp, BitSize, HeapVector, IntegerBitSize, MemoryAddress,
         Opcode as BrilligOpcode,
     },
+    circuit::brillig::BrilligFunctionId,
 };
 
 use crate::brillig::brillig_ir::artifact::GeneratedBrillig;
 
+/// Brillig calls such as for the Brillig std lib are resolved only after code generation is finished.
+/// This index should be used when adding a Brillig call during code generation.
+/// Code generation should then keep track of that unresolved call opcode which will be resolved with the
+/// correct function index after code generation.
+pub(crate) const PLACEHOLDER_BRILLIG_INDEX: BrilligFunctionId = BrilligFunctionId(0);
+
+#[derive(Debug, Clone)]
+pub(crate) struct BrilligStdLib<F> {
+    pub(crate) invert: GeneratedBrillig<F>,
+    pub(crate) quotient: GeneratedBrillig<F>,
+    pub(crate) to_le_bytes: GeneratedBrillig<F>,
+}
+
+impl<F: AcirField> Default for BrilligStdLib<F> {
+    fn default() -> Self {
+        Self {
+            invert: directive_invert(),
+            quotient: directive_quotient(),
+            to_le_bytes: directive_to_radix(),
+        }
+    }
+}
+
+impl<F: AcirField> BrilligStdLib<F> {
+    pub(crate) fn get_code(&self, func: BrilligStdlibFunc) -> &GeneratedBrillig<F> {
+        match func {
+            BrilligStdlibFunc::Inverse => &self.invert,
+            BrilligStdlibFunc::Quotient => &self.quotient,
+            BrilligStdlibFunc::ToLeBytes => &self.to_le_bytes,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
+pub(crate) enum BrilligStdlibFunc {
+    Inverse,
+    Quotient,
+    ToLeBytes,
+}
+
 /// Generates brillig bytecode which computes the inverse of its input if not null, and zero else.
 pub(crate) fn directive_invert<F: AcirField>() -> GeneratedBrillig<F> {
     //  We generate the following code:

compiler/noirc_evaluator/src/acir/acir_context/generated_acir/mod.rs

@@ -31,11 +31,7 @@ use num_bigint::BigUint;
 
 mod brillig_directive;
 
-/// Brillig calls such as for the Brillig std lib are resolved only after code generation is finished.
-/// This index should be used when adding a Brillig call during code generation.
-/// Code generation should then keep track of that unresolved call opcode which will be resolved with the
-/// correct function index after code generation.
-pub(super) const PLACEHOLDER_BRILLIG_INDEX: BrilligFunctionId = BrilligFunctionId(0);
+pub(crate) use brillig_directive::{BrilligStdLib, BrilligStdlibFunc, PLACEHOLDER_BRILLIG_INDEX};
 
 #[derive(Debug, Default)]
 /// The output of the Acir-gen pass, which should only be produced for entry point Acir functions
@@ -97,23 +93,6 @@ pub(crate) type BrilligOpcodeToLocationsMap = BTreeMap<BrilligOpcodeLocation, Ca
 
 pub(crate) type BrilligProcedureRangeMap = BTreeMap<ProcedureDebugId, (usize, usize)>;
 
-#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
-pub(crate) enum BrilligStdlibFunc {
-    Inverse,
-    Quotient,
-    ToLeBytes,
-}
-
-impl BrilligStdlibFunc {
-    pub(crate) fn get_generated_brillig<F: AcirField>(&self) -> GeneratedBrillig<F> {
-        match self {
-            BrilligStdlibFunc::Inverse => brillig_directive::directive_invert(),
-            BrilligStdlibFunc::Quotient => brillig_directive::directive_quotient(),
-            BrilligStdlibFunc::ToLeBytes => brillig_directive::directive_to_radix(),
-        }
-    }
-}
-
 impl<F: AcirField> GeneratedAcir<F> {
     /// Returns the current witness index.
     pub(crate) fn current_witness_index(&self) -> Witness {

compiler/noirc_evaluator/src/acir/acir_context/mod.rs

@@ -37,16 +37,16 @@ use super::{
     types::{AcirType, AcirVar},
 };
 use big_int::BigIntContext;
-use generated_acir::PLACEHOLDER_BRILLIG_INDEX;
 
-pub(crate) use generated_acir::{BrilligStdlibFunc, GeneratedAcir};
+pub(crate) use generated_acir::{BrilligStdLib, BrilligStdlibFunc, GeneratedAcir};
 
 #[derive(Debug, Default)]
 /// Context object which holds the relationship between
 /// `Variables`(AcirVar) and types such as `Expression` and `Witness`
 /// which are placed into ACIR.
 pub(crate) struct AcirContext<F: AcirField, B: BlackBoxFunctionSolver<F>> {
     pub(super) blackbox_solver: B,
+    brillig_stdlib: BrilligStdLib<F>,
 
     vars: HashMap<AcirVar, AcirVarData<F>>,
 
@@ -70,6 +70,19 @@ pub(crate) struct AcirContext<F: AcirField, B: BlackBoxFunctionSolver<F>> {
 }
 
 impl<F: AcirField, B: BlackBoxFunctionSolver<F>> AcirContext<F, B> {
+    pub(super) fn new(brillig_stdlib: BrilligStdLib<F>, blackbox_solver: B) -> Self {
+        AcirContext {
+            brillig_stdlib,
+            blackbox_solver,
+            vars: Default::default(),
+            constant_witnesses: Default::default(),
+            acir_ir: Default::default(),
+            big_int_ctx: Default::default(),
+            expression_width: Default::default(),
+            warnings: Default::default(),
+        }
+    }
+
     pub(crate) fn set_expression_width(&mut self, expression_width: ExpressionWidth) {
         self.expression_width = expression_width;
     }
@@ -273,18 +286,13 @@ impl<F: AcirField, B: BlackBoxFunctionSolver<F>> AcirContext<F, B> {
             return Ok(inverted_var);
         }
 
-        // Compute the inverse with brillig code
-        let inverse_code = BrilligStdlibFunc::Inverse.get_generated_brillig();
-
-        let results = self.brillig_call(
+        let results = self.stdlib_brillig_call(
             predicate,
-            &inverse_code,
+            BrilligStdlibFunc::Inverse,
+            &self.brillig_stdlib.get_code(BrilligStdlibFunc::Inverse).clone(),
             vec![AcirValue::Var(var, AcirType::field())],
             vec![AcirType::field()],
             true,
-            false,
-            PLACEHOLDER_BRILLIG_INDEX,
-            Some(BrilligStdlibFunc::Inverse),
         )?;
         let inverted_var = Self::expect_one_var(results);
 
@@ -799,18 +807,16 @@ impl<F: AcirField, B: BlackBoxFunctionSolver<F>> AcirContext<F, B> {
         };
 
         let [q_value, r_value]: [AcirValue; 2] = self
-            .brillig_call(
+            .stdlib_brillig_call(
                 predicate,
-                &BrilligStdlibFunc::Quotient.get_generated_brillig(),
+                BrilligStdlibFunc::Quotient,
+                &self.brillig_stdlib.get_code(BrilligStdlibFunc::Quotient).clone(),
                 vec![
                     AcirValue::Var(lhs, AcirType::unsigned(bit_size)),
                     AcirValue::Var(rhs, AcirType::unsigned(bit_size)),
                 ],
                 vec![AcirType::unsigned(max_q_bits), AcirType::unsigned(max_rhs_bits)],
                 true,
-                false,
-                PLACEHOLDER_BRILLIG_INDEX,
-                Some(BrilligStdlibFunc::Quotient),
             )?
             .try_into()
             .expect("quotient only returns two values");

compiler/noirc_evaluator/src/acir/mod.rs

@@ -55,11 +55,13 @@ use crate::ssa::{
     ssa_gen::Ssa,
 };
 pub(crate) use acir_context::GeneratedAcir;
-use acir_context::{AcirContext, BrilligStdlibFunc, power_of_two};
+use acir_context::{AcirContext, BrilligStdLib, BrilligStdlibFunc, power_of_two};
 use types::{AcirType, AcirVar};
 
 #[derive(Default)]
-struct SharedContext<F> {
+struct SharedContext<F: AcirField> {
+    brillig_stdlib: BrilligStdLib<F>,
+
     /// Final list of Brillig functions which will be part of the final program
     /// This is shared across `Context` structs as we want one list of Brillig
     /// functions across all ACIR artifacts
@@ -122,14 +124,14 @@ impl<F: AcirField> SharedContext<F> {
         {
             self.add_call_to_resolve(func_id, (opcode_location, generated_pointer));
         } else {
-            let code = brillig_stdlib_func.get_generated_brillig();
+            let code = self.brillig_stdlib.get_code(*brillig_stdlib_func);
             let generated_pointer = self.new_generated_pointer();
             self.insert_generated_brillig_stdlib(
                 *brillig_stdlib_func,
                 generated_pointer,
                 func_id,
                 opcode_location,
-                code,
+                code.clone(),
             );
         }
     }
@@ -224,9 +226,10 @@ impl<'a> Context<'a> {
         shared_context: &'a mut SharedContext<FieldElement>,
         expression_width: ExpressionWidth,
         brillig: &'a Brillig,
+        brillig_stdlib: BrilligStdLib<FieldElement>,
         brillig_options: &'a BrilligOptions,
     ) -> Context<'a> {
-        let mut acir_context = AcirContext::default();
+        let mut acir_context = AcirContext::new(brillig_stdlib, Bn254BlackBoxSolver::default());
         acir_context.set_expression_width(expression_width);
         let current_side_effects_enabled_var = acir_context.add_constant(FieldElement::one());
 
@@ -2760,7 +2763,7 @@ mod test {
             },
             circuit::{
                 ExpressionWidth, Opcode, OpcodeLocation,
-                brillig::{BrilligBytecode, BrilligFunctionId},
+                brillig::BrilligFunctionId,
                 opcodes::{AcirFunctionId, BlackBoxFuncCall},
             },
             native_types::{Witness, WitnessMap},
@@ -2773,8 +2776,8 @@ mod test {
     use std::collections::BTreeMap;
 
     use crate::{
-        acir::BrilligStdlibFunc,
-        brillig::{Brillig, BrilligOptions},
+        acir::{BrilligStdlibFunc, acir_context::BrilligStdLib, ssa::codegen_acir},
+        brillig::{Brillig, BrilligOptions, brillig_ir::artifact::GeneratedBrillig},
         ssa::{
             function_builder::FunctionBuilder,
             ir::{
@@ -3595,16 +3598,6 @@ mod test {
         }";
         let ssa = Ssa::from_str(src).unwrap();
 
-        let (acir_functions, mut brillig_functions, _, _) = ssa
-            .into_acir(&Brillig::default(), &BrilligOptions::default(), ExpressionWidth::default())
-            .expect("Should compile manually written SSA into ACIR");
-
-        assert_eq!(acir_functions.len(), 1);
-        // [`directive_quotient`, `directive_invert`]
-        assert_eq!(brillig_functions.len(), 2);
-
-        let main = &acir_functions[0];
-
         // Here we're attempting to perform a truncation of a `Field` type into 32 bits. We then do a euclidean
         // division `a/b` with `a` and `b` taking the values:
         //
@@ -3635,8 +3628,8 @@ mod test {
 
         // This brillig function replaces the standard implementation of `directive_quotient` with
         // an implementation which returns `(malicious_q, malicious_r)`.
-        let malicious_quotient = BrilligBytecode {
-            bytecode: vec![
+        let malicious_quotient = GeneratedBrillig {
+            byte_code: vec![
                 BrilligOpcode::Const {
                     destination: MemoryAddress::direct(10),
                     bit_size: BitSize::Integer(IntegerBitSize::U32),
@@ -3664,15 +3657,34 @@ mod test {
                     },
                 },
             ],
+            name: "malicious_directive_quotient".to_string(),
+            ..Default::default()
         };
-        let malicious_brillig = [malicious_quotient, brillig_functions.remove(1)];
+
+        let malicious_brillig_stdlib =
+            BrilligStdLib { quotient: malicious_quotient, ..BrilligStdLib::default() };
+
+        let (acir_functions, brillig_functions, _, _) = codegen_acir(
+            ssa,
+            &Brillig::default(),
+            malicious_brillig_stdlib,
+            &BrilligOptions::default(),
+            ExpressionWidth::default(),
+        )
+        .expect("Should compile manually written SSA into ACIR");
+
+        assert_eq!(acir_functions.len(), 1);
+        // [`malicious_directive_quotient`, `directive_invert`]
+        assert_eq!(brillig_functions.len(), 2);
+
+        let main = &acir_functions[0];
 
         let initial_witness = WitnessMap::from(BTreeMap::from([(Witness(0), input)]));
         let mut acvm = ACVM::new(
             &StubbedBlackBoxSolver(true),
             main.opcodes(),
             initial_witness,
-            &malicious_brillig,
+            &brillig_functions,
             &[],
         );