-
Notifications
You must be signed in to change notification settings - Fork 38
/
ast.go
117 lines (103 loc) · 2.16 KB
/
ast.go
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package lang
import (
"errors"
"fmt"
"io"
"strconv"
"strings"
"github.com/klothoplatform/klotho/pkg/multierr"
sitter "github.com/smacker/go-tree-sitter"
)
func WriteAST(node *sitter.Node, out io.Writer) error {
w := ASTWriter{
node: node,
out: out,
}
w.WriteAST()
return w.Err()
}
type ASTWriter struct {
out io.Writer
node *sitter.Node
errors multierr.Error
}
func (w *ASTWriter) Err() error {
return w.errors.ErrOrNil()
}
func (w *ASTWriter) indent(n *sitter.Node) {
for i := 0; i < w.level(n); i++ {
w.write(" ")
}
}
func (w *ASTWriter) writeLine(n *sitter.Node, s string) {
w.indent(n)
w.write(s)
if !strings.HasSuffix(s, "\n") {
_, err := w.out.Write([]byte("\n"))
if err != nil {
w.errors.Append(err)
}
}
}
func (w *ASTWriter) write(s string) {
var err error
if out, ok := w.out.(io.StringWriter); ok {
_, err = out.WriteString(s)
} else {
_, err = w.out.Write([]byte(s))
}
if err != nil {
w.errors.Append(err)
}
}
func (w *ASTWriter) nodeName(n *sitter.Node) string {
if n.Parent() == nil {
return ""
}
c := sitter.NewTreeCursor(n.Parent())
defer c.Close()
if !c.GoToFirstChild() {
return ""
}
for c.CurrentNode() != n {
if !c.GoToNextSibling() {
return ""
}
}
return c.CurrentFieldName()
}
func (w *ASTWriter) level(n *sitter.Node) int {
l := 0
for n.Parent() != nil && n.Parent() != w.node {
l++
n = n.Parent()
}
return l
}
func (w *ASTWriter) WriteAST() {
iterator := sitter.NewIterator(w.node, sitter.DFSMode)
w.writeLine(w.node, w.node.String())
w.writeLine(w.node, "=====")
//iterates over the entire AST in DFSMode
err := iterator.ForEach(func(n *sitter.Node) error {
name := w.nodeName(n)
if n.NamedChildCount() > 0 {
w.writeLine(n, fmt.Sprintf("%s: (%s)", name, n.Type()))
} else {
if n.Parent() != nil && n.Parent().NamedChildCount() > 0 {
content := n.Content()
if n.Type() == content {
w.writeLine(n, fmt.Sprintf("%s = %s", name, content))
} else {
w.writeLine(n, fmt.Sprintf("%s: (%s) = %s", name, n.Type(), strconv.Quote(content)))
}
} else {
return nil
}
}
return nil
})
if err != nil && !errors.Is(err, io.EOF) {
w.errors.Append(err)
}
}