Adding Augmentor

Author: grantnelson-wf

build/build.go

@@ -139,18 +139,16 @@ type overrideInfo struct {
 }
 
 // pkgOverrideInfo is the collection of overrides still needed for a package.
-//
-// Even after the overrides are applied to the files of a package and the
-// overrides map is empty, this will be kept to indicate that a package has
-// already been augmented so that the additional native information is not
-// re-applied when there is a file from this package parsed after all the
-// overrides are empty.
 type pkgOverrideInfo struct {
 	// overrides is a map of identifier to overrideInfo to override
 	// individual named structs, interfaces, functions, and methods.
-	// As identifiers are used, they will be removed from this map.
 	overrides map[string]overrideInfo
 
+	// overlayFiles are the files from the natives that still haven't been
+	// appended to a file from the package, typically the first file.
+	overlayFiles []*ast.File
+
+	// jsFiles are the additional JS files that are part of the natives.
 	jsFiles []JSFile
 }
 
@@ -165,40 +163,61 @@ type pkgOverrideInfo struct {
 //
 // The first file from a package will have any additional methods and
 // information from the natives injected into the AST. All files from a package
-// will be augmented by the overrides. After augmentation of a whole package,
-// the overrides should be empty, but might not be if the natives contain
-// overrides for methods that no longer exist.
+// will be augmented by the overrides.
 type Augmentor struct {
-	xctx XContext
-
 	// packages is a map of package import path to the package's override.
 	// This is used to keep track of the overrides for a package and indicate
 	// that additional files from the natives have already been applied.
 	packages map[string]*pkgOverrideInfo
 }
 
-func (aug *Augmentor) Augment(fileSet *token.FileSet, filename string, src *ast.File) error {
-	pkgName := src.Name.Name
-	importPath := pkgName // TODO: Determine unique import path for the package.
+func (aug *Augmentor) getPackageOverrides(xctx XContext, pkg *PackageData, fileSet *token.FileSet) *pkgOverrideInfo {
+	importPath := pkg.ImportPath
+	if pkgAug, ok := aug.packages[importPath]; ok {
+		return pkgAug
+	}
 
-	pkgAug, ok := aug.packages[importPath]
-	if !ok {
-		if aug.packages == nil {
-			aug.packages = map[string]*pkgOverrideInfo{}
-		}
-		jsFiles, overlayFiles := parseOverlayFiles(xctx, pkg, fileSet)
+	jsFiles, overlayFiles := parseOverlayFiles(xctx, pkg, fileSet)
 
-		overrides := make(map[string]overrideInfo)
-		for _, file := range overlayFiles {
-			augmentOverlayFile(file, overrides)
-		}
-		delete(overrides, `init`)
+	overrides := make(map[string]overrideInfo)
+	for _, file := range overlayFiles {
+		augmentOverlayFile(file, overrides)
+	}
+	delete(overrides, `init`)
+
+	pkgAug := &pkgOverrideInfo{
+		overrides:    overrides,
+		overlayFiles: overlayFiles,
+		jsFiles:      jsFiles,
+	}
+
+	if aug.packages == nil {
+		aug.packages = map[string]*pkgOverrideInfo{}
+	}
+	aug.packages[importPath] = pkgAug
+	return pkgAug
+}
+
+func (aug *Augmentor) Augment(xctx XContext, pkg *PackageData, fileSet *token.FileSet, file *ast.File) error {
+	pkgAug := aug.getPackageOverrides(xctx, pkg, fileSet)
+
+	augmentOriginalImports(pkg.ImportPath, file)
 
-		pkgAug = &pkgOverrideInfo{
-			overrides: overrides,
-			jsFiles:   jsFiles,
+	if len(pkgAug.overrides) > 0 {
+		augmentOriginalFile(file, pkgAug.overrides)
+	}
+
+	if len(pkgAug.overlayFiles) > 0 {
+		// Append the overlay files to the first file of the package.
+		// This is to ensure that the package is augmented with all the
+		// additional methods and information from the natives.
+		err := astutil.ConcatenateFiles(file, pkgAug.overlayFiles...)
+		if err != nil {
+			panic(fmt.Errorf("failed to concatenate overlay files onto %q: %w", fileSet.Position(file.Package).Filename, err))
 		}
-		aug.packages[importPath] = pkgAug
+		pkgAug.overlayFiles = nil
+
+		// TODO: Finish
 	}
 
 	return nil
@@ -384,8 +403,8 @@ func augmentOverlayFile(file *ast.File, overrides map[string]overrideInfo) {
 		}
 	}
 	if anyChange {
-		finalizeRemovals(file)
-		pruneImports(file)
+		astutil.FinalizeRemovals(file)
+		astutil.PruneImports(file)
 	}
 }
 
@@ -496,167 +515,9 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
 		}
 	}
 	if anyChange {
-		finalizeRemovals(file)
-		pruneImports(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
+		astutil.FinalizeRemovals(file)
+		astutil.PruneImports(file)
 	}
-	return true
-}
-
-// pruneImports will remove any unused imports from the file.
-//
-// 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.
-//
-// 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.
-		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 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 {
-				delete(unused, id.Name)
-			}
-		}
-		return len(unused) > 0
-	})
-	if len(unused) == 0 {
-		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) {
-			// since the import is otherwise unused set the name to blank.
-			in.Name = ast.NewIdent(`_`)
-			delete(unused, name)
-		}
-	}
-	if len(unused) == 0 {
-		return
-	}
-
-	// Remove all unused import specifications
-	isUnusedSpec := map[*ast.ImportSpec]bool{}
-	for _, index := range unused {
-		isUnusedSpec[file.Imports[index]] = true
-	}
-	for _, decl := range file.Decls {
-		if d, ok := decl.(*ast.GenDecl); ok {
-			for i, spec := range d.Specs {
-				if other, ok := spec.(*ast.ImportSpec); ok && isUnusedSpec[other] {
-					d.Specs[i] = nil
-				}
-			}
-		}
-	}
-
-	// Remove the unused import copies in the file
-	for _, index := range unused {
-		file.Imports[index] = nil
-	}
-
-	finalizeRemovals(file)
-}
-
-// finalizeRemovals fully removes any declaration, specification, imports
-// 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 {
-		switch d := decl.(type) {
-		case nil:
-			fileChanged = true
-		case *ast.GenDecl:
-			declChanged := false
-			for j, spec := range d.Specs {
-				switch s := spec.(type) {
-				case nil:
-					declChanged = true
-				case *ast.ValueSpec:
-					specChanged := false
-					for _, name := range s.Names {
-						if name == nil {
-							specChanged = true
-							break
-						}
-					}
-					if specChanged {
-						s.Names = astutil.Squeeze(s.Names)
-						s.Values = astutil.Squeeze(s.Values)
-						if len(s.Names) == 0 {
-							declChanged = true
-							d.Specs[j] = nil
-						}
-					}
-				}
-			}
-			if declChanged {
-				d.Specs = astutil.Squeeze(d.Specs)
-				if len(d.Specs) == 0 {
-					fileChanged = true
-					file.Decls[i] = nil
-				}
-			}
-		}
-	}
-	if fileChanged {
-		file.Decls = astutil.Squeeze(file.Decls)
-	}
-
-	file.Imports = astutil.Squeeze(file.Imports)
-
-	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.

build/build_test.go

@@ -7,6 +7,7 @@ import (
 	"strconv"
 	"testing"
 
+	"github.com/gopherjs/gopherjs/compiler/astutil"
 	"github.com/gopherjs/gopherjs/internal/srctesting"
 	"github.com/shurcooL/go/importgraphutil"
 )
@@ -423,7 +424,7 @@ func TestOverlayAugmentation(t *testing.T) {
 
 			overrides := map[string]overrideInfo{}
 			augmentOverlayFile(fileSrc, overrides)
-			pruneImports(fileSrc)
+			astutil.PruneImports(fileSrc)
 
 			got := srctesting.Format(t, f.FileSet, fileSrc)
 
@@ -724,7 +725,7 @@ func TestOriginalAugmentation(t *testing.T) {
 
 			augmentOriginalImports(importPath, fileSrc)
 			augmentOriginalFile(fileSrc, test.info)
-			pruneImports(fileSrc)
+			astutil.PruneImports(fileSrc)
 
 			got := srctesting.Format(t, f.FileSet, fileSrc)