gopls/internal/analysis/modernize: preserves comments in minmax

This CL changes the original deletion (from after rh0
to the end of the if stmt) into deletion from the start of assignment to
the end of if stmt), add all comments between them and create a new
assigment as-is.

This change preserves all comments inside if stmt and the comments
after the line of assignment and before if stmt, causing comments B,C,D
to be preserved and put on the top of min/max function call after fix.

- source:

	lhs0 = rhs0 // A
	// B
	if rhs0 < b {
	    // C
	    lhs0 = b
	    // D
	}
- fixed:
	// A
	// B
	// C
	// D
	lhs0 = max(rhs0,b)

Fixes golang/go#72727

Change-Id: I7c193711aac5834ebb0d5e8ae22c26ae7990c34f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/656655
Auto-Submit: Alan Donovan <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Robert Findley <[email protected]>

Author: xieyuschen

gopls/internal/analysis/modernize/minmax.go

@@ -9,6 +9,7 @@ import (
 	"go/ast"
 	"go/token"
 	"go/types"
+	"strings"
 
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/inspect"
@@ -32,7 +33,7 @@ func minmax(pass *analysis.Pass) {
 
 	// check is called for all statements of this form:
 	//   if a < b { lhs = rhs }
-	check := func(curIfStmt cursor.Cursor, compare *ast.BinaryExpr) {
+	check := func(file *ast.File, curIfStmt cursor.Cursor, compare *ast.BinaryExpr) {
 		var (
 			ifStmt  = curIfStmt.Node().(*ast.IfStmt)
 			tassign = ifStmt.Body.List[0].(*ast.AssignStmt)
@@ -44,6 +45,14 @@ func minmax(pass *analysis.Pass) {
 			sign    = isInequality(compare.Op)
 		)
 
+		allComments := func(file *ast.File, start, end token.Pos) string {
+			var buf strings.Builder
+			for co := range analysisinternal.Comments(file, start, end) {
+				_, _ = fmt.Fprintf(&buf, "%s\n", co.Text)
+			}
+			return buf.String()
+		}
+
 		if fblock, ok := ifStmt.Else.(*ast.BlockStmt); ok && isAssignBlock(fblock) {
 			fassign := fblock.List[0].(*ast.AssignStmt)
 
@@ -85,7 +94,8 @@ func minmax(pass *analysis.Pass) {
 							// Replace IfStmt with lhs = min(a, b).
 							Pos: ifStmt.Pos(),
 							End: ifStmt.End(),
-							NewText: fmt.Appendf(nil, "%s = %s(%s, %s)",
+							NewText: fmt.Appendf(nil, "%s%s = %s(%s, %s)",
+								allComments(file, ifStmt.Pos(), ifStmt.End()),
 								analysisinternal.Format(pass.Fset, lhs),
 								sym,
 								analysisinternal.Format(pass.Fset, a),
@@ -144,10 +154,13 @@ func minmax(pass *analysis.Pass) {
 					SuggestedFixes: []analysis.SuggestedFix{{
 						Message: fmt.Sprintf("Replace if/else with %s", sym),
 						TextEdits: []analysis.TextEdit{{
-							// Replace rhs0 and IfStmt with min(a, b)
-							Pos: rhs0.Pos(),
+							Pos: fassign.Pos(),
 							End: ifStmt.End(),
-							NewText: fmt.Appendf(nil, "%s(%s, %s)",
+							// Replace "x := a; if ... {}" with "x = min(...)", preserving comments.
+							NewText: fmt.Appendf(nil, "%s %s %s %s(%s, %s)",
+								allComments(file, fassign.Pos(), ifStmt.End()),
+								analysisinternal.Format(pass.Fset, lhs),
+								fassign.Tok.String(),
 								sym,
 								analysisinternal.Format(pass.Fset, a),
 								analysisinternal.Format(pass.Fset, b)),
@@ -161,16 +174,16 @@ func minmax(pass *analysis.Pass) {
 	// Find all "if a < b { lhs = rhs }" statements.
 	inspect := pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)
 	for curFile := range filesUsing(inspect, pass.TypesInfo, "go1.21") {
+		astFile := curFile.Node().(*ast.File)
 		for curIfStmt := range curFile.Preorder((*ast.IfStmt)(nil)) {
 			ifStmt := curIfStmt.Node().(*ast.IfStmt)
-
 			if compare, ok := ifStmt.Cond.(*ast.BinaryExpr); ok &&
 				ifStmt.Init == nil &&
 				isInequality(compare.Op) != 0 &&
 				isAssignBlock(ifStmt.Body) {
 
 				// Have: if a < b { lhs = rhs }
-				check(curIfStmt, compare)
+				check(astFile, curIfStmt, compare)
 			}
 		}
 	}

gopls/internal/analysis/modernize/testdata/src/minmax/minmax.go

@@ -1,9 +1,12 @@
 package minmax
 
 func ifmin(a, b int) {
-	x := a
+	x := a // A
+	// B
 	if a < b { // want "if statement can be modernized using max"
-		x = b
+		// C
+		x = b // D
+		// E
 	}
 	print(x)
 }
@@ -33,20 +36,30 @@ func ifmaxvariant(a, b int) {
 }
 
 func ifelsemin(a, b int) {
-	var x int
+	var x int // A
+	// B
 	if a <= b { // want "if/else statement can be modernized using min"
-		x = a
+		// C
+		x = a // D
+		// E
 	} else {
-		x = b
+		// F
+		x = b // G
+		// H
 	}
 	print(x)
 }
 
 func ifelsemax(a, b int) {
-	var x int
+	// A
+	var x int // B
+	// C
 	if a >= b { // want "if/else statement can be modernized using max"
-		x = a
+		// D
+		x = a // E
+		// F
 	} else {
+		// G
 		x = b
 	}
 	print(x)
@@ -115,3 +128,11 @@ func nopeHasElseBlock(x int) int {
 	}
 	return y
 }
+
+func fix72727(a, b int) {
+	o := a - 42
+	// some important comment. DO NOT REMOVE.
+	if o < b { // want "if statement can be modernized using max"
+		o = b
+	}
+}

gopls/internal/analysis/modernize/testdata/src/minmax/minmax.go.golden

@@ -1,33 +1,57 @@
 package minmax
 
 func ifmin(a, b int) {
+	// A
+	// B
+	// want "if statement can be modernized using max"
+	// C
+	// D
+	// E
 	x := max(a, b)
 	print(x)
 }
 
 func ifmax(a, b int) {
+	// want "if statement can be modernized using min"
 	x := min(a, b)
 	print(x)
 }
 
 func ifminvariant(a, b int) {
+	// want "if statement can be modernized using min"
 	x := min(a, b)
 	print(x)
 }
 
 func ifmaxvariant(a, b int) {
+	// want "if statement can be modernized using min"
 	x := min(a, b)
 	print(x)
 }
 
 func ifelsemin(a, b int) {
-	var x int
+	var x int // A
+	// B
+	// want "if/else statement can be modernized using min"
+	// C
+	// D
+	// E
+	// F
+	// G
+	// H
 	x = min(a, b)
 	print(x)
 }
 
 func ifelsemax(a, b int) {
-	var x int
+	// A
+	var x int // B
+	// C
+	// want "if/else statement can be modernized using max"
+	// D
+	// E
+	// F
+	// G
 	x = max(a, b)
 	print(x)
 }
@@ -55,6 +79,7 @@ func nopeIfStmtHasInitStmt() {
 // Regression test for a bug: fix was "y := max(x, y)".
 func oops() {
 	x := 1
+	// want "if statement can be modernized using max"
 	y := max(x, 2)
 	print(y)
 }
@@ -92,3 +117,9 @@ func nopeHasElseBlock(x int) int {
 	}
 	return y
 }
+
+func fix72727(a, b int) {
+	// some important comment. DO NOT REMOVE.
+	// want "if statement can be modernized using max"
+	o := max(a-42, b)
+}

internal/analysisinternal/analysis.go

@@ -15,6 +15,7 @@ import (
 	"go/scanner"
 	"go/token"
 	"go/types"
+	"iter"
 	pathpkg "path"
 	"slices"
 	"strings"
@@ -608,3 +609,24 @@ Outer:
 	}
 	return []analysis.TextEdit{edit}
 }
+
+// Comments returns an iterator over the comments overlapping the specified interval.
+func Comments(file *ast.File, start, end token.Pos) iter.Seq[*ast.Comment] {
+	// TODO(adonovan): optimize use binary O(log n) instead of linear O(n) search.
+	return func(yield func(*ast.Comment) bool) {
+		for _, cg := range file.Comments {
+			for _, co := range cg.List {
+				if co.Pos() > end {
+					return
+				}
+				if co.End() < start {
+					continue
+				}
+
+				if !yield(co) {
+					return
+				}
+			}
+		}
+	}
+}

internal/analysisinternal/analysis_test.go

@@ -8,6 +8,7 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
+	"slices"
 	"testing"
 
 	"golang.org/x/tools/go/ast/inspector"
@@ -253,3 +254,47 @@ func TestDeleteStmt(t *testing.T) {
 
 	}
 }
+
+func TestComments(t *testing.T) {
+	src := `
+package main
+
+// A
+func fn() { }`
+	var fset token.FileSet
+	f, err := parser.ParseFile(&fset, "", []byte(src), parser.ParseComments|parser.AllErrors)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	commentA := f.Comments[0].List[0]
+	commentAMidPos := (commentA.Pos() + commentA.End()) / 2
+
+	want := []*ast.Comment{commentA}
+	testCases := []struct {
+		name       string
+		start, end token.Pos
+		want       []*ast.Comment
+	}{
+		{name: "comment totally overlaps with given interval", start: f.Pos(), end: f.End(), want: want},
+		{name: "interval from file start to mid of comment A", start: f.Pos(), end: commentAMidPos, want: want},
+		{name: "interval from mid of comment A to file end", start: commentAMidPos, end: commentA.End(), want: want},
+		{name: "interval from start of comment A to mid of comment A", start: commentA.Pos(), end: commentAMidPos, want: want},
+		{name: "interval from mid of comment A to comment A end", start: commentAMidPos, end: commentA.End(), want: want},
+		{name: "interval at the start of comment A", start: commentA.Pos(), end: commentA.Pos(), want: want},
+		{name: "interval at the end of comment A", start: commentA.End(), end: commentA.End(), want: want},
+		{name: "interval from file start to the front of comment A start", start: f.Pos(), end: commentA.Pos() - 1, want: nil},
+		{name: "interval from the position after end of comment A to file end", start: commentA.End() + 1, end: f.End(), want: nil},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			var got []*ast.Comment
+			for co := range Comments(f, tc.start, tc.end) {
+				got = append(got, co)
+			}
+			if !slices.Equal(got, tc.want) {
+				t.Errorf("%s: got %v, want %v", tc.name, got, tc.want)
+			}
+		})
+	}
+}