fix(lsp): suggest all possible trait methods, but only visible ones (#7027)

Co-authored-by: jfecher <jake@aztecprotocol.com>

Author: asterite

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -1322,9 +1322,15 @@ impl<'context> Elaborator<'context> {
             let span = trait_impl.object_type.span;
             self.declare_methods_on_struct(Some(trait_id), &mut trait_impl.methods, span);
 
+            let trait_visibility = self.interner.get_trait(trait_id).visibility;
+
             let methods = trait_impl.methods.function_ids();
             for func_id in &methods {
                 self.interner.set_function_trait(*func_id, self_type.clone(), trait_id);
+
+                // A trait impl method has the same visibility as its trait
+                let modifiers = self.interner.function_modifiers_mut(func_id);
+                modifiers.visibility = trait_visibility;
             }
 
             let trait_generics = trait_impl.resolved_trait_generics.clone();

compiler/noirc_frontend/src/elaborator/traits.rs

@@ -56,6 +56,7 @@ impl<'context> Elaborator<'context> {
                 this.interner.update_trait(*trait_id, |trait_def| {
                     trait_def.set_trait_bounds(resolved_trait_bounds);
                     trait_def.set_where_clause(where_clause);
+                    trait_def.set_visibility(unresolved_trait.trait_def.visibility);
                 });
 
                 let methods = this.resolve_trait_methods(*trait_id, unresolved_trait);

compiler/noirc_frontend/src/hir/def_collector/dc_mod.rs

@@ -1228,9 +1228,10 @@ pub(crate) fn collect_trait_impl_items(
     for item in std::mem::take(&mut trait_impl.items) {
         match item.item.kind {
             TraitImplItemKind::Function(mut impl_method) => {
-                // Regardless of what visibility was on the source code, treat it as public
-                // (a warning is produced during parsing for this)
-                impl_method.def.visibility = ItemVisibility::Public;
+                // Set the impl method visibility as temporarily private.
+                // Eventually when we find out what trait is this impl for we'll set it
+                // to the trait's visibility.
+                impl_method.def.visibility = ItemVisibility::Private;
 
                 let func_id = interner.push_empty_fn();
                 let location = Location::new(impl_method.span(), file_id);

compiler/noirc_frontend/src/hir_def/traits.rs

@@ -1,7 +1,7 @@
 use iter_extended::vecmap;
 use rustc_hash::FxHashMap as HashMap;
 
-use crate::ast::{Ident, NoirFunction};
+use crate::ast::{Ident, ItemVisibility, NoirFunction};
 use crate::hir::type_check::generics::TraitGenerics;
 use crate::ResolvedGeneric;
 use crate::{
@@ -66,6 +66,7 @@ pub struct Trait {
     pub name: Ident,
     pub generics: Generics,
     pub location: Location,
+    pub visibility: ItemVisibility,
 
     /// When resolving the types of Trait elements, all references to `Self` resolve
     /// to this TypeVariable. Then when we check if the types of trait impl elements
@@ -160,6 +161,10 @@ impl Trait {
         self.where_clause = where_clause;
     }
 
+    pub fn set_visibility(&mut self, visibility: ItemVisibility) {
+        self.visibility = visibility;
+    }
+
     pub fn find_method(&self, name: &str) -> Option<TraitMethodId> {
         for (idx, method) in self.methods.iter().enumerate() {
             if &method.name == name {

compiler/noirc_frontend/src/node_interner.rs

@@ -729,6 +729,7 @@ impl NodeInterner {
             crate_id: unresolved_trait.crate_id,
             location: Location::new(unresolved_trait.trait_def.span, unresolved_trait.file_id),
             generics,
+            visibility: ItemVisibility::Private,
             self_type_typevar: TypeVariable::unbound(self.next_type_variable_id(), Kind::Normal),
             methods: Vec::new(),
             method_ids: unresolved_trait.method_ids.clone(),
@@ -2268,31 +2269,26 @@ impl Methods {
     }
 
     /// Iterate through each method, starting with the direct methods
-    pub fn iter(&self) -> impl Iterator<Item = (FuncId, &Option<Type>)> + '_ {
-        let trait_impl_methods = self.trait_impl_methods.iter().map(|m| (m.method, &m.typ));
-        let direct = self.direct.iter().copied().map(|func_id| {
-            let typ: &Option<Type> = &None;
-            (func_id, typ)
-        });
+    pub fn iter(&self) -> impl Iterator<Item = (FuncId, Option<&Type>, Option<TraitId>)> {
+        let trait_impl_methods =
+            self.trait_impl_methods.iter().map(|m| (m.method, m.typ.as_ref(), Some(m.trait_id)));
+        let direct = self.direct.iter().copied().map(|func_id| (func_id, None, None));
         direct.chain(trait_impl_methods)
     }
 
-    /// Select the 1 matching method with an object type matching `typ`
-    pub fn find_matching_method(
-        &self,
-        typ: &Type,
+    pub fn find_matching_methods<'a>(
+        &'a self,
+        typ: &'a Type,
         has_self_param: bool,
-        interner: &NodeInterner,
-    ) -> Option<FuncId> {
-        // When adding methods we always check they do not overlap, so there should be
-        // at most 1 matching method in this list.
-        for (method, method_type) in self.iter() {
+        interner: &'a NodeInterner,
+    ) -> impl Iterator<Item = (FuncId, Option<TraitId>)> + 'a {
+        self.iter().filter_map(move |(method, method_type, trait_id)| {
             if Self::method_matches(typ, has_self_param, method, method_type, interner) {
-                return Some(method);
+                Some((method, trait_id))
+            } else {
+                None
             }
-        }
-
-        None
+        })
     }
 
     pub fn find_direct_method(
@@ -2302,7 +2298,7 @@ impl Methods {
         interner: &NodeInterner,
     ) -> Option<FuncId> {
         for method in &self.direct {
-            if Self::method_matches(typ, has_self_param, *method, &None, interner) {
+            if Self::method_matches(typ, has_self_param, *method, None, interner) {
                 return Some(*method);
             }
         }
@@ -2320,7 +2316,7 @@ impl Methods {
 
         for trait_impl_method in &self.trait_impl_methods {
             let method = trait_impl_method.method;
-            let method_type = &trait_impl_method.typ;
+            let method_type = trait_impl_method.typ.as_ref();
             let trait_id = trait_impl_method.trait_id;
 
             if Self::method_matches(typ, has_self_param, method, method_type, interner) {
@@ -2335,7 +2331,7 @@ impl Methods {
         typ: &Type,
         has_self_param: bool,
         method: FuncId,
-        method_type: &Option<Type>,
+        method_type: Option<&Type>,
         interner: &NodeInterner,
     ) -> bool {
         match interner.function_meta(&method).typ.instantiate(interner).0 {

tooling/lsp/src/requests/completion.rs

@@ -28,6 +28,7 @@ use noirc_frontend::{
         def_map::{CrateDefMap, LocalModuleId, ModuleDefId, ModuleId},
         resolution::visibility::{
             item_in_module_is_visible, method_call_is_visible, struct_member_is_visible,
+            trait_member_is_visible,
         },
     },
     hir_def::traits::Trait,
@@ -245,7 +246,7 @@ impl<'a> NodeFinder<'a> {
         let mut idents: Vec<Ident> = Vec::new();
 
         // Find in which ident we are in, and in which part of it
-        // (it could be that we are completting in the middle of an ident)
+        // (it could be that we are completing in the middle of an ident)
         for segment in &path.segments {
             let ident = &segment.ident;
 
@@ -658,39 +659,47 @@ impl<'a> NodeFinder<'a> {
         let has_self_param = matches!(function_kind, FunctionKind::SelfType(..));
 
         for (name, methods) in methods_by_name {
-            let Some(func_id) =
-                methods.find_matching_method(typ, has_self_param, self.interner).or_else(|| {
-                    // Also try to find a method assuming typ is `&mut typ`:
-                    // we want to suggest methods that take `&mut self` even though a variable might not
-                    // be mutable, so a user can know they need to mark it as mutable.
-                    let typ = Type::MutableReference(Box::new(typ.clone()));
-                    methods.find_matching_method(&typ, has_self_param, self.interner)
-                })
-            else {
+            if !name_matches(name, prefix) {
                 continue;
-            };
-
-            if let Some(struct_id) = struct_id {
-                let modifiers = self.interner.function_modifiers(&func_id);
-                let visibility = modifiers.visibility;
-                if !struct_member_is_visible(struct_id, visibility, self.module_id, self.def_maps) {
-                    continue;
-                }
             }
 
-            if is_primitive
-                && !method_call_is_visible(
-                    typ,
-                    func_id,
-                    self.module_id,
-                    self.interner,
-                    self.def_maps,
-                )
+            for (func_id, trait_id) in
+                methods.find_matching_methods(typ, has_self_param, self.interner)
             {
-                continue;
-            }
+                if let Some(struct_id) = struct_id {
+                    let modifiers = self.interner.function_modifiers(&func_id);
+                    let visibility = modifiers.visibility;
+                    if !struct_member_is_visible(
+                        struct_id,
+                        visibility,
+                        self.module_id,
+                        self.def_maps,
+                    ) {
+                        continue;
+                    }
+                }
+
+                if let Some(trait_id) = trait_id {
+                    let modifiers = self.interner.function_modifiers(&func_id);
+                    let visibility = modifiers.visibility;
+                    if !trait_member_is_visible(trait_id, visibility, self.module_id, self.def_maps)
+                    {
+                        continue;
+                    }
+                }
+
+                if is_primitive
+                    && !method_call_is_visible(
+                        typ,
+                        func_id,
+                        self.module_id,
+                        self.interner,
+                        self.def_maps,
+                    )
+                {
+                    continue;
+                }
 
-            if name_matches(name, prefix) {
                 let completion_items = self.function_completion_items(
                     name,
                     func_id,

tooling/lsp/src/requests/completion/tests.rs

@@ -2879,4 +2879,55 @@ fn main() {
         let items = get_completions(src).await;
         assert_eq!(items.len(), 1);
     }
+
+    #[test]
+    async fn test_does_not_suggest_trait_function_not_visible() {
+        let src = r#"
+        mod moo {
+            trait Foo {
+                fn foobar();
+            }
+
+            impl Foo for Field {
+                fn foobar() {}
+            }
+        }
+
+        fn main() {
+            Field::fooba>|<
+        }
+
+        "#;
+        assert_completion(src, vec![]).await;
+    }
+
+    #[test]
+    async fn test_suggests_multiple_trait_methods() {
+        let src = r#"
+        mod moo {
+            pub trait Foo {
+                fn foobar();
+            }
+
+            impl Foo for Field {
+                fn foobar() {}
+            }
+
+            pub trait Bar {
+                fn foobar();
+            }
+
+            impl Bar for Field {
+                fn foobar() {}
+            }
+        }
+
+        fn main() {
+            Field::fooba>|<
+        }
+
+        "#;
+        let items = get_completions(src).await;
+        assert_eq!(items.len(), 2);
+    }
 }

tooling/lsp/src/requests/hover.rs

@@ -399,7 +399,10 @@ fn format_function(id: FuncId, args: &ProcessRequestCallbackArgs) -> String {
     }
     string.push_str("    ");
 
-    if func_modifiers.visibility != ItemVisibility::Private {
+    if func_modifiers.visibility != ItemVisibility::Private
+        && func_meta.trait_id.is_none()
+        && func_meta.trait_impl.is_none()
+    {
         string.push_str(&func_modifiers.visibility.to_string());
         string.push(' ');
     }
@@ -1154,7 +1157,7 @@ mod hover_tests {
         assert!(hover_text.starts_with(
             "    two
     impl<A> Bar<A, i32> for Foo<A>
-    pub fn bar_stuff(self)"
+    fn bar_stuff(self)"
         ));
     }