Rollup merge of #119668 - cjgillot:transform-promote, r=oli-obk

Simplify implementation of MIR promotion

Non-functional changes.
Best read ignoring whitespace.

compiler/rustc_const_eval/src/transform/mod.rs

@@ -1,3 +1,2 @@
 pub mod check_consts;
-pub mod promote_consts;
 pub mod validate;

compiler/rustc_mir_transform/src/lib.rs

@@ -1,5 +1,6 @@
 #![deny(rustc::untranslatable_diagnostic)]
 #![deny(rustc::diagnostic_outside_of_impl)]
+#![feature(assert_matches)]
 #![feature(box_patterns)]
 #![feature(cow_is_borrowed)]
 #![feature(decl_macro)]
@@ -94,6 +95,7 @@ mod multiple_return_terminators;
 mod normalize_array_len;
 mod nrvo;
 mod prettify;
+mod promote_consts;
 mod ref_prop;
 mod remove_noop_landing_pads;
 mod remove_storage_markers;
@@ -115,7 +117,6 @@ mod uninhabited_enum_branching;
 mod unreachable_prop;
 
 use rustc_const_eval::transform::check_consts::{self, ConstCx};
-use rustc_const_eval::transform::promote_consts;
 use rustc_const_eval::transform::validate;
 use rustc_mir_dataflow::rustc_peek;
 

compiler/rustc_mir_transform/src/promote_consts.rs

@@ -12,6 +12,7 @@
 //! initialization and can otherwise silence errors, if
 //! move analysis runs after promotion on broken MIR.
 
+use either::{Left, Right};
 use rustc_hir as hir;
 use rustc_middle::mir;
 use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
@@ -22,10 +23,11 @@ use rustc_span::Span;
 
 use rustc_index::{Idx, IndexSlice, IndexVec};
 
+use std::assert_matches::assert_matches;
 use std::cell::Cell;
 use std::{cmp, iter, mem};
 
-use crate::transform::check_consts::{qualifs, ConstCx};
+use rustc_const_eval::transform::check_consts::{qualifs, ConstCx};
 
 /// A `MirPass` for promotion.
 ///
