Add a fast path for lowering trivial consts

compiler/rustc_mir_transform/src/lib.rs

@@ -23,11 +23,11 @@ use rustc_hir::def::{CtorKind, DefKind};
 use rustc_hir::def_id::LocalDefId;
 use rustc_index::IndexVec;
 use rustc_middle::mir::{
-    AnalysisPhase, Body, CallSource, ClearCrossCrate, ConstOperand, ConstQualifs, LocalDecl,
-    MirPhase, Operand, Place, ProjectionElem, Promoted, RuntimePhase, Rvalue, START_BLOCK,
-    SourceInfo, Statement, StatementKind, TerminatorKind,
+    AnalysisPhase, Body, CallSource, ClearCrossCrate, ConstOperand, ConstQualifs, ConstValue,
+    LocalDecl, MirPhase, Operand, Place, ProjectionElem, Promoted, RETURN_PLACE, RuntimePhase,
+    Rvalue, START_BLOCK, SourceInfo, Statement, StatementKind, TerminatorKind,
 };
-use rustc_middle::ty::{self, TyCtxt, TypeVisitableExt};
+use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitableExt};
 use rustc_middle::util::Providers;
 use rustc_middle::{bug, query, span_bug};
 use rustc_mir_build::builder::build_mir;
@@ -226,6 +226,7 @@ pub fn provide(providers: &mut Providers) {
         promoted_mir,
         deduced_param_attrs: deduce_param_attrs::deduced_param_attrs,
         coroutine_by_move_body_def_id: coroutine::coroutine_by_move_body_def_id,
+        trivial_const: trivial_const_provider,
         ..providers.queries
     };
 }
@@ -376,9 +377,71 @@ fn mir_const_qualif(tcx: TyCtxt<'_>, def: LocalDefId) -> ConstQualifs {
     validator.qualifs_in_return_place()
 }
 
+fn def_kind_compatible_with_trivial_mir<'tcx>(tcx: TyCtxt<'tcx>, def: LocalDefId) -> bool {
+    // Static and InlineConst are the obvious additions, but
+    // * Statics need additional type-checking to taint `static A: _ = 0;`, currently we'd ICE.
+    // * The MIR for InlineConst is used by the borrow checker, and not easy to skip over.
+    matches!(tcx.def_kind(def), DefKind::AssocConst | DefKind::Const | DefKind::AnonConst)
+}
+
+fn trivial_const_provider<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def: LocalDefId,
+) -> Option<(ConstValue, Ty<'tcx>)> {
+    if def_kind_compatible_with_trivial_mir(tcx, def) {
+        trivial_const(&tcx.mir_built(def).borrow())
+    } else {
+        None
+    }
+}
+
+fn trivial_const<'tcx>(body: &Body<'tcx>) -> Option<(ConstValue, Ty<'tcx>)> {
+    if body.has_opaque_types() {
+        return None;
+    }
+
+    if body.basic_blocks.len() != 1 {
+        return None;
+    }
+
+    let block = &body.basic_blocks[START_BLOCK];
+    if block.statements.len() != 1 {
+        return None;
+    }
+
+    if block.terminator().kind != TerminatorKind::Return {
+        return None;
+    }
+
+    let StatementKind::Assign(box (place, rvalue)) = &block.statements[0].kind else {
+        return None;
+    };
+
+    if *place != Place::from(RETURN_PLACE) {
+        return None;
+    }
+
+    if let Rvalue::Use(Operand::Constant(c)) = rvalue {
+        if let rustc_middle::mir::Const::Val(v, ty) = c.const_ {
+            return Some((v, ty));
+        }
+    }
+
+    return None;
+}
+
 fn mir_built(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> {
     let mut body = build_mir(tcx, def);
 
+    // Identifying trivial consts based on their mir_built is easy, but a little wasteful.
+    // Trying to push this logic earlier in the compiler and never even produce the Body would
+    // probably improve compile time.
+    if def_kind_compatible_with_trivial_mir(tcx, def) && trivial_const(&body).is_some() {
+        let body = tcx.alloc_steal_mir(body);
+        pass_manager::dump_mir_for_phase_change(tcx, &body.borrow());
+        return body;
+    }
+
     pass_manager::dump_mir_for_phase_change(tcx, &body);
 
     pm::run_passes(
@@ -409,6 +472,8 @@ fn mir_promoted(
     tcx: TyCtxt<'_>,
     def: LocalDefId,
 ) -> (&Steal<Body<'_>>, &Steal<IndexVec<Promoted, Body<'_>>>) {
+    debug_assert!(!tcx.is_trivial_const(def), "Tried to get mir_promoted of a trivial const");
+
     // Ensure that we compute the `mir_const_qualif` for constants at
     // this point, before we steal the mir-const result.
     // Also this means promotion can rely on all const checks having been done.
@@ -436,6 +501,9 @@ fn mir_promoted(
         tcx.ensure_done().coroutine_by_move_body_def_id(def);
     }
 
+    // the `trivial_const` query uses mir_built, so make sure it is run.
+    tcx.ensure_done().trivial_const(def);
+
     let mut body = tcx.mir_built(def).steal();
     if let Some(error_reported) = const_qualifs.tainted_by_errors {
         body.tainted_by_errors = Some(error_reported);
@@ -463,6 +531,7 @@ fn mir_promoted(
 
 /// Compute the MIR that is used during CTFE (and thus has no optimizations run on it)
 fn mir_for_ctfe(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &Body<'_> {
+    debug_assert!(!tcx.is_trivial_const(def_id), "Tried to get mir_for_ctfe of a trivial const");
     tcx.arena.alloc(inner_mir_for_ctfe(tcx, def_id))
 }