From 7856a930233bf67ac66214941d7914dbc20ff085 Mon Sep 17 00:00:00 2001
From: fcarreiro <facundo@aztecprotocol.com>
Date: Fri, 8 Mar 2024 14:09:48 +0000
Subject: [PATCH] Fix some brillig problems

---
 avm-transpiler/src/transpile.rs               |  8 +++++
 .../src/brillig/brillig_gen/brillig_block.rs  |  7 +++--
 .../noirc_evaluator/src/brillig/brillig_ir.rs | 31 +++++++++----------
 3 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/avm-transpiler/src/transpile.rs b/avm-transpiler/src/transpile.rs
index a5ca11a5a470..2a25f8242f9e 100644
--- a/avm-transpiler/src/transpile.rs
+++ b/avm-transpiler/src/transpile.rs
@@ -62,6 +62,7 @@ pub fn brillig_to_avm(brillig: &Brillig) -> Vec<u8> {
                 lhs,
                 rhs,
             } => {
+                assert!(is_integral_bit_size(*bit_size), "BinaryIntOp::{:?} bit_size must be integral, got {:?}", op, bit_size);
                 let avm_opcode = match op {
                     BinaryIntOp::Add => AvmOpcode::ADD,
                     BinaryIntOp::Sub => AvmOpcode::SUB,
@@ -913,6 +914,13 @@ fn map_brillig_pcs_to_avm_pcs(initial_offset: usize, brillig: &Brillig) -> Vec<u
     pc_map
 }
 
+fn is_integral_bit_size(bit_size: u32) -> bool {
+    match bit_size {
+        1 | 8 | 16 | 32 | 64 | 128 => true,
+        _ => false,
+    }
+}
+
 fn tag_from_bit_size(bit_size: u32) -> AvmTypeTag {
     match bit_size {
         1 => AvmTypeTag::UINT8, // temp workaround
diff --git a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
index c04d8475f087..64a2b59daa98 100644
--- a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
+++ b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
@@ -831,14 +831,17 @@ impl<'block> BrilligBlock<'block> {
             _ => unreachable!("ICE: array set on non-array"),
         };
 
-        let one = self.brillig_context.make_usize_constant(1_usize.into());
+        // Here we want to compare the reference count against 1.
+        // I don't think we have the bit_size for the reference count available, so I'm hardcoding it.
+        let bit_size = usize::BITS;
+        let one = self.brillig_context.make_constant(1_usize.into(), bit_size);
         let condition = self.brillig_context.allocate_register();
 
         self.brillig_context.binary_instruction(
             reference_count,
             one,
             condition,
-            BrilligBinaryOp::Field { op: BinaryFieldOp::Equals },
+            BrilligBinaryOp::Integer { op: BinaryIntOp::Equals, bit_size },
         );
 
         self.brillig_context.branch_instruction(condition, |ctx, cond| {
diff --git a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir.rs b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir.rs
index 662dc074d980..57050cafaa50 100644
--- a/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir.rs
+++ b/noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir.rs
@@ -27,7 +27,6 @@ use acvm::{
     FieldElement,
 };
 use debug_show::DebugShow;
-use num_bigint::BigUint;
 
 /// The Brillig VM does not apply a limit to the memory address space,
 /// As a convention, we take use 64 bits. This means that we assume that
@@ -217,7 +216,10 @@ impl BrilligContext {
             array_ptr,
             index,
             index_of_element_in_memory,
-            BrilligBinaryOp::Field { op: BinaryFieldOp::Add },
+            BrilligBinaryOp::Integer {
+                op: BinaryIntOp::Add,
+                bit_size: BRILLIG_MEMORY_ADDRESSING_BIT_SIZE,
+            },
         );
 
         self.load_instruction(result, index_of_element_in_memory);
@@ -239,7 +241,10 @@ impl BrilligContext {
             array_ptr,
             index,
             index_of_element_in_memory,
-            BrilligBinaryOp::Field { op: BinaryFieldOp::Add },
+            BrilligBinaryOp::Integer {
+                op: BinaryIntOp::Add,
+                bit_size: BRILLIG_MEMORY_ADDRESSING_BIT_SIZE,
+            },
         );
 
         self.store_instruction(index_of_element_in_memory, value);
@@ -744,20 +749,12 @@ impl BrilligContext {
             value_to_truncate.bit_size
         );
 
-        let mask = BigUint::from(2_u32).pow(bit_size) - BigUint::from(1_u32);
-        let mask_constant = self.make_constant(
-            FieldElement::from_be_bytes_reduce(&mask.to_bytes_be()).into(),
-            value_to_truncate.bit_size,
-        );
-
-        self.binary_instruction(
-            value_to_truncate.address,
-            mask_constant,
-            destination_of_truncated_value.address,
-            BrilligBinaryOp::Integer { op: BinaryIntOp::And, bit_size: value_to_truncate.bit_size },
-        );
-
-        self.deallocate_register(mask_constant);
+        // We cast back and forth to ensure that the value is truncated.
+        let intermediate_register =
+            SingleAddrVariable { address: self.allocate_register(), bit_size };
+        self.cast_instruction(intermediate_register, value_to_truncate);
+        self.cast_instruction(destination_of_truncated_value, intermediate_register);
+        self.deallocate_register(intermediate_register.address);
     }
 
     /// Emits a stop instruction