feat: filter already present enum variants in match arms

8 files changed, 114 insertions, 16 deletions

diff --git a/crates/ide-completion/src/completions.rs b/crates/ide-completion/src/completions.rs
index cac5dccde2b..fddd02fc13e 100644
--- a/crates/ide-completion/src/completions.rs
+++ b/crates/ide-completion/src/completions.rs
@@ -23,7 +23,7 @@ pub(crate) mod env_vars;
 
 use std::iter;
 
-use hir::{known, ScopeDef};
+use hir::{known, ScopeDef, Variant};
 use ide_db::{imports::import_assets::LocatedImport, SymbolKind};
 use syntax::ast;
 
@@ -538,18 +538,25 @@ fn enum_variants_with_paths(
     enum_: hir::Enum,
     impl_: &Option<ast::Impl>,
     cb: impl Fn(&mut Completions, &CompletionContext<'_>, hir::Variant, hir::ModPath),
+    missing_variants: Option<Vec<Variant>>,
 ) {
-    let variants = enum_.variants(ctx.db);
+    let mut process_variant = |variant: Variant| {
+        let self_path = hir::ModPath::from_segments(
+            hir::PathKind::Plain,
+            iter::once(known::SELF_TYPE).chain(iter::once(variant.name(ctx.db))),
+        );
+
+        cb(acc, ctx, variant, self_path);
+    };
+
+    let variants = match missing_variants {
+        Some(missing_variants) => missing_variants,
+        None => enum_.variants(ctx.db),
+    };
 
     if let Some(impl_) = impl_.as_ref().and_then(|impl_| ctx.sema.to_def(impl_)) {
         if impl_.self_ty(ctx.db).as_adt() == Some(hir::Adt::Enum(enum_)) {
-            for &variant in &variants {
-                let self_path = hir::ModPath::from_segments(
-                    hir::PathKind::Plain,
-                    iter::once(known::SELF_TYPE).chain(iter::once(variant.name(ctx.db))),
-                );
-                cb(acc, ctx, variant, self_path);
-            }
+            variants.iter().for_each(|variant| process_variant(*variant));
         }
     }
 
diff --git a/crates/ide-completion/src/completions/expr.rs b/crates/ide-completion/src/completions/expr.rs
index 3192b21cfb2..8946011280a 100644
--- a/crates/ide-completion/src/completions/expr.rs
+++ b/crates/ide-completion/src/completions/expr.rs
@@ -208,6 +208,7 @@ pub(crate) fn complete_expr_path(
                             |acc, ctx, variant, path| {
                                 acc.add_qualified_enum_variant(ctx, path_ctx, variant, path)
                             },
+                            None,
                         );
                     }
                 }
diff --git a/crates/ide-completion/src/completions/flyimport.rs b/crates/ide-completion/src/completions/flyimport.rs
index 364969af9c9..0979f6a6dfc 100644
--- a/crates/ide-completion/src/completions/flyimport.rs
+++ b/crates/ide-completion/src/completions/flyimport.rs
@@ -5,10 +5,7 @@ use ide_db::imports::{
     insert_use::ImportScope,
 };
 use itertools::Itertools;
-use syntax::{
-    ast::{self},
-    AstNode, SyntaxNode, T,
-};
+use syntax::{ast, AstNode, SyntaxNode, T};
 
 use crate::{
     context::{
diff --git a/crates/ide-completion/src/completions/pattern.rs b/crates/ide-completion/src/completions/pattern.rs
index 58d5bf114cc..6ad6a06f11a 100644
--- a/crates/ide-completion/src/completions/pattern.rs
+++ b/crates/ide-completion/src/completions/pattern.rs
@@ -58,6 +58,7 @@ pub(crate) fn complete_pattern(
                 |acc, ctx, variant, path| {
                     acc.add_qualified_variant_pat(ctx, pattern_ctx, variant, path);
                 },
+                Some(pattern_ctx.missing_variants.clone()),
             );
         }
     }
