feat(avm): add AvmContextInputs (#5396)

NOTE: I don't know why this triggered a change in the snapshot.

---

This structure lets us easily pass things from the (TS) context to the constructor of the `AvmContext` in Noir, using `calldata` as a vehicle (effectively adding them as public inputs).

The choice to add the structure to the function arguments as LAST and not first is because adding them first would break any non-noir-generated bytecode (since they would have to shift their expected calldata by a magic number `N = sizeof(AvmContextInputs)`). Putting it last, makes it transparent for them. A calldatacopy would still work.

However, having this makes any external call always have `AvmContextInputs` in the calldata, bloating it (on chain) for non-noir-generated bytecode. This is not an issue now.

For the moment, this is temporary, but might be useful long term. (Sean mentioned passing maybe env getters like this).

---

Implemented `AvmContext.selector()` and `AvmContext.get_args_hash()` using the above.

noir-projects/aztec-nr/aztec/src/context/avm_context.nr

@@ -3,15 +3,17 @@ use dep::protocol_types::traits::Serialize;
 use dep::protocol_types::abis::function_selector::FunctionSelector;
 use dep::protocol_types::abis::public_circuit_public_inputs::PublicCircuitPublicInputs;
 use dep::protocol_types::constants::RETURN_VALUES_LENGTH;
-use crate::context::inputs::PublicContextInputs;
+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 {
-    pub fn new() -> Self {
-        AVMContext {}
+    pub fn new(inputs: AvmContextInputs) -> Self {
+        AVMContext { inputs }
     }
 
     pub fn origin(self) -> AztecAddress {
@@ -190,16 +192,14 @@ impl ContextInterface for AVMContext {
         version()
     }
     fn selector(self) -> FunctionSelector {
-        assert(false, "'selector' not implemented!");
-        FunctionSelector::zero()
+        FunctionSelector::from_field(self.inputs.selector)
     }
     fn get_header(self) -> Header {
         assert(false, "'get_header' not implemented!");
         Header::empty()
     }
     fn get_args_hash(self) -> Field {
-        assert(false, "'get_args_hash' not implemented!");
-        0
+        self.inputs.args_hash
     }
 }
 

noir-projects/aztec-nr/aztec/src/context/inputs.nr

@@ -1,5 +1,7 @@
 mod private_context_inputs;
 mod public_context_inputs;
+mod avm_context_inputs;
 
 use crate::context::inputs::private_context_inputs::PrivateContextInputs;
 use crate::context::inputs::public_context_inputs::PublicContextInputs;
+use crate::context::inputs::avm_context_inputs::AvmContextInputs;

noir-projects/aztec-nr/aztec/src/context/inputs/avm_context_inputs.nr

@@ -0,0 +1,4 @@
+struct AvmContextInputs {
+    selector: Field,
+    args_hash: Field,
+}

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

@@ -232,6 +232,16 @@ contract AvmTest {
     //     context.contract_call_depth()
     // }
 
+    #[aztec(public-vm)]
+    fn check_selector() {
+        assert(context.selector() == FunctionSelector::from_signature("check_selector()"));
+    }
+
+    #[aztec(public-vm)]
+    fn get_args_hash(_a: u8, _fields: [Field; 3]) -> pub Field {
+        context.get_args_hash()
+    }
+
     #[aztec(public-vm)]
     fn emit_unencrypted_log() {
         context.accumulate_unencrypted_logs(/*event_selector=*/ 5, /*message=*/ [10, 20, 30]);

noir/noir-repo/aztec_macros/src/transforms/functions.rs

@@ -116,6 +116,10 @@ pub fn transform_vm_function(
     let create_context = create_avm_context()?;
     func.def.body.0.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;
 
@@ -353,11 +357,14 @@ fn create_context(ty: &str, params: &[Param]) -> Result<Vec<Statement>, AztecMac
 ///     // ...
 /// }
 fn create_avm_context() -> Result<Statement, AztecMacroError> {
+    // Create the inputs to the context
+    let inputs_expression = variable("inputs");
+
     let let_context = mutable_assignment(
         "context", // Assigned to
         call(
             variable_path(chained_dep!("aztec", "context", "AVMContext", "new")), // Path
-            vec![],                                                               // args
+            vec![inputs_expression],                                              // args
         ),
     );
 

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

@@ -1,6 +1,6 @@
 import { AztecAddress, Fr } from '@aztec/circuits.js';
 
-import { allSameExcept, initContext } from './fixtures/index.js';
+import { allSameExcept, anyAvmContextInputs, initContext } from './fixtures/index.js';
 
 describe('Avm Context', () => {
   it('New call should fork context correctly', () => {
@@ -15,7 +15,8 @@ describe('Avm Context', () => {
       allSameExcept(context.environment, {
         address: newAddress,
         storageAddress: newAddress,
-        calldata: newCalldata,
+        // Calldata also includes AvmContextInputs
+        calldata: anyAvmContextInputs().concat(newCalldata),
         isStaticCall: false,
       }),
     );
@@ -41,7 +42,8 @@ describe('Avm Context', () => {
       allSameExcept(context.environment, {
         address: newAddress,
         storageAddress: newAddress,
-        calldata: newCalldata,
+        // Calldata also includes AvmContextInputs
+        calldata: anyAvmContextInputs().concat(newCalldata),
         isStaticCall: true,
       }),
     );

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

@@ -1,6 +1,6 @@
 import { Fr } from '@aztec/foundation/fields';
 
-import { allSameExcept, initExecutionEnvironment } from './fixtures/index.js';
+import { allSameExcept, anyAvmContextInputs, initExecutionEnvironment } from './fixtures/index.js';
 
 describe('Execution Environment', () => {
   const newAddress = new Fr(123456n);
@@ -14,7 +14,8 @@ describe('Execution Environment', () => {
       allSameExcept(executionEnvironment, {
         address: newAddress,
         storageAddress: newAddress,
-        calldata,
+        // Calldata also includes AvmContextInputs
+        calldata: anyAvmContextInputs().concat(calldata),
       }),
     );
   });
@@ -27,7 +28,8 @@ describe('Execution Environment', () => {
       allSameExcept(executionEnvironment, {
         address: newAddress,
         isDelegateCall: true,
-        calldata,
+        // Calldata also includes AvmContextInputs
+        calldata: anyAvmContextInputs().concat(calldata),
       }),
     );
   });
@@ -41,7 +43,8 @@ describe('Execution Environment', () => {
         address: newAddress,
         storageAddress: newAddress,
         isStaticCall: true,
-        calldata,
+        // Calldata also includes AvmContextInputs
+        calldata: anyAvmContextInputs().concat(calldata),
       }),
     );
   });

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

@@ -1,8 +1,19 @@
 import { FunctionSelector, GlobalVariables } from '@aztec/circuits.js';
 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';
 
+export class AvmContextInputs {
+  static readonly SIZE = 2;
+
+  constructor(private selector: Fr, private argsHash: Fr) {}
+
+  public toFields(): Fr[] {
+    return [this.selector, this.argsHash];
+  }
+}
+
 /**
  * Contains variables that remain constant during AVM execution
  * These variables are provided by the public kernel circuit
@@ -40,7 +51,15 @@ export class AvmExecutionEnvironment {
     // containing all functions, and function selector will become an application-level mechanism
     // (e.g. first few bytes of calldata + compiler-generated jump table)
     public readonly temporaryFunctionSelector: FunctionSelector,
-  ) {}
+  ) {
+    // 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.map(word => word.toBuffer())),
+    );
+    this.calldata = [...inputs.toFields(), ...calldata];
+  }
 
   public deriveEnvironmentForNestedCall(
     address: AztecAddress,