[WIP] Support const blocks in const_continue

Author: bjorn3

compiler/rustc_mir_build/src/builder/matches/mod.rs

@@ -11,15 +11,16 @@ use std::mem;
 use std::sync::Arc;
 
 use rustc_abi::VariantIdx;
-use rustc_ast::LitKind;
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_hir::{BindingMode, ByRef, LetStmt, LocalSource, Node};
-use rustc_middle::bug;
 use rustc_middle::middle::region;
 use rustc_middle::mir::{self, *};
 use rustc_middle::thir::{self, *};
-use rustc_middle::ty::{self, CanonicalUserTypeAnnotation, Ty};
+use rustc_middle::ty::{
+    self, CanonicalUserTypeAnnotation, Ty, TypeVisitableExt, ValTree, ValTreeKind,
+};
+use rustc_middle::{bug, span_bug};
 use rustc_pattern_analysis::rustc::{DeconstructedPat, RustcPatCtxt};
 use rustc_span::{BytePos, Pos, Span, Symbol, sym};
 use tracing::{debug, instrument};
@@ -2870,7 +2871,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     pub(crate) fn static_pattern_match(
         &self,
         cx: &RustcPatCtxt<'_, 'tcx>,
-        value: ExprId,
+        constant: ConstOperand<'tcx>,
         arms: &[ArmId],
         built_match_tree: &BuiltMatchTree<'tcx>,
     ) -> Option<BasicBlock> {
@@ -2889,56 +2890,81 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         .or_else(|| branch.sub_branches.last())
                         .unwrap();
 
-                    match self.static_pattern_match_help(value, &pat.pat) {
+                    match self.static_pattern_match_help(constant, &pat.pat) {
                         true => return Some(sub_branch.success_block),
                         false => continue,
                     }
                 }
-            } else if self.static_pattern_match_help(value, &pat) {
+            } else if self.static_pattern_match_help(constant, &pat) {
                 return Some(branch.sub_branches[0].success_block);
             }
         }
 
         None
     }
 
