feat(aztec-nr/public): dispatch function (#8821)

Towards [public dispatch functions](https://docs.google.com/document/d/14l0cgpAe8iMc075XkaIkh5qUhv-VR069upC0T3zriaw/). This PR adds a `public_dispatch` function to every contract with public functions, but nothing uses it yet.

Most of the work was done by @asterite.

The generated `public_dispatch` function for the Token contract is 50k, while the sum of every public function is 42k.

---

```
AuthWitTest::public_dispatch: bytecode is 1900 bytes
AvmInitializerTest::public_dispatch: bytecode is 3852 bytes
AuthRegistry::public_dispatch: bytecode is 9825 bytes
Auth::public_dispatch: bytecode is 16544 bytes
EasyPrivateVoting::public_dispatch: bytecode is 7569 bytes
Claim::public_dispatch: bytecode is 6163 bytes
Delegator::public_dispatch: bytecode is 2149 bytes
ImportTest::public_dispatch: bytecode is 1452 bytes
PriceFeed::public_dispatch: bytecode is 2758 bytes
FeeJuice::public_dispatch: bytecode is 5912 bytes
Router::public_dispatch: bytecode is 4332 bytes
FPC::public_dispatch: bytecode is 12260 bytes
AvmTest::public_dispatch: bytecode is 82982 bytes
Child::public_dispatch: bytecode is 5936 bytes
DelegatedOn::public_dispatch: bytecode is 1346 bytes
PrivateFPC::public_dispatch: bytecode is 5634 bytes
Parent::public_dispatch: bytecode is 11734 bytes
Crowdfunding::public_dispatch: bytecode is 8942 bytes
StaticParent::public_dispatch: bytecode is 8243 bytes
Lending::public_dispatch: bytecode is 59801 bytes
Benchmarking::public_dispatch: bytecode is 3607 bytes
AppSubscription::public_dispatch: bytecode is 8740 bytes
Uniswap::public_dispatch: bytecode is 36542 bytes
NFT::public_dispatch: bytecode is 26572 bytes
Spam::public_dispatch: bytecode is 2329 bytes
TokenBridge::public_dispatch: bytecode is 36421 bytes
StaticChild::public_dispatch: bytecode is 4107 bytes
InclusionProofs::public_dispatch: bytecode is 6132 bytes
TestLog::public_dispatch: bytecode is 3875 bytes
DocsExample::public_dispatch: bytecode is 6486 bytes
Test::public_dispatch: bytecode is 35885 bytes
StatefulTest::public_dispatch: bytecode is 9648 bytes
CardGame::public_dispatch: bytecode is 22558 bytes
TokenBlacklist::public_dispatch: bytecode is 93626 bytes
Token::public_dispatch: bytecode is 50054 bytes
```

avm-transpiler/src/transpile_contract.rs

@@ -91,7 +91,6 @@ impl From<CompiledAcirContractArtifact> for TranspiledContractArtifact {
         let mut functions: Vec<AvmOrAcirContractFunctionArtifact> = Vec::new();
 
         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()) {
                 info!("Transpiling AVM function {} on contract {}", function.name, contract.name);
                 // Extract Brillig Opcodes from acir

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

@@ -0,0 +1,154 @@
+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 mut parameters_size = 0;
+        for param in parameters {
+            parameters_size += size_in_fields(param.1);
+        }
+
+        let initial_read = if parameters.len() == 0 {
+            quote {}
+        } else {
+            // The initial calldata_copy offset is 1 to skip the Field selector
+            // The expected calldata is the serialization of
+            // - FunctionSelector: the selector of the function intended to dispatch
+            // - Parameters: the parameters of the function intended to dispatch
+            // That is, exactly what is expected for a call to the target function,
+            // but with a selector added at the beginning.
+            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 reads = parameters.map(|param: (Quoted, Type)| {
+            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(dep::aztec::protocol_types::traits::Deserialize::deserialize);
+            };
+            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;
+                // Force early return.
+                dep::aztec::context::public_context::avm_return([]);
+            }
+        } 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 {  });
+
+        let body = quote {
+            // We mark this as public because our whole system depends on public
+            // functions having this attribute. However, the public MACRO will
+            // handle the public_dispatch function specially and do nothing.
+            #[public]
+            unconstrained pub fn public_dispatch(selector: Field) {
+                $dispatch
+            }
+        };
+
+        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

@@ -225,6 +225,15 @@ pub comptime fn private(f: FunctionDefinition) -> Quoted {
 
 /// Public functions are executed sequencer-side and do not preserve privacy, similar to the EVM.
 pub comptime fn public(f: FunctionDefinition) -> Quoted {
+    // We don't want to transform the public_dispatch function.
+    if f.name() == quote { public_dispatch } {
+        quote {}
+    } else {
+        transform_public(f)
+    }
+}
+
+comptime fn transform_public(f: FunctionDefinition) -> Quoted {
     let fn_abi = create_fn_abi_export(f);
     let fn_stub = stub_fn(f);
     register_stub(f.module(), fn_stub);

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-protocol-circuits/crates/types/src/traits.nr

@@ -155,12 +155,6 @@ pub trait Serialize<let N: u32> {
 }
 // docs:end:serialize
 
-impl<let N: u32> Serialize<N> for [Field; N] {
-    fn serialize(self) -> [Field; N] {
-        self
-    }
-}
-
 impl<let N: u32> Serialize<N> for str<N> {
     fn serialize(self) -> [Field; N] {
         let bytes = self.as_bytes();
@@ -179,8 +173,8 @@ pub trait Deserialize<let N: u32> {
 }
 // docs:end:deserialize
 
-impl<let N: u32> Deserialize<N> for [Field; N] {
+impl <let N: u32> Deserialize<N> for str<N> {
     fn deserialize(fields: [Field; N]) -> Self {
-        fields
+        str<N>::from(fields.map(|value| value as u8))
     }
 }

noir-projects/noir-protocol-circuits/crates/types/src/type_serialization.nr

@@ -20,7 +20,7 @@ impl Deserialize<BOOL_SERIALIZED_LEN> for bool {
 }
 
 impl Serialize<U8_SERIALIZED_LEN> for u8 {
-    fn serialize(self) -> [Field; U32_SERIALIZED_LEN] {
+    fn serialize(self) -> [Field; U8_SERIALIZED_LEN] {
         [self as Field]
     }
 }
@@ -56,7 +56,7 @@ impl Deserialize<U64_SERIALIZED_LEN> for u64 {
 }
 
 impl Serialize<U128_SERIALIZED_LEN> for U128 {
-    fn serialize(self) -> [Field; 1] {
+    fn serialize(self) -> [Field; U128_SERIALIZED_LEN] {
         [self.to_integer()]
     }
 }
@@ -68,7 +68,7 @@ impl Deserialize<U128_SERIALIZED_LEN> for U128 {
 }
 
 impl Serialize<FIELD_SERIALIZED_LEN> for Field {
-    fn serialize(self) -> [Field; U32_SERIALIZED_LEN] {
+    fn serialize(self) -> [Field; FIELD_SERIALIZED_LEN] {
         [self]
     }
 }
@@ -78,3 +78,25 @@ impl Deserialize<FIELD_SERIALIZED_LEN> for Field {
         fields[0]
     }
 }
+
+impl <T, let N: u32, let M: u32> Serialize<N * M> for [T; N] where T: Serialize<M> {
+    fn serialize(self) -> [Field; N * M] {
+        let mut result: [Field; N * M] = std::mem::zeroed();
+        let mut serialized: [Field; M] = std::mem::zeroed();
+        for i in 0..N {
+            serialized = self[i].serialize();
+            for j in 0..M {
+                result[i * M + j] = serialized[j];
+            }
+        }
+        result
+    }
+}
+
+impl <T, let N: u32, let M: u32> Deserialize<N * M> for [T; N] where T: Deserialize<M> {
+    fn deserialize(fields: [Field; N * M]) -> Self {
+        let mut reader = crate::utils::reader::Reader::new(fields);
+        let mut result: [T; N] = std::mem::zeroed();
+        reader.read_struct_array::<T, M, N>(Deserialize::deserialize, result)
+    }
+}