chore: comptime optimizations

asterite · asterite · commit 4aa11bc978db · 2026-01-29T22:33:59.000Z
This PR includes some refactors to our comptime code so it's interpreted faster: 1. a `CHashMap` type is introduced. This is a dummy HashMap that does lookups by traversing an entire list. At first this might seem slower. However, `UHashMap` does a lot of things to get or insert a key (it hashes is, computing the hash using poseidon2, then does a linear probing, etc.), and most of that time is then spent in the interpreter and not actually in the logic. The linear logic ends up being much faster, even for big maps (and here the maps are small, but it still makes a difference). See also [this PR](noir-lang/noir#11384). 2. Use `CtString` more to build function signatures and so on. This is a bit more efficient than working with `Quoted` 3. Eventually I removed `AsStrQuote`. There are two places where we need the length of a `CtString` but then I inlined the code (it's a one-line thing). For `token_contract`, the time to re-check main in LSP goes from 270ms to 220ms. 50ms is not much, but from an IDE-perspective it is. Then `nargo check` for `token_contract` goes from 1.13 seconds to 1.07 seconds. Again, not a huge win, but it's noticeable in the IDE. (this should improve the LSP speed of all contracts, not just token_contract, but that's one of the biggest ones so I always use it to try things on)
diff --git a/noir-projects/aztec-nr/aztec/src/macros/authorization.nr b/noir-projects/aztec-nr/aztec/src/macros/authorization.nr
@@ -1,15 +1,16 @@
-use crate::macros::utils::{compute_struct_selector, derive_serialize_if_not_implemented, get_trait_impl_method};
-use poseidon::poseidon2::Poseidon2Hasher;
-use std::{collections::umap::UHashMap, hash::BuildHasherDefault};
+use crate::{
+    macros::utils::{compute_struct_selector, derive_serialize_if_not_implemented, get_trait_impl_method},
+    utils::cmap::CHashMap,
+};
 
 /// Hashmap that stores authorization structs indexed by their selectors, so no duplicates are inadvertedly created by
 /// developers
