Skip to main content

kevy_alloc/
heap_free.rs

1//! The free side of [`Heap`] (child module via `#[path]`, the house
2//! pattern) — claims-first recycling, the small free's routing, the
3//! local bitmap free, and the foreign-free drain. Split from `heap.rs`
4//! for the 500-LOC ceiling; the seam is real: everything here runs on
5//! release paths, nothing on allocation.
6
7use core::ptr::NonNull;
8
9use crate::class;
10use crate::segment::{self, NO_CLASS, Segment};
11
12use super::Heap;
13
14impl Heap {
15    /// The claims-first free: when the pointer lands in the class's
16    /// claimed word, recycle the bit without reading the segment header
17    /// at all. The match itself proves ownership — a claim only ever
18    /// covers this heap's own spans, and equal segment addresses mean
19    /// the same segment — so the owner check would confirm what the
20    /// compare already did. 99.86 % of collection-write frees take
21    /// this path (branch-rate probe, hset storm), and the header read
22    /// was the fast path's single foreign cache line.
23    #[inline]
24    fn try_free_claimed(&mut self, seg: NonNull<Segment>, ptr: NonNull<u8>, c: usize) -> bool {
25        let Some(cl) = &mut self.claims[c] else { return false };
26        if cl.seg != seg {
27            return false;
28        }
29        let ix = segment::span_index_of(ptr);
30        let slot = segment::slot_index_of(ptr, c);
31        if usize::from(cl.span_ix) != ix || slot / 64 != u32::from(cl.word) {
32            return false;
33        }
34        let bit = 1u64 << (slot % 64);
35        if cl.taken & bit == 0 {
36            return false;
37        }
38        cl.taken &= !bit;
39        true
40    }
41
42    /// # Safety
43    /// See [`Heap::dealloc`].
44    pub(super) unsafe fn dealloc_small(&mut self, ptr: NonNull<u8>, c: usize, size: usize) {
45        // SAFETY: a small allocation always lies inside a segment.
46        let seg = unsafe { segment::segment_of(ptr) };
47        if self.try_free_claimed(seg, ptr, c) {
48            self.live_bytes -= size as u64;
49            self.rounding_bytes -= (class::size_of(c) - size) as u64;
50            return;
51        }
52        // SAFETY: the mask lands on a live header for our own pointers.
53        let seg_ref = unsafe { seg.as_ref() };
54        debug_assert!(seg_ref.is_valid(), "pointer did not come from kevy-alloc");
55        if seg_ref.owner == self.id {
56            self.live_bytes -= size as u64;
57            self.rounding_bytes -= (class::size_of(c) - size) as u64;
58            // SAFETY: our own segment; exclusive access.
59            unsafe { self.free_local(seg, ptr, c) };
60        } else {
61            // Not ours to decrement. The bytes were counted on the
62            // allocating thread's heap, and a non-atomic counter over
63            // there is exactly what this design refuses to reach across
64            // for — the owner settles when it drains.
65            //
66            // Nor ours to touch the owner's segment per-op: M1 measured
67            // that bill at 18–39 % of cross-shard KV. The free lands in
68            // the local outbound ring — two plain stores — and crosses
69            // cores only when a whole batch ships.
70            if !self.outbound.push(ptr, size, c) {
71                self.outbound.flush();
72                let ok = self.outbound.push(ptr, size, c);
73                debug_assert!(ok, "a freshly flushed ring cannot be full");
74            }
75        }
76    }
77
78    /// Move every slot other shards freed back onto its own span's list.
79    pub fn drain_foreign(&mut self) {
80        let mut seg = self.segments;
81        while !seg.is_null() {
82            // SAFETY: live header from our own list.
83            let s = unsafe { &*seg };
84            let mut node = segment::take_foreign(s);
85            while !node.is_null() {
86                // SAFETY: foreign entries are slot addresses of this
87                // segment, linked through their first word.
88                let next = unsafe { node.cast::<*mut u8>().read() };
89                // SAFETY: non-null in this branch.
90                let p = unsafe { NonNull::new_unchecked(node) };
91                // SAFETY: still queued and untouched, so the size the
92                // freeing thread recorded is still there.
93                let requested = unsafe { segment::foreign_requested(p) };
94                let ix = segment::span_index_of(p);
95                // SAFETY: the span index came from the address itself.
96                let cls = unsafe { (*seg).spans[ix].class };
97                if cls != NO_CLASS {
98                    let c = cls as usize;
99                    self.live_bytes -= requested as u64;
100                    self.rounding_bytes -= (class::size_of(c) - requested) as u64;
101                    // SAFETY: our segment, exclusive access here.
102                    unsafe { self.free_local(NonNull::new_unchecked(seg), p, c) };
103                }
104                node = next;
105            }
106            seg = s.next;
107        }
108    }
109
110    /// Mark a slot free in its span's bitmap. Nothing is written into
111    /// the slot itself — that absence is what makes its pages
112    /// returnable. A span going full → partial is registered in the
113    /// class's partial ring so the slow path finds it in O(1).
114    ///
115    /// # Safety
116    /// `seg` must own `ptr`, and the caller must have exclusive access.
117    pub(crate) unsafe fn free_local(&mut self, seg: NonNull<Segment>, ptr: NonNull<u8>, c: usize) {
118        let ix = segment::span_index_of(ptr);
119        let slot = segment::slot_index_of(ptr, c);
120        // A free landing inside the class's claimed word recycles the
121        // bit heap-locally — no header touch at all. Collection writes
122        // are exactly this shape (several short-lived small allocations
123        // per op), which is where the far-line residual lived.
124        if let Some(cl) = &mut self.claims[c]
125            && cl.seg == seg
126            && usize::from(cl.span_ix) == ix
127            && slot / 64 == u32::from(cl.word)
128        {
129            let bit = 1u64 << (slot % 64);
130            if cl.taken & bit != 0 {
131                cl.taken &= !bit;
132                return;
133            }
134        }
135        // SAFETY: caller holds exclusive access to this segment.
136        let meta = unsafe { &mut (*seg.as_ptr()).spans[ix] };
137        let was_full = u32::from(meta.live) == meta.capacity();
138        meta.free_slot(slot);
139        if was_full {
140            self.partials[c].push(seg.as_ptr(), ix);
141        }
142    }
143}