Auto merge of #78407 - oli-obk:ub_checkable_ctfe, r=RalfJung,pnkfelix

Make CTFE able to check for UB...

... by not doing any optimizations on the `const fn` MIR used in CTFE. This means we duplicate all `const fn`'s MIR now, once for CTFE, once for runtime. This PR is for checking the perf effect, so we have some data when talking about https://github.com/rust-lang/const-eval/blob/master/rfcs/0000-const-ub.md

To do this, we now have two queries for obtaining mir: `optimized_mir` and `mir_for_ctfe`. It is now illegal to invoke `optimized_mir` to obtain the MIR of a const/static item's initializer, an array length, an inline const expression or an enum discriminant initializer. For `const fn`, both `optimized_mir` and `mir_for_ctfe` work, the former returning the MIR that LLVM should use if the function is called at runtime. Similarly it is illegal to invoke `mir_for_ctfe` on regular functions.

This is all checked via appropriate assertions and I don't think it is easy to get wrong, as there should be no `mir_for_ctfe` calls outside the const evaluator or metadata encoding. Almost all rustc devs should keep using `optimized_mir` (or `instance_mir` for that matter).

compiler/rustc_mir/src/const_eval/machine.rs

@@ -201,6 +201,22 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
 
     type MemoryExtra = MemoryExtra;
 
+    fn load_mir(
+        ecx: &InterpCx<'mir, 'tcx, Self>,
+        instance: ty::InstanceDef<'tcx>,
+    ) -> InterpResult<'tcx, &'tcx mir::Body<'tcx>> {
+        match instance {
+            ty::InstanceDef::Item(def) => {
+                if ecx.tcx.is_ctfe_mir_available(def.did) {
+                    Ok(ecx.tcx.mir_for_ctfe_opt_const_arg(def))
+                } else {
+                    throw_unsup!(NoMirFor(def.did))
+                }
+            }
+            _ => Ok(ecx.tcx.instance_mir(instance)),
+        }
+    }
+
     fn find_mir_or_eval_fn(
         ecx: &mut InterpCx<'mir, 'tcx, Self>,
         instance: ty::Instance<'tcx>,

compiler/rustc_mir/src/interpret/eval_context.rs

@@ -477,16 +477,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         if let Some(promoted) = promoted {
             return Ok(&self.tcx.promoted_mir_opt_const_arg(def)[promoted]);
         }
-        match instance {
-            ty::InstanceDef::Item(def) => {
-                if self.tcx.is_mir_available(def.did) {
-                    Ok(self.tcx.optimized_mir_opt_const_arg(def))
-                } else {
-                    throw_unsup!(NoMirFor(def.did))
-                }
-            }
-            _ => Ok(self.tcx.instance_mir(instance)),
-        }
+        M::load_mir(self, instance)
     }
 
     /// Call this on things you got out of the MIR (so it is as generic as the current

compiler/rustc_mir/src/interpret/machine.rs

@@ -132,6 +132,16 @@ pub trait Machine<'mir, 'tcx>: Sized {
     /// Whether to enforce the validity invariant
     fn enforce_validity(ecx: &InterpCx<'mir, 'tcx, Self>) -> bool;
 
+    /// Entry point for obtaining the MIR of anything that should get evaluated.
+    /// So not just functions and shims, but also const/static initializers, anonymous
+    /// constants, ...
+    fn load_mir(
+        ecx: &InterpCx<'mir, 'tcx, Self>,
+        instance: ty::InstanceDef<'tcx>,
+    ) -> InterpResult<'tcx, &'tcx mir::Body<'tcx>> {
+        Ok(ecx.tcx.instance_mir(instance))
+    }
+
     /// Entry point to all function calls.
     ///
     /// Returns either the mir to use for the call, or `None` if execution should

compiler/rustc_mir/src/monomorphize/polymorphize.rs

@@ -5,7 +5,7 @@
 //! generic parameters are unused (and eventually, in what ways generic parameters are used - only
 //! for their size, offset of a field, etc.).
 
