Compacting GC: merge threading and marking

rts/motoko-rts-tests/src/gc.rs

@@ -21,23 +21,49 @@ use std::fmt::Write;
 pub fn test() {
     println!("Testing garbage collection ...");
 
-    // TODO: Add more tests
-
     for test_heap in test_heaps() {
         test_gcs(&test_heap);
     }
 }
 
 fn test_heaps() -> Vec<TestHeap> {
-    vec![TestHeap {
-        heap: hashmap! {
-            0 => vec![0, 2],
-            2 => vec![0],
-            3 => vec![3],
+    vec![
+        // Just a random test that covers a bunch of cases:
+        // - Self references
+        // - Forward pointers
+        // - Backwards pointers
+        // - More than one fields in an object
+        TestHeap {
+            heap: hashmap! {
+                0 => vec![0, 2],
+                2 => vec![0],
+                3 => vec![3],
+            },
+            roots: vec![0, 2, 3],
+            closure_table: vec![0],
+        },
+        // Tests pointing to the same object in multiple fields of an object. Also has unreachable
+        // objects.
+        TestHeap {
+            heap: hashmap! {
+                0 => vec![],
+                1 => vec![],
+                2 => vec![],
+            },
+            roots: vec![1],
+            closure_table: vec![0, 0],
         },
-        roots: vec![0, 2, 3],
-        closure_table: vec![0],
-    }]
+        // Root points backwards in heap. Caught a bug in mark-compact collector.
+        TestHeap {
+            heap: hashmap! {
+                0 => vec![],
+                1 => vec![2],
+                2 => vec![1],
+            },
+            roots: vec![2],
+            closure_table: vec![],
+        },
+    ]
 }
 
 #[derive(Debug)]
@@ -67,17 +93,6 @@ fn test_gc(
 ) {
     let heap = MotokoHeap::new(refs, roots, closure_table, gc);
 
-    // Check `check_dynamic_heap` sanity
-    check_dynamic_heap(
-        refs,
-        roots,
-        closure_table,
-        &**heap.heap(),
-        heap.heap_base_offset(),
-        heap.heap_ptr_offset(),
-        heap.closure_table_ptr_offset(),
-    );
-
     for _ in 0..3 {
         gc.run(heap.clone());
 

rts/motoko-rts-tests/src/mark_stack.rs

@@ -32,12 +32,12 @@ fn test_<M: Memory>(mem: &mut M, n_objs: u32) -> TestCaseResult {
         alloc_mark_stack(mem);
 
         for obj in &objs {
-            push_mark_stack(mem, *obj as usize);
+            push_mark_stack(mem, *obj as usize, obj.wrapping_sub(1));
         }
 
-        for obj in objs.iter().rev() {
+        for obj in objs.iter().copied().rev() {
             let popped = pop_mark_stack();
-            if popped != Some(*obj as usize) {
+            if popped != Some((obj as usize, obj.wrapping_sub(1))) {
                 free_mark_stack();
                 return Err(TestCaseError::Fail(
                     format!(

rts/motoko-rts/src/debug.rs

@@ -33,11 +33,22 @@ pub unsafe fn dump_heap(
 }
 
 pub(crate) unsafe fn print_closure_table(closure_tbl_loc: *mut SkewedPtr) {
+    println!(
+        200,
+        "Closure table pointer location: {:#x}", closure_tbl_loc as usize
+    );
+
     if (*closure_tbl_loc).0 == 0 {
         println!(100, "Closure table not initialized");
         return;
     }
 
+    println!(
+        200,
+        "Closure table pointer: {:#x}",
+        (*closure_tbl_loc).unskew() as usize
+    );
+
     let arr = (*closure_tbl_loc).unskew() as *mut Array;
     let len = (*arr).len;
 

rts/motoko-rts/src/gc/copying.rs

@@ -146,7 +146,7 @@ unsafe fn evac<M: Memory>(
 unsafe fn scav<M: Memory>(mem: &mut M, begin_from_space: usize, begin_to_space: usize, obj: usize) {
     let obj = obj as *mut Obj;
 
-    crate::visitor::visit_pointer_fields(obj, begin_from_space, |field_addr| {
+    crate::visitor::visit_pointer_fields(obj, obj.tag(), begin_from_space, |field_addr| {
         evac(mem, begin_from_space, begin_to_space, field_addr as usize);
     });
 }

rts/motoko-rts/src/gc/mark_compact.rs

@@ -10,7 +10,7 @@ use crate::constants::WORD_SIZE;
 use crate::mem_utils::memcpy_words;
 use crate::memory::Memory;
 use crate::types::*;
-use crate::visitor::visit_pointer_fields;
+use crate::visitor::{pointer_to_dynamic_heap, visit_pointer_fields};
 
 use motoko_rts_macros::ic_mem_fn;
 
@@ -76,27 +76,21 @@ unsafe fn mark_compact<M: Memory, SetHp: Fn(u32)>(
     static_roots: SkewedPtr,
     closure_table_loc: *mut SkewedPtr,
 ) {
-    let mem_size = Bytes(heap_end - heap_base);
+    let heap_size = Bytes(heap_end - heap_base);
 
-    alloc_bitmap(mem, mem_size);
+    alloc_bitmap(mem, heap_size);
     alloc_mark_stack(mem);
 
     mark_static_roots(mem, static_roots, heap_base);
 
     if (*closure_table_loc).unskew() >= heap_base as usize {
         push_mark_stack(mem, *closure_table_loc, heap_base);
+        thread(closure_table_loc);
     }
 
     mark_stack(mem, heap_base);
 
-    thread_roots(static_roots, heap_base);
-
-    if (*closure_table_loc).unskew() >= heap_base as usize {
-        thread(closure_table_loc);
-    }
-
-    update_fwd_refs(heap_base);
-    update_bwd_refs(set_hp, heap_base);
+    update_refs(set_hp, heap_base);
 
     free_mark_stack();
     free_bitmap();
@@ -108,11 +102,15 @@ unsafe fn mark_static_roots<M: Memory>(mem: &mut M, static_roots: SkewedPtr, hea
     // Static objects are not in the dynamic heap so don't need marking.
     for i in 0..root_array.len() {
         let obj = root_array.get(i).unskew() as *mut Obj;
-        mark_fields(mem, obj, heap_base);
+        // Root array should only has pointers to other static MutBoxes
+        debug_assert_eq!(obj.tag(), TAG_MUTBOX); // check tag
+        debug_assert!((obj as u32) < heap_base); // check that MutBox is static
+        mark_root_mutbox_fields(mem, obj as *mut MutBox, heap_base);
     }
 }
 
 unsafe fn push_mark_stack<M: Memory>(mem: &mut M, obj: SkewedPtr, heap_base: u32) {
+    let obj_tag = obj.tag();
     let obj = obj.unskew() as u32;
 
     let obj_idx = (obj - heap_base) / WORD_SIZE;
@@ -123,70 +121,60 @@ unsafe fn push_mark_stack<M: Memory>(mem: &mut M, obj: SkewedPtr, heap_base: u32
     }
 
     set_bit(obj_idx);
-    mark_stack::push_mark_stack(mem, obj as usize);
+    mark_stack::push_mark_stack(mem, obj as usize, obj_tag);
 }
 
 unsafe fn mark_stack<M: Memory>(mem: &mut M, heap_base: u32) {
-    while let Some(obj) = pop_mark_stack() {
-        mark_fields(mem, obj as *mut Obj, heap_base);
+    while let Some((obj, tag)) = pop_mark_stack() {
+        mark_fields(mem, obj as *mut Obj, tag, heap_base);
     }
 }
 
-unsafe fn mark_fields<M: Memory>(mem: &mut M, obj: *mut Obj, heap_base: u32) {
-    visit_pointer_fields(obj, heap_base as usize, |field_addr| {
+unsafe fn mark_fields<M: Memory>(mem: &mut M, obj: *mut Obj, obj_tag: Tag, heap_base: u32) {
+    visit_pointer_fields(obj, obj_tag, heap_base as usize, |field_addr| {
         push_mark_stack(mem, *field_addr, heap_base);
+        thread(field_addr);
     });
 }
 
-unsafe fn thread_roots(static_roots: SkewedPtr, heap_base: u32) {
-    // Static roots
-    let root_array = static_roots.as_array();
-    for i in 0..root_array.len() {
-        thread_obj_fields(root_array.get(i).unskew() as *mut Obj, heap_base);
+/// Specialized version of `mark_fields` for root `MutBox`es.
+unsafe fn mark_root_mutbox_fields<M: Memory>(mem: &mut M, mutbox: *mut MutBox, heap_base: u32) {
+    let field_addr = &mut (*mutbox).field;
+    // TODO: Not sure if this check is necessary?
+    if pointer_to_dynamic_heap(field_addr, heap_base as usize) {
+        // TODO: We should be able to omit the "already marked" check here as no two root MutBox
+        // can point to the same object (I think)
+        push_mark_stack(mem, *field_addr, heap_base);
+        thread(field_addr);
     }
-    // No need to thread closure table here as it's on heap and we already marked it
 }
 
-/// Scan the heap, update forward references. At the end of this pass all fields will be threaded
-/// and forward references will be updated, pointing to the object's new location.
-unsafe fn update_fwd_refs(heap_base: u32) {
+/// Expects all fields to be threaded. First pass updates fields, second pass moves objects.
+unsafe fn update_refs<SetHp: Fn(u32)>(set_hp: SetHp, heap_base: u32) {
     let mut free = heap_base;
 
     let mut bitmap_iter = iter_bits();
     let mut bit = bitmap_iter.next();
     while bit != BITMAP_ITER_END {
         let p = (heap_base + (bit * WORD_SIZE)) as *mut Obj;
 
-        // Update forward references to the object to the object's new location and restore
-        // object header
+        // Update references to the object's new location and restore object header
         unthread(p, free);
 
-        // Get the size before threading the fields, to handle self references.
-        let size = object_size(p as usize).to_bytes().0;
-
-        // Thread fields
-        thread_obj_fields(p, heap_base);
+        // All references to the object now point to the new location (object only moved in next pass)
+        let p_size_words = object_size(p as usize);
 
-        free += size;
+        free += p_size_words.to_bytes().0;
 
         bit = bitmap_iter.next();
     }
-}
 
-/// Expects all fields to be threaded. Updates backward references and moves objects to their new
-/// locations.
-unsafe fn update_bwd_refs<SetHp: Fn(u32)>(set_hp: SetHp, heap_base: u32) {
     let mut free = heap_base;
-
     let mut bitmap_iter = iter_bits();
     let mut bit = bitmap_iter.next();
     while bit != BITMAP_ITER_END {
         let p = (heap_base + (bit * WORD_SIZE)) as *mut Obj;
 
-        // Update backward references to the object's new location and restore object header
-        unthread(p, free);
-
-        // All references to the object now point to the new location, move the object
         let p_size_words = object_size(p as usize);
         if free as usize != p as usize {
             memcpy_words(free as usize, p as usize, p_size_words);
@@ -200,10 +188,6 @@ unsafe fn update_bwd_refs<SetHp: Fn(u32)>(set_hp: SetHp, heap_base: u32) {
     set_hp(free);
 }
 
-unsafe fn thread_obj_fields(obj: *mut Obj, heap_base: u32) {
-    visit_pointer_fields(obj, heap_base as usize, |field_addr| thread(field_addr));
-}
-
 unsafe fn thread(field: *mut SkewedPtr) {
     // Store pointed object's header in the field, field address in the pointed object's header
     let pointed = (*field).unskew() as *mut Obj;

rts/motoko-rts/src/gc/mark_compact/mark_stack.rs

@@ -2,7 +2,7 @@
 //! otherwise things will break as we push. This invariant is checked in debug builds.
 
 use crate::memory::{alloc_blob, Memory};
-use crate::types::{Blob, Words};
+use crate::types::{Blob, Tag, Words};
 
 use core::ptr::null_mut;
 
@@ -50,20 +50,25 @@ unsafe fn grow_stack<M: Memory>(mem: &mut M) {
     (*STACK_BLOB_PTR).len = new_cap.to_bytes();
 }
 
-pub unsafe fn push_mark_stack<M: Memory>(mem: &mut M, obj: usize) {
+pub unsafe fn push_mark_stack<M: Memory>(mem: &mut M, obj: usize, obj_tag: Tag) {
+    // We add 2 words in a push, and `STACK_PTR` and `STACK_TOP` are both multiples of 2, so we can
+    // do simple equality check here
     if STACK_PTR == STACK_TOP {
         grow_stack(mem);
     }
 
     *STACK_PTR = obj;
-    STACK_PTR = STACK_PTR.add(1);
+    *(STACK_PTR.add(1)) = obj_tag as usize;
+    STACK_PTR = STACK_PTR.add(2);
 }
 
-pub unsafe fn pop_mark_stack() -> Option<usize> {
+pub unsafe fn pop_mark_stack() -> Option<(usize, Tag)> {
     if STACK_PTR == STACK_BASE {
         None
     } else {
-        STACK_PTR = STACK_PTR.sub(1);
-        Some(*STACK_PTR)
+        STACK_PTR = STACK_PTR.sub(2);
+        let p = *STACK_PTR;
+        let tag = *STACK_PTR.add(1);
+        Some((p, tag as u32))
     }
 }

rts/motoko-rts/src/visitor.rs

@@ -3,11 +3,15 @@ use crate::types::*;
 
 /// A visitor that passes field addresses of fields with pointers to dynamic heap to the given
 /// callback
-pub unsafe fn visit_pointer_fields<F>(obj: *mut Obj, heap_base: usize, mut visit_ptr_field: F)
-where
+pub unsafe fn visit_pointer_fields<F>(
+    obj: *mut Obj,
+    tag: Tag,
+    heap_base: usize,
+    mut visit_ptr_field: F,
+) where
     F: FnMut(*mut SkewedPtr),
 {
-    match obj.tag() {
+    match tag {
         TAG_OBJECT => {
             let obj = obj as *mut Object;
             let obj_payload = obj.payload_addr();
@@ -99,6 +103,6 @@ where
     }
 }
 
-unsafe fn pointer_to_dynamic_heap(field_addr: *mut SkewedPtr, heap_base: usize) -> bool {
+pub unsafe fn pointer_to_dynamic_heap(field_addr: *mut SkewedPtr, heap_base: usize) -> bool {
     (!(*field_addr).is_tagged_scalar()) && ((*field_addr).unskew() >= heap_base)
 }