From ff2d79c2cb08e0b410143e26f71a0e80e49d5b20 Mon Sep 17 00:00:00 2001
From: Pierre Neidhardt <ambrevar@gmail.com>
Date: Wed, 3 Aug 2016 12:25:38 +0530
Subject: [PATCH] Fix walk so that it does not step outside the root node

When passed a non-container root node, or when the visitor whould return
SkipChildren for the root node, the former algorithm would go on walking
down the rest of the tree beyond the root.

This commit also removes unneeded functions and variables.
---
 node.go | 37 ++++++++++++-------------------------
 1 file changed, 12 insertions(+), 25 deletions(-)

diff --git a/node.go b/node.go
index bcd2913..bb91123 100644
--- a/node.go
+++ b/node.go
@@ -270,16 +270,16 @@ type NodeVisitor func(node *Node, entering bool) WalkStatus
 
 // Walk is a convenience method that instantiates a walker and starts a
 // traversal of subtree rooted at n.
-func (n *Node) Walk(visitor NodeVisitor) {
-	walker := newNodeWalker(n)
-	node, entering := walker.next()
-	for node != nil {
-		status := visitor(node, entering)
+func (root *Node) Walk(visitor NodeVisitor) {
+	w := newNodeWalker(root)
+	for w.current != nil {
+		status := visitor(w.current, w.entering)
 		switch status {
 		case GoToNext:
-			node, entering = walker.next()
+			w.next()
 		case SkipChildren:
-			node, entering = walker.resumeAt(node, false)
+			w.entering = false
+			w.next()
 		case Terminate:
 			return
 		}
@@ -295,18 +295,15 @@ type nodeWalker struct {
 func newNodeWalker(root *Node) *nodeWalker {
 	return &nodeWalker{
 		current:  root,
-		root:     nil,
+		root:     root,
 		entering: true,
 	}
 }
 
-func (nw *nodeWalker) next() (*Node, bool) {
-	if nw.current == nil {
-		return nil, false
-	}
-	if nw.root == nil {
-		nw.root = nw.current
-		return nw.current, nw.entering
+func (nw *nodeWalker) next() {
+	if !nw.entering && nw.current == nw.root {
+		nw.current = nil
+		return
 	}
 	if nw.entering && nw.current.isContainer() {
 		if nw.current.FirstChild != nil {
@@ -322,16 +319,6 @@ func (nw *nodeWalker) next() (*Node, bool) {
 		nw.current = nw.current.Next
 		nw.entering = true
 	}
-	if nw.current == nw.root {
-		return nil, false
-	}
-	return nw.current, nw.entering
-}
-
-func (nw *nodeWalker) resumeAt(node *Node, entering bool) (*Node, bool) {
-	nw.current = node
-	nw.entering = entering
-	return nw.next()
 }
 
 func dump(ast *Node) {