wasm: add //go:wasmexport support (#4451)

This adds support for the `//go:wasmexport` pragma as proposed here:
https://github.com/golang/go/issues/65199

It is currently implemented only for wasip1 and wasm-unknown, but it is
certainly possible to extend it to other targets like GOOS=js and
wasip2.

src/runtime/runtime_tinygowasm.go

@@ -1,5 +1,8 @@
 //go:build tinygo.wasm && !wasm_unknown && !wasip2
 
+// This file is for wasm/wasip1 and for wasm/js, which both use much of the
+// WASIp1 API.
+
 package runtime
 
 import (

src/runtime/runtime_wasip1.go

@@ -13,15 +13,6 @@ type timeUnit int64
 //export __wasm_call_ctors
 func __wasm_call_ctors()
 
-//export _start
-func _start() {
-	// These need to be initialized early so that the heap can be initialized.
-	heapStart = uintptr(unsafe.Pointer(&heapStartSymbol))
-	heapEnd = uintptr(wasm_memory_size(0) * wasmPageSize)
-	run()
-	__stdio_exit()
-}
-
 // Read the command line arguments from WASI.
 // For example, they can be passed to a program with wasmtime like this:
 //
@@ -100,6 +91,10 @@ func ticks() timeUnit {
 	return timeUnit(nano)
 }
 
+func beforeExit() {
+	__stdio_exit()
+}
+
 // Implementations of WASI APIs
 
 //go:wasmimport wasi_snapshot_preview1 args_get

src/runtime/runtime_wasm_js_scheduler.go

@@ -14,7 +14,7 @@ func resume() {
 	}
 
 	wasmNested = true
-	scheduler()
+	scheduler(false)
 	wasmNested = false
 }
 
@@ -26,6 +26,6 @@ func go_scheduler() {
 	}
 
 	wasmNested = true
-	scheduler()
+	scheduler(false)
 	wasmNested = false
 }

src/runtime/runtime_wasm_unknown.go

@@ -2,7 +2,8 @@
 
 package runtime
 
-import "unsafe"
+// TODO: this is essentially reactor mode wasm. So we might want to support
+// -buildmode=c-shared (and default to it).
 
 type timeUnit int64
 
@@ -11,14 +12,6 @@ type timeUnit int64
 //export __wasm_call_ctors
 func __wasm_call_ctors()
 
-//export _initialize
-func _initialize() {
-	// These need to be initialized early so that the heap can be initialized.
-	heapStart = uintptr(unsafe.Pointer(&heapStartSymbol))
-	heapEnd = uintptr(wasm_memory_size(0) * wasmPageSize)
-	initAll()
-}
-
 func init() {
 	__wasm_call_ctors()
 }
@@ -40,3 +33,6 @@ func sleepTicks(d timeUnit) {
 func ticks() timeUnit {
 	return timeUnit(0)
 }
+
+func beforeExit() {
+}

src/runtime/runtime_wasmentry.go

@@ -0,0 +1,100 @@
+//go:build tinygo.wasm && !wasip2 && !js
+
+package runtime
+
+// Entry points for WebAssembly modules, and runtime support for
+// //go:wasmexport: runtime.wasmExport* function calls are inserted by the
+// compiler for //go:wasmexport support.
+
+import (
+	"internal/task"
+	"unsafe"
+)
+
+// This is the _start entry point, when using -buildmode=default.
+func wasmEntryCommand() {
+	// These need to be initialized early so that the heap can be initialized.
+	heapStart = uintptr(unsafe.Pointer(&heapStartSymbol))
+	heapEnd = uintptr(wasm_memory_size(0) * wasmPageSize)
+	wasmExportState = wasmExportStateInMain
+	run()
+	wasmExportState = wasmExportStateExited
+	beforeExit()
+}
+
+// This is the _initialize entry point, when using -buildmode=c-shared.
+func wasmEntryReactor() {
+	// This function is called before any //go:wasmexport functions are called
+	// to initialize everything. It must not block.
+
+	// Initialize the heap.
+	heapStart = uintptr(unsafe.Pointer(&heapStartSymbol))
+	heapEnd = uintptr(wasm_memory_size(0) * wasmPageSize)
+	initHeap()
+
+	if hasScheduler {
+		// A package initializer might do funky stuff like start a goroutine and
+		// wait until it completes, so we have to run package initializers in a
+		// goroutine.
+		go func() {
+			initAll()
+			wasmExportState = wasmExportStateReactor
+		}()
+		scheduler(true)
+		if wasmExportState != wasmExportStateReactor {
+			// Unlikely, but if package initializers do something blocking (like
+			// time.Sleep()), that's a bug.
+			runtimePanic("package initializer blocks")
+		}
+	} else {
+		// There are no goroutines (except for the main one, if you can call it
+		// that), so we can just run all the package initializers.
+		initAll()
+		wasmExportState = wasmExportStateReactor
+	}
+}
+
+// Track which state we're in: before (or during) init, running inside
+// main.main, after main.main returned, or reactor mode (after init).
+var wasmExportState uint8
+
+const (
+	wasmExportStateInit = iota
+	wasmExportStateInMain
+	wasmExportStateExited
+	wasmExportStateReactor
+)
+
+func wasmExportCheckRun() {
+	switch wasmExportState {
+	case wasmExportStateInit:
+		runtimePanic("//go:wasmexport function called before runtime initialization")
+	case wasmExportStateExited:
+		runtimePanic("//go:wasmexport function called after main.main returned")
+	}
+}
+
+// Called from within a //go:wasmexport wrapper (the one that's exported from
+// the wasm module) after the goroutine has been queued. Just run the scheduler,
+// and check that the goroutine finished when the scheduler is idle (as required
+// by the //go:wasmexport proposal).
+//
+// This function is not called when the scheduler is disabled.
+func wasmExportRun(done *bool) {
+	scheduler(true)
+	if !*done {
+		runtimePanic("//go:wasmexport function did not finish")
+	}
+}
+
+// Called from the goroutine wrapper for the //go:wasmexport function. It just
+// signals to the runtime that the //go:wasmexport call has finished, and can
+// switch back to the wasmExportRun function.
+//
+// This function is not called when the scheduler is disabled.
+func wasmExportExit() {
+	task.Pause()
+
+	// TODO: we could cache the allocated stack so we don't have to keep
+	// allocating a new stack on every //go:wasmexport call.
+}

src/runtime/scheduler.go

@@ -157,7 +157,13 @@ func removeTimer(tim *timer) bool {
 }
 
 // Run the scheduler until all tasks have finished.
-func scheduler() {
+// There are a few special cases:
+//   - When returnAtDeadlock is true, it also returns when there are no more
+//     runnable goroutines.
+//   - When using the asyncify scheduler, it returns when it has to wait
+//     (JavaScript uses setTimeout so the scheduler must return to the JS
+//     environment).
+func scheduler(returnAtDeadlock bool) {
 	// Main scheduler loop.
 	var now timeUnit
 	for !schedulerDone {
@@ -193,6 +199,9 @@ func scheduler() {
 		t := runqueue.Pop()
 		if t == nil {
 			if sleepQueue == nil && timerQueue == nil {
+				if returnAtDeadlock {
+					return
+				}
 				if asyncScheduler {
 					// JavaScript is treated specially, see below.
 					return

src/runtime/scheduler_any.go

@@ -25,7 +25,7 @@ func run() {
 		callMain()
 		schedulerDone = true
 	}()
-	scheduler()
+	scheduler(false)
 }
 
 const hasScheduler = true