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