testing: recognize helper functions that use range-over-function

This assumes the current behavior of the gc compiler:
range functions use a suffix of -rangeNNNN.

Fixes #72794

Change-Id: I3c10c60829853cf2cb4c17a75f6243def0313ae9
Reviewed-on: https://go-review.googlesource.com/c/go/+/656775
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: David Chase <drchase@google.com>
Auto-Submit: David Chase <drchase@google.com>

3 files changed, 64 insertions, 1 deletions

diff --git a/src/testing/helper_test.go b/src/testing/helper_test.go
index a698e79fa9..8efd89f217 100644
--- a/src/testing/helper_test.go
+++ b/src/testing/helper_test.go
@@ -85,6 +85,28 @@ func TestTBHelperParallel(t *testing.T) {
 	}
 }
 
+// Issue 72794.
+func TestHelperRange(t *testing.T) {
+	if os.Getenv("GO_WANT_HELPER_PROCESS") == "1" {
+		rangeHelperHelper(t)
+		return
+	}
+
+	t.Parallel()
+
+	cmd := testenv.Command(t, testenv.Executable(t), "-test.run=^TestHelperRange$")
+	cmd = testenv.CleanCmdEnv(cmd)
+	cmd.Env = append(cmd.Env, "GO_WANT_HELPER_PROCESS=1")
+	out, _ := cmd.CombinedOutput()
+	want := `--- FAIL: TestHelperRange \([^)]+\)
+    helperfuncs_test.go:139: range
+    helperfuncs_test.go:139: range
+`
+	if !regexp.MustCompile(want).Match(out) {
+		t.Errorf("got output:\n\n%s\nwant matching:\n\n%s", out, want)
+	}
+}
+
 func BenchmarkTBHelper(b *testing.B) {
 	f1 := func() {
 		b.Helper()
diff --git a/src/testing/helperfuncs_test.go b/src/testing/helperfuncs_test.go
index f0295f35df..d3dadde81e 100644
--- a/src/testing/helperfuncs_test.go
+++ b/src/testing/helperfuncs_test.go
@@ -122,3 +122,19 @@ func doPanic(t *testing.T, msg string) {
 	t.Helper()
 	panic(msg)
 }
+
+func rangeHelper(t *testing.T, msg string) {
+	t.Helper()
+	iter := func(yield func(int) bool) {
+		if yield(0) {
+			yield(1)
+		}
+	}
+	for range iter {
+		t.Error(msg)
+	}
+}
+
+func rangeHelperHelper(t *testing.T) {
+	rangeHelper(t, "range")
+}
diff --git a/src/testing/testing.go b/src/testing/testing.go
index 626772e57a..f6ecd5b901 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -816,8 +816,15 @@ func (c *common) frameSkip(skip int) runtime.Frame {
 	}
 	frames := runtime.CallersFrames(pc[:n])
 	var firstFrame, prevFrame, frame runtime.Frame
+	skipRange := false
 	for more := true; more; prevFrame = frame {
 		frame, more = frames.Next()
+		if skipRange {
+			// Skip the iterator function when a helper
+			// functions does a range over function.
+			skipRange = false
+			continue
+		}
 		if frame.Function == "runtime.gopanic" {
 			continue
 		}
@@ -861,7 +868,25 @@ func (c *common) frameSkip(skip int) runtime.Frame {
 				c.helperNames[pcToName(pc)] = struct{}{}
 			}
 		}
-		if _, ok := c.helperNames[frame.Function]; !ok {
+
+		fnName := frame.Function
+		// Ignore trailing -rangeN used for iterator functions.
+		const rangeSuffix = "-range"
+		if suffixIdx := strings.LastIndex(fnName, rangeSuffix); suffixIdx > 0 {
+			ok := true
+			for i := suffixIdx + len(rangeSuffix); i < len(fnName); i++ {
+				if fnName[i] < '0' || fnName[i] > '9' {
+					ok = false
+					break
+				}
+			}
+			if ok {
+				fnName = fnName[:suffixIdx]
+				skipRange = true
+			}
+		}
+
+		if _, ok := c.helperNames[fnName]; !ok {
 			// Found a frame that wasn't inside a helper function.
 			return frame
 		}