From 55144ad6085c5c0c28d154589c8a7f2f1568364a Mon Sep 17 00:00:00 2001
From: Phil Kedy <phil.kedy@gmail.com>
Date: Thu, 12 Sep 2019 02:35:43 -0400
Subject: [PATCH] WASM: Support for setting an imported function's module name
 (#455)

* wasm: add support for setting a function's Wasm import module name by using the //go:wasm-module comment.
---
 compiler/compiler.go        |  5 +++++
 ir/ir.go                    | 12 ++++++++++++
 src/examples/wasm/README.md |  4 +++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/compiler/compiler.go b/compiler/compiler.go
index adf3fef5..f64b6790 100644
--- a/compiler/compiler.go
+++ b/compiler/compiler.go
@@ -821,6 +821,11 @@ func (c *Compiler) parseFuncDecl(f *ir.Function) *Frame {
 	// External/exported functions may not retain pointer values.
 	// https://golang.org/cmd/cgo/#hdr-Passing_pointers
 	if f.IsExported() {
+		// Set the wasm-import-module attribute if the function's module is set.
+		if f.Module() != "" {
+			wasmImportModuleAttr := c.ctx.CreateStringAttribute("wasm-import-module", f.Module())
+			frame.fn.LLVMFn.AddFunctionAttr(wasmImportModuleAttr)
+		}
 		nocaptureKind := llvm.AttributeKindID("nocapture")
 		nocapture := c.ctx.CreateEnumAttribute(nocaptureKind, 0)
 		for i, typ := range paramTypes {
diff --git a/ir/ir.go b/ir/ir.go
index f6917356..f3327a79 100644
--- a/ir/ir.go
+++ b/ir/ir.go
@@ -28,6 +28,7 @@ type Program struct {
 type Function struct {
 	*ssa.Function
 	LLVMFn    llvm.Value
+	module    string     // go:wasm-module
 	linkName  string     // go:linkname, go:export, go:interrupt
 	exported  bool       // go:export
 	nobounds  bool       // go:nobounds
@@ -229,6 +230,12 @@ func (f *Function) parsePragmas() {
 				}
 				f.linkName = parts[1]
 				f.exported = true
+			case "//go:wasm-module":
+				// Alternative comment for setting the import module.
+				if len(parts) != 2 {
+					continue
+				}
+				f.module = parts[1]
 			case "//go:inline":
 				f.inline = InlineHint
 			case "//go:noinline":
@@ -291,6 +298,11 @@ func (f *Function) Inline() InlineType {
 	return f.inline
 }
 
+// Return the module name if not the default.
+func (f *Function) Module() string {
+	return f.module
+}
+
 // Return the link name for this function.
 func (f *Function) LinkName() string {
 	if f.linkName != "" {
diff --git a/src/examples/wasm/README.md b/src/examples/wasm/README.md
index 145cb537..a17653b6 100644
--- a/src/examples/wasm/README.md
+++ b/src/examples/wasm/README.md
@@ -3,7 +3,9 @@
 The examples here show two different ways of using WebAssembly with TinyGo:
 
 1. Defining and exporting functions via the `//go:export <name>` directive. See
-[the export folder](./export) for an example of this.
+[the export folder](./export) for an example of this.  Additionally, the Wasm
+module (which has a default value of `env`) can be specified using
+`//go:wasm-module <module>`.
 1. Defining and executing a `func main()`. This is similar to how the Go
 standard library implementation works. See [the main folder](./main) for an
 example of this.