6 files changed, 340 insertions, 22 deletions

diff --git a/crates/assists/src/handlers/add_lifetime_to_type.rs b/crates/assists/src/handlers/add_lifetime_to_type.rs
new file mode 100644
index 000000000..c1603e972
--- /dev/null
+++ b/crates/assists/src/handlers/add_lifetime_to_type.rs
@@ -0,0 +1,228 @@
+use ast::FieldList;
+use syntax::ast::{self, AstNode, GenericParamsOwner, NameOwner, RefType, Type};
+
+use crate::{AssistContext, AssistId, AssistKind, Assists};
+
+// Assist: add_lifetime_to_type
+//
+// Adds a new lifetime to a struct, enum or union.
+//
+// ```
+// struct Point {
+//     x: &$0u32,
+//     y: u32,
+// }
+// ```
+// ->
+// ```
+// struct Point<'a> {
+//     x: &'a u32,
+//     y: u32,
+// }
+// ```
+pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let ref_type_focused = ctx.find_node_at_offset::<ast::RefType>()?;
+    if ref_type_focused.lifetime().is_some() {
+        return None;
+    }
+
+    let node = ctx.find_node_at_offset::<ast::AdtDef>()?;
+    let has_lifetime = node
+        .generic_param_list()
+        .map(|gen_list| gen_list.lifetime_params().count() > 0)
+        .unwrap_or_default();
+
+    if has_lifetime {
+        return None;
+    }
+
+    let ref_types = fetch_borrowed_types(&node)?;
+    let target = node.syntax().text_range();
+
+    acc.add(
+        AssistId("add_lifetime_to_type", AssistKind::Generate),
+        "Add lifetime`",
+        target,
+        |builder| {
+            match node.generic_param_list() {
+                Some(gen_param) => {
+                    if let Some(left_angle) = gen_param.l_angle_token() {
+                        builder.insert(left_angle.text_range().end(), "'a, ");
+                    }
+                }
+                None => {
+                    if let Some(name) = node.name() {
+                        builder.insert(name.syntax().text_range().end(), "<'a>");
+                    }
+                }
+            }
+
+            for ref_type in ref_types {
+                if let Some(amp_token) = ref_type.amp_token() {
+                    builder.insert(amp_token.text_range().end(), "'a ");
+                }
+            }
+        },
+    )
+}
+
+fn fetch_borrowed_types(node: &ast::AdtDef) -> Option<Vec<RefType>> {
+    let ref_types: Vec<RefType> = match node {
+        ast::AdtDef::Enum(enum_) => {
+            let variant_list = enum_.variant_list()?;
+            variant_list
+                .variants()
+                .filter_map(|variant| {
+                    let field_list = variant.field_list()?;
+
+                    find_ref_types_from_field_list(&field_list)
+                })
+                .flatten()
+                .collect()
+        }
+        ast::AdtDef::Struct(strukt) => {
+            let field_list = strukt.field_list()?;
+            find_ref_types_from_field_list(&field_list)?
+        }
+        ast::AdtDef::Union(un) => {
+            let record_field_list = un.record_field_list()?;
+            record_field_list
+                .fields()
+                .filter_map(|r_field| {
+                    if let Type::RefType(ref_type) = r_field.ty()? {
+                        if ref_type.lifetime().is_none() {
+                            return Some(ref_type);
+                        }
+                    }
+
+                    None
+                })
+                .collect()
+        }
+    };
+
+    if ref_types.is_empty() {
+        None
+    } else {
+        Some(ref_types)
+    }
+}
+
+fn find_ref_types_from_field_list(field_list: &FieldList) -> Option<Vec<RefType>> {
+    let ref_types: Vec<RefType> = match field_list {
+        ast::FieldList::RecordFieldList(record_list) => record_list
+            .fields()
+            .filter_map(|f| {
+                if let Type::RefType(ref_type) = f.ty()? {
+                    if ref_type.lifetime().is_none() {
+                        return Some(ref_type);
+                    }
+                }
+
+                None
+            })
+            .collect(),
+        ast::FieldList::TupleFieldList(tuple_field_list) => tuple_field_list
+            .fields()
+            .filter_map(|f| {
+                if let Type::RefType(ref_type) = f.ty()? {
+                    if ref_type.lifetime().is_none() {
+                        return Some(ref_type);
+                    }
+                }
+
+                None
+            })
+            .collect(),
+    };
+
+    if ref_types.is_empty() {
+        None
+    } else {
+        Some(ref_types)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn add_lifetime_to_struct() {
+        check_assist(
+            add_lifetime_to_type,
+            "struct Foo { a: &$0i32 }",
+            "struct Foo<'a> { a: &'a i32 }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "struct Foo { a: &$0i32, b: &usize }",
+            "struct Foo<'a> { a: &'a i32, b: &'a usize }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "struct Foo { a: &$0i32, b: usize }",
+            "struct Foo<'a> { a: &'a i32, b: usize }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "struct Foo<T> { a: &$0T, b: usize }",
+            "struct Foo<'a, T> { a: &'a T, b: usize }",
+        );
+
+        check_assist_not_applicable(add_lifetime_to_type, "struct Foo<'a> { a: &$0'a i32 }");
+        check_assist_not_applicable(add_lifetime_to_type, "struct Foo { a: &'a$0 i32 }");
+    }
+
+    #[test]
+    fn add_lifetime_to_enum() {
+        check_assist(
+            add_lifetime_to_type,
+            "enum Foo { Bar { a: i32 }, Other, Tuple(u32, &$0u32)}",
+            "enum Foo<'a> { Bar { a: i32 }, Other, Tuple(u32, &'a u32)}",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "enum Foo { Bar { a: &$0i32 }}",
+            "enum Foo<'a> { Bar { a: &'a i32 }}",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "enum Foo<T> { Bar { a: &$0i32, b: &T }}",
+            "enum Foo<'a, T> { Bar { a: &'a i32, b: &'a T }}",
+        );
+
+        check_assist_not_applicable(add_lifetime_to_type, "enum Foo<'a> { Bar { a: &$0'a i32 }}");
+        check_assist_not_applicable(add_lifetime_to_type, "enum Foo { Bar, $0Misc }");
+    }
+
+    #[test]
+    fn add_lifetime_to_union() {
+        check_assist(
+            add_lifetime_to_type,
+            "union Foo { a: &$0i32 }",
+            "union Foo<'a> { a: &'a i32 }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "union Foo { a: &$0i32, b: &usize }",
+            "union Foo<'a> { a: &'a i32, b: &'a usize }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "union Foo<T> { a: &$0T, b: usize }",
+            "union Foo<'a, T> { a: &'a T, b: usize }",
+        );
+
+        check_assist_not_applicable(add_lifetime_to_type, "struct Foo<'a> { a: &'a $0i32 }");
+    }
+}
diff --git a/crates/assists/src/handlers/fill_match_arms.rs b/crates/assists/src/handlers/fill_match_arms.rs
index da47187e4..7663d211d 100644
--- a/crates/assists/src/handlers/fill_match_arms.rs
+++ b/crates/assists/src/handlers/fill_match_arms.rs
@@ -8,7 +8,7 @@ use syntax::ast::{self, make, AstNode, MatchArm, NameOwner, Pat};
 use test_utils::mark;
 
 use crate::{
-    utils::{render_snippet, Cursor},
+    utils::{does_pat_match_variant, render_snippet, Cursor},
     AssistContext, AssistId, AssistKind, Assists,
 };
 
@@ -147,25 +147,6 @@ fn is_variant_missing(existing_arms: &mut Vec<MatchArm>, var: &Pat) -> bool {
     })
 }
 
