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::()?; + 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> { + let ref_types: Vec = 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> { + let ref_types: Vec = 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$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$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$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 @@ pub(crate) fn all() -> &'static [Handler] { &[ // 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 @@ -103,6 +103,25 @@ fn foo(&self) -> u32 { ) } +#[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(