Rollup merge of #138780 - trifectatechfoundation:loop_match_attr, r=oli-obk,traviscross

Add `#[loop_match]` for improved DFA codegen

tracking issue: #132306
project goal: rust-lang/rust-project-goals#258

This PR adds the `#[loop_match]` attribute, which aims to improve code generation for state machines. For some (very exciting) benchmarks, see rust-lang/rust-project-goals#258 (comment)

Currently, a very restricted syntax pattern is accepted. We'd like to get feedback and merge this now before we go too far in a direction that others have concerns with.

## current state

We accept code that looks like this

```rust
#[loop_match]
loop {
    state = 'blk: {
        match state {
            State::A => {
                #[const_continue]
                break 'blk State::B
            }
            State::B => { /* ... */ }
            /* ... */
        }
    }
}
```

- a loop should have the same semantics with and without `#[loop_match]`: normal `continue` and `break` continue to work
- `#[const_continue]` is only allowed in loops annotated with `#[loop_match]`
- the loop body needs to have this particular shape (a single assignment to the match scrutinee, with the body a labelled block containing just a match)

## future work

- perform const evaluation on the `break` value
- support more state/scrutinee types

## maybe future work

- allow `continue 'label value` syntax, which `#[const_continue]` could then use.
- allow the match to be on an arbitrary expression (e.g. `State::Initial`)
- attempt to also optimize `break`/`continue` expressions that are not marked with `#[const_continue]`

r? ``@traviscross``

Author: workingjubilee

Cargo.lock

@@ -4116,6 +4116,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;
@@ -111,9 +112,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
+    }
 }