Skip to content

Commit

Permalink
Add cache for allocate_str
Browse files Browse the repository at this point in the history
  • Loading branch information
adwinwhite committed Jul 14, 2024
1 parent 03c2100 commit e595f3d
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 10 deletions.
14 changes: 13 additions & 1 deletion compiler/rustc_const_eval/src/interpret/place.rs
Original file line number Diff line number Diff line change
Expand Up @@ -995,13 +995,25 @@ where
}

/// 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<M::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))
Expand Down
17 changes: 11 additions & 6 deletions compiler/rustc_middle/src/mir/interpret/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -393,7 +393,6 @@ pub(crate) struct AllocMap<'tcx> {
alloc_map: FxHashMap<AllocId, GlobalAlloc<'tcx>>,

/// 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<GlobalAlloc<'tcx>, AllocId>,
Expand Down Expand Up @@ -433,13 +432,13 @@ impl<'tcx> TyCtxt<'tcx> {
}

/// 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;
Expand All @@ -451,6 +450,12 @@ impl<'tcx> TyCtxt<'tcx> {
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 {
Expand Down
5 changes: 3 additions & 2 deletions compiler/rustc_middle/src/ty/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1442,11 +1442,12 @@ impl<'tcx> TyCtxt<'tcx> {
}

/// 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
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_mir_build/src/build/expr/as_constant.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)) =>
Expand Down

0 comments on commit e595f3d

Please sign in to comment.