diff --git a/ast/helpers.go b/ast/helpers.go
index ceadd00..648140f 100644
--- a/ast/helpers.go
+++ b/ast/helpers.go
@@ -30,16 +30,16 @@ func AppendNode(n1, n2 parsley.Node) parsley.Node {
 // Walkable is a generic interface to allow to apply a function on the node
 // The Walk function should return true if the walk should be interrupted
 type Walkable interface {
-	Walk(f func(i int, n parsley.Node) bool) bool
+	Walk(f func(n parsley.Node) bool) bool
 }
 
 // WalkNode applies the given function to the node
-func WalkNode(node parsley.Node, f func(i int, n parsley.Node) bool) bool {
+func WalkNode(node parsley.Node, f func(n parsley.Node) bool) bool {
 	switch n := node.(type) {
 	case Walkable:
 		return n.Walk(f)
 	default:
-		return f(0, node)
+		return f(node)
 	}
 }
 
diff --git a/ast/helpers_test.go b/ast/helpers_test.go
index 4aa56d4..11584b5 100644
--- a/ast/helpers_test.go
+++ b/ast/helpers_test.go
@@ -64,8 +64,7 @@ var _ = Describe("WalkNode", func() {
 
 		nl := ast.NodeList([]parsley.Node{n1, n2})
 		called := []parsley.Node{}
-		f := func(i int, n parsley.Node) bool {
-			Expect(i).To(Equal(len(called)))
+		f := func(n parsley.Node) bool {
 			called = append(called, n)
 			return false
 		}
@@ -80,8 +79,7 @@ var _ = Describe("WalkNode", func() {
 		node := &parsleyfakes.FakeNode{}
 		node.TokenReturns("TEST")
 		called := []parsley.Node{}
-		f := func(i int, n parsley.Node) bool {
-			Expect(i).To(Equal(len(called)))
+		f := func(n parsley.Node) bool {
 			called = append(called, n)
 			return false
 		}
diff --git a/ast/node.go b/ast/node.go
index 63845df..b1b417a 100644
--- a/ast/node.go
+++ b/ast/node.go
@@ -172,9 +172,9 @@ func (n *NonTerminalNode) String() string {
 }
 
 // Walk runs the given function on all child nodes
-func (n *NonTerminalNode) Walk(f func(i int, n parsley.Node) bool) bool {
-	for i, node := range n.children {
-		if f(i, node) {
+func (n *NonTerminalNode) Walk(f func(n parsley.Node) bool) bool {
+	for _, n := range n.children {
+		if WalkNode(n, f) {
 			return true
 		}
 	}
@@ -231,9 +231,9 @@ func (nl *NodeList) Append(node parsley.Node) {
 }
 
 // Walk runs the given function on all nodes
-func (nl NodeList) Walk(f func(i int, n parsley.Node) bool) bool {
-	for i, node := range nl {
-		if f(i, node) {
+func (nl NodeList) Walk(f func(n parsley.Node) bool) bool {
+	for _, n := range nl {
+		if WalkNode(n, f) {
 			return true
 		}
 	}
diff --git a/ast/node_test.go b/ast/node_test.go
index 1a8a594..a83fc30 100644
--- a/ast/node_test.go
+++ b/ast/node_test.go
@@ -205,8 +205,7 @@ var _ = Describe("NonTerminalNode", func() {
 			Describe("Walk", func() {
 				It("should call the function with all children", func() {
 					called := []parsley.Node{}
-					f := func(i int, n parsley.Node) bool {
-						Expect(i).To(Equal(len(called)))
+					f := func(n parsley.Node) bool {
 						called = append(called, n)
 						return false
 					}
@@ -219,10 +218,9 @@ var _ = Describe("NonTerminalNode", func() {
 
 				It("should stop if the function returns with true", func() {
 					called := []parsley.Node{}
-					f := func(i int, n parsley.Node) bool {
-						Expect(i).To(Equal(len(called)))
+					f := func(n parsley.Node) bool {
 						called = append(called, n)
-						return i == 0
+						return true
 					}
 
 					res := node.Walk(f)
@@ -417,8 +415,7 @@ var _ = Describe("NodeList", func() {
 		It("should call the function with all children", func() {
 			nl = ast.NodeList([]parsley.Node{n1, n2})
 			called := []parsley.Node{}
-			f := func(i int, n parsley.Node) bool {
-				Expect(i).To(Equal(len(called)))
+			f := func(n parsley.Node) bool {
 				called = append(called, n)
 				return false
 			}
@@ -432,16 +429,15 @@ var _ = Describe("NodeList", func() {
 		It("should stop if the function returns with true", func() {
 			nl = ast.NodeList([]parsley.Node{n1, n2, n3})
 			called := []parsley.Node{}
-			f := func(i int, n parsley.Node) bool {
-				Expect(i).To(Equal(len(called)))
+			f := func(n parsley.Node) bool {
 				called = append(called, n)
-				return i > 0
+				return true
 			}
 
 			res := nl.Walk(f)
 			Expect(res).To(BeTrue())
 
-			Expect(called).To(Equal([]parsley.Node{n1, n2}))
+			Expect(called).To(Equal([]parsley.Node{n1}))
 		})
 	})
 })