diff --git a/crates/ide_assists/src/handlers/extract_function.rs b/crates/ide_assists/src/handlers/extract_function.rs
index 282f7a60af9..7ffb5728cc5 100644
--- a/crates/ide_assists/src/handlers/extract_function.rs
+++ b/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(())
}
"#,
);
diff --git a/crates/ide_db/src/helpers/famous_defs.rs b/crates/ide_db/src/helpers/famous_defs.rs
index 08bd8e0cba6..18524986e2d 100644
--- a/crates/ide_db/src/helpers/famous_defs.rs
+++ b/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")
}