-use rustc_hir::{def::DefKind, def_id::DefId};
+use rustc_hir::{def::DefKind, def_id::DefId, ConstContext};
 use rustc_index::bit_set::FiniteBitSet;
 use rustc_middle::mir::{
     visit::{TyContext, Visitor},
@@ -54,9 +54,17 @@ fn unused_generic_params(tcx: TyCtxt<'_>, def_id: DefId) -> FiniteBitSet<u32> {
     }
 
     // Exit early when there is no MIR available.
-    if !tcx.is_mir_available(def_id) {
-        debug!("unused_generic_params: (no mir available) def_id={:?}", def_id);
-        return FiniteBitSet::new_empty();
+    let context = tcx.hir().body_const_context(def_id.expect_local());
+    match context {
+        Some(ConstContext::ConstFn) | None if !tcx.is_mir_available(def_id) => {
+            debug!("unused_generic_params: (no mir available) def_id={:?}", def_id);
+            return FiniteBitSet::new_empty();
+        }
+        Some(_) if !tcx.is_ctfe_mir_available(def_id) => {
+            debug!("unused_generic_params: (no ctfe mir available) def_id={:?}", def_id);
+            return FiniteBitSet::new_empty();
+        }
+        _ => {}
     }
 
     // Create a bitset with N rightmost ones for each parameter.
@@ -69,7 +77,12 @@ fn unused_generic_params(tcx: TyCtxt<'_>, def_id: DefId) -> FiniteBitSet<u32> {
     debug!("unused_generic_params: (after default) unused_parameters={:?}", unused_parameters);
 
     // Visit MIR and accumululate used generic parameters.
-    let body = tcx.optimized_mir(def_id);
+    let body = match context {
+        // Const functions are actually called and should thus be considered for polymorphization
+        // via their runtime MIR
+        Some(ConstContext::ConstFn) | None => tcx.optimized_mir(def_id),
+        Some(_) => tcx.mir_for_ctfe(def_id),
+    };
     let mut vis = MarkUsedGenericParams { tcx, def_id, unused_parameters: &mut unused_parameters };
     vis.visit_body(body);
     debug!("unused_generic_params: (after visitor) unused_parameters={:?}", unused_parameters);

compiler/rustc_mir/src/transform/const_prop.rs

@@ -185,6 +185,13 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for ConstPropMachine<'mir, 'tcx>
 
     type MemoryExtra = ();
 
+    fn load_mir(
+        _ecx: &InterpCx<'mir, 'tcx, Self>,
+        _instance: ty::InstanceDef<'tcx>,
+    ) -> InterpResult<'tcx, &'tcx Body<'tcx>> {
+        throw_machine_stop_str!("calling functions isn't supported in ConstProp")
+    }
+
     fn find_mir_or_eval_fn(
         _ecx: &mut InterpCx<'mir, 'tcx, Self>,
         _instance: ty::Instance<'tcx>,

compiler/rustc_mir/src/transform/coverage/query.rs

@@ -4,7 +4,7 @@ use rustc_middle::mir::coverage::*;
 use rustc_middle::mir::visit::Visitor;
 use rustc_middle::mir::{self, Coverage, CoverageInfo, Location};
 use rustc_middle::ty::query::Providers;
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::{self, TyCtxt};
 use rustc_span::def_id::DefId;
 
 /// The `query` provider for `CoverageInfo`, requested by `codegen_coverage()` (to inject each
@@ -112,7 +112,7 @@ impl Visitor<'_> for CoverageVisitor {
 }
 
 fn coverageinfo_from_mir<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> CoverageInfo {
-    let mir_body = tcx.optimized_mir(def_id);
+    let mir_body = mir_body(tcx, def_id);
 
     let mut coverage_visitor = CoverageVisitor {
         // num_counters always has at least the `ZERO` counter.
@@ -129,8 +129,7 @@ fn coverageinfo_from_mir<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> CoverageInfo
 }
 
 fn covered_file_name<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Symbol> {
-    let mir_body = tcx.optimized_mir(def_id);
-    for bb_data in mir_body.basic_blocks().iter() {
+    for bb_data in mir_body(tcx, def_id).basic_blocks().iter() {
         for statement in bb_data.statements.iter() {
             if let StatementKind::Coverage(box ref coverage) = statement.kind {
                 if let Some(code_region) = coverage.code_region.as_ref() {
@@ -142,9 +141,17 @@ fn covered_file_name<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Symbol> {
     None
 }
 
+/// This function ensures we obtain the correct MIR for the given item irrespective of
+/// whether that means const mir or runtime mir. For `const fn` this opts for runtime
+/// mir.
+fn mir_body<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> &'tcx mir::Body<'tcx> {
+    let id = ty::WithOptConstParam::unknown(def_id);
+    let def = ty::InstanceDef::Item(id);
+    tcx.instance_mir(def)
+}
+
 fn covered_code_regions<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Vec<&'tcx CodeRegion> {
-    let mir_body: &'tcx mir::Body<'tcx> = tcx.optimized_mir(def_id);
-    mir_body
+    mir_body(tcx, def_id)
         .basic_blocks()
         .iter()
         .map(|data| {

compiler/rustc_mir/src/transform/mod.rs

@@ -71,9 +71,11 @@ pub(crate) fn provide(providers: &mut Providers) {
         },
         mir_promoted,
         mir_drops_elaborated_and_const_checked,
+        mir_for_ctfe,
+        mir_for_ctfe_of_const_arg,
         optimized_mir,
-        optimized_mir_of_const_arg,
         is_mir_available,
+        is_ctfe_mir_available: |tcx, did| is_mir_available(tcx, did),
         promoted_mir: |tcx, def_id| {
             let def_id = def_id.expect_local();
             if let Some(def) = ty::WithOptConstParam::try_lookup(def_id, tcx) {
@@ -278,6 +280,7 @@ fn mir_const<'tcx>(
     tcx.alloc_steal_mir(body)
 }
 
+/// Compute the main MIR body and the list of MIR bodies of the promoteds.
 fn mir_promoted(
     tcx: TyCtxt<'tcx>,
     def: ty::WithOptConstParam<LocalDefId>,
@@ -319,6 +322,87 @@ fn mir_promoted(
     (tcx.alloc_steal_mir(body), tcx.alloc_steal_promoted(promoted))
 }
 
+/// Compute the MIR that is used during CTFE (and thus has no optimizations run on it)
+fn mir_for_ctfe<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> &'tcx Body<'tcx> {
+    let did = def_id.expect_local();
+    if let Some(def) = ty::WithOptConstParam::try_lookup(did, tcx) {
+        tcx.mir_for_ctfe_of_const_arg(def)
+    } else {
+        tcx.arena.alloc(inner_mir_for_ctfe(tcx, ty::WithOptConstParam::unknown(did)))
+    }
+}
+
+/// Same as `mir_for_ctfe`, but used to get the MIR of a const generic parameter.
+/// The docs on `WithOptConstParam` explain this a bit more, but the TLDR is that
+/// we'd get cycle errors with `mir_for_ctfe`, because typeck would need to typeck
+/// the const parameter while type checking the main body, which in turn would try
+/// to type check the main body again.
+fn mir_for_ctfe_of_const_arg<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    (did, param_did): (LocalDefId, DefId),
+) -> &'tcx Body<'tcx> {
+    tcx.arena.alloc(inner_mir_for_ctfe(
+        tcx,
+        ty::WithOptConstParam { did, const_param_did: Some(param_did) },
+    ))
+}
+
+fn inner_mir_for_ctfe(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_> {
+    // FIXME: don't duplicate this between the optimized_mir/mir_for_ctfe queries
+    if tcx.is_constructor(def.did.to_def_id()) {
+        // There's no reason to run all of the MIR passes on constructors when
+        // we can just output the MIR we want directly. This also saves const
+        // qualification and borrow checking the trouble of special casing
+        // constructors.
+        return shim::build_adt_ctor(tcx, def.did.to_def_id());
+    }
+
+    let context = tcx
+        .hir()
+        .body_const_context(def.did)
+        .expect("mir_for_ctfe should not be used for runtime functions");
+
+    let mut body = tcx.mir_drops_elaborated_and_const_checked(def).borrow().clone();
+
+    match context {
+        // Do not const prop functions, either they get executed at runtime or exported to metadata,
+        // so we run const prop on them, or they don't, in which case we const evaluate some control
+        // flow paths of the function and any errors in those paths will get emitted as const eval
+        // errors.
+        hir::ConstContext::ConstFn => {}
+        // Static items always get evaluated, so we can just let const eval see if any erroneous
+        // control flow paths get executed.
+        hir::ConstContext::Static(_) => {}
+        // Associated constants get const prop run so we detect common failure situations in the
+        // crate that defined the constant.
+        // Technically we want to not run on regular const items, but oli-obk doesn't know how to
+        // conveniently detect that at this point without looking at the HIR.
+        hir::ConstContext::Const => {
+            #[rustfmt::skip]
+            let optimizations: &[&dyn MirPass<'_>] = &[
+                &const_prop::ConstProp,
+            ];
+
+            #[rustfmt::skip]
+            run_passes(
+                tcx,
+                &mut body,
+                MirPhase::Optimization,
+                &[
+                    optimizations,
+                ],
+            );
+        }
+    }
+
+    debug_assert!(!body.has_free_regions(), "Free regions in MIR for CTFE");
+
+    body
+}
+
+/// Obtain just the main MIR (no promoteds) and run some cleanups on it. This also runs
+/// mir borrowck *before* doing so in order to ensure that borrowck can be run and doesn't
+/// end up missing the source MIR due to stealing happening.
 fn mir_drops_elaborated_and_const_checked<'tcx>(
     tcx: TyCtxt<'tcx>,
     def: ty::WithOptConstParam<LocalDefId>,
@@ -456,35 +540,32 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     );
 }
 
+/// Optimize the MIR and prepare it for codegen.
 fn optimized_mir<'tcx>(tcx: TyCtxt<'tcx>, did: DefId) -> &'tcx Body<'tcx> {
     let did = did.expect_local();
-    if let Some(def) = ty::WithOptConstParam::try_lookup(did, tcx) {
-        tcx.optimized_mir_of_const_arg(def)
-    } else {
-        tcx.arena.alloc(inner_optimized_mir(tcx, ty::WithOptConstParam::unknown(did)))
-    }
+    assert_eq!(ty::WithOptConstParam::try_lookup(did, tcx), None);
+    tcx.arena.alloc(inner_optimized_mir(tcx, did))
 }
 
-fn optimized_mir_of_const_arg<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    (did, param_did): (LocalDefId, DefId),
-) -> &'tcx Body<'tcx> {
-    tcx.arena.alloc(inner_optimized_mir(
-        tcx,
-        ty::WithOptConstParam { did, const_param_did: Some(param_did) },
-    ))
-}
-
-fn inner_optimized_mir(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_> {
-    if tcx.is_constructor(def.did.to_def_id()) {
+fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
+    if tcx.is_constructor(did.to_def_id()) {
         // There's no reason to run all of the MIR passes on constructors when
         // we can just output the MIR we want directly. This also saves const
         // qualification and borrow checking the trouble of special casing
         // constructors.
-        return shim::build_adt_ctor(tcx, def.did.to_def_id());
+        return shim::build_adt_ctor(tcx, did.to_def_id());
     }
 
-    let mut body = tcx.mir_drops_elaborated_and_const_checked(def).steal();
+    match tcx.hir().body_const_context(did) {
+        // Run the `mir_for_ctfe` query, which depends on `mir_drops_elaborated_and_const_checked`
+        // which we are going to steal below. Thus we need to run `mir_for_ctfe` first, so it
+        // computes and caches its result.
+        Some(hir::ConstContext::ConstFn) => tcx.ensure().mir_for_ctfe(did),
+        None => {}
+        Some(other) => panic!("do not use `optimized_mir` for constants: {:?}", other),
+    }
+    let mut body =
+        tcx.mir_drops_elaborated_and_const_checked(ty::WithOptConstParam::unknown(did)).steal();
     run_optimization_passes(tcx, &mut body);
 
     debug_assert!(!body.has_free_regions(), "Free regions in optimized MIR");
@@ -492,6 +573,8 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>)
     body
 }
 
+/// Fetch all the promoteds of an item and prepare their MIR bodies to be ready for
+/// constant evaluation once all substitutions become known.
 fn promoted_mir<'tcx>(
     tcx: TyCtxt<'tcx>,
     def: ty::WithOptConstParam<LocalDefId>,
@@ -510,7 +593,6 @@ fn promoted_mir<'tcx>(
 
     for body in &mut promoted {
         run_post_borrowck_cleanup_passes(tcx, body);
-        run_optimization_passes(tcx, body);
     }
 
     debug_assert!(!promoted.has_free_regions(), "Free regions in promoted MIR");

compiler/rustc_mir/src/util/pretty.rs

@@ -273,8 +273,6 @@ pub fn write_mir_pretty<'tcx>(
 
     let mut first = true;
     for def_id in dump_mir_def_ids(tcx, single) {
-        let body = &tcx.optimized_mir(def_id);
-
         if first {
             first = false;
         } else {
@@ -282,11 +280,28 @@ pub fn write_mir_pretty<'tcx>(
             writeln!(w)?;
         }
 
-        write_mir_fn(tcx, body, &mut |_, _| Ok(()), w)?;
-
-        for body in tcx.promoted_mir(def_id) {
-            writeln!(w)?;
+        let render_body = |w: &mut dyn Write, body| -> io::Result<()> {
             write_mir_fn(tcx, body, &mut |_, _| Ok(()), w)?;
+
+            for body in tcx.promoted_mir(def_id) {
+                writeln!(w)?;
+                write_mir_fn(tcx, body, &mut |_, _| Ok(()), w)?;
+            }
+            Ok(())
+        };
+        match tcx.hir().body_const_context(def_id.expect_local()) {
+            None => render_body(w, tcx.optimized_mir(def_id))?,
+            // For `const fn` we want to render the optimized MIR. If you want the mir used in
+            // ctfe, you can dump the MIR after the `Deaggregator` optimization pass.
+            Some(rustc_hir::ConstContext::ConstFn) => {
+                render_body(w, tcx.optimized_mir(def_id))?;
+                writeln!(w)?;
+                writeln!(w, "// MIR FOR CTFE")?;
+                // Do not use `render_body`, as that would render the promoteds again, but these
+                // are shared between mir_for_ctfe and optimized_mir
+                write_mir_fn(tcx, tcx.mir_for_ctfe(def_id), &mut |_, _| Ok(()), w)?;
+            }
+            Some(_) => render_body(w, tcx.mir_for_ctfe(def_id))?,
         }
     }
     Ok(())