From 02aad23f00004ebcc4ce5f5b27a726787e147e48 Mon Sep 17 00:00:00 2001
From: Prashant Varanasi <prashant@prashantv.com>
Date: Sat, 17 Feb 2024 22:37:14 -0800
Subject: [PATCH] cmd/errtrace: Fix blank import handling (#92)

With toolexec mode, a blank import of errtrace is more likely to opt-in
to rewriting. However, this blank import can't be used to call errtrace
functions, so we need to track both:

 * Whether errtrace is imported (used to opt-in toolexec rewriting)
 * Whether errtrace needs to be imported for any Wrap calls.

Since this change modifies the rewriting logic, we bump the errtrace
toolexec version manually.
---
 cmd/errtrace/main.go                            | 16 +++++++++++-----
 cmd/errtrace/testdata/golden/imported_blank.go  | 17 +++++++++++++++++
 .../testdata/golden/imported_blank.go.golden    | 17 +++++++++++++++++
 cmd/errtrace/testdata/toolexec-test/errtrace.go |  4 ----
 cmd/errtrace/testdata/toolexec-test/main.go     |  1 +
 cmd/errtrace/toolexec.go                        |  2 +-
 6 files changed, 47 insertions(+), 10 deletions(-)
 create mode 100644 cmd/errtrace/testdata/golden/imported_blank.go
 create mode 100644 cmd/errtrace/testdata/golden/imported_blank.go.golden
 delete mode 100644 cmd/errtrace/testdata/toolexec-test/errtrace.go

diff --git a/cmd/errtrace/main.go b/cmd/errtrace/main.go
index 9f778f5..53adfe9 100644
--- a/cmd/errtrace/main.go
+++ b/cmd/errtrace/main.go
@@ -398,7 +398,7 @@ type parsedFile struct {
 	file *ast.File
 
 	errtracePkg     string
-	importsErrtrace bool
+	importsErrtrace bool // includes blank imports
 	inserts         []insert
 	unusedOptouts   []int // list of line numbers
 }
@@ -410,21 +410,27 @@ func (cmd *mainCmd) parseFile(filename string, src []byte) (parsedFile, error) {
 		return parsedFile{}, errtrace.Wrap(err)
 	}
 
-	errtracePkg := "errtrace" // name to use for errtrace package
-	var importsErrtrace bool  // whether the file imports errtrace already
+	errtracePkg := "errtrace"  // name to use for errtrace package
+	var importsErrtrace bool   // whether there's any errtrace import, including blank imports
+	needErrtraceImport := true // whether to add a new import.
 	for _, imp := range f.Imports {
 		if imp.Path.Value == `"braces.dev/errtrace"` {
 			importsErrtrace = true
 			if imp.Name != nil {
+				if imp.Name.Name == "_" {
+					// Can't use a blank import, keep processing imports.
+					continue
+				}
 				// If the file already imports errtrace,
 				// we'll want to use the name it's imported under.
 				errtracePkg = imp.Name.Name
 			}
+			needErrtraceImport = false
 			break
 		}
 	}
 
-	if !importsErrtrace {
+	if needErrtraceImport {
 		// If the file doesn't import errtrace already,
 		// do a quick check to find an unused identifier name.
 		idents := make(map[string]struct{})
@@ -475,7 +481,7 @@ func (cmd *mainCmd) parseFile(filename string, src []byte) (parsedFile, error) {
 	// If errtrace isn't imported, but at least one insert was made,
 	// we'll need to import errtrace.
 	// Add an import declaration to the file.
-	if !importsErrtrace && len(inserts) > 0 {
+	if needErrtraceImport && len(inserts) > 0 {
 		// We want to insert the import after the last existing import.
 		// If the last import is part of a group, we'll make it part of the group.
 		//
diff --git a/cmd/errtrace/testdata/golden/imported_blank.go b/cmd/errtrace/testdata/golden/imported_blank.go
new file mode 100644
index 0000000..d466ed7
--- /dev/null
+++ b/cmd/errtrace/testdata/golden/imported_blank.go
@@ -0,0 +1,17 @@
+//go:build ignore
+
+package foo
+
+import (
+	"strconv"
+
+	_ "braces.dev/errtrace"
+)
+
+func Unwrapped(s string) (int, error) {
+	i, err := strconv.Atoi(s)
+	if err != nil {
+		return 0, err
+	}
+	return i + 42, nil
+}
diff --git a/cmd/errtrace/testdata/golden/imported_blank.go.golden b/cmd/errtrace/testdata/golden/imported_blank.go.golden
new file mode 100644
index 0000000..fec3280
--- /dev/null
+++ b/cmd/errtrace/testdata/golden/imported_blank.go.golden
@@ -0,0 +1,17 @@
+//go:build ignore
+
+package foo
+
+import (
+	"strconv"
+
+	_ "braces.dev/errtrace"; "braces.dev/errtrace"
+)
+
+func Unwrapped(s string) (int, error) {
+	i, err := strconv.Atoi(s)
+	if err != nil {
+		return 0, errtrace.Wrap(err)
+	}
+	return i + 42, nil
+}
diff --git a/cmd/errtrace/testdata/toolexec-test/errtrace.go b/cmd/errtrace/testdata/toolexec-test/errtrace.go
deleted file mode 100644
index 3df1b34..0000000
--- a/cmd/errtrace/testdata/toolexec-test/errtrace.go
+++ /dev/null
@@ -1,4 +0,0 @@
-package main
-
-// Opt-in to errtrace wrapping with toolexec.
-import _ "braces.dev/errtrace"
diff --git a/cmd/errtrace/testdata/toolexec-test/main.go b/cmd/errtrace/testdata/toolexec-test/main.go
index 056c7cf..fa384ee 100644
--- a/cmd/errtrace/testdata/toolexec-test/main.go
+++ b/cmd/errtrace/testdata/toolexec-test/main.go
@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 
+	_ "braces.dev/errtrace" // Opt-in to errtrace wrapping with toolexec.
 	"braces.dev/errtrace/cmd/errtrace/testdata/toolexec-test/p1"
 )
 
diff --git a/cmd/errtrace/toolexec.go b/cmd/errtrace/toolexec.go
index 3197f8e..c05a9dd 100644
--- a/cmd/errtrace/toolexec.go
+++ b/cmd/errtrace/toolexec.go
@@ -53,7 +53,7 @@ func (cmd *mainCmd) toolExecVersion(args []string) int {
 	}
 
 	// TODO: This version number should change whenever the rewriting changes.
-	fmt.Fprintf(cmd.Stdout, "%s-errtrace0\n", strings.TrimSpace(stdout.String()))
+	fmt.Fprintf(cmd.Stdout, "%s-errtrace1\n", strings.TrimSpace(stdout.String()))
 	return 0
 }