Added more tests

Author: grantnelson-wf

compiler/astutil/astutil.go

@@ -224,14 +224,14 @@ func hasDirective(node ast.Node, directiveAction string) bool {
 // HasDirectivePrefix determines if any line in the given file
 // has the given directive prefix in it.
 func HasDirectivePrefix(file *ast.File, prefix string) bool {
-	for _, cg := range file.Comments {
-		for _, c := range cg.List {
-			if strings.HasPrefix(c.Text, prefix) {
-				return true
-			}
+	foundDirective := false
+	ast.Inspect(file, func(n ast.Node) bool {
+		if c, ok := n.(*ast.Comment); ok && strings.HasPrefix(c.Text, prefix) {
+			foundDirective = true
 		}
-	}
-	return false
+		return !foundDirective
+	})
+	return foundDirective
 }
 
 // FindLoopStmt tries to find the loop statement among the AST nodes in the
@@ -333,7 +333,7 @@ func PruneImports(file *ast.File) {
 		}
 	}
 
-	// Remove "unused imports" for any import which is used.
+	// Remove from "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 {
@@ -346,7 +346,7 @@ func PruneImports(file *ast.File) {
 		return
 	}
 
-	// Remove "unused imports" for any import used for a directive.
+	// Remove from "unused imports" for any import used for a directive.
 	directiveImports := map[string]string{
 		`unsafe`: `//go:linkname `,
 		`embed`:  `//go:embed `,
@@ -409,6 +409,7 @@ func squeeze[E ast.Node, S ~[]E](s S) S {
 // 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.
+// Comments that are floating and tied to a node will be lost.
 func FinalizeRemovals(file *ast.File) {
 	fileChanged := false
 	for i, decl := range file.Decls {

compiler/astutil/astutil_test.go

@@ -589,6 +589,124 @@ func TestSqueezeIdents(t *testing.T) {
 	}
 }
 
+func TestPruneImports(t *testing.T) {
+	tests := []struct {
+		name string
+		src  string
+		want string
+	}{
+		{
+			name: `no imports`,
+			src: `package testpackage
+				func foo() {}`,
+			want: `package testpackage
+				func foo() {}`,
+		}, {
+			name: `keep used imports`,
+			src: `package testpackage
+				import "fmt"
+				func foo() { fmt.Println("foo") }`,
+			want: `package testpackage
+				import "fmt"
+				func foo() { fmt.Println("foo") }`,
+		}, {
+			name: `remove imports that are not used`,
+			src: `package testpackage
+				import "fmt"
+				func foo() { }`,
+			want: `package testpackage
+				func foo() { }`,
+		}, {
+			name: `remove imports that are unused but masked by an object`,
+			src: `package testpackage
+				import "fmt"
+				var fmt = "format"
+				func foo() string { return fmt }`,
+			want: `package testpackage
+				var fmt = "format"
+				func foo() string { return fmt }`,
+		}, {
+			name: `remove imports from empty file`,
+			src: `package testpackage
+				import "fmt"
+				import _ "unsafe"`,
+			want: `package testpackage`,
+		}, {
+			name: `remove imports from empty file except for unsafe when linking`,
+			src: `package testpackage
+				import "fmt"
+				import "embed"
+
+				//go:linkname foo runtime.foo
+				import "unsafe"`,
+			want: `package testpackage
+
+				//go:linkname foo runtime.foo
+				import _ "unsafe"`,
+		}, {
+			name: `keep embed imports when embedding`,
+			src: `package testpackage
+				import "fmt"
+				import "embed"
+				import "unsafe"
+
+				//go:embed "foo.txt"
+				var foo string`,
+			want: `package testpackage
+				import _ "embed"
+
+				//go:embed "foo.txt"
+				var foo string`,
+		}, {
+			name: `keep imports that just needed an underscore`,
+			src: `package testpackage
+				import "embed"
+				//go:linkname foo runtime.foo
+				import "unsafe"
+				//go:embed "foo.txt"
+				var foo string`,
+			want: `package testpackage
+				import _ "embed"
+				//go:linkname foo runtime.foo
+				import _ "unsafe"
+				//go:embed "foo.txt"
+				var foo string`,
+		}, {
+			name: `keep imports without names`,
+			src: `package testpackage
+				import _ "fmt"
+				import "log"
+				import . "math"
+
+				var foo string`,
+			want: `package testpackage
+				import _ "fmt"
+
+				import . "math"
+
+				var foo string`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			st := srctesting.New(t)
+
+			srcFile := st.Parse(`testSrc.go`, test.src)
+			PruneImports(srcFile)
+			got := srctesting.Format(t, st.FileSet, srcFile)
+
+			// parse and format the expected result so that formatting matches
+			wantFile := st.Parse(`testWant.go`, test.want)
+			want := srctesting.Format(t, st.FileSet, wantFile)
+
+			if got != want {
+				t.Errorf("Unexpected resulting AST after PruneImports:\n\tgot:  %q\n\twant: %q", got, want)
+			}
+		})
+	}
+}
+
 func TestFinalizeRemovals(t *testing.T) {
 	tests := []struct {
 		name       string