gopls/internal/analysis/modernize: b.N -> b.Loop()

This CL adds a modernizer for go1.24's testing.B.Loop.
It replaces for/range loops over b.N by b.Loop() and
deletes any preceding calls to b.{Start,Stop,Reset}Timer.

Updates golang/go#70815

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

Author: adonovan

gopls/internal/analysis/modernize/bloop.go

@@ -0,0 +1,217 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package modernize
+
+import (
+	"fmt"
+	"go/ast"
+	"go/token"
+	"go/types"
+
+	"golang.org/x/tools/go/analysis"
+	"golang.org/x/tools/go/analysis/passes/inspect"
+	"golang.org/x/tools/go/ast/inspector"
+	"golang.org/x/tools/go/types/typeutil"
+	"golang.org/x/tools/internal/astutil/cursor"
+	"golang.org/x/tools/internal/typesinternal"
+	"golang.org/x/tools/internal/versions"
+)
+
+// bloop updates benchmarks that use "for range b.N", replacing it
+// with go1.24's b.Loop() and eliminating any preceding
+// b.{Start,Stop,Reset}Timer calls.
+//
+// Variants:
+//
+//	for i := 0; i < b.N; i++ {}  =>   for b.Loop() {}
+//	for range b.N {}
+func bloop(pass *analysis.Pass) {
+	if !_imports(pass.Pkg, "testing") {
+		return
+	}
+
+	info := pass.TypesInfo
+
+	// edits computes the text edits for a matched for/range loop
+	// at the specified cursor. b is the *testing.B value, and
+	// (start, end) is the portion using b.N to delete.
+	edits := func(cur cursor.Cursor, b ast.Expr, start, end token.Pos) (edits []analysis.TextEdit) {
+		// Within the same function, delete all calls to
+		// b.{Start,Stop,Timer} that precede the loop.
+		filter := []ast.Node{(*ast.ExprStmt)(nil), (*ast.FuncLit)(nil)}
+		fn, _ := enclosingFunc(cur)
+		fn.Inspect(filter, func(cur cursor.Cursor, push bool) (descend bool) {
+			if push {
+				node := cur.Node()
+				if is[*ast.FuncLit](node) {
+					return false // don't descend into FuncLits (e.g. sub-benchmarks)
+				}
+				stmt := node.(*ast.ExprStmt)
+				if stmt.Pos() > start {
+					return false // not preceding: stop
+				}
+				if call, ok := stmt.X.(*ast.CallExpr); ok {
+					fn := typeutil.StaticCallee(info, call)
+					if fn != nil &&
+						(isMethod(fn, "testing", "B", "StopTimer") ||
+							isMethod(fn, "testing", "B", "StartTimer") ||
+							isMethod(fn, "testing", "B", "ResetTimer")) {
+
+						// Delete call statement.
+						// TODO(adonovan): delete following newline, or
+						// up to start of next stmt? (May delete a comment.)
+						edits = append(edits, analysis.TextEdit{
+							Pos: stmt.Pos(),
+							End: stmt.End(),
+						})
+					}
+				}
+			}
+			return true
+		})
+
+		// Replace ...b.N... with b.Loop().
+		return append(edits, analysis.TextEdit{
+			Pos:     start,
+			End:     end,
+			NewText: fmt.Appendf(nil, "%s.Loop()", formatNode(pass.Fset, b)),
+		})
+	}
+
+	// Find all for/range statements.
+	inspect := pass.ResultOf[inspect.Analyzer].(*inspector.Inspector)
+	filter := []ast.Node{
+		(*ast.File)(nil),
+		(*ast.ForStmt)(nil),
+		(*ast.RangeStmt)(nil),
+	}
+	cursor.Root(inspect).Inspect(filter, func(cur cursor.Cursor, push bool) (descend bool) {
+		if push {
+			switch n := cur.Node().(type) {
+			case *ast.File:
+				if versions.Before(info.FileVersions[n], "go1.24") {
+					return false
+				}
+
+			case *ast.ForStmt:
+				// for _; i < b.N; _ {}
+				if cmp, ok := n.Cond.(*ast.BinaryExpr); ok && cmp.Op == token.LSS {
+					if sel, ok := cmp.Y.(*ast.SelectorExpr); ok &&
+						sel.Sel.Name == "N" &&
+						isPtrToNamed(info.TypeOf(sel.X), "testing", "B") {
+
+						delStart, delEnd := n.Cond.Pos(), n.Cond.End()
+
+						// Eliminate variable i if no longer needed:
+						//  for i := 0; i < b.N; i++ {
+						//    ...no references to i...
+						//  }
+						body, _ := cur.LastChild()
+						if assign, ok := n.Init.(*ast.AssignStmt); ok &&
+							assign.Tok == token.DEFINE &&
+							len(assign.Rhs) == 1 &&
+							isZeroLiteral(assign.Rhs[0]) &&
+							is[*ast.IncDecStmt](n.Post) &&
+							n.Post.(*ast.IncDecStmt).Tok == token.INC &&
+							equalSyntax(n.Post.(*ast.IncDecStmt).X, assign.Lhs[0]) &&
+							!uses(info, body, info.Defs[assign.Lhs[0].(*ast.Ident)]) {
+
+							delStart, delEnd = n.Init.Pos(), n.Post.End()
+						}
+
+						pass.Report(analysis.Diagnostic{
+							// Highlight "i < b.N".
+							Pos:      n.Cond.Pos(),
+							End:      n.Cond.End(),
+							Category: "bloop",
+							Message:  "b.N can be modernized using b.Loop()",
+							SuggestedFixes: []analysis.SuggestedFix{{
+								Message:   "Replace b.N with b.Loop()",
+								TextEdits: edits(cur, sel.X, delStart, delEnd),
+							}},
+						})
+					}
+				}
+
+			case *ast.RangeStmt:
+				// for range b.N {} -> for b.Loop() {}
+				//
+				// TODO(adonovan): handle "for i := range b.N".
+				if sel, ok := n.X.(*ast.SelectorExpr); ok &&
+					n.Key == nil &&
+					n.Value == nil &&
+					sel.Sel.Name == "N" &&
+					isPtrToNamed(info.TypeOf(sel.X), "testing", "B") {
+
+					pass.Report(analysis.Diagnostic{
+						// Highlight "range b.N".
+						Pos:      n.Range,
+						End:      n.X.End(),
+						Category: "bloop",
+						Message:  "b.N can be modernized using b.Loop()",
+						SuggestedFixes: []analysis.SuggestedFix{{
+							Message:   "Replace b.N with b.Loop()",
+							TextEdits: edits(cur, sel.X, n.Range, n.X.End()),
+						}},
+					})
+				}
+			}
+		}
+		return true
+	})
+}
+
+// isPtrToNamed reports whether t is type "*pkgpath.Name".
+func isPtrToNamed(t types.Type, pkgpath, name string) bool {
+	if ptr, ok := t.(*types.Pointer); ok {
+		named, ok := ptr.Elem().(*types.Named)
+		return ok &&
+			named.Obj().Name() == name &&
+			named.Obj().Pkg().Path() == pkgpath
+	}
+	return false
+}
+
+// uses reports whether the subtree cur contains a use of obj.
+func uses(info *types.Info, cur cursor.Cursor, obj types.Object) bool {
+	for curId := range cur.Preorder((*ast.Ident)(nil)) {
+		if info.Uses[curId.Node().(*ast.Ident)] == obj {
+			return true
+		}
+	}
+	return false
+}
+
+// isMethod reports whether fn is pkgpath.(T).Name.
+func isMethod(fn *types.Func, pkgpath, T, name string) bool {
+	if recv := fn.Signature().Recv(); recv != nil {
+		_, recvName := typesinternal.ReceiverNamed(recv)
+		return recvName != nil &&
+			isPackageLevel(recvName.Obj(), pkgpath, T) &&
+			fn.Name() == name
+	}
+	return false
+}
+
+// enclosingFunc returns the cursor for the innermost Func{Decl,Lit}
+// that encloses (or is) c, if any.
+//
+// TODO(adonovan): consider adding:
+//
+//	func (Cursor) AnyEnclosing(filter ...ast.Node) (Cursor bool)
+//	func (Cursor) Enclosing[N ast.Node]() (Cursor, bool)
+//
+// See comments at [cursor.Cursor.Stack].
+func enclosingFunc(c cursor.Cursor) (cursor.Cursor, bool) {
+	for {
+		switch c.Node().(type) {
+		case *ast.FuncLit, *ast.FuncDecl:
+			return c, true
+		case nil:
+			return cursor.Cursor{}, false
+		}
+		c = c.Parent()
+	}
+}

