Add #[loop_match] for improved DFA codegen

Co-authored-by: Folkert de Vries <folkert@folkertdev.nl>

Author: bjorn3

compiler/rustc_feature/src/builtin_attrs.rs

@@ -634,6 +634,19 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         EncodeCrossCrate::Yes, min_generic_const_args, experimental!(type_const),
     ),
 
+    // The `#[loop_match]` and `#[const_continue]` attributes are part of the
+    // lang experiment for RFC 3720 tracked in:
+    //
+    // - https://github.com/rust-lang/rust/issues/132306
+    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

@@ -557,6 +557,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(132306)),
     /// 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_hir_typeck/messages.ftl

@@ -79,6 +79,9 @@ hir_typeck_cast_unknown_pointer = cannot cast {$to ->
     .note = the type information given here is insufficient to check whether the pointer cast is valid
     .label_from = the type information given here is insufficient to check whether the pointer cast is valid
 
+hir_typeck_const_continue_bad_label =
+    `#[const_continue]` must break to a labeled block that participates in a `#[loop_match]`
+
 hir_typeck_const_select_must_be_const = this argument must be a `const fn`
     .help = consult the documentation on `const_eval_select` for more information
 

compiler/rustc_hir_typeck/src/errors.rs

@@ -1170,3 +1170,10 @@ pub(crate) struct AbiCustomCall {
     #[primary_span]
     pub span: Span,
 }
+
+#[derive(Diagnostic)]
+#[diag(hir_typeck_const_continue_bad_label)]
+pub(crate) struct ConstContinueBadLabel {
+    #[primary_span]
+    pub span: Span,
+}

compiler/rustc_hir_typeck/src/loops.rs

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 use std::fmt;
 
 use Context::*;
+use rustc_ast::Label;
 use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::LocalDefId;
@@ -11,11 +12,12 @@ use rustc_middle::hir::nested_filter;
 use rustc_middle::span_bug;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::hygiene::DesugaringKind;
-use rustc_span::{BytePos, Span};
+use rustc_span::{BytePos, Span, sym};
 
 use crate::errors::{
-    BreakInsideClosure, BreakInsideCoroutine, BreakNonLoop, ContinueLabeledBlock, OutsideLoop,
-    OutsideLoopSuggestion, UnlabeledCfInWhileCondition, UnlabeledInLabeledBlock,
+    BreakInsideClosure, BreakInsideCoroutine, BreakNonLoop, ConstContinueBadLabel,
+    ContinueLabeledBlock, OutsideLoop, OutsideLoopSuggestion, UnlabeledCfInWhileCondition,
+    UnlabeledInLabeledBlock,
 };
 
 /// The context in which a block is encountered.
@@ -37,6 +39,11 @@ enum Context {
     AnonConst,
     /// E.g. `const { ... }`.
     ConstBlock,
+    /// E.g. `#[loop_match] loop { state = 'label: { /* ... */ } }`.
+    LoopMatch {
+        /// The label of the labeled block (not of the loop itself).
+        labeled_block: Label,
+    },
 }
 
 #[derive(Clone)]
@@ -141,7 +148,12 @@ impl<'hir> Visitor<'hir> for CheckLoopVisitor<'hir> {
                 }
             }
             hir::ExprKind::Loop(ref b, _, source, _) => {
-                self.with_context(Loop(source), |v| v.visit_block(b));
+                let cx = match self.is_loop_match(e, b) {
+                    Some(labeled_block) => LoopMatch { labeled_block },
+                    None => Loop(source),
+                };
+
+                self.with_context(cx, |v| v.visit_block(b));
             }
             hir::ExprKind::Closure(&hir::Closure {
                 ref fn_decl, body, fn_decl_span, kind, ..
@@ -197,6 +209,24 @@ impl<'hir> Visitor<'hir> for CheckLoopVisitor<'hir> {
                     Err(hir::LoopIdError::UnresolvedLabel) => None,
                 };
 
+                // A `#[const_continue]` must break to a block in a `#[loop_match]`.
+                let attrs = self.tcx.hir_attrs(e.hir_id);
+                if attrs.iter().any(|attr| attr.has_name(sym::const_continue)) {
+                    if let Some(break_label) = break_label.label {
+                        let is_target_label = |cx: &Context| match cx {
+                            Context::LoopMatch { labeled_block } => {
+                                break_label.ident.name == labeled_block.ident.name
+                            }
+                            _ => false,
+                        };
+
+                        if !self.cx_stack.iter().rev().any(is_target_label) {
+                            let span = break_label.ident.span;
+                            self.tcx.dcx().emit_fatal(ConstContinueBadLabel { span });
+                        }
+                    }
+                }
+
                 if let Some(Node::Block(_)) = loop_id.map(|id| self.tcx.hir_node(id)) {
                     return;
                 }
@@ -299,7 +329,7 @@ impl<'hir> CheckLoopVisitor<'hir> {
         cx_pos: usize,
     ) {
         match self.cx_stack[cx_pos] {
-            LabeledBlock | Loop(_) => {}
+            LabeledBlock | Loop(_) | LoopMatch { .. } => {}
             Closure(closure_span) => {
                 self.tcx.dcx().emit_err(BreakInsideClosure {
                     span,
@@ -380,4 +410,36 @@ impl<'hir> CheckLoopVisitor<'hir> {
             });
         }
     }
+
+    /// Is this a loop annotated with `#[loop_match]` that looks syntactically sound?
+    fn is_loop_match(
+        &self,
+        e: &'hir hir::Expr<'hir>,
+        body: &'hir hir::Block<'hir>,
+    ) -> Option<Label> {
+        if !self.tcx.hir_attrs(e.hir_id).iter().any(|attr| attr.has_name(sym::loop_match)) {
+            return None;
+        }
+
+        // NOTE: Diagnostics are emitted during MIR construction.
+
+        // Accept either `state = expr` or `state = expr;`.
+        let loop_body_expr = match body.stmts {
+            [] => match body.expr {
+                Some(expr) => expr,
+                None => return None,
+            },
+            [single] if body.expr.is_none() => match single.kind {
+                hir::StmtKind::Expr(expr) | hir::StmtKind::Semi(expr) => expr,
+                _ => return None,
+            },
+            [..] => return None,
+        };
+
+        let hir::ExprKind::Assign(_, rhs_expr, _) = loop_body_expr.kind else { return None };
+
+        let hir::ExprKind::Block(_, label) = rhs_expr.kind else { return None };
+
+        label
+    }
 }

compiler/rustc_middle/src/thir.rs

@@ -378,6 +378,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 `&&`).
     ///
@@ -454,6 +462,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

@@ -83,7 +83,7 @@ pub fn walk_expr<'thir, 'tcx: 'thir, V: Visitor<'thir, 'tcx>>(
             visitor.visit_pat(pat);
         }
         Loop { body } => visitor.visit_expr(&visitor.thir()[body]),
-        Match { scrutinee, ref arms, .. } => {
+        LoopMatch { state: scrutinee, ref arms, .. } | Match { scrutinee, ref arms, .. } => {
             visitor.visit_expr(&visitor.thir()[scrutinee]);
             for &arm in &**arms {
                 visitor.visit_arm(&visitor.thir()[arm]);
@@ -108,6 +108,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,15 @@ 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_bad_const = could not determine the target branch for this `#[const_continue]`
+    .label = this value is too generic
+    .note = the value must be a literal or a monomorphic const
+
+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
+    .label = this value must be a literal or a monomorphic const
+
 mir_build_const_defined_here = constant defined here
 
 mir_build_const_param_in_pattern = constant parameters cannot be referenced in patterns
@@ -212,6 +221,30 @@ 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 = a local variable must be the scrutinee within a `#[loop_match]`
+
+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_loop_match_unsupported_type =
+    this `#[loop_match]` state value has type `{$ty}`, which is not supported
+    .note = only integers, floats, bool, char, and enums without fields are supported
+
 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

@@ -565,12 +565,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 { .. }