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_metadata/src/rmeta/decoder.rs

@@ -1160,6 +1160,10 @@ impl<'a, 'tcx> CrateMetadataRef<'a> {
         }
     }
 
+    fn is_ctfe_mir_available(&self, id: DefIndex) -> bool {
+        self.root.tables.mir_for_ctfe.get(self, id).is_some()
+    }
+
     fn is_item_mir_available(&self, id: DefIndex) -> bool {
         self.root.tables.mir.get(self, id).is_some()
     }
@@ -1183,6 +1187,17 @@ impl<'a, 'tcx> CrateMetadataRef<'a> {
             .decode((self, tcx))
     }
 
+    fn get_mir_for_ctfe(&self, tcx: TyCtxt<'tcx>, id: DefIndex) -> Body<'tcx> {
+        self.root
+            .tables
+            .mir_for_ctfe
+            .get(self, id)
+            .unwrap_or_else(|| {
+                bug!("get_mir_for_ctfe: missing MIR for `{:?}`", self.local_def_id(id))
+            })
+            .decode((self, tcx))
+    }
+
     fn get_mir_abstract_const(
         &self,
         tcx: TyCtxt<'tcx>,

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -115,6 +115,7 @@ provide! { <'tcx> tcx, def_id, other, cdata,
         })
     }
     optimized_mir => { tcx.arena.alloc(cdata.get_optimized_mir(tcx, def_id.index)) }
+    mir_for_ctfe => { tcx.arena.alloc(cdata.get_mir_for_ctfe(tcx, def_id.index)) }
     promoted_mir => { tcx.arena.alloc(cdata.get_promoted_mir(tcx, def_id.index)) }
     mir_abstract_const => { cdata.get_mir_abstract_const(tcx, def_id.index) }
     unused_generic_params => { cdata.get_unused_generic_params(def_id.index) }
