Skip to content

Commit d954310

Browse files
youknowoneyouknowone
committed
working
1 parent 1c30924 commit d954310

File tree

4 files changed

+97
-52
lines changed

4 files changed

+97
-52
lines changed

‎crates/common/src/refcount.rs‎

Lines changed: 27 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -27,18 +27,17 @@ pub const EPOCH: u64 = if EPOCH_WIDTH == 0 {
2727
};
2828
pub const DESTRUCTED: u64 = 1 << (EPOCH_MASK_HEIGHT - 1);
2929
pub const WEAKED: u64 = 1 << (EPOCH_MASK_HEIGHT - 2);
30-
pub const TOTAL_COUNT_WIDTH: u32 = u64::BITS - EPOCH_WIDTH - 2;
30+
/// LEAKED bit for interned objects (never deallocated)
31+
pub const LEAKED: u64 = 1 << (EPOCH_MASK_HEIGHT - 3);
32+
// 3 flag bits: DESTRUCTED, WEAKED, LEAKED
33+
pub const TOTAL_COUNT_WIDTH: u32 = u64::BITS - EPOCH_WIDTH - 3;
3134
pub const WEAK_WIDTH: u32 = TOTAL_COUNT_WIDTH / 2;
3235
pub const STRONG_WIDTH: u32 = TOTAL_COUNT_WIDTH - WEAK_WIDTH;
3336
pub const STRONG: u64 = (1 << STRONG_WIDTH) - 1;
3437
pub const WEAK: u64 = ((1 << WEAK_WIDTH) - 1) << STRONG_WIDTH;
3538
pub const COUNT: u64 = 1;
3639
pub const WEAK_COUNT: u64 = 1 << STRONG_WIDTH;
3740