gopls/internal/analysis/modernize/modernize.go

@@ -57,6 +57,7 @@ func run(pass *analysis.Pass) (any, error) {
 	sortslice(pass)
 	efaceany(pass)
 	appendclipped(pass)
+	bloop(pass)
 
 	// TODO(adonovan):
 	// - more modernizers here; see #70815.
@@ -106,3 +107,18 @@ func formatExprs(fset *token.FileSet, exprs []ast.Expr) string {
 	}
 	return buf.String()
 }
+
+// isZeroLiteral reports whether e is the literal 0.
+func isZeroLiteral(e ast.Expr) bool {
+	lit, ok := e.(*ast.BasicLit)
+	return ok && lit.Kind == token.INT && lit.Value == "0"
+}
+
+// isPackageLevel reports whether obj is the package-level symbol pkg.Name.
+func isPackageLevel(obj types.Object, pkgpath, name string) bool {
+	pkg := obj.Pkg()
+	return pkg != nil &&
+		obj.Parent() == pkg.Scope() &&
+		obj.Pkg().Path() == pkgpath &&
+		obj.Name() == name
+}

gopls/internal/analysis/modernize/modernize_test.go

@@ -16,5 +16,6 @@ func Test(t *testing.T) {
 		"minmax",
 		"sortslice",
 		"efaceany",
-		"appendclipped")
+		"appendclipped",
+		"bloop")
 }

