From 50d46863ef5e9bd3162e2dc5f6039f2ca27cf0ad Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Sat, 7 Aug 2021 21:56:42 +0200
Subject: [PATCH] Look for enum variants and trait assoc functions when looking
 for lang items

---
 crates/hir_def/src/lang_item.rs               | 34 +++++++++++++++----
 .../src/completions/flyimport.rs              |  2 +-
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/crates/hir_def/src/lang_item.rs b/crates/hir_def/src/lang_item.rs
index 3a45cbfa1217..6797a77de58d 100644
--- a/crates/hir_def/src/lang_item.rs
+++ b/crates/hir_def/src/lang_item.rs
@@ -8,8 +8,8 @@ use rustc_hash::FxHashMap;
 use syntax::SmolStr;
 
 use crate::{
-    db::DefDatabase, AdtId, AttrDefId, CrateId, EnumId, FunctionId, ImplId, ModuleDefId, StaticId,
-    StructId, TraitId,
+    db::DefDatabase, AdtId, AttrDefId, CrateId, EnumId, EnumVariantId, FunctionId, ImplId,
+    ModuleDefId, StaticId, StructId, TraitId,
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -20,6 +20,7 @@ pub enum LangItemTarget {
     StaticId(StaticId),
     StructId(StructId),
     TraitId(TraitId),
+    EnumVariantId(EnumVariantId),
 }
 
 impl LangItemTarget {
@@ -64,6 +65,13 @@ impl LangItemTarget {
             _ => None,
         }
     }
+
+    pub fn as_enum_variant(self) -> Option<EnumVariantId> {
+        match self {
+            LangItemTarget::EnumVariantId(id) => Some(id),
+            _ => None,
+        }
+    }
 }
 
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
@@ -92,19 +100,31 @@ impl LangItems {
             for def in module_data.scope.declarations() {
                 match def {
                     ModuleDefId::TraitId(trait_) => {
-                        lang_items.collect_lang_item(db, trait_, LangItemTarget::TraitId)
+                        lang_items.collect_lang_item(db, trait_, LangItemTarget::TraitId);
+                        db.trait_data(trait_).items.iter().for_each(|&(_, assoc_id)| {
+                            if let crate::AssocItemId::FunctionId(f) = assoc_id {
+                                lang_items.collect_lang_item(db, f, LangItemTarget::FunctionId);
+                            }
+                        });
                     }
                     ModuleDefId::AdtId(AdtId::EnumId(e)) => {
-                        lang_items.collect_lang_item(db, e, LangItemTarget::EnumId)
+                        lang_items.collect_lang_item(db, e, LangItemTarget::EnumId);
+                        db.enum_data(e).variants.iter().for_each(|(local_id, _)| {
+                            lang_items.collect_lang_item(
+                                db,
+                                EnumVariantId { parent: e, local_id },
+                                LangItemTarget::EnumVariantId,
+                            );
+                        });
                     }
                     ModuleDefId::AdtId(AdtId::StructId(s)) => {
-                        lang_items.collect_lang_item(db, s, LangItemTarget::StructId)
+                        lang_items.collect_lang_item(db, s, LangItemTarget::StructId);
                     }
                     ModuleDefId::FunctionId(f) => {
-                        lang_items.collect_lang_item(db, f, LangItemTarget::FunctionId)
+                        lang_items.collect_lang_item(db, f, LangItemTarget::FunctionId);
                     }
                     ModuleDefId::StaticId(s) => {
-                        lang_items.collect_lang_item(db, s, LangItemTarget::StaticId)
+                        lang_items.collect_lang_item(db, s, LangItemTarget::StaticId);
                     }
                     _ => {}
                 }
diff --git a/crates/ide_completion/src/completions/flyimport.rs b/crates/ide_completion/src/completions/flyimport.rs
index 2a5aa7e91fb8..cbfe89b911c6 100644
--- a/crates/ide_completion/src/completions/flyimport.rs
+++ b/crates/ide_completion/src/completions/flyimport.rs
@@ -696,8 +696,8 @@ fn main() {
 }
 "#,
             expect![[r#"
-                ct SPECIAL_CONST (use dep::test_mod::TestTrait) DEPRECATED
                 fn weird_function() (use dep::test_mod::TestTrait) fn() DEPRECATED
+                ct SPECIAL_CONST (use dep::test_mod::TestTrait) DEPRECATED
             "#]],
         );
     }