feat: let LSP suggest traits in trait bounds (#6370)

compiler/noirc_frontend/src/ast/visitor.rs

@@ -22,8 +22,8 @@ use crate::{
 
 use super::{
     ForBounds, FunctionReturnType, GenericTypeArgs, IntegerBitSize, ItemVisibility, Pattern,
-    Signedness, TraitImplItemKind, TypePath, UnresolvedGenerics, UnresolvedTraitConstraint,
-    UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
+    Signedness, TraitBound, TraitImplItemKind, TypePath, UnresolvedGenerics,
+    UnresolvedTraitConstraint, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
 };
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -438,6 +438,14 @@ pub trait Visitor {
         true
     }
 
+    fn visit_trait_bound(&mut self, _: &TraitBound) -> bool {
+        true
+    }
+
+    fn visit_unresolved_trait_constraint(&mut self, _: &UnresolvedTraitConstraint) -> bool {
+        true
+    }
+
     fn visit_pattern(&mut self, _: &Pattern) -> bool {
         true
     }
@@ -555,6 +563,12 @@ impl NoirFunction {
             param.typ.accept(visitor);
         }
 
+        self.def.return_type.accept(visitor);
+
+        for constraint in &self.def.where_clause {
+            constraint.accept(visitor);
+        }
+
         self.def.body.accept(None, visitor);
     }
 }
@@ -645,6 +659,14 @@ impl NoirTrait {
             attribute.accept(AttributeTarget::Trait, visitor);
         }
 
+        for bound in &self.bounds {
+            bound.accept(visitor);
+        }
+
+        for constraint in &self.where_clause {
+            constraint.accept(visitor);
+        }
+
         for item in &self.items {
             item.item.accept(visitor);
         }
