From 319df9bd59b63291d2ced3db5c070ae14f201f3f Mon Sep 17 00:00:00 2001
From: Phil Kedy <phil.kedy@gmail.com>
Date: Sat, 20 Jul 2019 00:31:38 -0400
Subject: [PATCH] Support for setting a function's WASM import module name.

---
 compiler/compiler.go        |  5 +++++
 ir/ir.go                    | 24 ++++++++++++++++++++++--
 src/examples/wasm/README.md |  5 +++--
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/compiler/compiler.go b/compiler/compiler.go
index 1800f462af..ab0bfa9009 100644
--- a/compiler/compiler.go
+++ b/compiler/compiler.go
@@ -791,6 +791,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 f69173568e..18c0f5f8d9 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:module, go:export
 	linkName  string     // go:linkname, go:export, go:interrupt
 	exported  bool       // go:export
 	nobounds  bool       // go:nobounds
@@ -223,11 +224,25 @@ func (f *Function) parsePragmas() {
 			}
 			parts := strings.Fields(text)
 			switch parts[0] {
-			case "//go:export":
+			case "//go:module":
+				// Alternative comment for setting the import module.
 				if len(parts) != 2 {
 					continue
 				}
-				f.linkName = parts[1]
+				f.module = parts[1]
+			case "//go:export":
+				// Accept:
+				// [module] linkName
+				// where module is an optional WASM import module name.
+				switch len(parts) {
+				case 2:
+					f.linkName = parts[1]
+				case 3:
+					f.module = parts[1]
+					f.linkName = parts[2]
+				default:
+					continue
+				}
 				f.exported = true
 			case "//go:inline":
 				f.inline = InlineHint
@@ -291,6 +306,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 8a50ba2527..12c29f92c8 100644
--- a/src/examples/wasm/README.md
+++ b/src/examples/wasm/README.md
@@ -2,8 +2,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.
+1. Defining and exporting functions via the `//go:export [module] <name>` directive. See
+[the export folder](./export) for an example of this.  `module` is optional and overrides
+the default module name of "env".
 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.