feat: LSP fields, functions and methods completion after "." and "::"

aztec_macros/src/transforms/contract_interface.rs

@@ -324,101 +324,105 @@ pub fn update_fn_signatures_in_contract_interface(
             });
 
         if let Some(interface_struct) = maybe_interface_struct {
-            let methods = context.def_interner.get_struct_methods(interface_struct.borrow().id);
-
-            for func_id in methods.iter().flat_map(|methods| methods.direct.iter()) {
-                let name = context.def_interner.function_name(func_id);
-                let fn_parameters = &context.def_interner.function_meta(func_id).parameters.clone();
-
-                if name == "at" || name == "interface" || name == "storage" {
-                    continue;
-                }
-
-                let fn_signature_hash = compute_fn_signature_hash(
-                    name,
-                    &fn_parameters
-                        .iter()
-                        .skip(1)
-                        .map(|(_, typ, _)| typ.clone())
-                        .collect::<Vec<Type>>(),
-                );
-                let hir_func = context.def_interner.function(func_id).block(&context.def_interner);
-
-                let function_selector_statement = context.def_interner.statement(
-                    hir_func.statements().get(hir_func.statements().len() - 2).ok_or((
-                        AztecMacroError::CouldNotGenerateContractInterface {
-                            secondary_message: Some(
-                                "Function signature statement not found, invalid body length"
-                                    .to_string(),
-                            ),
-                        },
-                        file_id,
-                    ))?,
-                );
-                let function_selector_expression_id = match function_selector_statement {
-                    HirStatement::Let(let_statement) => Ok(let_statement.expression),
-                    _ => Err((
-                        AztecMacroError::CouldNotGenerateContractInterface {
-                            secondary_message: Some(
-                                "Function selector statement must be an expression".to_string(),
-                            ),
-                        },
-                        file_id,
-                    )),
-                }?;
-                let function_selector_expression =
-                    context.def_interner.expression(&function_selector_expression_id);
-
-                let current_fn_signature_expression_id = match function_selector_expression {
-                    HirExpression::Call(call_expr) => Ok(call_expr.arguments[0]),
-                    _ => Err((
-                        AztecMacroError::CouldNotGenerateContractInterface {
-                            secondary_message: Some(
-                                "Function selector argument expression must be call expression"
-                                    .to_string(),
-                            ),
-                        },
-                        file_id,
-                    )),
-                }?;
-
-                let current_fn_signature_expression =
-                    context.def_interner.expression(&current_fn_signature_expression_id);
+            if let Some(methods) =
+                context.def_interner.get_struct_methods(interface_struct.borrow().id).cloned()
+            {
+                for func_id in methods.iter().flat_map(|(_name, methods)| methods.direct.iter()) {
+                    let name = context.def_interner.function_name(func_id);
+                    let fn_parameters =
+                        &context.def_interner.function_meta(func_id).parameters.clone();
+
+                    if name == "at" || name == "interface" || name == "storage" {
+                        continue;
+                    }
 
-                match current_fn_signature_expression {
-                    HirExpression::Literal(HirLiteral::Integer(value, _)) => {
-                        if !value.is_zero() {
-                            Err((
+                    let fn_signature_hash = compute_fn_signature_hash(
+                        name,
+                        &fn_parameters
+                            .iter()
+                            .skip(1)
+                            .map(|(_, typ, _)| typ.clone())
+                            .collect::<Vec<Type>>(),
+                    );
+                    let hir_func =
+                        context.def_interner.function(func_id).block(&context.def_interner);
+
+                    let function_selector_statement = context.def_interner.statement(
+                        hir_func.statements().get(hir_func.statements().len() - 2).ok_or((
+                            AztecMacroError::CouldNotGenerateContractInterface {
+                                secondary_message: Some(
+                                    "Function signature statement not found, invalid body length"
+                                        .to_string(),
+                                ),
+                            },
+                            file_id,
+                        ))?,
+                    );
+                    let function_selector_expression_id = match function_selector_statement {
+                        HirStatement::Let(let_statement) => Ok(let_statement.expression),
+                        _ => Err((
+                            AztecMacroError::CouldNotGenerateContractInterface {
+                                secondary_message: Some(
+                                    "Function selector statement must be an expression".to_string(),
+                                ),
+                            },
+                            file_id,
+                        )),
+                    }?;
+                    let function_selector_expression =
+                        context.def_interner.expression(&function_selector_expression_id);
+
+                    let current_fn_signature_expression_id = match function_selector_expression {
+                        HirExpression::Call(call_expr) => Ok(call_expr.arguments[0]),
+                        _ => Err((
+                            AztecMacroError::CouldNotGenerateContractInterface {
+                                secondary_message: Some(
+                                    "Function selector argument expression must be call expression"
+                                        .to_string(),
+                                ),
+                            },
+                            file_id,
+                        )),
+                    }?;
+
+                    let current_fn_signature_expression =
+                        context.def_interner.expression(&current_fn_signature_expression_id);
+
+                    match current_fn_signature_expression {
+                        HirExpression::Literal(HirLiteral::Integer(value, _)) => {
+                            if !value.is_zero() {
+                                Err((
                                 AztecMacroError::CouldNotGenerateContractInterface {
                                     secondary_message: Some(
                                         "Function signature argument must be a placeholder with value 0".to_string()),
                                 },
                                 file_id,
                             ))
-                        } else {
-                            Ok(())
+                            } else {
+                                Ok(())
+                            }
                         }
-                    }
-                    _ => Err((
-                        AztecMacroError::CouldNotGenerateContractInterface {
-                            secondary_message: Some(
-                                "Function signature argument must be a literal field element"
-                                    .to_string(),
-                            ),
+                        _ => Err((
+                            AztecMacroError::CouldNotGenerateContractInterface {
+                                secondary_message: Some(
+                                    "Function signature argument must be a literal field element"
+                                        .to_string(),
+                                ),
+                            },
+                            file_id,
+                        )),
+                    }?;
+
+                    context.def_interner.update_expression(
+                        current_fn_signature_expression_id,
+                        |expr| {
+                            *expr = HirExpression::Literal(HirLiteral::Integer(
+                                FieldElement::from(fn_signature_hash as u128),
+                                false,
+                            ))
                         },
-                        file_id,
-                    )),
-                }?;
-
-                context.def_interner.update_expression(
-                    current_fn_signature_expression_id,
-                    |expr| {
-                        *expr = HirExpression::Literal(HirLiteral::Integer(
-                            FieldElement::from(fn_signature_hash as u128),
-                            false,
-                        ))
-                    },
-                );
+                    );
+                }
             }
         }
     }

compiler/noirc_frontend/src/hir_def/types.rs

@@ -665,12 +665,12 @@
             Type::Function(args, ret, env) => {
                 let closure_env_text = match **env {
                     Type::Unit => "".to_string(),
-                    _ => format!(" with env {env}"),
+                    _ => format!("[{env}]"),
                 };
 
                 let args = vecmap(args.iter(), ToString::to_string);
 
-                write!(f, "fn({}) -> {ret}{closure_env_text}", args.join(", "))
+                write!(f, "fn{closure_env_text}({}) -> {ret}", args.join(", "))
             }
             Type::MutableReference(element) => {
                 write!(f, "&mut {element}")
@@ -1991,7 +1991,7 @@
         }
 
         let recur_on_binding = |id, replacement: &Type| {
             // Prevent recuring forever if there's a `T := T` binding
             if replacement.type_variable_id() == Some(id) {
                 replacement.clone()
             } else {
@@ -2062,7 +2062,7 @@
                 Type::Tuple(fields)
             }
             Type::Forall(typevars, typ) => {
                 // Trying to substitute_helper a variable de, substitute_bound_typevarsfined within a nested Forall
                 // is usually impossible and indicative of an error in the type checker somewhere.
                 for var in typevars {
                     assert!(!type_bindings.contains_key(&var.id()));
@@ -2219,7 +2219,7 @@
 
     /// Replace any `Type::NamedGeneric` in this type with a `Type::TypeVariable`
     /// using to the same inner `TypeVariable`. This is used during monomorphization
     /// to bind to named generics since they are unbindable during type checking.
     pub fn replace_named_generics_with_type_variables(&mut self) {
         match self {
             Type::FieldElement

compiler/noirc_frontend/src/node_interner.rs

@@ -179,10 +179,10 @@
     /// may have both `impl Struct<u32> { fn foo(){} }` and `impl Struct<u8> { fn foo(){} }`.
     /// If this happens, the returned Vec will have 2 entries and we'll need to further
     /// disambiguate them by checking the type of each function.
-    struct_methods: HashMap<(StructId, String), Methods>,
+    struct_methods: HashMap<StructId, HashMap<String, Methods>>,
 
     /// Methods on primitive types defined in the stdlib.
-    primitive_methods: HashMap<(TypeMethodKey, String), Methods>,
+    primitive_methods: HashMap<TypeMethodKey, HashMap<String, Methods>>,
 
     // For trait implementation functions, this is their self type and trait they belong to
     func_id_to_trait: HashMap<FuncId, (Type, TraitId)>,
@@ -194,13 +194,13 @@
     /// Stores the [Location] of a [Type] reference
     pub(crate) type_ref_locations: Vec<(Type, Location)>,
 
     /// In Noir's metaprogramming, a noir type has the type `Type`. When these are spliced
     /// into `quoted` expressions, we preserve the original type by assigning it a unique id
     /// and creating a `Token::QuotedType(id)` from this id. We cannot create a token holding
     /// the actual type since types do not implement Send or Sync.
     quoted_types: noirc_arena::Arena<Type>,
 
     /// Determins whether to run in LSP mode. In LSP mode references are tracked.
     pub(crate) lsp_mode: bool,
 
     /// Store the location of the references in the graph.
@@ -1131,22 +1131,26 @@
         self.structs[&id].clone()
     }
 
-    pub fn get_struct_methods(&self, id: StructId) -> Vec<Methods> {
-        self.struct_methods
-            .keys()
-            .filter_map(|(key_id, name)| {
-                if key_id == &id {
-                    Some(
-                        self.struct_methods
-                            .get(&(*key_id, name.clone()))
-                            .expect("get_struct_methods given invalid StructId")
-                            .clone(),
-                    )
-                } else {
-                    None
-                }
-            })
-            .collect()
+    pub fn get_struct_methods(&self, id: StructId) -> Option<&HashMap<String, Methods>> {
+        self.struct_methods.get(&id)
+    }
+
+    fn get_primitive_methods(&self, key: TypeMethodKey) -> Option<&HashMap<String, Methods>> {
+        self.primitive_methods.get(&key)
+    }
+
+    pub fn get_type_methods(&self, typ: &Type) -> Option<&HashMap<String, Methods>> {
+        match typ {
+            Type::Struct(struct_type, _) => {
+                let struct_type = struct_type.borrow();
+                self.get_struct_methods(struct_type.id)
+            }
+            Type::Alias(type_alias, _) => {
+                let type_alias = type_alias.borrow();
+                self.get_type_methods(&type_alias.typ)
+            }
+            _ => get_type_method_key(typ).and_then(|key| self.get_primitive_methods(key)),
+        }
     }
 
     pub fn get_trait(&self, id: TraitId) -> &Trait {
@@ -1300,8 +1304,12 @@
                     return Some(existing);
                 }
 
-                let key = (id, method_name);
-                self.struct_methods.entry(key).or_default().add_method(method_id, is_trait_method);
+                self.struct_methods
+                    .entry(id)
+                    .or_default()
+                    .entry(method_name)
+                    .or_default()
+                    .add_method(method_id, is_trait_method);
                 None
             }
             Type::Error => None,
@@ -1314,7 +1322,9 @@
                     unreachable!("Cannot add a method to the unsupported type '{}'", other)
                 });
                 self.primitive_methods
-                    .entry((key, method_name))
+                    .entry(key)
+                    .or_default()
+                    .entry(method_name)
                     .or_default()
                     .add_method(method_id, is_trait_method);
                 None
@@ -1648,7 +1658,7 @@
         method_name: &str,
         force_type_check: bool,
     ) -> Option<FuncId> {
-        let methods = self.struct_methods.get(&(id, method_name.to_owned()));
+        let methods = self.struct_methods.get(&id).and_then(|h| h.get(method_name));
 
         // If there is only one method, just return it immediately.
         // It will still be typechecked later.
@@ -1673,16 +1683,16 @@
         } else {
             // Failed to find a match for the type in question, switch to looking at impls
             // for all types `T`, e.g. `impl<T> Foo for T`
-            let key = &(TypeMethodKey::Generic, method_name.to_owned());
-            let global_methods = self.primitive_methods.get(key)?;
+            let global_methods =
+                self.primitive_methods.get(&TypeMethodKey::Generic)?.get(method_name)?;
             global_methods.find_matching_method(typ, self)
         }
     }
 
     /// Looks up a given method name on the given primitive type.
     pub fn lookup_primitive_method(&self, typ: &Type, method_name: &str) -> Option<FuncId> {
         let key = get_type_method_key(typ)?;
-        let methods = self.primitive_methods.get(&(key, method_name.to_owned()))?;
+        let methods = self.primitive_methods.get(&key)?.get(method_name)?;
         self.find_matching_method(typ, Some(methods), method_name)
     }
 
@@ -1803,6 +1813,11 @@
         self.prefix_operator_traits.insert(operator, trait_id);
     }
 
+    pub fn is_operator_trait(&self, trait_id: TraitId) -> bool {
+        self.infix_operator_traits.values().any(|id| *id == trait_id)
+            || self.prefix_operator_traits.values().any(|id| *id == trait_id)
+    }
+
     /// This function is needed when creating a NodeInterner for testing so that calls
     /// to `get_operator_trait` do not panic when the stdlib isn't present.
     #[cfg(test)]
@@ -2017,7 +2032,7 @@
     }
 
     /// Iterate through each method, starting with the direct methods
-    fn iter(&self) -> impl Iterator<Item = FuncId> + '_ {
+    pub fn iter(&self) -> impl Iterator<Item = FuncId> + '_ {
         self.direct.iter().copied().chain(self.trait_impl_methods.iter().copied())
     }
 

compiler/noirc_frontend/src/parser/errors.rs

@@ -16,6 +16,8 @@ pub enum ParserErrorReason {
     ExpectedFieldName(Token),
     #[error("expected a pattern but found a type - {0}")]
     ExpectedPatternButFoundType(Token),
+    #[error("expected an identifier after .")]
+    ExpectedIdentifierAfterDot,
     #[error("expected an identifier after ::")]
     ExpectedIdentifierAfterColons,
     #[error("Expected a ; separating these two statements")]