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 16, 2024
1 parent eb72697 commit 963e157
Showing 1 changed file with 20 additions and 60 deletions.
80 changes: 20 additions & 60 deletions compiler/rustc_middle/src/mir/interpret/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,7 @@ 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;
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,52 +208,28 @@ 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
}
let mut entry = self.state.decoding_state[idx].lock();
let alloc_id = match *entry {
State::Done(alloc_id) => {
return alloc_id;
}
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();
Some(alloc_id)
}
}
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);
AllocDiscriminant::Fn
| AllocDiscriminant::Static
| AllocDiscriminant::VTable => {
// Fns and statics cannot be cyclic, and their `AllocId`
// is determined later by interning.
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)
}
}
}
};

Expand Down Expand Up @@ -317,9 +279,7 @@ impl<'s> AllocDecodingSession<'s> {
}
});

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

alloc_id
}
Expand Down

0 comments on commit 963e157

Please sign in to comment.