add support for naming anon lifetimes in function return type

crates/ra_assists/src/handlers/change_lifetime_anon_to_named.rs

@@ -1,6 +1,10 @@
-use crate::{AssistContext, AssistId, Assists};
-use ra_syntax::{ast, ast::TypeParamsOwner, AstNode, SyntaxKind};
-use std::collections::HashSet;
+use crate::{assist_context::AssistBuilder, AssistContext, AssistId, Assists};
+use ast::{NameOwner, ParamList, TypeAscriptionOwner, TypeParamList, TypeRef};
+use ra_syntax::{ast, ast::TypeParamsOwner, AstNode, SyntaxKind, TextRange, TextSize};
+use rustc_hash::FxHashSet;
+
+static ASSIST_NAME: &str = "change_lifetime_anon_to_named";
+static ASSIST_LABEL: &str = "Give anonymous lifetime a name";
 
 // Assist: change_lifetime_anon_to_named
 //
@@ -26,59 +30,117 @@ use std::collections::HashSet;
 // }
 // ```
 // FIXME: How can we handle renaming any one of multiple anonymous lifetimes?
+// FIXME: should also add support for the case fun(f: &Foo) -> &<|>Foo
 pub(crate) fn change_lifetime_anon_to_named(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let lifetime_token = ctx.find_token_at_offset(SyntaxKind::LIFETIME)?;
-    let lifetime_arg = ast::LifetimeArg::cast(lifetime_token.parent())?;
-    if lifetime_arg.syntax().text() != "'_" {
-        return None;
-    }
-    let next_token = lifetime_token.next_token()?;
-    if next_token.kind() != SyntaxKind::R_ANGLE {
+    let lifetime_token = ctx
+        .find_token_at_offset(SyntaxKind::LIFETIME)
+        .filter(|lifetime| lifetime.text() == "'_")?;
+    if let Some(fn_def) = lifetime_token.ancestors().find_map(ast::FnDef::cast) {
+        generate_fn_def_assist(acc, &fn_def, lifetime_token.text_range())
+    } else if let Some(impl_def) = lifetime_token.ancestors().find_map(ast::ImplDef::cast) {
         // only allow naming the last anonymous lifetime
-        return None;
+        lifetime_token.next_token().filter(|tok| tok.kind() == SyntaxKind::R_ANGLE)?;
+        generate_impl_def_assist(acc, &impl_def, lifetime_token.text_range())
+    } else {
+        None
     }
-    let impl_def = lifetime_arg.syntax().ancestors().find_map(ast::ImplDef::cast)?;
-    // get the `impl` keyword so we know where to add the lifetime argument
-    let impl_kw = impl_def.syntax().first_child_or_token()?.into_token()?;
-    if impl_kw.kind() != SyntaxKind::IMPL_KW {
-        return None;
-    }
-    let new_lifetime_param = match impl_def.type_param_list() {
-        Some(type_params) => {
-            let used_lifetime_params: HashSet<_> = type_params
-                .lifetime_params()
-                .map(|p| {
-                    let mut param_name = p.syntax().text().to_string();
-                    param_name.remove(0);
-                    param_name
-                })
-                .collect();
-            (b'a'..=b'z')
-                .map(char::from)
-                .find(|c| !used_lifetime_params.contains(&c.to_string()))?
+}
+
+/// Generate the assist for the fn def case
+fn generate_fn_def_assist(
+    acc: &mut Assists,
+    fn_def: &ast::FnDef,
+    lifetime_loc: TextRange,
+) -> Option<()> {
+    let param_list: ParamList = fn_def.param_list()?;
+    let new_lifetime_param = generate_unique_lifetime_param_name(&fn_def.type_param_list())?;
+    let end_of_fn_ident = fn_def.name()?.ident_token()?.text_range().end();
+    let self_param =
+        // use the self if it's a reference and has no explicit lifetime
+        param_list.self_param().filter(|p| p.lifetime_token().is_none() && p.amp_token().is_some());
+    // compute the location which implicitly has the same lifetime as the anonymous lifetime
+    let loc_needing_lifetime = if let Some(self_param) = self_param {
+        // if we have a self reference, use that
+        Some(self_param.self_token()?.text_range().start())
+    } else {
+        // otherwise, if there's a single reference parameter without a named liftime, use that
+        let fn_params_without_lifetime: Vec<_> = param_list
+            .params()
+            .filter_map(|param| match param.ascribed_type() {
+                Some(TypeRef::ReferenceType(ascribed_type))
+                    if ascribed_type.lifetime_token() == None =>
+                {
+                    Some(ascribed_type.amp_token()?.text_range().end())
+                }
+                _ => None,
+            })
+            .collect();
+        match fn_params_without_lifetime.len() {
+            1 => Some(fn_params_without_lifetime.into_iter().nth(0)?),
+            0 => None,
+            // multiple unnnamed is invalid. assist is not applicable
+            _ => return None,
         }
-        None => 'a',
     };
-    acc.add(
-        AssistId("change_lifetime_anon_to_named"),
-        "Give anonymous lifetime a name",
-        lifetime_arg.syntax().text_range(),
-        |builder| {
-            match impl_def.type_param_list() {
-                Some(type_params) => {
-                    builder.insert(
-                        (u32::from(type_params.syntax().text_range().end()) - 1).into(),
-                        format!(", '{}", new_lifetime_param),
-                    );
-                }
-                None => {
-                    builder
-                        .insert(impl_kw.text_range().end(), format!("<'{}>", new_lifetime_param));
-                }
-            }
-            builder.replace(lifetime_arg.syntax().text_range(), format!("'{}", new_lifetime_param));
-        },
-    )
+    acc.add(AssistId(ASSIST_NAME), ASSIST_LABEL, lifetime_loc, |builder| {
+        add_lifetime_param(fn_def, builder, end_of_fn_ident, new_lifetime_param);
+        builder.replace(lifetime_loc, format!("'{}", new_lifetime_param));
+        loc_needing_lifetime.map(|loc| builder.insert(loc, format!("'{} ", new_lifetime_param)));
+    })
+}
+
+/// Generate the assist for the impl def case
+fn generate_impl_def_assist(
+    acc: &mut Assists,
+    impl_def: &ast::ImplDef,
+    lifetime_loc: TextRange,
+) -> Option<()> {
+    let new_lifetime_param = generate_unique_lifetime_param_name(&impl_def.type_param_list())?;
+    let end_of_impl_kw = impl_def.impl_token()?.text_range().end();
+    acc.add(AssistId(ASSIST_NAME), ASSIST_LABEL, lifetime_loc, |builder| {
+        add_lifetime_param(impl_def, builder, end_of_impl_kw, new_lifetime_param);
+        builder.replace(lifetime_loc, format!("'{}", new_lifetime_param));
+    })
+}
+
+/// Given a type parameter list, generate a unique lifetime parameter name
+/// which is not in the list
+fn generate_unique_lifetime_param_name(
+    existing_type_param_list: &Option<TypeParamList>,
+) -> Option<char> {
+    match existing_type_param_list {
+        Some(type_params) => {
+            let used_lifetime_params: FxHashSet<_> = type_params
+                .lifetime_params()
+                .map(|p| p.syntax().text().to_string()[1..].to_owned())
+                .collect();
+            (b'a'..=b'z').map(char::from).find(|c| !used_lifetime_params.contains(&c.to_string()))
+        }
+        None => Some('a'),
+    }
+}
+
+/// Add the lifetime param to `builder`. If there are type parameters in `type_params_owner`, add it to the end. Otherwise
+/// add new type params brackets with the lifetime parameter at `new_type_params_loc`.
+fn add_lifetime_param<TypeParamsOwner: ast::TypeParamsOwner>(
+    type_params_owner: &TypeParamsOwner,
+    builder: &mut AssistBuilder,
+    new_type_params_loc: TextSize,
+    new_lifetime_param: char,
+) {
+    match type_params_owner.type_param_list() {
+        // add the new lifetime parameter to an existing type param list
+        Some(type_params) => {
+            builder.insert(
+                (u32::from(type_params.syntax().text_range().end()) - 1).into(),
+                format!(", '{}", new_lifetime_param),
+            );
+        }
+        // create a new type param list containing only the new lifetime parameter
+        None => {
+            builder.insert(new_type_params_loc, format!("<'{}>", new_lifetime_param));
+        }
+    }
 }
 
 #[cfg(test)]
@@ -117,10 +179,36 @@ mod tests {
     }
 
     #[test]
-    fn test_not_applicable() {
+    fn test_example_case_cursor_after_tick() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"impl Cursor<'<|>_> {"#,
+            r#"impl<'a> Cursor<'a> {"#,
+        );
+    }
+
+    #[test]
+    fn test_example_case_cursor_before_tick() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"impl Cursor<<|>'_> {"#,
+            r#"impl<'a> Cursor<'a> {"#,
+        );
+    }
+
+    #[test]
+    fn test_not_applicable_cursor_position() {
         check_assist_not_applicable(change_lifetime_anon_to_named, r#"impl Cursor<'_><|> {"#);
         check_assist_not_applicable(change_lifetime_anon_to_named, r#"impl Cursor<|><'_> {"#);
+    }
+
+    #[test]
+    fn test_not_applicable_lifetime_already_name() {
         check_assist_not_applicable(change_lifetime_anon_to_named, r#"impl Cursor<'a<|>> {"#);
+        check_assist_not_applicable(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun<'a>() -> X<'a<|>>"#,
+        );
     }
 
     #[test]
@@ -140,4 +228,76 @@ mod tests {
             r#"impl<'a, 'b, 'c> Cursor<'a, 'b, 'c>"#,
         );
     }
+
+    #[test]
+    fn test_function_return_value_anon_lifetime_param() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun() -> X<'_<|>>"#,
+            r#"fn my_fun<'a>() -> X<'a>"#,
+        );
+    }
+
+    #[test]
+    fn test_function_return_value_anon_reference_lifetime() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun() -> &'_<|> X"#,
+            r#"fn my_fun<'a>() -> &'a X"#,
+        );
+    }
+
+    #[test]
+    fn test_function_param_anon_lifetime() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun(x: X<'_<|>>)"#,
+            r#"fn my_fun<'a>(x: X<'a>)"#,
+        );
+    }
+
+    #[test]
+    fn test_function_add_lifetime_to_params() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun(f: &Foo) -> X<'_<|>>"#,
+            r#"fn my_fun<'a>(f: &'a Foo) -> X<'a>"#,
+        );
+    }
+
+    #[test]
+    fn test_function_add_lifetime_to_params_in_presence_of_other_lifetime() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun<'other>(f: &Foo, b: &'other Bar) -> X<'_<|>>"#,
+            r#"fn my_fun<'other, 'a>(f: &'a Foo, b: &'other Bar) -> X<'a>"#,
+        );
+    }
+
+    #[test]
+    fn test_function_not_applicable_without_self_and_multiple_unnamed_param_lifetimes() {
+        // this is not permitted under lifetime elision rules
+        check_assist_not_applicable(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun(f: &Foo, b: &Bar) -> X<'_<|>>"#,
+        );
+    }
+
+    #[test]
+    fn test_function_add_lifetime_to_self_ref_param() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun<'other>(&self, f: &Foo, b: &'other Bar) -> X<'_<|>>"#,
+            r#"fn my_fun<'other, 'a>(&'a self, f: &Foo, b: &'other Bar) -> X<'a>"#,
+        );
+    }
+
+    #[test]
+    fn test_function_add_lifetime_to_param_with_non_ref_self() {
+        check_assist(
+            change_lifetime_anon_to_named,
+            r#"fn my_fun<'other>(self, f: &Foo, b: &'other Bar) -> X<'_<|>>"#,
+            r#"fn my_fun<'other, 'a>(self, f: &'a Foo, b: &'other Bar) -> X<'a>"#,
+        );
+    }
 }