fix: "Extract to function" assist preserves `break` and `continue` labels

Adds a label / lifetime parameter to `ide_assists::handlers::extract_function::FlowKind::{Break, Continue}`, adds support for emitting labels to `syntax::ast::make::{expr_break, expr_continue}`, and implements the required machinery to let `extract_function` make use of them.

This does modify the external API of the `syntax` crate, but the changes there are simple, not used outside `ide_assists`, and, well, we should probably support emitting `break` and `continue` labels through `syntax` anyways, they're part of the language spec.

Closes #11413.

4 files changed, 146 insertions, 26 deletions

diff --git a/crates/ide_assists/src/handlers/convert_to_guarded_return.rs b/crates/ide_assists/src/handlers/convert_to_guarded_return.rs
index a7c43413d2c..445567b7f90 100644
--- a/crates/ide_assists/src/handlers/convert_to_guarded_return.rs
+++ b/crates/ide_assists/src/handlers/convert_to_guarded_return.rs
@@ -97,7 +97,7 @@ pub(crate) fn convert_to_guarded_return(acc: &mut Assists, ctx: &AssistContext)
     let parent_container = parent_block.syntax().parent()?;
 
     let early_expression: ast::Expr = match parent_container.kind() {
-        WHILE_EXPR | LOOP_EXPR => make::expr_continue(),
+        WHILE_EXPR | LOOP_EXPR => make::expr_continue(None),
         FN => make::expr_return(None),
         _ => return None,
     };
diff --git a/crates/ide_assists/src/handlers/convert_while_to_loop.rs b/crates/ide_assists/src/handlers/convert_while_to_loop.rs
index 345d0c2063f..739097067d0 100644
--- a/crates/ide_assists/src/handlers/convert_while_to_loop.rs
+++ b/crates/ide_assists/src/handlers/convert_while_to_loop.rs
@@ -53,7 +53,7 @@ pub(crate) fn convert_while_to_loop(acc: &mut Assists, ctx: &AssistContext) -> O
             let while_indent_level = IndentLevel::from_node(while_expr.syntax());
 
             let break_block =