@@ -145,6 +146,7 @@ provide! { <'tcx> tcx, def_id, other, cdata,
     impl_parent => { cdata.get_parent_impl(def_id.index) }
     trait_of_item => { cdata.get_trait_of_item(def_id.index) }
     is_mir_available => { cdata.is_item_mir_available(def_id.index) }
+    is_ctfe_mir_available => { cdata.is_ctfe_mir_available(def_id.index) }
 
     dylib_dependency_formats => { cdata.get_dylib_dependency_formats(tcx) }
     is_panic_runtime => { cdata.root.panic_runtime }

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -758,8 +758,6 @@ impl EncodeContext<'a, 'tcx> {
         self.encode_generics(def_id);
         self.encode_explicit_predicates(def_id);
         self.encode_inferred_outlives(def_id);
-        self.encode_optimized_mir(def_id.expect_local());
-        self.encode_promoted_mir(def_id.expect_local());
     }
 
     fn encode_enum_variant_ctor(&mut self, def: &ty::AdtDef, index: VariantIdx) {
@@ -789,6 +787,7 @@ impl EncodeContext<'a, 'tcx> {
         self.encode_generics(def_id);
         self.encode_explicit_predicates(def_id);
         self.encode_inferred_outlives(def_id);
+        self.encode_mir_for_ctfe(def_id.expect_local());
         self.encode_optimized_mir(def_id.expect_local());
         self.encode_promoted_mir(def_id.expect_local());
     }
@@ -897,6 +896,7 @@ impl EncodeContext<'a, 'tcx> {
         self.encode_explicit_predicates(def_id);
         self.encode_inferred_outlives(def_id);
         self.encode_optimized_mir(def_id.expect_local());
+        self.encode_mir_for_ctfe(def_id.expect_local());
         self.encode_promoted_mir(def_id.expect_local());
     }
 
@@ -1015,8 +1015,21 @@ impl EncodeContext<'a, 'tcx> {
         self.encode_inferred_outlives(def_id);
 
         // This should be kept in sync with `PrefetchVisitor.visit_trait_item`.
-        self.encode_optimized_mir(def_id.expect_local());
-        self.encode_promoted_mir(def_id.expect_local());
+        match trait_item.kind {
+            ty::AssocKind::Type => {}
+            ty::AssocKind::Const => {
+                if self.tcx.mir_keys(LOCAL_CRATE).contains(&def_id.expect_local()) {
+                    self.encode_mir_for_ctfe(def_id.expect_local());
+                    self.encode_promoted_mir(def_id.expect_local());
+                }
+            }
+            ty::AssocKind::Fn => {
+                if self.tcx.mir_keys(LOCAL_CRATE).contains(&def_id.expect_local()) {
+                    self.encode_optimized_mir(def_id.expect_local());
+                    self.encode_promoted_mir(def_id.expect_local());
+                }
+            }
+        }
     }
 
     fn metadata_output_only(&self) -> bool {
@@ -1089,23 +1102,28 @@ impl EncodeContext<'a, 'tcx> {
 
         // The following part should be kept in sync with `PrefetchVisitor.visit_impl_item`.
 
-        let mir = match ast_item.kind {
-            hir::ImplItemKind::Const(..) => true,
+        let (mir, mir_const) = match ast_item.kind {
+            hir::ImplItemKind::Const(..) => (false, true),
             hir::ImplItemKind::Fn(ref sig, _) => {
                 let generics = self.tcx.generics_of(def_id);
                 let needs_inline = (generics.requires_monomorphization(self.tcx)
                     || tcx.codegen_fn_attrs(def_id).requests_inline())
                     && !self.metadata_output_only();
                 let is_const_fn = sig.header.constness == hir::Constness::Const;
                 let always_encode_mir = self.tcx.sess.opts.debugging_opts.always_encode_mir;
-                needs_inline || is_const_fn || always_encode_mir
+                (needs_inline || always_encode_mir, is_const_fn)
             }
-            hir::ImplItemKind::TyAlias(..) => false,
+            hir::ImplItemKind::TyAlias(..) => (false, false),
         };
         if mir {
             self.encode_optimized_mir(def_id.expect_local());
+        }
+        if mir || mir_const {
             self.encode_promoted_mir(def_id.expect_local());
         }
+        if mir_const {
+            self.encode_mir_for_ctfe(def_id.expect_local());
+        }
     }
 
     fn encode_fn_param_names_for_body(&mut self, body_id: hir::BodyId) -> Lazy<[Ident]> {
@@ -1116,28 +1134,34 @@ impl EncodeContext<'a, 'tcx> {
         self.lazy(param_names.iter())
     }
 
-    fn encode_optimized_mir(&mut self, def_id: LocalDefId) {
-        debug!("EntryBuilder::encode_mir({:?})", def_id);
-        if self.tcx.mir_keys(LOCAL_CRATE).contains(&def_id) {
-            record!(self.tables.mir[def_id.to_def_id()] <- self.tcx.optimized_mir(def_id));
+    fn encode_mir_for_ctfe(&mut self, def_id: LocalDefId) {
+        debug!("EntryBuilder::encode_mir_for_ctfe({:?})", def_id);
+        record!(self.tables.mir_for_ctfe[def_id.to_def_id()] <- self.tcx.mir_for_ctfe(def_id));
 
-            let unused = self.tcx.unused_generic_params(def_id);
-            if !unused.is_empty() {
-                record!(self.tables.unused_generic_params[def_id.to_def_id()] <- unused);
-            }
+        let unused = self.tcx.unused_generic_params(def_id);
+        if !unused.is_empty() {
+            record!(self.tables.unused_generic_params[def_id.to_def_id()] <- unused);
+        }
 
-            let abstract_const = self.tcx.mir_abstract_const(def_id);
-            if let Ok(Some(abstract_const)) = abstract_const {
-                record!(self.tables.mir_abstract_consts[def_id.to_def_id()] <- abstract_const);
-            }
+        let abstract_const = self.tcx.mir_abstract_const(def_id);
+        if let Ok(Some(abstract_const)) = abstract_const {
+            record!(self.tables.mir_abstract_consts[def_id.to_def_id()] <- abstract_const);
+        }
+    }
+
+    fn encode_optimized_mir(&mut self, def_id: LocalDefId) {
+        debug!("EntryBuilder::encode_optimized_mir({:?})", def_id);
+        record!(self.tables.mir[def_id.to_def_id()] <- self.tcx.optimized_mir(def_id));
+
+        let unused = self.tcx.unused_generic_params(def_id);
+        if !unused.is_empty() {
+            record!(self.tables.unused_generic_params[def_id.to_def_id()] <- unused);
         }
     }
 
     fn encode_promoted_mir(&mut self, def_id: LocalDefId) {
         debug!("EncodeContext::encode_promoted_mir({:?})", def_id);
-        if self.tcx.mir_keys(LOCAL_CRATE).contains(&def_id) {
-            record!(self.tables.promoted_mir[def_id.to_def_id()] <- self.tcx.promoted_mir(def_id));
-        }
+        record!(self.tables.promoted_mir[def_id.to_def_id()] <- self.tcx.promoted_mir(def_id));
     }
 
     // Encodes the inherent implementations of a structure, enumeration, or trait.
@@ -1406,22 +1430,31 @@ impl EncodeContext<'a, 'tcx> {
 
         // The following part should be kept in sync with `PrefetchVisitor.visit_item`.
 
-        let mir = match item.kind {
-            hir::ItemKind::Static(..) | hir::ItemKind::Const(..) => true,
+        let (mir, const_mir) = match item.kind {
+            hir::ItemKind::Static(..) | hir::ItemKind::Const(..) => (false, true),
             hir::ItemKind::Fn(ref sig, ..) => {
                 let generics = tcx.generics_of(def_id);
                 let needs_inline = (generics.requires_monomorphization(tcx)
                     || tcx.codegen_fn_attrs(def_id).requests_inline())
                     && !self.metadata_output_only();
+
+                let is_const_fn = sig.header.constness == hir::Constness::Const;
                 let always_encode_mir = self.tcx.sess.opts.debugging_opts.always_encode_mir;
-                needs_inline || sig.header.constness == hir::Constness::Const || always_encode_mir
+                let mir = needs_inline || always_encode_mir;
+                // We don't need the optimized MIR for const fns.
+                (mir, is_const_fn)
             }
-            _ => false,
+            _ => (false, false),
         };
         if mir {
             self.encode_optimized_mir(def_id.expect_local());
+        }
+        if mir || const_mir {
             self.encode_promoted_mir(def_id.expect_local());
         }
+        if const_mir {
+            self.encode_mir_for_ctfe(def_id.expect_local());
+        }
     }
 
     /// Serialize the text of exported macros
@@ -1486,7 +1519,7 @@ impl EncodeContext<'a, 'tcx> {
         self.encode_generics(def_id.to_def_id());
         self.encode_explicit_predicates(def_id.to_def_id());
         self.encode_inferred_outlives(def_id.to_def_id());
-        self.encode_optimized_mir(def_id);
+        self.encode_mir_for_ctfe(def_id);
         self.encode_promoted_mir(def_id);
     }
 
@@ -1951,6 +1984,12 @@ struct PrefetchVisitor<'tcx> {
 }
 
 impl<'tcx> PrefetchVisitor<'tcx> {
+    fn prefetch_ctfe_mir(&self, def_id: LocalDefId) {
+        if self.mir_keys.contains(&def_id) {
+            self.tcx.ensure().mir_for_ctfe(def_id);
+            self.tcx.ensure().promoted_mir(def_id);
+        }
+    }
     fn prefetch_mir(&self, def_id: LocalDefId) {
         if self.mir_keys.contains(&def_id) {
             self.tcx.ensure().optimized_mir(def_id);
@@ -1965,42 +2004,57 @@ impl<'tcx, 'v> ParItemLikeVisitor<'v> for PrefetchVisitor<'tcx> {
         let tcx = self.tcx;
         match item.kind {
             hir::ItemKind::Static(..) | hir::ItemKind::Const(..) => {
-                self.prefetch_mir(tcx.hir().local_def_id(item.hir_id))
+                self.prefetch_ctfe_mir(tcx.hir().local_def_id(item.hir_id))
             }
             hir::ItemKind::Fn(ref sig, ..) => {
                 let def_id = tcx.hir().local_def_id(item.hir_id);
                 let generics = tcx.generics_of(def_id.to_def_id());
                 let needs_inline = generics.requires_monomorphization(tcx)
                     || tcx.codegen_fn_attrs(def_id.to_def_id()).requests_inline();
-                if needs_inline || sig.header.constness == hir::Constness::Const {
+                if needs_inline {
                     self.prefetch_mir(def_id)
                 }
+                if sig.header.constness == hir::Constness::Const {
+                    self.prefetch_ctfe_mir(def_id);
+                }
             }
             _ => (),
         }
     }
 
     fn visit_trait_item(&self, trait_item: &'v hir::TraitItem<'v>) {
         // This should be kept in sync with `encode_info_for_trait_item`.
-        self.prefetch_mir(self.tcx.hir().local_def_id(trait_item.hir_id));
+        let def_id = self.tcx.hir().local_def_id(trait_item.hir_id);
+        match trait_item.kind {
+            hir::TraitItemKind::Type(..) => {}
+            hir::TraitItemKind::Const(..) => {
+                self.prefetch_ctfe_mir(def_id);
+            }
+            hir::TraitItemKind::Fn(..) => {
+                self.prefetch_mir(def_id);
+            }
+        }
     }
 
     fn visit_impl_item(&self, impl_item: &'v hir::ImplItem<'v>) {
         // This should be kept in sync with `encode_info_for_impl_item`.
         let tcx = self.tcx;
         match impl_item.kind {
             hir::ImplItemKind::Const(..) => {
-                self.prefetch_mir(tcx.hir().local_def_id(impl_item.hir_id))
+                self.prefetch_ctfe_mir(tcx.hir().local_def_id(impl_item.hir_id))
             }
             hir::ImplItemKind::Fn(ref sig, _) => {
                 let def_id = tcx.hir().local_def_id(impl_item.hir_id);
                 let generics = tcx.generics_of(def_id.to_def_id());
                 let needs_inline = generics.requires_monomorphization(tcx)
                     || tcx.codegen_fn_attrs(def_id.to_def_id()).requests_inline();
                 let is_const_fn = sig.header.constness == hir::Constness::Const;
-                if needs_inline || is_const_fn {
+                if needs_inline {
                     self.prefetch_mir(def_id)
                 }
+                if is_const_fn {
+                    self.prefetch_ctfe_mir(def_id);
+                }
             }
             hir::ImplItemKind::TyAlias(..) => (),
         }

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -302,6 +302,7 @@ define_tables! {
     // As an optimization, a missing entry indicates an empty `&[]`.
     explicit_item_bounds: Table<DefIndex, Lazy!([(ty::Predicate<'tcx>, Span)])>,
     mir: Table<DefIndex, Lazy!(mir::Body<'tcx>)>,
+    mir_for_ctfe: Table<DefIndex, Lazy!(mir::Body<'tcx>)>,
     promoted_mir: Table<DefIndex, Lazy!(IndexVec<mir::Promoted, mir::Body<'tcx>>)>,
     mir_abstract_consts: Table<DefIndex, Lazy!(&'tcx [mir::abstract_const::Node<'tcx>])>,
     unused_generic_params: Table<DefIndex, Lazy<FiniteBitSet<u32>>>,

compiler/rustc_middle/src/mir/query.rs

@@ -439,17 +439,26 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 
     #[inline]
-    pub fn optimized_mir_opt_const_arg(
+    pub fn optimized_mir_or_const_arg_mir(
         self,
         def: ty::WithOptConstParam<DefId>,
     ) -> &'tcx Body<'tcx> {
         if let Some((did, param_did)) = def.as_const_arg() {
-            self.optimized_mir_of_const_arg((did, param_did))
+            self.mir_for_ctfe_of_const_arg((did, param_did))
         } else {
             self.optimized_mir(def.did)
         }
     }
 
+    #[inline]
+    pub fn mir_for_ctfe_opt_const_arg(self, def: ty::WithOptConstParam<DefId>) -> &'tcx Body<'tcx> {
+        if let Some((did, param_did)) = def.as_const_arg() {
+            self.mir_for_ctfe_of_const_arg((did, param_did))
+        } else {
+            self.mir_for_ctfe(def.did)
+        }
+    }
+
     #[inline]
     pub fn mir_abstract_const_opt_const_arg(
         self,

compiler/rustc_middle/src/query/mod.rs

@@ -312,6 +312,20 @@ rustc_queries! {
             desc { |tcx| "elaborating drops for `{}`", tcx.def_path_str(key.did.to_def_id()) }
         }
 
+        query mir_for_ctfe(
+            key: DefId
+        ) -> &'tcx mir::Body<'tcx> {
+            desc { |tcx| "caching mir of `{}` for CTFE", tcx.def_path_str(key) }
+            cache_on_disk_if { key.is_local() }
+        }
+
+        query mir_for_ctfe_of_const_arg(key: (LocalDefId, DefId)) -> &'tcx mir::Body<'tcx> {
+            desc {
+                |tcx| "MIR for CTFE of the const argument `{}`",
+                tcx.def_path_str(key.0.to_def_id())
+            }
+        }
+
         query mir_promoted(key: ty::WithOptConstParam<LocalDefId>) ->
             (
                 &'tcx Steal<mir::Body<'tcx>>,
@@ -331,12 +345,6 @@ rustc_queries! {
             desc { |tcx| "optimizing MIR for `{}`", tcx.def_path_str(key) }
             cache_on_disk_if { key.is_local() }
         }
-        query optimized_mir_of_const_arg(key: (LocalDefId, DefId)) -> &'tcx mir::Body<'tcx> {
-            desc {
-                |tcx| "optimizing MIR for the const argument `{}`",
-                tcx.def_path_str(key.0.to_def_id())
-            }
-        }
 
         /// Returns coverage summary info for a function, after executing the `InstrumentCoverage`
         /// MIR pass (assuming the -Zinstrument-coverage option is enabled).
@@ -927,6 +935,9 @@ rustc_queries! {
     }
 
     Codegen {
+        query is_ctfe_mir_available(key: DefId) -> bool {
+            desc { |tcx| "checking if item has ctfe mir available: `{}`", tcx.def_path_str(key) }
+        }
         query is_mir_available(key: DefId) -> bool {
             desc { |tcx| "checking if item has mir available: `{}`", tcx.def_path_str(key) }
         }

compiler/rustc_middle/src/ty/mod.rs

@@ -3010,7 +3010,16 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Returns the possibly-auto-generated MIR of a `(DefId, Subst)` pair.
     pub fn instance_mir(self, instance: ty::InstanceDef<'tcx>) -> &'tcx Body<'tcx> {
         match instance {
-            ty::InstanceDef::Item(def) => self.optimized_mir_opt_const_arg(def),
+            ty::InstanceDef::Item(def) => match self.def_kind(def.did) {
+                DefKind::Const
+                | DefKind::Static
+                | DefKind::AssocConst
+                | DefKind::Ctor(..)
+                | DefKind::AnonConst => self.mir_for_ctfe_opt_const_arg(def),
+                // If the caller wants `mir_for_ctfe` of a function they should not be using
+                // `instance_mir`, so we'll assume const fn also wants the optimized version.
+                _ => self.optimized_mir_or_const_arg_mir(def),
+            },
             ty::InstanceDef::VtableShim(..)
             | ty::InstanceDef::ReifyShim(..)
             | ty::InstanceDef::Intrinsic(..)