Incremental mark&sweep in Rus

IncrementalMark.rs

#[derive(Clone, PartialEq, Eq, Copy, Debug)]
#[repr(u8)]
pub enum IncrementalState {
    Done,
    Mark,
    Sweep,
    Roots,
}

pub struct IncrementalMarkAndSweep<'a> {
    rootset: &'a [RootHandle],
    heap_region: Region,
    freelist: FreeList,
    scan_list: &'a mut LinkedList<GcHandle<dyn Trace>>,
    heap: &'a [GcHandle<dyn Trace>],
    new_rootset: Vec<RootHandle>,
    new_heap: Vec<GcHandle<dyn Trace>>,
    state: IncrementalState,
}

impl<'a> IncrementalMarkAndSweep<'a> {
    pub fn collect(&mut self, limit: usize) {
        let mut result = 0;
        while result < limit {
            result += self.step(limit);
            if self.state == IncrementalState::Done {
                break;
            }
        }
    }

    fn step(&mut self, limit: usize) -> usize {
        match &self.state {
            IncrementalState::Roots => {
                for root in self.rootset.iter() {
                    let root = Ptr(root.0);
                    if root.get().is_rooted() {
                        self.scan_list.push_front(GcHandle(root.get().inner()));
                        self.new_rootset.push(RootHandle(root.0));
                    }
                }
                return 0;
            }
            IncrementalState::Mark => {
                if !self.scan_list.is_empty() {
                    return self.mark(limit);
                } else {
                    self.state = IncrementalState::Sweep;
                    return 0;
                }
            }
            IncrementalState::Sweep => {
                self.sweep(limit);
                self.state = IncrementalState::Done;
                return 0;
            }
            _ => unreachable!(),
        }
    }

    fn sweep(&mut self, limit: usize) -> usize {
        let mut new_heap = vec![];
        let mut count = 0;
        let mut garbage_start = Address::null();

        for i in 0..self.heap.len() {
            if !(count < limit) {
                break;
            }
            unsafe {
                let object = self.heap.get_unchecked(i);
                let object = Ptr(object.0);
                if object.get().color() != GcColor::White {
                    self.add_freelist(garbage_start, Address::from_ptr(object.0 as *const u8));
                    garbage_start = Address::null();
                    object.get().set_color(GcColor::White);
                    new_heap.push(GcHandle(object.0));
                } else if garbage_start.is_non_null() {
                    object.get().value.finalize();
                } else {
                    object.get().value.finalize();
                    garbage_start = Address::from_ptr(object.0 as *const u8);
                }
                count += 1;
            }
        }
        self.new_heap = new_heap;
        self.add_freelist(garbage_start, self.heap_region.end);

        count
    }

    fn add_freelist(&mut self, start: Address, end: Address) {
        if start.is_null() {
            return;
        }

        let size = end.offset_from(start);
        self.freelist.add(start, size);
    }

    fn mark(&mut self, limit: usize) -> usize {
        if self.scan_list.is_empty() {
            assert_eq!(self.state, IncrementalState::Mark);
            self.state = IncrementalState::Sweep;
            return 0;
        }
        let mut count = 0;
        while let Some(object) = self.scan_list.pop_front() {
            if !(count < limit) {
                break;
            }
            let ptr = Ptr(object.0);
            ptr.get().set_color(GcColor::Black);
            count += 1;
            for field in ptr.get().value.fields() {
                let field = Ptr(field.inner());
                if field.get().color() == GcColor::White {
                    field.get().set_color(GcColor::Grey);
                    self.scan_list.push_front(GcHandle(field.0));
                }
            }
        }

        count
    }
}