Skip to content

Commit 388e7d7

Browse files
feat: visit API (#26)
1 parent aca7aac commit 388e7d7

File tree

4 files changed

+221
-0
lines changed

4 files changed

+221
-0
lines changed

README.md

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ Tools and IntelliSense for GLSL and WGSL.
1313
- [Minify](#minify)
1414
- [Parse](#parse)
1515
- [Generate](#generate)
16+
- [Visit](#visit)
1617
- [AST](#ast)
1718
- [Node Objects](#node-objects)
1819
- [Identifier](#identifier)
@@ -265,6 +266,29 @@ const code: string = generate(program: Program, {
265266
})
266267
```
267268

269+
## Visit
270+
271+
Recurses through an [AST](https://en.wikipedia.org/wiki/Abstract_syntax_tree), calling a visitor object on matching nodes.
272+
273+
```ts
274+
visit(
275+
program: Program,
276+
visitors: {
277+
Program: {
278+
enter(node, ancestors) {
279+
// Called before any descendant nodes are processed
280+
},
281+
exit(node, ancestors) {
282+
// Called after all nodes are processed
283+
}
284+
},
285+
Identifier(node, ancestors) {
286+
// Called before any descendant nodes are processed (alias to enter)
287+
}
288+
} satisfies Visitors
289+
)
290+
```
291+
268292
## AST
269293

270294
An [Abstract Syntax Tree](https://en.wikipedia.org/wiki/Abstract_syntax_tree) loosely based on [ESTree](https://github.com/estree/estree) for GLSL and WGSL grammars.

src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,4 @@ export * from './generator.js'
44
export * from './minifier.js'
55
export * from './parser.js'
66
export * from './tokenizer.js'
7+
export * from './visitor.js'

src/visitor.ts

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
import { type AST } from './ast'
2+
3+
export type Visitors = Partial<{
4+
[K in AST['type']]:
5+
| ((node: Extract<AST, { type: K }>, ancestors: AST[]) => void)
6+
| {
7+
enter?(node: Extract<AST, { type: K }>, ancestors: AST[]): void
8+
exit?(node: Extract<AST, { type: K }>, ancestors: AST[]): void
9+
}
10+
}>
11+
12+
/**
13+
* Recurses through an [AST](https://en.wikipedia.org/wiki/Abstract_syntax_tree), calling a visitor object on matching nodes.
14+
*/
15+
export function visit(node: AST, visitors: Visitors, ancestors: AST[] = []): void {
16+
const parentAncestors = ancestors
17+
const visitor = visitors[node.type]
18+
19+
// @ts-ignore
20+
;(visitor?.enter ?? visitor)?.(node, parentAncestors)
21+
22+
ancestors = [...ancestors, node]
23+
24+
switch (node.type) {
25+
case 'ArraySpecifier':
26+
visit(node.typeSpecifier, visitors, ancestors)
27+
for (const dimension of node.dimensions) if (dimension) visit(dimension, visitors, ancestors)
28+
break
29+
case 'ExpressionStatement':
30+
visit(node.expression, visitors, ancestors)
31+
break
32+
case 'BlockStatement':
33+
for (const statement of node.body) visit(statement, visitors, ancestors)
34+
break
35+
case 'PreprocessorStatement':
36+
if (node.value) for (const expression of node.value) visit(expression, visitors, ancestors)
37+
break
38+
case 'PrecisionStatement':
39+
visit(node.typeSpecifier, visitors, ancestors)
40+
break
41+
case 'ReturnStatement':
42+
if (node.argument) visit(node.argument, visitors, ancestors)
43+
break
44+
case 'IfStatement':
45+
visit(node.test, visitors, ancestors)
46+
visit(node.consequent, visitors, ancestors)
47+
if (node.alternate) visit(node.alternate, visitors, ancestors)
48+
break
49+
case 'SwitchStatement':
50+
visit(node.discriminant, visitors, ancestors)
51+
for (const kase of node.cases) visit(kase, visitors, ancestors)
52+
break
53+
case 'SwitchCase':
54+
if (node.test) visit(node.test, visitors, ancestors)
55+
for (const statement of node.consequent) visit(statement, visitors, ancestors)
56+
break
57+
case 'WhileStatement':
58+
case 'DoWhileStatement':
59+
visit(node.test, visitors, ancestors)
60+
visit(node.body, visitors, ancestors)
61+
break
62+
case 'ForStatement':
63+
if (node.init) visit(node.init, visitors, ancestors)
64+
if (node.test) visit(node.test, visitors, ancestors)
65+
if (node.update) visit(node.update, visitors, ancestors)
66+
visit(node.body, visitors, ancestors)
67+
break
68+
case 'FunctionDeclaration':
69+
visit(node.typeSpecifier, visitors, ancestors)
70+
visit(node.id, visitors, ancestors)
71+
if (node.body) visit(node.body, visitors, ancestors)
72+
break
73+
case 'FunctionParameter':
74+
visit(node.typeSpecifier, visitors, ancestors)
75+
visit(node.id, visitors, ancestors)
76+
break
77+
case 'VariableDeclaration':
78+
for (const declaration of node.declarations) visit(declaration, visitors, ancestors)
79+
break
80+
case 'VariableDeclarator':
81+
visit(node.typeSpecifier, visitors, ancestors)
82+
visit(node.id, visitors, ancestors)
83+
if (node.init) visit(node.init, visitors, ancestors)
84+
break
85+
case 'UniformDeclarationBlock':
86+
visit(node.typeSpecifier, visitors, ancestors)
87+
for (const member of node.members) visit(member, visitors, ancestors)
88+
if (node.id) visit(node.id, visitors, ancestors)
89+
break
90+
case 'StructDeclaration':
91+
visit(node.id, visitors, ancestors)
92+
for (const member of node.members) visit(member, visitors, ancestors)
93+
break
94+
case 'ArrayExpression':
95+
visit(node.typeSpecifier, visitors, ancestors)
96+
for (const element of node.elements) visit(element, visitors, ancestors)
97+
break
98+
case 'UnaryExpression':
99+
case 'UpdateExpression':
100+
visit(node.argument, visitors, ancestors)
101+
break
102+
case 'BinaryExpression':
103+
case 'AssignmentExpression':
104+
case 'LogicalExpression':
105+
visit(node.left, visitors, ancestors)
106+
visit(node.right, visitors, ancestors)
107+
break
108+
case 'MemberExpression':
109+
visit(node.object, visitors, ancestors)
110+
visit(node.property, visitors, ancestors)
111+
break
112+
case 'ConditionalExpression':
113+
visit(node.test, visitors, ancestors)
114+
visit(node.consequent, visitors, ancestors)
115+
visit(node.alternate, visitors, ancestors)
116+
break
117+
case 'CallExpression':
118+
visit(node.callee, visitors, ancestors)
119+
for (const argument of node.arguments) visit(argument, visitors, ancestors)
120+
break
121+
case 'Program':
122+
for (const statement of node.body) visit(statement, visitors, ancestors)
123+
break
124+
}
125+
126+
// @ts-ignore
127+
visitor?.exit?.(node, parentAncestors)
128+
}

tests/traversal.test.ts

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import { vi, describe, it, expect } from 'vitest'
2+
import { parse, visit, type Visitors, type ReturnStatement, type AST } from 'shaderkit'
3+
4+
describe('traversal', () => {
5+
it('calls visitor enter/exit methods', () => {
6+
const program = parse('return;')
7+
const statement = program.body[0] as ReturnStatement
8+
9+
const visitors = {
10+
Program: {
11+
enter: vi.fn(),
12+
exit: vi.fn(),
13+
},
14+
ReturnStatement: {
15+
enter: vi.fn(),
16+
},
17+
} satisfies Visitors
18+
visit(program, visitors)
19+
20+
expect(visitors.Program.enter).toHaveBeenCalledOnce()
21+
expect(visitors.Program.enter).toHaveBeenCalledWith(program, [])
22+
expect(visitors.ReturnStatement.enter).toHaveBeenCalledOnce()
23+
expect(visitors.ReturnStatement.enter).toHaveBeenCalledWith(statement, [program])
24+
expect(visitors.Program.exit).toHaveBeenCalledOnce()
25+
expect(visitors.Program.exit).toHaveBeenCalledWith(program, [])
26+
})
27+
28+
it('calls visitor default methods', () => {
29+
const program = parse('return;')
30+
const statement = program.body[0] as ReturnStatement
31+
32+
const visitors = {
33+
Program: vi.fn(),
34+
ReturnStatement: vi.fn(),
35+
} satisfies Visitors
36+
visit(program, visitors)
37+
38+
expect(visitors.Program).toHaveBeenCalledOnce()
39+
expect(visitors.Program).toHaveBeenCalledWith(program, [])
40+
expect(visitors.ReturnStatement).toHaveBeenCalledOnce()
41+
expect(visitors.ReturnStatement).toHaveBeenCalledWith(statement, [program])
42+
})
43+
44+
it('tracks node ancestors in parallel', () => {
45+
const program = parse('void a() {} void b() {}')
46+
let parent: AST | null = null
47+
48+
const visitors = {
49+
Program(node, ancestors) {
50+
expect(node).toBe(program)
51+
expect(ancestors).toStrictEqual([])
52+
parent = node
53+
},
54+
FunctionDeclaration(node, ancestors) {
55+
expect(ancestors).toStrictEqual([program])
56+
parent = node
57+
},
58+
BlockStatement(node, ancestors) {
59+
expect(parent).not.toBe(null)
60+
expect(parent).not.toBe(program)
61+
expect(parent).not.toBe(node)
62+
expect(ancestors).toStrictEqual([program, parent])
63+
parent = node
64+
},
65+
} satisfies Visitors
66+
visit(program, visitors)
67+
})
68+
})

0 commit comments

Comments
 (0)