From f46eecacd29786eab36c065ca621c79215a99568 Mon Sep 17 00:00:00 2001
From: Clifton Kaznocha <ckaznocha@users.noreply.github.com>
Date: Mon, 9 Sep 2024 20:27:01 -0700
Subject: [PATCH] Catch more patterns where range len could be simpler

---
 intrange.go             | 60 ++++++++++++++++++++++++++++++-----------
 testdata/main.go        | 22 +++++++++++++++
 testdata/main.go.golden | 22 +++++++++++++++
 3 files changed, 88 insertions(+), 16 deletions(-)

diff --git a/intrange.go b/intrange.go
index 33cddf3..44a1509 100644
--- a/intrange.go
+++ b/intrange.go
@@ -25,8 +25,9 @@ var (
 )
 
 const (
-	msg         = "for loop can be changed to use an integer range (Go 1.22+)"
-	msgLenRange = "for loop can be changed to `i := range %s`"
+	msg                = "for loop can be changed to use an integer range (Go 1.22+)"
+	msgLenRange        = "for loop can be changed to `%s := range %s`"
+	msgLenRangeNoIdent = "for loop can be changed to `range %s`"
 )
 
 func run(pass *analysis.Pass) (any, error) {
@@ -243,21 +244,26 @@ func checkForStmt(pass *analysis.Pass, forStmt *ast.ForStmt) {
 }
 
 func checkRangeStmt(pass *analysis.Pass, rangeStmt *ast.RangeStmt) {
-	if rangeStmt.Key == nil {
+	if rangeStmt.Value != nil {
 		return
 	}
 
-	ident, ok := rangeStmt.Key.(*ast.Ident)
-	if !ok {
-		return
-	}
+	startPos := rangeStmt.Range
+	usesKey := rangeStmt.Key != nil
+	identName := ""
 
-	if ident.Name == "_" {
-		return
-	}
+	if usesKey {
+		ident, ok := rangeStmt.Key.(*ast.Ident)
+		if !ok {
+			return
+		}
 
-	if rangeStmt.Value != nil {
-		return
+		if ident.Name == "_" {
+			usesKey = false
+		}
+
+		identName = ident.Name
+		startPos = ident.Pos()
 	}
 
 	if rangeStmt.X == nil {
@@ -295,18 +301,40 @@ func checkRangeStmt(pass *analysis.Pass, rangeStmt *ast.RangeStmt) {
 		return
 	}
 
+	if usesKey {
+		pass.Report(analysis.Diagnostic{
+			Pos:     startPos,
+			End:     x.End(),
+			Message: fmt.Sprintf(msgLenRange, identName, arg.Name),
+			SuggestedFixes: []analysis.SuggestedFix{
+				{
+					Message: fmt.Sprintf("Replace `len(%s)` with `%s`", arg.Name, arg.Name),
+					TextEdits: []analysis.TextEdit{
+						{
+							Pos:     x.Pos(),
+							End:     x.End(),
+							NewText: []byte(arg.Name),
+						},
+					},
+				},
+			},
+		})
+
+		return
+	}
+
 	pass.Report(analysis.Diagnostic{
-		Pos:     ident.Pos(),
+		Pos:     startPos,
 		End:     x.End(),
-		Message: fmt.Sprintf(msgLenRange, arg.Name),
+		Message: fmt.Sprintf(msgLenRangeNoIdent, arg.Name),
 		SuggestedFixes: []analysis.SuggestedFix{
 			{
 				Message: fmt.Sprintf("Replace `len(%s)` with `%s`", arg.Name, arg.Name),
 				TextEdits: []analysis.TextEdit{
 					{
-						Pos:     x.Pos(),
+						Pos:     startPos,
 						End:     x.End(),
-						NewText: []byte(arg.Name),
+						NewText: []byte(fmt.Sprintf("range %s", arg.Name)),
 					},
 				},
 			},
diff --git a/testdata/main.go b/testdata/main.go
index 1130c8e..4afde5a 100644
--- a/testdata/main.go
+++ b/testdata/main.go
@@ -234,6 +234,16 @@ func issue27() {
 		print(i)
 	}
 
+	for n := range len(someSlice) { // want "for loop can be changed to `n := range someSlice`"
+		print(n)
+	}
+
+	for _ = range len(someSlice) { // want "for loop can be changed to `range someSlice`"
+	}
+
+	for range len(someSlice) { // want "for loop can be changed to `range someSlice`"
+	}
+
 	for i := range notLen(someSlice) {
 		print(i)
 	}
@@ -256,6 +266,12 @@ func issue27() {
 		print(i)
 	}
 
+	for _ = range len(someArray) { // want "for loop can be changed to `range someArray`"
+	}
+
+	for range len(someArray) { // want "for loop can be changed to `range someArray`"
+	}
+
 	for i := range notLen(someArray) {
 		print(i)
 	}
@@ -277,6 +293,12 @@ func issue27() {
 		print(i)
 	}
 
+	for _ = range len(someMap) {
+	}
+
+	for range len(someMap) {
+	}
+
 	for i := range notLen(someMap) {
 		print(i)
 	}
diff --git a/testdata/main.go.golden b/testdata/main.go.golden
index e4e7ad5..2d9027e 100644
--- a/testdata/main.go.golden
+++ b/testdata/main.go.golden
@@ -234,6 +234,16 @@ func issue27() {
 		print(i)
 	}
 
+	for n := range someSlice { // want "for loop can be changed to `n := range someSlice`"
+		print(n)
+	}
+
+	for range someSlice { // want "for loop can be changed to `range someSlice`"
+	}
+
+	for range someSlice { // want "for loop can be changed to `range someSlice`"
+	}
+
 	for i := range notLen(someSlice) {
 		print(i)
 	}
@@ -256,6 +266,12 @@ func issue27() {
 		print(i)
 	}
 
+	for range someArray { // want "for loop can be changed to `range someArray`"
+	}
+
+	for range someArray { // want "for loop can be changed to `range someArray`"
+	}
+
 	for i := range notLen(someArray) {
 		print(i)
 	}
@@ -277,6 +293,12 @@ func issue27() {
 		print(i)
 	}
 
+	for _ = range len(someMap) {
+	}
+
+	for range len(someMap) {
+	}
+
 	for i := range notLen(someMap) {
 		print(i)
 	}