Skip to content

Commit a533e78

Browse files
authored
Merge pull request #1267 from Workiva/limitAugmentation
[go1.19] Limit augmentation to native files which need it
2 parents c852920 + 5d25518 commit a533e78

File tree

2 files changed

+67
-8
lines changed

2 files changed

+67
-8
lines changed

build/build.go

Lines changed: 34 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,12 @@ func parseAndAugment(xctx XContext, pkg *PackageData, isTest bool, fileSet *toke
180180

181181
for _, file := range originalFiles {
182182
augmentOriginalImports(pkg.ImportPath, file)
183-
augmentOriginalFile(file, overrides)
184-
pruneImports(file)
183+
}
184+
185+
if len(overrides) > 0 {
186+
for _, file := range originalFiles {
187+
augmentOriginalFile(file, overrides)
188+
}
185189
}
186190

187191
return append(overlayFiles, originalFiles...), jsFiles, nil
@@ -275,6 +279,7 @@ func parserOriginalFiles(pkg *PackageData, fileSet *token.FileSet) ([]*ast.File,
275279
// an overlay file AST to collect information such as compiler directives
276280
// and perform any initial augmentation needed to the overlay.
277281
func augmentOverlayFile(file *ast.File, overrides map[string]overrideInfo) {
282+
anyChange := false
278283
for i, decl := range file.Decls {
279284
purgeDecl := astutil.Purge(decl)
280285
switch d := decl.(type) {
@@ -302,15 +307,20 @@ func augmentOverlayFile(file *ast.File, overrides map[string]overrideInfo) {
302307
}
303308
}
304309
if purgeSpec {
310+
anyChange = true
305311
d.Specs[j] = nil
306312
}
307313
}
308314
}
309315
if purgeDecl {
316+
anyChange = true
310317
file.Decls[i] = nil
311318
}
312319
}
313-
finalizeRemovals(file)
320+
if anyChange {
321+
finalizeRemovals(file)
322+
pruneImports(file)
323+
}
314324
}
315325

316326
// augmentOriginalImports is the part of parseAndAugment that processes
@@ -334,10 +344,12 @@ func augmentOriginalImports(importPath string, file *ast.File) {
334344
// original file AST to augment the source code using the overrides from
335345
// the overlay files.
336346
func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
347+
anyChange := false
337348
for i, decl := range file.Decls {
338349
switch d := decl.(type) {
339350
case *ast.FuncDecl:
340351
if info, ok := overrides[astutil.FuncKey(d)]; ok {
352+
anyChange = true
341353
removeFunc := true
342354
if info.keepOriginal {
343355
// Allow overridden function calls
@@ -358,6 +370,7 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
358370
} else if recvKey := astutil.FuncReceiverKey(d); len(recvKey) > 0 {
359371
// check if the receiver has been purged, if so, remove the method too.
360372
if info, ok := overrides[recvKey]; ok && info.purgeMethods {
373+
anyChange = true
361374
file.Decls[i] = nil
362375
}
363376
}
@@ -366,6 +379,7 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
366379
switch s := spec.(type) {
367380
case *ast.TypeSpec:
368381
if _, ok := overrides[s.Name.Name]; ok {
382+
anyChange = true
369383
d.Specs[j] = nil
370384
}
371385
case *ast.ValueSpec:
@@ -378,6 +392,7 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
378392
// to be run, add the call into the overlay.
379393
for k, name := range s.Names {
380394
if _, ok := overrides[name.Name]; ok {
395+
anyChange = true
381396
s.Names[k] = nil
382397
s.Values[k] = nil
383398
}
@@ -405,6 +420,7 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
405420
}
406421
}
407422
if removeSpec {
423+
anyChange = true
408424
d.Specs[j] = nil
409425
}
410426
}
@@ -413,7 +429,10 @@ func augmentOriginalFile(file *ast.File, overrides map[string]overrideInfo) {
413429
}
414430
}
415431
}
416-
finalizeRemovals(file)
432+
if anyChange {
433+
finalizeRemovals(file)
434+
pruneImports(file)
435+
}
417436
}
418437

419438
// isOnlyImports determines if this file is empty except for imports.
@@ -437,6 +456,14 @@ func isOnlyImports(file *ast.File) bool {
437456
// If the removal of code causes an import to be removed, the init's from that
438457
// import may not be run anymore. If we still need to run an init for an import
439458
// which is no longer used, add it to the overlay as a blank (`_`) import.
459+
//
460+
// This uses the given name or guesses at the name using the import path,
461+
// meaning this doesn't work for packages which have a different package name
462+
// from the path, including those paths which are versioned
463+
// (e.g. `github.com/foo/bar/v2` where the package name is `bar`)
464+
// or if the import is defined using a relative path (e.g. `./..`).
465+
// Those cases don't exist in the native for Go, so we should only run
466+
// this pruning when we have native overlays, but not for unknown packages.
440467
func pruneImports(file *ast.File) {
441468
if isOnlyImports(file) && !astutil.HasDirectivePrefix(file, `//go:linkname `) {
442469
// The file is empty, remove all imports including any `.` or `_` imports.
@@ -478,12 +505,11 @@ func pruneImports(file *ast.File) {
478505
// since the import is otherwise unused set the name to blank.
479506
in.Name = ast.NewIdent(`_`)
480507
delete(unused, name)
481-
if len(unused) == 0 {
482-
return
483-
}
484-
break
485508
}
486509
}
510+
if len(unused) == 0 {
511+
return
512+
}
487513

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

build/build_test.go

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -680,6 +680,39 @@ func TestOriginalAugmentation(t *testing.T) {
680680
import _ "unsafe"`,
681681
want: `//go:linkname foo bar
682682
import _ "unsafe"`,
683+
}, {
684+
desc: `multiple imports for directives`,
685+
info: map[string]overrideInfo{
686+
`A`: {},
687+
`C`: {},
688+
},
689+
src: `import "unsafe"
690+
import "embed"
691+
692+
//go:embed hello.txt
693+
var A embed.FS
694+
695+
//go:embed goodbye.txt
696+
var B string
697+
698+
var C unsafe.Pointer
699+
700+
// override Now with hardcoded time for testing
701+
//go:linkname timeNow time.Now
702+
func timeNow() time.Time {
703+
return time.Date(2012, 8, 6, 0, 0, 0, 0, time.UTC)
704+
}`,
705+
want: `import _ "unsafe"
706+
import _ "embed"
707+
708+
//go:embed goodbye.txt
709+
var B string
710+
711+
// override Now with hardcoded time for testing
712+
//go:linkname timeNow time.Now
713+
func timeNow() time.Time {
714+
return time.Date(2012, 8, 6, 0, 0, 0, 0, time.UTC)
715+
}`,
683716
},
684717
}
685718

0 commit comments

Comments
 (0)