diff --git a/baml_language/crates/bex_heap/src/chunked_vec.rs b/baml_language/crates/bex_heap/src/chunked_vec.rs index 1c378393a9c..809f9f0965e 100644 --- a/baml_language/crates/bex_heap/src/chunked_vec.rs +++ b/baml_language/crates/bex_heap/src/chunked_vec.rs @@ -275,21 +275,34 @@ impl ChunkedVec { /// every chunk makes generation checks expensive as the old heap grows. pub fn contains_ptr(&self, ptr: *const T) -> bool { let _read = self.chunks_lock.read(); - // SAFETY: the read lock keeps the outer Vec and its chunk descriptors - // stable. Each chunk has CHUNK_SIZE initialized cells; we only compare - // addresses, never read elements or dereference the supplied pointer. + // SAFETY: the read lock excludes concurrent growth for the scan. + unsafe { self.locate_unsynchronized(ptr).is_some() } + } + + /// Find `ptr` in the backing storage, including unused capacity, as + /// `(chunk_idx, offset)`. Compares addresses only; never dereferences + /// `ptr` or reads elements. + /// + /// # Safety + /// + /// Takes no lock: nothing may grow, clear, swap or drop `self` during the + /// call. The heap's Gen1 and Gen2 meet this outside collection, because + /// only the collector changes them and it runs with every mutator parked. + #[inline] + pub unsafe fn locate_unsynchronized(&self, ptr: *const T) -> Option<(usize, usize)> { + // SAFETY: the caller excludes concurrent changes to the outer Vec. + // Each chunk has CHUNK_SIZE initialized cells. unsafe { - let chunks = self.chunks.get(); - let data = (*chunks).as_ptr(); - for index in 0..(*chunks).len() { - let start = (*data.add(index)).as_ptr().cast::(); + let chunks = &*self.chunks.get(); + for (index, chunk) in chunks.iter().enumerate() { + let start = chunk.as_ptr().cast::(); let end = start.add(CHUNK_SIZE); if ptr >= start && ptr < end { - return true; + return Some((index, ptr.offset_from(start) as usize)); } } } - false + None } /// Calculate chunk index and offset within chunk for a given index. diff --git a/baml_language/crates/bex_heap/src/heap.rs b/baml_language/crates/bex_heap/src/heap.rs index d571fcda079..b8d0107c6df 100644 --- a/baml_language/crates/bex_heap/src/heap.rs +++ b/baml_language/crates/bex_heap/src/heap.rs @@ -831,13 +831,15 @@ impl BexHeap { /// /// # Safety /// - /// When `vec` comes from a generation's `UnsafeCell`, the caller must keep - /// that generation from being swapped/cleared by GC while it is borrowed - /// (a heap permit or exclusive GC access). `ChunkedVec` itself synchronizes - /// access to its storage descriptors, including concurrent growth. + /// Takes no lock (see [`ChunkedVec::locate_unsynchronized`]): nothing may + /// grow, clear or swap `vec` during the call. Gen1, Gen2 and the inactive + /// space change only inside a collection, so this holds for a mutator + /// holding a heap permit and for the collector itself; it does not hold + /// for Gen0, which grows on any VM's TLAB refill. #[inline] unsafe fn ptr_in_chunked_vec(vec: &ChunkedVec, raw_ptr: *const Object) -> bool { - vec.contains_ptr(raw_ptr) + // SAFETY: forwarded from the caller. + unsafe { vec.locate_unsynchronized(raw_ptr).is_some() } } /// Bug H, check 3 helper (heap_debug only): is `ptr` inside the @@ -889,26 +891,15 @@ impl BexHeap { /// /// # Safety /// - /// Must only be called at safepoints or with appropriate external synchronization. + /// Same as [`Self::ptr_in_chunked_vec`]: takes no lock, so `vec` must not + /// change during the call. #[inline] pub(crate) unsafe fn locate_in_chunked_vec( vec: &ChunkedVec, raw_ptr: *const Object, ) -> Option<(usize, usize)> { - // `num_chunks` and `chunk_start_ptr` now serialize on the - // ChunkedVec's internal RwLock; concurrent growers are excluded. - let num_chunks = vec.num_chunks(); - for chunk_idx in 0..num_chunks { - // SAFETY: `chunk_idx < num_chunks` by loop bound. - let chunk_start = unsafe { vec.chunk_start_ptr(chunk_idx) }; - let chunk_end = unsafe { chunk_start.add(ChunkedVec::::CHUNK_SIZE) }; - if raw_ptr >= chunk_start && raw_ptr < chunk_end { - // SAFETY: Both pointers are within the same allocated chunk. - let offset = unsafe { raw_ptr.offset_from(chunk_start) as usize }; - return Some((chunk_idx, offset)); - } - } - None + // SAFETY: forwarded from the caller. + unsafe { vec.locate_unsynchronized(raw_ptr) } } /// Monotonic change token for GC roots disappearing, including after