15 files changed, 490 insertions, 425 deletions

diff --git a/crates/assists/src/handlers/add_turbo_fish.rs b/crates/assists/src/handlers/add_turbo_fish.rs
index e3d84d698..1f486c013 100644
--- a/crates/assists/src/handlers/add_turbo_fish.rs
+++ b/crates/assists/src/handlers/add_turbo_fish.rs
@@ -25,7 +25,7 @@ use crate::{
 // }
 // ```
 pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let ident = ctx.find_token_at_offset(SyntaxKind::IDENT).or_else(|| {
+    let ident = ctx.find_token_syntax_at_offset(SyntaxKind::IDENT).or_else(|| {
         let arg_list = ctx.find_node_at_offset::<ast::ArgList>()?;
         if arg_list.args().count() > 0 {
             return None;
diff --git a/crates/assists/src/handlers/change_return_type_to_result.rs b/crates/assists/src/handlers/change_return_type_to_result.rs
index be480943c..76f33a5b6 100644
--- a/crates/assists/src/handlers/change_return_type_to_result.rs
+++ b/crates/assists/src/handlers/change_return_type_to_result.rs
@@ -2,7 +2,7 @@ use std::iter;
 
 use syntax::{
     ast::{self, make, BlockExpr, Expr, LoopBodyOwner},
-    AstNode, SyntaxNode,
+    match_ast, AstNode, SyntaxNode,
 };
 use test_utils::mark;
 
@@ -21,8 +21,18 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // ```
 pub(crate) fn change_return_type_to_result(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let ret_type = ctx.find_node_at_offset::<ast::RetType>()?;
-    // FIXME: extend to lambdas as well
-    let fn_def = ret_type.syntax().parent().and_then(ast::Fn::cast)?;
+    let parent = ret_type.syntax().parent()?;
+    let block_expr = match_ast! {
+        match parent {
+            ast::Fn(func) => func.body()?,
+            ast::ClosureExpr(closure) => match closure.body()? {
+                Expr::BlockExpr(block) => block,
+                // closures require a block when a return type is specified
+                _ => return None,
+            },
+            _ => return None,
+        }
+    };
 
     let type_ref = &ret_type.ty()?;
     let ret_type_str = type_ref.syntax().text().to_string();
@@ -34,16 +44,14 @@ pub(crate) fn change_return_type_to_result(acc: &mut Assists, ctx: &AssistContex
         }
     }
 
-    let block_expr = &fn_def.body()?;
-
     acc.add(
         AssistId("change_return_type_to_result", AssistKind::RefactorRewrite),
         "Wrap return type in Result",
         type_ref.syntax().text_range(),
         |builder| {
             let mut tail_return_expr_collector = TailReturnCollector::new();
-            tail_return_expr_collector.collect_jump_exprs(block_expr, false);
-            tail_return_expr_collector.collect_tail_exprs(block_expr);
+            tail_return_expr_collector.collect_jump_exprs(&block_expr, false);
+            tail_return_expr_collector.collect_tail_exprs(&block_expr);
 
             for ret_expr_arg in tail_return_expr_collector.exprs_to_wrap {
                 let ok_wrapped = make::expr_call(
@@ -285,16 +293,20 @@ mod tests {
     }
 
     #[test]
-    fn change_return_type_to_result_simple_return_type() {
+    fn change_return_type_to_result_simple_closure() {
         check_assist(
             change_return_type_to_result,
-            r#"fn foo() -> i32<|> {
-                let test = "test";
-                return 42i32;
+            r#"fn foo() {
+                || -> i32<|> {
+                    let test = "test";
+                    return 42i32;
+                };
             }"#,
-            r#"fn foo() -> Result<i32, ${0:_}> {
-                let test = "test";
-                return Ok(42i32);
+            r#"fn foo() {
+                || -> Result<i32, ${0:_}> {
+                    let test = "test";
+                    return Ok(42i32);
+                };
             }"#,
         );
     }
@@ -311,6 +323,29 @@ mod tests {
     }
 
     #[test]
+    fn change_return_type_to_result_simple_return_type_bad_cursor_closure() {
+        check_assist_not_applicable(
+            change_return_type_to_result,
+            r#"fn foo() {
+                || -> i32 {
+                    let test = "test";<|>
+                    return 42i32;
+                };
+            }"#,
+        );
+    }
+
+    #[test]
+    fn change_return_type_to_result_closure_non_block() {
+        check_assist_not_applicable(
+            change_return_type_to_result,
+            r#"fn foo() {
+                || -> i<|>32 3;
+            }"#,
+        );
+    }
+
+    #[test]
     fn change_return_type_to_result_simple_return_type_already_result_std() {
         check_assist_not_applicable(
             change_return_type_to_result,
@@ -334,6 +369,19 @@ mod tests {
     }
 
     #[test]
+    fn change_return_type_to_result_simple_return_type_already_result_closure() {
+        check_assist_not_applicable(
+            change_return_type_to_result,
+            r#"fn foo() {
+                || -> Result<i32<|>, String> {
+                    let test = "test";
+                    return 42i32;
+                };
+            }"#,
+        );
+    }
+
+    #[test]
     fn change_return_type_to_result_simple_with_cursor() {
         check_assist(
             change_return_type_to_result,
@@ -364,6 +412,25 @@ mod tests {
     }
 
     #[test]
+    fn change_return_type_to_result_simple_with_tail_closure() {
+        check_assist(
+            change_return_type_to_result,
+            r#"fn foo() {
+                || -><|> i32 {
+                    let test = "test";
+                    42i32
+                };
+            }"#,
+            r#"fn foo() {
+                || -> Result<i32, ${0:_}> {
+                    let test = "test";
+                    Ok(42i32)
+                };
+            }"#,
+        );
+    }
+
+    #[test]
     fn change_return_type_to_result_simple_with_tail_only() {
         check_assist(
             change_return_type_to_result,
@@ -375,6 +442,7 @@ mod tests {
             }"#,
         );
     }
+
     #[test]
     fn change_return_type_to_result_simple_with_tail_block_like() {
         check_assist(
@@ -397,6 +465,31 @@ mod tests {
     }
 
     #[test]
+    fn change_return_type_to_result_simple_without_block_closure() {
+        check_assist(
+            change_return_type_to_result,
+            r#"fn foo() {
+                || -> i32<|> {
+                    if true {
+                        42i32
+                    } else {
+                        24i32
+                    }
+                };
+            }"#,
+            r#"fn foo() {
+                || -> Result<i32, ${0:_}> {
+                    if true {
+                        Ok(42i32)
+                    } else {
+                        Ok(24i32)
+                    }
+                };
+            }"#,
+        );
+    }
+
+    #[test]
     fn change_return_type_to_result_simple_with_nested_if() {
         check_assist(
             change_return_type_to_result,
diff --git a/crates/assists/src/handlers/convert_integer_literal.rs b/crates/assists/src/handlers/convert_integer_literal.rs
index c8af80701..667115382 100644
--- a/crates/assists/src/handlers/convert_integer_literal.rs
+++ b/crates/assists/src/handlers/convert_integer_literal.rs
@@ -1,4 +1,4 @@
-use syntax::{ast, ast::Radix, AstNode};
+use syntax::{ast, ast::Radix, AstToken};
 
 use crate::{AssistContext, AssistId, AssistKind, Assists, GroupLabel};
 
@@ -15,14 +15,16 @@ use crate::{AssistContext, AssistId, AssistKind, Assists, GroupLabel};
 // ```
 pub(crate) fn convert_integer_literal(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let literal = ctx.find_node_at_offset::<ast::Literal>()?;
-    let (radix, value) = literal.int_value()?;
+    let literal = match literal.kind() {
+        ast::LiteralKind::IntNumber(it) => it,
+        _ => return None,
+    };
+    let radix = literal.radix();
+    let value = literal.value()?;
+    let suffix = literal.suffix();
 
     let range = literal.syntax().text_range();
     let group_id = GroupLabel("Convert integer base".into());
-    let suffix = match literal.kind() {
-        ast::LiteralKind::IntNumber { suffix } => suffix,
-        _ => return None,
-    };
 
     for &target_radix in Radix::ALL {
         if target_radix == radix {
@@ -36,16 +38,11 @@ pub(crate) fn convert_integer_literal(acc: &mut Assists, ctx: &AssistContext) ->
             Radix::Hexadecimal => format!("0x{:X}", value),
         };
 
-        let label = format!(
-            "Convert {} to {}{}",
-            literal,
-            converted,
-            suffix.as_deref().unwrap_or_default()
-        );
+        let label = format!("Convert {} to {}{}", literal, converted, suffix.unwrap_or_default());
 
         // Appends the type suffix back into the new literal if it exists.
-        if let Some(suffix) = &suffix {
-            converted.push_str(&suffix);
+        if let Some(suffix) = suffix {
+            converted.push_str(suffix);
         }
 
         acc.add_group(
@@ -132,34 +129,6 @@ mod tests {
         );
     }
 
-    // Decimal numbers under 3 digits have a special case where they return early because we can't fit a
-    // other base's prefix, so we have a separate test for that.
-    #[test]
-    fn convert_small_decimal_integer() {
-        let before = "const _: i32 = 10<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b1010;",
-            "Convert 10 to 0b1010",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o12;",
-            "Convert 10 to 0o12",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xA;",
-            "Convert 10 to 0xA",
-        );
-    }
-
     #[test]
     fn convert_hexadecimal_integer() {
         let before = "const _: i32 = 0xFF<|>;";
@@ -239,7 +208,7 @@ mod tests {
     }
 
     #[test]
-    fn convert_decimal_integer_with_underscores() {
+    fn convert_integer_with_underscores() {
         let before = "const _: i32 = 1_00_0<|>;";
 
         check_assist_by_label(
@@ -265,111 +234,7 @@ mod tests {
     }
 
     #[test]
-    fn convert_small_decimal_integer_with_underscores() {
-        let before = "const _: i32 = 1_0<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b1010;",
-            "Convert 1_0 to 0b1010",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o12;",
-            "Convert 1_0 to 0o12",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xA;",
-            "Convert 1_0 to 0xA",
-        );
-    }
-
-    #[test]
-    fn convert_hexadecimal_integer_with_underscores() {
-        let before = "const _: i32 = 0x_F_F<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b11111111;",
-            "Convert 0x_F_F to 0b11111111",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o377;",
-            "Convert 0x_F_F to 0o377",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255;",
-            "Convert 0x_F_F to 255",
-        );
-    }
-
-    #[test]
-    fn convert_binary_integer_with_underscores() {
-        let before = "const _: i32 = 0b1111_1111<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o377;",
-            "Convert 0b1111_1111 to 0o377",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255;",
-            "Convert 0b1111_1111 to 255",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xFF;",
-            "Convert 0b1111_1111 to 0xFF",
-        );
-    }
-
-    #[test]
-    fn convert_octal_integer_with_underscores() {
-        let before = "const _: i32 = 0o3_77<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b11111111;",
-            "Convert 0o3_77 to 0b11111111",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255;",
-            "Convert 0o3_77 to 255",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xFF;",
-            "Convert 0o3_77 to 0xFF",
-        );
-    }
-
-    #[test]
-    fn convert_decimal_integer_with_suffix() {
+    fn convert_integer_with_suffix() {
         let before = "const _: i32 = 1000i32<|>;";
 
         check_assist_by_label(
@@ -395,240 +260,6 @@ mod tests {
     }
 
     #[test]
-    fn convert_small_decimal_integer_with_suffix() {
-        let before = "const _: i32 = 10i32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b1010i32;",
-            "Convert 10i32 to 0b1010i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o12i32;",
-            "Convert 10i32 to 0o12i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xAi32;",
-            "Convert 10i32 to 0xAi32",
-        );
-    }
-
-    #[test]
-    fn convert_hexadecimal_integer_with_suffix() {
-        let before = "const _: i32 = 0xFFi32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b11111111i32;",
-            "Convert 0xFFi32 to 0b11111111i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o377i32;",
-            "Convert 0xFFi32 to 0o377i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255i32;",
-            "Convert 0xFFi32 to 255i32",
-        );
-    }
-
-    #[test]
-    fn convert_binary_integer_with_suffix() {
-        let before = "const _: i32 = 0b11111111i32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o377i32;",
-            "Convert 0b11111111i32 to 0o377i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255i32;",
-            "Convert 0b11111111i32 to 255i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xFFi32;",
-            "Convert 0b11111111i32 to 0xFFi32",
-        );
-    }
-
-    #[test]
-    fn convert_octal_integer_with_suffix() {
-        let before = "const _: i32 = 0o377i32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b11111111i32;",
-            "Convert 0o377i32 to 0b11111111i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255i32;",
-            "Convert 0o377i32 to 255i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xFFi32;",
-            "Convert 0o377i32 to 0xFFi32",
-        );
-    }
-
-    #[test]
-    fn convert_decimal_integer_with_underscores_and_suffix() {
-        let before = "const _: i32 = 1_00_0i32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b1111101000i32;",
-            "Convert 1_00_0i32 to 0b1111101000i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o1750i32;",
-            "Convert 1_00_0i32 to 0o1750i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0x3E8i32;",
-            "Convert 1_00_0i32 to 0x3E8i32",
-        );
-    }
-
-    #[test]
-    fn convert_small_decimal_integer_with_underscores_and_suffix() {
-        let before = "const _: i32 = 1_0i32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b1010i32;",
-            "Convert 1_0i32 to 0b1010i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o12i32;",
-            "Convert 1_0i32 to 0o12i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xAi32;",
-            "Convert 1_0i32 to 0xAi32",
-        );
-    }
-
-    #[test]
-    fn convert_hexadecimal_integer_with_underscores_and_suffix() {
-        let before = "const _: i32 = 0x_F_Fi32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b11111111i32;",
-            "Convert 0x_F_Fi32 to 0b11111111i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o377i32;",
-            "Convert 0x_F_Fi32 to 0o377i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255i32;",
-            "Convert 0x_F_Fi32 to 255i32",
-        );
-    }
-
-    #[test]
-    fn convert_binary_integer_with_underscores_and_suffix() {
-        let before = "const _: i32 = 0b1111_1111i32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0o377i32;",
-            "Convert 0b1111_1111i32 to 0o377i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255i32;",
-            "Convert 0b1111_1111i32 to 255i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xFFi32;",
-            "Convert 0b1111_1111i32 to 0xFFi32",
-        );
-    }
-
-    #[test]
-    fn convert_octal_integer_with_underscores_and_suffix() {
-        let before = "const _: i32 = 0o3_77i32<|>;";
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0b11111111i32;",
-            "Convert 0o3_77i32 to 0b11111111i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 255i32;",
-            "Convert 0o3_77i32 to 255i32",
-        );
-
-        check_assist_by_label(
-            convert_integer_literal,
-            before,
-            "const _: i32 = 0xFFi32;",
-            "Convert 0o3_77i32 to 0xFFi32",
-        );
-    }
-
-    #[test]
     fn convert_overflowing_literal() {
         let before = "const _: i32 =
             111111111111111111111111111111111111111111111111111111111111111111111111<|>;";
diff --git a/crates/assists/src/handlers/expand_glob_import.rs b/crates/assists/src/handlers/expand_glob_import.rs
index 316a58d88..853266395 100644
--- a/crates/assists/src/handlers/expand_glob_import.rs
+++ b/crates/assists/src/handlers/expand_glob_import.rs
@@ -41,7 +41,7 @@ use crate::{
 // fn qux(bar: Bar, baz: Baz) {}
 // ```
 pub(crate) fn expand_glob_import(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let star = ctx.find_token_at_offset(T![*])?;
+    let star = ctx.find_token_syntax_at_offset(T![*])?;
     let (parent, mod_path) = find_parent_and_path(&star)?;
     let target_module = match ctx.sema.resolve_path(&mod_path)? {
         PathResolution::Def(ModuleDef::Module(it)) => it,
diff --git a/crates/assists/src/handlers/flip_comma.rs b/crates/assists/src/handlers/flip_comma.rs
index 5c69db53e..64b4b1a76 100644
--- a/crates/assists/src/handlers/flip_comma.rs
+++ b/crates/assists/src/handlers/flip_comma.rs
@@ -18,7 +18,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // }
 // ```
 pub(crate) fn flip_comma(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let comma = ctx.find_token_at_offset(T![,])?;
+    let comma = ctx.find_token_syntax_at_offset(T![,])?;
     let prev = non_trivia_sibling(comma.clone().into(), Direction::Prev)?;
     let next = non_trivia_sibling(comma.clone().into(), Direction::Next)?;
 
diff --git a/crates/assists/src/handlers/flip_trait_bound.rs b/crates/assists/src/handlers/flip_trait_bound.rs
index 347e79b1d..92ee42181 100644
--- a/crates/assists/src/handlers/flip_trait_bound.rs
+++ b/crates/assists/src/handlers/flip_trait_bound.rs
@@ -20,7 +20,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 pub(crate) fn flip_trait_bound(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     // We want to replicate the behavior of `flip_binexpr` by only suggesting
     // the assist when the cursor is on a `+`
-    let plus = ctx.find_token_at_offset(T![+])?;
+    let plus = ctx.find_token_syntax_at_offset(T![+])?;
 
     // Make sure we're in a `TypeBoundList`
     if ast::TypeBoundList::cast(plus.parent()).is_none() {
diff --git a/crates/assists/src/handlers/infer_function_return_type.rs b/crates/assists/src/handlers/infer_function_return_type.rs
new file mode 100644
index 000000000..520d07ae0
--- /dev/null
+++ b/crates/assists/src/handlers/infer_function_return_type.rs
@@ -0,0 +1,337 @@
+use hir::HirDisplay;
+use syntax::{ast, AstNode, TextRange, TextSize};
+use test_utils::mark;
+
+use crate::{AssistContext, AssistId, AssistKind, Assists};
+
+// Assist: infer_function_return_type
+//
+// Adds the return type to a function or closure inferred from its tail expression if it doesn't have a return
+// type specified. This assists is useable in a functions or closures tail expression or return type position.
+//
+// ```
+// fn foo() { 4<|>2i32 }
+// ```
+// ->
+// ```
+// fn foo() -> i32 { 42i32 }
+// ```
+pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let (tail_expr, builder_edit_pos, wrap_expr) = extract_tail(ctx)?;
+    let module = ctx.sema.scope(tail_expr.syntax()).module()?;
+    let ty = ctx.sema.type_of_expr(&tail_expr)?;
+    if ty.is_unit() {
+        return None;
+    }
+    let ty = ty.display_source_code(ctx.db(), module.into()).ok()?;
+
+    acc.add(
+        AssistId("infer_function_return_type", AssistKind::RefactorRewrite),
+        "Add this function's return type",
+        tail_expr.syntax().text_range(),
+        |builder| {
+            match builder_edit_pos {
+                InsertOrReplace::Insert(insert_pos) => {
+                    builder.insert(insert_pos, &format!("-> {} ", ty))
+                }
+                InsertOrReplace::Replace(text_range) => {
+                    builder.replace(text_range, &format!("-> {}", ty))
+                }
+            }
+            if wrap_expr {
+                mark::hit!(wrap_closure_non_block_expr);
+                // `|x| x` becomes `|x| -> T x` which is invalid, so wrap it in a block
+                builder.replace(tail_expr.syntax().text_range(), &format!("{{{}}}", tail_expr));
+            }
+        },
+    )
+}
+
+enum InsertOrReplace {
+    Insert(TextSize),
+    Replace(TextRange),
+}
+
+/// Check the potentially already specified return type and reject it or turn it into a builder command
+/// if allowed.
+fn ret_ty_to_action(ret_ty: Option<ast::RetType>, insert_pos: TextSize) -> Option<InsertOrReplace> {
+    match ret_ty {
+        Some(ret_ty) => match ret_ty.ty() {
+            Some(ast::Type::InferType(_)) | None => {
+                mark::hit!(existing_infer_ret_type);
+                mark::hit!(existing_infer_ret_type_closure);
+                Some(InsertOrReplace::Replace(ret_ty.syntax().text_range()))
+            }
+            _ => {
+                mark::hit!(existing_ret_type);
+                mark::hit!(existing_ret_type_closure);
+                None
+            }
+        },
+        None => Some(InsertOrReplace::Insert(insert_pos + TextSize::from(1))),
+    }
+}
+
+fn extract_tail(ctx: &AssistContext) -> Option<(ast::Expr, InsertOrReplace, bool)> {
+    let (tail_expr, return_type_range, action, wrap_expr) =
+        if let Some(closure) = ctx.find_node_at_offset::<ast::ClosureExpr>() {
+            let rpipe_pos = closure.param_list()?.syntax().last_token()?.text_range().end();
+            let action = ret_ty_to_action(closure.ret_type(), rpipe_pos)?;
+
+            let body = closure.body()?;
+            let body_start = body.syntax().first_token()?.text_range().start();
+            let (tail_expr, wrap_expr) = match body {
+                ast::Expr::BlockExpr(block) => (block.expr()?, false),
+                body => (body, true),
+            };
+
+            let ret_range = TextRange::new(rpipe_pos, body_start);
+            (tail_expr, ret_range, action, wrap_expr)
+        } else {
+            let func = ctx.find_node_at_offset::<ast::Fn>()?;
+            let rparen_pos = func.param_list()?.r_paren_token()?.text_range().end();
+            let action = ret_ty_to_action(func.ret_type(), rparen_pos)?;
+
+            let body = func.body()?;
+            let tail_expr = body.expr()?;
+
+            let ret_range_end = body.l_curly_token()?.text_range().start();
+            let ret_range = TextRange::new(rparen_pos, ret_range_end);
+            (tail_expr, ret_range, action, false)
+        };
+    let frange = ctx.frange.range;
+    if return_type_range.contains_range(frange) {
+        mark::hit!(cursor_in_ret_position);
+        mark::hit!(cursor_in_ret_position_closure);
+    } else if tail_expr.syntax().text_range().contains_range(frange) {
+        mark::hit!(cursor_on_tail);
+        mark::hit!(cursor_on_tail_closure);
+    } else {
+        return None;
+    }
+    Some((tail_expr, action, wrap_expr))
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn infer_return_type_specified_inferred() {
+        mark::check!(existing_infer_ret_type);
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() -> <|>_ {
+    45
+}"#,
+            r#"fn foo() -> i32 {
+    45
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type_specified_inferred_closure() {
+        mark::check!(existing_infer_ret_type_closure);
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() {
+    || -> _ {<|>45};
+}"#,
+            r#"fn foo() {
+    || -> i32 {45};
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type_cursor_at_return_type_pos() {
+        mark::check!(cursor_in_ret_position);
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() <|>{
+    45
+}"#,
+            r#"fn foo() -> i32 {
+    45
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type_cursor_at_return_type_pos_closure() {
+        mark::check!(cursor_in_ret_position_closure);
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() {
+    || <|>45
+}"#,
+            r#"fn foo() {
+    || -> i32 {45}
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type() {
+        mark::check!(cursor_on_tail);
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() {
+    45<|>
+}"#,
+            r#"fn foo() -> i32 {
+    45
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type_nested() {
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() {
+    if true {
+        3<|>
+    } else {
+        5
+    }
+}"#,
+            r#"fn foo() -> i32 {
+    if true {
+        3
+    } else {
+        5
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn not_applicable_ret_type_specified() {
+        mark::check!(existing_ret_type);
+        check_assist_not_applicable(
+            infer_function_return_type,
+            r#"fn foo() -> i32 {
+    ( 45<|> + 32 ) * 123
+}"#,
+        );
+    }
+
+    #[test]
+    fn not_applicable_non_tail_expr() {
+        check_assist_not_applicable(
+            infer_function_return_type,
+            r#"fn foo() {
+    let x = <|>3;
+    ( 45 + 32 ) * 123
+}"#,
+        );
+    }
+
+    #[test]
+    fn not_applicable_unit_return_type() {
+        check_assist_not_applicable(
+            infer_function_return_type,
+            r#"fn foo() {
+    (<|>)
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type_closure_block() {
+        mark::check!(cursor_on_tail_closure);
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() {
+    |x: i32| {
+        x<|>
+    };
+}"#,
+            r#"fn foo() {
+    |x: i32| -> i32 {
+        x
+    };
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type_closure() {
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() {
+    |x: i32| { x<|> };
+}"#,
+            r#"fn foo() {
+    |x: i32| -> i32 { x };
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type_closure_wrap() {
+        mark::check!(wrap_closure_non_block_expr);
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() {
+    |x: i32| x<|>;
+}"#,
+            r#"fn foo() {
+    |x: i32| -> i32 {x};
+}"#,
+        );
+    }
+
+    #[test]
+    fn infer_return_type_nested_closure() {
+        check_assist(
+            infer_function_return_type,
+            r#"fn foo() {
+    || {
+        if true {
+            3<|>
+        } else {
+            5
+        }
+    }
+}"#,
+            r#"fn foo() {
+    || -> i32 {
+        if true {
+            3
+        } else {
+            5
+        }
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn not_applicable_ret_type_specified_closure() {
+        mark::check!(existing_ret_type_closure);
+        check_assist_not_applicable(
+            infer_function_return_type,
+            r#"fn foo() {
+    || -> i32 { 3<|> }
+}"#,
+        );
+    }
+
+    #[test]
+    fn not_applicable_non_tail_expr_closure() {
+        check_assist_not_applicable(
+            infer_function_return_type,
+            r#"fn foo() {
+    || -> i32 {
+        let x = 3<|>;
+        6
+    }
+}"#,
+        );
+    }
+}
diff --git a/crates/assists/src/handlers/introduce_named_lifetime.rs b/crates/assists/src/handlers/introduce_named_lifetime.rs
index 5f623e5f7..4cc8dae65 100644
--- a/crates/assists/src/handlers/introduce_named_lifetime.rs
+++ b/crates/assists/src/handlers/introduce_named_lifetime.rs
@@ -36,7 +36,7 @@ static ASSIST_LABEL: &str = "Introduce named lifetime";
 // FIXME: should also add support for the case fun(f: &Foo) -> &<|>Foo
 pub(crate) fn introduce_named_lifetime(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
     let lifetime_token = ctx
-        .find_token_at_offset(SyntaxKind::LIFETIME)
+        .find_token_syntax_at_offset(SyntaxKind::LIFETIME)
         .filter(|lifetime| lifetime.text() == "'_")?;
     if let Some(fn_def) = lifetime_token.ancestors().find_map(ast::Fn::cast) {
         generate_fn_def_assist(acc, &fn_def, lifetime_token.text_range())
diff --git a/crates/assists/src/handlers/invert_if.rs b/crates/assists/src/handlers/invert_if.rs
index 461fcf862..ea722b91b 100644
--- a/crates/assists/src/handlers/invert_if.rs
+++ b/crates/assists/src/handlers/invert_if.rs
@@ -29,7 +29,7 @@ use crate::{
 // ```
 
 pub(crate) fn invert_if(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let if_keyword = ctx.find_token_at_offset(T![if])?;
+    let if_keyword = ctx.find_token_syntax_at_offset(T![if])?;
     let expr = ast::IfExpr::cast(if_keyword.parent())?;
     let if_range = if_keyword.text_range();
     let cursor_in_range = if_range.contains_range(ctx.frange.range);
diff --git a/crates/assists/src/handlers/raw_string.rs b/crates/assists/src/handlers/raw_string.rs
index 9ddd116e0..4c759cc25 100644
--- a/crates/assists/src/handlers/raw_string.rs
+++ b/crates/assists/src/handlers/raw_string.rs
@@ -1,11 +1,6 @@
 use std::borrow::Cow;
 
-use syntax::{
-    ast::{self, HasQuotes, HasStringValue},
-    AstToken,
-    SyntaxKind::{RAW_STRING, STRING},
-    TextRange, TextSize,
-};
+use syntax::{ast, AstToken, TextRange, TextSize};
 use test_utils::mark;
 
 use crate::{AssistContext, AssistId, AssistKind, Assists};
@@ -26,7 +21,10 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // }
 // ```
 pub(crate) fn make_raw_string(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let token = ctx.find_token_at_offset(STRING).and_then(ast::String::cast)?;
+    let token = ctx.find_token_at_offset::<ast::String>()?;
+    if token.is_raw() {
+        return None;
+    }
     let value = token.value()?;
     let target = token.syntax().text_range();
     acc.add(
@@ -65,7 +63,10 @@ pub(crate) fn make_raw_string(acc: &mut Assists, ctx: &AssistContext) -> Option<
 // }
 // ```
 pub(crate) fn make_usual_string(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let token = ctx.find_token_at_offset(RAW_STRING).and_then(ast::RawString::cast)?;
+    let token = ctx.find_token_at_offset::<ast::String>()?;
+    if !token.is_raw() {
+        return None;
+    }
     let value = token.value()?;
     let target = token.syntax().text_range();
     acc.add(
@@ -104,11 +105,15 @@ pub(crate) fn make_usual_string(acc: &mut Assists, ctx: &AssistContext) -> Optio
 // }
 // ```
 pub(crate) fn add_hash(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let token = ctx.find_token_at_offset(RAW_STRING)?;
-    let target = token.text_range();
+    let token = ctx.find_token_at_offset::<ast::String>()?;
+    if !token.is_raw() {
+        return None;
+    }
+    let text_range = token.syntax().text_range();
+    let target = text_range;
     acc.add(AssistId("add_hash", AssistKind::Refactor), "Add #", target, |edit| {
-        edit.insert(token.text_range().start() + TextSize::of('r'), "#");
-        edit.insert(token.text_range().end(), "#");
+        edit.insert(text_range.start() + TextSize::of('r'), "#");
+        edit.insert(text_range.end(), "#");
     })
 }
 
@@ -128,7 +133,10 @@ pub(crate) fn add_hash(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
 // }
 // ```
 pub(crate) fn remove_hash(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let token = ctx.find_token_at_offset(RAW_STRING).and_then(ast::RawString::cast)?;
+    let token = ctx.find_token_at_offset::<ast::String>()?;
+    if !token.is_raw() {
+        return None;
+    }
 
     let text = token.text().as_str();
     if !text.starts_with("r#") && text.ends_with('#') {
diff --git a/crates/assists/src/handlers/remove_mut.rs b/crates/assists/src/handlers/remove_mut.rs
index 44f41daa9..575b271f7 100644
--- a/crates/assists/src/handlers/remove_mut.rs
+++ b/crates/assists/src/handlers/remove_mut.rs
@@ -18,7 +18,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // }
 // ```
 pub(crate) fn remove_mut(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let mut_token = ctx.find_token_at_offset(T![mut])?;
+    let mut_token = ctx.find_token_syntax_at_offset(T![mut])?;
     let delete_from = mut_token.text_range().start();
     let delete_to = match mut_token.next_token() {
         Some(it) if it.kind() == SyntaxKind::WHITESPACE => it.text_range().end(),
diff --git a/crates/assists/src/handlers/replace_let_with_if_let.rs b/crates/assists/src/handlers/replace_let_with_if_let.rs
index a5bcbda24..69d3b08d3 100644
--- a/crates/assists/src/handlers/replace_let_with_if_let.rs
+++ b/crates/assists/src/handlers/replace_let_with_if_let.rs
@@ -37,7 +37,7 @@ use ide_db::ty_filter::TryEnum;
 // fn compute() -> Option<i32> { None }
 // ```
 pub(crate) fn replace_let_with_if_let(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let let_kw = ctx.find_token_at_offset(T![let])?;
+    let let_kw = ctx.find_token_syntax_at_offset(T![let])?;
     let let_stmt = let_kw.ancestors().find_map(ast::LetStmt::cast)?;
     let init = let_stmt.initializer()?;
     let original_pat = let_stmt.pat()?;
diff --git a/crates/assists/src/handlers/replace_string_with_char.rs b/crates/assists/src/handlers/replace_string_with_char.rs
index 4ca87a8ec..b4b898846 100644
--- a/crates/assists/src/handlers/replace_string_with_char.rs
+++ b/crates/assists/src/handlers/replace_string_with_char.rs
@@ -1,8 +1,4 @@
-use syntax::{
-    ast::{self, HasStringValue},
-    AstToken,
-    SyntaxKind::STRING,
-};
+use syntax::{ast, AstToken, SyntaxKind::STRING};
 
 use crate::{AssistContext, AssistId, AssistKind, Assists};
 
@@ -22,7 +18,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // }
 // ```
 pub(crate) fn replace_string_with_char(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let token = ctx.find_token_at_offset(STRING).and_then(ast::String::cast)?;
+    let token = ctx.find_token_syntax_at_offset(STRING).and_then(ast::String::cast)?;
     let value = token.value()?;
     let target = token.syntax().text_range();
 
diff --git a/crates/assists/src/handlers/split_import.rs b/crates/assists/src/handlers/split_import.rs
index 15e67eaa1..ef1f6b8a1 100644
--- a/crates/assists/src/handlers/split_import.rs
+++ b/crates/assists/src/handlers/split_import.rs
@@ -16,7 +16,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
 // use std::{collections::HashMap};
 // ```
 pub(crate) fn split_import(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let colon_colon = ctx.find_token_at_offset(T![::])?;
+    let colon_colon = ctx.find_token_syntax_at_offset(T![::])?;
     let path = ast::Path::cast(colon_colon.parent())?.qualifier()?;
     let top_path = successors(Some(path.clone()), |it| it.parent_path()).last()?;
 
diff --git a/crates/assists/src/handlers/unwrap_block.rs b/crates/assists/src/handlers/unwrap_block.rs
index 3851aeb3e..36ef871b9 100644
--- a/crates/assists/src/handlers/unwrap_block.rs
+++ b/crates/assists/src/handlers/unwrap_block.rs
@@ -29,7 +29,7 @@ pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     let assist_id = AssistId("unwrap_block", AssistKind::RefactorRewrite);
     let assist_label = "Unwrap block";
 
-    let l_curly_token = ctx.find_token_at_offset(T!['{'])?;
+    let l_curly_token = ctx.find_token_syntax_at_offset(T!['{'])?;
     let mut block = ast::BlockExpr::cast(l_curly_token.parent())?;
     let mut parent = block.syntax().parent()?;
     if ast::MatchArm::can_cast(parent.kind()) {