From 8da82a3535b6945028236f0390f7161f07a60547 Mon Sep 17 00:00:00 2001
From: Clifton Kaznocha <ckaznocha@gmail.com>
Date: Sat, 23 Mar 2024 11:27:47 -0700
Subject: [PATCH] Don't warn when loop Op needs to be re-evaluated (#15)

---
 intrange.go      | 76 +++++++++++++++++++++++++++++++++++++++---------
 testdata/main.go | 52 +++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+), 14 deletions(-)

diff --git a/intrange.go b/intrange.go
index 6f1f4b7..f8d037d 100644
--- a/intrange.go
+++ b/intrange.go
@@ -84,12 +84,16 @@ func check(pass *analysis.Pass) func(node ast.Node) {
 			return
 		}
 
+		var nExpr ast.Expr
+
 		switch cond.Op {
-		case token.LSS: // ;i < x;
+		case token.LSS: // ;i < n;
 			if isBenchmark(cond.Y) {
 				return
 			}
 
+			nExpr = findNExpr(cond.Y)
+
 			x, ok := cond.X.(*ast.Ident)
 			if !ok {
 				return
@@ -98,11 +102,13 @@ func check(pass *analysis.Pass) func(node ast.Node) {
 			if x.Name != initIdent.Name {
 				return
 			}
-		case token.GTR: // ;x > i;
+		case token.GTR: // ;n > i;
 			if isBenchmark(cond.X) {
 				return
 			}
 
+			nExpr = findNExpr(cond.X)
+
 			y, ok := cond.Y.(*ast.Ident)
 			if !ok {
 				return
@@ -209,6 +215,7 @@ func check(pass *analysis.Pass) func(node ast.Node) {
 
 		bc := &bodyChecker{
 			initIdent: initIdent,
+			nExpr:     nExpr,
 		}
 
 		ast.Inspect(forStmt.Body, bc.check)
@@ -224,6 +231,29 @@ func check(pass *analysis.Pass) func(node ast.Node) {
 	}
 }
 
+func findNExpr(expr ast.Expr) ast.Expr {
+	switch e := expr.(type) {
+	case *ast.CallExpr:
+		if e.Fun.(*ast.Ident).Name != "len" {
+			return nil
+		}
+
+		if len(e.Args) != 1 {
+			return nil
+		}
+
+		return findNExpr(e.Args[0])
+	case *ast.BasicLit:
+		return nil
+	case *ast.Ident:
+		return e
+	case *ast.SelectorExpr:
+		return e
+	default:
+		return nil
+	}
+}
+
 func isBenchmark(expr ast.Expr) bool {
 	selectorExpr, ok := expr.(*ast.SelectorExpr)
 	if !ok {
@@ -246,8 +276,36 @@ func isBenchmark(expr ast.Expr) bool {
 	return false
 }
 
+func identEqual(a, b ast.Expr) bool {
+	if a == nil || b == nil {
+		return false
+	}
+
+	switch aT := a.(type) {
+	case *ast.Ident:
+		identB, ok := b.(*ast.Ident)
+		if !ok {
+			return false
+		}
+
+		return aT.Name == identB.Name
+	case *ast.SelectorExpr:
+		selectorB, ok := b.(*ast.SelectorExpr)
+		if !ok {
+			return false
+		}
+
+		return identEqual(aT.Sel, selectorB.Sel) && identEqual(aT.X, selectorB.X)
+	case *ast.IndexExpr:
+		return identEqual(aT.X, b)
+	default:
+		return false
+	}
+}
+
 type bodyChecker struct {
 	initIdent *ast.Ident
+	nExpr     ast.Expr
 	modified  bool
 }
 
@@ -255,24 +313,14 @@ func (b *bodyChecker) check(n ast.Node) bool {
 	switch stmt := n.(type) {
 	case *ast.AssignStmt:
 		for _, lhs := range stmt.Lhs {
-			ident, ok := lhs.(*ast.Ident)
-			if !ok {
-				continue
-			}
-
-			if b.initIdent.Name == ident.Name {
+			if identEqual(lhs, b.initIdent) || identEqual(lhs, b.nExpr) {
 				b.modified = true
 
 				return false
 			}
 		}
 	case *ast.IncDecStmt:
-		ident, ok := stmt.X.(*ast.Ident)
-		if !ok {
-			return true
-		}
-
-		if b.initIdent.Name == ident.Name {
+		if identEqual(stmt.X, b.initIdent) || identEqual(stmt.X, b.nExpr) {
 			b.modified = true
 
 			return false
diff --git a/testdata/main.go b/testdata/main.go
index e5ba782..72292e5 100644
--- a/testdata/main.go
+++ b/testdata/main.go
@@ -134,4 +134,56 @@ func main() {
 
 	for i := 0; b.N >= i; i++ {
 	}
+
+	var n int
+	for i := 0; i < n; i++ {
+		n--
+	}
+
+	for i := 0; i < n; i++ {
+		n++
+	}
+
+	// Example from https://github.com/ckaznocha/intrange/issues/12
+	var what string
+	for i := 0; i < len(what); i++ {
+		if what[i] == 'v' && i+1 < len(what) && what[i+1] >= '0' && what[i+1] <= '9' {
+			what = what[:i] + what[i+1:]
+		}
+	}
+
+	for i := 0; i < len(what); i++ { // want `for loop can be changed to use an integer range \(Go 1\.22\+\)`
+	}
+
+	var t struct{ n int }
+	for i := 0; i < t.n; i++ { // want `for loop can be changed to use an integer range \(Go 1\.22\+\)`
+	}
+
+	for i := 0; i < t.n; i++ {
+		t.n++
+	}
+
+	var s []int
+	for i := 0; i < len(s); i++ { // want `for loop can be changed to use an integer range \(Go 1\.22\+\)`
+	}
+
+	for i := 0; i < len(s); i++ {
+		s = append(s, 4)
+	}
+
+	var m map[int]int
+	for i := 0; i < len(m); i++ { // want `for loop can be changed to use an integer range \(Go 1\.22\+\)`
+	}
+
+	for i := 0; i < len(m); i++ {
+		m[4] = 4
+	}
+
+	var t2 struct{ m map[int]int }
+	for i := 0; i < len(t2.m); i++ { // want `for loop can be changed to use an integer range \(Go 1\.22\+\)`
+	}
+
+	for i := 0; i < len(t2.m); i++ {
+		t2.m[4] = 4
+	}
 }