[release-branch.r60] template: range over channel

««« CL 4951046 / 379452c085ab
template: range over channel

R=golang-dev, dsymonds
CC=golang-dev
https://golang.org/cl/4951046
»»»

R=dsymonds
CC=golang-dev
https://golang.org/cl/4983052

2 files changed, 52 insertions, 22 deletions

diff --git a/src/pkg/template/exec.go b/src/pkg/template/exec.go
index eaa57ae818..dbe6541dbf 100644
--- a/src/pkg/template/exec.go
+++ b/src/pkg/template/exec.go
@@ -196,23 +196,25 @@ func (s *state) walkRange(dot reflect.Value, r *parse.RangeNode) {
 	val, _ := indirect(s.evalPipeline(dot, r.Pipe))
 	// mark top of stack before any variables in the body are pushed.
 	mark := s.mark()
+	oneIteration := func(index, elem reflect.Value) {
+		// Set top var (lexically the second if there are two) to the element.
+		if len(r.Pipe.Decl) > 0 {
+			s.setVar(1, elem)
+		}
+		// Set next var (lexically the first if there are two) to the index.
+		if len(r.Pipe.Decl) > 1 {
+			s.setVar(2, index)
+		}
+		s.walk(elem, r.List)
+		s.pop(mark)
+	}
 	switch val.Kind() {
 	case reflect.Array, reflect.Slice:
 		if val.Len() == 0 {
 			break
 		}
 		for i := 0; i < val.Len(); i++ {
-			elem := val.Index(i)
-			// Set top var (lexically the second if there are two) to the element.
-			if len(r.Pipe.Decl) > 0 {
-				s.setVar(1, elem)
-			}
-			// Set next var (lexically the first if there are two) to the index.
-			if len(r.Pipe.Decl) > 1 {
-				s.setVar(2, reflect.ValueOf(i))
-			}
-			s.walk(elem, r.List)
-			s.pop(mark)
+			oneIteration(reflect.ValueOf(i), val.Index(i))
 		}
 		return
 	case reflect.Map:
@@ -220,17 +222,23 @@ func (s *state) walkRange(dot reflect.Value, r *parse.RangeNode) {
 			break
 		}
 		for _, key := range val.MapKeys() {
-			elem := val.MapIndex(key)
-			// Set top var (lexically the second if there are two) to the element.
-			if len(r.Pipe.Decl) > 0 {
-				s.setVar(1, elem)
-			}
-			// Set next var (lexically the first if there are two) to the key.
-			if len(r.Pipe.Decl) > 1 {
-				s.setVar(2, key)
+			oneIteration(key, val.MapIndex(key))
+		}
+		return
+	case reflect.Chan:
+		if val.IsNil() {
+			break
+		}
+		i := 0
+		for ; ; i++ {
+			elem, ok := val.Recv()
+			if !ok {
+				break
 			}
-			s.walk(elem, r.List)
-			s.pop(mark)
+			oneIteration(reflect.ValueOf(i), elem)
+		}
+		if i == 0 {
+			break
 		}
 		return
 	case reflect.Invalid:
diff --git a/src/pkg/template/exec_test.go b/src/pkg/template/exec_test.go
index 82f56e13cb..7e07e8c2db 100644
--- a/src/pkg/template/exec_test.go
+++ b/src/pkg/template/exec_test.go
@@ -391,6 +391,8 @@ var execTests = []execTest{
 	{"range $x $y MSIone", "{{range $x, $y := .MSIone}}<{{$x}}={{$y}}>{{end}}", "<one=1>", tVal, true},
 	{"range $x PSI", "{{range $x := .PSI}}<{{$x}}>{{end}}", "<21><22><23>", tVal, true},
 	{"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},
 
 	// Cute examples.
 	{"or as if true", `{{or .SI "slice is empty"}}`, "[3 4 5]", tVal, true},
@@ -424,9 +426,29 @@ func oneArg(a string) string {
 	return "oneArg=" + a
 }
 
+// count returns a channel that will deliver n sequential 1-letter strings starting at "a"
+func count(n int) chan string {
+	if n == 0 {
+		return nil
+	}
+	c := make(chan string)
+	go func() {
+		for i := 0; i < n; i++ {
+			c <- "abcdefghijklmnop"[i : i+1]
+		}
+		close(c)
+	}()
+	return c
+}
+
 func testExecute(execTests []execTest, set *Set, t *testing.T) {
 	b := new(bytes.Buffer)
-	funcs := FuncMap{"zeroArgs": zeroArgs, "oneArg": oneArg, "typeOf": typeOf}
+	funcs := FuncMap{
+		"count":    count,
+		"oneArg":   oneArg,
+		"typeOf":   typeOf,
+		"zeroArgs": zeroArgs,
+	}
 	for _, test := range execTests {
 		tmpl := New(test.name).Funcs(funcs)
 		_, err := tmpl.ParseInSet(test.input, set)