From 15a736f88da05e95b6f7fbda6faf224bb084aad3 Mon Sep 17 00:00:00 2001
From: Clifton Kaznocha <ckaznocha@users.noreply.github.com>
Date: Thu, 21 Mar 2024 09:31:43 -0700
Subject: [PATCH] Don't warn on benchmarks

---
 intrange.go      | 30 ++++++++++++++++++++++++++++++
 testdata/main.go |  9 +++++++++
 2 files changed, 39 insertions(+)

diff --git a/intrange.go b/intrange.go
index 0a7ae3d..6f1f4b7 100644
--- a/intrange.go
+++ b/intrange.go
@@ -86,6 +86,10 @@ func check(pass *analysis.Pass) func(node ast.Node) {
 
 		switch cond.Op {
 		case token.LSS: // ;i < x;
+			if isBenchmark(cond.Y) {
+				return
+			}
+
 			x, ok := cond.X.(*ast.Ident)
 			if !ok {
 				return
@@ -95,6 +99,10 @@ func check(pass *analysis.Pass) func(node ast.Node) {
 				return
 			}
 		case token.GTR: // ;x > i;
+			if isBenchmark(cond.X) {
+				return
+			}
+
 			y, ok := cond.Y.(*ast.Ident)
 			if !ok {
 				return
@@ -216,6 +224,28 @@ func check(pass *analysis.Pass) func(node ast.Node) {
 	}
 }
 
+func isBenchmark(expr ast.Expr) bool {
+	selectorExpr, ok := expr.(*ast.SelectorExpr)
+	if !ok {
+		return false
+	}
+
+	if selectorExpr.Sel.Name != "N" {
+		return false
+	}
+
+	ident, ok := selectorExpr.X.(*ast.Ident)
+	if !ok {
+		return false
+	}
+
+	if ident.Name == "b" {
+		return true
+	}
+
+	return false
+}
+
 type bodyChecker struct {
 	initIdent *ast.Ident
 	modified  bool
diff --git a/testdata/main.go b/testdata/main.go
index 4c87f33..e5ba782 100644
--- a/testdata/main.go
+++ b/testdata/main.go
@@ -1,5 +1,7 @@
 package main
 
+import "testing"
+
 func main() {
 	for i := 2; i < 10; i++ {
 	}
@@ -125,4 +127,11 @@ func main() {
 
 	for i := 0; x > i; i = 0x1 + i { // want `for loop can be changed to use an integer range \(Go 1\.22\+\)`
 	}
+
+	var b *testing.B
+	for i := 0; i < b.N; i++ {
+	}
+
+	for i := 0; b.N >= i; i++ {
+	}
 }