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.

6 files changed, 104 insertions, 80 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)
 	}
 }
diff --git a/pkg/compiler/compiler.go b/pkg/compiler/compiler.go
index 3901747a6..6019bda94 100644
--- a/pkg/compiler/compiler.go
+++ b/pkg/compiler/compiler.go
@@ -47,7 +47,7 @@ func Compile(desc *ast.Description, consts map[string]uint64, target *targets.Ta
 		eh = ast.LoggingHandler
 	}
 	comp := &compiler{
-		desc:         ast.Clone(desc),
+		desc:         desc.Clone(),
 		target:       target,
 		eh:           eh,
 		ptrSize:      target.PtrSize,
diff --git a/pkg/compiler/consts.go b/pkg/compiler/consts.go
index 338dea4b1..f2e4d4850 100644
--- a/pkg/compiler/consts.go
+++ b/pkg/compiler/consts.go
@@ -43,8 +43,8 @@ func ExtractConsts(desc *ast.Description, target *targets.Target, eh0 ast.ErrorH
 	incdirMap := make(map[string]bool)
 	constMap := make(map[string]bool)
 
-	ast.Walk(desc, func(n1 ast.Node) {
-		switch n := n1.(type) {
+	desc.Walk(ast.Recursive(func(n0 ast.Node) {
+		switch n := n0.(type) {
 		case *ast.Include:
 			file := n.File.Value
 			if includeMap[file] {
@@ -85,7 +85,7 @@ func ExtractConsts(desc *ast.Description, target *targets.Target, eh0 ast.ErrorH
 		case *ast.Int:
 			constMap[n.Ident] = true
 		}
-	})
+	}))
 
 	if errors != 0 {
 		return nil
@@ -179,7 +179,7 @@ func (comp *compiler) patchConsts(consts map[string]uint64) {
 		case *ast.Resource, *ast.Struct, *ast.Call, *ast.TypeDef:
 			// Walk whole tree and replace consts in Int's and Type's.
 			missing := ""
-			ast.WalkNode(decl, func(n0 ast.Node) {
+			decl.Walk(ast.Recursive(func(n0 ast.Node) {
 				switch n := n0.(type) {
 				case *ast.Int:
 					comp.patchIntConst(n.Pos, &n.Value, &n.Ident, consts, &missing)
@@ -193,7 +193,7 @@ func (comp *compiler) patchConsts(consts map[string]uint64) {
 						}
 					}
 				}
-			})
+			}))
 			if missing == "" {
 				top = append(top, decl)
 				continue