internal: use ControlFlow in "extract function" assist

crates/ide_assists/src/handlers/extract_function.rs

@@ -2,10 +2,15 @@ use std::{hash::BuildHasherDefault, iter};
 
 use ast::make;
 use either::Either;
-use hir::{HirDisplay, InFile, Local, Semantics, TypeInfo};
+use hir::{HirDisplay, InFile, Local, ModuleDef, Semantics, TypeInfo};
 use ide_db::{
     defs::{Definition, NameRefClass},
-    helpers::node_ext::{preorder_expr, walk_expr, walk_pat, walk_patterns_in_expr},
+    helpers::{
+        insert_use::{insert_use, ImportScope},
+        mod_path_to_ast,
+        node_ext::{preorder_expr, walk_expr, walk_pat, walk_patterns_in_expr},
+        FamousDefs,
+    },
     search::{FileReference, ReferenceCategory, SearchScope},
     RootDatabase,
 };
@@ -86,6 +91,8 @@ pub(crate) fn extract_function(acc: &mut Assists, ctx: &AssistContext) -> Option
 
     let target_range = body.text_range();
 
+    let scope = ImportScope::find_insert_use_container_with_macros(&node, &ctx.sema)?;
+
     acc.add(
         AssistId("extract_function", crate::AssistKind::RefactorExtract),
         "Extract into function",
@@ -118,10 +125,34 @@ pub(crate) fn extract_function(acc: &mut Assists, ctx: &AssistContext) -> Option
 
             let fn_def = format_function(ctx, module, &fun, old_indent, new_indent);
             let insert_offset = insert_after.text_range().end();
+
+            if fn_def.contains("ControlFlow") {
+                let scope = match scope {
+                    ImportScope::File(it) => ImportScope::File(builder.make_mut(it)),
+                    ImportScope::Module(it) => ImportScope::Module(builder.make_mut(it)),
+                    ImportScope::Block(it) => ImportScope::Block(builder.make_mut(it)),
+                };
+
+                let control_flow_enum =
+                    FamousDefs(&ctx.sema, Some(module.krate())).core_ops_ControlFlow();
+
+                if let Some(control_flow_enum) = control_flow_enum {
+                    let mod_path = module.find_use_path_prefixed(
+                        ctx.sema.db,
+                        ModuleDef::from(control_flow_enum),
+                        ctx.config.insert_use.prefix_kind,
+                    );
+
+                    if let Some(mod_path) = mod_path {
+                        insert_use(&scope, mod_path_to_ast(&mod_path), &ctx.config.insert_use);
+                    }
+                }
+            }
+
             match ctx.config.snippet_cap {
                 Some(cap) => builder.insert_snippet(cap, insert_offset, fn_def),
                 None => builder.insert(insert_offset, fn_def),
-            }
+            };
         },
     )
 }
@@ -1184,7 +1215,17 @@ impl FlowHandler {
                 let action = action.make_result_handler(None);
                 let stmt = make::expr_stmt(action);
                 let block = make::block_expr(iter::once(stmt.into()), None);
-                let condition = make::condition(call_expr, None);
+                let controlflow_break_path = make::path_from_text("ControlFlow::Break");
+                let condition = make::condition(
+                    call_expr,
+                    Some(
+                        make::tuple_struct_pat(
+                            controlflow_break_path,
+                            iter::once(make::wildcard_pat().into()),
+                        )
+                        .into(),
+                    ),
+                );
                 make::expr_if(condition, block, None)
             }
             FlowHandler::IfOption { action } => {
@@ -1326,7 +1367,7 @@ impl Function {
                     .unwrap_or_else(make::ty_placeholder);
                 make::ext::ty_result(fun_ty.make_ty(ctx, module), handler_ty)
             }
-            FlowHandler::If { .. } => make::ext::ty_bool(),
+            FlowHandler::If { .. } => make::ty("ControlFlow<()>"),
             FlowHandler::IfOption { action } => {
                 let handler_ty = action
                     .expr_ty(ctx)
@@ -1461,8 +1502,11 @@ fn make_body(
             })
         }
         FlowHandler::If { .. } => {
-            let lit_false = make::expr_literal("false");
-            with_tail_expr(block, lit_false.into())
+            let controlflow_continue = make::expr_call(
+                make::expr_path(make::path_from_text("ControlFlow::Continue")),
+                make::arg_list(iter::once(make::expr_unit())),
+            );
+            with_tail_expr(block, controlflow_continue.into())
         }
         FlowHandler::IfOption { .. } => {
             let none = make::expr_path(make::ext::ident_path("None"));
@@ -1638,7 +1682,10 @@ fn update_external_control_flow(handler: &FlowHandler, syntax: &SyntaxNode) {
 fn make_rewritten_flow(handler: &FlowHandler, arg_expr: Option<ast::Expr>) -> Option<ast::Expr> {
     let value = match handler {
         FlowHandler::None | FlowHandler::Try { .. } => return None,
-        FlowHandler::If { .. } => make::expr_literal("true").into(),
+        FlowHandler::If { .. } => make::expr_call(
+            make::expr_path(make::path_from_text("ControlFlow::Break")),
+            make::arg_list(iter::once(make::expr_unit())),
+        ),
         FlowHandler::IfOption { .. } => {
             let expr = arg_expr.unwrap_or_else(|| make::expr_tuple(Vec::new()));
             let args = make::arg_list(iter::once(expr));
@@ -3270,6 +3317,7 @@ fn foo() {
         check_assist(
             extract_function,
             r#"
+//- minicore: try
 fn foo() {
     loop {
         let mut n = 1;
@@ -3281,21 +3329,23 @@ fn foo() {
 }
 "#,
             r#"
+use core::ops::ControlFlow;
+
 fn foo() {
     loop {
         let mut n = 1;
-        if fun_name(&mut n) {
+        if let ControlFlow::Break(_) = fun_name(&mut n) {
             break;
         }
         let h = 1 + n;
     }
 }
 
-fn $0fun_name(n: &mut i32) -> bool {
+fn $0fun_name(n: &mut i32) -> ControlFlow<()> {
     let m = *n + 1;
-    return true;
+    return ControlFlow::Break(());
     *n += m;
-    false
+    ControlFlow::Continue(())
 }
 "#,
         );
@@ -3306,6 +3356,7 @@ fn $0fun_name(n: &mut i32) -> bool {
         check_assist(
             extract_function,
             r#"
+//- minicore: try
 fn foo() {
     loop {
         let mut n = 1;
@@ -3318,22 +3369,24 @@ fn foo() {
 }
 "#,
             r#"
+use core::ops::ControlFlow;
+
 fn foo() {
     loop {
         let mut n = 1;
-        if fun_name(n) {
+        if let ControlFlow::Break(_) = fun_name(n) {
             break;
         }
         let h = 1;
     }
 }
 
-fn $0fun_name(n: i32) -> bool {
+fn $0fun_name(n: i32) -> ControlFlow<()> {
     let m = n + 1;
     if m == 42 {
-        return true;
+        return ControlFlow::Break(());
     }
-    false
+    ControlFlow::Continue(())
 }
 "#,
         );

crates/ide_db/src/helpers/famous_defs.rs

@@ -68,6 +68,10 @@ impl FamousDefs<'_, '_> {
         self.find_trait("core:ops:Deref")
     }
 
+    pub fn core_ops_ControlFlow(&self) -> Option<Enum> {
+        self.find_enum("core:ops:ControlFlow")
+    }
+
     pub fn alloc(&self) -> Option<Crate> {
         self.find_crate("alloc")
     }