gopls/internal/analysis/modernize/slices.go

@@ -9,7 +9,6 @@ package modernize
 import (
 	"fmt"
 	"go/ast"
-	"go/token"
 	"go/types"
 	"slices"
 
@@ -197,10 +196,6 @@ func isClippedSlice(info *types.Info, e ast.Expr) (clipped, empty bool) {
 	switch e := e.(type) {
 	case *ast.SliceExpr:
 		// x[:0:0], x[:len(x):len(x)], x[:k:k], x[:0]
-		isZeroLiteral := func(e ast.Expr) bool {
-			lit, ok := e.(*ast.BasicLit)
-			return ok && lit.Kind == token.INT && lit.Value == "0"
-		}
 		clipped = e.Slice3 && e.High != nil && e.Max != nil && equalSyntax(e.High, e.Max) // x[:k:k]
 		empty = e.High != nil && isZeroLiteral(e.High)                                    // x[:0:*]
 		return

gopls/internal/analysis/modernize/sortslice.go

@@ -116,13 +116,7 @@ func isQualifiedIdent(info *types.Info, e ast.Expr, pkgpath, name string) bool {
 		return false
 	}
 	obj, ok := info.Uses[id]
-	return ok && isPackageLevel(obj) && obj.Pkg().Path() == pkgpath && id.Name == name
-}
-
-// isPackageLevel reports whether obj is a package-level func/var/const/type.
-func isPackageLevel(obj types.Object) bool {
-	pkg := obj.Pkg()
-	return pkg != nil && obj.Parent() == pkg.Scope()
+	return ok && isPackageLevel(obj, pkgpath, name)
 }
 
 // enclosingFile returns the file enclosing pos.

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

@@ -0,0 +1 @@
+package bloop

gopls/internal/analysis/modernize/testdata/src/bloop/bloop_test.go

@@ -0,0 +1,69 @@
+//go:build go1.24
+
+package bloop
+
+import "testing"
+
+func BenchmarkA(b *testing.B) {
+	println("slow")
+	b.ResetTimer()
+
+	for range b.N { // want "b.N can be modernized using b.Loop.."
+	}
+}
+
+func BenchmarkB(b *testing.B) {
+	// setup
+	{
+		b.StopTimer()
+		println("slow")
+		b.StartTimer()
+	}
+
+	for i := range b.N { // Nope. Should we change this to "for i := 0; b.Loop(); i++"?
+		print(i)
+	}
+
+	b.StopTimer()
+	println("slow")
+}
+
+func BenchmarkC(b *testing.B) {
+	// setup
+	{
+		b.StopTimer()
+		println("slow")
+		b.StartTimer()
+	}
+
+	for i := 0; i < b.N; i++ { // want "b.N can be modernized using b.Loop.."
+		println("no uses of i")
+	}
+
+	b.StopTimer()
+	println("slow")
+}
+
+func BenchmarkD(b *testing.B) {
+	for i := 0; i < b.N; i++ { // want "b.N can be modernized using b.Loop.."
+		println(i)
+	}
+}
+
+func BenchmarkE(b *testing.B) {
+	b.Run("sub", func(b *testing.B) {
+		b.StopTimer() // not deleted
+		println("slow")
+		b.StartTimer() // not deleted
+
+		// ...
+	})
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ { // want "b.N can be modernized using b.Loop.."
+		println("no uses of i")
+	}
+
+	b.StopTimer()
+	println("slow")
+}