From 0f63406c4d667eaa5865dd42a745e205fd4f7981 Mon Sep 17 00:00:00 2001
From: Joe <56809242+Jo3-L@users.noreply.github.com>
Date: Fri, 16 Apr 2021 10:50:59 -0700
Subject: [PATCH] all: add break & continue loop actions

This is mostly taken from https://go-review.googlesource.com/c/go/+/66410/, with some edits to support while actions. The above commit was reverted due to issues with html/template (see golang/go#23683). However, as we don't have such issues, there should be no issue adding it to this fork.
---
 doc.go              |  6 +++
 exec.go             | 90 +++++++++++++++++++++++++++++++++------------
 exec_test.go        |  8 ++++
 parse/lex.go        |  4 ++
 parse/lex_test.go   |  8 +++-
 parse/node.go       | 48 ++++++++++++++++++++++++
 parse/parse.go      | 34 +++++++++++++++++
 parse/parse_test.go | 21 +++++++++++
 8 files changed, 195 insertions(+), 24 deletions(-)

diff --git a/doc.go b/doc.go
index c502b8f..2285dda 100644
--- a/doc.go
+++ b/doc.go
@@ -119,6 +119,12 @@ data, defined in detail in the corresponding sections that follow.
 		Execute T1 while the value of the pipeline is not empty. If the initial
 		value of the pipeline was empty, evaluate T0. Dot is unaffected.
 
+	{{break}}
+		Break out of the surrounding range or while loop.
+
+	{{continue}}
+		Begin the next iteration of the surrounding range or while loop.
+
 	{{template "name"}}
 		The template with the specified name is executed with nil data.
 
