ide_assists/handlers/

desugar_try_expr.rs

use std::iter;

use ide_db::{assists::AssistId, ty_filter::TryEnum};
use syntax::{
    AstNode, T,
    ast::{
        self,
        edit::{AstNodeEdit, IndentLevel},
        syntax_factory::SyntaxFactory,
    },
};

use crate::assist_context::{AssistContext, Assists};

// Assist: desugar_try_expr_match
//
// Replaces a `try` expression with a `match` expression.
//
// ```
// # //- minicore: try, option
// fn handle() {
//     let pat = Some(true)$0?;
// }
// ```
// ->
// ```
// fn handle() {
//     let pat = match Some(true) {
//         Some(it) => it,
//         None => return None,
//     };
// }
// ```

// Assist: desugar_try_expr_let_else
//
// Replaces a `try` expression with a `let else` statement.
//
// ```
// # //- minicore: try, option
// fn handle() {
//     let pat = Some(true)$0?;
// }
// ```
// ->
// ```
// fn handle() {
//     let Some(pat) = Some(true) else {
//         return None;
//     };
// }
// ```
pub(crate) fn desugar_try_expr(acc: &mut Assists, ctx: &AssistContext<'_, '_>) -> Option<()> {
    let question_tok = ctx.find_token_syntax_at_offset(T![?])?;
    let try_expr = question_tok.parent().and_then(ast::TryExpr::cast)?;

    let expr = try_expr.expr()?;
    let expr_type_info = ctx.sema.type_of_expr(&expr)?;

    let try_enum = TryEnum::from_ty(&ctx.sema, &expr_type_info.original)?;

    let target = try_expr.syntax().text_range();
    acc.add(
        AssistId::refactor_rewrite("desugar_try_expr_match"),
        "Replace try expression with match",
        target,
        |builder| {
            let editor = builder.make_editor(try_expr.syntax());
            let make = editor.make();
            let sad_pat = match try_enum {
                TryEnum::Option => make.path_pat(make.ident_path("None")),
                TryEnum::Result => make
                    .tuple_struct_pat(
                        make.ident_path("Err"),
                        iter::once(make.path_pat(make.ident_path("err"))),
                    )
                    .into(),
            };
            let sad_expr = make.expr_return(Some(sad_expr(try_enum, make, || {
                make.expr_path(make.ident_path("err"))
            })));

            let happy_arm = make.match_arm(
                try_enum.happy_pattern(make.ident_pat(false, false, make.name("it")).into()),
                None,
                make.expr_path(make.ident_path("it")),
            );
            let sad_arm = make.match_arm(sad_pat, None, sad_expr.into());

            let match_arm_list = make.match_arm_list([happy_arm, sad_arm]);

            let expr_match = make
                .expr_match(expr.clone(), match_arm_list)
                .indent(IndentLevel::from_node(try_expr.syntax()));

            editor.replace(try_expr.syntax(), expr_match.syntax());
            builder.add_file_edits(ctx.vfs_file_id(), editor);
        },
    );

    if let Some(let_stmt) = try_expr.syntax().parent().and_then(ast::LetStmt::cast)
        && let_stmt.let_else().is_none()
    {
        let pat = let_stmt.pat()?;
        acc.add(
            AssistId::refactor_rewrite("desugar_try_expr_let_else"),
            "Replace try expression with let else",
            target,
            |builder| {
                let editor = builder.make_editor(let_stmt.syntax());
                let make = editor.make();

                let indent_level = IndentLevel::from_node(let_stmt.syntax());
                let fill_expr = || crate::utils::expr_fill_default(ctx.config);
                let new_let_stmt = make.let_else_stmt(
                    try_enum.happy_pattern(pat),
                    let_stmt.ty().map(|ty| match try_enum {
                        TryEnum::Option => make.ty_option(ty).into(),
                        TryEnum::Result => make.ty_result(ty, make.ty_infer().into()).into(),
                    }),
                    expr,
                    make.block_expr(
                        iter::once(
                            make.expr_stmt(
                                make.expr_return(Some(sad_expr(try_enum, make, fill_expr))).into(),
                            )
                            .into(),
                        ),
                        None,
                    )
                    .indent(indent_level),
                );
                editor.replace(let_stmt.syntax(), new_let_stmt.syntax());
                builder.add_file_edits(ctx.vfs_file_id(), editor);
            },
        );
    }
    Some(())
}

fn sad_expr(try_enum: TryEnum, make: &SyntaxFactory, err: impl Fn() -> ast::Expr) -> ast::Expr {
    match try_enum {
        TryEnum::Option => make.expr_path(make.ident_path("None")),
        TryEnum::Result => make
            .expr_call(make.expr_path(make.ident_path("Err")), make.arg_list(iter::once(err())))
            .into(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use crate::tests::{check_assist, check_assist_by_label, check_assist_not_applicable};

    #[test]
    fn test_desugar_try_expr_not_applicable() {
        check_assist_not_applicable(
            desugar_try_expr,
            r#"
                fn test() {
                    let pat: u32 = 25$0;
                }
            "#,
        );
    }

    #[test]
    fn test_desugar_try_expr_option() {
        check_assist(
            desugar_try_expr,
            r#"
//- minicore: try, option
fn test() {
    let pat = Some(true)$0?;
}
            "#,
            r#"
fn test() {
    let pat = match Some(true) {
        Some(it) => it,
        None => return None,
    };
}
            "#,
        );
    }

    #[test]
    fn test_desugar_try_expr_result() {
        check_assist(
            desugar_try_expr,
            r#"
//- minicore: try, from, result
fn test() {
    let pat = Ok(true)$0?;
}
            "#,
            r#"
fn test() {
    let pat = match Ok(true) {
        Ok(it) => it,
        Err(err) => return Err(err),
    };
}
            "#,
        );
    }

    #[test]
    fn test_desugar_try_expr_option_let_else() {
        check_assist_by_label(
            desugar_try_expr,
            r#"
//- minicore: try, option
fn test() {
    let pat = Some(true)$0?;
}
            "#,
            r#"
fn test() {
    let Some(pat) = Some(true) else {
        return None;
    };
}
            "#,
            "Replace try expression with let else",
        );
    }

    #[test]
    fn test_desugar_try_expr_result_let_else() {
        check_assist_by_label(
            desugar_try_expr,
            r#"
//- minicore: try, from, result
fn test() {
    let pat = Ok(true)$0?;
}
            "#,
            r#"
fn test() {
    let Ok(pat) = Ok(true) else {
        return Err(todo!());
    };
}
            "#,
            "Replace try expression with let else",
        );
    }

    #[test]
    fn test_desugar_try_expr_option_let_else_with_type() {
        check_assist_by_label(
            desugar_try_expr,
            r#"
//- minicore: try, option
fn test() {
    let pat: bool = Some(true)$0?;
}
            "#,
            r#"
fn test() {
    let Some(pat): Option<bool> = Some(true) else {
        return None;
    };
}
            "#,
            "Replace try expression with let else",
        );
    }

    #[test]
    fn test_desugar_try_expr_result_let_else_with_type() {
        check_assist_by_label(
            desugar_try_expr,
            r#"
//- minicore: try, result
fn test() {
    let pat: bool = Ok(true)$0?;
}
            "#,
            r#"
fn test() {
    let Ok(pat): Result<bool, _> = Ok(true) else {
        return Err(todo!());
    };
}
            "#,
            "Replace try expression with let else",
        );
    }
}