feat: Brillig heterogeneous memory cells (#5608)

Since memory in brillig is typed, we can store data in a different way
depending on the bit size. This PR makes brillig store memory cells as
field only if they are field elements and as BigUints otherwise,
avoiding conversions when doing integer operations to improve
performance.

---------

Co-authored-by: TomAFrench <tom@tomfren.ch>
Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com>
Co-authored-by: Maxim Vezenov <mvezenov@gmail.com>

noir/noir-repo/acvm-repo/acvm/src/pwg/brillig.rs

@@ -206,13 +206,13 @@ impl<'b, B: BlackBoxFunctionSolver> BrilligSolver<'b, B> {
         for output in brillig.outputs.iter() {
             match output {
                 BrilligOutputs::Simple(witness) => {
-                    insert_value(witness, memory[current_ret_data_idx].value, witness_map)?;
+                    insert_value(witness, memory[current_ret_data_idx].to_field(), witness_map)?;
                     current_ret_data_idx += 1;
                 }
                 BrilligOutputs::Array(witness_arr) => {
                     for witness in witness_arr.iter() {
-                        let value = memory[current_ret_data_idx];
-                        insert_value(witness, value.value, witness_map)?;
+                        let value = &memory[current_ret_data_idx];
+                        insert_value(witness, value.to_field(), witness_map)?;
                         current_ret_data_idx += 1;
                     }
                 }

noir/noir-repo/acvm-repo/brillig_vm/src/arithmetic.rs

@@ -1,16 +1,19 @@
 use acir::brillig::{BinaryFieldOp, BinaryIntOp};
 use acir::FieldElement;
 use num_bigint::BigUint;
-use num_traits::{One, ToPrimitive, Zero};
+use num_traits::ToPrimitive;
+use num_traits::{One, Zero};
 
-use crate::memory::MemoryValue;
+use crate::memory::{MemoryTypeError, MemoryValue};
 
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum BrilligArithmeticError {
     #[error("Bit size for lhs {lhs_bit_size} does not match op bit size {op_bit_size}")]
     MismatchedLhsBitSize { lhs_bit_size: u32, op_bit_size: u32 },
     #[error("Bit size for rhs {rhs_bit_size} does not match op bit size {op_bit_size}")]
     MismatchedRhsBitSize { rhs_bit_size: u32, op_bit_size: u32 },
+    #[error("Integer operation BinaryIntOp::{op:?} is not supported on FieldElement")]
+    IntegerOperationOnField { op: BinaryIntOp },
     #[error("Shift with bit size {op_bit_size} is invalid")]
     InvalidShift { op_bit_size: u32 },
 }
@@ -21,21 +24,19 @@ pub(crate) fn evaluate_binary_field_op(
     lhs: MemoryValue,
     rhs: MemoryValue,
 ) -> Result<MemoryValue, BrilligArithmeticError> {
-    if lhs.bit_size != FieldElement::max_num_bits() {
+    let MemoryValue::Field(a) = lhs else {
         return Err(BrilligArithmeticError::MismatchedLhsBitSize {
-            lhs_bit_size: lhs.bit_size,
+            lhs_bit_size: lhs.bit_size(),
             op_bit_size: FieldElement::max_num_bits(),
         });
-    }
-    if rhs.bit_size != FieldElement::max_num_bits() {
-        return Err(BrilligArithmeticError::MismatchedRhsBitSize {
-            rhs_bit_size: rhs.bit_size,
+    };
+    let MemoryValue::Field(b) = rhs else {
+        return Err(BrilligArithmeticError::MismatchedLhsBitSize {
+            lhs_bit_size: rhs.bit_size(),
             op_bit_size: FieldElement::max_num_bits(),
         });
-    }
+    };
 
-    let a = lhs.value;
-    let b = rhs.value;
     Ok(match op {
         // Perform addition, subtraction, multiplication, and division based on the BinaryOp variant.
         BinaryFieldOp::Add => (a + b).into(),
@@ -62,21 +63,26 @@ pub(crate) fn evaluate_binary_int_op(
     rhs: MemoryValue,
     bit_size: u32,
 ) -> Result<MemoryValue, BrilligArithmeticError> {
-    if lhs.bit_size != bit_size {
-        return Err(BrilligArithmeticError::MismatchedLhsBitSize {
-            lhs_bit_size: lhs.bit_size,
-            op_bit_size: bit_size,
-        });
-    }
-    if rhs.bit_size != bit_size {
-        return Err(BrilligArithmeticError::MismatchedRhsBitSize {
-            rhs_bit_size: rhs.bit_size,
-            op_bit_size: bit_size,
-        });
-    }
+    let lhs = lhs.expect_integer_with_bit_size(bit_size).map_err(|err| match err {
+        MemoryTypeError::MismatchedBitSize { value_bit_size, expected_bit_size } => {
+            BrilligArithmeticError::MismatchedLhsBitSize {
+                lhs_bit_size: value_bit_size,
+                op_bit_size: expected_bit_size,
+            }
+        }
+    })?;
+    let rhs = rhs.expect_integer_with_bit_size(bit_size).map_err(|err| match err {
+        MemoryTypeError::MismatchedBitSize { value_bit_size, expected_bit_size } => {
+            BrilligArithmeticError::MismatchedRhsBitSize {
+                rhs_bit_size: value_bit_size,
+                op_bit_size: expected_bit_size,
+            }
+        }
+    })?;
 
-    let lhs = BigUint::from_bytes_be(&lhs.value.to_be_bytes());
-    let rhs = BigUint::from_bytes_be(&rhs.value.to_be_bytes());
+    if bit_size == FieldElement::max_num_bits() {
+        return Err(BrilligArithmeticError::IntegerOperationOnField { op: *op });
+    }
 
     let bit_modulo = &(BigUint::one() << bit_size);
     let result = match op {
@@ -136,13 +142,11 @@ pub(crate) fn evaluate_binary_int_op(
         }
     };
 
-    let result_as_field = FieldElement::from_be_bytes_reduce(&result.to_bytes_be());
-
     Ok(match op {
         BinaryIntOp::Equals | BinaryIntOp::LessThan | BinaryIntOp::LessThanEquals => {
-            MemoryValue::new(result_as_field, 1)
+            MemoryValue::new_integer(result, 1)
         }
-        _ => MemoryValue::new(result_as_field, bit_size),
+        _ => MemoryValue::new_integer(result, bit_size),
     })
 }
 
@@ -159,13 +163,13 @@ mod tests {
     fn evaluate_u128(op: &BinaryIntOp, a: u128, b: u128, bit_size: u32) -> u128 {
         let result_value = evaluate_binary_int_op(
             op,
-            MemoryValue::new(a.into(), bit_size),
-            MemoryValue::new(b.into(), bit_size),
+            MemoryValue::new_integer(a.into(), bit_size),
+            MemoryValue::new_integer(b.into(), bit_size),
             bit_size,
         )
         .unwrap();
         // Convert back to u128
-        result_value.value.to_u128()
+        result_value.to_field().to_u128()
     }
 
     fn to_negative(a: u128, bit_size: u32) -> u128 {

noir/noir-repo/acvm-repo/brillig_vm/src/black_box.rs

@@ -20,7 +20,7 @@ fn read_heap_array<'a>(memory: &'a Memory, array: &HeapArray) -> &'a [MemoryValu
 /// Extracts the last byte of every value
 fn to_u8_vec(inputs: &[MemoryValue]) -> Vec<u8> {
     let mut result = Vec::with_capacity(inputs.len());
-    for &input in inputs {
+    for input in inputs {
         result.push(input.try_into().unwrap());
     }
     result
@@ -63,7 +63,7 @@ pub(crate) fn evaluate_black_box<Solver: BlackBoxFunctionSolver>(
         BlackBoxOp::Keccakf1600 { message, output } => {
             let state_vec: Vec<u64> = read_heap_vector(memory, message)
                 .iter()
-                .map(|&memory_value| memory_value.try_into().unwrap())
+                .map(|memory_value| memory_value.try_into().unwrap())
                 .collect();
             let state: [u64; 25] = state_vec.try_into().unwrap();
 
@@ -151,7 +151,7 @@ pub(crate) fn evaluate_black_box<Solver: BlackBoxFunctionSolver>(
         }
         BlackBoxOp::PedersenCommitment { inputs, domain_separator, output } => {
             let inputs: Vec<FieldElement> =
-                read_heap_vector(memory, inputs).iter().map(|&x| x.try_into().unwrap()).collect();
+                read_heap_vector(memory, inputs).iter().map(|x| x.try_into().unwrap()).collect();
             let domain_separator: u32 =
                 memory.read(*domain_separator).try_into().map_err(|_| {
                     BlackBoxResolutionError::Failed(
@@ -165,7 +165,7 @@ pub(crate) fn evaluate_black_box<Solver: BlackBoxFunctionSolver>(
         }
         BlackBoxOp::PedersenHash { inputs, domain_separator, output } => {
             let inputs: Vec<FieldElement> =
-                read_heap_vector(memory, inputs).iter().map(|&x| x.try_into().unwrap()).collect();
+                read_heap_vector(memory, inputs).iter().map(|x| x.try_into().unwrap()).collect();
             let domain_separator: u32 =
                 memory.read(*domain_separator).try_into().map_err(|_| {
                     BlackBoxResolutionError::Failed(
@@ -185,7 +185,7 @@ pub(crate) fn evaluate_black_box<Solver: BlackBoxFunctionSolver>(
         BlackBoxOp::BigIntToLeBytes { .. } => todo!(),
         BlackBoxOp::Poseidon2Permutation { message, output, len } => {
             let input = read_heap_vector(memory, message);
-            let input: Vec<FieldElement> = input.iter().map(|&x| x.try_into().unwrap()).collect();
+            let input: Vec<FieldElement> = input.iter().map(|x| x.try_into().unwrap()).collect();
             let len = memory.read(*len).try_into().unwrap();
             let result = solver.poseidon2_permutation(&input, len)?;
             let mut values = Vec::new();
@@ -204,7 +204,7 @@ pub(crate) fn evaluate_black_box<Solver: BlackBoxFunctionSolver>(
                     format!("Expected 16 inputs but encountered {}", &inputs.len()),
                 ));
             }
-            for (i, &input) in inputs.iter().enumerate() {
+            for (i, input) in inputs.iter().enumerate() {
                 message[i] = input.try_into().unwrap();
             }
             let mut state = [0; 8];
@@ -215,7 +215,7 @@ pub(crate) fn evaluate_black_box<Solver: BlackBoxFunctionSolver>(
                     format!("Expected 8 values but encountered {}", &values.len()),
                 ));
             }
-            for (i, &value) in values.iter().enumerate() {
+            for (i, value) in values.iter().enumerate() {
                 state[i] = value.try_into().unwrap();
             }