Engage!

1 files changed, 79 insertions, 0 deletions

diff --git a/vendor/github.com/mitjafelicijan/go-tree-sitter/iter.go b/vendor/github.com/mitjafelicijan/go-tree-sitter/iter.go
new file mode 100644
index 0000000..5b7d903
--- /dev/null
+++ b/vendor/github.com/mitjafelicijan/go-tree-sitter/iter.go
@@ -0,0 +1,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
+		}
+	}
+}