pkg/ast: refactor Walk

Refactor Walk so that it's possible to abort or wrap walk of child nodes.
Will be needed for future changes.

4 files changed, 98 insertions, 74 deletions

diff --git a/pkg/ast/ast.go b/pkg/ast/ast.go
index 9e87a1b81..454f28b37 100644
--- a/pkg/ast/ast.go
+++ b/pkg/ast/ast.go
@@ -24,6 +24,9 @@ type Node interface {
 	// If newPos is not zero, sets Pos of all nodes to newPos.
 	// If newPos is zero, Pos of nodes is left intact.
 	Clone(newPos Pos) Node
+	// Walk calls callback cb for all child nodes of this node.
+	// Note: it's not recursive. Use Recursive helper for recursive walk.
+	Walk(cb func(Node))
 }
 
 // Top-level AST nodes:
diff --git a/pkg/ast/clone.go b/pkg/ast/clone.go
index 044b6c17d..dcd715c0a 100644
--- a/pkg/ast/clone.go
+++ b/pkg/ast/clone.go
@@ -3,7 +3,7 @@
 
 package ast
 
-func Clone(desc *Description) *Description {
+func (desc *Description) Clone() *Description {
 	desc1 := &Description{}
 	for _, n := range desc.Nodes {
 		desc1.Nodes = append(desc1.Nodes, n.Clone(Pos{}))
diff --git a/pkg/ast/parser_test.go b/pkg/ast/parser_test.go
index c7018b243..f1458e3e7 100644
--- a/pkg/ast/parser_test.go
+++ b/pkg/ast/parser_test.go
@@ -48,7 +48,7 @@ func TestParseAll(t *testing.T) {
 					t.Fatalf("formatting changed code:\n%#v\nvs:\n%#v", n1, n2)
 				}
 			}
-			data3 := Format(Clone(desc))
+			data3 := Format(desc.Clone())
 			if !bytes.Equal(data, data3) {
 				t.Fatalf("Clone lost data")
 			}
diff --git a/pkg/ast/walk.go b/pkg/ast/walk.go
index 7bc34410e..fd5065013 100644
--- a/pkg/ast/walk.go
+++ b/pkg/ast/walk.go
@@ -3,81 +3,102 @@
 
 package ast
 
-import (
-	"fmt"
-)
-
-// Walk calls callback cb for every node in AST.
-func Walk(desc *Description, cb func(n Node)) {
+// Walk calls callback cb for every top-level node in description.
+// Note: it's not recursive. Use Recursive helper for recursive walk.
+func (desc *Description) Walk(cb func(Node)) {
 	for _, n := range desc.Nodes {
-		WalkNode(n, cb)
+		cb(n)
+	}
+}
+
+func Recursive(cb func(Node)) func(Node) {
+	var rec func(Node)
+	rec = func(n Node) {
+		cb(n)
+		n.Walk(rec)
+	}
+	return rec
+}
+
+func (n *NewLine) Walk(cb func(Node)) {}
+func (n *Comment) Walk(cb func(Node)) {}
+func (n *Ident) Walk(cb func(Node))   {}
+func (n *String) Walk(cb func(Node))  {}
+func (n *Int) Walk(cb func(Node))     {}
+
+func (n *Include) Walk(cb func(Node)) {
+	cb(n.File)
+}
+
+func (n *Incdir) Walk(cb func(Node)) {
+	cb(n.Dir)
+}
+
+func (n *Define) Walk(cb func(Node)) {
+	cb(n.Name)
+	cb(n.Value)
+}
+
+func (n *Resource) Walk(cb func(Node)) {
+	cb(n.Name)
+	cb(n.Base)
+	for _, v := range n.Values {
+		cb(v)
+	}
+}
+
+func (n *TypeDef) Walk(cb func(Node)) {
+	cb(n.Name)
+	cb(n.Type)
+}
+
+func (n *Call) Walk(cb func(Node)) {
+	cb(n.Name)
+	for _, f := range n.Args {
+		cb(f)
+	}
+	if n.Ret != nil {
+		cb(n.Ret)
+	}
+}
+
+func (n *Struct) Walk(cb func(Node)) {
+	cb(n.Name)
+	for _, f := range n.Fields {
+		cb(f)
+	}
+	for _, a := range n.Attrs {
+		cb(a)
+	}
+	for _, c := range n.Comments {
+		cb(c)
+	}
+}
+
+func (n *IntFlags) Walk(cb func(Node)) {
+	cb(n.Name)
+	for _, v := range n.Values {
+		cb(v)
+	}
+}
+
+func (n *StrFlags) Walk(cb func(Node)) {
+	cb(n.Name)
+	for _, v := range n.Values {
+		cb(v)
+	}
+}
+
+func (n *Type) Walk(cb func(Node)) {
+	for _, t := range n.Args {
+		cb(t)
 	}
 }
 
-func WalkNode(n0 Node, cb func(n Node)) {
-	cb(n0)
-	switch n := n0.(type) {
-	case *NewLine:
-	case *Comment:
-	case *Include:
-		WalkNode(n.File, cb)
-	case *Incdir:
-		WalkNode(n.Dir, cb)
-	case *Define:
-		WalkNode(n.Name, cb)
-		WalkNode(n.Value, cb)
-	case *Resource:
-		WalkNode(n.Name, cb)
-		WalkNode(n.Base, cb)
-		for _, v := range n.Values {
-			WalkNode(v, cb)
-		}
-	case *TypeDef:
-		WalkNode(n.Name, cb)
-		WalkNode(n.Type, cb)
-	case *Call:
-		WalkNode(n.Name, cb)
-		for _, f := range n.Args {
-			WalkNode(f, cb)
-		}
-		if n.Ret != nil {
-			WalkNode(n.Ret, cb)
-		}
-	case *Struct:
-		WalkNode(n.Name, cb)
-		for _, f := range n.Fields {
-			WalkNode(f, cb)
-		}
-		for _, a := range n.Attrs {
-			WalkNode(a, cb)
-		}
-		for _, c := range n.Comments {
-			WalkNode(c, cb)
-		}
-	case *IntFlags:
-		WalkNode(n.Name, cb)
-		for _, v := range n.Values {
-			WalkNode(v, cb)
-		}
-	case *StrFlags:
-		WalkNode(n.Name, cb)
-		for _, v := range n.Values {
-			WalkNode(v, cb)
-		}
-	case *Ident:
-	case *String:
-	case *Int:
-	case *Type:
-		for _, t := range n.Args {
-			WalkNode(t, cb)
-		}
-	case *Field:
-		WalkNode(n.Name, cb)
-		WalkNode(n.Type, cb)
-		for _, c := range n.Comments {
-			WalkNode(c, cb)
-		}
-	default:
-		panic(fmt.Sprintf("unknown AST node: %#v", n))
+func (n *Field) Walk(cb func(Node)) {
+	cb(n.Name)
+	cb(n.Type)
+	for _, c := range n.Comments {
+		cb(c)
 	}
 }