From e595f3d13f0491186577d7b6e0290f31791f5059 Mon Sep 17 00:00:00 2001 From: Adwin White Date: Fri, 12 Jul 2024 13:08:05 +0800 Subject: [PATCH] Add cache for `allocate_str` --- .../rustc_const_eval/src/interpret/place.rs | 14 +++++++++++++- compiler/rustc_middle/src/mir/interpret/mod.rs | 17 +++++++++++------ compiler/rustc_middle/src/ty/context.rs | 5 +++-- .../src/build/expr/as_constant.rs | 2 +- 4 files changed, 28 insertions(+), 10 deletions(-) diff --git a/compiler/rustc_const_eval/src/interpret/place.rs b/compiler/rustc_const_eval/src/interpret/place.rs index baaee67e787..33c25b746cc 100644 --- a/compiler/rustc_const_eval/src/interpret/place.rs +++ b/compiler/rustc_const_eval/src/interpret/place.rs @@ -995,13 +995,25 @@ pub fn allocate( } /// Returns a wide MPlace of type `str` to a new 1-aligned allocation. + /// Immutable strings are deduplicated and stored in global memory. pub fn allocate_str( &mut self, str: &str, kind: MemoryKind, mutbl: Mutability, ) -> InterpResult<'tcx, MPlaceTy<'tcx, M::Provenance>> { - let ptr = self.allocate_bytes_ptr(str.as_bytes(), Align::ONE, kind, mutbl)?; + let tcx = self.tcx.tcx; + + // Use cache for immutable strings. + let ptr = if mutbl.is_not() { + // Use dedup'd allocation function. + let id = tcx.allocate_bytes_dedup(str.as_bytes()); + + // Turn untagged "global" pointers (obtained via `tcx`) into the machine pointer to the allocation. + M::adjust_alloc_root_pointer(&self, Pointer::from(id), Some(kind))? + } else { + self.allocate_bytes_ptr(str.as_bytes(), Align::ONE, kind, mutbl)? + }; let meta = Scalar::from_target_usize(u64::try_from(str.len()).unwrap(), self); let layout = self.layout_of(self.tcx.types.str_).unwrap(); Ok(self.ptr_with_meta_to_mplace(ptr.into(), MemPlaceMeta::Meta(meta), layout)) diff --git a/compiler/rustc_middle/src/mir/interpret/mod.rs b/compiler/rustc_middle/src/mir/interpret/mod.rs index 4e95e600b5a..bdd1eb11a38 100644 --- a/compiler/rustc_middle/src/mir/interpret/mod.rs +++ b/compiler/rustc_middle/src/mir/interpret/mod.rs @@ -393,7 +393,6 @@ pub(crate) struct AllocMap<'tcx> { alloc_map: FxHashMap>, /// Used to ensure that statics and functions only get one associated `AllocId`. - /// Should never contain a `GlobalAlloc::Memory`! // // FIXME: Should we just have two separate dedup maps for statics and functions each? dedup: FxHashMap, AllocId>, @@ -433,13 +432,13 @@ pub fn reserve_alloc_id(self) -> AllocId { } /// Reserves a new ID *if* this allocation has not been dedup-reserved before. - /// Should only be used for "symbolic" allocations (function pointers, vtables, statics), we - /// don't want to dedup IDs for "real" memory! + /// Should not be used for mutable memory. fn reserve_and_set_dedup(self, alloc: GlobalAlloc<'tcx>) -> AllocId { let mut alloc_map = self.alloc_map.lock(); - match alloc { - GlobalAlloc::Function { .. } | GlobalAlloc::Static(..) | GlobalAlloc::VTable(..) => {} - GlobalAlloc::Memory(..) => bug!("Trying to dedup-reserve memory with real data!"), + if let GlobalAlloc::Memory(mem) = alloc { + if mem.inner().mutability.is_mut() { + bug!("trying to dedup-reserve mutable memory"); + } } if let Some(&alloc_id) = alloc_map.dedup.get(&alloc) { return alloc_id; @@ -451,6 +450,12 @@ fn reserve_and_set_dedup(self, alloc: GlobalAlloc<'tcx>) -> AllocId { id } + /// Generates an `AllocId` for a memory allocation. If the exact same memory has been + /// allocated before, this will return the same `AllocId`. + pub fn reserve_and_set_memory_dedup(self, mem: ConstAllocation<'tcx>) -> AllocId { + self.reserve_and_set_dedup(GlobalAlloc::Memory(mem)) + } + /// Generates an `AllocId` for a static or return a cached one in case this function has been /// called on the same static before. pub fn reserve_and_set_static_alloc(self, static_id: DefId) -> AllocId { diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs index 9e24ea485b2..25070e6b042 100644 --- a/compiler/rustc_middle/src/ty/context.rs +++ b/compiler/rustc_middle/src/ty/context.rs @@ -1442,11 +1442,12 @@ pub fn mk_adt_def( } /// Allocates a read-only byte or string literal for `mir::interpret`. - pub fn allocate_bytes(self, bytes: &[u8]) -> interpret::AllocId { + /// Returns the same `AllocId` if called again with the same bytes. + pub fn allocate_bytes_dedup(self, bytes: &[u8]) -> interpret::AllocId { // Create an allocation that just contains these bytes. let alloc = interpret::Allocation::from_bytes_byte_aligned_immutable(bytes); let alloc = self.mk_const_alloc(alloc); - self.reserve_and_set_memory_alloc(alloc) + self.reserve_and_set_memory_dedup(alloc) } /// Returns a range of the start/end indices specified with the diff --git a/compiler/rustc_mir_build/src/build/expr/as_constant.rs b/compiler/rustc_mir_build/src/build/expr/as_constant.rs index 3b69058d3cb..be62a3d3736 100644 --- a/compiler/rustc_mir_build/src/build/expr/as_constant.rs +++ b/compiler/rustc_mir_build/src/build/expr/as_constant.rs @@ -140,7 +140,7 @@ fn lit_to_mir_constant<'tcx>( ConstValue::Slice { data: allocation, meta: allocation.inner().size().bytes() } } (ast::LitKind::ByteStr(data, _), ty::Ref(_, inner_ty, _)) if inner_ty.is_array() => { - let id = tcx.allocate_bytes(data); + let id = tcx.allocate_bytes_dedup(data); ConstValue::Scalar(Scalar::from_pointer(id.into(), &tcx)) } (ast::LitKind::CStr(data, _), ty::Ref(_, inner_ty, _)) if matches!(inner_ty.kind(), ty::Adt(def, _) if tcx.is_lang_item(def.did(), LangItem::CStr)) =>