chore(avm-simulator): make sure we support Map storage

avm-transpiler/src/opcodes.rs

@@ -1,6 +1,6 @@
 /// All AVM opcodes
 /// Keep updated with TS and yellow paper!
-#[derive(PartialEq, Copy, Clone)]
+#[derive(PartialEq, Copy, Clone, Debug)]
 pub enum AvmOpcode {
     // Compute
     ADD,

avm-transpiler/src/transpile.rs

@@ -58,6 +58,10 @@ pub fn brillig_to_avm(brillig: &Brillig) -> Vec<u8> {
                         },
                     ],
                 });
+                // Brillig currently expects comparison instructions to return an u1 (for us, an u8).
+                if avm_opcode == AvmOpcode::EQ {
+                    avm_instrs.push(generate_cast_instruction(destination.to_usize() as u32, destination.to_usize() as u32, AvmTypeTag::UINT8));
+                }
             }
             BrilligOpcode::BinaryIntOp {
                 destination,
@@ -114,6 +118,10 @@ pub fn brillig_to_avm(brillig: &Brillig) -> Vec<u8> {
                         },
                     ],
                 });
+                // Brillig currently expects comparison instructions to return an u1 (for us, an u8).
+                if avm_opcode == AvmOpcode::EQ || avm_opcode == AvmOpcode::LT || avm_opcode == AvmOpcode::LTE {
+                    avm_instrs.push(generate_cast_instruction(destination.to_usize() as u32, destination.to_usize() as u32, AvmTypeTag::UINT8));
+                }
             }
             BrilligOpcode::CalldataCopy { destination_address, size, offset } => {
                 avm_instrs.push(AvmInstruction {
@@ -1005,6 +1013,15 @@ fn map_brillig_pcs_to_avm_pcs(initial_offset: usize, brillig: &Brillig) -> Vec<u
     for i in 0..brillig.bytecode.len() - 1 {
         let num_avm_instrs_for_this_brillig_instr = match &brillig.bytecode[i] {
             BrilligOpcode::Const { bit_size: 254, .. } => 2,
+            // Brillig currently expects comparison instructions to return an u1 (for us, an u8).
+            BrilligOpcode::BinaryIntOp {
+                op: BinaryIntOp::Equals | BinaryIntOp::LessThan | BinaryIntOp::LessThanEquals,
+                ..
+            } => 2,
+            BrilligOpcode::BinaryFieldOp {
+                op: BinaryFieldOp::Equals,
+                ..
+            } => 2,
             _ => 1,
         };
         // next Brillig pc will map to an AVM pc offset by the

noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr

@@ -21,9 +21,11 @@ contract AvmTest {
     use crate::Note;
 
     // Libs
+    use dep::aztec::prelude::Map;
     use dep::aztec::state_vars::PublicMutable;
     use dep::aztec::protocol_types::{address::{AztecAddress, EthAddress}, constants::L1_TO_L2_MESSAGE_LENGTH};
     use dep::aztec::protocol_types::abis::function_selector::FunctionSelector;
+    use dep::aztec::protocol_types::{hash::pedersen_hash, traits::{ToField}};
     use dep::compressed_string::CompressedString;
 
     // avm lib
@@ -35,6 +37,7 @@ contract AvmTest {
     struct Storage {
         single: PublicMutable<Field>,
         list: PublicMutable<Note>,
+        map: Map<AztecAddress, PublicMutable<u32>>,
     }
 
     #[aztec(public-vm)]
@@ -64,6 +67,26 @@ contract AvmTest {
         note.serialize()
     }
 
+    #[aztec(public-vm)]
+    fn setStorageMap(to: AztecAddress, amount: u32) -> pub Field {
+        storage.map.at(to).write(amount);
+        // returns storage slot for key
+        pedersen_hash([storage.map.storage_slot, to.to_field()], 0)
+    }
+
+    #[aztec(public-vm)]
+    fn addStorageMap(to: AztecAddress, amount: u32) -> pub Field {
+        let new_balance = storage.map.at(to).read().add(amount);
+        storage.map.at(to).write(new_balance);
+        // returns storage slot for key
+        pedersen_hash([storage.map.storage_slot, to.to_field()], 0)
+    }
+
+    #[aztec(public-vm)]
+    fn readStorageMap(address: AztecAddress) -> pub u32 {
+        storage.map.at(address).read()
+    }
+
     #[aztec(public-vm)]
     fn addArgsReturn(argA: Field, argB: Field) -> pub Field {
         argA + argB

yarn-project/simulator/src/avm/avm_simulator.test.ts

@@ -499,7 +499,7 @@ describe('AVM simulator', () => {
     });
 
     describe('Storage accesses', () => {
-      it('Should set a single value in storage', async () => {
+      it('Should set value in storage (single)', async () => {
         const slot = 1n;
         const address = AztecAddress.fromField(new Fr(420));
         const value = new Fr(88);
@@ -525,7 +525,7 @@ describe('AVM simulator', () => {
         expect(slotTrace).toEqual([value]);
       });
 
-      it('Should read a single value in storage', async () => {
+      it('Should read value in storage (single)', async () => {
         const slot = 1n;
         const value = new Fr(12345);
         const address = AztecAddress.fromField(new Fr(420));
@@ -551,7 +551,7 @@ describe('AVM simulator', () => {
         expect(slotTrace).toEqual([value]);
       });
 
-      it('Should set and read a value from storage', async () => {
+      it('Should set and read a value from storage (single)', async () => {
         const slot = 1n;
         const value = new Fr(12345);
         const address = AztecAddress.fromField(new Fr(420));
@@ -578,7 +578,7 @@ describe('AVM simulator', () => {
         expect(slotWriteTrace).toEqual([value]);
       });
 
-      it('Should set multiple values in storage', async () => {
+      it('Should set a value in storage (list)', async () => {
         const slot = 2n;
         const sender = AztecAddress.fromField(new Fr(1));
         const address = AztecAddress.fromField(new Fr(420));
@@ -603,7 +603,7 @@ describe('AVM simulator', () => {
         expect(storageTrace.get(slot + 1n)).toEqual([calldata[1]]);
       });
 
-      it('Should read multiple values in storage', async () => {
+      it('Should read a value in storage (list)', async () => {
         const slot = 2n;
         const address = AztecAddress.fromField(new Fr(420));
         const values = [new Fr(1), new Fr(2)];
@@ -630,6 +630,80 @@ describe('AVM simulator', () => {
         expect(storageTrace.get(slot)).toEqual([values[0]]);
         expect(storageTrace.get(slot + 1n)).toEqual([values[1]]);
       });
+
+      it('Should set a value in storage (map)', async () => {
+        const address = AztecAddress.fromField(new Fr(420));
+        const value = new Fr(12345);
+        const calldata = [address.toField(), value];
+
+        const context = initContext({
+          env: initExecutionEnvironment({ address, calldata, storageAddress: address }),
+        });
+        const bytecode = getAvmTestContractBytecode('avm_setStorageMap');
+        const results = await new AvmSimulator(context).executeBytecode(bytecode);
+
+        expect(results.reverted).toBe(false);
+        // returns the storage slot for modified key
+        const slotNumber = results.output[0].toBigInt();
+
+        const worldState = context.persistableState.flush();
+        const storageSlot = worldState.currentStorageValue.get(address.toBigInt())!;
+        expect(storageSlot.get(slotNumber)).toEqual(value);
+
+        // Tracing
+        const storageTrace = worldState.storageWrites.get(address.toBigInt())!;
+        expect(storageTrace.get(slotNumber)).toEqual([value]);
+      });
+
+      it('Should read-add-set a value in storage (map)', async () => {
+        const address = AztecAddress.fromField(new Fr(420));
+        const value = new Fr(12345);
+        const calldata = [address.toField(), value];
+
+        const context = initContext({
+          env: initExecutionEnvironment({ address, calldata, storageAddress: address }),
+        });
+        const bytecode = getAvmTestContractBytecode('avm_addStorageMap');
+        const results = await new AvmSimulator(context).executeBytecode(bytecode);
+
+        expect(results.reverted).toBe(false);
+        // returns the storage slot for modified key
+        const slotNumber = results.output[0].toBigInt();
+
+        const worldState = context.persistableState.flush();
+        const storageSlot = worldState.currentStorageValue.get(address.toBigInt())!;
+        expect(storageSlot.get(slotNumber)).toEqual(value);
+
+        // Tracing
+        const storageReadTrace = worldState.storageReads.get(address.toBigInt())!;
+        expect(storageReadTrace.get(slotNumber)).toEqual([new Fr(0)]);
+        const storageWriteTrace = worldState.storageWrites.get(address.toBigInt())!;
+        expect(storageWriteTrace.get(slotNumber)).toEqual([value]);
+      });
+
+      it('Should read value in storage (map)', async () => {
+        const value = new Fr(12345);
+        const address = AztecAddress.fromField(new Fr(420));
+        const calldata = [address.toField()];
+
+        const context = initContext({
+          env: initExecutionEnvironment({ calldata, address, storageAddress: address }),
+        });
+        jest
+          .spyOn(context.persistableState.hostStorage.publicStateDb, 'storageRead')
+          .mockReturnValue(Promise.resolve(value));
+        const bytecode = getAvmTestContractBytecode('avm_readStorageMap');
+        const results = await new AvmSimulator(context).executeBytecode(bytecode);
+
+        // Get contract function artifact
+        expect(results.reverted).toBe(false);
+        expect(results.output).toEqual([value]);
+
+        // Tracing
+        const worldState = context.persistableState.flush();
+        const storageTrace = worldState.storageReads.get(address.toBigInt())!;
+        expect([...storageTrace.values()]).toEqual([[value]]);
+      });
     });
   });
 });