diff --git a/crates/ide-completion/src/context.rs b/crates/ide-completion/src/context.rs
index aa77f449530..954e88e093e 100644
--- a/crates/ide-completion/src/context.rs
+++ b/crates/ide-completion/src/context.rs
@@ -220,6 +220,8 @@ pub(super) struct PatternContext {
     /// The record pattern this name or ref is a field of
     pub(super) record_pat: Option<ast::RecordPat>,
     pub(super) impl_: Option<ast::Impl>,
+    /// List of missing variants in a match expr
+    pub(super) missing_variants: Vec<hir::Variant>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
diff --git a/crates/ide-completion/src/context/analysis.rs b/crates/ide-completion/src/context/analysis.rs
index c142a7305f9..73375250b5f 100644
--- a/crates/ide-completion/src/context/analysis.rs
+++ b/crates/ide-completion/src/context/analysis.rs
@@ -1,7 +1,7 @@
 //! Module responsible for analyzing the code surrounding the cursor for completion.
 use std::iter;
 
-use hir::{Semantics, Type, TypeInfo};
+use hir::{Semantics, Type, TypeInfo, Variant};
 use ide_db::{active_parameter::ActiveParameter, RootDatabase};
 use syntax::{
     algo::{find_node_at_offset, non_trivia_sibling},
@@ -1111,6 +1111,9 @@ fn pattern_context_for(
     pat: ast::Pat,
 ) -> PatternContext {
     let mut param_ctx = None;
+
+    let mut missing_variants = vec![];
+
     let (refutability, has_type_ascription) =
     pat
         .syntax()
@@ -1140,7 +1143,52 @@ fn pattern_context_for(
                         })();
                         return (PatternRefutability::Irrefutable, has_type_ascription)
                     },
-                    ast::MatchArm(_) => PatternRefutability::Refutable,
+                    ast::MatchArm(match_arm) => {
+                       let missing_variants_opt = match_arm
+                            .syntax()
+                            .parent()
+                            .and_then(ast::MatchArmList::cast)
+                            .and_then(|match_arm_list| {
+                                match_arm_list
+                                .syntax()
+                                .parent()
+                                .and_then(ast::MatchExpr::cast)
+                                .and_then(|match_expr| {
+                                    let expr_opt = find_opt_node_in_file(&original_file, match_expr.expr());
+
+                                    expr_opt.and_then(|expr| {
+                                        sema.type_of_expr(&expr)?
+                                        .adjusted()
+                                        .autoderef(sema.db)
+                                        .find_map(|ty| match ty.as_adt() {
+                                            Some(hir::Adt::Enum(e)) => Some(e),
+                                            _ => None,
+                                        }).and_then(|enum_| {
+                                            Some(enum_.variants(sema.db))
+                                        })
+                                    })
+                                }).and_then(|variants| {
+                                   Some(variants.iter().filter_map(|variant| {
+                                        let variant_name = variant.name(sema.db).to_string();
+
+                                        let variant_already_present = match_arm_list.arms().any(|arm| {
+                                            arm.pat().and_then(|pat| {
+                                                let pat_already_present = pat.syntax().to_string().contains(&variant_name);
+                                                pat_already_present.then(|| pat_already_present)
+                                            }).is_some()
+                                        });
+
+                                        (!variant_already_present).then_some(variant.clone())
+                                    }).collect::<Vec<Variant>>())
+                                })
+                        });
+
+                        if let Some(missing_variants_) = missing_variants_opt {
+                            missing_variants = missing_variants_;
+                        };
+
+                        PatternRefutability::Refutable
+                    },
                     ast::LetExpr(_) => PatternRefutability::Refutable,
                     ast::ForExpr(_) => PatternRefutability::Irrefutable,
                     _ => PatternRefutability::Irrefutable,
@@ -1162,6 +1210,7 @@ fn pattern_context_for(
         ref_token,
         record_pat: None,
         impl_: fetch_immediate_impl(sema, original_file, pat.syntax()),
+        missing_variants,
     }
 }
 
diff --git a/crates/ide-completion/src/render/pattern.rs b/crates/ide-completion/src/render/pattern.rs
index edb727b57d2..9cf766ce66c 100644
--- a/crates/ide-completion/src/render/pattern.rs
+++ b/crates/ide-completion/src/render/pattern.rs
@@ -38,7 +38,7 @@ pub(crate) fn render_struct_pat(
     let lookup = format_literal_lookup(name.as_str(), kind);
     let pat = render_pat(&ctx, pattern_ctx, &escaped_name, kind, &visible_fields, fields_omitted)?;
 
-    Some(build_completion(ctx, label, lookup, pat, strukt, false))
+    Some(build_completion(ctx, label, lookup, pat, strukt, true))
 }
 
 pub(crate) fn render_variant_pat(
diff --git a/crates/ide-completion/src/tests/record.rs b/crates/ide-completion/src/tests/record.rs
index 328faaa060f..8b8c56d1d5a 100644
--- a/crates/ide-completion/src/tests/record.rs
+++ b/crates/ide-completion/src/tests/record.rs
@@ -47,6 +47,47 @@ fn foo(s: Struct) {
 }
 
 #[test]
+fn record_pattern_field_enum() {
+    check(
+        r#"
+enum Baz { FOO, BAR }
+
+fn foo(baz: Baz) {
+    match baz {
+        Baz::FOO => (),
+        $0
+    }
+}
+"#,
+        expect![[r#"
+            en Baz
+            bn Baz::BAR Baz::BAR$0
+            kw mut
+            kw ref
+        "#]],
+    );
+
+    check(
+        r#"
+enum Baz { FOO, BAR }
+
+fn foo(baz: Baz) {
+    match baz {
+        FOO => (),
+        $0
+    }
+}
+"#,
+        expect![[r#"
+            en Baz
+            bn Baz::BAR Baz::BAR$0
+            kw mut
+            kw ref
+        "#]],
+    );
+}
+
+#[test]
 fn pattern_enum_variant() {
     check(
         r#"