Merge remote-tracking branch 'origin/ab/contract-entrypoint-dispatch'…

… into fc/avm-dispatch-changes

avm-transpiler/src/transpile_contract.rs

@@ -92,7 +92,9 @@ impl From<CompiledAcirContractArtifact> for TranspiledContractArtifact {
 
         for function in contract.functions {
             // TODO(4269): once functions are tagged for transpilation to AVM, check tag
-            if function.custom_attributes.contains(&"public".to_string()) {
+            if function.custom_attributes.contains(&"public".to_string())
+                || function.name == "public_dispatch"
+            {
                 info!("Transpiling AVM function {} on contract {}", function.name, contract.name);
                 // Extract Brillig Opcodes from acir
                 let acir_program = function.bytecode;

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

@@ -201,6 +201,7 @@ impl<let N: u32, T, Env> PublicCallInterface<N, T, Env> {
         returns.deserialize_into()
     }
 
+    // TODO(facundo): Function selector should be part of the args.
     pub fn enqueue(self, context: &mut PrivateContext) {
         let args_hash = hash_args(self.args);
         pack_arguments(self.args);

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

@@ -70,11 +70,15 @@ impl PublicContext {
         args: [Field],
         gas_opts: GasOpts
     ) -> FunctionReturns<RETURNS_COUNT> {
+        let args_with_selector = &[function_selector.to_field()].append(args);
+        let public_dispatch_selector = comptime {
+            FunctionSelector::from_signature("public_dispatch(Field)").to_field()
+        };
         let results = call(
             gas_for_call(gas_opts),
             contract_address,
-            args,
-            function_selector.to_field()
+            args_with_selector,
+            public_dispatch_selector
         );
         let data_to_return: [Field; RETURNS_COUNT] = results.0;
         let success: u8 = results.1;
@@ -90,11 +94,15 @@ impl PublicContext {
         args: [Field],
         gas_opts: GasOpts
     ) -> FunctionReturns<RETURNS_COUNT> {
+        let args_with_selector = &[function_selector.to_field()].append(args);
+        let public_dispatch_selector = comptime {
+            FunctionSelector::from_signature("public_dispatch(Field)").to_field()
+        };
         let (data_to_return, success): ([Field; RETURNS_COUNT], u8) = call_static(
             gas_for_call(gas_opts),
             contract_address,
-            args,
-            function_selector.to_field()
+            args_with_selector,
+            public_dispatch_selector
         );
 
         assert(success == 1, "Nested static call failed!");

noir-projects/aztec-nr/aztec/src/macros/dispatch/mod.nr

@@ -0,0 +1,152 @@
+use super::utils::compute_fn_selector;
+
+/// Returns an `fn public_dispatch(...)` function for the given module that's assumed to be an Aztec contract.
+pub comptime fn generate_public_dispatch(m: Module) -> Quoted {
+    let functions = m.functions();
+    let functions = functions.filter(|function: FunctionDefinition| function.has_named_attribute("public"));
+
+    let unit = get_type::<()>();
+
+    let ifs = functions.map(
+        |function: FunctionDefinition| {
+        let name = function.name();
+        let parameters = function.parameters();
+        let return_type = function.return_type();
+
+        let selector: Field = compute_fn_selector(function);
+
+        let param_sizes = parameters.map(|param: (Quoted, Type)| {
+            size_in_fields(param.1)
+        });
+
+        let mut parameters_size = 0;
+        for param_size in param_sizes {
+            parameters_size += param_size;
+        }
+
+        let initial_read = if parameters.len() == 0 {
+            quote {}
+        } else {
+            // The initial calldata_copy offset is 1 to skip the Field selector
+            quote {
+                let input_calldata: [Field; $parameters_size] = dep::aztec::context::public_context::calldata_copy(1, $parameters_size);
+                let mut reader = dep::aztec::protocol_types::utils::reader::Reader::new(input_calldata);
+            }
+        };
+
+        let mut parameter_index = 0;
+        let mut offset = 0;
+        let reads = parameters.map(|param: (Quoted, Type)| {
+            // let param_size = param_sizes[parameter_index];
+            let param_name = f"arg{parameter_index}".quoted_contents();
+            let param_type = param.1;
+            let read = quote {
+                // let $param_name: $param_type = reader.read_struct((($param_type)::deserialize));
+                let $param_name: $param_type = reader.read_struct(dep::aztec::protocol_types::traits::Deserialize::deserialize);
+                // let ($param_name, _) = dep::aztec::protocol_types::traits::FromCalldata::from_calldata(input_calldata, $offset);
+            };
+            // offset += param_size;
+            parameter_index += 1;
+            quote { $read }
+        });
+        let read = reads.join(quote { });
+
+        let mut args = &[];
+        for parameter_index in 0..parameters.len() {
+            let param_name = f"arg{parameter_index}".quoted_contents();
+            args = args.push_back(quote { $param_name });
+        }
+
+        let args = args.join(quote { , });
+        let call = quote { $name($args) };
+
+        let return_code = if return_type == unit {
+            quote { $call }
+        } else {
+            quote {
+                let return_value = dep::aztec::protocol_types::traits::Serialize::serialize($call);
+                dep::aztec::context::public_context::avm_return(return_value);
+            }
+        };
+
+        let if_ = quote {
+            if selector == $selector {
+                $initial_read
+                $read
+                $return_code
+            }
+        };
+        if_
+    }
+    );
+
+    if ifs.len() == 0 {
+        // No dispatch function if there are no public functions
+        quote {}
+    } else {
+        let ifs = ifs.push_back(quote { { panic(f"Unknown selector") } });
+        let dispatch = ifs.join(quote { else });
+
+        let body = quote {
+            unconstrained pub fn public_dispatch(selector: Field) {
+                $dispatch
+            }
+        };
+
+        // println(body);
+
+        body
+    }
+}
+
+comptime fn size_in_fields(typ: Type) -> u32 {
+    if typ.as_slice().is_some() {
+        panic(f"Can't determine size in fields of Slice type")
+    } else {
+        let size = array_size_in_fields(typ);
+        let size = size.or_else(|| struct_size_in_fields(typ));
+        let size = size.or_else(|| tuple_size_in_fields(typ));
+        size.unwrap_or(1)
+    }
+}
+
+comptime fn array_size_in_fields(typ: Type) -> Option<u32> {
+    typ.as_array().and_then(
+        |typ: (Type, Type)| {
+        let (typ, element_size) = typ;
+        element_size.as_constant().map(|x: u32| {
+            x * size_in_fields(typ)
+        })
+    }
+    )
+}
+
+comptime fn struct_size_in_fields(typ: Type) -> Option<u32> {
+    typ.as_struct().map(
+        |typ: (StructDefinition, [Type])| {
+        let struct_type = typ.0;
+        let mut size = 0;
+        for field in struct_type.fields() {
+            size += size_in_fields(field.1);
+        }
+        size
+    }
+    )
+}
+
+comptime fn tuple_size_in_fields(typ: Type) -> Option<u32> {
+    typ.as_tuple().map(
+        |types: [Type]| {
+        let mut size = 0;
+        for typ in types {
+            size += size_in_fields(typ);
+        }
+        size
+    }
+    )
+}
+
+comptime fn get_type<T>() -> Type {
+    let t: T = std::mem::zeroed();
+    std::meta::type_of(t)
+}

noir-projects/aztec-nr/aztec/src/macros/functions/mod.nr

@@ -237,10 +237,14 @@ pub comptime fn public(f: FunctionDefinition) -> Quoted {
     let args_len = original_params.map(|(name, typ): (Quoted, Type)| flatten_to_fields(name, typ, &[]).0.len()).fold(0, |acc: u32, val: u32| acc + val);
 
     // Unlike in the private case, in public the `context` does not need to receive the hash of the original params.
-    let context_creation = quote { let mut context = dep::aztec::context::public_context::PublicContext::new(|| {
-        let serialized_args : [Field; $args_len] = dep::aztec::context::public_context::calldata_copy(0, $args_len);
-        dep::aztec::hash::hash_args_array(serialized_args)
-     }); };
+    let context_creation = quote {
+        let mut context = dep::aztec::context::public_context::PublicContext::new(|| {
+            // The +1 is due to the function selector in the dispatch function.
+            // This function, as all public functions, is only expected to be called through "public_dispatch".
+            let serialized_args : [Field; $args_len] = dep::aztec::context::public_context::calldata_copy(0, $args_len + 1);
+            dep::aztec::hash::hash_args_array(serialized_args)
+        });
+    };
 
     // Modifications introduced by the different marker attributes.
     let internal_check = if is_fn_internal(f) {

noir-projects/aztec-nr/aztec/src/macros/mod.nr

@@ -1,3 +1,4 @@
+mod dispatch;
 mod functions;
 mod utils;
 mod notes;
@@ -10,26 +11,26 @@ use notes::{NOTES, generate_note_export};
 
 use functions::transform_unconstrained;
 use utils::module_has_storage;
+use dispatch::generate_public_dispatch;
 
 /// Marks a contract as an Aztec contract, generating the interfaces for its functions and notes, as well as injecting
 /// the `compute_note_hash_and_optionally_a_nullifier` function PXE requires in order to validate notes.
 pub comptime fn aztec(m: Module) -> Quoted {
     let interface = generate_contract_interface(m);
-
     let unconstrained_functions = m.functions().filter(
         | f: FunctionDefinition | f.is_unconstrained() & !f.has_named_attribute("test") & !f.has_named_attribute("public")
     );
     for f in unconstrained_functions {
         transform_unconstrained(f);
     }
-
     let compute_note_hash_and_optionally_a_nullifier = generate_compute_note_hash_and_optionally_a_nullifier();
     let note_exports = generate_note_exports();
-
+    let public_dispatch = generate_public_dispatch(m);
     quote {
         $note_exports
         $interface
         $compute_note_hash_and_optionally_a_nullifier
+        $public_dispatch
     }
 }
 

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

@@ -564,4 +564,4 @@ contract AvmTest {
         dep::aztec::oracle::debug_log::debug_log("nested_static_call_to_add");
         let _ = nested_static_call_to_add(1, 2);
     }
-}
+}

noir-projects/noir-protocol-circuits/crates/types/src/address/aztec_address.nr

@@ -1,11 +1,12 @@
 use crate::{
     crate::address::{partial_address::PartialAddress, public_keys_hash::PublicKeysHash},
     constants::{AZTEC_ADDRESS_LENGTH, GENERATOR_INDEX__CONTRACT_ADDRESS_V1},
-    hash::poseidon2_hash_with_separator, traits::{Empty, FromField, ToField, Serialize, Deserialize},
-    utils
+    hash::poseidon2_hash_with_separator,
+    traits::{Empty, FromField, ToField, Serialize, Deserialize, FromCalldata}, utils
 };
 
 // Aztec address
+#[derive(FromCalldata)]
 pub struct AztecAddress {
     inner : Field
 }

noir-projects/noir-protocol-circuits/crates/types/src/address/eth_address.nr

@@ -1,5 +1,6 @@
-use crate::{constants::ETH_ADDRESS_LENGTH, traits::{Empty, ToField, Serialize, Deserialize}, utils};
+use crate::{constants::ETH_ADDRESS_LENGTH, traits::{Empty, ToField, Serialize, Deserialize, FromCalldata}, utils};
 
+#[derive(FromCalldata)]
 pub struct EthAddress{
     inner : Field
 }