Mark-compact GC: update STACK_TOP when growing mark stack

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

@@ -1,16 +1,24 @@
 use crate::memory::TestMemory;
 
 use motoko_rts::gc::mark_compact::mark_stack::{
-    alloc_mark_stack, free_mark_stack, pop_mark_stack, push_mark_stack,
+    alloc_mark_stack, free_mark_stack, grow_stack, pop_mark_stack, push_mark_stack,
+    INIT_STACK_SIZE, STACK_BASE, STACK_PTR, STACK_TOP,
 };
 use motoko_rts::memory::Memory;
-use motoko_rts::types::Words;
+use motoko_rts::types::{size_of, Blob, Words};
 
 use proptest::test_runner::{Config, TestCaseError, TestCaseResult, TestRunner};
 
 pub unsafe fn test() {
     println!("Testing mark stack ...");
 
+    test_push_pop();
+    test_grow_stack();
+}
+
+fn test_push_pop() {
+    println!("  Testing push/pop");
+
     let mut proptest_runner = TestRunner::new(Config {
         cases: 100,
         failure_persistence: None,
@@ -55,3 +63,28 @@ fn test_<M: Memory>(mem: &mut M, n_objs: u32) -> TestCaseResult {
 
     Ok(())
 }
+
+unsafe fn test_grow_stack() {
+    println!("  Testing grow_stack");
+
+    // Allow doubling twice
+    let mut mem = TestMemory::new(
+        size_of::<Blob>() + INIT_STACK_SIZE + INIT_STACK_SIZE + INIT_STACK_SIZE * 2,
+    );
+
+    alloc_mark_stack(&mut mem);
+
+    let mut current_size = INIT_STACK_SIZE.as_usize();
+    assert_eq!(STACK_BASE.add(current_size), STACK_TOP);
+    assert_eq!(STACK_BASE, STACK_PTR);
+
+    grow_stack(&mut mem);
+    current_size *= 2;
+    assert_eq!(STACK_BASE.add(current_size), STACK_TOP);
+    assert_eq!(STACK_BASE, STACK_PTR);
+
+    grow_stack(&mut mem);
+    current_size *= 2;
+    assert_eq!(STACK_BASE.add(current_size), STACK_TOP);
+    assert_eq!(STACK_BASE, STACK_PTR);
+}

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

@@ -13,13 +13,13 @@ pub const INIT_STACK_SIZE: Words<u32> = Words(64);
 static mut STACK_BLOB_PTR: *mut Blob = null_mut();
 
 /// Bottom of the mark stack
-static mut STACK_BASE: *mut usize = null_mut();
+pub static mut STACK_BASE: *mut usize = null_mut();
 
 /// Top of the mark stack
-static mut STACK_TOP: *mut usize = null_mut();
+pub static mut STACK_TOP: *mut usize = null_mut();
 
 /// Next free slot in the mark stack
-static mut STACK_PTR: *mut usize = null_mut();
+pub static mut STACK_PTR: *mut usize = null_mut();
 
 pub unsafe fn alloc_mark_stack<M: Memory>(mem: &mut M) {
     debug_assert!(STACK_BLOB_PTR.is_null());
@@ -39,7 +39,7 @@ pub unsafe fn free_mark_stack() {
 }
 
 /// Doubles the stack size
-unsafe fn grow_stack<M: Memory>(mem: &mut M) {
+pub unsafe fn grow_stack<M: Memory>(mem: &mut M) {
     let stack_cap: Words<u32> = STACK_BLOB_PTR.len().to_words();
     let p = mem.alloc_words(stack_cap).unskew() as *mut usize;
 
@@ -48,6 +48,7 @@ unsafe fn grow_stack<M: Memory>(mem: &mut M) {
 
     let new_cap: Words<u32> = Words(stack_cap.0 * 2);
     (*STACK_BLOB_PTR).len = new_cap.to_bytes();
+    STACK_TOP = STACK_BASE.add(new_cap.as_usize());
 }
 
 pub unsafe fn push_mark_stack<M: Memory>(mem: &mut M, obj: usize, obj_tag: Tag) {

rts/motoko-rts/src/types.rs

@@ -17,6 +17,10 @@ impl Words<u32> {
     pub fn to_bytes(self) -> Bytes<u32> {
         Bytes(self.0 * WORD_SIZE)
     }
+
+    pub fn as_usize(self) -> usize {
+        self.0 as usize
+    }
 }
 
 impl<A: Add<Output = A>> Add for Words<A> {