gopherjs · nevkontakte · Jan 19, 2024 · Jan 16, 2024 · Jan 19, 2024
diff --git a/build/build.go b/build/build.go
@@ -130,6 +130,11 @@ type overrideInfo struct {
 	// If the method is defined in the overlays and therefore has its
 	// own overrides, this will be ignored.
 	purgeMethods bool
+
+	// overrideSignature is the function definition given in the overlays
+	// that should be used to replace the signature in the originals.
+	// Only receivers, type parameters, parameters, and results will be used.
+	overrideSignature *ast.FuncDecl
 }
 
 // parseAndAugment parses and returns all .go files of given pkg.
@@ -145,11 +150,16 @@ type overrideInfo struct {
 //   - For function identifiers that exist in the original and the overrides
 //     and have the directive `gopherjs:keep-original`, the original identifier
 //     in the AST gets prefixed by `_gopherjs_original_`.
-//   - For identifiers that exist in the original and the overrides and have
+//   - For identifiers that exist in the original and the overrides, and have
 //     the directive `gopherjs:purge`, both the original and override are
 //     removed. This is for completely removing something which is currently
 //     invalid for GopherJS. For any purged types any methods with that type as
 //     the receiver are also removed.
+//   - For function identifiers that exist in the original and the overrides,
+//     and have the directive `gopherjs:override-signature`, the overridden
+//     function is removed and the original function's signature is changed
+//     to match the overridden function signature. This allows the receiver,
+//     type parameters, parameter, and return values to be modified as needed.
 //   - Otherwise for identifiers that exist in the original and the overrides,
 //     the original is removed.
 //   - New identifiers that don't exist in original package get added.
@@ -270,9 +280,14 @@ func augmentOverlayFile(file *ast.File, overrides map[string]overrideInfo) {
 		switch d := decl.(type) {
 		case *ast.FuncDecl:
 			k := astutil.FuncKey(d)
-			overrides[k] = overrideInfo{
+			oi := overrideInfo{
 				keepOriginal: astutil.KeepOriginal(d),
 			}
+			if astutil.OverrideSignature(d) {
+				oi.overrideSignature = d
+				purgeDecl = true
+			}
+			overrides[k] = oi
 		case *ast.GenDecl:
 			for j, spec := range d.Specs {
 				purgeSpec := purgeDecl || astutil.Purge(spec)
@@ -323,11 +338,21 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
 		switch d := decl.(type) {
 		case *ast.FuncDecl:
 			if info, ok := overrides[astutil.FuncKey(d)]; ok {
+				removeFunc := true
 				if info.keepOriginal {
 					// Allow overridden function calls
 					// The standard library implementation of foo() becomes _gopherjs_original_foo()
 					d.Name.Name = "_gopherjs_original_" + d.Name.Name
-				} else {
+					removeFunc = false
+				}
+				if overSig := info.overrideSignature; overSig != nil {
+					d.Recv = overSig.Recv
+					d.Type.TypeParams = overSig.Type.TypeParams
+					d.Type.Params = overSig.Type.Params
+					d.Type.Results = overSig.Type.Results
+					removeFunc = false
+				}
+				if removeFunc {
 					file.Decls[i] = nil
 				}
 			} else if recvKey := astutil.FuncReceiverKey(d); len(recvKey) > 0 {
@@ -391,21 +416,43 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
 	finalizeRemovals(file)
 }
 
+// isOnlyImports determines if this file is empty except for imports.
+func isOnlyImports(file *ast.File) bool {
+	for _, decl := range file.Decls {
+		if gen, ok := decl.(*ast.GenDecl); ok && gen.Tok == token.IMPORT {
+			continue
+		}
+
+		// The decl was either a FuncDecl or a non-import GenDecl.
+		return false
+	}
+	return true
+}
+
 // pruneImports will remove any unused imports from the file.
 //
-// This will not remove any dot (`.`) or blank (`_`) imports.
+// This will not remove any dot (`.`) or blank (`_`) imports, unless
+// there are no declarations or directives meaning that all the imports
+// should be cleared.
 // If the removal of code causes an import to be removed, the init's from that
 // import may not be run anymore. If we still need to run an init for an import
 // which is no longer used, add it to the overlay as a blank (`_`) import.
 func pruneImports(file *ast.File) {
+	if isOnlyImports(file) && !astutil.HasDirectivePrefix(file, `//go:linkname `) {
+		// The file is empty, remove all imports including any `.` or `_` imports.
+		file.Imports = nil
+		file.Decls = nil
+		return
+	}
+
 	unused := make(map[string]int, len(file.Imports))
 	for i, in := range file.Imports {
 		if name := astutil.ImportName(in); len(name) > 0 {
 			unused[name] = i
 		}
 	}
 
-	// Remove "unused import" for any import which is used.
+	// Remove "unused imports" for any import which is used.
 	ast.Inspect(file, func(n ast.Node) bool {
 		if sel, ok := n.(*ast.SelectorExpr); ok {
 			if id, ok := sel.X.(*ast.Ident); ok && id.Obj == nil {
@@ -418,6 +465,24 @@ func pruneImports(file *ast.File) {
 		return
 	}
 
+	// Remove "unused imports" for any import used for a directive.
+	directiveImports := map[string]string{
+		`unsafe`: `//go:linkname `,
+		`embed`:  `//go:embed `,
+	}
+	for name, index := range unused {
+		in := file.Imports[index]
+		path, _ := strconv.Unquote(in.Path.Value)
+		directivePrefix, hasPath := directiveImports[path]
+		if hasPath && astutil.HasDirectivePrefix(file, directivePrefix) {
+			delete(unused, name)
+			if len(unused) == 0 {
+				return
+			}
+			break
+		}
+	}
+
 	// Remove all unused import specifications
 	isUnusedSpec := map[*ast.ImportSpec]bool{}
 	for _, index := range unused {
@@ -442,9 +507,8 @@ func pruneImports(file *ast.File) {
 }
 
 // finalizeRemovals fully removes any declaration, specification, imports
-// that have been set to nil. This will also remove the file's top-level
-// comment group to remove any unassociated comments, including the comments
-// from removed code.
+// that have been set to nil. This will also remove any unassociated comment
+// groups, including the comments from removed code.
 func finalizeRemovals(file *ast.File) {
 	fileChanged := false
 	for i, decl := range file.Decls {
@@ -487,8 +551,18 @@ func finalizeRemovals(file *ast.File) {
 	if fileChanged {
 		file.Decls = astutil.Squeeze(file.Decls)
 	}
+
 	file.Imports = astutil.Squeeze(file.Imports)
-	file.Comments = nil
+
+	file.Comments = nil // clear this first so ast.Inspect doesn't walk it.
+	remComments := []*ast.CommentGroup{}
+	ast.Inspect(file, func(n ast.Node) bool {
+		if cg, ok := n.(*ast.CommentGroup); ok {
+			remComments = append(remComments, cg)
+		}
+		return true
+	})
+	file.Comments = remComments
 }
 
 // Options controls build process behavior.

diff --git a/build/build_test.go b/build/build_test.go
@@ -371,6 +371,36 @@ func TestOverlayAugmentation(t *testing.T) {
 				`Sort`:    {},
 				`Equal`:   {},
 			},
+		}, {
+			desc: `remove unsafe and embed if not needed`,
+			src: `import "unsafe"
+				import "embed"
+
+				//gopherjs:purge
+				var eFile embed.FS
+
+				//gopherjs:purge
+				func SwapPointer(addr *unsafe.Pointer, new unsafe.Pointer) (old unsafe.Pointer)`,
+			want: ``,
+			expInfo: map[string]overrideInfo{
+				`SwapPointer`: {},
+				`eFile`:       {},
+			},
+		}, {
+			desc: `keep unsafe and embed for directives`,
+			src: `import "unsafe"
+				import "embed"
+
+				//go:embed hello.txt
+				var eFile embed.FS
+
+				//go:linkname runtimeNano runtime.nanotime
+				func runtimeNano() int64`,
+			noCodeChange: true,
+			expInfo: map[string]overrideInfo{
+				`eFile`:       {},
+				`runtimeNano`: {},
+			},
 		},
 	}
 
@@ -539,14 +569,18 @@ func TestOriginalAugmentation(t *testing.T) {
 				`Equal`:   {},
 			},
 			src: `import "cmp"
+
+				// keeps the isOnlyImports from skipping what is being tested.
+				func foo() {}
 
 				type Pointer[T any] struct {}
 
 				func Sort[S ~[]E, E cmp.Ordered](x S) {}
 
 				// overlay had stub "func Equal() {}"
 				func Equal[S ~[]E, E any](s1, s2 S) bool {}`,
-			want: ``,
+			want: `// keeps the isOnlyImports from skipping what is being tested.
+				func foo() {}`,
 		}, {
 			desc: `purge generics`,
 			info: map[string]overrideInfo{
@@ -556,6 +590,9 @@ func TestOriginalAugmentation(t *testing.T) {
 			},
 			src: `import "cmp"
 
+				// keeps the isOnlyImports from skipping what is being tested.
+				func foo() {}
+
 				type Pointer[T any] struct {}
 				func (x *Pointer[T]) Load() *T {}
 				func (x *Pointer[T]) Store(val *T) {}
@@ -564,12 +601,78 @@ func TestOriginalAugmentation(t *testing.T) {
 
 				// overlay had stub "func Equal() {}"
 				func Equal[S ~[]E, E any](s1, s2 S) bool {}`,
-			want: ``,
+			want: `// keeps the isOnlyImports from skipping what is being tested.
+				func foo() {}`,
 		}, {
 			desc: `prune an unused import`,
 			info: map[string]overrideInfo{},
-			src:  `import foo "some/other/bar"`,
+			src: `import foo "some/other/bar"
+
+				// keeps the isOnlyImports from skipping what is being tested.
+				func foo() {}`,
+			want: `// keeps the isOnlyImports from skipping what is being tested.
+				func foo() {}`,
+		}, {
+			desc: `override signature of function`,
+			info: map[string]overrideInfo{
+				`Foo`: {
+					overrideSignature: srctesting.ParseFuncDecl(t,
+						`package whatever
+						func Foo(a, b any) (any, bool) {}`),
+				},
+			},
+			src: `func Foo[T comparable](a, b T) (T, bool) {
+					if a == b {
+						return a, true
+					}
+					return b, false
+				}`,
+			want: `func Foo(a, b any) (any, bool) {
+				if a == b {
+					return a, true
+				}
+				return b, false
+			}`,
+		}, {
+			desc: `override signature of method`,
+			info: map[string]overrideInfo{
+				`Foo.Bar`: {
+					overrideSignature: srctesting.ParseFuncDecl(t,
+						`package whatever
+						func (r *Foo) Bar(a, b any) (any, bool) {}`),
+				},
+			},
+			src: `func (r *Foo[T]) Bar(a, b T) (T, bool) {
+					if r.isSame(a, b) {
+						return a, true
+					}
+					return b, false
+				}`,
+			want: `func (r *Foo) Bar(a, b any) (any, bool) {
+					if r.isSame(a, b) {
+						return a, true
+					}
+					return b, false
+				}`,
+		}, {
+			desc: `empty file removes all imports`,
+			info: map[string]overrideInfo{
+				`foo`: {},
+			},
+			src: `import . "math/rand"
+				func foo() int {
+					return Int()
+				}`,
 			want: ``,
+		}, {
+			desc: `empty file with directive`,
+			info: map[string]overrideInfo{
+				`foo`: {},
+			},
+			src: `//go:linkname foo bar
+				import _ "unsafe"`,
+			want: `//go:linkname foo bar
+				import _ "unsafe"`,
 		},
 	}
 

diff --git a/compiler/astutil/astutil.go b/compiler/astutil/astutil.go
@@ -9,6 +9,7 @@ import (
 	"reflect"
 	"regexp"
 	"strconv"
+	"strings"
 )
 
 func RemoveParens(e ast.Expr) ast.Expr {
@@ -148,6 +149,24 @@ func Purge(d ast.Node) bool {
 	return hasDirective(d, `purge`)
 }
 
+// OverrideSignature returns true if gopherjs:override-signature directive is
+// present on a function.
+//
+// `//gopherjs:override-signature` is a GopherJS-specific directive, which can
+// be applied in native overlays and will instruct the augmentation logic to
+// replace the original function signature which has the same FuncKey with the
+// signature defined in the native overlays.
+// This directive can be used to remove generics from a function signature or
+// to replace a receiver of a function with another one. The given native
+// overlay function will be removed, so no method body is needed in the overlay.
+//
+// The new signature may not contain types which require a new import since
+// the imports will not be automatically added when needed, only removed.
+// Use a type alias in the overlay to deal manage imports.
+func OverrideSignature(d *ast.FuncDecl) bool {
+	return hasDirective(d, `override-signature`)
+}
+
 // directiveMatcher is a regex which matches a GopherJS directive
 // and finds the directive action.
 var directiveMatcher = regexp.MustCompile(`^\/(?:\/|\*)gopherjs:([\w-]+)`)
@@ -179,6 +198,19 @@ func hasDirective(node ast.Node, directiveAction string) bool {
 	return foundDirective
 }
 
+// HasDirectivePrefix determines if any line in the given file
+// has the given directive prefix in it.
+func HasDirectivePrefix(file *ast.File, prefix string) bool {
+	for _, cg := range file.Comments {
+		for _, c := range cg.List {
+			if strings.HasPrefix(c.Text, prefix) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // FindLoopStmt tries to find the loop statement among the AST nodes in the
 // |stack| that corresponds to the break/continue statement represented by
 // branch.