/
ast.go
executable file
·162 lines (142 loc) · 4.7 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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
// Copyright (c) 2024 Dez Little <deslittle@gmail.com>
// All rights reserved. Use of this source code is governed by a LGPL v3
// license that can be found in the LICENSE file.
// ast.go implements an Abstract Syntax Tree for use by the DSL parser.
// The user tells the AST to add nodes and tokens inside the user parse
// function using three basic functions; p.AddNode(), p.AddToken() and
// p.WalkUp(). AST node types are defined by the user.
//
// The AST is made up of nodes (more accurately Node pointers), each of
// which contains a slice of Node children and a reference to it's parent.
// The nodes are made available to the user so they can walk up and down
// the tree once it is returned from the parser.
package dsl
import (
"fmt"
)
// RootNode is the entry point to the tree. curNode is used internally
// to keep track of where the next node should be added.
type AST struct {
ns NodeSet `json:"-"`
RootNode *Node `json:"root"`
curNode *Node `json:"-"`
}
// A Node can contain multiple Tokens which can be useful if the user knows how
// many Tokens belong to a particular Node type. Otherwise, the user should only
// add one token per node.
type Node struct {
Type NodeType `json:"type"`
Tokens []Token `json:"tokens"`
Parent *Node `json:"-"`
Children []Node `json:"children"`
}
type NodeSet map[NodeType]int
type NodeType string
func (nt NodeType) String() string {
return string(nt)
}
const (
NODE_ROOT NodeType = "ROOT"
)
func NewNodeSet(userTypes ...NodeType) NodeSet {
ns := make(map[NodeType]int)
ns[NODE_ROOT] = 1
for i, id := range userTypes {
ns[id] = i + 2
}
return ns
}
// newAST returns a new instance of AST. The RootNode has the
// builtin node type AST_ROOT.
func newAST(ns NodeSet) AST {
rootNode := &Node{Type: NODE_ROOT}
return AST{ns: ns, RootNode: rootNode, curNode: rootNode}
}
// ---------------------------------------------------------------------------------------------------------
// The AST is made up entirely of Node instances connected in a tree pattern.
// The AST ensures the Parent and Children references are set correctly as it
// is being constructed.
//
// Prints the entire AST tree. It does so by recursively calling Print() on
// each node in the tree in a depth first approach.
//
// Inspect traverses an AST in depth-first order: It starts by calling
// f(node);
func (a *AST) Inspect(fn func(*Node)) {
visit(a.RootNode, fn)
}
func visit(node *Node, fn func(*Node)) {
for _, child := range node.Children {
visit(&child, fn)
}
fn(node)
}
// Prints the entire AST tree. It does so by recursively calling Print() on
// each node in the tree in a depth first approach.
func (a *AST) Print() {
a.RootNode.Print("", true)
fmt.Println()
}
// Called by Parser.AddNode() in the user parse function. Creates a new node and
// builds the two-way reference to its parent. Also moves the AST curNode
// down the tree to the new node.
func (a *AST) addNode(nt NodeType) {
a.curNode.Children = append(a.curNode.Children, Node{Type: nt, Parent: a.curNode})
a.curNode = &a.curNode.Children[len(a.curNode.Children)-1]
}
// Called by Parser.AddToken() in the user parse function. Adds a token to the
// end of the Token slice belonging to the current node.
//
// If Parser.AddToken() is called without any tokens available on the Parser.toks buffer
// the call to AddToken will be logged but no tokens will be added to the node.
func (a *AST) addToken(toks []Token) {
if toks != nil {
tokens := append(a.curNode.Tokens, toks...)
a.curNode.Tokens = tokens
}
}
// Called by Parser.WalkUp() in the user parse function. Moves the AST
// curNode to its parent.
func (a *AST) walkUp() {
if a.ns[a.curNode.Type] != a.ns[NODE_ROOT] {
a.curNode = a.curNode.Parent
}
}
// Print is a recursive function that keeps track of where in the tree the node
// belongs to so it can print a pretty prefix. The prefix indicates how deep the
// node is and if it is the last node at that level.
//
// A user can print the entire tree using AST.Print() or only print a sub-branch
// by calling Print() on any node in the tree.
func (n *Node) Print(prefix string, isTail bool) {
fmt.Printf("\n%v", prefix)
if isTail {
fmt.Printf("└── ")
} else {
fmt.Printf("├── ")
}
fmt.Printf("%v - ", n.Type)
for _, token := range n.Tokens {
for _, rn := range token.Literal {
fmt.Print(string(rn))
}
fmt.Print(", ")
}
numNodes := len(n.Children)
if numNodes > 0 {
for _, node := range n.Children[:numNodes-1] {
if isTail {
node.Print(prefix+" ", false)
} else {
node.Print(prefix+"│ ", false)
}
}
}
if numNodes > 0 {
if isTail {
n.Children[numNodes-1].Print(prefix+" ", true)
} else {
n.Children[numNodes-1].Print(prefix+"│ ", true)
}
}
}