Updated to use Inspect and srctesting

grantnelson-wf · grantnelson-wf · commit 7758c0d142d8 · 2023-12-19T16:46:10.000-07:00
diff --git a/build/build.go b/build/build.go
@@ -406,14 +406,14 @@ func pruneImports(file *ast.File) {
 	}
 
 	// Remove "unused import" for any import which is used.
-	ast.Walk(astutil.NewCallbackVisitor(func(n ast.Node) bool {
+	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
-	}), file)
+	})
 	if len(unused) == 0 {
 		return
 	}
diff --git a/build/build_test.go b/build/build_test.go
@@ -1,10 +1,8 @@
 package build
 
 import (
-	"bytes"
 	"fmt"
 	gobuild "go/build"
-	"go/printer"
 	"go/token"
 	"strconv"
 	"testing"
@@ -390,16 +388,11 @@ func TestOverlayAugmentation(t *testing.T) {
 			augmentOverlayFile(fileSrc, overrides)
 			pruneImports(fileSrc)
 
-			buf := &bytes.Buffer{}
-			_ = printer.Fprint(buf, fsetSrc, fileSrc)
-			got := buf.String()
+			got := srctesting.Format(t, fsetSrc, fileSrc)
 
 			fsetWant := token.NewFileSet()
 			fileWant := srctesting.Parse(t, fsetWant, pkgName+test.want)
-
-			buf.Reset()
-			_ = printer.Fprint(buf, fsetWant, fileWant)
-			want := buf.String()
+			want := srctesting.Format(t, fsetWant, fileWant)
 
 			if got != want {
 				t.Errorf("augmentOverlayFile and pruneImports got unexpected code:\n"+
@@ -536,7 +529,7 @@ func TestOriginalAugmentation(t *testing.T) {
 
 				func NewFoo(bar int) *Foo { return &Foo{bar: bar} }`,
 			// NewFoo is not removed automatically since
-			// only functions with Foo as a receiver is removed.
+			// only functions with Foo as a receiver are removed.
 			want: `func NewFoo(bar int) *Foo { return &Foo{bar: bar} }`,
 		}, {
 			desc: `remove generics`,
@@ -591,16 +584,11 @@ func TestOriginalAugmentation(t *testing.T) {
 			augmentOriginalFile(fileSrc, test.info)
 			pruneImports(fileSrc)
 
-			buf := &bytes.Buffer{}
-			_ = printer.Fprint(buf, fsetSrc, fileSrc)
-			got := buf.String()
+			got := srctesting.Format(t, fsetSrc, fileSrc)
 
 			fsetWant := token.NewFileSet()
 			fileWant := srctesting.Parse(t, fsetWant, pkgName+test.want)
-
-			buf.Reset()
-			_ = printer.Fprint(buf, fsetWant, fileWant)
-			want := buf.String()
+			want := srctesting.Format(t, fsetWant, fileWant)
 
 			if got != want {
 				t.Errorf("augmentOriginalImports, augmentOriginalFile, and pruneImports got unexpected code:\n"+
diff --git a/compiler/astutil/astutil.go b/compiler/astutil/astutil.go
@@ -251,19 +251,3 @@ func Squeeze[E ast.Node, S ~[]E](s S) S {
 	}
 	return s[:dest]
 }
-
-type CallbackVisitor struct {
-	predicate func(node ast.Node) bool
-}
-
-func NewCallbackVisitor(predicate func(node ast.Node) bool) *CallbackVisitor {
-	return &CallbackVisitor{predicate: predicate}
-}
-
-func (v *CallbackVisitor) Visit(node ast.Node) ast.Visitor {
-	if v.predicate != nil && v.predicate(node) {
-		return v
-	}
-	v.predicate = nil
-	return nil
-}

Original file line number	Diff line number	Diff line change
`@@ -406,14 +406,14 @@ func pruneImports(file *ast.File) {`
`406`	`406`	`}`
`407`	`407`
`408`	`408`	`// Remove "unused import" for any import which is used.`
`409`		`- ast.Walk(astutil.NewCallbackVisitor(func(n ast.Node) bool {`
	`409`	`+ ast.Inspect(file, func(n ast.Node) bool {`
`410`	`410`	`if sel, ok := n.(*ast.SelectorExpr); ok {`
`411`	`411`	`if id, ok := sel.X.(*ast.Ident); ok && id.Obj == nil {`
`412`	`412`	`delete(unused, id.Name)`
`413`	`413`	`}`
`414`	`414`	`}`
`415`	`415`	`return len(unused) > 0`
`416`		`- }), file)`
	`416`	`+ })`
`417`	`417`	`if len(unused) == 0 {`
`418`	`418`	`return`
`419`	`419`	`}`