1package sitter
 2
 3import "io"
 4
 5type IterMode int
 6
 7const (
 8	DFSMode IterMode = iota
 9	BFSMode
10)
11
12// Iterator for a tree of nodes
13type Iterator struct {
14	named bool
15	mode  IterMode
16
17	nodesToVisit []*Node
18}
19
20// NewIterator takes a node and mode (DFS/BFS) and returns iterator over children of the node
21func NewIterator(n *Node, mode IterMode) *Iterator {
22	return &Iterator{
23		named:        false,
24		mode:         mode,
25		nodesToVisit: []*Node{n},
26	}
27}
28
29// NewNamedIterator takes a node and mode (DFS/BFS) and returns iterator over named children of the node
30func NewNamedIterator(n *Node, mode IterMode) *Iterator {
31	return &Iterator{
32		named:        true,
33		mode:         mode,
34		nodesToVisit: []*Node{n},
35	}
36}
37
38func (iter *Iterator) Next() (*Node, error) {
39	if len(iter.nodesToVisit) == 0 {
40		return nil, io.EOF
41	}
42
43	var n *Node
44	n, iter.nodesToVisit = iter.nodesToVisit[0], iter.nodesToVisit[1:]
45
46	var children []*Node
47	if iter.named {
48		for i := 0; i < int(n.NamedChildCount()); i++ {
49			children = append(children, n.NamedChild(i))
50		}
51	} else {
52		for i := 0; i < int(n.ChildCount()); i++ {
53			children = append(children, n.Child(i))
54		}
55	}
56
57	switch iter.mode {
58	case DFSMode:
59		iter.nodesToVisit = append(children, iter.nodesToVisit...)
60	case BFSMode:
61		iter.nodesToVisit = append(iter.nodesToVisit, children...)
62	default:
63		panic("not implemented")
64	}
65	return n, nil
66}
67
68func (iter *Iterator) ForEach(fn func(*Node) error) error {
69	for {
70		n, err := iter.Next()
71		if err != nil {
72			return err
73		}
74		err = fn(n)
75		if err != nil {
76			return err
77		}
78	}
79}