runtime/interrupt: add Checkpoint type

This type can be used to jump back to a previous position in a program
from inside an interrupt. This is useful for baremetal systems that
implement wfi but not wfe, and therefore have no easy (race-free) way to
wait until a flag gets changed inside an interrupt. This is an issue on
RISC-V, where this is racy (the interrupt might happen after the check
but before the wfi instruction):

    configureInterrupt()
    for flag.Load() != 0 {
        riscv.Asm("wfi")
    }

compiler/compiler.go

@@ -1895,6 +1895,8 @@ func (b *builder) createFunctionCall(instr *ssa.CallCommon) (llvm.Value, error)
 		case name == "runtime.exportedFuncPtr":
 			_, ptr := b.getFunction(instr.Args[0].(*ssa.Function))
 			return b.CreatePtrToInt(ptr, b.uintptrType, ""), nil
+		case name == "(*runtime/interrupt.Checkpoint).Save":
+			return b.createInterruptCheckpoint(instr.Args[0]), nil
 		case name == "internal/abi.FuncPCABI0":
 			retval := b.createDarwinFuncPCABI0Call(instr)
 			if !retval.IsNil() {

compiler/defer.go

@@ -103,10 +103,11 @@ func (b *builder) createLandingPad() {
 	b.CreateBr(b.blockEntries[b.fn.Recover])
 }
 
-// createInvokeCheckpoint saves the function state at the given point, to
-// continue at the landing pad if a panic happened. This is implemented using a
-// setjmp-like construct.
-func (b *builder) createInvokeCheckpoint() {
+// Create a checkpoint (similar to setjmp). This emits inline assembly that
+// stores the current program counter inside the ptr address (actually
+// ptr+sizeof(ptr)) and then returns a boolean indicating whether this is the
+// normal flow (false) or we jumped here from somewhere else (true).
+func (b *builder) createCheckpoint(ptr llvm.Value) llvm.Value {
 	// Construct inline assembly equivalents of setjmp.
 	// The assembly works as follows:
 	//   * All registers (both callee-saved and caller saved) are clobbered
@@ -217,11 +218,19 @@ li a0, 0
 		// This case should have been handled by b.supportsRecover().
 		b.addError(b.fn.Pos(), "unknown architecture for defer: "+b.archFamily())
 	}
-	asmType := llvm.FunctionType(resultType, []llvm.Type{b.deferFrame.Type()}, false)
+	asmType := llvm.FunctionType(resultType, []llvm.Type{b.dataPtrType}, false)
 	asm := llvm.InlineAsm(asmType, asmString, constraints, false, false, 0, false)
-	result := b.CreateCall(asmType, asm, []llvm.Value{b.deferFrame}, "setjmp")
+	result := b.CreateCall(asmType, asm, []llvm.Value{ptr}, "setjmp")
 	result.AddCallSiteAttribute(-1, b.ctx.CreateEnumAttribute(llvm.AttributeKindID("returns_twice"), 0))
 	isZero := b.CreateICmp(llvm.IntEQ, result, llvm.ConstInt(resultType, 0, false), "setjmp.result")
+	return isZero
+}
+
+// createInvokeCheckpoint saves the function state at the given point, to
+// continue at the landing pad if a panic happened. This is implemented using a
+// setjmp-like construct.
+func (b *builder) createInvokeCheckpoint() {
+	isZero := b.createCheckpoint(b.deferFrame)
 	continueBB := b.insertBasicBlock("")
 	b.CreateCondBr(isZero, continueBB, b.landingpad)
 	b.SetInsertPointAtEnd(continueBB)

compiler/inlineasm.go

@@ -249,3 +249,15 @@ func (b *builder) emitCSROperation(call *ssa.CallCommon) (llvm.Value, error) {
 		return llvm.Value{}, b.makeError(call.Pos(), "unknown CSR operation: "+name)
 	}
 }
+
+// Implement runtime/interrupt.Checkpoint.Save. It needs to be implemented
+// directly at the call site. If it isn't implemented directly at the call site
+// (but instead through a function call), it might result in an overwritten
+// stack in the non-jump return case.
+func (b *builder) createInterruptCheckpoint(ptr ssa.Value) llvm.Value {
+	addr := b.getValue(ptr, ptr.Pos())
+	b.createNilCheck(ptr, addr, "deref")
+	stackPointer := b.readStackPointer()
+	b.CreateStore(stackPointer, addr)
+	return b.createCheckpoint(addr)
+}

src/runtime/asm_riscv.S

@@ -50,3 +50,12 @@ tinygo_longjmp:
     lw sp, 0(a0)       // jumpSP
     lw a1, REGSIZE(a0) // jumpPC
     jr a1
+
+.section .text.tinygo_checkpointJump
+.global tinygo_checkpointJump
+tinygo_checkpointJump:
+    // Note: the code we jump to assumes a0 is non-zero, which is already the
+    // case because that's the stack pointer.
+    mv sp, a0          // jumpSP
+    csrw mepc, a1      // update MEPC value, so we resume there after the mret
+    mret               // jump to jumpPC

src/runtime/interrupt/checkpoint.go

@@ -0,0 +1,62 @@
+package interrupt
+
+// A checkpoint is a setjmp like buffer, that can be used as a flag for
+// interrupts.
+//
+// It can be used as follows:
+//
+//	// global var
+//	var c Checkpoint
+//
+//	// to set up the checkpoint and wait for it
+//	if c.Save() {
+//		setupInterrupt()
+//		for {
+//			waitForInterrupt()
+//		}
+//	}
+//
+//	// Inside the interrupt handler:
+//	if c.Saved() {
+//		c.Jump()
+//	}
+type Checkpoint struct {
+	jumpSP uintptr
+	jumpPC uintptr
+}
+
+// Save the execution state in the given checkpoint, overwriting a previous
+// saved checkpoint.
+//
+// This function returns twice: once the normal way after saving (returning
+// true) and once after jumping (returning false).
+//
+// This function is a compiler intrinsic, it is not implemented in Go.
+func (c *Checkpoint) Save() bool
+
+// Returns whether a jump point was saved (and not erased due to a jump).
+func (c *Checkpoint) Saved() bool {
+	return c.jumpPC != 0
+}
+
+// Jump to the point where the execution state was saved, and erase the saved
+// jump point. This must *only* be called from inside an interrupt.
+//
+// This method does not return in the conventional way, it resumes execution at
+// the last point a checkpoint was saved.
+func (c *Checkpoint) Jump() {
+	if !c.Saved() {
+		panic("runtime/interrupt: no checkpoint was saved")
+	}
+	jumpPC := c.jumpPC
+	jumpSP := c.jumpSP
+	c.jumpPC = 0
+	c.jumpSP = 0
+	if jumpPC == 0 {
+		panic("jumping to 0")
+	}
+	checkpointJump(jumpSP, jumpPC)
+}
+
+//export tinygo_checkpointJump
+func checkpointJump(jumpSP, jumpPC uintptr)