Skip to content

Commit

Permalink
Remove in-progress allocation decoding states
Browse files Browse the repository at this point in the history
  • Loading branch information
saethlin committed Jul 17, 2024
1 parent eb72697 commit 3623c76
Showing 1 changed file with 13 additions and 89 deletions.
102 changes: 13 additions & 89 deletions compiler/rustc_middle/src/mir/interpret/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,13 @@ use std::fmt;
use std::io;
use std::io::{Read, Write};
use std::num::NonZero;
use std::sync::atomic::{AtomicU32, Ordering};

use smallvec::{smallvec, SmallVec};
use tracing::{debug, trace};

use rustc_ast::LitKind;
use rustc_attr::InlineAttr;
use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::sync::{HashMapExt, Lock};
use rustc_data_structures::sync::Lock;
use rustc_errors::ErrorGuaranteed;
use rustc_hir::def_id::{DefId, LocalDefId};
use rustc_macros::{HashStable, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable};
Expand Down Expand Up @@ -159,14 +157,9 @@ pub fn specialized_encode_alloc_id<'tcx, E: TyEncoder<I = TyCtxt<'tcx>>>(
}
}

// Used to avoid infinite recursion when decoding cyclic allocations.
type DecodingSessionId = NonZero<u32>;

#[derive(Clone)]
enum State {
Empty,
InProgressNonAlloc(SmallVec<[DecodingSessionId; 1]>),
InProgress(SmallVec<[DecodingSessionId; 1]>, AllocId),
Done(AllocId),
}

Expand All @@ -180,13 +173,7 @@ pub struct AllocDecodingState {
impl AllocDecodingState {
#[inline]
pub fn new_decoding_session(&self) -> AllocDecodingSession<'_> {
static DECODER_SESSION_ID: AtomicU32 = AtomicU32::new(0);
let counter = DECODER_SESSION_ID.fetch_add(1, Ordering::SeqCst);

// Make sure this is never zero.
let session_id = DecodingSessionId::new((counter & 0x7FFFFFFF) + 1).unwrap();

AllocDecodingSession { state: self, session_id }
AllocDecodingSession { state: self }
}

pub fn new(data_offsets: Vec<u64>) -> Self {
Expand All @@ -200,7 +187,6 @@ impl AllocDecodingState {
#[derive(Copy, Clone)]
pub struct AllocDecodingSession<'s> {
state: &'s AllocDecodingState,
session_id: DecodingSessionId,
}

impl<'s> AllocDecodingSession<'s> {
Expand All @@ -222,104 +208,48 @@ impl<'s> AllocDecodingSession<'s> {

// Check the decoding state to see if it's already decoded or if we should
// decode it here.
let alloc_id = {
let mut entry = self.state.decoding_state[idx].lock();

match *entry {
State::Done(alloc_id) => {
return alloc_id;
}
ref mut entry @ State::Empty => {
// We are allowed to decode.
match alloc_kind {
AllocDiscriminant::Alloc => {
// If this is an allocation, we need to reserve an
// `AllocId` so we can decode cyclic graphs.
let alloc_id = decoder.interner().reserve_alloc_id();
*entry = State::InProgress(smallvec![self.session_id], alloc_id);
Some(alloc_id)
}
AllocDiscriminant::Fn
| AllocDiscriminant::Static
| AllocDiscriminant::VTable => {
// Fns and statics cannot be cyclic, and their `AllocId`
// is determined later by interning.
*entry = State::InProgressNonAlloc(smallvec![self.session_id]);
None
}
}
}
State::InProgressNonAlloc(ref mut sessions) => {
if sessions.contains(&self.session_id) {
bug!("this should be unreachable");
} else {
// Start decoding concurrently.
sessions.push(self.session_id);
None
}
}
State::InProgress(ref mut sessions, alloc_id) => {
if sessions.contains(&self.session_id) {
// Don't recurse.
return alloc_id;
} else {
// Start decoding concurrently.
sessions.push(self.session_id);
Some(alloc_id)
}
}
}
};
let mut entry = self.state.decoding_state[idx].lock();
if let State::Done(alloc_id) = *entry {
return alloc_id;
}

// Now decode the actual data.
let alloc_id = decoder.with_position(pos, |decoder| {
match alloc_kind {
AllocDiscriminant::Alloc => {
trace!("creating memory alloc ID");
let alloc = <ConstAllocation<'tcx> as Decodable<_>>::decode(decoder);
// We already have a reserved `AllocId`.
let alloc_id = alloc_id.unwrap();
trace!("decoded alloc {:?}: {:#?}", alloc_id, alloc);
decoder.interner().set_alloc_id_same_memory(alloc_id, alloc);
alloc_id
trace!("decoded alloc {:?}", alloc);
decoder.interner().reserve_and_set_memory_alloc(alloc)
}
AllocDiscriminant::Fn => {
assert!(alloc_id.is_none());
trace!("creating fn alloc ID");
let instance = ty::Instance::decode(decoder);
trace!("decoded fn alloc instance: {:?}", instance);
let unique = bool::decode(decoder);
// Here we cannot call `reserve_and_set_fn_alloc` as that would use a query, which
// is not possible in this context. That's why the allocation stores
// whether it is unique or not.
let alloc_id =
decoder.interner().reserve_and_set_fn_alloc_internal(instance, unique);
alloc_id
decoder.interner().reserve_and_set_fn_alloc_internal(instance, unique)
}
AllocDiscriminant::VTable => {
assert!(alloc_id.is_none());
trace!("creating vtable alloc ID");
let ty = <Ty<'_> as Decodable<D>>::decode(decoder);
let poly_trait_ref =
<Option<ty::PolyExistentialTraitRef<'_>> as Decodable<D>>::decode(decoder);
trace!("decoded vtable alloc instance: {ty:?}, {poly_trait_ref:?}");
let alloc_id =
decoder.interner().reserve_and_set_vtable_alloc(ty, poly_trait_ref);
alloc_id
decoder.interner().reserve_and_set_vtable_alloc(ty, poly_trait_ref)
}
AllocDiscriminant::Static => {
assert!(alloc_id.is_none());
trace!("creating extern static alloc ID");
let did = <DefId as Decodable<D>>::decode(decoder);
trace!("decoded static def-ID: {:?}", did);
let alloc_id = decoder.interner().reserve_and_set_static_alloc(did);
alloc_id
decoder.interner().reserve_and_set_static_alloc(did)
}
}
});

self.state.decoding_state[idx].with_lock(|entry| {
*entry = State::Done(alloc_id);
});
*entry = State::Done(alloc_id);

alloc_id
}
Expand Down Expand Up @@ -558,12 +488,6 @@ impl<'tcx> TyCtxt<'tcx> {
bug!("tried to set allocation ID {id:?}, but it was already existing as {old:#?}");
}
}

/// Freezes an `AllocId` created with `reserve` by pointing it at an `Allocation`. May be called
/// twice for the same `(AllocId, Allocation)` pair.
fn set_alloc_id_same_memory(self, id: AllocId, mem: ConstAllocation<'tcx>) {
self.alloc_map.lock().alloc_map.insert_same(id, GlobalAlloc::Memory(mem));
}
}

////////////////////////////////////////////////////////////////////////////////
Expand Down

0 comments on commit 3623c76

Please sign in to comment.