wasm: fix functions exported through //export

When a function is exported using //export, but also had a
//go:wasm-module pragma, the //export name was ignored. The
//go:wasm-module doesn't actually do anything besides breaking the
export (exported functions don't have a module name).

I've refactored and cleaned up the code, and in the process removed this
weird edge case.

3 files changed, 39 insertions, 35 deletions

diff --git a/compiler/symbol.go b/compiler/symbol.go
index 6426604ec..3e782dae8 100644
--- a/compiler/symbol.go
+++ b/compiler/symbol.go
@@ -23,9 +23,9 @@ import (
 // The linkName value contains a valid link name, even if //go:linkname is not
 // present.
 type functionInfo struct {
-	module     string     // go:wasm-module
-	importName string     // go:linkname, go:export - The name the developer assigns
-	linkName   string     // go:linkname, go:export - The name that we map for the particular module -> importName
+	wasmModule string     // go:wasm-module
+	wasmName   string     // wasm-export-name or wasm-import-name in the IR
+	linkName   string     // go:linkname, go:export - the IR function name
 	section    string     // go:section - object file section name
 	exported   bool       // go:export, CGo
 	interrupt  bool       // go:interrupt
@@ -194,20 +194,14 @@ func (c *compilerContext) getFunction(fn *ssa.Function) (llvm.Type, llvm.Value)
 	// External/exported functions may not retain pointer values.
 	// https://golang.org/cmd/cgo/#hdr-Passing_pointers
 	if info.exported {
-		if c.archFamily() == "wasm32" {
+		if c.archFamily() == "wasm32" && len(fn.Blocks) == 0 {
 			// We need to add the wasm-import-module and the wasm-import-name
 			// attributes.
-			module := info.module
-			if module == "" {
-				module = "env"
+			if info.wasmModule != "" {
+				llvmFn.AddFunctionAttr(c.ctx.CreateStringAttribute("wasm-import-module", info.wasmModule))
 			}
-			llvmFn.AddFunctionAttr(c.ctx.CreateStringAttribute("wasm-import-module", module))
 
-			name := info.importName
-			if name == "" {
-				name = info.linkName
-			}
-			llvmFn.AddFunctionAttr(c.ctx.CreateStringAttribute("wasm-import-name", name))
+			llvmFn.AddFunctionAttr(c.ctx.CreateStringAttribute("wasm-import-name", info.wasmName))
 		}
 		nocaptureKind := llvm.AttributeKindID("nocapture")
 		nocapture := c.ctx.CreateEnumAttribute(nocaptureKind, 0)
@@ -260,10 +254,6 @@ func (c *compilerContext) parsePragmas(info *functionInfo, f *ssa.Function) {
 		return
 	}
 	if decl, ok := f.Syntax().(*ast.FuncDecl); ok && decl.Doc != nil {
-
-		// Our importName for a wasm module (if we are compiling to wasm), or llvm link name
-		var importName string
-
 		for _, comment := range decl.Doc.List {
 			text := comment.Text
 			if strings.HasPrefix(text, "//export ") {
@@ -281,7 +271,8 @@ func (c *compilerContext) parsePragmas(info *functionInfo, f *ssa.Function) {
 					continue
 				}
 
-				importName = parts[1]
+				info.linkName = parts[1]
+				info.wasmName = info.linkName
 				info.exported = true
 			case "//go:interrupt":
 				if hasUnsafeImport(f.Pkg.Pkg) {
@@ -289,10 +280,11 @@ func (c *compilerContext) parsePragmas(info *functionInfo, f *ssa.Function) {
 				}
 			case "//go:wasm-module":
 				// Alternative comment for setting the import module.
+				// This is deprecated, use //go:wasmimport instead.
 				if len(parts) != 2 {
 					continue
 				}
-				info.module = parts[1]
+				info.wasmModule = parts[1]
 			case "//go:wasmimport":
 				// Import a WebAssembly function, for example a WASI function.
 				// Original proposal: https://github.com/golang/go/issues/38248
@@ -302,8 +294,8 @@ func (c *compilerContext) parsePragmas(info *functionInfo, f *ssa.Function) {
 				}
 				c.checkWasmImport(f, comment.Text)
 				info.exported = true
-				info.module = parts[1]
-				info.importName = parts[2]
+				info.wasmModule = parts[1]
+				info.wasmName = parts[2]
 			case "//go:inline":
 				info.inline = inlineHint
 			case "//go:noinline":
@@ -347,17 +339,6 @@ func (c *compilerContext) parsePragmas(info *functionInfo, f *ssa.Function) {
 				}
 			}
 		}
-
-		// Set the importName for our exported function if we have one
-		if importName != "" {
-			if info.module == "" {
-				info.linkName = importName
-			} else {
-				// WebAssembly import
-				info.importName = importName
-			}
-		}
-
 	}
 }
 
diff --git a/compiler/testdata/pragma.go b/compiler/testdata/pragma.go
index 0a4be11ae..fa1d4b0e9 100644
--- a/compiler/testdata/pragma.go
+++ b/compiler/testdata/pragma.go
@@ -62,6 +62,19 @@ func exportedFunctionInSection() {
 //go:wasmimport modulename import1
 func declaredImport()
 
+// Legacy way of importing a function.
+//
+//go:wasm-module foobar
+//export imported
+func foobarImport()
+
+// The wasm-module pragma is not functional here, but it should be safe.
+//
+//go:wasm-module foobar
+//export exported
+func foobarExportModule() {
+}
+
 // This function should not: it's only a declaration and not a definition.
 //
 //go:section .special_function_section
diff --git a/compiler/testdata/pragma.ll b/compiler/testdata/pragma.ll
index 78cc29457..397772e5d 100644
--- a/compiler/testdata/pragma.ll
+++ b/compiler/testdata/pragma.ll
@@ -6,7 +6,7 @@ target triple = "wasm32-unknown-wasi"
 @extern_global = external global [0 x i8], align 1
 @main.alignedGlobal = hidden global [4 x i32] zeroinitializer, align 32
 @main.alignedGlobal16 = hidden global [4 x i32] zeroinitializer, align 16
[email protected] = appending global [2 x ptr] [ptr @extern_func, ptr @exportedFunctionInSection]
[email protected] = appending global [3 x ptr] [ptr @extern_func, ptr @exportedFunctionInSection, ptr @exported]
 @main.globalInSection = hidden global i32 0, section ".special_global_section", align 4
 @undefinedGlobalNotInSection = external global i32, align 4
 @main.multipleGlobalPragmas = hidden global i32 0, section ".global_section", align 1024
@@ -62,13 +62,23 @@ entry:
 
 declare void @main.declaredImport() #7
 
+declare void @imported() #8
+
+; Function Attrs: nounwind
+define void @exported() #9 {
+entry:
+  ret void
+}
+
 declare void @main.undefinedFunctionNotInSection(ptr) #1
 
 attributes #0 = { allockind("alloc,zeroed") allocsize(0) "alloc-family"="runtime.alloc" "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
 attributes #1 = { "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
 attributes #2 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
-attributes #3 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-name"="extern_func" "wasm-import-module"="env" "wasm-import-name"="extern_func" }
+attributes #3 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-name"="extern_func" }
 attributes #4 = { inlinehint nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
 attributes #5 = { noinline nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" }
-attributes #6 = { noinline nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-name"="exportedFunctionInSection" "wasm-import-module"="env" "wasm-import-name"="exportedFunctionInSection" }
+attributes #6 = { noinline nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-name"="exportedFunctionInSection" }
 attributes #7 = { "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-import-module"="modulename" "wasm-import-name"="import1" }
+attributes #8 = { "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-import-module"="foobar" "wasm-import-name"="imported" }
+attributes #9 = { nounwind "target-features"="+bulk-memory,+nontrapping-fptoint,+sign-ext" "wasm-export-name"="exported" }