Add #[loop_match] for improved DFA codegen

compiler/rustc_feature/src/builtin_attrs.rs

@@ -583,6 +583,16 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         EncodeCrossCrate::Yes, min_generic_const_args, experimental!(type_const),
     ),
 
+    // #[loop_match] and #[const_continue]
+    gated!(
+        const_continue, Normal, template!(Word), ErrorFollowing,
+        EncodeCrossCrate::No, loop_match, experimental!(const_continue)
+    ),
+    gated!(
+        loop_match, Normal, template!(Word), ErrorFollowing,
+        EncodeCrossCrate::No, loop_match, experimental!(loop_match)
+    ),
+
     // ==========================================================================
     // Internal attributes: Stability, deprecation, and unsafe:
     // ==========================================================================

compiler/rustc_feature/src/unstable.rs

@@ -546,6 +546,8 @@ declare_features! (
     /// Allows using `#[link(kind = "link-arg", name = "...")]`
     /// to pass custom arguments to the linker.
     (unstable, link_arg_attribute, "1.76.0", Some(99427)),
+    /// Allows fused `loop`/`match` for direct intraprocedural jumps.
+    (incomplete, loop_match, "CURRENT_RUSTC_VERSION", Some(138777)),
     /// Give access to additional metadata about declarative macro meta-variables.
     (unstable, macro_metavar_expr, "1.61.0", Some(83527)),
     /// Provides a way to concatenate identifiers using metavariable expressions.

compiler/rustc_middle/src/thir.rs

@@ -375,6 +375,14 @@ pub enum ExprKind<'tcx> {
     Loop {
         body: ExprId,
     },
+    /// A `#[loop_match] loop { state = 'blk: { match state { ... } } }` expression.
+    LoopMatch {
+        /// The state variable that is updated, and also the scrutinee of the match
+        state: ExprId,
+        region_scope: region::Scope,
+        arms: Box<[ArmId]>,
+        match_span: Span,
+    },
     /// Special expression representing the `let` part of an `if let` or similar construct
     /// (including `if let` guards in match arms, and let-chains formed by `&&`).
     ///
@@ -451,6 +459,11 @@ pub enum ExprKind<'tcx> {
     Continue {
         label: region::Scope,
     },
+    /// A `#[const_continue] break` expression.
+    ConstContinue {
+        label: region::Scope,
+        value: ExprId,
+    },
     /// A `return` expression.
     Return {
         value: Option<ExprId>,

compiler/rustc_middle/src/thir/visit.rs

@@ -79,6 +79,12 @@ pub fn walk_expr<'thir, 'tcx: 'thir, V: Visitor<'thir, 'tcx>>(
             visitor.visit_pat(pat);
         }
         Loop { body } => visitor.visit_expr(&visitor.thir()[body]),
+        LoopMatch { state, ref arms, .. } => {
+            visitor.visit_expr(&visitor.thir()[state]);
+            for &arm in &**arms {
+                visitor.visit_arm(&visitor.thir()[arm]);
+            }
+        }
         Match { scrutinee, ref arms, .. } => {
             visitor.visit_expr(&visitor.thir()[scrutinee]);
             for &arm in &**arms {
@@ -104,6 +110,7 @@ pub fn walk_expr<'thir, 'tcx: 'thir, V: Visitor<'thir, 'tcx>>(
             }
         }
         Continue { label: _ } => {}
+        ConstContinue { value, label: _ } => visitor.visit_expr(&visitor.thir()[value]),
         Return { value } => {
             if let Some(value) = value {
                 visitor.visit_expr(&visitor.thir()[value])

compiler/rustc_mir_build/messages.ftl

@@ -84,6 +84,11 @@ mir_build_call_to_unsafe_fn_requires_unsafe_unsafe_op_in_unsafe_fn_allowed =
 
 mir_build_confused = missing patterns are not covered because `{$variable}` is interpreted as a constant pattern, not a new variable
 
+mir_build_const_continue_missing_value = a `#[const_continue]` must break to a label with a value
+
+mir_build_const_continue_unknown_jump_target = the target of this `#[const_continue]` is not statically known
+    .note = this value must be an integer or enum literal
+
 mir_build_const_defined_here = constant defined here
 
 mir_build_const_param_in_pattern = constant parameters cannot be referenced in patterns
@@ -212,6 +217,26 @@ mir_build_literal_in_range_out_of_bounds =
     literal out of range for `{$ty}`
     .label = this value does not fit into the type `{$ty}` whose range is `{$min}..={$max}`
 
+mir_build_loop_match_arm_with_guard =
+    match arms that are part of a `#[loop_match]` cannot have guards
+
+mir_build_loop_match_bad_rhs =
+    this expression must be a single `match` wrapped in a labeled block
+
+mir_build_loop_match_bad_statements =
+    statements are not allowed in this position within a `#[loop_match]`
+
+mir_build_loop_match_invalid_match =
+    invalid match on `loop_match` state
+    .note = only matches on local variables are valid
+
+mir_build_loop_match_invalid_update =
+    invalid update of the `loop_match` state
+    .label = the assignment must update this variable
+
+mir_build_loop_match_missing_assignment =
+    expected a single assignment expression
+
 mir_build_lower_range_bound_must_be_less_than_or_equal_to_upper =
     lower range bound must be less than or equal to upper
     .label = lower bound larger than upper bound

compiler/rustc_mir_build/src/builder/expr/as_place.rs

@@ -562,12 +562,14 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             | ExprKind::Match { .. }
             | ExprKind::If { .. }
             | ExprKind::Loop { .. }
+            | ExprKind::LoopMatch { .. }
             | ExprKind::Block { .. }
             | ExprKind::Let { .. }
             | ExprKind::Assign { .. }
             | ExprKind::AssignOp { .. }
             | ExprKind::Break { .. }
             | ExprKind::Continue { .. }
+            | ExprKind::ConstContinue { .. }
             | ExprKind::Return { .. }
             | ExprKind::Become { .. }
             | ExprKind::Literal { .. }

compiler/rustc_mir_build/src/builder/expr/as_rvalue.rs

@@ -538,6 +538,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             | ExprKind::RawBorrow { .. }
             | ExprKind::Adt { .. }
             | ExprKind::Loop { .. }
+            | ExprKind::LoopMatch { .. }
             | ExprKind::LogicalOp { .. }
             | ExprKind::Call { .. }
             | ExprKind::Field { .. }
@@ -548,6 +549,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             | ExprKind::UpvarRef { .. }
             | ExprKind::Break { .. }
             | ExprKind::Continue { .. }
+            | ExprKind::ConstContinue { .. }
             | ExprKind::Return { .. }
             | ExprKind::Become { .. }
             | ExprKind::InlineAsm { .. }

compiler/rustc_mir_build/src/builder/expr/category.rs

@@ -83,9 +83,11 @@ impl Category {
             | ExprKind::NamedConst { .. } => Some(Category::Constant),
 
             ExprKind::Loop { .. }
+            | ExprKind::LoopMatch { .. }
             | ExprKind::Block { .. }
             | ExprKind::Break { .. }
             | ExprKind::Continue { .. }
+            | ExprKind::ConstContinue { .. }
             | ExprKind::Return { .. }
             | ExprKind::Become { .. } =>
             // FIXME(#27840) these probably want their own

compiler/rustc_mir_build/src/builder/expr/into.rs

@@ -15,8 +15,9 @@ use rustc_trait_selection::infer::InferCtxtExt;
 use tracing::{debug, instrument};
 
 use crate::builder::expr::category::{Category, RvalueFunc};
-use crate::builder::matches::DeclareLetBindings;
+use crate::builder::matches::{DeclareLetBindings, HasMatchGuard};
 use crate::builder::{BlockAnd, BlockAndExtension, BlockFrame, Builder, NeedsTemporary};
+use crate::errors::LoopMatchArmWithGuard;
 
 impl<'a, 'tcx> Builder<'a, 'tcx> {
     /// Compile `expr`, storing the result into `destination`, which
@@ -244,6 +245,103 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     None
                 })
             }
+            ExprKind::LoopMatch { state, region_scope, match_span, ref arms } => {
+                // Intuitively, this is a combination of a loop containing a labeled block
+                // containing a match.
+                //
+                // The only new bit here is that the lowering of the match is wrapped in a
+                // `in_const_continuable_scope`, which makes the match arms and their target basic
+                // block available to the lowering of `#[const_continue]`.
+
+                let loop_block = this.cfg.start_new_block();
+
+                // Start the loop.
+                this.cfg.goto(block, source_info, loop_block);
+
+                this.in_breakable_scope(Some(loop_block), destination, expr_span, |this| {
+                    // Logic for `loop`.
+                    let mut body_block = this.cfg.start_new_block();
+                    this.cfg.terminate(
+                        loop_block,
+                        source_info,
+                        TerminatorKind::FalseUnwind {
+                            real_target: body_block,
+                            unwind: UnwindAction::Continue,
+                        },
+                    );
+                    this.diverge_from(loop_block);
+
+                    // Logic for `match`.
+                    let scrutinee_place_builder =
+                        unpack!(body_block = this.as_place_builder(body_block, state));
+                    let scrutinee_span = this.thir.exprs[state].span;
+                    let match_start_span = match_span.shrink_to_lo().to(scrutinee_span);
+
+                    let mut patterns = Vec::with_capacity(arms.len());
+                    for &arm_id in arms.iter() {
+                        let arm = &this.thir[arm_id];
+
+                        if let Some(guard) = arm.guard {
+                            let span = this.thir.exprs[guard].span;
+                            this.tcx.dcx().emit_fatal(LoopMatchArmWithGuard { span })
+                        }
+
+                        patterns.push((&*arm.pattern, HasMatchGuard::No));
+                    }
+
+                    // The `built_tree` maps match arms to their basic block (where control flow
+                    // jumps to when a value matches the arm). This structure is stored so that a
+                    // #[const_continue] can figure out what basic block to jump to.
+                    let built_tree = this.lower_match_tree(
+                        body_block,
+                        scrutinee_span,
+                        &scrutinee_place_builder,
+                        match_start_span,
+                        patterns,
+                        false,
+                    );
+
+                    let state_place = scrutinee_place_builder.to_place(this);
+
+                    // This is logic for the labeled block: a block is a drop scope, hence
+                    // `in_scope`, and a labeled block can be broken out of with a `break 'label`,
+                    // hence the `in_breakable_scope`.
+                    //
+                    // Then `in_const_continuable_scope` stores information for the lowering of
+                    // #[const_continue], and finally the match is lowered in the standard way.
+                    unpack!(
+                        body_block = this.in_scope(
+                            (region_scope, source_info),
+                            LintLevel::Inherited,
+                            move |this| {
+                                this.in_breakable_scope(None, state_place, expr_span, |this| {
+                                    Some(this.in_const_continuable_scope(
+                                        arms.clone(),
+                                        built_tree.clone(),
+                                        state_place,
+                                        expr_span,
+                                        |this| {
+                                            this.lower_match_arms(
+                                                destination,
+                                                scrutinee_place_builder,
+                                                scrutinee_span,
+                                                arms,
+                                                built_tree,
+                                                this.source_info(match_span),
+                                            )
+                                        },
+                                    ))
+                                })
+                            }
+                        )
+                    );
+
+                    this.cfg.goto(body_block, source_info, loop_block);
+
+                    // Loops are only exited by `break` expressions.
+                    None
+                })
+            }
             ExprKind::Call { ty: _, fun, ref args, from_hir_call, fn_span } => {
                 let fun = unpack!(block = this.as_local_operand(block, fun));
                 let args: Box<[_]> = args
@@ -601,6 +699,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
 
             ExprKind::Continue { .. }
+            | ExprKind::ConstContinue { .. }
             | ExprKind::Break { .. }
             | ExprKind::Return { .. }
             | ExprKind::Become { .. } => {

compiler/rustc_mir_build/src/builder/expr/stmt.rs

@@ -92,6 +92,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             ExprKind::Break { label, value } => {
                 this.break_scope(block, value, BreakableTarget::Break(label), source_info)
             }
+            ExprKind::ConstContinue { label, value } => {
+                this.break_const_continuable_scope(block, value, label, source_info)
+            }
             ExprKind::Return { value } => {
                 this.break_scope(block, value, BreakableTarget::Return, source_info)
             }