38-
/// LEAKED bit for interned objects (never deallocated)
39-
/// Position: just after WEAKED bit
40-
pub const LEAKED: u64 = 1 << (EPOCH_MASK_HEIGHT - 3);
41-
4241
/// State wraps reference count + flags in a single 64-bit word
4342
#[derive(Clone, Copy)]
4443
pub struct State {
@@ -206,27 +205,38 @@ thread_local! {
206205
static DEFERRED_QUEUE: RefCell<Vec<Box<dyn FnOnce()>>> = const { RefCell::new(Vec::new()) };
207206
}
208207

208+
/// RAII guard for deferred drop context.
209+
/// Restores the previous context state on drop, even if a panic occurs.
210+
struct DeferredDropGuard {
211+
was_in_context: bool,
212+
}
213+
214+
impl Drop for DeferredDropGuard {
215+
fn drop(&mut self) {
216+
IN_DEFERRED_CONTEXT.with(|in_ctx| {
217+
in_ctx.set(self.was_in_context);
218+
});
219+
// Only flush if we're the outermost context
220+
if !self.was_in_context {
221+
flush_deferred_drops();
222+
}
223+
}
224+
}
225+
209226
/// Execute a function within a deferred drop context.
210227
/// Any calls to `try_defer_drop` within this context will be queued
211-
/// and executed when the context exits.
228+
/// and executed when the context exits (even on panic).
212229
#[inline]
213230
pub fn with_deferred_drops<F, R>(f: F) -> R
214231
where
215232
F: FnOnce() -> R,
216233
{
217-
IN_DEFERRED_CONTEXT.with(|in_ctx| {
234+
let _guard = IN_DEFERRED_CONTEXT.with(|in_ctx| {
218235
let was_in_context = in_ctx.get();
219236
in_ctx.set(true);
220-
let result = f();
221-
in_ctx.set(was_in_context);
222-
223-
// Only flush if we're the outermost context
224-
if !was_in_context {
225-
flush_deferred_drops();
226-
}
227-
228-
result
229-
})
237+
DeferredDropGuard { was_in_context }
238+
});
239+
f()
230240
}
231241

232242
/// Try to defer a drop-related operation.

‎crates/stdlib/src/gc.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ mod gc {
5050
fn collect(args: CollectArgs, vm: &VirtualMachine) -> PyResult<i32> {
5151
let generation = args.generation;
5252
let generation_num = generation.unwrap_or(2);
53-
if generation_num < 0 || generation_num > 2 {
53+
if !(0..=2).contains(&generation_num) {
5454
return Err(vm.new_value_error("invalid generation".to_owned()));
5555
}
5656

@@ -164,12 +164,12 @@ mod gc {
164164
#[pyfunction]
165165
fn get_objects(args: GetObjectsArgs, vm: &VirtualMachine) -> PyResult<PyListRef> {
166166
let generation_opt = args.generation.flatten();
167-
if let Some(g) = generation_opt {
168-
if g < 0 || g > 2 {
169-
return Err(
170-
vm.new_value_error(format!("generation must be in range(0, 3), not {}", g))
171-
);
172-
}
167+
if let Some(g) = generation_opt
168+
&& !(0..=2).contains(&g)
169+
{
170+
return Err(
171+
vm.new_value_error(format!("generation must be in range(0, 3), not {}", g))
172+
);
173173
}
174174
let objects = gc_state::gc_state().get_objects(generation_opt);
175175
Ok(vm.ctx.new_list(objects))

‎crates/vm/src/gc_state.rs‎

Lines changed: 62 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,8 @@ pub struct GcState {
101101
/// Per-generation object tracking (for correct gc_refs algorithm)
102102
/// Objects start in gen0, survivors move to gen1, then gen2
103103
generation_objects: [RwLock<HashSet<GcObjectPtr>>; 3],
104+
/// Frozen/permanent objects (excluded from normal GC)
105+
permanent_objects: RwLock<HashSet<GcObjectPtr>>,
104106
/// Debug flags
105107
pub debug: AtomicU32,
106108
/// gc.garbage list (uncollectable objects with __del__)
@@ -145,6 +147,7 @@ impl GcState {
145147
RwLock::new(HashSet::new()),
146148
RwLock::new(HashSet::new()),
147149
],
150+
permanent_objects: RwLock::new(HashSet::new()),
148151
debug: AtomicU32::new(0),
149152
garbage: PyMutex::new(Vec::new()),
150153
callbacks: PyMutex::new(Vec::new()),
@@ -225,12 +228,14 @@ impl GcState {
225228
/// obj must be a valid pointer to a PyObject
226229
pub unsafe fn track_object(&self, obj: NonNull<PyObject>) {
227230
let gc_ptr = GcObjectPtr(obj);
228-
self.generations[0].count.fetch_add(1, Ordering::SeqCst);
229-
self.alloc_count.fetch_add(1, Ordering::SeqCst);
230231

231-
// Add to generation 0 tracking (for correct gc_refs algorithm)
232-
if let Ok(mut gen0) = self.generation_objects[0].write() {
233-
gen0.insert(gc_ptr);
232+
// Add to generation 0 tracking first (for correct gc_refs algorithm)
233+
// Only increment count if we successfully add to the set
234+
if let Ok(mut gen0) = self.generation_objects[0].write()
235+
&& gen0.insert(gc_ptr)
236+
{
237+
self.generations[0].count.fetch_add(1, Ordering::SeqCst);
238+
self.alloc_count.fetch_add(1, Ordering::SeqCst);
234239
}
235240

236241
// Also add to global tracking (for get_objects, etc.)
@@ -247,23 +252,20 @@ impl GcState {
247252
pub unsafe fn untrack_object(&self, obj: NonNull<PyObject>) {
248253
let gc_ptr = GcObjectPtr(obj);
249254

250-
// Remove from all generation tracking lists
251-
for generation in &self.generation_objects {
252-
if generation
253-
.write()
254-
.ok()
255-
.is_some_and(|mut gen_set| gen_set.remove(&gc_ptr))
255+
// Remove from generation tracking lists and decrement the correct generation's count
256+
for (gen_idx, generation) in self.generation_objects.iter().enumerate() {
257+
if let Ok(mut gen_set) = generation.write()
258+
&& gen_set.remove(&gc_ptr)
256259
{
260+
// Decrement count for the generation we removed from
261+
let count = self.generations[gen_idx].count.load(Ordering::SeqCst);
262+
if count > 0 {
263+
self.generations[gen_idx].count.fetch_sub(1, Ordering::SeqCst);
264+
}
257265
break; // Object can only be in one generation
258266
}
259267
}
260268

261-
// Update counts (simplified - just decrement gen0 for now)
262-
let count = self.generations[0].count.load(Ordering::SeqCst);
263-
if count > 0 {
264-
self.generations[0].count.fetch_sub(1, Ordering::SeqCst);
265-
}
266-
267269
// Remove from global tracking
268270
if let Ok(mut tracked) = self.tracked_objects.write() {
269271
tracked.remove(&gc_ptr);
@@ -712,14 +714,21 @@ impl GcState {
712714
for &ptr in survivors {
713715
// Remove from current generation
714716
for gen_idx in 0..=from_gen {
715-
if self.generation_objects[gen_idx]
716-
.write()
717-
.ok()
718-
.is_some_and(|mut gen_set| gen_set.remove(&ptr))
717+
if let Ok(mut gen_set) = self.generation_objects[gen_idx].write()
718+
&& gen_set.remove(&ptr)
719719
{
720+
// Decrement count for source generation
721+
let count = self.generations[gen_idx].count.load(Ordering::SeqCst);
722+
if count > 0 {
723+
self.generations[gen_idx].count.fetch_sub(1, Ordering::SeqCst);
724+
}
725+
720726
// Add to next generation
721-
if let Ok(mut next_set) = self.generation_objects[next_gen].write() {
722-
next_set.insert(ptr);
727+
if let Ok(mut next_set) = self.generation_objects[next_gen].write()
728+
&& next_set.insert(ptr)
729+
{
730+
// Increment count for target generation
731+
self.generations[next_gen].count.fetch_add(1, Ordering::SeqCst);
723732
}
724733
break;
725734
}
@@ -735,16 +744,42 @@ impl GcState {
735744
/// Freeze all tracked objects (move to permanent generation)
736745
pub fn freeze(&self) {
737746
// Move all objects from gen0-2 to permanent
738-
for generation in &self.generations {
739-
let count = generation.count.swap(0, Ordering::SeqCst);
747+
let mut objects_to_freeze: Vec<GcObjectPtr> = Vec::new();
748+
749+
for (gen_idx, generation) in self.generation_objects.iter().enumerate() {
750+
if let Ok(mut gen_set) = generation.write() {
751+
objects_to_freeze.extend(gen_set.drain());
752+
self.generations[gen_idx].count.store(0, Ordering::SeqCst);
753+
}
754+
}
755+
756+
// Add to permanent set
757+
if let Ok(mut permanent) = self.permanent_objects.write() {
758+
let count = objects_to_freeze.len();
759+
for ptr in objects_to_freeze {
760+
permanent.insert(ptr);
761+
}
740762
self.permanent.count.fetch_add(count, Ordering::SeqCst);
741763
}
742764
}
743765

744766
/// Unfreeze all objects (move from permanent to gen2)
745767
pub fn unfreeze(&self) {
746-
let count = self.permanent.count.swap(0, Ordering::SeqCst);
747-
self.generations[2].count.fetch_add(count, Ordering::SeqCst);
768+
let mut objects_to_unfreeze: Vec<GcObjectPtr> = Vec::new();
769+
770+
if let Ok(mut permanent) = self.permanent_objects.write() {
771+
objects_to_unfreeze.extend(permanent.drain());
772+
self.permanent.count.store(0, Ordering::SeqCst);
773+
}
774+
775+
// Add to generation 2
776+
if let Ok(mut gen2) = self.generation_objects[2].write() {
777+
let count = objects_to_unfreeze.len();
778+
for ptr in objects_to_unfreeze {
779+
gen2.insert(ptr);
780+
}
781+
self.generations[2].count.fetch_add(count, Ordering::SeqCst);
782+
}
748783
}
749784
}
750785

‎crates/vm/src/stdlib/thread.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -523,7 +523,7 @@ pub(crate) mod _thread {
523523
let mut handles = vm.state.shutdown_handles.lock();
524524
// Clean up finished entries
525525
handles.retain(|(inner_weak, _): &ShutdownEntry| {
526-
inner_weak.upgrade().map_or(false, |inner| {
526+
inner_weak.upgrade().is_some_and(|inner| {
527527
let guard = inner.lock();
528528
guard.state != ThreadHandleState::Done && guard.ident != current_ident
529529
})

0 commit comments

Comments
 (0)