-    fn static_pattern_match_help(&self, value: ExprId, pat: &DeconstructedPat<'_, 'tcx>) -> bool {
-        use rustc_middle::thir::ExprKind;
+    fn static_pattern_match_help(
+        &self,
+        constant: ConstOperand<'tcx>,
+        pat: &DeconstructedPat<'_, 'tcx>,
+    ) -> bool {
         use rustc_pattern_analysis::constructor::{IntRange, MaybeInfiniteInt};
         use rustc_pattern_analysis::rustc::Constructor;
 
-        match pat.ctor() {
-            Constructor::Variant(variant_index) => match &self.thir[value].kind {
-                ExprKind::Adt(value_adt) => {
-                    return *variant_index == value_adt.variant_index;
+        // Based on eval_unevaluated_mir_constant_to_valtree
+        let (valtree, ty) = 'a: {
+            assert!(!constant.const_.ty().has_param());
+            let (uv, ty) = match constant.const_ {
+                mir::Const::Unevaluated(uv, ty) => (uv.shrink(), ty),
+                mir::Const::Ty(_, c) => match c.kind() {
+                    // A constant that came from a const generic but was then used as an argument to
+                    // old-style simd_shuffle (passing as argument instead of as a generic param).
+                    ty::ConstKind::Value(cv) => break 'a (cv.valtree, cv.ty),
+                    other => span_bug!(constant.span, "{other:#?}"),
+                },
+                mir::Const::Val(mir::ConstValue::Scalar(mir::interpret::Scalar::Int(val)), ty) => {
+                    break 'a (ValTree::from_scalar_int(self.tcx, val), ty);
                 }
-                other => todo!("{other:?}"),
-            },
-            Constructor::IntRange(int_range) => match &self.thir[value].kind {
-                ExprKind::Literal { lit, neg } => match &lit.node {
-                    LitKind::Int(n, _) => {
-                        let n = if pat.ty().is_signed() {
-                            let size = pat.ty().primitive_size(self.tcx);
-                            MaybeInfiniteInt::new_finite_int(
-                                if *neg {
-                                    size.truncate((n.get() as i128).overflowing_neg().0 as u128)
-                                } else {
-                                    n.get()
-                                },
-                                size.bits(),
-                            )
-                        } else {
-                            MaybeInfiniteInt::new_finite_uint(n.get())
-                        };
-
-                        return IntRange::from_singleton(n).is_subrange(int_range);
-                    }
+                // We should never encounter `Const::Val` unless MIR opts (like const prop) evaluate
+                // a constant and write that value back into `Operand`s. This could happen, but is
+                // unlikely. Also: all users of `simd_shuffle` are on unstable and already need to take
+                // a lot of care around intrinsics. For an issue to happen here, it would require a
+                // macro expanding to a `simd_shuffle` call without wrapping the constant argument in a
+                // `const {}` block, but the user pass through arbitrary expressions.
+                // FIXME(oli-obk): replace the magic const generic argument of `simd_shuffle` with a
+                // real const generic, and get rid of this entire function.
+                other => span_bug!(constant.span, "{other:#?}"),
+            };
+            (
+                self.tcx
+                    .const_eval_resolve_for_typeck(self.typing_env(), uv, constant.span)
+                    .unwrap()
+                    .unwrap(),
+                ty,
+            )
+        };
+        assert!(!ty.has_param());
 
-                    other => todo!("{other:?}"),
-                },
+        match pat.ctor() {
+            Constructor::Variant(variant_index) => match *valtree {
+                ValTreeKind::Branch(box [actual_variant_idx]) => {
+                    *variant_index
+                        == VariantIdx::from_u32(actual_variant_idx.unwrap_leaf().to_u32())
+                }
                 other => todo!("{other:?}"),
             },
-            Constructor::Wildcard => return true,
+            Constructor::IntRange(int_range) => {
+                let size = pat.ty().primitive_size(self.tcx);
+                let actual_int = valtree.unwrap_leaf().to_bits(size);
+                let actual_int = if pat.ty().is_signed() {
+                    MaybeInfiniteInt::new_finite_int(actual_int, size.bits())
+                } else {
+                    MaybeInfiniteInt::new_finite_uint(actual_int)
+                };
+                IntRange::from_singleton(actual_int).is_subrange(int_range)
+            }
+            Constructor::Wildcard => true,
+            // FIXME error out before static_pattern_match gets run and replace this with bug!()
             _ => false,
         }
     }

compiler/rustc_mir_build/src/builder/scope.rs

@@ -827,6 +827,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             span_bug!(span, "break value must be a scope")
         };
 
+        // FIXME accept bare MyEnum::Foo as constant
+        let constant = self.as_constant(&self.thir[value]);
+
         let break_index = self
             .scopes
             .const_continuable_scopes
@@ -868,7 +871,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         };
 
         let Some(real_target) =
-            self.static_pattern_match(&cx, value, &*scope.arms, &scope.built_match_tree)
+            self.static_pattern_match(&cx, constant, &*scope.arms, &scope.built_match_tree)
         else {
             self.tcx.dcx().emit_fatal(ConstContinueUnknownJumpTarget { span })
         };

tests/ui/feature-gates/feature-gate-loop-match.rs

@@ -13,12 +13,12 @@ fn main() {
                 State::A => {
                     #[const_continue]
                     //~^ ERROR the `#[const_continue]` attribute is an experimental feature
-                    break 'blk State::B;
+                    break 'blk const { State::B };
                 }
                 State::B => {
                     #[const_continue]
                     //~^ ERROR the `#[const_continue]` attribute is an experimental feature
-                    break 'blk State::C;
+                    break 'blk const { State::C };
                 }
                 State::C => break 'a,
             }

tests/ui/loop-match/loop-match.rs

@@ -19,7 +19,7 @@ fn main() {
             match state {
                 State::A => {
                     #[const_continue]
-                    break 'blk State::B;
+                    break 'blk const { State::B };
                 }
                 State::B => {
                     // without special logic, the compiler believes this is a re-assignment to
@@ -29,10 +29,10 @@ fn main() {
 
                     if true {
                         #[const_continue]
-                        break 'blk State::C;
+                        break 'blk const { State::C };
                     } else {
                         #[const_continue]
-                        break 'blk State::A;
+                        break 'blk const { State::A };
                     }
                 }
                 State::C => break 'a,

tests/ui/loop-match/nested.rs

@@ -37,7 +37,7 @@ fn run() -> String {
                 State1::A => {
                     accum.push('a');
                     #[const_continue]
-                    break 'blk1 State1::B;
+                    break 'blk1 const { State1::B };
                 }
                 State1::B => {
                     accum.push('b');
@@ -48,22 +48,22 @@ fn run() -> String {
                                 State2::X => {
                                     accum.push('x');
                                     #[const_continue]
-                                    break 'blk2 State2::Y;
+                                    break 'blk2 const { State2::Y };
                                 }
                                 State2::Y => {
                                     accum.push('y');
                                     #[const_continue]
-                                    break 'blk2 State2::Z;
+                                    break 'blk2 const { State2::Z };
                                 }
                                 State2::Z => {
                                     accum.push('z');
                                     if first {
                                         first = false;
                                         #[const_continue]
-                                        break 'blk2 State2::X;
+                                        break 'blk2 const { State2::X };
                                     } else {
                                         #[const_continue]
-                                        break 'blk1 State1::C;
+                                        break 'blk1 const { State1::C };
                                     }
                                 }
                             }

tests/ui/loop-match/or-patterns.rs

@@ -25,21 +25,21 @@ fn main() {
                     states.push(state);
                     if first {
                         #[const_continue]
-                        break 'blk State::B;
+                        break 'blk const { State::B };
                     } else {
                         #[const_continue]
-                        break 'blk State::D;
+                        break 'blk const { State::D };
                     }
                 }
                 State::B | State::D => {
                     states.push(state);
                     if first {
                         first = false;
                         #[const_continue]
-                        break 'blk State::A;
+                        break 'blk const { State::A };
                     } else {
                         #[const_continue]
-                        break 'blk State::C;
+                        break 'blk const { State::C };
                     }
                 }
                 State::C => {