Adding Augmentor

Author: grantnelson-wf

build/build.go

@@ -138,16 +138,68 @@ type overrideInfo struct {
 	overrideSignature *ast.FuncDecl
 }
 
+// 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
+
+	jsFiles []JSFile
 }
 
+// Augmentor is an on-the-fly package augmentor.
+//
+// When a file from a package is being parsed, the Augmentor will augment
+// the AST with the changes loaded from the native overrides.
+// The augmentor will hold onto the override information for additional files
+// that come from the same package. This is designed to be used with
+// `x/tools/go/packages.Load` as a middleware in the parse file step via
+// `Config.ParseFile`.
+//
+// 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.
 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(fset *token.FileSet, filename string, src *ast.File) error {
+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.
+
+	pkgAug, ok := aug.packages[importPath]
+	if !ok {
+		if aug.packages == nil {
+			aug.packages = map[string]*pkgOverrideInfo{}
+		}
+		jsFiles, overlayFiles := parseOverlayFiles(xctx, pkg, fileSet)
+
+		overrides := make(map[string]overrideInfo)
+		for _, file := range overlayFiles {
+			augmentOverlayFile(file, overrides)
+		}
+		delete(overrides, `init`)
+
+		pkgAug = &pkgOverrideInfo{
+			overrides: overrides,
+			jsFiles:   jsFiles,
+		}
+		aug.packages[importPath] = pkgAug
+	}
 
 	return nil
 }
@@ -180,14 +232,12 @@ func (aug *Augmentor) Augment(fset *token.FileSet, filename string, src *ast.Fil
 //     the original is removed.
 //   - New identifiers that don't exist in original package get added.
 func parseAndAugment(xctx XContext, pkg *PackageData, fileSet *token.FileSet) ([]*ast.File, []JSFile, error) {
-	isTest := pkg.IsTest
-	jsFiles, overlayFiles := parseOverlayFiles(xctx, pkg, isTest, fileSet)
-
 	originalFiles, err := parserOriginalFiles(pkg, fileSet)
 	if err != nil {
 		return nil, nil, err
 	}
 
+	jsFiles, overlayFiles := parseOverlayFiles(xctx, pkg, fileSet)
 	overrides := make(map[string]overrideInfo)
 	for _, file := range overlayFiles {
 		augmentOverlayFile(file, overrides)
@@ -209,7 +259,7 @@ func parseAndAugment(xctx XContext, pkg *PackageData, fileSet *token.FileSet) ([
 
 // parseOverlayFiles loads and parses overlay files
 // to augment the original files with.
-func parseOverlayFiles(xctx XContext, pkg *PackageData, isTest bool, fileSet *token.FileSet) ([]JSFile, []*ast.File) {
+func parseOverlayFiles(xctx XContext, pkg *PackageData, fileSet *token.FileSet) ([]JSFile, []*ast.File) {
 	importPath := pkg.ImportPath
 	isXTest := strings.HasSuffix(importPath, "_test")
 	if isXTest {
@@ -225,7 +275,7 @@ func parseOverlayFiles(xctx XContext, pkg *PackageData, isTest bool, fileSet *to
 	jsFiles := nativesPkg.JSFiles
 	var files []*ast.File
 	names := nativesPkg.GoFiles
-	if isTest {
+	if pkg.IsTest {
 		names = append(names, nativesPkg.TestGoFiles...)
 	}
 	if isXTest {