gopls/internal/analysis/simplifyrange: more precise fix

This CL reduces the size of the fix offered by simplifyrange,
which makes the cursor jump less. It's also simpler, and
handles a missing case.

Change-Id: I8dbff96158e442b2073e86694b61ea1f0b1ea704
Reviewed-on: https://go-review.googlesource.com/c/tools/+/649355
Reviewed-by: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: adonovan

gopls/internal/analysis/simplifyrange/simplifyrange.go

@@ -5,10 +5,8 @@
 package simplifyrange
 
 import (
-	"bytes"
 	_ "embed"
 	"go/ast"
-	"go/printer"
 	"go/token"
 
 	"golang.org/x/tools/go/analysis"
@@ -42,73 +40,48 @@ func run(pass *analysis.Pass) (interface{}, error) {
 		(*ast.RangeStmt)(nil),
 	}
 	inspect.Preorder(nodeFilter, func(n ast.Node) {
-		if _, ok := generated[pass.Fset.File(n.Pos())]; ok {
-			return // skip checking if it's generated code
-		}
+		rng := n.(*ast.RangeStmt)
 
-		var copy *ast.RangeStmt // shallow-copy the AST before modifying
-		{
-			x := *n.(*ast.RangeStmt)
-			copy = &x
-		}
-		end := newlineIndex(pass.Fset, copy)
+		kblank := isBlank(rng.Key)
+		vblank := isBlank(rng.Value)
+		var start, end token.Pos
+		switch {
+		case kblank && (rng.Value == nil || vblank):
+			// for _    = range x {}
+			// for _, _ = range x {}
+			//     ^^^^^^^
+			start, end = rng.Key.Pos(), rng.Range
 
-		// Range statements of the form: for i, _ := range x {}
-		var old ast.Expr
-		if isBlank(copy.Value) {
-			old = copy.Value
-			copy.Value = nil
-		}
-		// Range statements of the form: for _ := range x {}
-		if isBlank(copy.Key) && copy.Value == nil {
-			old = copy.Key
-			copy.Key = nil
+		case vblank:
+			// for k, _ := range x {}
+			//      ^^^
+			start, end = rng.Key.End(), rng.Value.End()
+
+		default:
+			return
 		}
-		// Return early if neither if condition is met.
-		if old == nil {
+
+		if generated[pass.Fset.File(n.Pos())] {
 			return
 		}
+
 		pass.Report(analysis.Diagnostic{
-			Pos:            old.Pos(),
-			End:            old.End(),
-			Message:        "simplify range expression",
-			SuggestedFixes: suggestedFixes(pass.Fset, copy, end),
+			Pos:     start,
+			End:     end,
+			Message: "simplify range expression",
+			SuggestedFixes: []analysis.SuggestedFix{{
+				Message: "Remove empty value",
+				TextEdits: []analysis.TextEdit{{
+					Pos: start,
+					End: end,
+				}},
+			}},
 		})
 	})
 	return nil, nil
 }
 
-func suggestedFixes(fset *token.FileSet, rng *ast.RangeStmt, end token.Pos) []analysis.SuggestedFix {
-	var b bytes.Buffer
-	printer.Fprint(&b, fset, rng)
-	stmt := b.Bytes()
-	index := bytes.Index(stmt, []byte("\n"))
-	// If there is a new line character, then don't replace the body.
-	if index != -1 {
-		stmt = stmt[:index]
-	}
-	return []analysis.SuggestedFix{{
-		Message: "Remove empty value",
-		TextEdits: []analysis.TextEdit{{
-			Pos:     rng.Pos(),
-			End:     end,
-			NewText: stmt[:index],
-		}},
-	}}
-}
-
-func newlineIndex(fset *token.FileSet, rng *ast.RangeStmt) token.Pos {
-	var b bytes.Buffer
-	printer.Fprint(&b, fset, rng)
-	contents := b.Bytes()
-	index := bytes.Index(contents, []byte("\n"))
-	if index == -1 {
-		return rng.End()
-	}
-	return rng.Pos() + token.Pos(index)
-}
-
-func isBlank(x ast.Expr) bool {
-	ident, ok := x.(*ast.Ident)
-	return ok && ident.Name == "_"
+func isBlank(e ast.Expr) bool {
+	id, ok := e.(*ast.Ident)
+	return ok && id.Name == "_"
 }

gopls/internal/analysis/simplifyrange/simplifyrange_test.go

@@ -5,18 +5,15 @@
 package simplifyrange_test
 
 import (
-	"go/build"
-	"slices"
 	"testing"
 
 	"golang.org/x/tools/go/analysis/analysistest"
 	"golang.org/x/tools/gopls/internal/analysis/simplifyrange"
 )
 
 func Test(t *testing.T) {
-	testdata := analysistest.TestData()
-	analysistest.RunWithSuggestedFixes(t, testdata, simplifyrange.Analyzer, "a", "generatedcode")
-	if slices.Contains(build.Default.ReleaseTags, "go1.23") { // uses iter.Seq
-		analysistest.RunWithSuggestedFixes(t, testdata, simplifyrange.Analyzer, "rangeoverfunc")
-	}
+	analysistest.RunWithSuggestedFixes(t, analysistest.TestData(), simplifyrange.Analyzer,
+		"a",
+		"generatedcode",
+		"rangeoverfunc")
 }

gopls/internal/analysis/simplifyrange/testdata/src/a/a.go

@@ -13,4 +13,11 @@ func m() {
 	}
 	for _ = range maps { // want "simplify range expression"
 	}
+	for _, _ = range maps { // want "simplify range expression"
+	}
+	for _, v := range maps { // nope
+		println(v)
+	}
+	for range maps { // nope
+	}
 }

gopls/internal/analysis/simplifyrange/testdata/src/a/a.go.golden

@@ -13,4 +13,11 @@ func m() {
 	}
 	for range maps { // want "simplify range expression"
 	}
+	for range maps { // want "simplify range expression"
+	}
+	for _, v := range maps { // nope
+		println(v)
+	}
+	for range maps { // nope
+	}
 }