From 54aee58952b2433ccad83f1b5fc3088957b10fbb Mon Sep 17 00:00:00 2001 From: Facundo Date: Wed, 27 Mar 2024 16:06:07 +0100 Subject: [PATCH] refactor(avm): unify noir macros flow (#5461) Step towards integrating with initializers, etc. However in the future we might want to separate private and public-vm, many things might end up being different enough (see IFs in code). --- noir-projects/aztec-nr/aztec/src/context.nr | 14 +-- .../aztec-nr/aztec/src/context/avm_context.nr | 10 +- .../aztec-nr/aztec/src/context/interface.nr | 4 +- .../contracts/avm_test_contract/src/main.nr | 14 ++- noir/noir-repo/aztec_macros/src/lib.rs | 15 +-- .../aztec_macros/src/transforms/functions.rs | 97 ++++++++----------- .../end-to-end/src/e2e_avm_simulator.test.ts | 5 + .../src/avm/avm_execution_environment.ts | 4 +- .../simulator/src/avm/avm_simulator.test.ts | 12 +-- 9 files changed, 91 insertions(+), 84 deletions(-) diff --git a/noir-projects/aztec-nr/aztec/src/context.nr b/noir-projects/aztec-nr/aztec/src/context.nr index 4c87c33339d..ce4c7eea3c2 100644 --- a/noir-projects/aztec-nr/aztec/src/context.nr +++ b/noir-projects/aztec-nr/aztec/src/context.nr @@ -9,28 +9,28 @@ mod interface; use interface::ContextInterface; use private_context::PrivateContext; use public_context::PublicContext; -use avm_context::AVMContext; +use avm_context::AvmContext; struct Context { private: Option<&mut PrivateContext>, public: Option<&mut PublicContext>, - public_vm: Option<&mut AVMContext>, + avm: Option<&mut AvmContext>, } impl Context { pub fn private(context: &mut PrivateContext) -> Context { - Context { private: Option::some(context), public: Option::none(), public_vm: Option::none() } + Context { private: Option::some(context), public: Option::none(), avm: Option::none() } } pub fn public(context: &mut PublicContext) -> Context { - Context { public: Option::some(context), private: Option::none(), public_vm: Option::none() } + Context { public: Option::some(context), private: Option::none(), avm: Option::none() } } - pub fn public_vm(context: &mut AVMContext) -> Context { - Context { public_vm: Option::some(context), public: Option::none(), private: Option::none() } + pub fn avm(context: &mut AvmContext) -> Context { + Context { avm: Option::some(context), public: Option::none(), private: Option::none() } } pub fn none() -> Context { - Context { public: Option::none(), private: Option::none(), public_vm: Option::none() } + Context { public: Option::none(), private: Option::none(), avm: Option::none() } } } diff --git a/noir-projects/aztec-nr/aztec/src/context/avm_context.nr b/noir-projects/aztec-nr/aztec/src/context/avm_context.nr index 34c3e964000..69d34d8d622 100644 --- a/noir-projects/aztec-nr/aztec/src/context/avm_context.nr +++ b/noir-projects/aztec-nr/aztec/src/context/avm_context.nr @@ -7,13 +7,13 @@ use crate::context::inputs::avm_context_inputs::AvmContextInputs; use crate::context::interface::ContextInterface; use crate::context::interface::PublicContextInterface; -struct AVMContext { +struct AvmContext { inputs: AvmContextInputs, } -impl AVMContext { +impl AvmContext { pub fn new(inputs: AvmContextInputs) -> Self { - AVMContext { inputs } + AvmContext { inputs } } pub fn origin(self) -> AztecAddress { @@ -73,7 +73,7 @@ impl AVMContext { } } -impl PublicContextInterface for AVMContext { +impl PublicContextInterface for AvmContext { fn block_number(self) -> Field { block_number() } @@ -168,7 +168,7 @@ impl PublicContextInterface for AVMContext { } } -impl ContextInterface for AVMContext { +impl ContextInterface for AvmContext { fn push_new_note_hash(&mut self, note_hash: Field) { emit_note_hash(note_hash); } diff --git a/noir-projects/aztec-nr/aztec/src/context/interface.nr b/noir-projects/aztec-nr/aztec/src/context/interface.nr index c536ad02c08..29260eb5018 100644 --- a/noir-projects/aztec-nr/aztec/src/context/interface.nr +++ b/noir-projects/aztec-nr/aztec/src/context/interface.nr @@ -17,8 +17,8 @@ trait ContextInterface { } // TEMPORARY: This trait is to promote sharing of the current public context -// and the upcoming AVMContext. This will be removed once the AVMContext is the default. -// If you are adding something here, then you should also implement it in the AVMContext. +// and the upcoming AvmContext. This will be removed once the AvmContext is the default. +// If you are adding something here, then you should also implement it in the AvmContext. trait PublicContextInterface { fn block_number(self) -> Field; fn timestamp(self) -> u64; 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 f6f628e848e..eff362ec1d6 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 @@ -25,7 +25,7 @@ contract AvmTest { // Libs use dep::aztec::prelude::Map; - use dep::aztec::state_vars::PublicMutable; + use dep::aztec::state_vars::{PublicImmutable, 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::traits::ToField; @@ -39,11 +39,23 @@ contract AvmTest { single: PublicMutable, list: PublicMutable, map: Map>, + immutable: PublicImmutable, } /************************************************************************ * Storage ************************************************************************/ + // FIX: calls unsupported getNullifierMembershipWitness. + // #[aztec(public-vm)] + // #[aztec(initializer)] + // fn constructor() { + // storage.immutable.initialize(42); + // } + + unconstrained fn view_storage_immutable() -> pub Field { + storage.immutable.read() + } + unconstrained fn view_storage_single() -> pub Field { storage.single.read() } diff --git a/noir/noir-repo/aztec_macros/src/lib.rs b/noir/noir-repo/aztec_macros/src/lib.rs index 0fe450e6cb1..3ee6f9c21b9 100644 --- a/noir/noir-repo/aztec_macros/src/lib.rs +++ b/noir/noir-repo/aztec_macros/src/lib.rs @@ -4,7 +4,7 @@ mod utils; use transforms::{ compute_note_hash_and_nullifier::inject_compute_note_hash_and_nullifier, events::{generate_selector_impl, transform_events}, - functions::{transform_function, transform_unconstrained, transform_vm_function}, + functions::{transform_function, transform_unconstrained}, note_interface::generate_note_interface_impl, storage::{ assign_storage_slots, check_for_storage_definition, check_for_storage_implementation, @@ -138,9 +138,15 @@ fn transform_module(module: &mut SortedModule) -> Result } // Apply transformations to the function based on collected attributes - if is_private || is_public { + if is_private || is_public || is_public_vm { transform_function( - if is_private { "Private" } else { "Public" }, + if is_private { + "Private" + } else if is_public_vm { + "Avm" + } else { + "Public" + }, func, storage_defined, is_initializer, @@ -148,9 +154,6 @@ fn transform_module(module: &mut SortedModule) -> Result is_internal, )?; has_transformed_module = true; - } else if is_public_vm { - transform_vm_function(func, storage_defined)?; - has_transformed_module = true; } else if storage_defined && func.def.is_unconstrained { transform_unconstrained(func); has_transformed_module = true; diff --git a/noir/noir-repo/aztec_macros/src/transforms/functions.rs b/noir/noir-repo/aztec_macros/src/transforms/functions.rs index 2d84c2c8c44..9844abc30fe 100644 --- a/noir/noir-repo/aztec_macros/src/transforms/functions.rs +++ b/noir/noir-repo/aztec_macros/src/transforms/functions.rs @@ -35,6 +35,7 @@ pub fn transform_function( let context_name = format!("{}Context", ty); let inputs_name = format!("{}ContextInputs", ty); let return_type_name = format!("{}CircuitPublicInputs", ty); + let is_avm = ty == "Avm"; // Add check that msg sender equals this address and flag function as internal if is_internal { @@ -60,7 +61,11 @@ pub fn transform_function( } // Insert the context creation as the first action - let create_context = create_context(&context_name, &func.def.parameters)?; + let create_context = if !is_avm { + create_context(&context_name, &func.def.parameters)? + } else { + create_context_avm()? + }; func.def.body.statements.splice(0..0, (create_context).iter().cloned()); // Add the inputs to the params @@ -68,12 +73,14 @@ pub fn transform_function( func.def.parameters.insert(0, input); // Abstract return types such that they get added to the kernel's return_values - if let Some(return_values) = abstract_return_values(func) { - // In case we are pushing return values to the context, we remove the statement that originated it - // This avoids running duplicate code, since blocks like if/else can be value returning statements - func.def.body.statements.pop(); - // Add the new return statement - func.def.body.statements.push(return_values); + if !is_avm { + if let Some(return_values) = abstract_return_values(func) { + // In case we are pushing return values to the context, we remove the statement that originated it + // This avoids running duplicate code, since blocks like if/else can be value returning statements + func.def.body.statements.pop(); + // Add the new return statement + func.def.body.statements.push(return_values); + } } // Before returning mark the contract as initialized @@ -83,49 +90,29 @@ pub fn transform_function( } // Push the finish method call to the end of the function - let finish_def = create_context_finish(); - func.def.body.statements.push(finish_def); + if !is_avm { + let finish_def = create_context_finish(); + func.def.body.statements.push(finish_def); + } - let return_type = create_return_type(&return_type_name); - func.def.return_type = return_type; - func.def.return_visibility = Visibility::Public; + // The AVM doesn't need a return type yet. + if !is_avm { + let return_type = create_return_type(&return_type_name); + func.def.return_type = return_type; + func.def.return_visibility = Visibility::Public; + } // Distinct return types are only required for private functions // Public functions should have unconstrained auto-inferred match ty { "Private" => func.def.return_distinctness = Distinctness::Distinct, - "Public" => func.def.is_unconstrained = true, + "Public" | "Avm" => func.def.is_unconstrained = true, _ => (), } Ok(()) } -/// Transform a function to work with AVM bytecode -pub fn transform_vm_function( - func: &mut NoirFunction, - storage_defined: bool, -) -> Result<(), AztecMacroError> { - // Create access to storage - if storage_defined { - let storage = abstract_storage("public_vm", true); - func.def.body.statements.insert(0, storage); - } - - // Push Avm context creation to the beginning of the function - let create_context = create_avm_context()?; - func.def.body.statements.insert(0, create_context); - - // Add the inputs to the params (first!) - let input = create_inputs("AvmContextInputs"); - func.def.parameters.insert(0, input); - - // We want the function to be seen as a public function - func.def.is_unconstrained = true; - - Ok(()) -} - /// Transform Unconstrained /// /// Inserts the following code at the beginning of an unconstrained function @@ -339,36 +326,36 @@ fn create_context(ty: &str, params: &[Param]) -> Result, AztecMac Ok(injected_expressions) } -/// Creates an mutable avm context +/// Creates the private context object to be accessed within the function, the parameters need to be extracted to be +/// appended into the args hash object. /// +/// The replaced code: /// ```noir -/// /// Before /// #[aztec(public-vm)] -/// fn foo() -> Field { -/// let mut context = aztec::context::AVMContext::new(); -/// let timestamp = context.timestamp(); -/// // ... -/// } -/// -/// /// After -/// #[aztec(private)] -/// fn foo() -> Field { -/// let mut timestamp = context.timestamp(); -/// // ... +/// fn foo(inputs: AvmContextInputs, ...) -> Field { +/// let mut context = AvmContext::new(inputs); /// } -fn create_avm_context() -> Result { +/// ``` +fn create_context_avm() -> Result, AztecMacroError> { + let mut injected_expressions: Vec = vec![]; + // Create the inputs to the context + let ty = "AvmContext"; let inputs_expression = variable("inputs"); + let path_snippet = ty.to_case(Case::Snake); // e.g. private_context + // let mut context = {ty}::new(inputs, hash); let let_context = mutable_assignment( "context", // Assigned to call( - variable_path(chained_dep!("aztec", "context", "AVMContext", "new")), // Path - vec![inputs_expression], // args + variable_path(chained_dep!("aztec", "context", &path_snippet, ty, "new")), // Path + vec![inputs_expression], // args ), ); + injected_expressions.push(let_context); - Ok(let_context) + // Return all expressions that will be injected by the hasher + Ok(injected_expressions) } /// Abstract Return Type diff --git a/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts b/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts index 3dce96dd2a3..99032a48228 100644 --- a/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts +++ b/yarn-project/end-to-end/src/e2e_avm_simulator.test.ts @@ -25,6 +25,11 @@ describe('e2e_avm_simulator', () => { }, 50_000); describe('Storage', () => { + // FIX: Enable once the contract function works. + // it('Read immutable (initialized) storage (Field)', async () => { + // expect(await avmContact.methods.view_storage_immutable().view()).toEqual(42n); + // }); + it('Modifies storage (Field)', async () => { await avmContact.methods.set_storage_single(20n).send().wait(); expect(await avmContact.methods.view_storage_single().view()).toEqual(20n); diff --git a/yarn-project/simulator/src/avm/avm_execution_environment.ts b/yarn-project/simulator/src/avm/avm_execution_environment.ts index 77233c5b56b..2381cf2f16b 100644 --- a/yarn-project/simulator/src/avm/avm_execution_environment.ts +++ b/yarn-project/simulator/src/avm/avm_execution_environment.ts @@ -1,6 +1,6 @@ import { FunctionSelector, GlobalVariables } from '@aztec/circuits.js'; +import { computeVarArgsHash } from '@aztec/circuits.js/hash'; import { AztecAddress } from '@aztec/foundation/aztec-address'; -import { pedersenHash } from '@aztec/foundation/crypto'; import { EthAddress } from '@aztec/foundation/eth-address'; import { Fr } from '@aztec/foundation/fields'; @@ -54,7 +54,7 @@ export class AvmExecutionEnvironment { ) { // We encode some extra inputs (AvmContextInputs) in calldata. // This will have to go once we move away from one proof per call. - const inputs = new AvmContextInputs(temporaryFunctionSelector.toField(), pedersenHash(calldata)); + const inputs = new AvmContextInputs(temporaryFunctionSelector.toField(), computeVarArgsHash(calldata)); this.calldata = [...inputs.toFields(), ...calldata]; } diff --git a/yarn-project/simulator/src/avm/avm_simulator.test.ts b/yarn-project/simulator/src/avm/avm_simulator.test.ts index f1932f03be9..287ad4c8acd 100644 --- a/yarn-project/simulator/src/avm/avm_simulator.test.ts +++ b/yarn-project/simulator/src/avm/avm_simulator.test.ts @@ -1,4 +1,5 @@ import { UnencryptedL2Log } from '@aztec/circuit-types'; +import { computeVarArgsHash } from '@aztec/circuits.js/hash'; import { EventSelector, FunctionSelector } from '@aztec/foundation/abi'; import { AztecAddress } from '@aztec/foundation/aztec-address'; import { keccak, pedersenHash, poseidonHash, sha256 } from '@aztec/foundation/crypto'; @@ -20,22 +21,20 @@ import { initL1ToL2MessageOracleInput, initMachineState, } from './fixtures/index.js'; -import { Add, CalldataCopy, Instruction, Return } from './opcodes/index.js'; +import { Add, CalldataCopy, Return } from './opcodes/index.js'; import { encodeToBytecode } from './serialization/bytecode_serialization.js'; describe('AVM simulator: injected bytecode', () => { let calldata: Fr[]; - let ops: Instruction[]; let bytecode: Buffer; beforeAll(() => { calldata = [new Fr(1), new Fr(2)]; - ops = [ + bytecode = encodeToBytecode([ new CalldataCopy(/*indirect=*/ 0, /*cdOffset=*/ adjustCalldataIndex(0), /*copySize=*/ 2, /*dstOffset=*/ 0), new Add(/*indirect=*/ 0, TypeTag.FIELD, /*aOffset=*/ 0, /*bOffset=*/ 1, /*dstOffset=*/ 2), new Return(/*indirect=*/ 0, /*returnOffset=*/ 2, /*copySize=*/ 1), - ]; - bytecode = encodeToBytecode(ops); + ]); }); it('Should execute bytecode that performs basic addition', async () => { @@ -247,7 +246,7 @@ describe('AVM simulator: transpiled Noir contracts', () => { const results = await new AvmSimulator(context).executeBytecode(bytecode); expect(results.reverted).toBe(false); - expect(results.output).toEqual([pedersenHash(calldata)]); + expect(results.output).toEqual([computeVarArgsHash(calldata)]); }); }); @@ -462,6 +461,7 @@ describe('AVM simulator: transpiled Noir contracts', () => { const results = await new AvmSimulator(context).executeBytecode(callBytecode); + expect(results.revertReason).toBeUndefined(); expect(results.reverted).toBe(false); expect(results.output).toEqual([new Fr(3)]); });