Merge #3689

3689: implement fill match arm assist for tuple of enums r=matklad a=JoshMcguigan

This updates the fill match arm assist to work in cases where the user is matching on a tuple of enums. 

Note, for now this does not apply when some match arms exist (other than the trivial `_`), but I think this could be added in the future.

I think this also lays the groundwork for filling match arms when matching on tuples of non-enum values, for example a tuple of an enum and a boolean.

Co-authored-by: Josh Mcguigan <[email protected]>

4 files changed, 257 insertions, 14 deletions

diff --git a/Cargo.lock b/Cargo.lock
index f8d26192d..eb5247b69 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -876,6 +876,7 @@ name = "ra_assists"
 version = "0.1.0"
 dependencies = [
  "format-buf",
+ "itertools",
  "join_to_string",
  "ra_db",
  "ra_fmt",
diff --git a/crates/ra_assists/Cargo.toml b/crates/ra_assists/Cargo.toml
index d314dc8e6..85adddb5b 100644
--- a/crates/ra_assists/Cargo.toml
+++ b/crates/ra_assists/Cargo.toml
@@ -11,6 +11,7 @@ doctest = false
 format-buf = "1.0.0"
 join_to_string = "0.1.3"
 rustc-hash = "1.1.0"
+itertools = "0.8.2"
 
 ra_syntax = { path = "../ra_syntax" }
 ra_text_edit = { path = "../ra_text_edit" }
diff --git a/crates/ra_assists/src/handlers/fill_match_arms.rs b/crates/ra_assists/src/handlers/fill_match_arms.rs
index fbd6a3ec3..7463b2af7 100644
--- a/crates/ra_assists/src/handlers/fill_match_arms.rs
+++ b/crates/ra_assists/src/handlers/fill_match_arms.rs
@@ -3,6 +3,7 @@
 use std::iter;
 
 use hir::{Adt, HasSource, Semantics};
+use itertools::Itertools;
 use ra_ide_db::RootDatabase;
 
 use crate::{Assist, AssistCtx, AssistId};
@@ -39,13 +40,6 @@ pub(crate) fn fill_match_arms(ctx: AssistCtx) -> Option<Assist> {
     let match_arm_list = match_expr.match_arm_list()?;
 
     let expr = match_expr.expr()?;
-    let enum_def = resolve_enum_def(&ctx.sema, &expr)?;
-    let module = ctx.sema.scope(expr.syntax()).module()?;
-
-    let variants = enum_def.variants(ctx.db);
-    if variants.is_empty() {
-        return None;
-    }
 
     let mut arms: Vec<MatchArm> = match_arm_list.arms().collect();
     if arms.len() == 1 {
@@ -54,13 +48,49 @@ pub(crate) fn fill_match_arms(ctx: AssistCtx) -> Option<Assist> {
         }
     }
 
-    let db = ctx.db;
-    let missing_arms: Vec<MatchArm> = variants
-        .into_iter()
-        .filter_map(|variant| build_pat(db, module, variant))
-        .filter(|variant_pat| is_variant_missing(&mut arms, variant_pat))
-        .map(|pat| make::match_arm(iter::once(pat), make::expr_unit()))
-        .collect();
+    let module = ctx.sema.scope(expr.syntax()).module()?;
+
+    let missing_arms: Vec<MatchArm> = if let Some(enum_def) = resolve_enum_def(&ctx.sema, &expr) {
+        let variants = enum_def.variants(ctx.db);
+
+        variants
+            .into_iter()
+            .filter_map(|variant| build_pat(ctx.db, module, variant))
+            .filter(|variant_pat| is_variant_missing(&mut arms, variant_pat))
+            .map(|pat| make::match_arm(iter::once(pat), make::expr_unit()))
+            .collect()
+    } else if let Some(enum_defs) = resolve_tuple_of_enum_def(&ctx.sema, &expr) {
+        // Partial fill not currently supported for tuple of enums.
+        if !arms.is_empty() {
+            return None;
+        }
+
+        // We do not currently support filling match arms for a tuple
+        // containing a single enum.
+        if enum_defs.len() < 2 {
+            return None;
+        }
+
+        // When calculating the match arms for a tuple of enums, we want
+        // to create a match arm for each possible combination of enum
+        // values. The `multi_cartesian_product` method transforms
+        // Vec<Vec<EnumVariant>> into Vec<(EnumVariant, .., EnumVariant)>
+        // where each tuple represents a proposed match arm.
+        enum_defs
+            .into_iter()
+            .map(|enum_def| enum_def.variants(ctx.db))
+            .multi_cartesian_product()
+            .map(|variants| {
+                let patterns =
+                    variants.into_iter().filter_map(|variant| build_pat(ctx.db, module, variant));
+                ast::Pat::from(make::tuple_pat(patterns))
+            })
+            .filter(|variant_pat| is_variant_missing(&mut arms, variant_pat))
+            .map(|pat| make::match_arm(iter::once(pat), make::expr_unit()))
+            .collect()
+    } else {
+        return None;
+    };
 
     if missing_arms.is_empty() {
         return None;
@@ -104,6 +134,25 @@ fn resolve_enum_def(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<
     })
 }
 
+fn resolve_tuple_of_enum_def(
+    sema: &Semantics<RootDatabase>,
+    expr: &ast::Expr,
+) -> Option<Vec<hir::Enum>> {
+    sema.type_of_expr(&expr)?
+        .tuple_fields(sema.db)
+        .iter()
+        .map(|ty| {
+            ty.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
+                Some(Adt::Enum(e)) => Some(e),
+                // For now we only handle expansion for a tuple of enums. Here
+                // we map non-enum items to None and rely on `collect` to
+                // convert Vec<Option<hir::Enum>> into Option<Vec<hir::Enum>>.
+                _ => None,
+            })
+        })
+        .collect()
+}
+
 fn build_pat(db: &RootDatabase, module: hir::Module, var: hir::EnumVariant) -> Option<ast::Pat> {
     let path = crate::ast_transform::path_to_ast(module.find_use_path(db, var.into())?);
 
@@ -152,6 +201,21 @@ mod tests {
     }
 
     #[test]
+    fn tuple_of_non_enum() {
+        // for now this case is not handled, although it potentially could be
+        // in the future
+        check_assist_not_applicable(
+            fill_match_arms,
+            r#"
+            fn main() {
+                match (0, false)<|> {
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
     fn partial_fill_record_tuple() {
         check_assist(
             fill_match_arms,
@@ -308,6 +372,169 @@ mod tests {
     }
 
     #[test]
+    fn fill_match_arms_tuple_of_enum() {
+        check_assist(
+            fill_match_arms,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            enum B {
+                One,
+                Two,
+            }
+
+            fn main() {
+                let a = A::One;
+                let b = B::One;
+                match (a<|>, b) {}
+            }
+            "#,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            enum B {
+                One,
+                Two,
+            }
+
+            fn main() {
+                let a = A::One;
+                let b = B::One;
+                match <|>(a, b) {
+                    (A::One, B::One) => (),
+                    (A::One, B::Two) => (),
+                    (A::Two, B::One) => (),
+                    (A::Two, B::Two) => (),
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn fill_match_arms_tuple_of_enum_ref() {
+        check_assist(
+            fill_match_arms,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            enum B {
+                One,
+                Two,
+            }
+
+            fn main() {
+                let a = A::One;
+                let b = B::One;
+                match (&a<|>, &b) {}
+            }
+            "#,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            enum B {
+                One,
+                Two,
+            }
+
+            fn main() {
+                let a = A::One;
+                let b = B::One;
+                match <|>(&a, &b) {
+                    (A::One, B::One) => (),
+                    (A::One, B::Two) => (),
+                    (A::Two, B::One) => (),
+                    (A::Two, B::Two) => (),
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn fill_match_arms_tuple_of_enum_partial() {
+        check_assist_not_applicable(
+            fill_match_arms,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            enum B {
+                One,
+                Two,
+            }
+
+            fn main() {
+                let a = A::One;
+                let b = B::One;
+                match (a<|>, b) {
+                    (A::Two, B::One) => (),
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn fill_match_arms_tuple_of_enum_not_applicable() {
+        check_assist_not_applicable(
+            fill_match_arms,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+            enum B {
+                One,
+                Two,
+            }
+
+            fn main() {
+                let a = A::One;
+                let b = B::One;
+                match (a<|>, b) {
+                    (A::Two, B::One) => (),
+                    (A::One, B::One) => (),
+                    (A::One, B::Two) => (),
+                    (A::Two, B::Two) => (),
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
+    fn fill_match_arms_single_element_tuple_of_enum() {
+        // For now we don't hande the case of a single element tuple, but
+        // we could handle this in the future if `make::tuple_pat` allowed
+        // creating a tuple with a single pattern.
+        check_assist_not_applicable(
+            fill_match_arms,
+            r#"
+            enum A {
+                One,
+                Two,
+            }
+
+            fn main() {
+                let a = A::One;
+                match (a<|>, ) {
+                }
+            }
+            "#,
+        );
+    }
+
+    #[test]
     fn test_fill_match_arm_refs() {
         check_assist(
             fill_match_arms,
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index 9f6f1cc53..9257ccd1a 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -136,6 +136,20 @@ pub fn placeholder_pat() -> ast::PlaceholderPat {
     }
 }
 
+/// Creates a tuple of patterns from an interator of patterns.
+///
+/// Invariant: `pats` must be length > 1
+///
+/// FIXME handle `pats` length == 1
+pub fn tuple_pat(pats: impl IntoIterator<Item = ast::Pat>) -> ast::TuplePat {
+    let pats_str = pats.into_iter().map(|p| p.to_string()).join(", ");
+    return from_text(&format!("({})", pats_str));
+
+    fn from_text(text: &str) -> ast::TuplePat {
+        ast_from_text(&format!("fn f({}: ())", text))
+    }
+}
+
 pub fn tuple_struct_pat(
     path: ast::Path,
     pats: impl IntoIterator<Item = ast::Pat>,