text/template: support range-over-func

For #66107

Change-Id: I2fcd04bebe80346dbd244ab7ea09cbe6010b9d8e
GitHub-Last-Rev: 5ebf615db5889a04738c555c651e07c1fd287748
GitHub-Pull-Request: golang/go#68329
Reviewed-on: https://go-review.googlesource.com/c/go/+/596956
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Carlos Amedee <carlos@golang.org>
Reviewed-by: Ian Lance Taylor <iant@google.com>
Auto-Submit: Ian Lance Taylor <iant@google.com>

3 files changed, 73 insertions, 2 deletions

diff --git a/src/text/template/doc.go b/src/text/template/doc.go
index 12f6fe0d1c..847f96b725 100644
--- a/src/text/template/doc.go
+++ b/src/text/template/doc.go
@@ -98,7 +98,8 @@ data, defined in detail in the corresponding sections that follow.
 			{{if pipeline}} T1 {{else}}{{if pipeline}} T0 {{end}}{{end}}
 
 	{{range pipeline}} T1 {{end}}
-		The value of the pipeline must be an array, slice, map, or channel.
+		The value of the pipeline must be an array, slice, map, iter.Seq,
+		iter.Seq2 or channel.
 		If the value of the pipeline has length zero, nothing is output;
 		otherwise, dot is set to the successive elements of the array,
 		slice, or map and T1 is executed. If the value is a map and the
@@ -106,7 +107,8 @@ data, defined in detail in the corresponding sections that follow.
 		visited in sorted key order.
 
 	{{range pipeline}} T1 {{else}} T0 {{end}}
-		The value of the pipeline must be an array, slice, map, or channel.
+		The value of the pipeline must be an array, slice, map, iter.Seq,
+		iter.Seq2 or channel.
 		If the value of the pipeline has length zero, dot is unaffected and
 		T0 is executed; otherwise, dot is set to the successive elements
 		of the array, slice, or map and T1 is executed.
diff --git a/src/text/template/exec.go b/src/text/template/exec.go
index 5b35b3e5a8..96d2f50ef8 100644
--- a/src/text/template/exec.go
+++ b/src/text/template/exec.go
@@ -434,6 +434,43 @@ func (s *state) walkRange(dot reflect.Value, r *parse.RangeNode) {
 		return
 	case reflect.Invalid:
 		break // An invalid value is likely a nil map, etc. and acts like an empty map.
+	case reflect.Func:
+		if val.Type().CanSeq() {
+			if len(r.Pipe.Decl) > 1 {
+				s.errorf("can't use %s iterate over more than one variable", val)
+				break
+			}
+			run := false
+			for v := range val.Seq() {
+				run = true
+				// Pass element as second value,
+				// as we do for channels.
+				oneIteration(reflect.Value{}, v)
+			}
+			if !run {
+				break
+			}
+			return
+		}
+		if val.Type().CanSeq2() {
+			run := false
+			for i, v := range val.Seq2() {
+				run = true
+				if len(r.Pipe.Decl) > 1 {
+					oneIteration(i, v)
+				} else {
+					// If there is only one range variable,
+					// oneIteration will use the
+					// second value.
+					oneIteration(reflect.Value{}, i)
+				}
+			}
+			if !run {
+				break
+			}
+			return
+		}
+		fallthrough
 	default:
 		s.errorf("range can't iterate over %v", val)
 	}
diff --git a/src/text/template/exec_test.go b/src/text/template/exec_test.go
index 9903e17d0e..b84e278c12 100644
--- a/src/text/template/exec_test.go
+++ b/src/text/template/exec_test.go
@@ -10,6 +10,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"iter"
 	"reflect"
 	"strings"
 	"sync"
@@ -601,6 +602,17 @@ var execTests = []execTest{
 	{"declare in range", "{{range $x := .PSI}}<{{$foo:=$x}}{{$x}}>{{end}}", "<21><22><23>", tVal, true},
 	{"range count", `{{range $i, $x := count 5}}[{{$i}}]{{$x}}{{end}}`, "[0]a[1]b[2]c[3]d[4]e", tVal, true},
 	{"range nil count", `{{range $i, $x := count 0}}{{else}}empty{{end}}`, "empty", tVal, true},
+	{"range iter.Seq[int]", `{{range $i := .}}{{$i}}{{end}}`, "01", fVal1(2), true},
+	{"i = range iter.Seq[int]", `{{$i := 0}}{{range $i = .}}{{$i}}{{end}}`, "01", fVal1(2), true},
+	{"range iter.Seq[int] over two var", `{{range $i, $c := .}}{{$c}}{{end}}`, "", fVal1(2), false},
+	{"i, c := range iter.Seq2[int,int]", `{{range $i, $c := .}}{{$i}}{{$c}}{{end}}`, "0112", fVal2(2), true},
+	{"i, c = range iter.Seq2[int,int]", `{{$i := 0}}{{$c := 0}}{{range $i, $c = .}}{{$i}}{{$c}}{{end}}`, "0112", fVal2(2), true},
+	{"i = range iter.Seq2[int,int]", `{{$i := 0}}{{range $i = .}}{{$i}}{{end}}`, "01", fVal2(2), true},
+	{"i := range iter.Seq2[int,int]", `{{range $i := .}}{{$i}}{{end}}`, "01", fVal2(2), true},
+	{"i,c,x range iter.Seq2[int,int]", `{{$i := 0}}{{$c := 0}}{{$x := 0}}{{range $i, $c = .}}{{$i}}{{$c}}{{end}}`, "0112", fVal2(2), true},
+	{"i,x range iter.Seq[int]", `{{$i := 0}}{{$x := 0}}{{range $i = .}}{{$i}}{{end}}`, "01", fVal1(2), true},
+	{"range iter.Seq[int] else", `{{range $i := .}}{{$i}}{{else}}empty{{end}}`, "empty", fVal1(0), true},
+	{"range iter.Seq2[int,int] else", `{{range $i := .}}{{$i}}{{else}}empty{{end}}`, "empty", fVal2(0), true},
 
 	// Cute examples.
 	{"or as if true", `{{or .SI "slice is empty"}}`, "[3 4 5]", tVal, true},
@@ -705,6 +717,26 @@ var execTests = []execTest{
 	{"issue60801", "{{$k := 0}}{{$v := 0}}{{range $k, $v = .AI}}{{$k}}={{$v}} {{end}}", "0=3 1=4 2=5 ", tVal, true},
 }
 
+func fVal1(i int) iter.Seq[int] {
+	return func(yield func(int) bool) {
+		for v := range i {
+			if !yield(v) {
+				break
+			}
+		}
+	}
+}
+
+func fVal2(i int) iter.Seq2[int, int] {
+	return func(yield func(int, int) bool) {
+		for v := range i {
+			if !yield(v, v+1) {
+				break
+			}
+		}
+	}
+}
+
 func zeroArgs() string {
 	return "zeroArgs"
 }