1package ast
2
3import "fmt"
4
5// WalkFunc describes a function to be called for each node during a Walk. The
6// returned node can be used to rewrite the AST. Walking stops the returned
7// bool is false.
8type WalkFunc func(Node) (Node, bool)
9
10// Walk traverses an AST in depth-first order: It starts by calling fn(node);
11// node must not be nil. If fn returns true, Walk invokes fn recursively for
12// each of the non-nil children of node, followed by a call of fn(nil). The
13// returned node of fn can be used to rewrite the passed node to fn.
14func Walk(node Node, fn WalkFunc) Node {
15	rewritten, ok := fn(node)
16	if !ok {
17		return rewritten
18	}
19
20	switch n := node.(type) {
21	case *File:
22		n.Node = Walk(n.Node, fn)
23	case *ObjectList:
24		for i, item := range n.Items {
25			n.Items[i] = Walk(item, fn).(*ObjectItem)
26		}
27	case *ObjectKey:
28		// nothing to do
29	case *ObjectItem:
30		for i, k := range n.Keys {
31			n.Keys[i] = Walk(k, fn).(*ObjectKey)
32		}
33
34		if n.Val != nil {
35			n.Val = Walk(n.Val, fn)
36		}
37	case *LiteralType:
38		// nothing to do
39	case *ListType:
40		for i, l := range n.List {
41			n.List[i] = Walk(l, fn)
42		}
43	case *ObjectType:
44		n.List = Walk(n.List, fn).(*ObjectList)
45	default:
46		// should we panic here?
47		fmt.Printf("unknown type: %T\n", n)
48	}
49
50	fn(nil)
51	return rewritten
52}
53