<commit_before>use hir::HirDisplay;
use syntax::{ast, AstNode, SyntaxToken, 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.
//
// ```
// fn foo() { 4<|>2i32 }
// ```
// ->
// ```
// fn foo() -> i32 { 42i32 }
// ```
pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let expr = ctx.find_node_at_offset::<ast::Expr>()?;
    let (tail_expr, insert_pos) = extract_tail(expr)?;
    let module = ctx.sema.scope(tail_expr.syntax()).module()?;
    let ty = ctx.sema.type_of_expr(&tail_expr).filter(|ty| !ty.is_unit())?;
    let ty = ty.display_source_code(ctx.db(), module.into()).ok()?;

    acc.add(
        AssistId("change_return_type_to_result", AssistKind::RefactorRewrite),
        "Wrap return type in Result",
        tail_expr.syntax().text_range(),
        |builder| {
            let insert_pos = insert_pos.text_range().end() + TextSize::from(1);
            builder.insert(insert_pos, &format!("-> {} ", ty));
        },
    )
}

fn extract_tail(expr: ast::Expr) -> Option<(ast::Expr, SyntaxToken)> {
    let (ret_ty, tail_expr, insert_pos) =
        if let Some(closure) = expr.syntax().ancestors().find_map(ast::ClosureExpr::cast) {
            let tail_expr = match closure.body()? {
                ast::Expr::BlockExpr(block) => block.expr()?,
                body => body,
            };
            let ret_ty = closure.ret_type();
            let rpipe = closure.param_list()?.syntax().last_token()?;
            (ret_ty, tail_expr, rpipe)
        } else {
            let func = expr.syntax().ancestors().find_map(ast::Fn::cast)?;
            let tail_expr = func.body()?.expr()?;
            let ret_ty = func.ret_type();
            let rparen = func.param_list()?.r_paren_token()?;
            (ret_ty, tail_expr, rparen)
        };
    if ret_ty.is_some() {
        mark::hit!(existing_ret_type);
        mark::hit!(existing_ret_type_closure);
        return None;
    }
    // check whether the expr we were at is indeed the tail expression
    if !tail_expr.syntax().text_range().contains_range(expr.syntax().text_range()) {
        mark::hit!(not_tail_expr);
        return None;
    }
    Some((tail_expr, insert_pos))
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn infer_return_type() {
        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() {
        mark::check!(not_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() {
        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_nested_closure() {
        check_assist(
            infer_function_return_type,
            r#"fn foo() {
                || {<commit_msg>Wrap non-block expressions in closures with a block<commit_after>use hir::HirDisplay;
use syntax::{ast, AstNode, SyntaxToken, 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.
//
// ```
// fn foo() { 4<|>2i32 }
// ```
// ->
// ```
// fn foo() -> i32 { 42i32 }
// ```
pub(crate) fn infer_function_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let expr = ctx.find_node_at_offset::<ast::Expr>()?;
    let (tail_expr, insert_pos, wrap_expr) = extract_tail(expr)?;
    let module = ctx.sema.scope(tail_expr.syntax()).module()?;
    let ty = ctx.sema.type_of_expr(&tail_expr).filter(|ty| !ty.is_unit())?;
    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| {
            let insert_pos = insert_pos.text_range().end() + TextSize::from(1);
            builder.insert(insert_pos, &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));
            }
        },
    )
}

fn extract_tail(expr: ast::Expr) -> Option<(ast::Expr, SyntaxToken, bool)> {
    let (ret_ty, tail_expr, insert_pos, wrap_expr) =
        if let Some(closure) = expr.syntax().ancestors().find_map(ast::ClosureExpr::cast) {
            let (tail_expr, wrap_expr) = match closure.body()? {
                ast::Expr::BlockExpr(block) => (block.expr()?, false),
                body => (body, true),
            };
            let ret_ty = closure.ret_type();
            let rpipe = closure.param_list()?.syntax().last_token()?;
            (ret_ty, tail_expr, rpipe, wrap_expr)
        } else {
            let func = expr.syntax().ancestors().find_map(ast::Fn::cast)?;
            let tail_expr = func.body()?.expr()?;
            let ret_ty = func.ret_type();
            let rparen = func.param_list()?.r_paren_token()?;
            (ret_ty, tail_expr, rparen, false)
        };
    if ret_ty.is_some() {
        mark::hit!(existing_ret_type);
        mark::hit!(existing_ret_type_closure);
        return None;
    }
    // check whether the expr we were at is indeed the tail expression
    if !tail_expr.syntax().text_range().contains_range(expr.syntax().text_range()) {
        mark::hit!(not_tail_expr);
        return None;
    }
    Some((tail_expr, insert_pos, wrap_expr))
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn infer_return_type() {
        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() {
        mark::check!(not_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() {
        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() {
                || {
{'fraction_non_alphanumeric': 0.1439248754311997, 'fraction_numerical': 0.010540436949022614, 'mean_word_length': 1.5988545816733069, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 52, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1560888', 'n_tokens_mistral': 3462, 'n_tokens_neox': 3263, 'n_words': 596}