feat(avm-transpiler): implement tags for SET

avm-transpiler/src/instructions.rs

@@ -20,10 +20,9 @@ pub struct AvmInstruction {
     /// The 0th bit corresponds to an instruction's 0th offset arg, 1st to 1st, etc...
     pub indirect: Option<u8>,
 
-    /// Some instructions have a destination or input tag
-    // TODO(4271): add in_tag alongside its support in TS
-    //pub in_tag: Option<AvmTypeTag>,
-    pub dst_tag: Option<AvmTypeTag>,
+    /// Some instructions have a destination xor input tag
+    /// Its usage will depend on the instruction.
+    pub tag: Option<AvmTypeTag>,
 
     /// Different instructions have different numbers of operands
     pub operands: Vec<AvmOperand>,
@@ -35,9 +34,9 @@ impl Display for AvmInstruction {
         if let Some(indirect) = self.indirect {
             write!(f, ", indirect: {}", indirect)?;
         }
-        // TODO(4271): add in_tag alongside its support in TS
-        if let Some(dst_tag) = self.dst_tag {
-            write!(f, ", dst_tag: {}", dst_tag as u8)?;
+        // This will be either inTag or dstTag depending on the operation
+        if let Some(dst_tag) = self.tag {
+            write!(f, ", tag: {}", dst_tag as u8)?;
         }
         if !self.operands.is_empty() {
             write!(f, ", operands: [")?;
@@ -58,10 +57,9 @@ impl AvmInstruction {
         if let Some(indirect) = self.indirect {
             bytes.push(indirect);
         }
-        // TODO(4271): add in_tag alongside its support in TS
-        if let Some(dst_tag) = self.dst_tag {
-            // TODO(4271): make 8 bits when TS supports deserialization of 8 bit flags
-            bytes.extend_from_slice(&(dst_tag as u8).to_be_bytes());
+        // This will be either inTag or dstTag depending on the operation
+        if let Some(tag) = self.tag {
+            bytes.extend_from_slice(&(tag as u8).to_be_bytes());
         }
         for operand in &self.operands {
             bytes.extend_from_slice(&operand.to_be_bytes());
@@ -82,14 +80,14 @@ impl Default for AvmInstruction {
             opcode: AvmOpcode::ADD,
             // TODO(4266): default to Some(0), since all instructions have indirect flag except jumps
             indirect: None,
-            dst_tag: None,
+            tag: None,
             operands: vec![],
         }
     }
 }
 
 /// AVM instructions may include a type tag
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Debug)]
 pub enum AvmTypeTag {
     UNINITIALIZED,
     UINT8,
@@ -105,16 +103,20 @@ pub enum AvmTypeTag {
 /// Constants (as used by the SET instruction) can have size
 /// different from 32 bits
 pub enum AvmOperand {
+    U8 { value: u8 },
+    U16 { value: u16 },
     U32 { value: u32 },
-    // TODO(4267): Support operands of size other than 32 bits (for SET)
+    U64 { value: u64 },
     U128 { value: u128 },
 }
 
 impl Display for AvmOperand {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
+            AvmOperand::U8 { value } => write!(f, " U8:{}", value),
+            AvmOperand::U16 { value } => write!(f, " U16:{}", value),
             AvmOperand::U32 { value } => write!(f, " U32:{}", value),
-            // TODO(4267): Support operands of size other than 32 bits (for SET)
+            AvmOperand::U64 { value } => write!(f, " U64:{}", value),
             AvmOperand::U128 { value } => write!(f, " U128:{}", value),
         }
     }
@@ -123,8 +125,10 @@ impl Display for AvmOperand {
 impl AvmOperand {
     pub fn to_be_bytes(&self) -> Vec<u8> {
         match self {
+            AvmOperand::U8 { value } => value.to_be_bytes().to_vec(),
+            AvmOperand::U16 { value } => value.to_be_bytes().to_vec(),
             AvmOperand::U32 { value } => value.to_be_bytes().to_vec(),
-            // TODO(4267): Support operands of size other than 32 bits (for SET)
+            AvmOperand::U64 { value } => value.to_be_bytes().to_vec(),
             AvmOperand::U128 { value } => value.to_be_bytes().to_vec(),
         }
     }

avm-transpiler/src/transpile.rs

@@ -1,7 +1,7 @@
 use acvm::acir::brillig::Opcode as BrilligOpcode;
 use acvm::acir::circuit::brillig::Brillig;
 
-use acvm::brillig_vm::brillig::{BinaryFieldOp, BinaryIntOp, ValueOrArray};
+use acvm::brillig_vm::brillig::{BinaryFieldOp, BinaryIntOp, MemoryAddress, Value, ValueOrArray};
 
 use crate::instructions::{
     AvmInstruction, AvmOperand, AvmTypeTag, FIRST_OPERAND_INDIRECT, ZEROTH_OPERAND_INDIRECT,
@@ -35,12 +35,10 @@ pub fn brillig_to_avm(brillig: &Brillig) -> Vec<u8> {
                     BinaryFieldOp::Div => AvmOpcode::DIV,
                     BinaryFieldOp::Equals => AvmOpcode::EQ,
                 };
-                // TODO(4268): set in_tag to `field`
                 avm_instrs.push(AvmInstruction {
                     opcode: avm_opcode,
                     indirect: Some(0),
-                    // TODO(4268): TEMPORARY - typescript wireFormat expects this
-                    dst_tag: Some(AvmTypeTag::UINT32),
+                    tag: Some(AvmTypeTag::FIELD),
                     operands: vec![
                         AvmOperand::U32 {
                             value: lhs.to_usize() as u32,
@@ -57,7 +55,7 @@ pub fn brillig_to_avm(brillig: &Brillig) -> Vec<u8> {
             BrilligOpcode::BinaryIntOp {
                 destination,
                 op,
-                bit_size: _, // TODO(4268): support u8..u128 and use in_tag
+                bit_size,
                 lhs,
                 rhs,
             } => {
@@ -79,10 +77,10 @@ pub fn brillig_to_avm(brillig: &Brillig) -> Vec<u8> {
                         brillig_instr
                     ),
                 };
-                // TODO(4268): support u8..u128 and use in_tag
                 avm_instrs.push(AvmInstruction {
                     opcode: avm_opcode,
                     indirect: Some(0),
+                    tag: Some(tag_from_bit_size(bit_size)),
                     operands: vec![
                         AvmOperand::U32 {
                             value: lhs.to_usize() as u32,
@@ -141,23 +139,8 @@ pub fn brillig_to_avm(brillig: &Brillig) -> Vec<u8> {
                     ..Default::default()
                 });
             }
-            BrilligOpcode::Const { destination, value, bit_size:_ } => {
-                avm_instrs.push(AvmInstruction {
-                    opcode: AvmOpcode::SET,
-                    indirect: Some(0),
-                    dst_tag: Some(AvmTypeTag::UINT128),
-                    operands: vec![
-                        // TODO(4267): support u8..u128 and use dst_tag
-                        // value - temporarily as u128 - matching wireFormat in typescript
-                        AvmOperand::U128 {
-                            value: value.to_usize() as u128,
-                        },
-                        // dest offset
-                        AvmOperand::U32 {
-                            value: destination.to_usize() as u32,
-                        },
-                    ],
-                });
+            BrilligOpcode::Const { destination, value, bit_size } => {
+                handle_set(&mut avm_instrs, destination, value, bit_size);
             }
             BrilligOpcode::Mov {
                 destination,
@@ -321,6 +304,84 @@ fn handle_foreign_call(
     });
 }
 
+fn set_for_uint(tag: AvmTypeTag, dest: &MemoryAddress, value: u128) -> AvmInstruction {
+    AvmInstruction {
+        opcode: AvmOpcode::SET,
+        indirect: Some(0),
+        tag: Some(tag),
+        operands: vec![
+            // const
+            match tag {
+                AvmTypeTag::UINT8 => AvmOperand::U8 { value: value as u8 },
+                AvmTypeTag::UINT16 => AvmOperand::U16 {
+                    value: value as u16,
+                },
+                AvmTypeTag::UINT32 => AvmOperand::U32 {
+                    value: value as u32,
+                },
+                AvmTypeTag::UINT64 => AvmOperand::U64 {
+                    value: value as u64,
+                },
+                AvmTypeTag::UINT128 => AvmOperand::U128 { value: value },
+                _ => panic!("Invalid type tag {:?} for set", tag),
+            },
+            // dest offset
+            AvmOperand::U32 {
+                value: dest.to_usize() as u32,
+            },
+        ],
+    }
+}
+
+fn cast(
+    source: &MemoryAddress,
+    destination: &MemoryAddress,
+    destination_tag: AvmTypeTag,
+) -> AvmInstruction {
+    AvmInstruction {
+        opcode: AvmOpcode::CAST,
+        indirect: Some(0),
+        tag: Some(destination_tag),
+        operands: vec![
+            AvmOperand::U32 {
+                value: source.to_usize() as u32,
+            },
+            AvmOperand::U32 {
+                value: destination.to_usize() as u32,
+            },
+        ],
+    }
+}
+
+fn handle_set(
+    avm_instrs: &mut Vec<AvmInstruction>,
+    destination: &MemoryAddress,
+    value: &Value,
+    bit_size: &u32,
+) {
+    let tag = tag_from_bit_size(bit_size);
+
+    if !matches!(tag, AvmTypeTag::FIELD) {
+        avm_instrs.push(set_for_uint(tag, destination, value.to_u128()));
+    } else {
+        // Handling fields is a bit more complex since we cannot fit a field in a single instruction.
+        // We need to split the field into 128-bit chunks and set them individually.
+        let field = value.to_field();
+        if !field.fits_in_u128() {
+            // TODO: if the field doesn't fit in 128 bits, we need scratch space. That's not trivial.
+            // Will this ever happen? ACIR supports up to 126 bit fields.
+            // However, it might be needed _inside_ the unconstrained function.
+            panic!("SET: Field value doesn't fit in 128 bits, that's not supported yet!");
+        }
+        avm_instrs.extend([
+            // SET
+            set_for_uint(AvmTypeTag::UINT128, destination, field.to_u128()),
+            // CAST
+            cast(destination, destination, AvmTypeTag::FIELD),
+        ]);
+    }
+}
+
 /// Compute an array that maps each Brillig pc to an AVM pc.
 /// This must be done before transpiling to properly transpile jump destinations.
 /// This is necessary for two reasons:
@@ -339,6 +400,10 @@ fn map_brillig_pcs_to_avm_pcs(initial_offset: usize, brillig: &Brillig) -> Vec<u
         let num_avm_instrs_for_this_brillig_instr = match &brillig.bytecode[i] {
             BrilligOpcode::Load { .. } => 2,
             BrilligOpcode::Store { .. } => 2,
+            BrilligOpcode::Const { bit_size, .. } => match bit_size {
+                254 => 2, // Field.
+                _ => 1,
+            },
             _ => 1,
         };
         // next Brillig pc will map to an AVM pc offset by the
@@ -347,3 +412,15 @@ fn map_brillig_pcs_to_avm_pcs(initial_offset: usize, brillig: &Brillig) -> Vec<u
     }
     pc_map
 }
+
+fn tag_from_bit_size(bit_size: &u32) -> AvmTypeTag {
+    match bit_size {
+        8 => AvmTypeTag::UINT8,
+        16 => AvmTypeTag::UINT16,
+        32 => AvmTypeTag::UINT32,
+        64 => AvmTypeTag::UINT64,
+        128 => AvmTypeTag::UINT128,
+        254 => AvmTypeTag::FIELD,
+        _ => panic!("The AVM doesn't support integer bit size {:?}", bit_size),
+    }
+}

avm-transpiler/src/utils.rs

@@ -11,7 +11,7 @@ use crate::instructions::AvmInstruction;
 /// assuming the 0th ACIR opcode is the wrapper.
 pub fn extract_brillig_from_acir(opcodes: &Vec<Opcode>) -> &Brillig {
     if opcodes.len() != 1 {
-        panic!("There should only be one brillig opcode");
+        panic!("An AVM program should be contained entirely in only a single ACIR opcode flagged as 'Brillig'");
     }
     let opcode = &opcodes[0];
     match opcode {

yarn-project/noir-contracts/contracts/avm_test_contract/src/main.nr

@@ -19,6 +19,44 @@ contract AvmTest {
         argA + argB
     }
 
+    // TODO: once #4534 goes in, try to just return uint.
+    // Right now it makes noir add extra ACIR range checks.
+    #[aztec(public-vm)]
+    fn setOpcodeUint8() -> pub Field {
+        let mut l = 8 as u8;
+        l = l * 2;
+        l as Field
+    }
+
+    // TODO: Noir currently generates set and compares with different bit sizes.
+    // It also optimizes away the computations.
+    // #[aztec(public-vm)]
+    // fn setOpcodeUint16() -> pub Field {
+    //     let mut l = 60000 as u16;
+    //     l as Field
+    // }
+
+    // #[aztec(public-vm)]
+    // fn setOpcodeUint32() -> pub Field {
+    //     let mut l = 6 as u32;
+    //     l = l * 100000;
+    //     l as Field
+    // }
+
+    // #[aztec(public-vm)]
+    // fn setOpcodeUint64() -> pub Field {
+    //     let mut l = 60000 as u64;
+    //     l as Field
+    // }
+
+    // Field should fit in 128 bits
+    // ACIR only supports fields of up to 126 bits!
+    // Same with internal fields for unconstrained functions, apprently.
+    #[aztec(public-vm)]
+    fn setOpcodeSmallField() -> pub Field {
+        200 as Field
+    }
+
     /************************************************************************
      * AvmContext functions
      ************************************************************************/

yarn-project/simulator/src/avm/avm_memory_types.ts

@@ -335,7 +335,7 @@ export class TaggedMemory {
 
   // Truncates the value to fit the type.
   public static integralFromTag(v: bigint | number, tag: TypeTag): IntegralValue {
-    v = v as bigint;
+    v = BigInt(v);
     switch (tag) {
       case TypeTag.UINT8:
         return new Uint8(v & ((1n << 8n) - 1n));