Limit augmentation and prune imports to natives that need it

gopherjs · nevkontakte · Feb 24, 2024 · Nov 12, 2023 · Nov 14, 2023 · Nov 14, 2023
commit 02aea33b76248f0a35eb7c7044c880fdadf8ee6c
diff --git a/build/build.go b/build/build.go
@@ -180,8 +180,12 @@ func parseAndAugment(xctx XContext, pkg *PackageData, isTest bool, fileSet *toke

 	for _, file := range originalFiles {
 		augmentOriginalImports(pkg.ImportPath, file)
-		augmentOriginalFile(file, overrides)
-		pruneImports(file)
+	}
+
+	if len(overrides) > 0 {
+		for _, file := range originalFiles {
+			augmentOriginalFile(file, overrides)
+		}
 	}

 	return append(overlayFiles, originalFiles...), jsFiles, nil
@@ -275,6 +279,7 @@ func parserOriginalFiles(pkg *PackageData, fileSet *token.FileSet) ([]*ast.File,
 // an overlay file AST to collect information such as compiler directives
 // and perform any initial augmentation needed to the overlay.
 func augmentOverlayFile(file *ast.File, overrides map[string]overrideInfo) {
+	anyChange := false
 	for i, decl := range file.Decls {
 		purgeDecl := astutil.Purge(decl)
 		switch d := decl.(type) {
@@ -302,15 +307,20 @@ func augmentOverlayFile(file *ast.File, overrides map[string]overrideInfo) {
 					}
 				}
 				if purgeSpec {
+					anyChange = true
 					d.Specs[j] = nil
 				}
 			}
 		}
 		if purgeDecl {
+			anyChange = true
 			file.Decls[i] = nil
 		}
 	}
-	finalizeRemovals(file)
+	if anyChange {
+		finalizeRemovals(file)
+		pruneImports(file)
+	}
 }

 // augmentOriginalImports is the part of parseAndAugment that processes
@@ -334,10 +344,12 @@ func augmentOriginalImports(importPath string, file *ast.File) {
 // original file AST to augment the source code using the overrides from
 // the overlay files.
 func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
+	anyChange := false
 	for i, decl := range file.Decls {
 		switch d := decl.(type) {
 		case *ast.FuncDecl:
 			if info, ok := overrides[astutil.FuncKey(d)]; ok {
+				anyChange = true
 				removeFunc := true
 				if info.keepOriginal {
 					// Allow overridden function calls
@@ -358,6 +370,7 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
 			} else if recvKey := astutil.FuncReceiverKey(d); len(recvKey) > 0 {
 				// check if the receiver has been purged, if so, remove the method too.
 				if info, ok := overrides[recvKey]; ok && info.purgeMethods {
+					anyChange = true
 					file.Decls[i] = nil
 				}
 			}
@@ -366,6 +379,7 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
 				switch s := spec.(type) {
 				case *ast.TypeSpec:
 					if _, ok := overrides[s.Name.Name]; ok {
+						anyChange = true
 						d.Specs[j] = nil
 					}
 				case *ast.ValueSpec:
@@ -378,6 +392,7 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
 						// to be run, add the call into the overlay.
 						for k, name := range s.Names {
 							if _, ok := overrides[name.Name]; ok {
+								anyChange = true
 								s.Names[k] = nil
 								s.Values[k] = nil
 							}
@@ -405,6 +420,7 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
 								}
 							}
 							if removeSpec {
+								anyChange = true
 								d.Specs[j] = nil
 							}
 						}
@@ -413,7 +429,10 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
 			}
 		}
 	}
-	finalizeRemovals(file)
+	if anyChange {
+		finalizeRemovals(file)
+		pruneImports(file)
+	}
 }

 // isOnlyImports determines if this file is empty except for imports.
@@ -437,6 +456,14 @@ func isOnlyImports(file *ast.File) bool {
 // 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.
+//
+// This uses the given name or guesses at the name using the import path,
+// meaning this doesn't work for packages which have a different package name
+// from the path, including those paths which are versioned
+// (e.g. `github.com/foo/bar/v2` where the package name is `bar`)
+// or if the import is defined using a relative path (e.g. `./..`).
+// Those cases don't exist in the native for Go, so we should only run
+// this pruning when we have native overlays, but not for unknown packages.
 func pruneImports(file *ast.File) {
 	if isOnlyImports(file) && !astutil.HasDirectivePrefix(file, `//go:linkname `) {
 		// The file is empty, remove all imports including any `.` or `_` imports.
@@ -478,12 +505,11 @@ func pruneImports(file *ast.File) {
 			// since the import is otherwise unused set the name to blank.
 			in.Name = ast.NewIdent(`_`)
 			delete(unused, name)
-			if len(unused) == 0 {
-				return
-			}
-			break
 		}
 	}
+	if len(unused) == 0 {
+		return
+	}

 	// Remove all unused import specifications
 	isUnusedSpec := map[*ast.ImportSpec]bool{}

diff --git a/build/build_test.go b/build/build_test.go
@@ -680,6 +680,39 @@ func TestOriginalAugmentation(t *testing.T) {
 				import _ "unsafe"`,
 			want: `//go:linkname foo bar
 				import _ "unsafe"`,
+		}, {
+			desc: `multiple imports for directives`,
+			info: map[string]overrideInfo{
+				`A`: {},
+				`C`: {},
+			},
+			src: `import "unsafe"
+				import "embed"
+
+				//go:embed hello.txt
+				var A embed.FS
+
+				//go:embed goodbye.txt
+				var B string
+
+				var C unsafe.Pointer
+
+				// override Now with hardcoded time for testing
+				//go:linkname timeNow time.Now
+				func timeNow() time.Time {
+					return time.Date(2012, 8, 6, 0, 0, 0, 0, time.UTC)
+				}`,
+			want: `import _ "unsafe"
+				import _ "embed"
+
+				//go:embed goodbye.txt
+				var B string
+
+				// override Now with hardcoded time for testing
+				//go:linkname timeNow time.Now
+				func timeNow() time.Time {
+					return time.Date(2012, 8, 6, 0, 0, 0, 0, time.UTC)
+				}`,
 		},
 	}