-                make::block_expr(once(make::expr_stmt(make::expr_break(None)).into()), None)
+                make::block_expr(once(make::expr_stmt(make::expr_break(None, None)).into()), None)
                     .indent(while_indent_level);
             let block_expr = if is_pattern_cond(while_cond.clone()) {
                 let if_expr = make::expr_if(while_cond, while_body, Some(break_block.into()));
diff --git a/crates/ide_assists/src/handlers/extract_function.rs b/crates/ide_assists/src/handlers/extract_function.rs
index e48a1cd9ef2..53cdeaaf925 100644
--- a/crates/ide_assists/src/handlers/extract_function.rs
+++ b/crates/ide_assists/src/handlers/extract_function.rs
@@ -287,10 +287,10 @@ enum FlowKind {
     Try {
         kind: TryKind,
     },
-    /// Break with value (`break $expr;`)
-    Break(Option<ast::Expr>),
-    /// Continue
-    Continue,
+    /// Break with label and value (`break 'label $expr;`)
+    Break(Option<ast::Lifetime>, Option<ast::Expr>),
+    /// Continue with label (`continue 'label;`)
+    Continue(Option<ast::Lifetime>),
 }
 
 #[derive(Debug, Clone)]
@@ -433,21 +433,21 @@ impl FlowKind {
     fn make_result_handler(&self, expr: Option<ast::Expr>) -> ast::Expr {
         match self {
             FlowKind::Return(_) => make::expr_return(expr),
-            FlowKind::Break(_) => make::expr_break(expr),
+            FlowKind::Break(label, _) => make::expr_break(label.clone(), expr),
             FlowKind::Try { .. } => {
                 stdx::never!("cannot have result handler with try");
                 expr.unwrap_or_else(|| make::expr_return(None))
             }
-            FlowKind::Continue => {
+            FlowKind::Continue(label) => {
                 stdx::always!(expr.is_none(), "continue with value is not possible");
-                make::expr_continue()
+                make::expr_continue(label.clone())
             }
         }
     }
 
     fn expr_ty(&self, ctx: &AssistContext) -> Option<hir::Type> {
         match self {
-            FlowKind::Return(Some(expr)) | FlowKind::Break(Some(expr)) => {
+            FlowKind::Return(Some(expr)) | FlowKind::Break(_, Some(expr)) => {
                 ctx.sema.type_of_expr(expr).map(TypeInfo::adjusted)
             }
             FlowKind::Try { .. } => {
@@ -839,8 +839,8 @@ impl FunctionBody {
                 cov_mark::hit!(external_control_flow_break_and_continue);
                 return None;
             }
-            (None, None, Some(b), None) => Some(FlowKind::Break(b.expr())),
-            (None, None, None, Some(_)) => Some(FlowKind::Continue),
+            (None, None, Some(b), None) => Some(FlowKind::Break(b.lifetime(), b.expr())),
+            (None, None, None, Some(c)) => Some(FlowKind::Continue(c.lifetime())),
             (None, None, None, None) => None,
         };
 
@@ -1185,20 +1185,20 @@ impl FlowHandler {
                 let action = flow_kind.clone();
                 if *ret_ty == FunType::Unit {
                     match flow_kind {
-                        FlowKind::Return(None) | FlowKind::Break(None) | FlowKind::Continue => {
-                            FlowHandler::If { action }
-                        }
-                        FlowKind::Return(_) | FlowKind::Break(_) => {
+                        FlowKind::Return(None)
+                        | FlowKind::Break(_, None)
+                        | FlowKind::Continue(_) => FlowHandler::If { action },
+                        FlowKind::Return(_) | FlowKind::Break(_, _) => {
                             FlowHandler::IfOption { action }
                         }
                         FlowKind::Try { kind } => FlowHandler::Try { kind: kind.clone() },
                     }
                 } else {
                     match flow_kind {
-                        FlowKind::Return(None) | FlowKind::Break(None) | FlowKind::Continue => {
-                            FlowHandler::MatchOption { none: action }
-                        }
-                        FlowKind::Return(_) | FlowKind::Break(_) => {
+                        FlowKind::Return(None)
+                        | FlowKind::Break(_, None)
+                        | FlowKind::Continue(_) => FlowHandler::MatchOption { none: action },
+                        FlowKind::Return(_) | FlowKind::Break(_, _) => {
                             FlowHandler::MatchResult { err: action }
                         }
                         FlowKind::Try { kind } => FlowHandler::Try { kind: kind.clone() },
@@ -3405,6 +3405,76 @@ fn $0fun_name(n: i32) -> ControlFlow<()> {
     }
 
     #[test]
+    fn break_loop_nested_labeled() {
+        check_assist(
+            extract_function,
+            r#"
+//- minicore: try
+fn foo() {
+    'bar: loop {
+        loop {
+            $0break 'bar;$0
+        }
+    }
+}
+"#,
+            r#"
+use core::ops::ControlFlow;
+
+fn foo() {
+    'bar: loop {
+        loop {
+            if let ControlFlow::Break(_) = fun_name() {
+                break 'bar;
+            }
+        }
+    }
+}
+
+fn $0fun_name() -> ControlFlow<()> {
+    return ControlFlow::Break(());
+    ControlFlow::Continue(())
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn continue_loop_nested_labeled() {
+        check_assist(
+            extract_function,
+            r#"
+//- minicore: try
+fn foo() {
+    'bar: loop {
+        loop {
+            $0continue 'bar;$0
+        }
+    }
+}
+"#,
+            r#"
+use core::ops::ControlFlow;
+
+fn foo() {
+    'bar: loop {
+        loop {
+            if let ControlFlow::Break(_) = fun_name() {
+                continue 'bar;
+            }
+        }
+    }
+}
+
+fn $0fun_name() -> ControlFlow<()> {
+    return ControlFlow::Break(());
+    ControlFlow::Continue(())
+}
+"#,
+        );
+    }
+
+    #[test]
     fn return_from_nested_loop() {
         check_assist(
             extract_function,
@@ -3609,6 +3679,46 @@ fn $0fun_name() -> Option<i32> {
     }
 
     #[test]
+    fn break_with_value_and_label() {
+        check_assist(
+            extract_function,
+            r#"
+fn foo() -> i32 {
+    'bar: loop {
+        let n = 1;
+        $0let k = 1;
+        if k == 42 {
+            break 'bar 3;
+        }
+        let m = k + 1;$0
+        let h = 1;
+    }
+}
+"#,
+            r#"
+fn foo() -> i32 {
+    'bar: loop {
+        let n = 1;
+        if let Some(value) = fun_name() {
+            break 'bar value;
+        }
+        let h = 1;
+    }
+}
+
+fn $0fun_name() -> Option<i32> {
+    let k = 1;
+    if k == 42 {
+        return Some(4);
+    }
+    let m = k + 1;
+    None
+}
+"#,
+        );
+    }
+
+    #[test]
     fn break_with_value_and_return() {
         check_assist(
             extract_function,
diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs
index a69a77c7ee2..d11d658a273 100644
--- a/crates/syntax/src/ast/make.rs
+++ b/crates/syntax/src/ast/make.rs
@@ -368,18 +368,28 @@ pub fn expr_empty_block() -> ast::Expr {
 pub fn expr_path(path: ast::Path) -> ast::Expr {
     expr_from_text(&path.to_string())
 }
-pub fn expr_continue() -> ast::Expr {
-    expr_from_text("continue")
+pub fn expr_continue(label: Option<ast::Lifetime>) -> ast::Expr {
+    match label {
+        Some(label) => expr_from_text(&format!("continue {}", label)),
+        None => expr_from_text("continue"),
+    }
 }
 // Consider `op: SyntaxKind` instead for nicer syntax at the call-site?
 pub fn expr_bin_op(lhs: ast::Expr, op: ast::BinaryOp, rhs: ast::Expr) -> ast::Expr {
     expr_from_text(&format!("{} {} {}", lhs, op, rhs))
 }
-pub fn expr_break(expr: Option<ast::Expr>) -> ast::Expr {
-    match expr {
-        Some(expr) => expr_from_text(&format!("break {}", expr)),
-        None => expr_from_text("break"),
+pub fn expr_break(label: Option<ast::Lifetime>, expr: Option<ast::Expr>) -> ast::Expr {
+    let mut s = String::from("break");
+
+    if let Some(label) = label {
+        format_to!(s, " {}", label);
     }
+
+    if let Some(expr) = expr {
+        format_to!(s, " {}", expr);
+    }
+
+    expr_from_text(&s)
 }
 pub fn expr_return(expr: Option<ast::Expr>) -> ast::Expr {
     match expr {