Auto merge of #157275 - RalfJung:interpret-layout-cache, r=<try>

interpret: add per-interpreter layout cache

Author: bors

compiler/rustc_const_eval/src/interpret/eval_context.rs

@@ -1,5 +1,9 @@
+use std::cell::RefCell;
+use std::collections::hash_map::Entry;
+
 use either::{Left, Right};
 use rustc_abi::{Align, HasDataLayout, Size, TargetDataLayout};
+use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::def_id::DefId;
 use rustc_hir::limit::Limit;
 use rustc_middle::mir::interpret::{ErrorHandled, InvalidMetaKind, ReportedErrorInfo};
@@ -39,6 +43,9 @@ pub struct InterpCx<'tcx, M: Machine<'tcx>> {
     /// polymorphic context. This always uses `ty::TypingMode::PostAnalysis`.
     pub(super) typing_env: ty::TypingEnv<'tcx>,
 
+    /// The query cache is slow so we have our own cache in front of it.
+    pub(super) layout_cache: RefCell<FxHashMap<Ty<'tcx>, rustc_abi::Layout<'tcx>>>,
+
     /// The virtual memory system.
     pub memory: Memory<'tcx, M>,
 
@@ -131,10 +138,19 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
     /// This inherent method takes priority over the trait method with the same name in LayoutOf,
     /// and allows wrapping the actual [LayoutOf::layout_of] with a tracing span.
     /// See [LayoutOf::layout_of] for the original documentation.
-    #[inline(always)]
+    #[inline]
     pub fn layout_of(&self, ty: Ty<'tcx>) -> Result<TyAndLayout<'tcx>, InterpErrorKind<'tcx>> {
-        let _trace = enter_trace_span!(M, layouting::layout_of, ty = ?ty.kind());
-        LayoutOf::layout_of(self, ty)
+        match self.layout_cache.borrow_mut().entry(ty) {
+            Entry::Occupied(occupied_entry) => {
+                Ok(TyAndLayout { ty, layout: *occupied_entry.get() })
+            }
+            Entry::Vacant(vacant_entry) => {
+                let _trace = enter_trace_span!(M, layouting::layout_of, ty = ?ty.kind());
+                let layout = LayoutOf::layout_of(self, ty)?;
+                vacant_entry.insert(layout.layout);
+                Ok(layout)
+            }
+        }
     }
 
     /// This inherent method takes priority over the trait method with the same name in FnAbiOf,
@@ -258,6 +274,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             machine,
             tcx: tcx.at(root_span),
             typing_env,
+            layout_cache: RefCell::new(FxHashMap::default()),
             memory: Memory::new(),
             recursion_limit: tcx.recursion_limit(),
         }

compiler/rustc_const_eval/src/interpret/stack.rs

@@ -1,6 +1,5 @@
 //! Manages the low-level pushing and popping of stack frames and the (de)allocation of local variables.
 //! For handling of argument passing and return values, see the `call` module.
-use std::cell::Cell;
 use std::{fmt, mem};
 
 use either::{Either, Left, Right};
@@ -90,7 +89,7 @@ pub struct Frame<'tcx, Prov: Provenance = CtfeProvenance, Extra = ()> {
     /// can either directly contain `Scalar` or refer to some part of an `Allocation`.
     ///
     /// Do *not* access this directly; always go through the machine hook!
-    pub locals: IndexVec<mir::Local, LocalState<'tcx, Prov>>,
+    pub locals: IndexVec<mir::Local, LocalState<Prov>>,
 
     /// The complete variable argument list of this frame. Its elements must be dropped when the
     /// frame is popped.
@@ -129,19 +128,13 @@ pub enum ReturnContinuation {
 
 /// State of a local variable including a memoized layout
 #[derive(Clone)]
-pub struct LocalState<'tcx, Prov: Provenance = CtfeProvenance> {
+pub struct LocalState<Prov: Provenance = CtfeProvenance> {
     value: LocalValue<Prov>,
-    /// Don't modify if `Some`, this is only used to prevent computing the layout twice.
-    /// Avoids computing the layout of locals that are never actually initialized.
-    layout: Cell<Option<TyAndLayout<'tcx>>>,
 }
 
-impl<Prov: Provenance> std::fmt::Debug for LocalState<'_, Prov> {
+impl<Prov: Provenance> std::fmt::Debug for LocalState<Prov> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("LocalState")
-            .field("value", &self.value)
-            .field("ty", &self.layout.get().map(|l| l.ty))
-            .finish()
+        f.debug_struct("LocalState").field("value", &self.value).finish()
     }
 }
 
@@ -161,7 +154,7 @@ pub(super) enum LocalValue<Prov: Provenance = CtfeProvenance> {
     Live(Operand<Prov>),
 }
 
-impl<'tcx, Prov: Provenance> LocalState<'tcx, Prov> {
+impl<'tcx, Prov: Provenance> LocalState<Prov> {
     pub fn make_live_uninit(&mut self) {
         self.value = LocalValue::Live(Operand::Immediate(Immediate::Uninit));
     }
@@ -377,7 +370,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
 
         // First push a stack frame so we have access to `instantiate_from_current_frame` and other
         // `self.frame()`-based functions.
-        let dead_local = LocalState { value: LocalValue::Dead, layout: Cell::new(None) };
+        let dead_local = LocalState { value: LocalValue::Dead };
         let locals = IndexVec::from_elem(dead_local, &body.local_decls);
         let pre_frame = Frame {
             body,
@@ -613,21 +606,12 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         local: mir::Local,
         layout: Option<TyAndLayout<'tcx>>,
     ) -> InterpResult<'tcx, TyAndLayout<'tcx>> {
-        let state = &frame.locals[local];
-        if let Some(layout) = state.layout.get() {
-            return interp_ok(layout);
-        }
-
-        let layout = from_known_layout(self.tcx, self.typing_env, layout, || {
+        from_known_layout(self.tcx, self.typing_env, layout, || {
             let local_ty = frame.body.local_decls[local].ty;
             let local_ty =
                 self.instantiate_from_frame_and_normalize_erasing_regions(frame, local_ty)?;
             self.layout_of(local_ty).into()
-        })?;
-
-        // Layouts of locals are requested a lot, so we cache them.
-        state.layout.set(Some(layout));
-        interp_ok(layout)
+        })
     }
 }
 
@@ -683,7 +667,7 @@ impl<'a, 'tcx: 'a, M: Machine<'tcx>> InterpCx<'tcx, M> {
     }
 }
 
-impl<'tcx, Prov: Provenance> LocalState<'tcx, Prov> {
+impl<Prov: Provenance> LocalState<Prov> {
     pub(super) fn print(
         &self,
         allocs: &mut Vec<Option<AllocId>>,