Rollup merge of #55016 - oli-obk:vtables💥_vtables_everywhere, r=RalfJung

Deduplicate some code and compile-time values around vtables

r? @RalfJung

5 files changed, 36 insertions, 19 deletions

diff --git a/src/librustc/ty/query/mod.rs b/src/librustc/ty/query/mod.rs
index c4f39ffcd20..14e4ddfcdd6 100644
--- a/src/librustc/ty/query/mod.rs
+++ b/src/librustc/ty/query/mod.rs
@@ -369,16 +369,16 @@ define_queries! { <'tcx>
             -> Lrc<specialization_graph::Graph>,
         [] fn is_object_safe: ObjectSafety(DefId) -> bool,
 
-        // Get the ParameterEnvironment for a given item; this environment
-        // will be in "user-facing" mode, meaning that it is suitabe for
-        // type-checking etc, and it does not normalize specializable
-        // associated types. This is almost always what you want,
-        // unless you are doing MIR optimizations, in which case you
-        // might want to use `reveal_all()` method to change modes.
+        /// Get the ParameterEnvironment for a given item; this environment
+        /// will be in "user-facing" mode, meaning that it is suitabe for
+        /// type-checking etc, and it does not normalize specializable
+        /// associated types. This is almost always what you want,
+        /// unless you are doing MIR optimizations, in which case you
+        /// might want to use `reveal_all()` method to change modes.
         [] fn param_env: ParamEnv(DefId) -> ty::ParamEnv<'tcx>,
 
-        // Trait selection queries. These are best used by invoking `ty.moves_by_default()`,
-        // `ty.is_copy()`, etc, since that will prune the environment where possible.
+        /// Trait selection queries. These are best used by invoking `ty.moves_by_default()`,
+        /// `ty.is_copy()`, etc, since that will prune the environment where possible.
         [] fn is_copy_raw: is_copy_dep_node(ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool,
         [] fn is_sized_raw: is_sized_dep_node(ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool,
         [] fn is_freeze_raw: is_freeze_dep_node(ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool,
diff --git a/src/librustc_codegen_llvm/meth.rs b/src/librustc_codegen_llvm/meth.rs
index db06b87f44e..29c2e71960c 100644
--- a/src/librustc_codegen_llvm/meth.rs
+++ b/src/librustc_codegen_llvm/meth.rs
@@ -94,6 +94,10 @@ pub fn get_vtable(
     });
 
     let (size, align) = cx.size_and_align_of(ty);
+    // /////////////////////////////////////////////////////////////////////////////////////////////
+    // If you touch this code, be sure to also make the corresponding changes to
+    // `get_vtable` in rust_mir/interpret/traits.rs
+    // /////////////////////////////////////////////////////////////////////////////////////////////
     let components: Vec<_> = [
         callee::get_fn(cx, monomorphize::resolve_drop_in_place(cx.tcx, ty)),
         C_usize(cx, size.bytes()),
diff --git a/src/librustc_mir/interpret/cast.rs b/src/librustc_mir/interpret/cast.rs
index f4ddfa5293e..f5e824b7628 100644
--- a/src/librustc_mir/interpret/cast.rs
+++ b/src/librustc_mir/interpret/cast.rs
@@ -327,12 +327,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M>
             }
             (_, &ty::Dynamic(ref data, _)) => {
                 // Initial cast from sized to dyn trait
-                let trait_ref = data.principal().with_self_ty(
-                    *self.tcx,
-                    src_pointee_ty,
-                );
-                let trait_ref = self.tcx.erase_regions(&trait_ref);
-                let vtable = self.get_vtable(src_pointee_ty, trait_ref)?;
+                let vtable = self.get_vtable(src_pointee_ty, data.principal())?;
                 let ptr = self.read_value(src)?.to_scalar_ptr()?;
                 let val = Value::new_dyn_trait(ptr, vtable);
                 self.write_value(val, dest)
diff --git a/src/librustc_mir/interpret/eval_context.rs b/src/librustc_mir/interpret/eval_context.rs
index 85a8376134a..cf5358a9896 100644
--- a/src/librustc_mir/interpret/eval_context.rs
+++ b/src/librustc_mir/interpret/eval_context.rs
@@ -27,6 +27,7 @@ use rustc::mir::interpret::{
     EvalResult, EvalErrorKind,
     truncate, sign_extend,
 };
+use rustc_data_structures::fx::FxHashMap;
 
 use syntax::source_map::{self, Span};
 
@@ -50,6 +51,9 @@ pub struct EvalContext<'a, 'mir, 'tcx: 'a + 'mir, M: Machine<'a, 'mir, 'tcx>> {
 
     /// The virtual call stack.
     pub(crate) stack: Vec<Frame<'mir, 'tcx, M::PointerTag>>,
+
+    /// A cache for deduplicating vtables
+    pub(super) vtables: FxHashMap<(Ty<'tcx>, ty::PolyExistentialTraitRef<'tcx>), AllocId>,
 }
 
 /// A stack frame.
@@ -209,6 +213,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tc
             param_env,
             memory: Memory::new(tcx, memory_data),
             stack: Vec::new(),
+            vtables: FxHashMap::default(),
         }
     }
 
diff --git a/src/librustc_mir/interpret/traits.rs b/src/librustc_mir/interpret/traits.rs
index 227c85772d2..2b0febc1ce7 100644
--- a/src/librustc_mir/interpret/traits.rs
+++ b/src/librustc_mir/interpret/traits.rs
@@ -24,20 +24,32 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M>
     pub fn get_vtable(
         &mut self,
         ty: Ty<'tcx>,
-        trait_ref: ty::PolyTraitRef<'tcx>,
+        poly_trait_ref: ty::PolyExistentialTraitRef<'tcx>,
     ) -> EvalResult<'tcx, Pointer<M::PointerTag>> {
-        debug!("get_vtable(trait_ref={:?})", trait_ref);
+        debug!("get_vtable(trait_ref={:?})", poly_trait_ref);
 
-        // FIXME: Cache this!
+        let (ty, poly_trait_ref) = self.tcx.erase_regions(&(ty, poly_trait_ref));
 
-        let layout = self.layout_of(trait_ref.self_ty())?;
+        if let Some(&vtable) = self.vtables.get(&(ty, poly_trait_ref)) {
+            return Ok(Pointer::from(vtable).with_default_tag());
+        }
+
+        let trait_ref = poly_trait_ref.with_self_ty(*self.tcx, ty);
+        let trait_ref = self.tcx.erase_regions(&trait_ref);
+
+        let methods = self.tcx.vtable_methods(trait_ref);
+
+        let layout = self.layout_of(ty)?;
         assert!(!layout.is_unsized(), "can't create a vtable for an unsized type");
         let size = layout.size.bytes();
         let align = layout.align.abi();
 
         let ptr_size = self.pointer_size();
         let ptr_align = self.tcx.data_layout.pointer_align;
-        let methods = self.tcx.vtable_methods(trait_ref);
+        // /////////////////////////////////////////////////////////////////////////////////////////
+        // If you touch this code, be sure to also make the corresponding changes to
+        // `get_vtable` in rust_codegen_llvm/meth.rs
+        // /////////////////////////////////////////////////////////////////////////////////////////
         let vtable = self.memory.allocate(
             ptr_size * (3 + methods.len() as u64),
             ptr_align,
@@ -64,6 +76,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M>
         }
 
         self.memory.mark_immutable(vtable.alloc_id)?;
+        assert!(self.vtables.insert((ty, poly_trait_ref), vtable.alloc_id).is_none());
 
         Ok(vtable)
     }