From 08835f99e11c479cb498b411b15a16305695039f Mon Sep 17 00:00:00 2001 From: Facundo Date: Thu, 14 Mar 2024 15:00:01 +0000 Subject: [PATCH] chore(avm-simulator): make sure we support Map storage (#5207) Brillig expects comparison operations to return `u1`. The AVM returns the type of the input operands. I'm fixing this by adding a cast in the transpiler but I think that longer term we should also return u1 (well, u8) in the AVM. --- avm-transpiler/src/opcodes.rs | 2 +- avm-transpiler/src/transpile.rs | 17 ++++ .../contracts/avm_test_contract/src/main.nr | 23 +++++ .../simulator/src/avm/avm_simulator.test.ts | 84 +++++++++++++++++-- 4 files changed, 120 insertions(+), 6 deletions(-) diff --git a/avm-transpiler/src/opcodes.rs b/avm-transpiler/src/opcodes.rs index 49de1fb366a..e5c01bf3c49 100644 --- a/avm-transpiler/src/opcodes.rs +++ b/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, diff --git a/avm-transpiler/src/transpile.rs b/avm-transpiler/src/transpile.rs index 59cba34d907..b1dadccad57 100644 --- a/avm-transpiler/src/transpile.rs +++ b/avm-transpiler/src/transpile.rs @@ -58,6 +58,10 @@ pub fn brillig_to_avm(brillig: &Brillig) -> Vec { }, ], }); + // 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 { }, ], }); + // 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 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 diff --git a/noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr b/noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr index 2b06b4b8e80..675b8c40867 100644 --- a/noir-projects/noir-contracts/contracts/avm_test_contract/src/main.nr +++ b/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 @@ -32,6 +34,7 @@ contract AvmTest { struct Storage { single: PublicMutable, list: PublicMutable, + map: Map>, } #[aztec(public-vm)] @@ -61,6 +64,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 diff --git a/yarn-project/simulator/src/avm/avm_simulator.test.ts b/yarn-project/simulator/src/avm/avm_simulator.test.ts index c78050e2d7f..b8e4708d2aa 100644 --- a/yarn-project/simulator/src/avm/avm_simulator.test.ts +++ b/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]]); + }); }); }); });