chore: Address review comments

Signed-off-by: Kemal Akkoyun <[email protected]>

Author: kakkoyun

errorlint/analysis.go

@@ -40,7 +40,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 	var lints []analysis.Diagnostic
 	extInfo := newTypesInfoExt(pass)
 	if checkComparison {
-		l := LintErrorComparisons(extInfo)
+		l := LintErrorComparisons(pass.Fset, extInfo)
 		lints = append(lints, l...)
 	}
 	if checkAsserts {

errorlint/analysis_test.go

@@ -58,3 +58,8 @@ func TestIssueFixRegressions(t *testing.T) {
 	analyzer := NewAnalyzer()
 	analysistest.RunWithSuggestedFixes(t, analysistest.TestData(), analyzer, "issues/fix")
 }
+
+func TestASTManipulationFixes(t *testing.T) {
+	analyzer := NewAnalyzer()
+	analysistest.RunWithSuggestedFixes(t, analysistest.TestData(), analyzer, "astmanipulation")
+}

errorlint/lint.go

@@ -90,7 +90,7 @@ func extractComparison(cond ast.Expr) *ast.BinaryExpr {
 
 func buildVarDeclaration(assertion typeAssertion) string {
 	targetTypeStr := exprToString(assertion.targetType)
-	if baseType, found := strings.CutPrefix(targetTypeStr, "*") ; found {
+	if baseType, found := strings.CutPrefix(targetTypeStr, "*"); found {
 		return fmt.Sprintf("%s := &%s{}", assertion.varName, baseType)
 	}
 	return fmt.Sprintf("var %s %s", assertion.varName, targetTypeStr)
@@ -133,7 +133,7 @@ func groupDiagnosticsByIfStmt(lints []analysis.Diagnostic, extInfo *TypesInfoExt
 			continue
 		}
 
-		ifStmt := containingIf(extInfo, node)
+		ifStmt := statementContainsIf(extInfo, node)
 		if ifStmt == nil {
 			otherLints = append(otherLints, lint)
 			continue
@@ -297,7 +297,7 @@ func isFmtErrorfCallExpr(info types.Info, expr ast.Expr) (*ast.CallExpr, bool) {
 	return nil, false
 }
 
-func LintErrorComparisons(info *TypesInfoExt) []analysis.Diagnostic {
+func LintErrorComparisons(fset *token.FileSet, info *TypesInfoExt) []analysis.Diagnostic {
 	var lints []analysis.Diagnostic
 
 	// Check for error comparisons.
@@ -353,13 +353,20 @@ func LintErrorComparisons(info *TypesInfoExt) []analysis.Diagnostic {
 			replacement = "!" + replacement
 		}
 
+		textEdits := []analysis.TextEdit{{
+			Pos:     binExpr.Pos(),
+			End:     binExpr.End(),
+			NewText: []byte(replacement),
+		}}
+
+		if file := findContainingFile(info, binExpr); file != nil && needsErrorsImport(file) {
+			importFix := createImportFix(file)
+			textEdits = append(textEdits, importFix.TextEdits...)
+		}
+
 		diagnostic.SuggestedFixes = []analysis.SuggestedFix{{
-			Message: "Use errors.Is() to compare errors",
-			TextEdits: []analysis.TextEdit{{
-				Pos:     binExpr.Pos(),
-				End:     binExpr.End(),
-				NewText: []byte(replacement),
-			}},
+			Message:   "Use errors.Is() to compare errors",
+			TextEdits: textEdits,
 		}}
 
 		lints = append(lints, diagnostic)
@@ -489,7 +496,49 @@ func LintErrorComparisons(info *TypesInfoExt) []analysis.Diagnostic {
 	return lints
 }
 
-// exprToString converts an expression to its string representation.
+func findContainingFile(info *TypesInfoExt, node ast.Node) *ast.File {
+	current := node
+	for current != nil {
+		if file, ok := current.(*ast.File); ok {
+			return file
+		}
+		current = info.NodeParent[current]
+	}
+	return nil
+}
+
+func needsErrorsImport(file *ast.File) bool {
+	for _, imp := range file.Imports {
+		if imp.Path.Value == `"errors"` {
+			return false
+		}
+	}
+	return true
+}
+
+func createImportFix(file *ast.File) analysis.SuggestedFix {
+	var importPos token.Pos
+	var newText string
+
+	if len(file.Imports) > 0 {
+		lastImport := file.Imports[len(file.Imports)-1]
+		importPos = lastImport.End()
+		newText = "\n\t\"errors\""
+	} else {
+		importPos = file.Name.End()
+		newText = "\n\nimport \"errors\""
+	}
+
+	return analysis.SuggestedFix{
+		Message: "Add missing errors import",
+		TextEdits: []analysis.TextEdit{{
+			Pos:     importPos,
+			End:     importPos,
+			NewText: []byte(newText),
+		}},
+	}
+}
+
 func exprToString(expr ast.Expr) string {
 	switch e := expr.(type) {
 	case *ast.Ident:
@@ -1123,6 +1172,9 @@ func combinedFix(ifStmt *ast.IfStmt, extInfo *TypesInfoExt) *analysis.SuggestedF
 	// Check if this is an else-if statement
 	components.context.isElseIf = isElseIfStatement(ifStmt, extInfo)
 
+	// Handle else-if transformations carefully
+	// Note: else-if transformations change structure but preserve functionality
+
 	// Build the replacement text using the extracted components
 	replacement := buildReplacement(components)
 	if replacement == "" {
@@ -1137,13 +1189,21 @@ func combinedFix(ifStmt *ast.IfStmt, extInfo *TypesInfoExt) *analysis.SuggestedF
 		endPos = ifStmt.Body.End()
 	}
 
+	textEdits := []analysis.TextEdit{{
+		Pos:     ifStmt.Pos(),
+		End:     endPos,
+		NewText: []byte(replacement),
+	}}
+
+	// Add errors import if needed
+	if file := findContainingFile(extInfo, ifStmt); file != nil && needsErrorsImport(file) {
+		importFix := createImportFix(file)
+		textEdits = append(textEdits, importFix.TextEdits...)
+	}
+
 	return &analysis.SuggestedFix{
-		Message: "Use errors.As and errors.Is for error handling",
-		TextEdits: []analysis.TextEdit{{
-			Pos:     ifStmt.Pos(),
-			End:     endPos,
-			NewText: []byte(replacement),
-		}},
+		Message:   "Use errors.As and errors.Is for error handling",
+		TextEdits: textEdits,
 	}
 }
 
@@ -1177,8 +1237,8 @@ type comparison struct {
 	negated bool
 }
 
-// context holds if statement context information
-type context struct {
+// controlFlowContext holds if statement controlFlowContext information
+type controlFlowContext struct {
 	isElseIf  bool
 	bodyStmts []ast.Stmt
 }
@@ -1188,7 +1248,7 @@ type context struct {
 type ifComponents struct {
 	assertion  typeAssertion
 	comparison comparison
-	context    context
+	context    controlFlowContext
 }
 
 // parseIfComponents extracts the components of the if statement pattern
@@ -1224,7 +1284,7 @@ func parseIfComponents(ifStmt *ast.IfStmt) *ifComponents {
 			target:  rightBinExpr.Y,
 			negated: rightBinExpr.Op == token.NEQ,
 		},
-		context: context{
+		context: controlFlowContext{
 			isElseIf:  false,            // Will be set by the calling function if needed
 			bodyStmts: ifStmt.Body.List, // Capture body statements
 		},

errorlint/testdata/src/astmanipulation/comments.go

@@ -0,0 +1,39 @@
+package astmanipulation
+
+import (
+	"os"
+	"syscall"
+)
+
+// Test inline comment preservation
+func InlineComments() {
+	var err error
+
+	if err == syscall.ENOENT { // want "comparing with == will fail on wrapped errors. Use errors.Is to check for a specific error"
+		// This block comment should be preserved
+		println("File not found")
+	}
+}
+
+// Test else-if comment preservation (main PR issue)
+func ElseIfComments() {
+	var err error
+
+	if err != nil {
+		println("error occurred")
+	} else if e, ok := err.(*os.PathError); ok && e.Err == syscall.ESRCH { // want "type assertion and error comparison will fail on wrapped errors. Use errors.As and errors.Is to check for specific errors"
+		// If the process exits while reading its /proc/$PID/maps, the kernel will
+		// return ESRCH. Handle it as if the process did not exist.
+		println("Process not found")
+	}
+}
+
+// Test block comments
+func BlockComments() {
+	var err error
+
+	/* Pre-condition block comment */
+	if e, ok := err.(*os.PathError); ok { // want "type assertion on error will fail on wrapped errors. Use errors.As to check for specific errors"
+		/* Inline block comment */ println("Path error:", e.Path)
+	}
+}

errorlint/testdata/src/astmanipulation/comments.go.golden

@@ -0,0 +1,42 @@
+package astmanipulation
+
+import (
+	"errors"
+	"os"
+	"syscall"
+)
+
+// Test inline comment preservation
+func InlineComments() {
+	var err error
+
+	if errors.Is(err, syscall.ENOENT) { // want "comparing with == will fail on wrapped errors. Use errors.Is to check for a specific error"
+		// This block comment should be preserved
+		println("File not found")
+	}
+}
+
+// Test else-if comment preservation (main PR issue)
+func ElseIfComments() {
+	var err error
+
+	if err != nil {
+		println("error occurred")
+	} else {
+		e := &os.PathError{}
+		if errors.As(err, &e) && errors.Is(e.Err, syscall.ESRCH) {
+			println("Process not found")
+		}
+	}
+}
+
+// Test block comments
+func BlockComments() {
+	var err error
+
+	/* Pre-condition block comment */
+	e := &os.PathError{}
+	if errors.As(err, &e) { // want "type assertion on error will fail on wrapped errors. Use errors.As to check for specific errors"
+		/* Inline block comment */ println("Path error:", e.Path)
+	}
+}

errorlint/testdata/src/astmanipulation/imports.go

@@ -0,0 +1,25 @@
+package astmanipulation
+
+import (
+	"io"
+	"os"
+)
+
+// Test missing errors import with simple comparison
+func MissingImportSimple() {
+	var err error
+	var n int
+
+	if n == 0 || (err != nil && err != io.EOF) { // want "comparing with != will fail on wrapped errors. Use errors.Is to check for a specific error"
+		return
+	}
+}
+
+// Test missing errors import with type assertion
+func MissingImportTypeAssertion() {
+	var err error
+
+	if e, ok := err.(*os.PathError); ok { // want "type assertion on error will fail on wrapped errors. Use errors.As to check for specific errors"
+		println("Path error:", e.Path)
+	}
+}

errorlint/testdata/src/astmanipulation/imports.go.golden

@@ -0,0 +1,27 @@
+package astmanipulation
+
+import (
+	"errors"
+	"io"
+	"os"
+)
+
+// Test missing errors import with simple comparison
+func MissingImportSimple() {
+	var err error
+	var n int
+	
+	if n == 0 || (err != nil && !errors.Is(err, io.EOF)) { // want "comparing with != will fail on wrapped errors. Use errors.Is to check for a specific error"
+		return
+	}
+}
+
+// Test missing errors import with type assertion
+func MissingImportTypeAssertion() {
+	var err error
+	
+	e := &os.PathError{}
+	if errors.As(err, &e) { // want "type assertion on error will fail on wrapped errors. Use errors.As to check for specific errors"
+		println("Path error:", e.Path)
+	}
+}

errorlint/testdata/src/issues/fix/github-100.go.golden

@@ -1,6 +1,7 @@
 package issues
 
 import (
+	"errors"
 	"os"
 	"syscall"
 )