From 9da90c5929930e100e0e36619d52e1afa0d6f7bd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vytautas=20=C5=A0altenis?= <vytas@rtfb.lt>
Date: Tue, 5 Apr 2016 11:10:46 +0300
Subject: [PATCH] Allow NodeVisitor to have some control over traversal

Make NodeVisitor return status and decide upon it which node to go to
next. So far, this allows to skip subtrees and quit early.
---
 html.go     | 10 ++++++----
 markdown.go |  9 ++++++---
 node.go     | 27 +++++++++++++++++++++++----
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/html.go b/html.go
index 567d43a..10d6159 100644
--- a/html.go
+++ b/html.go
@@ -1100,7 +1100,7 @@ func (r *HTML) cr(w io.Writer) {
 	}
 }
 
-func (r *HTML) RenderNode(w io.Writer, node *Node, entering bool) {
+func (r *HTML) RenderNode(w io.Writer, node *Node, entering bool) WalkStatus {
 	attrs := []string{}
 	switch node.Type {
 	case Text:
@@ -1392,6 +1392,7 @@ func (r *HTML) RenderNode(w io.Writer, node *Node, entering bool) {
 	default:
 		panic("Unknown node type " + node.Type.String())
 	}
+	return GoToNext
 }
 
 func (r *HTML) writeDocumentHeader(w *bytes.Buffer, sr *SPRenderer) {
@@ -1447,7 +1448,7 @@ func (r *HTML) Render(ast *Node) []byte {
 	//dump(ast)
 	// Run Smartypants if it's enabled or simply escape text if not
 	sr := NewSmartypantsRenderer(r.extensions)
-	ast.Walk(func(node *Node, entering bool) {
+	ast.Walk(func(node *Node, entering bool) WalkStatus {
 		if node.Type == Text {
 			if r.extensions&Smartypants != 0 {
 				node.Literal = sr.Process(node.Literal)
@@ -1455,11 +1456,12 @@ func (r *HTML) Render(ast *Node) []byte {
 				node.Literal = esc(node.Literal, false)
 			}
 		}
+		return GoToNext
 	})
 	var buff bytes.Buffer
 	r.writeDocumentHeader(&buff, sr)
-	ast.Walk(func(node *Node, entering bool) {
-		r.RenderNode(&buff, node, entering)
+	ast.Walk(func(node *Node, entering bool) WalkStatus {
+		return r.RenderNode(&buff, node, entering)
 	})
 	r.writeDocumentFooter(&buff)
 	return buff.Bytes()
diff --git a/markdown.go b/markdown.go
index d30bdc5..97c2211 100644
--- a/markdown.go
+++ b/markdown.go
@@ -452,12 +452,13 @@ func Parse(input []byte, opts Options) *Node {
 		p.finalize(p.tip)
 	}
 	// Walk the tree again and process inline markdown in each block
-	p.doc.Walk(func(node *Node, entering bool) {
+	p.doc.Walk(func(node *Node, entering bool) WalkStatus {
 		if node.Type == Paragraph || node.Type == Header || node.Type == TableCell {
 			p.currBlock = node
 			p.inline(node.content)
 			node.content = nil
 		}
+		return GoToNext
 	})
 	p.parseRefsToAST()
 	p.generateTOC()
@@ -477,7 +478,7 @@ func (p *parser) generateTOC() {
 	var lastItem *Node
 	headerCount := 0
 	var currentLevel uint32
-	p.doc.Walk(func(node *Node, entering bool) {
+	p.doc.Walk(func(node *Node, entering bool) WalkStatus {
 		if entering && node.Type == Header {
 			if node.Level > currentLevel {
 				currentLevel++
@@ -504,6 +505,7 @@ func (p *parser) generateTOC() {
 			lastItem.appendChild(anchorNode)
 			anchorNode.appendChild(text(node.FirstChild.Literal))
 		}
+		return GoToNext
 	})
 	firstChild := p.doc.FirstChild
 	// Insert TOC only if there is anything to insert
@@ -558,12 +560,13 @@ func (p *parser) parseRefsToAST() {
 	finalizeList(block)
 	p.tip = above
 	finalizeHtmlBlock(p.addBlock(HTMLBlock, []byte("</div>")))
-	block.Walk(func(node *Node, entering bool) {
+	block.Walk(func(node *Node, entering bool) WalkStatus {
 		if node.Type == Paragraph || node.Type == Header {
 			p.currBlock = node
 			p.inline(node.content)
 			node.content = nil
 		}
+		return GoToNext
 	})
 }
 
diff --git a/node.go b/node.go
index c565d5c..a78449a 100644
--- a/node.go
+++ b/node.go
@@ -230,17 +230,35 @@ func (n *Node) canContain(t NodeType) bool {
 	return false
 }
 
+// WalkStatus allows NodeVisitor to have some control over the tree traversal.
+// It is returned from NodeVisitor and different values allow Node.Walk to
+// decide which node to go to next.
+type WalkStatus int
+
+const (
+	GoToNext     WalkStatus = iota // The default traversal of every node.
+	SkipChildren                   // Skips all children of current node.
+	Terminate                      // Terminates the traversal.
+)
+
 // NodeVisitor is a callback to be called when traversing the syntax tree.
 // Called twice for every node: once with entering=true when the branch is
 // first visited, then with entering=false after all the children are done.
-type NodeVisitor func(node *Node, entering bool)
+type NodeVisitor func(node *Node, entering bool) WalkStatus
 
 func (root *Node) Walk(visitor NodeVisitor) {
 	walker := NewNodeWalker(root)
 	node, entering := walker.next()
 	for node != nil {
-		visitor(node, entering)
-		node, entering = walker.next()
+		status := visitor(node, entering)
+		switch status {
+		case GoToNext:
+			node, entering = walker.next()
+		case SkipChildren:
+			node, entering = walker.resumeAt(node, false)
+		case Terminate:
+			return
+		}
 	}
 }
 
@@ -286,9 +304,10 @@ func (nw *NodeWalker) next() (*Node, bool) {
 	return nw.current, nw.entering
 }
 
-func (nw *NodeWalker) resumeAt(node *Node, entering bool) {
+func (nw *NodeWalker) resumeAt(node *Node, entering bool) (*Node, bool) {
 	nw.current = node
 	nw.entering = entering
+	return nw.next()
 }
 
 func dump(ast *Node) {