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

@@ -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()

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
 	})
 }
 

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) {