1- use syntax:: ast:: { self , AstNode , HasGenericParams , HasName } ;
1+ use syntax:: {
2+ SyntaxKind , SyntaxNode , SyntaxToken ,
3+ ast:: { self , AstNode , HasGenericParams , HasName } ,
4+ } ;
25
36use crate :: { AssistContext , AssistId , Assists } ;
47
@@ -21,7 +24,7 @@ use crate::{AssistContext, AssistId, Assists};
2124// ```
2225pub ( crate ) fn add_lifetime_to_type ( acc : & mut Assists , ctx : & AssistContext < ' _ > ) -> Option < ( ) > {
2326 let ref_type_focused = ctx. find_node_at_offset :: < ast:: RefType > ( ) ?;
24- if ref_type_focused. lifetime ( ) . is_some ( ) {
27+ if ref_type_focused. lifetime ( ) . is_some_and ( |lifetime| lifetime . text ( ) != "'_" ) {
2528 return None ;
2629 }
2730
@@ -34,10 +37,10 @@ pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_>) -
3437 return None ;
3538 }
3639
37- let ref_types = fetch_borrowed_types ( & node) ?;
40+ let changes = fetch_borrowed_types ( & node) ?;
3841 let target = node. syntax ( ) . text_range ( ) ;
3942
40- acc. add ( AssistId :: generate ( "add_lifetime_to_type" ) , "Add lifetime" , target, |builder| {
43+ acc. add ( AssistId :: quick_fix ( "add_lifetime_to_type" ) , "Add lifetime" , target, |builder| {
4144 match node. generic_param_list ( ) {
4245 Some ( gen_param) => {
4346 if let Some ( left_angle) = gen_param. l_angle_token ( ) {
@@ -51,16 +54,21 @@ pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_>) -
5154 }
5255 }
5356
54- for ref_type in ref_types {
55- if let Some ( amp_token) = ref_type. amp_token ( ) {
56- builder. insert ( amp_token. text_range ( ) . end ( ) , "'a " ) ;
57+ for change in changes {
58+ match change {
59+ Change :: Replace ( it) => {
60+ builder. replace ( it. text_range ( ) , "'a" ) ;
61+ }
62+ Change :: Insert ( it) => {
63+ builder. insert ( it. text_range ( ) . end ( ) , "'a " ) ;
64+ }
5765 }
5866 }
5967 } )
6068}
6169
62- fn fetch_borrowed_types ( node : & ast:: Adt ) -> Option < Vec < ast :: RefType > > {
63- let ref_types: Vec < ast :: RefType > = match node {
70+ fn fetch_borrowed_types ( node : & ast:: Adt ) -> Option < Vec < Change > > {
71+ let ref_types: Vec < _ > = match node {
6472 ast:: Adt :: Enum ( enum_) => {
6573 let variant_list = enum_. variant_list ( ) ?;
6674 variant_list
@@ -79,55 +87,50 @@ fn fetch_borrowed_types(node: &ast::Adt) -> Option<Vec<ast::RefType>> {
7987 }
8088 ast:: Adt :: Union ( un) => {
8189 let record_field_list = un. record_field_list ( ) ?;
82- record_field_list
83- . fields ( )
84- . filter_map ( |r_field| {
85- if let ast:: Type :: RefType ( ref_type) = r_field. ty ( ) ?
86- && ref_type. lifetime ( ) . is_none ( )
87- {
88- return Some ( ref_type) ;
89- }
90-
91- None
92- } )
93- . collect ( )
90+ find_ref_types_from_field_list ( & record_field_list. into ( ) ) ?
9491 }
9592 } ;
9693
9794 if ref_types. is_empty ( ) { None } else { Some ( ref_types) }
9895}
9996
100- fn find_ref_types_from_field_list ( field_list : & ast:: FieldList ) -> Option < Vec < ast:: RefType > > {
101- let ref_types: Vec < ast:: RefType > = match field_list {
102- ast:: FieldList :: RecordFieldList ( record_list) => record_list
103- . fields ( )
104- . filter_map ( |f| {
105- if let ast:: Type :: RefType ( ref_type) = f. ty ( ) ?
106- && ref_type. lifetime ( ) . is_none ( )
107- {
108- return Some ( ref_type) ;
109- }
110-
111- None
112- } )
113- . collect ( ) ,
114- ast:: FieldList :: TupleFieldList ( tuple_field_list) => tuple_field_list
115- . fields ( )
116- . filter_map ( |f| {
117- if let ast:: Type :: RefType ( ref_type) = f. ty ( ) ?
118- && ref_type. lifetime ( ) . is_none ( )
119- {
120- return Some ( ref_type) ;
121- }
122-
123- None
124- } )
125- . collect ( ) ,
97+ fn find_ref_types_from_field_list ( field_list : & ast:: FieldList ) -> Option < Vec < Change > > {
98+ let ref_types: Vec < _ > = match field_list {
99+ ast:: FieldList :: RecordFieldList ( record_list) => {
100+ record_list. fields ( ) . flat_map ( |f| infer_lifetimes ( f. syntax ( ) ) ) . collect ( )
101+ }
102+ ast:: FieldList :: TupleFieldList ( tuple_field_list) => {
103+ tuple_field_list. fields ( ) . flat_map ( |f| infer_lifetimes ( f. syntax ( ) ) ) . collect ( )
104+ }
126105 } ;
127106
128107 if ref_types. is_empty ( ) { None } else { Some ( ref_types) }
129108}
130109
110+ enum Change {
111+ Replace ( SyntaxToken ) ,
112+ Insert ( SyntaxToken ) ,
113+ }
114+
115+ fn infer_lifetimes ( node : & SyntaxNode ) -> Vec < Change > {
116+ node. children ( )
117+ . filter ( |it| !matches ! ( it. kind( ) , SyntaxKind :: FN_PTR_TYPE | SyntaxKind :: TYPE_BOUND_LIST ) )
118+ . flat_map ( |it| {
119+ infer_lifetimes ( & it)
120+ . into_iter ( )
121+ . chain ( ast:: Lifetime :: cast ( it. clone ( ) ) . and_then ( |lt| {
122+ lt. lifetime_ident_token ( ) . filter ( |lt| lt. text ( ) == "'_" ) . map ( Change :: Replace )
123+ } ) )
124+ . chain (
125+ ast:: RefType :: cast ( it)
126+ . filter ( |ty| ty. lifetime ( ) . is_none ( ) )
127+ . and_then ( |ty| ty. amp_token ( ) )
128+ . map ( Change :: Insert ) ,
129+ )
130+ } )
131+ . collect ( )
132+ }
133+
131134#[ cfg( test) ]
132135mod tests {
133136 use crate :: tests:: { check_assist, check_assist_not_applicable} ;
@@ -164,6 +167,24 @@ mod tests {
164167 check_assist_not_applicable ( add_lifetime_to_type, r#"struct Foo { a: &'a$0 i32 }"# ) ;
165168 }
166169
170+ #[ test]
171+ fn add_lifetime_to_nested_types ( ) {
172+ check_assist (
173+ add_lifetime_to_type,
174+ r#"struct Foo { a: &$0i32, b: &(&i32, fn(&str) -> &str) }"# ,
175+ r#"struct Foo<'a> { a: &'a i32, b: &'a (&'a i32, fn(&str) -> &str) }"# ,
176+ ) ;
177+ }
178+
179+ #[ test]
180+ fn add_lifetime_to_explicit_infer_lifetime ( ) {
181+ check_assist (
182+ add_lifetime_to_type,
183+ r#"struct Foo { a: &'_ $0i32, b: &'_ (&'_ i32, fn(&str) -> &str) }"# ,
184+ r#"struct Foo<'a> { a: &'a i32, b: &'a (&'a i32, fn(&str) -> &str) }"# ,
185+ ) ;
186+ }
187+
167188 #[ test]
168189 fn add_lifetime_to_enum ( ) {
169190 check_assist (
0 commit comments