Skip to main content

shape_value/
heap_header.rs

1//! Unified heap object header (v2 runtime spec).
2//!
3//! `HeapHeader` is a `#[repr(C)]` 8-byte struct that prefixes every heap-allocated
4//! object, giving the JIT a stable memory layout to read kind/flags and perform
5//! atomic reference counting without depending on Rust's enum discriminant layout.
6//!
7//! ## Memory layout (8 bytes)
8//!
9//! ```text
10//! Offset  Size  Field
11//! ------  ----  -----
12//!   0       4   refcount (AtomicU32)
13//!   4       2   kind (HeapKind as u16)
14//!   6       1   flags (bitfield: MARKED, PINNED, READONLY, etc.)
15//!   7       1   _pad (reserved, always 0)
16//! ```
17//!
18//! Data starts at offset 8 (`DATA_OFFSET`).
19//!
20//! Clone = `atomic_fetch_add([ptr+0], 1, Relaxed)`.
21//! Drop  = `atomic_fetch_sub([ptr+0], 1, Release)`.
22
23use std::sync::atomic::{AtomicU32, Ordering};
24
25use crate::heap_value::HeapKind;
26
27/// Flag: object has been marked by the GC during a collection cycle.
28pub const FLAG_MARKED: u8 = 0b0000_0001;
29/// Flag: object is pinned and must not be relocated by the GC.
30pub const FLAG_PINNED: u8 = 0b0000_0010;
31/// Flag: object is read-only (immutable after construction).
32pub const FLAG_READONLY: u8 = 0b0000_0100;
33
34/// Byte offset from the start of a heap allocation where payload data begins
35/// (immediately after the 8-byte HeapHeader).
36pub const DATA_OFFSET: usize = 8;
37
38/// Fixed-layout header for heap-allocated objects (v2 runtime spec).
39///
40/// This struct is designed to be readable by JIT-generated code at known offsets.
41/// The refcount lives at offset 0 for single-cycle atomic access.  Kind and flags
42/// follow at offsets 4 and 6 respectively.
43#[repr(C)]
44pub struct HeapHeader {
45    /// Reference count. Starts at 1 on allocation.
46    /// Clone: `fetch_add(1, Relaxed)`.  Drop: `fetch_sub(1, Release)`.
47    pub refcount: AtomicU32,
48    /// Object type discriminator (matches `HeapKind` and `HEAP_KIND_*` constants).
49    pub kind: u16,
50    /// Bitfield flags (FLAG_MARKED, FLAG_PINNED, FLAG_READONLY).
51    pub flags: u8,
52    /// Padding byte to reach 8-byte total size. Must be zero.
53    pub _pad: u8,
54}
55
56/// Compile-time size and offset assertions.
57const _: () = {
58    assert!(std::mem::size_of::<HeapHeader>() == 8);
59    assert!(std::mem::align_of::<HeapHeader>() == 4);
60    assert!(DATA_OFFSET == 8);
61};
62
63impl HeapHeader {
64    /// Byte offset of the `refcount` field (AtomicU32, 4 bytes).
65    pub const OFFSET_REFCOUNT: usize = 0;
66    /// Byte offset of the `kind` field (u16, 2 bytes).
67    pub const OFFSET_KIND: usize = 4;
68    /// Byte offset of the `flags` field (u8, 1 byte).
69    pub const OFFSET_FLAGS: usize = 6;
70
71    /// Byte offset where payload data starts, immediately after the header.
72    pub const DATA_OFFSET: usize = DATA_OFFSET;
73
74    /// Create a new HeapHeader with the given kind. Refcount starts at 1,
75    /// flags and padding are zeroed.
76    #[inline]
77    pub fn new(kind: u16) -> Self {
78        Self {
79            refcount: AtomicU32::new(1),
80            kind,
81            flags: 0,
82            _pad: 0,
83        }
84    }
85
86    /// Increment the reference count (clone semantics).
87    ///
88    /// Uses `Relaxed` ordering — the caller is responsible for establishing
89    /// a happens-before relationship when sharing the pointer across threads.
90    #[inline]
91    pub fn retain(&self) {
92        self.refcount.fetch_add(1, Ordering::Relaxed);
93    }
94
95    /// Decrement the reference count (drop semantics).
96    ///
97    /// Returns `true` if the refcount reached zero, meaning the caller should
98    /// deallocate the object.  Uses `Release` ordering on the decrement and
99    /// an `Acquire` fence when the count reaches zero, matching the
100    /// Arc drop protocol.
101    #[inline]
102    pub fn release(&self) -> bool {
103        let prev = self.refcount.fetch_sub(1, Ordering::Release);
104        if prev == 1 {
105            // Ensure all prior writes to the object are visible before we
106            // read/deallocate it.
107            std::sync::atomic::fence(Ordering::Acquire);
108            true
109        } else {
110            false
111        }
112    }
113
114    /// Get the current reference count (for debugging / testing only).
115    #[inline]
116    pub fn refcount(&self) -> u32 {
117        self.refcount.load(Ordering::Relaxed)
118    }
119
120    /// Get the HeapKind from this header.
121    #[inline]
122    pub fn heap_kind(&self) -> Option<HeapKind> {
123        HeapKind::from_u16(self.kind)
124    }
125
126    /// Check if a flag is set.
127    #[inline]
128    pub fn has_flag(&self, flag: u8) -> bool {
129        self.flags & flag != 0
130    }
131
132    /// Set a flag.
133    #[inline]
134    pub fn set_flag(&mut self, flag: u8) {
135        self.flags |= flag;
136    }
137
138    /// Clear a flag.
139    #[inline]
140    pub fn clear_flag(&mut self, flag: u8) {
141        self.flags &= !flag;
142    }
143}
144
145/// HeapHeader contains an AtomicU32 which is not Clone/Copy. We provide a
146/// manual Debug impl since we cannot derive it on the atomic field cleanly.
147impl std::fmt::Debug for HeapHeader {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        f.debug_struct("HeapHeader")
150            .field("refcount", &self.refcount.load(Ordering::Relaxed))
151            .field("kind", &self.kind)
152            .field("flags", &self.flags)
153            .finish()
154    }
155}
156
157impl HeapKind {
158    /// The last (highest-numbered) variant in HeapKind.
159    /// IMPORTANT: Update this when adding new HeapKind variants.
160    pub const MAX_VARIANT: Self = HeapKind::HashMap;
161
162    /// Convert a u16 discriminant to a HeapKind, returning None if out of range.
163    #[inline]
164    pub fn from_u16(v: u16) -> Option<Self> {
165        if v <= Self::MAX_VARIANT as u16 {
166            // Safety: HeapKind is repr(u8) with contiguous variants from 0..=MAX_VARIANT.
167            // We checked the range, and u16 fits in u8 for valid values.
168            Some(unsafe { std::mem::transmute(v as u8) })
169        } else {
170            None
171        }
172    }
173
174    /// Convert a u8 discriminant to a HeapKind, returning None if out of range.
175    #[inline]
176    pub fn from_u8(v: u8) -> Option<Self> {
177        Self::from_u16(v as u16)
178    }
179}
180
181/// Static assertion: HeapKind must be repr(u8), i.e. 1 byte.
182const _: () = {
183    assert!(
184        std::mem::size_of::<HeapKind>() == 1,
185        "HeapKind must be repr(u8) — transmute in from_u16 depends on this"
186    );
187};
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn test_header_size_and_alignment() {
195        assert_eq!(std::mem::size_of::<HeapHeader>(), 8);
196        // Natural alignment of AtomicU32 is 4.
197        assert_eq!(std::mem::align_of::<HeapHeader>(), 4);
198    }
199
200    #[test]
201    fn test_header_field_offsets_via_pointer_arithmetic() {
202        let h = HeapHeader::new(HeapKind::String as u16);
203        let base = &h as *const _ as usize;
204
205        let refcount_offset = &h.refcount as *const _ as usize - base;
206        let kind_offset = &h.kind as *const _ as usize - base;
207        let flags_offset = &h.flags as *const _ as usize - base;
208        let pad_offset = &h._pad as *const _ as usize - base;
209
210        assert_eq!(refcount_offset, 0, "refcount must be at offset 0");
211        assert_eq!(kind_offset, 4, "kind must be at offset 4");
212        assert_eq!(flags_offset, 6, "flags must be at offset 6");
213        assert_eq!(pad_offset, 7, "_pad must be at offset 7");
214    }
215
216    #[test]
217    fn test_header_offset_constants() {
218        assert_eq!(HeapHeader::OFFSET_REFCOUNT, 0);
219        assert_eq!(HeapHeader::OFFSET_KIND, 4);
220        assert_eq!(HeapHeader::OFFSET_FLAGS, 6);
221        assert_eq!(HeapHeader::DATA_OFFSET, 8);
222        assert_eq!(DATA_OFFSET, 8);
223    }
224
225    #[test]
226    fn test_new_header() {
227        let h = HeapHeader::new(HeapKind::TypedObject as u16);
228        assert_eq!(h.refcount(), 1);
229        assert_eq!(h.kind, HeapKind::TypedObject as u16);
230        assert_eq!(h.flags, 0);
231        assert_eq!(h._pad, 0);
232    }
233
234    #[test]
235    fn test_retain_increments_refcount() {
236        let h = HeapHeader::new(HeapKind::String as u16);
237        assert_eq!(h.refcount(), 1);
238        h.retain();
239        assert_eq!(h.refcount(), 2);
240        h.retain();
241        assert_eq!(h.refcount(), 3);
242    }
243
244    #[test]
245    fn test_release_decrements_refcount() {
246        let h = HeapHeader::new(HeapKind::String as u16);
247        h.retain(); // refcount = 2
248        h.retain(); // refcount = 3
249
250        assert!(!h.release()); // 3 -> 2, not zero
251        assert_eq!(h.refcount(), 2);
252
253        assert!(!h.release()); // 2 -> 1, not zero
254        assert_eq!(h.refcount(), 1);
255
256        assert!(h.release()); // 1 -> 0, reached zero!
257        assert_eq!(h.refcount(), 0);
258    }
259
260    #[test]
261    fn test_release_returns_true_on_last_drop() {
262        let h = HeapHeader::new(HeapKind::TypedObject as u16);
263        // refcount starts at 1; single release should return true
264        assert!(h.release());
265    }
266
267    #[test]
268    fn test_data_offset_after_header() {
269        // Verify that DATA_OFFSET equals the size of the header.
270        assert_eq!(
271            DATA_OFFSET,
272            std::mem::size_of::<HeapHeader>(),
273            "DATA_OFFSET must equal sizeof(HeapHeader)"
274        );
275    }
276
277    #[test]
278    fn test_heap_kind_roundtrip() {
279        assert_eq!(HeapKind::from_u16(0), Some(HeapKind::String));
280        assert_eq!(HeapKind::from_u16(1), Some(HeapKind::TypedObject));
281        assert_eq!(HeapKind::from_u16(2), Some(HeapKind::Closure));
282        assert_eq!(
283            HeapKind::from_u16(HeapKind::DataTable as u16),
284            Some(HeapKind::DataTable)
285        );
286        assert_eq!(
287            HeapKind::from_u16(HeapKind::TypedArray as u16),
288            Some(HeapKind::TypedArray)
289        );
290        assert_eq!(
291            HeapKind::from_u16(HeapKind::Char as u16),
292            Some(HeapKind::Char)
293        );
294        // One past the last variant must return None
295        assert_eq!(
296            HeapKind::from_u16(HeapKind::MAX_VARIANT as u16 + 1),
297            None
298        );
299        assert_eq!(HeapKind::from_u16(255), None);
300    }
301
302    #[test]
303    fn test_heap_kind_from_u8() {
304        assert_eq!(HeapKind::from_u8(0), Some(HeapKind::String));
305        assert_eq!(
306            HeapKind::from_u8(HeapKind::TypedArray as u8),
307            Some(HeapKind::TypedArray)
308        );
309        assert_eq!(
310            HeapKind::from_u8(HeapKind::Char as u8),
311            Some(HeapKind::Char)
312        );
313        assert_eq!(HeapKind::from_u8(200), None);
314    }
315
316    /// Validates that every HeapKind discriminant from 0..=MAX_VARIANT round-trips
317    /// through the unsafe transmute in `from_u16`. This catches holes in the enum
318    /// (e.g. if someone inserts a variant mid-enum or reorders them).
319    #[test]
320    fn test_heap_kind_all_variants_roundtrip_through_transmute() {
321        let max = HeapKind::MAX_VARIANT as u16;
322        for i in 0..=max {
323            let kind = HeapKind::from_u16(i)
324                .unwrap_or_else(|| panic!("HeapKind::from_u16({i}) returned None — gap in contiguous repr(u8) enum"));
325            assert_eq!(
326                kind as u16, i,
327                "HeapKind variant at discriminant {i} round-tripped to {}",
328                kind as u16
329            );
330        }
331    }
332
333    #[test]
334    fn test_flags() {
335        let mut h = HeapHeader::new(HeapKind::TypedObject as u16);
336        assert!(!h.has_flag(FLAG_MARKED));
337        assert!(!h.has_flag(FLAG_PINNED));
338
339        h.set_flag(FLAG_MARKED);
340        assert!(h.has_flag(FLAG_MARKED));
341        assert!(!h.has_flag(FLAG_PINNED));
342
343        h.set_flag(FLAG_PINNED);
344        assert!(h.has_flag(FLAG_MARKED));
345        assert!(h.has_flag(FLAG_PINNED));
346
347        h.clear_flag(FLAG_MARKED);
348        assert!(!h.has_flag(FLAG_MARKED));
349        assert!(h.has_flag(FLAG_PINNED));
350    }
351
352    #[test]
353    fn test_heap_kind_accessor() {
354        let h = HeapHeader::new(HeapKind::Closure as u16);
355        assert_eq!(h.heap_kind(), Some(HeapKind::Closure));
356
357        let h2 = HeapHeader::new(0xFFFF);
358        assert_eq!(h2.heap_kind(), None);
359    }
360
361    #[test]
362    fn test_debug_impl() {
363        let h = HeapHeader::new(HeapKind::String as u16);
364        let dbg = format!("{:?}", h);
365        assert!(dbg.contains("HeapHeader"));
366        assert!(dbg.contains("refcount"));
367        assert!(dbg.contains("kind"));
368    }
369}