diff --git a/exec.go b/exec.go
index 2bf19dd..d5bb774 100644
--- a/exec.go
+++ b/exec.go
@@ -38,6 +38,7 @@ type state struct {
 	node       parse.Node // current node, for errors
 	vars       []variable // push-down stack of variable values.
 	depth      int        // the height of the stack of executing templates.
+	loopDepth  int        // nesting level of range/while loops.
 	operations int
 
 	parent *state
@@ -247,9 +248,17 @@ func (t *Template) DefinedTemplates() string {
 	return s
 }
 
+type loopControl int8
+
+const (
+	loopNone     loopControl = iota // no action.
+	loopBreak                       // break out of loop.
+	loopContinue                    // continues next loop iteration.
+)
+
 // Walk functions step through the major pieces of the template structure,
 // generating output as they go.
-func (s *state) walk(dot reflect.Value, node parse.Node) {
+func (s *state) walk(dot reflect.Value, node parse.Node) loopControl {
 	s.at(node)
 	switch node := node.(type) {
 	case *parse.ActionNode:
@@ -260,15 +269,17 @@ func (s *state) walk(dot reflect.Value, node parse.Node) {
 			s.printValue(node, val)
 		}
 	case *parse.IfNode:
-		s.walkIfOrWith(parse.NodeIf, dot, node.Pipe, node.List, node.ElseList)
+		return s.walkIfOrWith(parse.NodeIf, dot, node.Pipe, node.List, node.ElseList)
 	case *parse.ListNode:
 		for _, node := range node.Nodes {
-			s.walk(dot, node)
+			if c := s.walk(dot, node); c != loopNone {
+				return c
+			}
 		}
 	case *parse.RangeNode:
-		s.walkRange(dot, node)
+		return s.walkRange(dot, node)
 	case *parse.WhileNode:
-		s.walkWhile(dot, node)
+		return s.walkWhile(dot, node)
 	case *parse.TemplateNode:
 		s.walkTemplate(dot, node)
 	case *parse.TextNode:
@@ -276,15 +287,26 @@ func (s *state) walk(dot reflect.Value, node parse.Node) {
 			s.writeError(err)
 		}
 	case *parse.WithNode:
-		s.walkIfOrWith(parse.NodeWith, dot, node.Pipe, node.List, node.ElseList)
+		return s.walkIfOrWith(parse.NodeWith, dot, node.Pipe, node.List, node.ElseList)
+	case *parse.BreakNode:
+		if s.loopDepth == 0 {
+			s.errorf("invalid break outside of loop")
+		}
+		return loopBreak
+	case *parse.ContinueNode:
+		if s.loopDepth == 0 {
+			s.errorf("invalid continue outside of loop")
+		}
+		return loopContinue
 	default:
 		s.errorf("unknown node: %s", node)
 	}
+	return loopNone
 }
 
 // walkIfOrWith walks an 'if' or 'with' node. The two control structures
 // are identical in behavior except that 'with' sets dot.
-func (s *state) walkIfOrWith(typ parse.NodeType, dot reflect.Value, pipe *parse.PipeNode, list, elseList *parse.ListNode) {
+func (s *state) walkIfOrWith(typ parse.NodeType, dot reflect.Value, pipe *parse.PipeNode, list, elseList *parse.ListNode) loopControl {
 	defer s.pop(s.mark())
 	val := s.evalPipeline(dot, pipe)
 	truth, ok := isTrue(indirectInterface(val))
@@ -293,13 +315,15 @@ func (s *state) walkIfOrWith(typ parse.NodeType, dot reflect.Value, pipe *parse.
 	}
 	if truth {
 		if typ == parse.NodeWith {
-			s.walk(val, list)
+			return s.walk(val, list)
 		} else {
-			s.walk(dot, list)
+			return s.walk(dot, list)
 		}
 	} else if elseList != nil {
-		s.walk(dot, elseList)
+		return s.walk(dot, elseList)
 	}
+
+	return loopNone
 }
 
 // IsTrue reports whether the value is 'true', in the sense of not the zero of its type,
@@ -337,7 +361,7 @@ func isTrue(val reflect.Value) (truth, ok bool) {
 	return truth, true
 }
 
-func (s *state) walkRange(dot reflect.Value, r *parse.RangeNode) {
+func (s *state) walkRange(dot reflect.Value, r *parse.RangeNode) loopControl {
 	s.incrOPs(1)
 
 	s.at(r)
@@ -345,7 +369,8 @@ 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) {
+	s.loopDepth++
+	oneIteration := func(index, elem reflect.Value) loopControl {
 		s.incrOPs(1)
 
 		// Set top var (lexically the second if there are two) to the element.
@@ -356,8 +381,9 @@ func (s *state) walkRange(dot reflect.Value, r *parse.RangeNode) {
 		if len(r.Pipe.Decl) > 1 {
 			s.setTopVar(2, index)
 		}
-		s.walk(elem, r.List)
+		ctrl := s.walk(elem, r.List)
 		s.pop(mark)
+		return ctrl
 	}
 	switch val.Kind() {
 	case reflect.Array, reflect.Slice:
@@ -365,17 +391,23 @@ func (s *state) walkRange(dot reflect.Value, r *parse.RangeNode) {
 			break
 		}
 		for i := 0; i < val.Len(); i++ {
-			oneIteration(reflect.ValueOf(i), val.Index(i))
+			if ctrl := oneIteration(reflect.ValueOf(i), val.Index(i)); ctrl == loopBreak {
+				break
+			}
 		}
-		return
+		s.loopDepth--
+		return loopNone
 	case reflect.Map:
 		if val.Len() == 0 {
 			break
 		}
 		for _, key := range sortKeys(val.MapKeys()) {
-			oneIteration(key, val.MapIndex(key))
+			if ctrl := oneIteration(key, val.MapIndex(key)); ctrl == loopBreak {
+				break
+			}
 		}
-		return
+		s.loopDepth--
+		return loopNone
 	case reflect.Chan:
 		if val.IsNil() {
 			break
@@ -386,29 +418,35 @@ func (s *state) walkRange(dot reflect.Value, r *parse.RangeNode) {
 			if !ok {
 				break
 			}
-			oneIteration(reflect.ValueOf(i), elem)
+			if ctrl := oneIteration(reflect.ValueOf(i), elem); ctrl == loopBreak {
+				break
+			}
 		}
 		if i == 0 {
 			break
 		}
-		return
+		s.loopDepth--
+		return loopNone
 	case reflect.Invalid:
 		break // An invalid value is likely a nil map, etc. and acts like an empty map.
 	default:
 		s.errorf("range can't iterate over %v", val)
 	}
+	s.loopDepth--
 	if r.ElseList != nil {
-		s.walk(dot, r.ElseList)
+		return s.walk(dot, r.ElseList)
 	}
+	return loopNone
 }
 
-func (s *state) walkWhile(dot reflect.Value, w *parse.WhileNode) {
+func (s *state) walkWhile(dot reflect.Value, w *parse.WhileNode) loopControl {
 	s.incrOPs(1)
 
 	s.at(w)
 	defer s.pop(s.mark())
 	// mark top of stack before any variables in the body are pushed.
 	mark := s.mark()
+	s.loopDepth++
 
 	i := 0
 	for ; ; i++ {
@@ -423,13 +461,19 @@ func (s *state) walkWhile(dot reflect.Value, w *parse.WhileNode) {
 			break
 		}
 
-		s.walk(dot, w.List)
+		ctrl := s.walk(dot, w.List)
 		s.pop(mark)
+		if ctrl == loopBreak {
+			break
+		}
 	}
 
+	s.loopDepth--
 	if i == 0 && w.ElseList != nil {
-		s.walk(dot, w.ElseList)
+		return s.walk(dot, w.ElseList)
 	}
+
+	return loopNone
 }
 
 func (s *state) walkTemplate(dot reflect.Value, t *parse.TemplateNode) {
diff --git a/exec_test.go b/exec_test.go
index aebb197..8a4b97b 100644
--- a/exec_test.go
+++ b/exec_test.go
@@ -548,6 +548,10 @@ 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 quick break", `{{range .SI}}{{break}}{{.}}{{end}}`, "", tVal, true},
+	{"range break after two", `{{range $i, $x := .SI}}{{if ge $i 2}}{{break}}{{end}}{{.}}{{end}}`, "34", tVal, true},
+	{"range continue", `{{range .SI}}{{continue}}{{.}}{{end}}`, "", tVal, true},
+	{"range continue condition", `{{range .SI}}{{if eq . 3 }}{{continue}}{{end}}{{.}}{{end}}`, "45", tVal, true},
 
 	// While.
 	{"while number", "{{$i := 0}}{{while lt $i 5}}<{{$i}}>{{$i = add $i 1}}{{end}}", "<0><1><2><3><4>", tVal, true},
@@ -557,6 +561,10 @@ var execTests = []execTest{
 	{"while declaration 03", "{{$i := 0}}{{$x := 7}}{{while lt $i 5}}{{$i = add $i 1}}{{$x = 5}}{{end}}{{$x}}", "5", tVal, true},
 	{"while falsey value", "{{while .MSIEmpty}}test{{end}}", "", tVal, true},
 	{"while falsey value else", "{{while .MSIEmpty}}test{{else}}falsey{{end}}", "falsey", tVal, true},
+	{"while quick break", "{{while true}}{{break}}1{{end}}", "", tVal, true},
+	{"while break at three", "{{$i := 0}}{{while lt $i 5}}{{if eq $i 3}}{{break}}{{end}}<{{$i}}>{{$i = add $i 1}}{{end}}", "<0><1><2>", tVal, true},
+	{"while range break", "{{$i := 0}}{{while lt $i 5}}<{{$i}}>{{range .SI}}{{break}}{{end}}{{$i = add $i 1}}{{end}}", "<0><1><2><3><4>", tVal, true},
+	{"while continue", "{{$i := 0}}{{while lt $i 5}}{{$i = add $i 1}}{{continue}}<{{$i}}>{{end}}", "", tVal, true},
 	// should be stopped by MaxOps
 	{"while infinite loop", "{{while true}}{{end}}", "", tVal, false},
 
diff --git a/parse/lex.go b/parse/lex.go
index c79f277..361b5de 100644
--- a/parse/lex.go
+++ b/parse/lex.go
@@ -61,6 +61,8 @@ const (
 	// Keywords appear after all the rest.
 	itemKeyword  // used only to delimit the keywords
 	itemBlock    // block keyword
+	itemBreak    // break keyword
+	itemContinue // continue keyword
 	itemDot      // the cursor, spelled '.'
 	itemDefine   // define keyword
 	itemElse     // else keyword
@@ -76,6 +78,8 @@ const (
 var key = map[string]itemType{
 	".":        itemDot,
 	"block":    itemBlock,
+	"break":    itemBreak,
+	"continue": itemContinue,
 	"define":   itemDefine,
 	"else":     itemElse,
 	"end":      itemEnd,
diff --git a/parse/lex_test.go b/parse/lex_test.go
index 4c3417c..61b16a9 100644
--- a/parse/lex_test.go
+++ b/parse/lex_test.go
@@ -33,7 +33,9 @@ var itemName = map[itemType]string{
 
 	// keywords
 	itemDot:      ".",
+	itemBreak:    "break",
 	itemBlock:    "block",
+	itemContinue: "continue",
 	itemDefine:   "define",
 	itemElse:     "else",
 	itemIf:       "if",
@@ -203,7 +205,7 @@ var lexTests = []lexTest{
 		tRight,
 		tEOF,
 	}},
-	{"keywords", "{{range if else end with while}}", []item{
+	{"keywords", "{{range if else end with break continue while}}", []item{
 		tLeft,
 		mkItem(itemRange, "range"),
 		tSpace,
@@ -215,6 +217,10 @@ var lexTests = []lexTest{
 		tSpace,
 		mkItem(itemWith, "with"),
 		tSpace,
+		mkItem(itemBreak, "break"),
+		tSpace,
+		mkItem(itemContinue, "continue"),
+		tSpace,
 		mkItem(itemWhile, "while"),
 		tRight,
 		tEOF,
diff --git a/parse/node.go b/parse/node.go
index 253b3eb..3b55676 100644
--- a/parse/node.go
+++ b/parse/node.go
@@ -52,8 +52,10 @@ const (
 	NodeText       NodeType = iota // Plain text.
 	NodeAction                     // A non-control action such as a field evaluation.
 	NodeBool                       // A boolean constant.
+	NodeBreak                      // A break action.
 	NodeChain                      // A sequence of field accesses.
 	NodeCommand                    // An element of a pipeline.
+	NodeContinue                   // A continue action.
 	NodeDot                        // The cursor, dot.
 	nodeElse                       // An else action. Not added to tree.
 	nodeEnd                        // An end action. Not added to tree.
@@ -829,6 +831,52 @@ func (w *WhileNode) Copy() Node {
 	return w.tr.newWhile(w.Pos, w.Line, w.Pipe.CopyPipe(), w.List.CopyList(), w.ElseList.CopyList())
 }
 
+// BreakNode represents a {{break}} action.
+type BreakNode struct {
+	NodeType
+	Pos
+	tr *Tree
+}
+
+func (t *Tree) newBreak(pos Pos) *BreakNode {
+	return &BreakNode{NodeType: NodeBreak, Pos: pos, tr: t}
+}
+
+func (b *BreakNode) String() string {
+	return "{{break}}"
+}
+
+func (b *BreakNode) Copy() Node {
+	return b.tr.newBreak(b.Pos)
+}
+
+func (b *BreakNode) tree() *Tree {
+	return b.tr
+}
+
+// ContinueNode represents a {{continue}} action.
+type ContinueNode struct {
+	NodeType
+	Pos
+	tr *Tree
+}
+
+func (t *Tree) newContinue(pos Pos) *ContinueNode {
+	return &ContinueNode{NodeType: NodeContinue, Pos: pos, tr: t}
+}
+
+func (c *ContinueNode) String() string {
+	return "{{continue}}"
+}
+
+func (c *ContinueNode) Copy() Node {
+	return c.tr.newContinue(c.Pos)
+}
+
+func (c *ContinueNode) tree() *Tree {
+	return c.tr
+}
+
 // TemplateNode represents a {{template}} action.
 type TemplateNode struct {
 	NodeType
diff --git a/parse/parse.go b/parse/parse.go
index 46a757d..8cb982c 100644
--- a/parse/parse.go
+++ b/parse/parse.go
@@ -29,6 +29,7 @@ type Tree struct {
 	peekCount int
 	vars      []string // variables defined at the moment.
 	treeSet   map[string]*Tree
+	loopDepth int // nesting level of range/while loops.
 }
 
 // Copy returns a copy of the Tree. Any parsing state is discarded.
@@ -208,6 +209,7 @@ func (t *Tree) startParse(funcs []map[string]interface{}, lex *lexer, treeSet ma
 	t.vars = []string{"$"}
 	t.funcs = funcs
 	t.treeSet = treeSet
+	t.loopDepth = 0
 }
 
 // stopParse terminates parsing.
@@ -264,6 +266,8 @@ func IsEmptyTree(n Node) bool {
 	case *TemplateNode:
 	case *TextNode:
 		return len(bytes.TrimSpace(n.Text)) == 0
+	case *BreakNode:
+	case *ContinueNode:
 	case *WithNode:
 	default:
 		panic("unknown node: " + n.String())
@@ -357,8 +361,12 @@ func (t *Tree) textOrAction() Node {
 // First word could be a keyword such as range.
 func (t *Tree) action() (n Node) {
 	switch token := t.nextNonSpace(); token.typ {
+	case itemBreak:
+		return t.breakControl()
 	case itemBlock:
 		return t.blockControl()
+	case itemContinue:
+		return t.continueControl()
 	case itemElse:
 		return t.elseControl()
 	case itemEnd:
@@ -466,7 +474,13 @@ func (t *Tree) parseControl(allowElseIf bool, context string) (pos Pos, line int
 	defer t.popVars(len(t.vars))
 	pipe = t.pipeline(context)
 	var next Node
+	if context == "range" || context == "while" {
+		t.loopDepth++
+	}
 	list, next = t.itemList()
+	if context == "range" || context == "while" {
+		t.loopDepth--
+	}
 	switch next.Type() {
 	case nodeEnd: //done
 	case nodeElse:
@@ -548,6 +562,16 @@ func (t *Tree) elseControl() Node {
 	return t.newElse(token.pos, token.line)
 }
 
+// Break:
+// 	{{break}}
+// Break keyword is past.
+func (t *Tree) breakControl() Node {
+	if t.loopDepth == 0 {
+		t.errorf("unexpected break outside of loop")
+	}
+	return t.newBreak(t.expect(itemRightDelim, "break").pos)
+}
+
 // Block:
 //	{{block stringValue pipeline}}
 // Block keyword is past.
@@ -575,6 +599,16 @@ func (t *Tree) blockControl() Node {
 	return t.newTemplate(token.pos, token.line, name, pipe)
 }
 
+// Continue
+// 	{{continue}}
+// Continue keyword is past.
+func (t *Tree) continueControl() Node {
+	if t.loopDepth == 0 {
+		t.errorf("unexpected continue outside of range")
+	}
+	return t.newContinue(t.expect(itemRightDelim, "continue").pos)
+}
+
 // Template:
 //	{{template stringValue pipeline}}
 // Template keyword is past. The name must be something that can evaluate
diff --git a/parse/parse_test.go b/parse/parse_test.go
index f0c0c49..11a38fe 100644
--- a/parse/parse_test.go
+++ b/parse/parse_test.go
@@ -230,6 +230,12 @@ var parseTests = []parseTest{
 		`{{range $x := .SI}}{{.}}{{end}}`},
 	{"range 2 vars", "{{range $x, $y := .SI}}{{.}}{{end}}", noError,
 		`{{range $x, $y := .SI}}{{.}}{{end}}`},
+	{"range []int with break", "{{range .SI}}{{break}}{{.}}{{end}}", noError,
+		`{{range .SI}}{{break}}{{.}}{{end}}`},
+	{"range []int with break in else", "{{range .SI}}{{range .SI}}{{.}}{{else}}{{break}}{{end}}{{end}}", noError,
+		`{{range .SI}}{{range .SI}}{{.}}{{else}}{{break}}{{end}}{{end}}`},
+	{"range []int with continue", "{{range .SI}}{{continue}}{{.}}{{end}}", noError,
+		`{{range .SI}}{{continue}}{{.}}{{end}}`},
 	{"simple while", "{{while .X}}hello{{end}}", noError,
 		`{{while .X}}"hello"{{end}}`},
 	{"chained field while", "{{while .X.Y.Z}}hello{{end}}", noError,
@@ -244,6 +250,12 @@ var parseTests = []parseTest{
 		`{{while .B}}{{.}}{{end}}`},
 	{"while 1 var", "{{while $x := .I}}hello{{end}}", noError,
 		`{{while $x := .I}}"hello"{{end}}`},
+	{"while 1 var with break", "{{while $i := .I}}{{break}}{{.}}{{end}}", noError,
+		`{{while $i := .I}}{{break}}{{.}}{{end}}`},
+	{"while 1 var with break in else", "{{while $x := .I}}{{while $y := .I}}{{.}}{{else}}{{break}}{{end}}{{end}}", noError,
+		`{{while $x := .I}}{{while $y := .I}}{{.}}{{else}}{{break}}{{end}}{{end}}`},
+	{"while 1 var with continue", "{{while $i := .I}}{{continue}}{{.}}{{end}}", noError,
+		`{{while $i := .I}}{{continue}}{{.}}{{end}}`},
 	{"constants", "{{range .SI 1 -3.2i true false 'a' nil}}{{end}}", noError,
 		`{{range .SI 1 -3.2i true false 'a' nil}}{{end}}`},
 	{"template", "{{template `x`}}", noError,
@@ -314,6 +326,15 @@ var parseTests = []parseTest{
 	{"empty pipeline", `{{printf "%d" ( ) }}`, hasError, ""},
 	// Missing pipeline in block
 	{"block definition", `{{block "foo"}}hello{{end}}`, hasError, ""},
+	// Invalid loop control
+	{"break outside of loop", `{{break}}`, hasError, ""},
+	{"break in range else, outside of range", `{{range .}}{{.}}{{else}}{{break}{{end}}`, hasError, ""},
+	{"break in while else, outside of while", `{{while $i := .I}}{{.}}{{else}}{{break}}{{end}}`, hasError, ""},
+	{"continue outside of loop", `{{continue}}`, hasError, ""},
+	{"continue in range else, outside of range", `{{range .}}{{.}}{{else}}{{continue}}{{end}}`, hasError, ""},
+	{"continue in while else, outside of while", `{{while true}}{{.}}{{else}}{{continue}}{{end}}`, hasError, ""},
+	{"additional break data", `{{range .}}{{break label}}{{end}}`, hasError, ""},
+	{"additional continue data", `{{range .}}{{continue label}}{{end}}`, hasError, ""},
 }
 
 var builtins = map[string]interface{}{