-fn does_pat_match_variant(pat: &Pat, var: &Pat) -> bool {
-    let first_node_text = |pat: &Pat| pat.syntax().first_child().map(|node| node.text());
-
-    let pat_head = match pat {
-        Pat::IdentPat(bind_pat) => {
-            if let Some(p) = bind_pat.pat() {
-                first_node_text(&p)
-            } else {
-                return false;
-            }
-        }
-        pat => first_node_text(pat),
-    };
-
-    let var_head = first_node_text(var);
-
-    pat_head == var_head
-}
-
 fn resolve_enum_def(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<hir::Enum> {
     sema.type_of_expr(&expr)?.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
         Some(Adt::Enum(e)) => Some(e),
diff --git a/crates/assists/src/handlers/replace_if_let_with_match.rs b/crates/assists/src/handlers/replace_if_let_with_match.rs
index aee3397ab..aee880625 100644
--- a/crates/assists/src/handlers/replace_if_let_with_match.rs
+++ b/crates/assists/src/handlers/replace_if_let_with_match.rs
@@ -10,7 +10,10 @@ use syntax::{
     AstNode,
 };
 
-use crate::{utils::unwrap_trivial_block, AssistContext, AssistId, AssistKind, Assists};
+use crate::{
+    utils::{does_pat_match_variant, unwrap_trivial_block},
+    AssistContext, AssistId, AssistKind, Assists,
+};
 
 // Assist: replace_if_let_with_match
 //
@@ -66,7 +69,13 @@ pub(crate) fn replace_if_let_with_match(acc: &mut Assists, ctx: &AssistContext)
                         .sema
                         .type_of_pat(&pat)
                         .and_then(|ty| TryEnum::from_ty(&ctx.sema, &ty))
-                        .map(|it| it.sad_pattern())
+                        .map(|it| {
+                            if does_pat_match_variant(&pat, &it.sad_pattern()) {
+                                it.happy_pattern()
+                            } else {
+                                it.sad_pattern()
+                            }
+                        })
                         .unwrap_or_else(|| make::wildcard_pat().into());
                     let else_expr = unwrap_trivial_block(else_block);
                     make::match_arm(vec![pattern], else_expr)
@@ -279,6 +288,36 @@ fn foo(x: Option<i32>) {
     }
 
     #[test]
+    fn special_case_inverted_option() {
+        check_assist(
+            replace_if_let_with_match,
+            r#"
+enum Option<T> { Some(T), None }
+use Option::*;
+
+fn foo(x: Option<i32>) {
+    $0if let None = x {
+        println!("none")
+    } else {
+        println!("some")
+    }
+}
+           "#,
+            r#"
+enum Option<T> { Some(T), None }
+use Option::*;
+
+fn foo(x: Option<i32>) {
+    match x {
+        None => println!("none"),
+        Some(_) => println!("some"),
+    }
+}
+           "#,
+        );
+    }
+
+    #[test]
     fn special_case_result() {
         check_assist(
             replace_if_let_with_match,
@@ -309,6 +348,36 @@ fn foo(x: Result<i32, ()>) {
     }
 
     #[test]
+    fn special_case_inverted_result() {
+        check_assist(
+            replace_if_let_with_match,
+            r#"
+enum Result<T, E> { Ok(T), Err(E) }
+use Result::*;
+
+fn foo(x: Result<i32, ()>) {
+    $0if let Err(x) = x {
+        println!("{}", x)
+    } else {
+        println!("ok")
+    }
+}
+           "#,
+            r#"
+enum Result<T, E> { Ok(T), Err(E) }
+use Result::*;
+
+fn foo(x: Result<i32, ()>) {
+    match x {
+        Err(x) => println!("{}", x),
+        Ok(_) => println!("ok"),
+    }
+}
+           "#,
+        );
+    }
+
+    #[test]
     fn nested_indent() {
         check_assist(
             replace_if_let_with_match,
diff --git a/crates/assists/src/lib.rs b/crates/assists/src/lib.rs
index 14178a651..559b9651e 100644
--- a/crates/assists/src/lib.rs
+++ b/crates/assists/src/lib.rs
@@ -108,6 +108,7 @@ mod handlers {
     pub(crate) type Handler = fn(&mut Assists, &AssistContext) -> Option<()>;
 
     mod add_explicit_type;
+    mod add_lifetime_to_type;
     mod add_missing_impl_members;
     mod add_turbo_fish;
     mod apply_demorgan;
@@ -164,6 +165,7 @@ mod handlers {
         &[
             // These are alphabetic for the foolish consistency
             add_explicit_type::add_explicit_type,
+            add_lifetime_to_type::add_lifetime_to_type,
             add_turbo_fish::add_turbo_fish,
             apply_demorgan::apply_demorgan,
             auto_import::auto_import,
diff --git a/crates/assists/src/tests/generated.rs b/crates/assists/src/tests/generated.rs
index d48d063b4..9aa807f10 100644
--- a/crates/assists/src/tests/generated.rs
+++ b/crates/assists/src/tests/generated.rs
@@ -104,6 +104,25 @@ impl Trait<u32> for () {
 }
 
 #[test]
+fn doctest_add_lifetime_to_type() {
+    check_doc_test(
+        "add_lifetime_to_type",
+        r#####"
+struct Point {
+    x: &$0u32,
+    y: u32,
+}
+"#####,
+        r#####"
+struct Point<'a> {
+    x: &'a u32,
+    y: u32,
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_add_turbo_fish() {
     check_doc_test(
         "add_turbo_fish",
diff --git a/crates/assists/src/utils.rs b/crates/assists/src/utils.rs
index 44c35bafa..4e762e18b 100644
--- a/crates/assists/src/utils.rs
+++ b/crates/assists/src/utils.rs
@@ -248,3 +248,22 @@ fn invert_special_case(expr: &ast::Expr) -> Option<ast::Expr> {
 pub(crate) fn next_prev() -> impl Iterator<Item = Direction> {
     [Direction::Next, Direction::Prev].iter().copied()
 }
+
+pub(crate) fn does_pat_match_variant(pat: &ast::Pat, var: &ast::Pat) -> bool {
+    let first_node_text = |pat: &ast::Pat| pat.syntax().first_child().map(|node| node.text());
+
+    let pat_head = match pat {
+        ast::Pat::IdentPat(bind_pat) => {
+            if let Some(p) = bind_pat.pat() {
+                first_node_text(&p)
+            } else {
+                return pat.syntax().text() == var.syntax().text();
+            }
+        }
+        pat => first_node_text(pat),
+    };
+
+    let var_head = first_node_text(var);
+
+    pat_head == var_head
+}