feat: implement many featuers for the co-brillig rep3 backend

co-noir/co-acvm/src/mpc.rs

@@ -25,11 +25,12 @@ pub trait NoirWitnessExtensionProtocol<F: PrimeField> {
 
     type BrilligDriver: BrilligDriver<F>;
 
-    fn init_brillig_driver(&self) -> Self::BrilligDriver;
+    fn init_brillig_driver(&mut self) -> std::io::Result<Self::BrilligDriver>;
 
-    fn from_brillig_result(
+    fn parse_brillig_result(
+        &mut self,
         brillig_result: Vec<<Self::BrilligDriver as BrilligDriver<F>>::BrilligType>,
-    ) -> Vec<Self::AcvmType>;
+    ) -> eyre::Result<Vec<Self::AcvmType>>;
 
     /// Returns F::zero() as a ACVM-type. The default implementation uses the `Default` trait. If `Default` does not return 0, this function has to be overwritten.
     fn public_zero() -> Self::AcvmType {

co-noir/co-acvm/src/mpc/plain.rs

@@ -30,12 +30,15 @@ impl<F: PrimeField> NoirWitnessExtensionProtocol<F> for PlainAcvmSolver<F> {
 
     type BrilligDriver = PlainBrilligDriver<F>;
 
-    fn init_brillig_driver(&self) -> Self::BrilligDriver {
-        PlainBrilligDriver::default()
+    fn init_brillig_driver(&mut self) -> std::io::Result<Self::BrilligDriver> {
+        Ok(PlainBrilligDriver::default())
     }
 
-    fn from_brillig_result(brillig_result: Vec<PlainBrilligType<F>>) -> Vec<Self::AcvmType> {
-        brillig_result.into_iter().map(|v| v.into_field()).collect()
+    fn parse_brillig_result(
+        &mut self,
+        brillig_result: Vec<PlainBrilligType<F>>,
+    ) -> eyre::Result<Vec<Self::AcvmType>> {
+        Ok(brillig_result.into_iter().map(|v| v.into_field()).collect())
     }
 
     fn is_public_zero(a: &Self::AcvmType) -> bool {

co-noir/co-acvm/src/mpc/rep3.rs

@@ -106,14 +106,26 @@ impl<F: PrimeField> From<ArithmeticShare<F>> for Rep3AcvmType<F> {
 }
 
 impl<F: PrimeField> From<Rep3AcvmType<F>> for Rep3BrilligType<F> {
-    fn from(_val: Rep3AcvmType<F>) -> Self {
-        todo!()
+    fn from(val: Rep3AcvmType<F>) -> Self {
+        match val {
+            Rep3AcvmType::Public(public) => Rep3BrilligType::public_field(public),
+            Rep3AcvmType::Shared(share) => Rep3BrilligType::shared_field(share),
+        }
     }
 }
 
-impl<F: PrimeField> From<Rep3BrilligType<F>> for Rep3AcvmType<F> {
-    fn from(_value: Rep3BrilligType<F>) -> Self {
-        todo!()
+impl<F: PrimeField> Rep3AcvmType<F> {
+    fn from_brillig_type<N: Rep3Network>(
+        value: Rep3BrilligType<F>,
+        io_context: &mut IoContext<N>,
+    ) -> eyre::Result<Self> {
+        match value {
+            Rep3BrilligType::Public(public) => Ok(Rep3AcvmType::Public(public.into_field())),
+            Rep3BrilligType::Shared(shared) => {
+                let shared = Rep3BrilligType::into_arithmetic_share(io_context, shared)?;
+                Ok(Rep3AcvmType::Shared(shared))
+            }
+        }
     }
 }
 
@@ -124,14 +136,20 @@ impl<F: PrimeField, N: Rep3Network> NoirWitnessExtensionProtocol<F> for Rep3Acvm
 
     type AcvmType = Rep3AcvmType<F>;
 
-    type BrilligDriver = Rep3BrilligDriver<F>;
+    type BrilligDriver = Rep3BrilligDriver<F, N>;
 
-    fn init_brillig_driver(&self) -> Self::BrilligDriver {
-        Rep3BrilligDriver::default()
+    fn init_brillig_driver(&mut self) -> std::io::Result<Self::BrilligDriver> {
+        Ok(Rep3BrilligDriver::with_io_context(self.io_context.fork()?))
     }
 
-    fn from_brillig_result(_brillig_result: Vec<Rep3BrilligType<F>>) -> Vec<Self::AcvmType> {
-        todo!()
+    fn parse_brillig_result(
+        &mut self,
+        brillig_result: Vec<Rep3BrilligType<F>>,
+    ) -> eyre::Result<Vec<Self::AcvmType>> {
+        brillig_result
+            .into_iter()
+            .map(|value| Rep3AcvmType::from_brillig_type(value, &mut self.io_context))
+            .collect()
     }
 
     fn is_public_zero(a: &Self::AcvmType) -> bool {

co-noir/co-acvm/src/mpc/shamir.rs

@@ -86,14 +86,20 @@ impl<F: PrimeField> From<ShamirPrimeFieldShare<F>> for ShamirAcvmType<F> {
 }
 
 impl<F: PrimeField> From<ShamirAcvmType<F>> for ShamirBrilligType<F> {
-    fn from(_val: ShamirAcvmType<F>) -> Self {
-        todo!()
+    fn from(val: ShamirAcvmType<F>) -> Self {
+        match val {
+            ShamirAcvmType::Public(public) => ShamirBrilligType::from(public),
+            ShamirAcvmType::Shared(share) => ShamirBrilligType::Shared(share),
+        }
     }
 }
 
 impl<F: PrimeField> From<ShamirBrilligType<F>> for ShamirAcvmType<F> {
-    fn from(_value: ShamirBrilligType<F>) -> Self {
-        todo!()
+    fn from(value: ShamirBrilligType<F>) -> Self {
+        match value {
+            ShamirBrilligType::Public(public) => ShamirAcvmType::Public(public.into_field()),
+            ShamirBrilligType::Shared(shared) => ShamirAcvmType::Shared(shared),
+        }
     }
 }
 
@@ -104,18 +110,26 @@ impl<F: PrimeField, N: ShamirNetwork> NoirWitnessExtensionProtocol<F> for Shamir
 
     type AcvmType = ShamirAcvmType<F>;
 
-    type BrilligDriver = ShamirBrilligDriver<F>;
+    type BrilligDriver = ShamirBrilligDriver<F, N>;
 
-    fn init_brillig_driver(&self) -> Self::BrilligDriver {
-        todo!()
+    fn init_brillig_driver(&mut self) -> std::io::Result<Self::BrilligDriver> {
+        Ok(ShamirBrilligDriver::with_protocol(
+            self.protocol.fork_with_pairs(0)?, // TODO maybe have some pairs here
+        ))
     }
 
     fn public_zero() -> Self::AcvmType {
         Self::AcvmType::default()
     }
 
-    fn from_brillig_result(_brillig_result: Vec<ShamirBrilligType<F>>) -> Vec<Self::AcvmType> {
-        todo!()
+    fn parse_brillig_result(
+        &mut self,
+        brillig_result: Vec<ShamirBrilligType<F>>,
+    ) -> eyre::Result<Vec<Self::AcvmType>> {
+        Ok(brillig_result
+            .into_iter()
+            .map(ShamirAcvmType::from)
+            .collect())
     }
 
     fn is_public_zero(a: &Self::AcvmType) -> bool {

co-noir/co-acvm/src/solver.rs

@@ -152,7 +152,7 @@ where
     }
 
     pub fn new_bn254<P>(
-        driver: T,
+        mut driver: T,
         compiled_program: ProgramArtifact,
         prover_path: P,
     ) -> eyre::Result<Self>
@@ -164,7 +164,7 @@ where
         witness_map[Self::DEFAULT_FUNCTION_INDEX] =
             Self::read_abi_bn254(prover_path, &compiled_program.abi)?;
         let brillig = CoBrilligVM::init(
-            driver.init_brillig_driver(),
+            driver.init_brillig_driver()?,
             compiled_program.bytecode.unconstrained_functions,
         );
         Ok(Self {
@@ -185,7 +185,7 @@ where
     }
 
     pub fn new_bn254_with_witness(
-        driver: T,
+        mut driver: T,
         compiled_program: ProgramArtifact,
         witness: WitnessMap<T::AcvmType>,
     ) -> eyre::Result<Self> {
@@ -194,7 +194,7 @@ where
         witness_map[Self::DEFAULT_FUNCTION_INDEX] = witness;
 
         let brillig = CoBrilligVM::init(
-            driver.init_brillig_driver(),
+            driver.init_brillig_driver()?,
             compiled_program.bytecode.unconstrained_functions,
         );
         Ok(Self {

co-noir/co-acvm/src/solver/brillig_call_solver.rs

@@ -63,7 +63,8 @@ where
                 BrilligInputs::MemoryArray(_) => todo!("memory array calldata TODO"),
             }
         }
-        let brillig_result = T::from_brillig_result(self.brillig.run(id, calldata)?);
+        let brillig_result =
+            T::parse_brillig_result(&mut self.driver, self.brillig.run(id, calldata)?)?;
         self.fill_output(brillig_result, outputs);
         Ok(())
     }

co-noir/co-brillig/Cargo.toml

@@ -23,6 +23,7 @@ noirc-abi.workspace = true
 noirc-artifacts.workspace = true
 num-bigint.workspace = true
 num-traits.workspace = true
+rand.workspace = true
 rayon.workspace = true
 serde.workspace = true
 thiserror.workspace = true

co-noir/co-brillig/src/mpc.rs

@@ -51,7 +51,11 @@ pub trait BrilligDriver<F: PrimeField> {
     /// Casts the provided value to the provided bit size. This includes upcasts
     /// and downcasts between integer types, but also between fields to integers
     /// and vice verca.
-    fn cast(&self, src: Self::BrilligType, bit_size: BitSize) -> eyre::Result<Self::BrilligType>;
+    fn cast(
+        &mut self,
+        src: Self::BrilligType,
+        bit_size: BitSize,
+    ) -> eyre::Result<Self::BrilligType>;
 
     /// Tries to convert the provided value to a `usize`. Returns an error
     /// if it is not possible (e.g., is a shared value).
@@ -80,7 +84,7 @@ pub trait BrilligDriver<F: PrimeField> {
     /// This operation returns an error if the provided inputs
     /// are not the same type.
     fn sub(
-        &mut self,
+        &self,
         lhs: Self::BrilligType,
         rhs: Self::BrilligType,
     ) -> eyre::Result<Self::BrilligType>;
@@ -140,8 +144,11 @@ pub trait BrilligDriver<F: PrimeField> {
     ///
     /// Similar to Noir's brillig-VM, this method compares fields
     /// by casting them `to u128` and compares the integer values.
-    fn lt(&self, lhs: Self::BrilligType, rhs: Self::BrilligType)
-        -> eyre::Result<Self::BrilligType>;
+    fn lt(
+        &mut self,
+        lhs: Self::BrilligType,
+        rhs: Self::BrilligType,
+    ) -> eyre::Result<Self::BrilligType>;
 
     /// Checks whether `lhs <= rhs`. The result
     /// is a brillig integer type with bit size 1 (BitSize = U1).
@@ -152,7 +159,7 @@ pub trait BrilligDriver<F: PrimeField> {
     /// Similar to Noir's brillig-VM, this method compares fields
     /// by casting them `to u128` and compares the integer values.
     fn le(
-        &self,
+        &mut self,
         lhs: Self::BrilligType,
         rhs: Self::BrilligType,
     ) -> eyre::Result<Self::BrilligType> {
@@ -169,7 +176,7 @@ pub trait BrilligDriver<F: PrimeField> {
     /// Similar to Noir's brillig-VM, this method compares fields
     /// by casting them `to u128` and compares the integer values.
     fn gt(
-        &self,
+        &mut self,
         lhs: Self::BrilligType,
         rhs: Self::BrilligType,
     ) -> eyre::Result<Self::BrilligType> {
@@ -186,7 +193,7 @@ pub trait BrilligDriver<F: PrimeField> {
     /// Similar to Noir's brillig-VM, this method compares fields
     /// by casting them `to u128` and compares the integer values.
     fn ge(
-        &self,
+        &mut self,
         lhs: Self::BrilligType,
         rhs: Self::BrilligType,
     ) -> eyre::Result<Self::BrilligType> {
@@ -209,7 +216,7 @@ pub trait BrilligDriver<F: PrimeField> {
     /// Similar to Noir's brillig-VM, this method compares fields
     /// by casting them to `u128` and compares the integer values.
     fn to_radix(
-        &self,
+        &mut self,
         val: Self::BrilligType,
         radix: Self::BrilligType,
         output_size: usize,