Compute missing arms lazily.

1 files changed, 24 insertions, 23 deletions

diff --git a/crates/ide_assists/src/handlers/fill_match_arms.rs b/crates/ide_assists/src/handlers/fill_match_arms.rs
index 58b001050..12ea17d7a 100644
--- a/crates/ide_assists/src/handlers/fill_match_arms.rs
+++ b/crates/ide_assists/src/handlers/fill_match_arms.rs
@@ -1,4 +1,4 @@
-use std::iter;
+use std::iter::{self, Peekable};
 
 use either::Either;
 use hir::{Adt, HasSource, ModuleDef, Semantics};
@@ -63,33 +63,33 @@ pub(crate) fn fill_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<
 
     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 mut missing_pats: Peekable<Box<dyn Iterator<Item = ast::Pat>>> = if let Some(enum_def) =
+        resolve_enum_def(&ctx.sema, &expr)
+    {
         let variants = enum_def.variants(ctx.db());
 
-        let mut variants = variants
+        let missing_pats = variants
             .into_iter()
             .filter_map(|variant| build_pat(ctx.db(), module, variant))
-            .filter(|variant_pat| is_variant_missing(&top_lvl_pats, variant_pat))
-            .map(|pat| make::match_arm(iter::once(pat), make::expr_empty_block()))
-            .map(|it| it.clone_for_update())
-            .collect::<Vec<_>>();
-        if Some(enum_def)
-            == FamousDefs(&ctx.sema, Some(module.krate()))
-                .core_option_Option()
-                .map(|x| lift_enum(x))
+            .filter(|variant_pat| is_variant_missing(&top_lvl_pats, variant_pat));
+
+        let missing_pats: Box<dyn Iterator<Item = _>> = if Some(enum_def)
+            == FamousDefs(&ctx.sema, Some(module.krate())).core_option_Option().map(lift_enum)
         {
             // Match `Some` variant first.
             cov_mark::hit!(option_order);
-            variants.reverse()
-        }
-        variants
+            Box::new(missing_pats.rev())
+        } else {
+            Box::new(missing_pats)
+        };
+        missing_pats.peekable()
     } else if let Some(enum_defs) = resolve_tuple_of_enum_def(&ctx.sema, &expr) {
         // 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
+        let missing_pats = enum_defs
             .into_iter()
             .map(|enum_def| enum_def.variants(ctx.db()))
             .multi_cartesian_product()
@@ -98,15 +98,13 @@ pub(crate) fn fill_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<
                     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(&top_lvl_pats, variant_pat))
-            .map(|pat| make::match_arm(iter::once(pat), make::expr_empty_block()))
-            .map(|it| it.clone_for_update())
-            .collect()
+            .filter(|variant_pat| is_variant_missing(&top_lvl_pats, variant_pat));
+        (Box::new(missing_pats) as Box<dyn Iterator<Item = _>>).peekable()
     } else {
         return None;
     };
 
-    if missing_arms.is_empty() {
+    if missing_pats.peek().is_none() {
         return None;
     }
 
@@ -117,6 +115,9 @@ pub(crate) fn fill_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<
         target,
         |builder| {
             let new_match_arm_list = match_arm_list.clone_for_update();
+            let missing_arms = missing_pats
+                .map(|pat| make::match_arm(iter::once(pat), make::expr_empty_block()))
+                .map(|it| it.clone_for_update());
 
             let catch_all_arm = new_match_arm_list
                 .arms()
@@ -167,13 +168,13 @@ fn does_pat_match_variant(pat: &Pat, var: &Pat) -> bool {
     }
 }
 
-#[derive(Eq, PartialEq, Clone)]
+#[derive(Eq, PartialEq, Clone, Copy)]
 enum ExtendedEnum {
     Bool,
     Enum(hir::Enum),
 }
 
-#[derive(Eq, PartialEq, Clone)]
+#[derive(Eq, PartialEq, Clone, Copy)]
 enum ExtendedVariant {
     True,
     False,
@@ -185,7 +186,7 @@ fn lift_enum(e: hir::Enum) -> ExtendedEnum {
 }
 
 impl ExtendedEnum {
-    fn variants(&self, db: &RootDatabase) -> Vec<ExtendedVariant> {
+    fn variants(self, db: &RootDatabase) -> Vec<ExtendedVariant> {
         match self {
             ExtendedEnum::Enum(e) => {
                 e.variants(db).into_iter().map(|x| ExtendedVariant::Variant(x)).collect::<Vec<_>>()