Add `#[loop_match]` for improved DFA codegen #138780

folkertdev · 2025-03-21T11:40:03Z

tracking issue: #138777
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

#[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
report proper errors when the jump target could not be determined statically
rework the pattern matching logic so hopefully it can use more existing code

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

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

rustbot · 2025-03-21T11:40:10Z

Some changes occurred in match checking

cc @Nadrieril

Some changes occurred in compiler/rustc_passes/src/check_attr.rs

cc @jdonszelmann

Some changes occurred in rustc_ty_utils::consts.rs

cc @BoxyUwU

traviscross

Thanks @folkertdev for putting up this PR. The big picture looks right, in terms of the behavior of the tests and how to approach the experiment in terms of starting with the attributes for thiis.

This is a first partial pass on the details.

@rustbot author

compiler/rustc_feature/src/unstable.rs

compiler/rustc_mir_build/messages.ftl

tests/ui/loop-match/break-to-block.rs

traviscross · 2025-03-21T23:43:58Z

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

@@ -244,6 +248,188 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                    None
                })
            }
+            ExprKind::LoopMatch { state, region_scope, ref arms, .. } => {


Probably the body of this match arm should be broken out and located elsewhere, in the same way that match_expr is, and should have a substantial doc comment above it, like match_expr does.

traviscross · 2025-03-21T23:53:05Z

compiler/rustc_mir_build/src/builder/scope.rs

+                let rustc_middle::thir::ExprKind::Scope { value, .. } =
+                    self.thir[value.unwrap()].kind
+                else {
+                    panic!();
+                };


If this is unreachable, it's worth marking it that way. Either way, it's worth a comment about this.

traviscross · 2025-03-21T23:53:56Z

compiler/rustc_mir_build/src/builder/scope.rs

+                    //let is_coroutine = this.coroutine.is_some();
+
+                    /*// Link the exit drop tree to unwind drop tree.
+                    if drops.drops.iter().any(|drop_node| drop_node.data.kind == DropKind::Value) {
+                        let unwind_target = this.diverge_cleanup_target(region_scope, span);
+                        let mut unwind_indices = IndexVec::from_elem_n(unwind_target, 1);
+                        for (drop_idx, drop_node) in drops.drops.iter_enumerated().skip(1) {
+                            match drop_node.data.kind {
+                                DropKind::Storage | DropKind::ForLint => {
+                                    if is_coroutine {
+                                        let unwind_drop = this.scopes.unwind_drops.add_drop(
+                                            drop_node.data,
+                                            unwind_indices[drop_node.next],
+                                        );
+                                        unwind_indices.push(unwind_drop);
+                                    } else {
+                                        unwind_indices.push(unwind_indices[drop_node.next]);
+                                    }
+                                }
+                                DropKind::Value => {
+                                    let unwind_drop = this
+                                        .scopes
+                                        .unwind_drops
+                                        .add_drop(drop_node.data, unwind_indices[drop_node.next]);
+                                    this.scopes.unwind_drops.add_entry_point(
+                                        blocks[drop_idx].unwrap(),
+                                        unwind_indices[drop_node.next],
+                                    );
+                                    unwind_indices.push(unwind_drop);
+                                }
+                            }
+                        }
+                    }*/


What's the story here?

Dropping the rest of the local variables when unwinding out of drops on #[const_continue] isn't implemented yet.

compiler/rustc_passes/messages.ftl

traviscross · 2025-03-22T00:00:58Z

compiler/rustc_middle/src/thir.rs

+    /// A `#[loop_match] loop { state = 'blk: { match state { ... } } }` expression.
+    LoopMatch {
+        state: ExprId,
+
+        region_scope: region::Scope,
+        arms: Box<[ArmId]>,
+    },


Please add the relevant pretty test for this and the other to ensure these get pretty printed correctly.

what sort of a test should I add for this? would //@ pretty-mode:hir,typed exercise the relevant logic (we only add this node for thir, not for hir).

We don't have pretty printing support for THIR. Only -Zunpretty=thir-tree and -Zunpretty=thir-flat exist, which both use Debug impls.

Co-authored-by: Travis Cross <tc@traviscross.com>

folkertdev

Thanks for the detailed review!

I've fixed a bunch of the low-hanging fruit (e.g. in the tests). For the actual pattern matching logic, I have a branch with what I believe is a better solution that re-uses more existing pattern matching infra. We'll come back to that here once björn has had a chance to look at it.

folkertdev · 2025-03-22T12:15:39Z

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

+                // FIXME do we need the breakable scope?
+                this.in_breakable_scope(Some(loop_block), destination, expr_span, |this| {


that breakable scope is definitely needed (to break out of the loop)

folkertdev · 2025-03-22T12:16:38Z

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

+
+                // FIXME do we need the breakable scope?
+                this.in_breakable_scope(Some(loop_block), destination, expr_span, |this| {
+                    // conduct the test, if necessary


Noted. In this case I think the comment is not really in the best place.

folkertdev · 2025-03-22T12:34:10Z

compiler/rustc_middle/src/thir.rs

+    /// A `#[loop_match] loop { state = 'blk: { match state { ... } } }` expression.
+    LoopMatch {
+        state: ExprId,
+
+        region_scope: region::Scope,
+        arms: Box<[ArmId]>,
+    },


what sort of a test should I add for this? would //@ pretty-mode:hir,typed exercise the relevant logic (we only add this node for thir, not for hir).

rustbot · 2025-03-24T10:46:56Z

Some changes occurred in exhaustiveness checking

cc @Nadrieril

Some changes occurred in match lowering

cc @Nadrieril

Add #[loop_match] for improved DFA codegen

8676f7c

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

rustbot assigned traviscross Mar 21, 2025

rustbot added A-attributes Area: Attributes (`#[…]`, `#![…]`) S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Mar 21, 2025

traviscross reviewed Mar 22, 2025

View reviewed changes

rustbot added S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Mar 22, 2025

Apply suggestions from code review

a6179ae

Co-authored-by: Travis Cross <tc@traviscross.com>

This comment has been minimized.

Sign in to view

add comments to tests

2793d38

folkertdev commented Mar 22, 2025

View reviewed changes

folkertdev added 2 commits March 24, 2025 11:46

remove some comments that are now inaccurate

cbc478d

static pattern matching when lowering #[const_continue]

90ef42f

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `#[loop_match]` for improved DFA codegen #138780

Add `#[loop_match]` for improved DFA codegen #138780

folkertdev commented Mar 21, 2025

rustbot commented Mar 21, 2025

traviscross left a comment

traviscross Mar 21, 2025

traviscross Mar 21, 2025

traviscross Mar 21, 2025

bjorn3 Mar 22, 2025 •

edited

Loading

traviscross Mar 22, 2025

folkertdev Mar 22, 2025

bjorn3 Mar 24, 2025

This comment has been minimized.

folkertdev left a comment

folkertdev Mar 22, 2025

folkertdev Mar 22, 2025

folkertdev Mar 22, 2025

rustbot commented Mar 24, 2025

		// FIXME do we need the breakable scope?
		this.in_breakable_scope(Some(loop_block), destination, expr_span, \|this\| {

Add #[loop_match] for improved DFA codegen #138780

Are you sure you want to change the base?

Add #[loop_match] for improved DFA codegen #138780

Conversation

folkertdev commented Mar 21, 2025

current state

future work

maybe future work

rustbot commented Mar 21, 2025

traviscross left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bjorn3 Mar 22, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

This comment has been minimized.

folkertdev left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rustbot commented Mar 24, 2025

Add `#[loop_match]` for improved DFA codegen #138780

Add `#[loop_match]` for improved DFA codegen #138780

bjorn3 Mar 22, 2025 •

edited

Loading