Connecting up the augmentor

build/augmentor.go

@@ -65,7 +65,7 @@ type Augmentor struct {
 	packages map[string]*pkgOverrideInfo
 }
 
-func (aug *Augmentor) Augment(xctx XContext, pkg *PackageData, fileSet *token.FileSet, file *ast.File) error {
+func (aug *Augmentor) Augment(xctx XContext, pkg *PackageData, fileSet *token.FileSet, file *ast.File) {
 	pkgAug := aug.getPackageOverrides(xctx, pkg, fileSet)
 
 	augmentOriginalImports(pkg.ImportPath, file)
@@ -84,8 +84,6 @@ func (aug *Augmentor) Augment(xctx XContext, pkg *PackageData, fileSet *token.Fi
 		}
 		pkgAug.overlayFiles = nil
 	}
-
-	return nil
 }
 
 // getPackageOverrides looks up an already loaded package override

build/build.go

@@ -149,24 +149,13 @@ func parseAndAugment(xctx XContext, pkg *PackageData, fileSet *token.FileSet) ([
 		return nil, nil, err
 	}
 
-	jsFiles, overlayFiles := parseOverlayFiles(xctx, pkg, fileSet)
-	overrides := make(map[string]overrideInfo)
-	for _, file := range overlayFiles {
-		augmentOverlayFile(file, overrides)
-	}
-	delete(overrides, `init`)
-
+	aug := &Augmentor{}
 	for _, file := range originalFiles {
-		augmentOriginalImports(pkg.ImportPath, file)
-	}
-
-	if len(overrides) > 0 {
-		for _, file := range originalFiles {
-			augmentOriginalFile(file, overrides)
-		}
+		aug.Augment(xctx, pkg, fileSet, file)
 	}
 
-	return append(overlayFiles, originalFiles...), jsFiles, nil
+	jsFiles := aug.packages[pkg.ImportPath].jsFiles
+	return originalFiles, jsFiles, nil
 }
 
 // parserOriginalFiles loads and parses the original files to augment.

compiler/astutil/astutil.go

@@ -406,6 +406,9 @@ func squeeze[E ast.Node, S ~[]E](s S) S {
 	return s[:dest]
 }
 
+// updateFileComments rebuilds the file comments by reading the comments
+// off of the nodes in the file. Any comments that are not associated with
+// a node will be lost.
 func updateFileComments(file *ast.File) {
 	file.Comments = nil // clear this first so ast.Inspect doesn't walk it.
 	remComments := []*ast.CommentGroup{}