Add #[loop_match] for improved DFA codegen

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

Cargo.lock

@@ -4117,6 +4117,7 @@ dependencies = [
  "rustc_apfloat",
  "rustc_arena",
  "rustc_ast",
+ "rustc_attr_data_structures",
  "rustc_data_structures",
  "rustc_errors",
  "rustc_fluent_macro",

compiler/rustc_attr_data_structures/src/attributes.rs

@@ -212,6 +212,9 @@ pub enum AttributeKind {
         first_span: Span,
     },
 
+    /// Represents `#[const_continue]`.
+    ConstContinue(Span),
+
     /// Represents `#[rustc_const_stable]` and `#[rustc_const_unstable]`.
     ConstStability {
         stability: PartialConstStability,
@@ -231,6 +234,9 @@ pub enum AttributeKind {
     /// Represents `#[inline]` and `#[rustc_force_inline]`.
     Inline(InlineAttr, Span),
 
+    /// Represents `#[loop_match]`.
+    LoopMatch(Span),
+
     /// Represents `#[rustc_macro_transparency]`.
     MacroTransparency(Transparency),
 

compiler/rustc_attr_parsing/src/attributes/loop_match.rs

@@ -0,0 +1,31 @@
+use rustc_attr_data_structures::AttributeKind;
+use rustc_feature::{AttributeTemplate, template};
+use rustc_span::{Symbol, sym};
+
+use crate::attributes::{AttributeOrder, OnDuplicate, SingleAttributeParser};
+use crate::context::{AcceptContext, Stage};
+use crate::parser::ArgParser;
+
+pub(crate) struct LoopMatchParser;
+impl<S: Stage> SingleAttributeParser<S> for LoopMatchParser {
+    const PATH: &[Symbol] = &[sym::loop_match];
+    const ATTRIBUTE_ORDER: AttributeOrder = AttributeOrder::KeepFirst;
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Warn;
+    const TEMPLATE: AttributeTemplate = template!(Word);
+
+    fn convert(cx: &mut AcceptContext<'_, '_, S>, _args: &ArgParser<'_>) -> Option<AttributeKind> {
+        Some(AttributeKind::LoopMatch(cx.attr_span))
+    }
+}
+
+pub(crate) struct ConstContinueParser;
+impl<S: Stage> SingleAttributeParser<S> for ConstContinueParser {
+    const PATH: &[Symbol] = &[sym::const_continue];
+    const ATTRIBUTE_ORDER: AttributeOrder = AttributeOrder::KeepFirst;
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Warn;
+    const TEMPLATE: AttributeTemplate = template!(Word);
+
+    fn convert(cx: &mut AcceptContext<'_, '_, S>, _args: &ArgParser<'_>) -> Option<AttributeKind> {
+        Some(AttributeKind::ConstContinue(cx.attr_span))
+    }
+}

compiler/rustc_attr_parsing/src/attributes/mod.rs

@@ -32,6 +32,7 @@ pub(crate) mod confusables;
 pub(crate) mod deprecation;
 pub(crate) mod inline;
 pub(crate) mod lint_helpers;
+pub(crate) mod loop_match;
 pub(crate) mod must_use;
 pub(crate) mod repr;
 pub(crate) mod semantics;

compiler/rustc_attr_parsing/src/context.rs

@@ -20,6 +20,7 @@ use crate::attributes::confusables::ConfusablesParser;
 use crate::attributes::deprecation::DeprecationParser;
 use crate::attributes::inline::{InlineParser, RustcForceInlineParser};
 use crate::attributes::lint_helpers::{AsPtrParser, PubTransparentParser};
+use crate::attributes::loop_match::{ConstContinueParser, LoopMatchParser};
 use crate::attributes::must_use::MustUseParser;
 use crate::attributes::repr::{AlignParser, ReprParser};
 use crate::attributes::semantics::MayDangleParser;
@@ -110,9 +111,11 @@ attribute_parsers!(
         // tidy-alphabetical-start
         Single<AsPtrParser>,
         Single<ColdParser>,
+        Single<ConstContinueParser>,
         Single<ConstStabilityIndirectParser>,
         Single<DeprecationParser>,
         Single<InlineParser>,
+        Single<LoopMatchParser>,
         Single<MayDangleParser>,
         Single<MustUseParser>,
         Single<NoMangleParser>,

compiler/rustc_feature/src/builtin_attrs.rs

@@ -657,6 +657,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

@@ -1167,3 +1167,10 @@ pub(crate) struct AbiCannotBeCalled {
     pub span: Span,
     pub abi: ExternAbi,
 }
+
+#[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,8 @@ use std::collections::BTreeMap;
 use std::fmt;
 
 use Context::*;
+use rustc_ast::Label;
+use rustc_attr_data_structures::{AttributeKind, find_attr};
 use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::LocalDefId;
@@ -14,8 +16,9 @@ use rustc_span::hygiene::DesugaringKind;
 use rustc_span::{BytePos, Span};
 
 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 +40,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 +149,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 +210,23 @@ impl<'hir> Visitor<'hir> for CheckLoopVisitor<'hir> {
                     Err(hir::LoopIdError::UnresolvedLabel) => None,
                 };
 
+                // A `#[const_continue]` must break to a block in a `#[loop_match]`.
+                if find_attr!(self.tcx.hir_attrs(e.hir_id), AttributeKind::ConstContinue(_)) {
+                    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 !find_attr!(self.tcx.hir_attrs(e.hir_id), AttributeKind::LoopMatch(_)) {
+            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/Cargo.toml

@@ -11,6 +11,7 @@ rustc_abi = { path = "../rustc_abi" }
 rustc_apfloat = "0.2.0"
 rustc_arena = { path = "../rustc_arena" }
 rustc_ast = { path = "../rustc_ast" }
+rustc_attr_data_structures = { path = "../rustc_attr_data_structures" }
 rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_errors = { path = "../rustc_errors" }
 rustc_fluent_macro = { path = "../rustc_fluent_macro" }

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 { .. }

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