Add support for using wizer to handle running initAll()

Author: dgryski

Diff for: builder/build.go

@@ -834,6 +834,33 @@ func Build(pkgName, outpath, tmpdir string, config *compileopts.Config) (BuildRe
 				}
 			}
 
+			if config.Options.WizerInit {
+				var args []string
+
+				resultWizer := result.Executable + "-wizer"
+
+				args = append(args,
+					"--allow-wasi",
+					"--wasm-bulk-memory=true",
+					"-f", "runtime.wizerInit",
+					result.Executable,
+					"-o", resultWizer,
+				)
+
+				cmd := exec.Command(goenv.Get("WIZER"), args...)
+				cmd.Stdout = os.Stdout
+				cmd.Stderr = os.Stderr
+
+				err := cmd.Run()
+				if err != nil {
+					return fmt.Errorf("wizer failed: %w", err)
+				}
+
+				if err := os.Rename(resultWizer, result.Executable); err != nil {
+					return fmt.Errorf("rename failed: %w", err)
+				}
+			}
+
 			// Print code size if requested.
 			if config.Options.PrintSizes == "short" || config.Options.PrintSizes == "full" {
 				packagePathMap := make(map[string]string, len(lprogram.Packages))
@@ -1036,9 +1063,10 @@ func createEmbedObjectFile(data, hexSum, sourceFile, sourceDir, tmpdir string, c
 // needed to convert a program to its final form. Some transformations are not
 // optional and must be run as the compiler expects them to run.
 func optimizeProgram(mod llvm.Module, config *compileopts.Config) error {
-	err := interp.Run(mod, config.Options.InterpTimeout, config.DumpSSA())
-	if err != nil {
-		return err
+	if !config.Options.WizerInit {
+		if err := interp.Run(mod, config.Options.InterpTimeout, config.DumpSSA()); err != nil {
+			return err
+		}
 	}
 	if config.VerifyIR() {
 		// Only verify if we really need it.
@@ -1054,7 +1082,7 @@ func optimizeProgram(mod llvm.Module, config *compileopts.Config) error {
 	}
 
 	// Insert values from -ldflags="-X ..." into the IR.
-	err = setGlobalValues(mod, config.Options.GlobalValues)
+	err := setGlobalValues(mod, config.Options.GlobalValues)
 	if err != nil {
 		return err
 	}

Diff for: compileopts/options.go

@@ -52,6 +52,7 @@ type Options struct {
 	Monitor         bool
 	BaudRate        int
 	Timeout         time.Duration
+	WizerInit       bool
 }
 
 // Verify performs a validation on the given options, raising an error if options are not valid.

Diff for: goenv/goenv.go

@@ -107,6 +107,35 @@ func Get(name string) string {
 		}
 
 		return findWasmOpt()
+
+	case "WIZER":
+		if path := os.Getenv("WIZER"); path != "" {
+			wizerBin := "wizer"
+			if runtime.GOOS == "windows" {
+				wizerBin += ".exe"
+			}
+
+			_, err := exec.LookPath(wizerBin)
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "cannot use %q as wasm-opt (from WASMOPT environment variable): %s", path, err.Error())
+				os.Exit(1)
+			}
+
+			return path
+		}
+
+		wizerBin := "wizer"
+		if runtime.GOOS == "windows" {
+			wizerBin += ".exe"
+		}
+
+		path, err := exec.LookPath(wizerBin)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "cannot use %q as wizer: %s", path, err.Error())
+			os.Exit(1)
+		}
+		return path
+
 	default:
 		return ""
 	}

Diff for: main.go

