Add enum variant references CodeLens.

4 files changed, 60 insertions, 12 deletions

diff --git a/crates/ide/src/annotations.rs b/crates/ide/src/annotations.rs
index 472b396ac5d..6c75bf34923 100644
--- a/crates/ide/src/annotations.rs
+++ b/crates/ide/src/annotations.rs
@@ -5,7 +5,7 @@ use ide_db::{
     helpers::visit_file_defs,
     RootDatabase,
 };
-use syntax::{ast::HasName, AstNode, TextRange};
+use syntax::{ast::HasName, AstNode, TextRange, TextSize};
 
 use crate::{
     fn_references::find_all_methods,
@@ -40,6 +40,7 @@ pub struct AnnotationConfig {
     pub annotate_impls: bool,
     pub annotate_references: bool,
     pub annotate_method_references: bool,
+    pub annotate_enum_variant_references: bool,
 }
 
 pub(crate) fn annotations(
@@ -63,18 +64,33 @@ pub(crate) fn annotations(
 
     visit_file_defs(&Semantics::new(db), file_id, &mut |def| match def {
         Either::Left(def) => {
-            let range = match def {
+            let (range, ranges_variants) = match def {
                 hir::ModuleDef::Const(konst) => {
-                    konst.source(db).and_then(|node| name_range(&node, file_id))
+                    (konst.source(db).and_then(|node| name_range(&node, file_id)), vec![None])
                 }
                 hir::ModuleDef::Trait(trait_) => {
-                    trait_.source(db).and_then(|node| name_range(&node, file_id))
+                    (trait_.source(db).and_then(|node| name_range(&node, file_id)), vec![None])
                 }
-                hir::ModuleDef::Adt(adt) => {
-                    adt.source(db).and_then(|node| name_range(&node, file_id))
-                }
-                _ => None,
+                hir::ModuleDef::Adt(adt) => match adt {
+                    hir::Adt::Enum(enum_) => (
+                        enum_.source(db).and_then(|node| name_range(&node, file_id)),
+                        if config.annotate_enum_variant_references {
+                            enum_
+                                .variants(db)
+                                .into_iter()
+                                .map(|variant| {
+                                    variant.source(db).and_then(|node| name_range(&node, file_id))
+                                })
+                                .collect()
+                        } else {
+                            vec![None]
+                        },
+                    ),
+                    _ => (adt.source(db).and_then(|node| name_range(&node, file_id)), vec![None]),
+                },
+                _ => (None, vec![None]),
             };
+
             let (range, offset) = match range {
                 Some(range) => (range, range.start()),
                 None => return,
@@ -99,6 +115,26 @@ pub(crate) fn annotations(
                 });
             }
 
+            if config.annotate_enum_variant_references {
+                let mut variants_metadata: Vec<(TextRange, TextSize)> = Vec::new();
+                for range_variant in ranges_variants.into_iter() {
+                    let (range, offset) = match range_variant {
+                        Some(range) => (range, range.start()),
+                        None => return,
+                    };
+                    variants_metadata.push((range, offset))
+                }
+                for variant_metadata in variants_metadata.into_iter() {
+                    annotations.push(Annotation {
+                        range: variant_metadata.0,
+                        kind: AnnotationKind::HasReferences {
+                            position: FilePosition { file_id, offset: variant_metadata.1 },
+                            data: None,
+                        },
+                    });
+                }
+            }
+
             fn name_range<T: HasName>(node: &InFile<T>, file_id: FileId) -> Option<TextRange> {
                 if node.file_id == file_id.into() {
                     node.value.name().map(|it| it.syntax().text_range())
@@ -173,6 +209,7 @@ mod tests {
                     annotate_impls: true,
                     annotate_references: true,
                     annotate_method_references: true,
+                    annotate_enum_variant_references: true,
                 },
                 file_id,
             )
diff --git a/crates/rust-analyzer/src/config.rs b/crates/rust-analyzer/src/config.rs
index 6c098333683..505123fc6af 100644
--- a/crates/rust-analyzer/src/config.rs
+++ b/crates/rust-analyzer/src/config.rs
@@ -225,9 +225,12 @@ config_data! {
         /// Whether to show `Method References` lens. Only applies when
         /// `#rust-analyzer.lens.enable#` is set.
         lens_methodReferences: bool = "false",
-        /// Whether to show `References` lens. Only applies when
-        /// `#rust-analyzer.lens.enable#` is set.
+        /// Whether to show `References` lens for Struct, Enum, Union and Trait.
+        /// Only applies when `#rust-analyzer.lens.enable#` is set.
         lens_references: bool = "false",
+        /// Whether to show `References` lens for Enum Variants.
+        /// Only applies when `#rust-analyzer.lens.enable#` is set.
+        lens_enumVariantReferences: bool = "false",
         /// Internal config: use custom client-side commands even when the
         /// client doesn't set the corresponding capability.
         lens_forceCustomCommands: bool = "true",
@@ -323,6 +326,7 @@ pub struct LensConfig {
     pub implementations: bool,
     pub method_refs: bool,
     pub refs: bool, // for Struct, Enum, Union and Trait
+    pub enum_variant_refs: bool,
 }
 
 impl LensConfig {
@@ -339,7 +343,7 @@ impl LensConfig {
     }
 
     pub fn references(&self) -> bool {
-        self.method_refs || self.refs
+        self.method_refs || self.refs || self.enum_variant_refs
     }
 }
 
@@ -805,6 +809,7 @@ impl Config {
             implementations: self.data.lens_enable && self.data.lens_implementations,
             method_refs: self.data.lens_enable && self.data.lens_methodReferences,
             refs: self.data.lens_enable && self.data.lens_references,
+            enum_variant_refs: self.data.lens_enable && self.data.lens_enumVariantReferences,
         }
     }
     pub fn hover_actions(&self) -> HoverActionsConfig {
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index e62bb9499fa..412a52d1c36 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -1135,6 +1135,7 @@ pub(crate) fn handle_code_lens(
             annotate_impls: lens_config.implementations,
             annotate_references: lens_config.refs,
             annotate_method_references: lens_config.method_refs,
+            annotate_enum_variant_references: lens_config.enum_variant_refs,
         },
         file_id,
     )?;
diff --git a/editors/code/package.json b/editors/code/package.json
index 4ca93b17bc3..9e08e638bbe 100644
--- a/editors/code/package.json
+++ b/editors/code/package.json
@@ -793,7 +793,12 @@
                     "type": "boolean"
                 },
                 "rust-analyzer.lens.references": {
-                    "markdownDescription": "Whether to show `References` lens. Only applies when\n`#rust-analyzer.lens.enable#` is set.",
+                    "markdownDescription": "Whether to show `References` lens for Struct, Enum, Union and Trait. Only applies when\n`#rust-analyzer.lens.enable#` is set.",
+                    "default": false,
+                    "type": "boolean"
+                },
+                "rust-analyzer.lens.enumVariantReferences": {
+                    "markdownDescription": "Whether to show `References` lens for Enum Variants. Only applies when\n`#rust-analyzer.lens.enable#` is set.",
                     "default": false,
                     "type": "boolean"
                 },