Compacting GC: thread backwards pointers in marking

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

@@ -29,15 +29,43 @@ pub fn test() {
 }
 
 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],
         },
-        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],
+        },
+        // 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)]

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

@@ -32,12 +32,13 @@ 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);
+            // Pushing a dummy argument derived from `obj` for tag
+            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/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

@@ -1,4 +1,6 @@
-//! Implements "threaded compaction" as described in The Garbage Collection Handbook section 3.3.
+//! Implements threaded compaction as described in "High-Performance Garbage Collection for
+//! Memory-Constrained Environments" section 5.1.2, which is an improved version of the original
+//! threaded compaction algorithm described in The Garbage Collection Handbook section 3.3.
 
 pub mod bitmap;
 pub mod mark_stack;
@@ -10,7 +12,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;
 
@@ -46,7 +48,7 @@ pub unsafe fn compacting_gc_internal<
     get_hp: GetHp,
     set_hp: SetHp,
     static_roots: SkewedPtr,
-    closure_table_loc: *mut SkewedPtr,
+    closure_table_ptr_loc: *mut SkewedPtr,
     note_live_size: NoteLiveSize,
     note_reclaimed: NoteReclaimed,
 ) {
@@ -58,7 +60,7 @@ pub unsafe fn compacting_gc_internal<
         heap_base,
         old_hp,
         static_roots,
-        closure_table_loc,
+        closure_table_ptr_loc,
     );
 
     let reclaimed = old_hp - (get_hp() as u32);
@@ -74,7 +76,7 @@ unsafe fn mark_compact<M: Memory, SetHp: Fn(u32)>(
     heap_base: u32,
     heap_end: u32,
     static_roots: SkewedPtr,
-    closure_table_loc: *mut SkewedPtr,
+    closure_table_ptr_loc: *mut SkewedPtr,
 ) {
     let mem_size = Bytes(heap_end - heap_base);
 
@@ -83,20 +85,17 @@ unsafe fn mark_compact<M: Memory, SetHp: Fn(u32)>(
 
     mark_static_roots(mem, static_roots, heap_base);
 
-    if (*closure_table_loc).unskew() >= heap_base as usize {
-        mark_object(mem, *closure_table_loc, heap_base);
+    if (*closure_table_ptr_loc).unskew() >= heap_base as usize {
+        // TODO: No need to check if closure table is already marked
+        mark_object(mem, *closure_table_ptr_loc, heap_base);
+        // Similar to `mark_root_mutbox_fields`, `closure_table_ptr_loc` is in static heap so it
+        // will be readable when we unthread closure table
+        thread(closure_table_ptr_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 +107,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 mark_object<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,87 +126,89 @@ unsafe fn mark_object<M: Memory>(mem: &mut M, obj: SkewedPtr, heap_base: u32) {
     }
 
     set_bit(obj_idx);
-    push_mark_stack(mem, obj as usize);
+    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| {
-        mark_object(mem, *field_addr, heap_base);
-    });
-}
+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| {
+        let field_value = *field_addr;
+        mark_object(mem, field_value, heap_base);
 
-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);
-    }
-    // No need to thread closure table here as it's on heap and we already marked it
+        // Thread if backwards pointer
+        if field_value.unskew() < obj as usize {
+            thread(field_addr);
+        }
+    });
 }
 
-/// 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) {
-    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
-        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);
-
-        free += size;
-
-        bit = bitmap_iter.next();
+/// 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)
+        mark_object(mem, *field_addr, heap_base);
+        // It's OK to thread forward pointers here as the static objects won't be moved, so we will
+        // be able to unthread objects pointed by these fields later.
+        thread(field_addr);
     }
 }
 
-/// 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) {
+/// Linearly scan the heap, for each live object:
+///
+/// - Mark step threads all backwards pointers and pointers from roots, so unthread to update those
+///   pointers to the objects new location.
+///
+/// - Move the object
+///
+/// - Thread forward pointers of the object
+///
+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;
+        let p_new = free;
 
-        // Update backward references to the object's new location and restore object header
-        unthread(p, free);
+        // Update backwards references to the object's new location and restore object header
+        unthread(p, p_new);
 
-        // All references to the object now point to the new location, move the object
+        // 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);
+        if p_new as usize != p as usize {
+            memcpy_words(p_new as usize, p as usize, p_size_words);
         }
 
         free += p_size_words.to_bytes().0;
 
+        // Thread forward pointers of the object
+        thread_fwd_pointers(p_new as *mut Obj, heap_base);
+
         bit = bitmap_iter.next();
     }
 
     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));
+/// Thread forwards pointers in object
+unsafe fn thread_fwd_pointers(obj: *mut Obj, heap_base: u32) {
+    visit_pointer_fields(obj, obj.tag(), heap_base as usize, |field_addr| {
+        if (*field_addr).unskew() > field_addr as usize {
+            thread(field_addr)
+        }
+    });
 }
 
+/// Thread a pointer field
 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;
@@ -212,7 +217,7 @@ unsafe fn thread(field: *mut SkewedPtr) {
     (*pointed).tag = field as u32;
 }
 
-/// Unthread all references, replacing with `new_loc`
+/// Unthread all references at given header, replacing with `new_loc`. Restores object header.
 unsafe fn unthread(obj: *mut Obj, new_loc: u32) {
     // NOTE: For this to work heap addresses need to be greater than the largest value for object
     // headers. Currently this holds. TODO: Document this better.

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)
 }