@@ -1420,6 +1420,7 @@ func main() {
 	cpuprofile := flag.String("cpuprofile", "", "cpuprofile output")
 	monitor := flag.Bool("monitor", false, "enable serial monitor")
 	baudrate := flag.Int("baudrate", 115200, "baudrate of serial monitor")
+	wizerInit := flag.Bool("wizer-init", false, "use wizer for package initialization")
 
 	var flagJSON, flagDeps, flagTest bool
 	if command == "help" || command == "list" || command == "info" || command == "build" {
@@ -1481,6 +1482,11 @@ func main() {
 		ocdCommands = strings.Split(*ocdCommandsString, ",")
 	}
 
+	if *wizerInit && *target != "wasi" {
+		fmt.Fprintf(os.Stderr, "-wizer-init only makes sense with -target=wasi, got -target=%q\n", *target)
+		os.Exit(1)
+	}
+
 	options := &compileopts.Options{
 		GOOS:            goenv.Get("GOOS"),
 		GOARCH:          goenv.Get("GOARCH"),
@@ -1512,7 +1518,9 @@ func main() {
 		Monitor:         *monitor,
 		BaudRate:        *baudrate,
 		Timeout:         *timeout,
+		WizerInit:       *wizerInit,
 	}
+
 	if *printCommands {
 		options.PrintCommands = printCommand
 	}

Diff for: src/runtime/runtime.go

@@ -12,6 +12,12 @@ const Compiler = "tinygo"
 // package.
 func initAll()
 
+var runtimeInitialized = false
+
+func runtime_is_initialized() bool {
+	return runtimeInitialized
+}
+
 //go:linkname callMain main.main
 func callMain()
 

Diff for: src/runtime/runtime_wasm_wasi.go

@@ -15,10 +15,29 @@ func __wasm_call_ctors()
 
 //export _start
 func _start() {
+	runtimeInitialize()
+	run()
+}
+
+//export runtime.wizerInit
+func runtimeInitialize() {
+	if runtimeInitialized {
+		// Second time initialization is happening. Refresh environment
+		// to whatever our current host gives us instead of whatever
+		// libc cached.
+		reset_libc_environment()
+		return
+	}
 	// 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()
+	initHeap()
+	if hasScheduler {
+		go initAll()
+	} else {
+		initAll()
+	}
+	runtimeInitialized = true
 }
 
 // Read the command line arguments from WASI.
@@ -29,6 +48,17 @@ func init() {
 	__wasm_call_ctors()
 }
 
+//export __wasilibc_deinitialize_environ
+func __wasilibc_deinitialize_environ()
+
+//export __wasilibc_initialize_environ
+func __wasilibc_initialize_environ()
+
+func reset_libc_environment() {
+	__wasilibc_deinitialize_environ()
+	__wasilibc_initialize_environ()
+}
+
 var args []string
 
 //go:linkname os_runtime_args os.runtime_args
@@ -39,7 +69,10 @@ func os_runtime_args() []string {
 		var argc, argv_buf_size uint32
 		args_sizes_get(&argc, &argv_buf_size)
 		if argc == 0 {
-			return nil
+			// Most things expect os.Args to have at least the
+			// program name.  We don't have one, but also don't
+			// return just a nil slice.
+			return []string{""}
 		}
 
 		// Obtain the command line arguments

Diff for: src/runtime/scheduler_any.go

@@ -19,9 +19,14 @@ func sleep(duration int64) {
 // run is called by the program entry point to execute the go program.
 // With a scheduler, init and the main function are invoked in a goroutine before starting the scheduler.
 func run() {
-	initHeap()
+	if !runtimeInitialized {
+		initHeap()
+	}
 	go func() {
-		initAll()
+		if !runtimeInitialized {
+			initAll()
+			runtimeInitialized = true
+		}
 		callMain()
 		schedulerDone = true
 	}()

Diff for: src/runtime/scheduler_none.go

@@ -20,8 +20,11 @@ func getSystemStackPointer() uintptr {
 // run is called by the program entry point to execute the go program.
 // With the "none" scheduler, init and the main function are invoked directly.
 func run() {
-	initHeap()
-	initAll()
+	if !runtimeInitialized {
+		initHeap()
+		initAll()
+		runtimeInitialized = true
+	}
 	callMain()
 }
 

Diff for: src/syscall/build_asserts.go

@@ -0,0 +1,7 @@
+//go:build syscall_asserts
+
+package syscall
+
+// enable assertions for checking Environ during runtime preinitialization.
+// This is to catch people using os.Environ() during package inits with wizer.
+const panicOnEnvironDuringInitAll = true

Diff for: src/syscall/build_noasserts.go

@@ -0,0 +1,7 @@
+//go:build !syscall_asserts
+
+package syscall
+
+// enable assertions for checking Environ during runtime preinitialization.
+// This is to catch people using os.Environ() during package init with wizer.
+const panicOnEnvironDuringInitAll = false

Diff for: src/syscall/syscall_libc.go

@@ -243,7 +243,13 @@ func Mprotect(b []byte, prot int) (err error) {
 	return
 }
 
+//go:linkname runtime_is_initialized runtime.runtime_is_initialized
+func runtime_is_initialized() bool
+
 func Environ() []string {
+	if panicOnEnvironDuringInitAll && !runtime_is_initialized() {
+		panic("syscall.Environ() called during runtime preinitialization")
+	}
 
 	// This function combines all the environment into a single allocation.
 	// While this optimizes for memory usage and garbage collector