@@ -64,7 +66,7 @@ impl<'tcx> MirPass<'tcx> for PromoteTemps<'tcx> {
 
 /// State of a temporary during collection and promotion.
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
-pub enum TempState {
+enum TempState {
     /// No references to this temp.
     Undefined,
     /// One direct assignment and any number of direct uses.
@@ -78,18 +80,11 @@ pub enum TempState {
     PromotedOut,
 }
 
-impl TempState {
-    pub fn is_promotable(&self) -> bool {
-        debug!("is_promotable: self={:?}", self);
-        matches!(self, TempState::Defined { .. })
-    }
-}
-
 /// A "root candidate" for promotion, which will become the
 /// returned value in a promoted MIR, unless it's a subset
 /// of a larger candidate.
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
-pub struct Candidate {
+struct Candidate {
     location: Location,
 }
 
@@ -123,16 +118,15 @@ impl<'tcx> Visitor<'tcx> for Collector<'_, 'tcx> {
 
         let temp = &mut self.temps[index];
         debug!("visit_local: temp={:?}", temp);
-        if *temp == TempState::Undefined {
+        *temp = match *temp {
-            match context {
+            TempState::Undefined => match context {
                 PlaceContext::MutatingUse(MutatingUseContext::Store)
                 | PlaceContext::MutatingUse(MutatingUseContext::Call) => {
-                    *temp = TempState::Defined { location, uses: 0, valid: Err(()) };
+                    TempState::Defined { location, uses: 0, valid: Err(()) }
-                    return;
                 }
-                _ => { /* mark as unpromotable below */ }
+                _ => TempState::Unpromotable,
-            }
+            },
-        } else if let TempState::Defined { uses, .. } = temp {
+            TempState::Defined { ref mut uses, .. } => {
                 // We always allow borrows, even mutable ones, as we need
                 // to promote mutable borrows of some ZSTs e.g., `&mut []`.
                 let allowed_use = match context {
@@ -145,24 +139,22 @@ impl<'tcx> Visitor<'tcx> for Collector<'_, 'tcx> {
                     *uses += 1;
                     return;
                 }
-            /* mark as unpromotable below */
+                TempState::Unpromotable
             }
-        *temp = TempState::Unpromotable;
+            TempState::Unpromotable | TempState::PromotedOut => TempState::Unpromotable,
+        };
     }
 
     fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) {
         self.super_rvalue(rvalue, location);
 
-        match *rvalue {
+        if let Rvalue::Ref(..) = *rvalue {
-            Rvalue::Ref(..) => {
             self.candidates.push(Candidate { location });
         }
-            _ => {}
-        }
     }
 }
 
-pub fn collect_temps_and_candidates<'tcx>(
+fn collect_temps_and_candidates<'tcx>(
     ccx: &ConstCx<'_, 'tcx>,
 ) -> (IndexVec<Local, TempState>, Vec<Candidate>) {
     let mut collector = Collector {
@@ -196,10 +188,9 @@ struct Unpromotable;
 
 impl<'tcx> Validator<'_, 'tcx> {
     fn validate_candidate(&mut self, candidate: Candidate) -> Result<(), Unpromotable> {
-        let loc = candidate.location;
+        let Left(statement) = self.body.stmt_at(candidate.location) else { bug!() };
-        let statement = &self.body[loc.block].statements[loc.statement_index];
+        let Some((_, Rvalue::Ref(_, kind, place))) = statement.kind.as_assign() else { bug!() };
-        match &statement.kind {
+
-            StatementKind::Assign(box (_, Rvalue::Ref(_, kind, place))) => {
         // We can only promote interior borrows of promotable temps (non-temps
         // don't get promoted anyway).
         self.validate_local(place.local)?;
@@ -216,104 +207,93 @@ impl<'tcx> Validator<'_, 'tcx> {
 
         Ok(())
     }
-            _ => bug!(),
-        }
-    }
 
     // FIXME(eddyb) maybe cache this?
     fn qualif_local<Q: qualifs::Qualif>(&mut self, local: Local) -> bool {
-        if let TempState::Defined { location: loc, .. } = self.temps[local] {
+        let TempState::Defined { location: loc, .. } = self.temps[local] else {
-            let num_stmts = self.body[loc.block].statements.len();
+            return false;
+        };
 
-            if loc.statement_index < num_stmts {
+        let stmt_or_term = self.body.stmt_at(loc);
-                let statement = &self.body[loc.block].statements[loc.statement_index];
+        match stmt_or_term {
-                match &statement.kind {
+            Left(statement) => {
-                    StatementKind::Assign(box (_, rhs)) => qualifs::in_rvalue::<Q, _>(
+                let Some((_, rhs)) = statement.kind.as_assign() else {
-                        self.ccx,
+                    span_bug!(statement.source_info.span, "{:?} is not an assignment", statement)
-                        &mut |l| self.qualif_local::<Q>(l),
+                };
-                        rhs,
+                qualifs::in_rvalue::<Q, _>(self.ccx, &mut |l| self.qualif_local::<Q>(l), rhs)
-                    ),
-                    _ => {
-                        span_bug!(
-                            statement.source_info.span,
-                            "{:?} is not an assignment",
-                            statement
-                        );
             }
-                }
+            Right(terminator) => {
-            } else {
+                assert_matches!(terminator.kind, TerminatorKind::Call { .. });
-                let terminator = self.body[loc.block].terminator();
-                match &terminator.kind {
-                    TerminatorKind::Call { .. } => {
                 let return_ty = self.body.local_decls[local].ty;
                 Q::in_any_value_of_ty(self.ccx, return_ty)
             }
-                    kind => {
-                        span_bug!(terminator.source_info.span, "{:?} not promotable", kind);
-                    }
-                }
-            }
-        } else {
-            false
         }
     }
 
     fn validate_local(&mut self, local: Local) -> Result<(), Unpromotable> {
-        if let TempState::Defined { location: loc, uses, valid } = self.temps[local] {
+        let TempState::Defined { location: loc, uses, valid } = self.temps[local] else {
-            // We cannot promote things that need dropping, since the promoted value
+            return Err(Unpromotable);
-            // would not get dropped.
+        };
+
+        // We cannot promote things that need dropping, since the promoted value would not get
+        // dropped.
         if self.qualif_local::<qualifs::NeedsDrop>(local) {
             return Err(Unpromotable);
         }
-            valid.or_else(|_| {
-                let ok = {
-                    let block = &self.body[loc.block];
-                    let num_stmts = block.statements.len();
 
-                    if loc.statement_index < num_stmts {
+        if valid.is_ok() {
-                        let statement = &block.statements[loc.statement_index];
+            return Ok(());
-                        match &statement.kind {
+        }
-                            StatementKind::Assign(box (_, rhs)) => self.validate_rvalue(rhs),
+
-                            _ => {
+        let ok = {
+            let stmt_or_term = self.body.stmt_at(loc);
+            match stmt_or_term {
+                Left(statement) => {
+                    let Some((_, rhs)) = statement.kind.as_assign() else {
                         span_bug!(
                             statement.source_info.span,
                             "{:?} is not an assignment",
                             statement
-                                );
+                        )
-                            }
+                    };
-                        }
+                    self.validate_rvalue(rhs)
-                    } else {
-                        let terminator = block.terminator();
-                        match &terminator.kind {
-                            TerminatorKind::Call { func, args, .. } => {
-                                self.validate_call(func, args)
                 }
+                Right(terminator) => match &terminator.kind {
+                    TerminatorKind::Call { func, args, .. } => self.validate_call(func, args),
                     TerminatorKind::Yield { .. } => Err(Unpromotable),
                     kind => {
                         span_bug!(terminator.source_info.span, "{:?} not promotable", kind);
                     }
-                        }
+                },
             }
         };
+
         self.temps[local] = match ok {
             Ok(()) => TempState::Defined { location: loc, uses, valid: Ok(()) },
             Err(_) => TempState::Unpromotable,
         };
+
         ok
-            })
-        } else {
-            Err(Unpromotable)
-        }
     }
 
     fn validate_place(&mut self, place: PlaceRef<'tcx>) -> Result<(), Unpromotable> {
-        match place.last_projection() {
+        let Some((place_base, elem)) = place.last_projection() else {
-            None => self.validate_local(place.local),
+            return self.validate_local(place.local);
-            Some((place_base, elem)) => {
+        };
+
         // Validate topmost projection, then recurse.
         match elem {
+            // Recurse directly.
+            ProjectionElem::ConstantIndex { .. }
+            | ProjectionElem::Subtype(_)
+            | ProjectionElem::Subslice { .. } => {}
+
+            // Never recurse.
+            ProjectionElem::OpaqueCast(..) | ProjectionElem::Downcast(..) => {
+                return Err(Unpromotable);
+            }
+
             ProjectionElem::Deref => {
-                        let mut promotable = false;
                 // When a static is used by-value, that gets desugared to `*STATIC_ADDR`,
                 // and we need to be able to promote this. So check if this deref matches
                 // that specific pattern.
@@ -321,91 +301,40 @@ impl<'tcx> Validator<'_, 'tcx> {
                 // We need to make sure this is a `Deref` of a local with no further projections.
                 // Discussion can be found at
                 // https://github.com/rust-lang/rust/pull/74945#discussion_r463063247
-                        if let Some(local) = place_base.as_local() {
+                if let Some(local) = place_base.as_local()
-                            if let TempState::Defined { location, .. } = self.temps[local] {
+                    && let TempState::Defined { location, .. } = self.temps[local]
-                                let def_stmt = self.body[location.block]
+                    && let Left(def_stmt) = self.body.stmt_at(location)
-                                    .statements
+                    && let Some((_, Rvalue::Use(Operand::Constant(c)))) = def_stmt.kind.as_assign()
-                                    .get(location.statement_index);
+                    && let Some(did) = c.check_static_ptr(self.tcx)
-                                if let Some(Statement {
-                                    kind:
-                                        StatementKind::Assign(box (
-                                            _,
-                                            Rvalue::Use(Operand::Constant(c)),
-                                        )),
-                                    ..
-                                }) = def_stmt
-                                {
-                                    if let Some(did) = c.check_static_ptr(self.tcx) {
                     // Evaluating a promoted may not read statics except if it got
                     // promoted from a static (this is a CTFE check). So we
                     // can only promote static accesses inside statics.
-                                        if let Some(hir::ConstContext::Static(..)) = self.const_kind
+                    && let Some(hir::ConstContext::Static(..)) = self.const_kind
+                    && !self.tcx.is_thread_local_static(did)
                 {
-                                            if !self.tcx.is_thread_local_static(did) {
+                    // Recurse.
-                                                promotable = true;
+                } else {
-                                            }
-                                        }
-                                    }
-                                }
-                            }
-                        }
-                        if !promotable {
                     return Err(Unpromotable);
                 }
             }
-                    ProjectionElem::OpaqueCast(..) | ProjectionElem::Downcast(..) => {
-                        return Err(Unpromotable);
-                    }
-
-                    ProjectionElem::ConstantIndex { .. }
-                    | ProjectionElem::Subtype(_)
-                    | ProjectionElem::Subslice { .. } => {}
-
             ProjectionElem::Index(local) => {
-                        let mut promotable = false;
                 // Only accept if we can predict the index and are indexing an array.
-                        let val = if let TempState::Defined { location: loc, .. } =
+                if let TempState::Defined { location: loc, .. } = self.temps[local]
-                            self.temps[local]
+                    && let Left(statement) =  self.body.stmt_at(loc)
-                        {
+                    && let Some((_, Rvalue::Use(Operand::Constant(c)))) = statement.kind.as_assign()
-                            let block = &self.body[loc.block];
+                    && let Some(idx) = c.const_.try_eval_target_usize(self.tcx, self.param_env)
-                            if loc.statement_index < block.statements.len() {
-                                let statement = &block.statements[loc.statement_index];
-                                match &statement.kind {
-                                    StatementKind::Assign(box (
-                                        _,
-                                        Rvalue::Use(Operand::Constant(c)),
-                                    )) => c.const_.try_eval_target_usize(self.tcx, self.param_env),
-                                    _ => None,
-                                }
-                            } else {
-                                None
-                            }
-                        } else {
-                            None
-                        };
-                        if let Some(idx) = val {
                     // Determine the type of the thing we are indexing.
-                            let ty = place_base.ty(self.body, self.tcx).ty;
+                    && let ty::Array(_, len) = place_base.ty(self.body, self.tcx).ty.kind()
-                            match ty.kind() {
-                                ty::Array(_, len) => {
                     // It's an array; determine its length.
-                                    if let Some(len) =
+                    && let Some(len) = len.try_eval_target_usize(self.tcx, self.param_env)
-                                        len.try_eval_target_usize(self.tcx, self.param_env)
-                                    {
                     // If the index is in-bounds, go ahead.
-                                        if idx < len {
+                    && idx < len
-                                            promotable = true;
+                {
-                                        }
+                    self.validate_local(local)?;
-                                    }
+                    // Recurse.
-                                }
+                } else {
-                                _ => {}
-                            }
-                        }
-                        if !promotable {
                     return Err(Unpromotable);
                 }
-
-                        self.validate_local(local)?;
             }
 
             ProjectionElem::Field(..) => {
@@ -419,8 +348,6 @@ impl<'tcx> Validator<'_, 'tcx> {
 
         self.validate_place(place_base)
     }
-        }
-    }
 
     fn validate_operand(&mut self, operand: &Operand<'tcx>) -> Result<(), Unpromotable> {
         match operand {
@@ -676,7 +603,7 @@ impl<'tcx> Validator<'_, 'tcx> {
 }
 
 // FIXME(eddyb) remove the differences for promotability in `static`, `const`, `const fn`.
-pub fn validate_candidates(
+fn validate_candidates(
     ccx: &ConstCx<'_, '_>,
     temps: &mut IndexSlice<Local, TempState>,
     candidates: &[Candidate],
@@ -930,7 +857,7 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Promoter<'a, 'tcx> {
     }
 }
 
-pub fn promote_candidates<'tcx>(
+fn promote_candidates<'tcx>(
     body: &mut Body<'tcx>,
     tcx: TyCtxt<'tcx>,
     mut temps: IndexVec<Local, TempState>,