Add assist: add lifetime to type #7200

Signed-off-by: Benjamin Coenen <5719034+bnjjj@users.noreply.github.com>

3 files changed, 238 insertions, 0 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 00000000000..c7af8470465
--- /dev/null
+++ b/crates/assists/src/handlers/add_lifetime_to_type.rs
@@ -0,0 +1,217 @@
+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$0 {
+//     x: &u32,
+//     y: u32,
+// }
+// ```
+// ->
+// ```
+// struct Point<'a> {
+//     x: &'a u32,
+//     y: u32,
+// }
+// ```
+pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    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$0 { a: &i32 }",
+            "struct Foo<'a> { a: &'a i32 }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "struct Foo$0 { a: &i32, b: &usize }",
+            "struct Foo<'a> { a: &'a i32, b: &'a usize }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "struct Foo<T>$0 { a: &T, b: usize }",
+            "struct Foo<'a, T> { a: &'a T, b: usize }",
+        );
+
+        check_assist_not_applicable(add_lifetime_to_type, "struct Foo<'a>$0 { a: &'a i32 }");
+        check_assist_not_applicable(add_lifetime_to_type, "struct Foo$0 { a: &'a i32 }");
+    }
+
+    #[test]
+    fn add_lifetime_to_enum() {
+        check_assist(
+            add_lifetime_to_type,
+            "enum Foo$0 { Bar { a: i32 }, Other, Tuple(u32, &u32)}",
+            "enum Foo<'a> { Bar { a: i32 }, Other, Tuple(u32, &'a u32)}",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "enum Foo$0 { Bar { a: &i32 }}",
+            "enum Foo<'a> { Bar { a: &'a i32 }}",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "enum Foo<T>$0 { Bar { a: &i32, b: &T }}",
+            "enum Foo<'a, T> { Bar { a: &'a i32, b: &'a T }}",
+        );
+
+        check_assist_not_applicable(add_lifetime_to_type, "enum Foo<'a>$0 { Bar { a: &'a i32 }}");
+        check_assist_not_applicable(add_lifetime_to_type, "enum Foo$0 { Bar, Misc }");
+    }
+
+    #[test]
+    fn add_lifetime_to_union() {
+        check_assist(
+            add_lifetime_to_type,
+            "union Foo$0 { a: &i32 }",
+            "union Foo<'a> { a: &'a i32 }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "union Foo$0 { a: &i32, b: &usize }",
+            "union Foo<'a> { a: &'a i32, b: &'a usize }",
+        );
+
+        check_assist(
+            add_lifetime_to_type,
+            "union Foo<T>$0 { a: &T, b: usize }",
+            "union Foo<'a, T> { a: &'a T, b: usize }",
+        );
+
+        check_assist_not_applicable(add_lifetime_to_type, "struct Foo<'a>$0 { a: &'a i32 }");
+    }
+}
diff --git a/crates/assists/src/lib.rs b/crates/assists/src/lib.rs
index 14178a6510f..559b9651e80 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 d48d063b407..4b254d463ca 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$0 {
+    x: &u32,
+    y: u32,
+}
+"#####,
+        r#####"
+struct Point<'a> {
+    x: &'a u32,
+    y: u32,
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_add_turbo_fish() {
     check_doc_test(
         "add_turbo_fish",