use std::alloc::{Layout, alloc, dealloc};
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use crate::gc::heap::OldSpace;
use crate::gc::trace::{Tracer, trace_heap_object};
use crate::objects::heap_object::HeapObject;
pub const BLOCK_SIZE: usize = 32 * 1024;
pub const LINE_SIZE: usize = 128;
pub const LINES_PER_BLOCK: usize = BLOCK_SIZE / LINE_SIZE;
pub const EVACUATION_THRESHOLD_PCT: u8 = 50;
const DEFAULT_BLOCK_COUNT: usize = 128;
#[derive(Clone)]
pub struct LineMap {
bits: [u8; LINES_PER_BLOCK / 8],
}
impl LineMap {
pub fn new() -> Self {
Self {
bits: [0u8; LINES_PER_BLOCK / 8],
}
}
pub fn mark(&mut self, index: usize) {
assert!(index < LINES_PER_BLOCK, "line index out of range");
self.bits[index / 8] |= 1 << (index % 8);
}
pub fn is_marked(&self, index: usize) -> bool {
assert!(index < LINES_PER_BLOCK, "line index out of range");
(self.bits[index / 8] & (1 << (index % 8))) != 0
}
pub fn clear(&mut self) {
self.bits = [0u8; LINES_PER_BLOCK / 8];
}
pub fn live_line_count(&self) -> usize {
self.bits.iter().map(|b| b.count_ones() as usize).sum()
}
pub fn occupancy_percent(&self) -> u8 {
let live = self.live_line_count();
((live * 100) / LINES_PER_BLOCK) as u8
}
}
impl Default for LineMap {
fn default() -> Self {
Self::new()
}
}
pub struct ImmixBlock {
base: *mut u8,
cursor: usize,
line_map: LineMap,
evacuate: bool,
}
unsafe impl Send for ImmixBlock {}
impl ImmixBlock {
pub fn new() -> Self {
let layout = Layout::from_size_align(BLOCK_SIZE, 8).expect("valid layout");
let base = unsafe { alloc(layout) };
assert!(!base.is_null(), "Immix block allocation failed");
Self {
base,
cursor: 0,
line_map: LineMap::new(),
evacuate: false,
}
}
pub fn bump_alloc(&mut self, layout: Layout) -> Option<*mut u8> {
let align = layout.align();
let size = layout.size();
let current = self.base as usize + self.cursor;
let aligned = current.checked_add(align - 1)? & !(align - 1);
let end = aligned.checked_add(size)?;
let new_cursor = end - self.base as usize;
if new_cursor > BLOCK_SIZE {
return None;
}
self.cursor = new_cursor;
Some(aligned as *mut u8)
}
pub fn contains(&self, ptr: *const u8) -> bool {
let addr = ptr as usize;
let base = self.base as usize;
addr >= base && addr < base + BLOCK_SIZE
}
pub fn line_index_of(&self, ptr: *const u8) -> usize {
let offset = ptr as usize - self.base as usize;
assert!(offset < BLOCK_SIZE, "pointer not within this block");
offset / LINE_SIZE
}
pub fn mark_lines_for_object(&mut self, ptr: *const u8, size: usize) {
let start = self.line_index_of(ptr);
let end_byte = (ptr as usize + size).saturating_sub(1);
let end = (end_byte - self.base as usize) / LINE_SIZE;
for i in start..=end.min(LINES_PER_BLOCK - 1) {
self.line_map.mark(i);
}
}
pub fn should_evacuate(&self) -> bool {
self.line_map.occupancy_percent() < EVACUATION_THRESHOLD_PCT
}
pub fn line_map(&self) -> &LineMap {
&self.line_map
}
pub fn line_map_mut(&mut self) -> &mut LineMap {
&mut self.line_map
}
pub fn base_ptr(&self) -> *mut u8 {
self.base
}
pub fn cursor(&self) -> usize {
self.cursor
}
pub fn reset(&mut self) {
self.cursor = 0;
self.line_map.clear();
self.evacuate = false;
}
pub fn is_evacuate_candidate(&self) -> bool {
self.evacuate
}
pub fn set_evacuate(&mut self, flag: bool) {
self.evacuate = flag;
}
pub fn used(&self) -> usize {
self.cursor
}
pub fn remaining(&self) -> usize {
BLOCK_SIZE - self.cursor
}
}
impl Default for ImmixBlock {
fn default() -> Self {
Self::new()
}
}
impl Drop for ImmixBlock {
fn drop(&mut self) {
if !self.base.is_null() {
let layout = Layout::from_size_align(BLOCK_SIZE, 8).expect("valid layout");
unsafe { dealloc(self.base, layout) };
}
}
}
pub struct ImmixSpace {
full_blocks: Vec<ImmixBlock>,
free_blocks: Vec<ImmixBlock>,
max_blocks: usize,
}
impl ImmixSpace {
pub fn new(max_blocks: usize) -> Self {
Self {
full_blocks: Vec::new(),
free_blocks: Vec::new(),
max_blocks,
}
}
pub fn with_defaults() -> Self {
Self::new(DEFAULT_BLOCK_COUNT)
}
pub fn obtain_block(&mut self) -> Option<ImmixBlock> {
if let Some(block) = self.free_blocks.pop() {
return Some(block);
}
let total = self.full_blocks.len() + self.free_blocks.len() + 1;
if total <= self.max_blocks {
return Some(ImmixBlock::new());
}
None
}
pub fn return_full_block(&mut self, block: ImmixBlock) {
self.full_blocks.push(block);
}
pub fn return_free_block(&mut self, block: ImmixBlock) {
self.free_blocks.push(block);
}
pub fn used(&self) -> usize {
self.full_blocks.iter().map(|b| b.used()).sum()
}
pub fn capacity(&self) -> usize {
self.max_blocks * BLOCK_SIZE
}
pub fn full_block_count(&self) -> usize {
self.full_blocks.len()
}
pub fn free_block_count(&self) -> usize {
self.free_blocks.len()
}
pub fn contains(&self, ptr: *const u8) -> bool {
self.full_blocks.iter().any(|b| b.contains(ptr))
}
pub fn recycle_blocks(&mut self) -> usize {
let mut keep = Vec::new();
let mut recycled = 0usize;
for mut block in self.full_blocks.drain(..) {
if block.should_evacuate() {
block.reset();
self.free_blocks.push(block);
recycled += 1;
} else {
block.line_map_mut().clear();
keep.push(block);
}
}
self.full_blocks = keep;
recycled
}
pub(crate) fn drain_full_blocks(&mut self) -> Vec<ImmixBlock> {
self.full_blocks.drain(..).collect()
}
pub(crate) fn return_blocks(&mut self, full: Vec<ImmixBlock>, free: Vec<ImmixBlock>) {
self.full_blocks.extend(full);
self.free_blocks.extend(free);
}
}
impl Default for ImmixSpace {
fn default() -> Self {
Self::with_defaults()
}
}
pub struct Tlab {
current_block: Option<ImmixBlock>,
bytes_allocated: usize,
}
impl Tlab {
pub fn new() -> Self {
Self {
current_block: None,
bytes_allocated: 0,
}
}
pub fn allocate(&mut self, layout: Layout, space: &mut ImmixSpace) -> Option<*mut u8> {
if let Some(ref mut block) = self.current_block
&& let Some(ptr) = block.bump_alloc(layout)
{
self.bytes_allocated += layout.size();
return Some(ptr);
}
if let Some(block) = self.current_block.take() {
space.return_full_block(block);
}
let mut new_block = space.obtain_block()?;
let ptr = new_block.bump_alloc(layout)?;
self.bytes_allocated += layout.size();
self.current_block = Some(new_block);
Some(ptr)
}
pub fn flush(&mut self, space: &mut ImmixSpace) {
if let Some(block) = self.current_block.take() {
space.return_full_block(block);
}
}
pub fn bytes_allocated(&self) -> usize {
self.bytes_allocated
}
pub fn has_block(&self) -> bool {
self.current_block.is_some()
}
pub fn remaining(&self) -> usize {
self.current_block.as_ref().map_or(0, |b| b.remaining())
}
}
impl Default for Tlab {
fn default() -> Self {
Self::new()
}
}
pub struct ImmixCollector {
mark_set: HashSet<usize>,
}
impl ImmixCollector {
pub fn new() -> Self {
Self {
mark_set: HashSet::new(),
}
}
fn mark_object(&mut self, ptr: *mut HeapObject) {
self.mark_set.insert(ptr as usize);
}
pub fn is_marked(&self, ptr: *mut HeapObject) -> bool {
self.mark_set.contains(&(ptr as usize))
}
pub unsafe fn mark(&mut self, roots: &[*mut *mut HeapObject], blocks: &mut [ImmixBlock]) {
let mut grey_stack: Vec<*mut HeapObject> = Vec::new();
for &slot in roots {
let ptr = unsafe { *slot };
if ptr.is_null() {
continue;
}
if !self.is_marked(ptr) {
self.mark_object(ptr);
grey_stack.push(ptr);
let obj_size = unsafe { (*ptr).alloc_size() } as usize;
for block in blocks.iter_mut() {
if block.contains(ptr as *const u8) {
block.mark_lines_for_object(ptr as *const u8, obj_size);
break;
}
}
}
}
while let Some(obj) = grey_stack.pop() {
let mut tracer = Tracer::new();
unsafe { trace_heap_object(obj, &mut tracer) };
for child_raw in tracer.gray_stack {
let child = child_raw as *mut HeapObject;
if child.is_null() || self.is_marked(child) {
continue;
}
self.mark_object(child);
grey_stack.push(child);
let child_size = unsafe { (*child).alloc_size() } as usize;
for block in blocks.iter_mut() {
if block.contains(child as *const u8) {
block.mark_lines_for_object(child as *const u8, child_size);
break;
}
}
}
}
}
pub fn select_evacuation_candidates(&self, blocks: &mut [ImmixBlock]) -> usize {
let mut count = 0;
for block in blocks.iter_mut() {
if block.should_evacuate() && block.used() > 0 {
block.set_evacuate(true);
count += 1;
}
}
count
}
pub unsafe fn evacuate(
&self,
blocks: &mut Vec<ImmixBlock>,
roots: &mut [*mut *mut HeapObject],
space: &mut ImmixSpace,
) -> Vec<ImmixBlock> {
let mut evacuated: Vec<ImmixBlock> = Vec::new();
let mut target_block: Option<ImmixBlock> = None;
let mut forwarding: Vec<(usize, usize)> = Vec::new();
let mut kept = Vec::new();
for block in blocks.drain(..) {
if block.is_evacuate_candidate() {
let base = block.base_ptr();
let used = block.cursor();
let mut offset = 0usize;
while offset < used {
let old_ptr = unsafe { base.add(offset) } as *mut HeapObject;
let size = unsafe { (*old_ptr).alloc_size() } as usize;
if size == 0 {
break;
}
if self.is_marked(old_ptr) {
let layout = Layout::from_size_align(size, 8).expect("valid layout");
let dest = loop {
if let Some(ref mut tb) = target_block {
if let Some(ptr) = tb.bump_alloc(layout) {
break ptr;
}
kept.push(target_block.take().unwrap());
}
target_block = space.obtain_block();
if target_block.is_none() {
break std::ptr::null_mut();
}
};
if dest.is_null() {
break;
}
unsafe {
std::ptr::copy_nonoverlapping(old_ptr as *const u8, dest, size);
}
forwarding.push((old_ptr as usize, dest as usize));
}
offset += size;
}
evacuated.push(block);
} else {
kept.push(block);
}
}
if let Some(tb) = target_block {
if tb.used() > 0 {
kept.push(tb);
} else {
space.return_free_block(tb);
}
}
*blocks = kept;
for slot in roots.iter_mut() {
let old_addr = unsafe { *(*slot) } as usize;
if let Some(&(_, new_addr)) = forwarding.iter().find(|(o, _)| *o == old_addr) {
unsafe { **slot = new_addr as *mut HeapObject };
}
}
evacuated
}
pub unsafe fn collect(&mut self, roots: &mut [*mut *mut HeapObject], space: &mut ImmixSpace) {
let mut blocks = space.drain_full_blocks();
unsafe { self.mark(roots, &mut blocks) };
let _candidates = self.select_evacuation_candidates(&mut blocks);
let evacuated = unsafe { self.evacuate(&mut blocks, roots, space) };
let mut free = Vec::new();
for mut block in evacuated {
block.reset();
free.push(block);
}
for block in &mut blocks {
block.line_map_mut().clear();
block.set_evacuate(false);
}
space.return_blocks(blocks, free);
self.mark_set.clear();
}
}
impl Default for ImmixCollector {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct ConcurrentMarkResult {
pub live_set: HashSet<usize>,
pub marked_count: usize,
}
pub struct ConcurrentMarker {
result: Arc<Mutex<Option<ConcurrentMarkResult>>>,
}
impl ConcurrentMarker {
pub fn new() -> Self {
Self {
result: Arc::new(Mutex::new(None)),
}
}
pub unsafe fn mark_old_generation(&self, roots: &[*const HeapObject], old_space: &OldSpace) {
let mut live_set = HashSet::new();
let mut grey_stack: Vec<*const HeapObject> = Vec::new();
for &ptr in roots {
if ptr.is_null() {
continue;
}
if old_space.contains(ptr as *mut u8) && !live_set.contains(&(ptr as usize)) {
live_set.insert(ptr as usize);
grey_stack.push(ptr);
}
}
while let Some(obj) = grey_stack.pop() {
let mut tracer = Tracer::new();
unsafe { trace_heap_object(obj as *mut HeapObject, &mut tracer) };
for child_raw in tracer.gray_stack {
let child = child_raw as *const HeapObject;
if child.is_null() {
continue;
}
let addr = child as usize;
if old_space.contains(child as *mut u8) && !live_set.contains(&addr) {
live_set.insert(addr);
grey_stack.push(child);
}
}
}
let marked_count = live_set.len();
let result = ConcurrentMarkResult {
live_set,
marked_count,
};
if let Ok(mut guard) = self.result.lock() {
*guard = Some(result);
}
}
pub fn take_result(&self) -> Option<ConcurrentMarkResult> {
self.result.lock().ok().and_then(|mut guard| guard.take())
}
pub fn result_channel(&self) -> Arc<Mutex<Option<ConcurrentMarkResult>>> {
Arc::clone(&self.result)
}
}
impl Default for ConcurrentMarker {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::alloc::Layout;
#[test]
fn test_line_map_mark_and_query() {
let mut lm = LineMap::new();
assert!(!lm.is_marked(0));
lm.mark(0);
assert!(lm.is_marked(0));
assert!(!lm.is_marked(1));
}
#[test]
fn test_line_map_clear_resets_all() {
let mut lm = LineMap::new();
lm.mark(0);
lm.mark(100);
lm.mark(255);
assert_eq!(lm.live_line_count(), 3);
lm.clear();
assert_eq!(lm.live_line_count(), 0);
}
#[test]
fn test_line_map_occupancy() {
let mut lm = LineMap::new();
for i in 0..LINES_PER_BLOCK / 2 {
lm.mark(i);
}
assert_eq!(lm.occupancy_percent(), 50);
}
#[test]
fn test_line_map_full_occupancy() {
let mut lm = LineMap::new();
for i in 0..LINES_PER_BLOCK {
lm.mark(i);
}
assert_eq!(lm.live_line_count(), LINES_PER_BLOCK);
assert_eq!(lm.occupancy_percent(), 100);
}
#[test]
fn test_block_bump_alloc_and_contains() {
let mut block = ImmixBlock::new();
let layout = Layout::from_size_align(64, 8).unwrap();
let ptr = block.bump_alloc(layout).expect("block has space");
assert!(block.contains(ptr));
assert_eq!(block.used(), 64);
}
#[test]
fn test_block_exhaustion() {
let mut block = ImmixBlock::new();
let layout = Layout::from_size_align(BLOCK_SIZE, 8).unwrap();
assert!(block.bump_alloc(layout).is_some());
let small = Layout::from_size_align(1, 1).unwrap();
assert!(block.bump_alloc(small).is_none());
}
#[test]
fn test_block_reset() {
let mut block = ImmixBlock::new();
let layout = Layout::from_size_align(128, 8).unwrap();
block.bump_alloc(layout).unwrap();
block.line_map_mut().mark(0);
block.reset();
assert_eq!(block.used(), 0);
assert_eq!(block.line_map().live_line_count(), 0);
}
#[test]
fn test_block_line_marking() {
let mut block = ImmixBlock::new();
let layout = Layout::from_size_align(256, 8).unwrap();
let ptr = block.bump_alloc(layout).unwrap();
block.mark_lines_for_object(ptr, 256);
assert!(block.line_map().is_marked(0));
assert!(block.line_map().is_marked(1));
assert!(!block.line_map().is_marked(2));
}
#[test]
fn test_block_evacuation_threshold() {
let mut block = ImmixBlock::new();
assert!(block.should_evacuate());
for i in 0..LINES_PER_BLOCK / 2 {
block.line_map_mut().mark(i);
}
assert!(!block.should_evacuate());
}
#[test]
fn test_block_remaining() {
let mut block = ImmixBlock::new();
assert_eq!(block.remaining(), BLOCK_SIZE);
let layout = Layout::from_size_align(100, 8).unwrap();
block.bump_alloc(layout).unwrap();
assert_eq!(block.remaining(), BLOCK_SIZE - 100);
}
#[test]
fn test_space_obtain_and_return() {
let mut space = ImmixSpace::new(4);
let block = space.obtain_block().expect("should get a block");
assert_eq!(block.used(), 0);
space.return_full_block(block);
assert_eq!(space.full_block_count(), 1);
}
#[test]
fn test_space_exhaustion() {
let mut space = ImmixSpace::new(2);
let b1 = space.obtain_block().expect("first block");
space.return_full_block(b1);
let b2 = space.obtain_block().expect("second block");
space.return_full_block(b2);
assert!(space.obtain_block().is_none());
}
#[test]
fn test_space_recycle() {
let mut space = ImmixSpace::new(4);
let mut block = ImmixBlock::new();
let layout = Layout::from_size_align(64, 8).unwrap();
block.bump_alloc(layout).unwrap();
block.line_map_mut().mark(0);
space.return_full_block(block);
let recycled = space.recycle_blocks();
assert_eq!(recycled, 1);
assert_eq!(space.full_block_count(), 0);
assert_eq!(space.free_block_count(), 1);
}
#[test]
fn test_space_capacity() {
let space = ImmixSpace::new(4);
assert_eq!(space.capacity(), 4 * BLOCK_SIZE);
}
#[test]
fn test_tlab_allocate_within_block() {
let mut space = ImmixSpace::new(4);
let mut tlab = Tlab::new();
let layout = Layout::from_size_align(64, 8).unwrap();
let ptr = tlab.allocate(layout, &mut space);
assert!(ptr.is_some());
assert!(tlab.has_block());
assert_eq!(tlab.bytes_allocated(), 64);
}
#[test]
fn test_tlab_block_transition() {
let mut space = ImmixSpace::new(4);
let mut tlab = Tlab::new();
let layout = Layout::from_size_align(BLOCK_SIZE, 8).unwrap();
tlab.allocate(layout, &mut space).expect("first block");
let small = Layout::from_size_align(64, 8).unwrap();
let ptr = tlab.allocate(small, &mut space);
assert!(ptr.is_some());
assert_eq!(space.full_block_count(), 1);
}
#[test]
fn test_tlab_flush() {
let mut space = ImmixSpace::new(4);
let mut tlab = Tlab::new();
let layout = Layout::from_size_align(64, 8).unwrap();
tlab.allocate(layout, &mut space).unwrap();
assert!(tlab.has_block());
tlab.flush(&mut space);
assert!(!tlab.has_block());
assert_eq!(space.full_block_count(), 1);
}
#[test]
fn test_tlab_remaining() {
let mut space = ImmixSpace::new(4);
let mut tlab = Tlab::new();
assert_eq!(tlab.remaining(), 0);
let layout = Layout::from_size_align(64, 8).unwrap();
tlab.allocate(layout, &mut space).unwrap();
assert_eq!(tlab.remaining(), BLOCK_SIZE - 64);
}
fn alloc_in_block(block: &mut ImmixBlock) -> *mut HeapObject {
let layout = Layout::new::<HeapObject>();
let raw = block.bump_alloc(layout).expect("block has space");
unsafe { std::ptr::write_bytes(raw, 0, layout.size()) };
let obj = raw as *mut HeapObject;
unsafe { (*obj).init_alloc_size(layout.size() as u32) };
obj
}
#[test]
fn test_collector_mark_roots() {
let mut block = ImmixBlock::new();
let obj = alloc_in_block(&mut block);
let mut root: *mut HeapObject = obj;
let roots = [&raw mut root as *mut *mut HeapObject];
let mut collector = ImmixCollector::new();
let mut blocks = [block];
unsafe { collector.mark(&roots, &mut blocks) };
assert!(collector.is_marked(obj));
assert!(blocks[0].line_map().is_marked(0));
}
#[test]
fn test_collector_skips_null_roots() {
let mut collector = ImmixCollector::new();
let mut root: *mut HeapObject = std::ptr::null_mut();
let roots = [&raw mut root as *mut *mut HeapObject];
let mut blocks: [ImmixBlock; 0] = [];
unsafe { collector.mark(&roots, &mut blocks) };
assert!(collector.mark_set.is_empty());
}
#[test]
fn test_collector_evacuation_candidates() {
let mut block = ImmixBlock::new();
let _obj = alloc_in_block(&mut block);
block.line_map_mut().mark(0);
let collector = ImmixCollector::new();
let mut blocks = vec![block];
let count = collector.select_evacuation_candidates(&mut blocks);
assert_eq!(count, 1);
assert!(blocks[0].is_evacuate_candidate());
}
#[test]
fn test_collector_full_cycle() {
let mut space = ImmixSpace::new(8);
let mut tlab = Tlab::new();
let layout = Layout::new::<HeapObject>();
let raw = tlab.allocate(layout, &mut space).expect("alloc succeeds");
unsafe { std::ptr::write_bytes(raw, 0, layout.size()) };
let obj = raw as *mut HeapObject;
unsafe { (*obj).init_alloc_size(layout.size() as u32) };
let mut root: *mut HeapObject = obj;
let mut roots = [&raw mut root as *mut *mut HeapObject];
tlab.flush(&mut space);
let mut collector = ImmixCollector::new();
unsafe { collector.collect(&mut roots, &mut space) };
assert!(!root.is_null(), "root must survive collection");
}
#[test]
fn test_collector_unreachable_block_recycled() {
let mut space = ImmixSpace::new(8);
let mut block = ImmixBlock::new();
let _dead = alloc_in_block(&mut block);
space.return_full_block(block);
let mut roots: [*mut *mut HeapObject; 0] = [];
let mut collector = ImmixCollector::new();
unsafe { collector.collect(&mut roots, &mut space) };
assert_eq!(space.full_block_count(), 0);
assert!(space.free_block_count() > 0);
}
#[test]
fn test_concurrent_marker_marks_roots() {
let mut old = OldSpace::new(65536);
let layout = Layout::new::<HeapObject>();
let raw = old.bump_alloc(layout).expect("old space has room");
unsafe { std::ptr::write_bytes(raw, 0, layout.size()) };
let obj = raw as *mut HeapObject;
unsafe { (*obj).init_alloc_size(layout.size() as u32) };
let marker = ConcurrentMarker::new();
let roots = [obj as *const HeapObject];
unsafe { marker.mark_old_generation(&roots, &old) };
let result = marker.take_result().expect("result should be available");
assert_eq!(result.marked_count, 1);
assert!(result.live_set.contains(&(obj as usize)));
}
#[test]
fn test_concurrent_marker_no_roots() {
let old = OldSpace::new(65536);
let marker = ConcurrentMarker::new();
let roots: &[*const HeapObject] = &[];
unsafe { marker.mark_old_generation(roots, &old) };
let result = marker.take_result().expect("result should be available");
assert_eq!(result.marked_count, 0);
}
#[test]
fn test_concurrent_marker_take_result_is_consume() {
let old = OldSpace::new(65536);
let marker = ConcurrentMarker::new();
let roots: &[*const HeapObject] = &[];
unsafe { marker.mark_old_generation(roots, &old) };
assert!(marker.take_result().is_some());
assert!(
marker.take_result().is_none(),
"second take must return None"
);
}
#[test]
fn test_concurrent_marker_result_channel() {
let marker = ConcurrentMarker::new();
let channel = marker.result_channel();
let guard = channel.lock().unwrap();
assert!(guard.is_none());
}
}