gopls/internal/analysis/modernize/forvar: provide fix for second loop var

Fixes golang/go#74849

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

Author: skewb1k

gopls/internal/analysis/modernize/forvar.go

@@ -30,66 +30,50 @@ import (
 // is declared implicitly before executing the post statement and initialized to the
 // value of the previous iteration's variable at that moment.")
 func forvar(pass *analysis.Pass) {
-	info := pass.TypesInfo
-
 	inspect := pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)
-	for curFile := range filesUsing(inspect, info, "go1.22") {
+	for curFile := range filesUsing(inspect, pass.TypesInfo, "go1.22") {
+		astFile := curFile.Node().(*ast.File)
 		for curLoop := range curFile.Preorder((*ast.RangeStmt)(nil)) {
-			// in a range loop. Is the first statement var := var?
-			// if so, is var one of the range vars, and is it defined
-			// in the for statement?
-			// If so, decide how much to delete.
 			loop := curLoop.Node().(*ast.RangeStmt)
 			if loop.Tok != token.DEFINE {
 				continue
 			}
-			v, stmt := loopVarRedecl(loop.Body)
-			if v == nil {
-				continue // index is not redeclared
-			}
-			if (loop.Key == nil || !equalSyntax(loop.Key, v)) &&
-				(loop.Value == nil || !equalSyntax(loop.Value, v)) {
-				continue
+			isLoopVarRedecl := func(assign *ast.AssignStmt) bool {
+				for i, lhs := range assign.Lhs {
+					if !(equalSyntax(lhs, assign.Rhs[i]) &&
+						(equalSyntax(lhs, loop.Key) || equalSyntax(lhs, loop.Value))) {
+						return false
+					}
+				}
+				return true
 			}
-			astFile := curFile.Node().(*ast.File)
-			edits := analysisinternal.DeleteStmt(pass.Fset, astFile, stmt, bug.Reportf)
-			if len(edits) == 0 {
-				bug.Reportf("forvar failed to delete statement")
-				continue
-			}
-			remove := edits[0]
-			diag := analysis.Diagnostic{
-				Pos:      remove.Pos,
-				End:      remove.End,
-				Category: "forvar",
-				Message:  "copying variable is unneeded",
-				SuggestedFixes: []analysis.SuggestedFix{{
-					Message:   "Remove unneeded redeclaration",
-					TextEdits: []analysis.TextEdit{remove},
-				}},
+			// Have: for k, v := range x { stmts }
+			//
+			// Delete the prefix of stmts that are
+			// of the form k := k; v := v; k, v := k, v; v, k := v, k.
+			for _, stmt := range loop.Body.List {
+				if assign, ok := stmt.(*ast.AssignStmt); ok &&
+					assign.Tok == token.DEFINE &&
+					len(assign.Lhs) == len(assign.Rhs) &&
+					isLoopVarRedecl(assign) {
+
+					edits := analysisinternal.DeleteStmt(pass.Fset, astFile, stmt, bug.Reportf)
+					if len(edits) > 0 {
+						pass.Report(analysis.Diagnostic{
+							Pos:      stmt.Pos(),
+							End:      stmt.End(),
+							Category: "forvar",
+							Message:  "copying variable is unneeded",
+							SuggestedFixes: []analysis.SuggestedFix{{
+								Message:   "Remove unneeded redeclaration",
+								TextEdits: edits,
+							}},
+						})
+					}
+				} else {
+					break // stop at first other statement
+				}
 			}
-			pass.Report(diag)
 		}
 	}
 }
-
-// if the first statement is var := var, return var and the stmt
-func loopVarRedecl(body *ast.BlockStmt) (*ast.Ident, *ast.AssignStmt) {
-	if len(body.List) < 1 {
-		return nil, nil
-	}
-	stmt, ok := body.List[0].(*ast.AssignStmt)
-	if !ok || !isSimpleAssign(stmt) || stmt.Tok != token.DEFINE {
-		return nil, nil
-	}
-	if _, ok := stmt.Lhs[0].(*ast.Ident); !ok {
-		return nil, nil
-	}
-	if _, ok := stmt.Rhs[0].(*ast.Ident); !ok {
-		return nil, nil
-	}
-	if stmt.Lhs[0].(*ast.Ident).Name == stmt.Rhs[0].(*ast.Ident).Name {
-		return stmt.Lhs[0].(*ast.Ident), stmt
-	}
-	return nil, nil
-}

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

@@ -12,12 +12,24 @@ func _(m map[int]int, s []int) {
 	}
 	for k, v := range m {
 		k := k // want "copying variable is unneeded"
-		v := v // nope: report only the first redeclaration
+		v := v // want "copying variable is unneeded"
 		go f(k)
 		go f(v)
 	}
-	for _, v := range m {
+	for k, v := range m {
 		v := v // want "copying variable is unneeded"
+		k := k // want "copying variable is unneeded"
+		go f(k)
+		go f(v)
+	}
+	for k, v := range m {
+		k, v := k, v // want "copying variable is unneeded"
+		go f(k)
+		go f(v)
+	}
+	for k, v := range m {
+		v, k := v, k // want "copying variable is unneeded"
+		go f(k)
 		go f(v)
 	}
 	for i := range s {
@@ -53,10 +65,25 @@ func _(m map[int]int, s []int) {
 		v := i
 		go f(v)
 	}
-	for i := range s {
+	for k, v := range m { // nope, LHS and RHS differ
+		v, k := k, v
+		go f(k)
+		go f(v)
+	}
+	for k, v := range m { // nope, not a simple redecl
+		k, v, x := k, v, 1
+		go f(k)
+		go f(v)
+		go f(x)
+	}
+	for i := range s { // nope, not a simple redecl
 		i := (i)
 		go f(i)
 	}
+	for i := range s { // nope, not a simple redecl
+		i := i + 1
+		go f(i)
+	}
 }
 
 func f(n int) {}

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

@@ -12,12 +12,24 @@ func _(m map[int]int, s []int) {
 	}
 	for k, v := range m {
 		// want "copying variable is unneeded"
-		v := v // nope: report only the first redeclaration
+		// want "copying variable is unneeded"
+		go f(k)
+		go f(v)
+	}
+	for k, v := range m {
+		// want "copying variable is unneeded"
+		// want "copying variable is unneeded"
 		go f(k)
 		go f(v)
 	}
-	for _, v := range m {
+	for k, v := range m {
 		// want "copying variable is unneeded"
+		go f(k)
+		go f(v)
+	}
+	for k, v := range m {
+		// want "copying variable is unneeded"
+		go f(k)
 		go f(v)
 	}
 	for i := range s {
@@ -53,10 +65,25 @@ func _(m map[int]int, s []int) {
 		v := i
 		go f(v)
 	}
-	for i := range s {
+	for k, v := range m { // nope, LHS and RHS differ
+		v, k := k, v
+		go f(k)
+		go f(v)
+	}
+	for k, v := range m { // nope, not a simple redecl
+		k, v, x := k, v, 1
+		go f(k)
+		go f(v)
+		go f(x)
+	}
+	for i := range s { // nope, not a simple redecl
 		i := (i)
 		go f(i)
 	}
+	for i := range s { // nope, not a simple redecl
+		i := i + 1
+		go f(i)
+	}
 }
 
 func f(n int) {}