runtime (gc): scan callee-saved registers while marking stack

src/internal/task/task_coroutine.go

@@ -95,3 +95,9 @@ func fake() {
 	go func() {}()
 	Pause()
 }
+
+// OnSystemStack returns whether the caller is running on the system stack.
+func OnSystemStack() bool {
+	// This scheduler does not do any stack switching.
+	return true
+}

src/internal/task/task_none.go

@@ -27,3 +27,9 @@ type state struct{}
 func (t *Task) Resume() {
 	runtimePanic("scheduler is disabled")
 }
+
+// OnSystemStack returns whether the caller is running on the system stack.
+func OnSystemStack() bool {
+	// This scheduler does not do any stack switching.
+	return true
+}

src/internal/task/task_stack.go

@@ -72,3 +72,9 @@ func start(fn uintptr, args unsafe.Pointer) {
 	t.state.initialize(fn, args)
 	runqueuePushBack(t)
 }
+
+// OnSystemStack returns whether the caller is running on the system stack.
+func OnSystemStack() bool {
+	// If there is not an active goroutine, then this must be running on the system stack.
+	return Current() == nil
+}

src/runtime/gc_stack_raw.go

@@ -3,12 +3,38 @@
 
 package runtime
 
+import "internal/task"
+
 // markStack marks all root pointers found on the stack.
 //
 // This implementation is conservative and relies on the stack top (provided by
 // the linker) and getting the current stack pointer from a register. Also, it
 // assumes a descending stack. Thus, it is not very portable.
 func markStack() {
+	// Scan the current stack, and all current registers.
+	scanCurrentStack()
+
+	if !task.OnSystemStack() {
 		// Mark system stack.
 		markRoots(getSystemStackPointer(), stackTop)
+	}
+}
+
+//go:export tinygo_scanCurrentStack
+func scanCurrentStack()
+
+//go:export tinygo_scanstack
+func scanstack(sp uintptr) {
+	// Mark current stack.
+	// This function is called by scanCurrentStack, after pushing all registers onto the stack.
+	// Callee-saved registers have been pushed onto stack by tinygo_localscan, so this will scan them too.
+	if task.OnSystemStack() {
+		// This is the system stack.
+		// Scan all words on the stack.
+		markRoots(sp, stackTop)
+	} else {
+		// This is a goroutine stack.
+		// It is an allocation, so scan it as if it were a value in a global.
+		markRoot(0, sp)
+	}
 }

src/runtime/scheduler_avr.S

@@ -187,3 +187,55 @@ tinygo_switchToScheduler:
 
     // Return into the scheduler, as if tinygo_switchToTask was a regular call.
     ret
+
+.global tinygo_scanCurrentStack
+.type tinygo_scanCurrentStack, %function
+tinygo_scanCurrentStack:
+    // Save callee-saved registers.
+    push r29 // Y
+    push r28 // Y
+    push r17
+    push r16
+    push r15
+    push r14
+    push r13
+    push r12
+    push r11
+    push r10
+    push r9
+    push r8
+    push r7
+    push r6
+    push r5
+    push r4
+    push r3
+    push r2
+
+    // Scan the stack.
+    in  r24, 0x3d; SPL
+    in  r25, 0x3e; SPH
+#if __AVR_ARCH__ == 2 || __AVR_ARCH__ == 25
+    rcall tinygo_scanstack
+#else
+    call tinygo_scanstack
+#endif
+
+    // Restore callee-saved registers.
+    pop r2
+    pop r3
+    pop r4
+    pop r5
+    pop r6
+    pop r7
+    pop r8
+    pop r9
+    pop r10
+    pop r11
+    pop r12
+    pop r13
+    pop r14
+    pop r15
+    pop r16
+    pop r17
+    pop r28 // Y
+    pop r29 // Y

src/runtime/scheduler_cortexm.S

@@ -111,3 +111,26 @@ tinygo_swapTask:
     mov r11, r3
     pop {pc}
     #endif
+
+.global  tinygo_scanCurrentStack
+.type    tinygo_scanCurrentStack, %function
+tinygo_scanCurrentStack:
+    // Save callee-saved registers onto the stack.
+    #if defined(__thumb2__)
+    push {r4-r11, lr}
+    #else
+    mov r0, r8
+    mov r1, r9
+    mov r2, r10
+    mov r3, r11
+    push {r0-r3, lr}
+    push {r4-r7}
+    #endif
+
+    // Scan the stack.
+    mov r0, sp
+    bl tinygo_scanstack
+
+    // Restore stack state and return.
+    add sp, #32
+    pop {pc}

src/runtime/scheduler_tinygoriscv.S

@@ -0,0 +1,29 @@
+.section .text.tinygo_scanCurrentStack
+.global  tinygo_scanCurrentStack
+.type    tinygo_scanCurrentStack, %function
+tinygo_scanCurrentStack:
+   // Push callee-saved registers onto the stack.
+   addi sp, sp, -64
+   sw ra,  60(sp)
+   sw s11, 56(sp)
+   sw s10, 52(sp)
+   sw s9,  48(sp)
+   sw s8,  44(sp)
+   sw s7,  40(sp)
+   sw s6,  36(sp)
+   sw s5,  32(sp)
+   sw s4,  28(sp)
+   sw s3,  24(sp)
+   sw s2,  20(sp)
+   sw s1,  16(sp)
+   sw s0,  12(sp)
+
+   // Scan the stack.
+   mv a0, sp
+   call tinygo_scanstack
+
+   // Restore stack state.
+   addi sp, sp, 64
+
+   // Return to the caller.
+   ret

targets/riscv.json

@@ -22,7 +22,8 @@
 		"--gc-sections"
 	],
 	"extra-files": [
-		"src/device/riscv/start.S"
+		"src/device/riscv/start.S",
+		"src/runtime/scheduler_tinygoriscv.S"
 	],
 	"gdb": "riscv64-unknown-elf-gdb"
 }