@@ -686,7 +708,7 @@ impl TraitItem {
                     return_type.accept(visitor);
 
                     for unresolved_trait_constraint in where_clause {
-                        unresolved_trait_constraint.typ.accept(visitor);
+                        unresolved_trait_constraint.accept(visitor);
                     }
 
                     if let Some(body) = body {
@@ -1346,6 +1368,32 @@ impl FunctionReturnType {
     }
 }
 
+impl TraitBound {
+    pub fn accept(&self, visitor: &mut impl Visitor) {
+        if visitor.visit_trait_bound(self) {
+            self.accept_children(visitor);
+        }
+    }
+
+    pub fn accept_children(&self, visitor: &mut impl Visitor) {
+        self.trait_path.accept(visitor);
+        self.trait_generics.accept(visitor);
+    }
+}
+
+impl UnresolvedTraitConstraint {
+    pub fn accept(&self, visitor: &mut impl Visitor) {
+        if visitor.visit_unresolved_trait_constraint(self) {
+            self.accept_children(visitor);
+        }
+    }
+
+    pub fn accept_children(&self, visitor: &mut impl Visitor) {
+        self.typ.accept(visitor);
+        self.trait_bound.accept(visitor);
+    }
+}
+
 impl Pattern {
     pub fn accept(&self, visitor: &mut impl Visitor) {
         if visitor.visit_pattern(self) {

tooling/lsp/src/requests/completion.rs

@@ -19,9 +19,9 @@ use noirc_frontend::{
         Expression, ExpressionKind, ForLoopStatement, GenericTypeArgs, Ident, IfExpression,
         IntegerBitSize, ItemVisibility, LValue, Lambda, LetStatement, MemberAccessExpression,
         MethodCallExpression, NoirFunction, NoirStruct, NoirTraitImpl, Path, PathKind, Pattern,
-        Signedness, Statement, TraitImplItemKind, TypeImpl, TypePath, UnresolvedGeneric,
-        UnresolvedGenerics, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression, UseTree,
-        UseTreeKind, Visitor,
+        Signedness, Statement, TraitBound, TraitImplItemKind, TypeImpl, TypePath,
+        UnresolvedGeneric, UnresolvedGenerics, UnresolvedType, UnresolvedTypeData,
+        UnresolvedTypeExpression, UseTree, UseTreeKind, Visitor,
     },
     graph::{CrateId, Dependency},
     hir::{
@@ -384,7 +384,7 @@ impl<'a> NodeFinder<'a> {
                     self.builtin_types_completion(&prefix);
                     self.type_parameters_completion(&prefix);
                 }
-                RequestedItems::OnlyAttributeFunctions(..) => (),
+                RequestedItems::OnlyTraits | RequestedItems::OnlyAttributeFunctions(..) => (),
             }
             self.complete_auto_imports(&prefix, requested_items, function_completion_kind);
         }
@@ -1124,6 +1124,10 @@ impl<'a> Visitor for NodeFinder<'a> {
 
         noir_function.def.return_type.accept(self);
 
+        for constraint in &noir_function.def.where_clause {
+            constraint.accept(self);
+        }
+
         self.local_variables.clear();
         for param in &noir_function.def.parameters {
             self.collect_local_variables(&param.pattern);
@@ -1231,7 +1235,7 @@ impl<'a> Visitor for NodeFinder<'a> {
         return_type.accept(self);
 
         for unresolved_trait_constraint in where_clause {
-            unresolved_trait_constraint.typ.accept(self);
+            unresolved_trait_constraint.accept(self);
         }
 
         if let Some(body) = body {
@@ -1702,6 +1706,12 @@ impl<'a> Visitor for NodeFinder<'a> {
             last_was_dollar = false;
         }
     }
+
+    fn visit_trait_bound(&mut self, trait_bound: &TraitBound) -> bool {
+        self.find_in_path(&trait_bound.trait_path, RequestedItems::OnlyTraits);
+        trait_bound.trait_generics.accept(self);
+        false
+    }
 }
 
 fn get_field_type(typ: &Type, name: &str) -> Option<Type> {

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

@@ -35,6 +35,14 @@ impl<'a> NodeFinder<'a> {
                 | ModuleDefId::TypeAliasId(_)
                 | ModuleDefId::TraitId(_) => (),
             },
+            RequestedItems::OnlyTraits => match module_def_id {
+                ModuleDefId::FunctionId(_) | ModuleDefId::GlobalId(_) | ModuleDefId::TypeId(_) => {
+                    return Vec::new()
+                }
+                ModuleDefId::ModuleId(_)
+                | ModuleDefId::TypeAliasId(_)
+                | ModuleDefId::TraitId(_) => (),
+            },
             RequestedItems::OnlyAttributeFunctions(..) => {
                 if !matches!(module_def_id, ModuleDefId::FunctionId(..)) {
                     return Vec::new();

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

@@ -25,8 +25,10 @@ pub(super) enum FunctionKind<'a> {
 pub(super) enum RequestedItems {
     // Suggest any items (types, functions, etc.).
     AnyItems,
-    // Only suggest types.
+    // Only suggest types (and modules, because they can contain types).
     OnlyTypes,
+    // Only suggest traits (and modules, because they can contain traits).
+    OnlyTraits,
     // Only attribute functions
     OnlyAttributeFunctions(AttributeTarget),
 }

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

@@ -2719,4 +2719,42 @@ fn main() {
 
         assert_completion(src, Vec::new()).await;
     }
+
+    #[test]
+    async fn test_suggests_trait_in_trait_parent_bounds() {
+        let src = r#"
+        trait Foobar {}
+        struct Foobarbaz {}
+
+        trait Bar: Foob>|< {}
+        "#;
+        assert_completion(
+            src,
+            vec![simple_completion_item(
+                "Foobar",
+                CompletionItemKind::INTERFACE,
+                Some("Foobar".to_string()),
+            )],
+        )
+        .await;
+    }
+
+    #[test]
+    async fn test_suggests_trait_in_function_where_clause() {
+        let src = r#"
+        trait Foobar {}
+        struct Foobarbaz {}
+
+        fn foo<T>() where T: Foob>|< {}
+        "#;
+        assert_completion(
+            src,
+            vec![simple_completion_item(
+                "Foobar",
+                CompletionItemKind::INTERFACE,
+                Some("Foobar".to_string()),
+            )],
+        )
+        .await;
+    }
 }