-comptime mut global AUTH_TYPES: UHashMap<Field, TypeDefinition, BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
+comptime mut global AUTH_TYPES: CHashMap<Field, TypeDefinition> = CHashMap::new();
 
 comptime fn register_authorization(authorization_selector: Field, authorization_struct: TypeDefinition) {
-    if AUTH_TYPES.contains_key(authorization_selector) {
-        let existing_authorization = AUTH_TYPES.get(authorization_selector).unwrap().name();
+    let existing_authorization = AUTH_TYPES.get(authorization_selector);
+    if existing_authorization.is_some() {
+        let existing_authorization = existing_authorization.unwrap().name();
         let authorization_name = authorization_struct.name();
         panic(
             f"Selector collision detected between authorizations '{authorization_name}' and '{existing_authorization}'",
diff --git a/noir-projects/aztec-nr/aztec/src/macros/calls_generation/external_functions_stubs.nr b/noir-projects/aztec-nr/aztec/src/macros/calls_generation/external_functions_stubs.nr
@@ -2,7 +2,7 @@
 // of manually serializing arguments and creating call interfaces, stubs allow natural syntax, e.g. for ! enqueuing
 // calls to public functions: ! !   ExternalContract.at(address).some_method(arg1, arg2).enqueue()
 
-use crate::macros::utils::{AsStrQuote, compute_fn_selector};
+use crate::macros::utils::compute_fn_selector;
 use crate::protocol::meta::utils::derive_serialization_quotes;
 use std::meta::unquote;
 
@@ -36,7 +36,7 @@ comptime fn create_stub_base(f: FunctionDefinition) -> (Quoted, Quoted, Quoted,
         derive_serialization_quotes(fn_parameters, false);
     let serialized_args_array_len: u32 = unquote!(quote { ($serialized_args_array_len_quote) as u32 });
 
-    let (fn_name_str, _) = fn_name.as_str_quote();
+    let fn_name_str = f"\"{fn_name}\"".quoted_contents();
     let fn_name_len: u32 = unquote!(quote { $fn_name_str.as_bytes().len()});
     let fn_selector: Field = compute_fn_selector(f);
 
diff --git a/noir-projects/aztec-nr/aztec/src/macros/dispatch.nr b/noir-projects/aztec-nr/aztec/src/macros/dispatch.nr
@@ -1,16 +1,16 @@
 use crate::macros::internals_functions_generation::external_functions_registry::get_public_functions;
 use crate::protocol::meta::utils::get_params_len_quote;
+use crate::utils::cmap::CHashMap;
 use super::utils::compute_fn_selector;
-use poseidon::poseidon2::Poseidon2Hasher;
-use std::{collections::umap::UHashMap, hash::BuildHasherDefault, panic};
+use std::panic;
 
 /// 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 = get_public_functions(m);
 
     let unit = get_type::<()>();
 
-    let seen_selectors = &mut UHashMap::<Field, Quoted, BuildHasherDefault<Poseidon2Hasher>>::default();
+    let seen_selectors = &mut CHashMap::<Field, Quoted>::new();
 
     let ifs = functions.map(|function: FunctionDefinition| {
         let parameters = function.parameters();
@@ -21,8 +21,9 @@ pub comptime fn generate_public_dispatch(m: Module) -> Quoted {
 
         // Since function selectors are computed as the first 4 bytes of the hash of the function signature, it's
         // possible to have collisions. With the following check, we ensure it doesn't happen within the same contract.
-        if seen_selectors.contains_key(selector) {
-            let existing_fn = seen_selectors.get(selector).unwrap();
+        let existing_fn = seen_selectors.get(selector);
+        if existing_fn.is_some() {
+            let existing_fn = existing_fn.unwrap();
             panic(
                 f"Public function selector collision detected between functions '{fn_name}' and '{existing_fn}'",
             );
diff --git a/noir-projects/aztec-nr/aztec/src/macros/events.nr b/noir-projects/aztec-nr/aztec/src/macros/events.nr
@@ -1,11 +1,12 @@
-use crate::macros::utils::{compute_struct_selector, derive_serialize_if_not_implemented, get_trait_impl_method};
-use poseidon::poseidon2::Poseidon2Hasher;
-use std::{collections::umap::UHashMap, hash::BuildHasherDefault, panic};
+use crate::{
+    macros::utils::{compute_struct_selector, derive_serialize_if_not_implemented, get_trait_impl_method},
+    utils::cmap::CHashMap,
+};
+use std::panic;
 
 /// A map from event selector to event name indicating whether the event selector has already been seen during the
 /// contract compilation - prevents event selector collisions.
-pub comptime mut global EVENT_SELECTORS: UHashMap<Field, Quoted, BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
+pub comptime mut global EVENT_SELECTORS: CHashMap<Field, Quoted> = CHashMap::new();
 
 comptime fn generate_event_interface_and_get_selector(s: TypeDefinition) -> (Quoted, Field) {
     let name = s.name();
@@ -31,8 +32,9 @@ comptime fn generate_event_interface_and_get_selector(s: TypeDefinition) -> (Quo
 }
 
 comptime fn register_event_selector(event_selector: Field, event_name: Quoted) {
-    if EVENT_SELECTORS.contains_key(event_selector) {
-        let existing_event = EVENT_SELECTORS.get(event_selector).unwrap();
+    let existing_event = EVENT_SELECTORS.get(event_selector);
+    if existing_event.is_some() {
+        let existing_event = existing_event.unwrap();
         panic(
             f"Event selector collision detected between events '{event_name}' and '{existing_event}'",
         );
diff --git a/noir-projects/aztec-nr/aztec/src/macros/functions/auth_registry.nr b/noir-projects/aztec-nr/aztec/src/macros/functions/auth_registry.nr
@@ -1,10 +1,9 @@
-use poseidon::poseidon2::Poseidon2Hasher;
-use std::{collections::umap::UHashMap, hash::BuildHasherDefault};
+use crate::utils::cmap::CHashMap;
 
 /// Registers the functions that have the `#[authorize_once(...)]` macro, and the parameters the macro was invoked
 /// with. This is used to later inject the authorization check (see ./utils.nr -> create_authorize_once_check) via the
 /// #[external("private")] or #[external("public")] macros. The `#[authorize_once(...)]` macro is not used directly to
 /// inject the check because it has to be placed after context instantiation but before the function body, which only
 /// the aforementioned macros can do.
-pub(crate) comptime mut global AUTHORIZE_ONCE_REGISTRY: UHashMap<FunctionDefinition, (CtString, CtString), BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
+pub(crate) comptime mut global AUTHORIZE_ONCE_REGISTRY: CHashMap<FunctionDefinition, (CtString, CtString)> =
+    CHashMap::new();
diff --git a/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/helpers.nr b/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/helpers.nr
@@ -2,7 +2,6 @@ use crate::macros::{
     functions::auth_registry::AUTHORIZE_ONCE_REGISTRY,
     utils::{is_fn_initializer, is_fn_only_self, is_fn_view},
 };
-use std::meta::ctstring::AsCtString;
 
 /// Gathers all attributes relevant to the function's ABI and returns a quote that can be applied to the newly
 /// generated function. We apply the abi marker attributes instead of the original ones (e.g. abi_view instead of view)
@@ -98,7 +97,6 @@ pub(crate) comptime fn create_authorize_once_check(f: FunctionDefinition, is_pri
         quote { aztec::authwit::auth::assert_current_call_valid_authwit_public }
     };
     let invalid_nonce_message = f"Invalid authwit nonce. When '{from_arg_name}' and 'msg_sender' are the same, '{nonce_arg_name}' must be zero"
-        .as_ctstring()
         .as_quoted_str();
     quote {         
         if (!$from_arg_name_quoted.eq(self.msg_sender())) {
diff --git a/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/private.nr b/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/private.nr
@@ -6,7 +6,7 @@ use crate::macros::{
     },
 };
 use crate::protocol::meta::utils::derive_serialization_quotes;
-use std::meta::{ctstring::AsCtString, type_of};
+use std::meta::type_of;
 
 pub(crate) comptime fn generate_private_external(f: FunctionDefinition) -> Quoted {
     let module_has_initializer = module_has_initializer(f.module());
@@ -70,8 +70,7 @@ pub(crate) comptime fn generate_private_external(f: FunctionDefinition) -> Quote
     };
 
     let view_check = if is_fn_view(f) {
-        let assertion_message =
-            f"Function {original_function_name} can only be called statically".as_ctstring().as_quoted_str();
+        let assertion_message = f"Function {original_function_name} can only be called statically".as_quoted_str();
         quote { assert(self.context.inputs.call_context.is_static_call, $assertion_message); }
     } else {
         quote {}
diff --git a/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/public.nr b/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/public.nr
@@ -5,7 +5,6 @@ use crate::macros::{
         module_has_initializer, module_has_storage,
     },
 };
-use std::meta::ctstring::AsCtString;
 
 pub(crate) comptime fn generate_public_external(f: FunctionDefinition) -> Quoted {
     let module_has_initializer = module_has_initializer(f.module());
@@ -71,8 +70,7 @@ pub(crate) comptime fn generate_public_external(f: FunctionDefinition) -> Quoted
     };
 
     let view_check = if is_fn_view(f) {
-        let assertion_message =
-            f"Function {original_function_name} can only be called statically".as_ctstring().as_quoted_str();
+        let assertion_message = f"Function {original_function_name} can only be called statically".as_quoted_str();
         quote { assert(self.context.is_static_call(), $assertion_message); }
     } else {
         quote {}
diff --git a/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external_functions_registry.nr b/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external_functions_registry.nr
@@ -6,21 +6,14 @@
 //! macro, because at that point we do not have the information about whether a given external function is private or
 //! public. We only see the internal attribute itself, not the "private" or "public" argument.
 
-use poseidon::poseidon2::Poseidon2Hasher;
-use std::{collections::umap::UHashMap, hash::BuildHasherDefault};
+use crate::utils::cmap::CHashMap;
 
 // Key is a contract module, value is an array of function definitions.
-comptime mut global PRIVATE_REGISTRY: UHashMap<Module, [FunctionDefinition], BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
-comptime mut global PUBLIC_REGISTRY: UHashMap<Module, [FunctionDefinition], BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
-comptime mut global UTILITY_REGISTRY: UHashMap<Module, [FunctionDefinition], BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
-
-comptime fn add_to_registry(
-    registry: &mut UHashMap<Module, [FunctionDefinition], BuildHasherDefault<Poseidon2Hasher>>,
-    f: FunctionDefinition,
-) {
+comptime mut global PRIVATE_REGISTRY: CHashMap<Module, [FunctionDefinition]> = CHashMap::new();
+comptime mut global PUBLIC_REGISTRY: CHashMap<Module, [FunctionDefinition]> = CHashMap::new();
+comptime mut global UTILITY_REGISTRY: CHashMap<Module, [FunctionDefinition]> = CHashMap::new();
+
+comptime fn add_to_registry(registry: &mut CHashMap<Module, [FunctionDefinition]>, f: FunctionDefinition) {
     let module = f.module();
     let current_functions = registry.get(module);
     let functions_to_insert = if current_functions.is_some() {
diff --git a/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/internal_functions_registry.nr b/noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/internal_functions_registry.nr
@@ -5,19 +5,13 @@
 //! macro, because at that point we do not have the information about whether a given internal function is private or
 //! public. We only see the internal attribute itself, not the "private" or "public" argument.
 
-use poseidon::poseidon2::Poseidon2Hasher;
-use std::{collections::umap::UHashMap, hash::BuildHasherDefault};
+use crate::utils::cmap::CHashMap;
 
 // Key is a contract module, value is an array of function definitions.
-comptime mut global PRIVATE_INTERNAL_REGISTRY: UHashMap<Module, [FunctionDefinition], BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
-comptime mut global PUBLIC_INTERNAL_REGISTRY: UHashMap<Module, [FunctionDefinition], BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
+comptime mut global PRIVATE_INTERNAL_REGISTRY: CHashMap<Module, [FunctionDefinition]> = CHashMap::new();
+comptime mut global PUBLIC_INTERNAL_REGISTRY: CHashMap<Module, [FunctionDefinition]> = CHashMap::new();
 
-comptime fn add_to_registry(
-    registry: &mut UHashMap<Module, [FunctionDefinition], BuildHasherDefault<Poseidon2Hasher>>,
-    f: FunctionDefinition,
-) {
+comptime fn add_to_registry(registry: &mut CHashMap<Module, [FunctionDefinition]>, f: FunctionDefinition) {
     let module = f.module();
     let current_functions = registry.get(module);
     let functions_to_insert = if current_functions.is_some() {
diff --git a/noir-projects/aztec-nr/aztec/src/macros/notes.nr b/noir-projects/aztec-nr/aztec/src/macros/notes.nr
@@ -1,5 +1,5 @@
 use crate::note::note_getter_options::PropertySelector;
-use std::{collections::bounded_vec::BoundedVec, meta::{ctstring::AsCtString, type_of}};
+use std::{collections::bounded_vec::BoundedVec, meta::type_of};
 
 /// Maximum number of note types within 1 contract.
 comptime global MAX_NOTE_TYPES: u32 = 128;
@@ -34,7 +34,7 @@ comptime fn get_next_note_type_id() -> Field {
 comptime fn generate_note_type_impl(s: TypeDefinition, note_type_id: Field) -> Quoted {
     let name = s.name();
     let typ = s.as_type();
-    let note_type_name: str<_> = f"{name}".as_ctstring().as_quoted_str!();
+    let note_type_name: str<_> = f"{name}".as_quoted_str!();
     let max_note_packed_len = crate::messages::logs::note::MAX_NOTE_PACKED_LEN;
 
     quote {
diff --git a/noir-projects/aztec-nr/aztec/src/macros/storage.nr b/noir-projects/aztec-nr/aztec/src/macros/storage.nr
@@ -1,12 +1,9 @@
-use poseidon::poseidon2::Poseidon2Hasher;
-use std::{collections::umap::UHashMap, hash::BuildHasherDefault};
-
-use super::utils::AsStrQuote;
+use crate::utils::cmap::CHashMap;
+use std::meta::unquote;
 
 /// Stores a map from a module to the name of the struct that describes its storage layout. This is then used when
 /// generating a `storage_layout()` getter on the contract struct.
-pub comptime mut global STORAGE_LAYOUT_NAME: UHashMap<Module, Quoted, BuildHasherDefault<Poseidon2Hasher>> =
-    UHashMap::default();
+pub comptime mut global STORAGE_LAYOUT_NAME: CHashMap<Module, Quoted> = CHashMap::new();
 
 /// This function
 /// - marks the contract as having storage, so that `macros::utils::module_has_storage` will return true,
@@ -83,7 +80,8 @@ pub comptime fn storage(s: TypeDefinition) -> Quoted {
     let module = s.module();
     let module_name = module.name();
     let storage_layout_name = f"STORAGE_LAYOUT_{module_name}".quoted_contents();
-    let (module_name_str, module_name_len) = module_name.as_str_quote();
+    let module_name_str = f"{module_name}".as_quoted_str();
+    let module_name_len: u32 = unquote!(quote { $module_name_str.as_bytes().len()});
     STORAGE_LAYOUT_NAME.insert(module, storage_layout_name);
 
     quote {
diff --git a/noir-projects/aztec-nr/aztec/src/macros/utils.nr b/noir-projects/aztec-nr/aztec/src/macros/utils.nr
@@ -1,5 +1,5 @@
 use crate::protocol::meta::derive_serialize;
-use std::meta::unquote;
+use std::meta::{ctstring::AsCtString, unquote};
 
 pub(crate) comptime fn is_fn_external(f: FunctionDefinition) -> bool {
     f.has_named_attribute("external")
@@ -41,70 +41,50 @@ pub(crate) comptime fn fn_has_authorize_once(f: FunctionDefinition) -> bool {
     f.has_named_attribute("authorize_once")
 }
 
-comptime fn signature_of_type(typ: Type) -> Quoted {
+comptime fn signature_of_type(typ: Type) -> CtString {
     if typ.is_field() {
-        quote {Field}
+        "Field".as_ctstring()
     } else if typ.as_integer().is_some() {
         let (is_signed, bit_size) = typ.as_integer().unwrap();
         if is_signed {
-            f"i{bit_size}".quoted_contents()
+            f"i{bit_size}".as_ctstring()
         } else {
-            f"u{bit_size}".quoted_contents()
+            f"u{bit_size}".as_ctstring()
         }
     } else if typ.is_bool() {
-        quote {bool}
+        "bool".as_ctstring()
     } else if typ.as_str().is_some() {
         let str_len_typ = typ.as_str().unwrap();
         let str_len = str_len_typ.as_constant().unwrap();
-        f"str<{str_len}>".quoted_contents()
+        f"str<{str_len}>".as_ctstring()
     } else if typ.as_array().is_some() {
         let (element_type, array_len) = typ.as_array().unwrap();
         let array_len = if array_len.as_constant().is_some() {
             let const_len = array_len.as_constant().unwrap();
-            f"{const_len}".quoted_contents()
+            f"{const_len}".as_ctstring()
         } else {
             // If array length is not a constant, chances are it's a numeric generic This will make the signature of a
             // type like Struct<N> be the same regardless of the value of N, but at least this fn will not crash. It is
             // up to the caller to decide whether this is acceptable.
-            f"{array_len}".quoted_contents()
+            f"{array_len}".as_ctstring()
         };
 
         let element_typ_quote = signature_of_type(element_type);
-        f"[{element_typ_quote};{array_len}]".quoted_contents()
+        f"[{element_typ_quote};{array_len}]".as_ctstring()
     } else if typ.as_data_type().is_some() {
         let (s, generics) = typ.as_data_type().unwrap();
-        let field_signatures = s.fields(generics).map(|(_, typ, _)| signature_of_type(typ)).join(quote {,});
-        f"({field_signatures})".quoted_contents()
+        let field_signatures = s.fields(generics).map(|(_, typ, _)| signature_of_type(typ)).join(",".as_ctstring());
+        f"({field_signatures})".as_ctstring()
     } else if typ.as_tuple().is_some() {
         // Note that tuples are handled the same way as structs
         let types = typ.as_tuple().unwrap();
-        let field_signatures = types.map(|typ: Type| signature_of_type(typ)).join(quote {,});
-        f"({field_signatures})".quoted_contents()
+        let field_signatures = types.map(|typ: Type| signature_of_type(typ)).join(",".as_ctstring());
+        f"({field_signatures})".as_ctstring()
     } else {
         panic(f"Unsupported type {typ}")
     }
 }
 
-pub trait AsStrQuote {
-    fn as_str_quote(self) -> (Self, u32);
-}
-
-impl AsStrQuote for Quoted {
-    // Used to convert an arbitrary quoted type into a quoted string, removing whitespace between tokens
-    comptime fn as_str_quote(self) -> (Quoted, u32) {
-        let tokens = self.tokens();
-        let mut string = CtString::new();
-        for token in tokens {
-            let token_as_fmt_str = f"{token}";
-            string = string.append_fmtstr(token_as_fmt_str);
-        }
-        let string: str<_> = string.as_quoted_str!();
-        let array = string.as_bytes();
-        let total_len = array.len();
-        (quote { $string }, total_len)
-    }
-}
-
 pub(crate) comptime fn compute_fn_selector(f: FunctionDefinition) -> Field {
     // The function selector is computed from the function signature, which is made up of the function name and types
     // of parameters, but not including the return type. For example, given:
@@ -113,12 +93,11 @@ pub(crate) comptime fn compute_fn_selector(f: FunctionDefinition) -> Field {
     //
     // The signature will be "foo(Field,AztecAddress)".
     let fn_name = f.name();
-    let args_signatures = f.parameters().map(|(_, typ)| signature_of_type(typ)).join(quote {,});
-    let signature_quote = quote { $fn_name($args_signatures) };
-    let (signature_str_quote, _) = signature_quote.as_str_quote();
+    let args_signatures = f.parameters().map(|(_, typ)| signature_of_type(typ)).join(",".as_ctstring());
+    let signature_quote = f"{fn_name}({args_signatures})".as_ctstring();
 
     let computation_quote = quote {
-        crate::protocol::traits::ToField::to_field(crate::protocol::abis::function_selector::FunctionSelector::from_signature($signature_str_quote))
+        crate::protocol::traits::ToField::to_field(crate::protocol::abis::function_selector::FunctionSelector::from_signature($signature_quote))
     };
     unquote!(computation_quote)
 }
@@ -140,12 +119,11 @@ pub(crate) comptime fn compute_struct_selector(s: TypeDefinition, from_signature
             // handled here!
             signature_of_type(typ)
         })
-        .join(quote {,});
-    let signature_quote = quote { $struct_name($args_signatures) };
-    let (signature_str_quote, _) = signature_quote.as_str_quote();
+        .join(",".as_ctstring());
+    let signature_quote = f"{struct_name}({args_signatures})".as_ctstring();
 
     let computation_quote = quote {
-        crate::protocol::traits::ToField::to_field($from_signature_fn($signature_str_quote))
+        crate::protocol::traits::ToField::to_field($from_signature_fn($signature_quote))
     };
     unquote!(computation_quote)
 }
diff --git a/noir-projects/aztec-nr/aztec/src/utils/cmap.nr b/noir-projects/aztec-nr/aztec/src/utils/cmap.nr
diff --git a/noir-projects/aztec-nr/aztec/src/utils/mod.nr b/noir-projects/aztec-nr/aztec/src/utils/mod.nr
