Skip to content

Commit

Permalink
Fix issues with the ast walkArray function (#1347)
Browse files Browse the repository at this point in the history 
* Fix issues with the ast walkArray function

* Fix missing node visitor issues
  • Loading branch information
@TwitchBronBron
TwitchBronBron authored Nov 21, 2024
1 parent d2d08a7 commit 3d998ef
Show file tree
Hide file tree
Showing 3 changed files with 260 additions and 25 deletions.
217 changes: 210 additions & 7 deletions src/astUtils/visitors.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,16 @@ import type { BrsFile } from '../files/BrsFile';
import type { FunctionStatement } from '../parser/Statement';
import { PrintStatement, Block, ReturnStatement, ExpressionStatement } from '../parser/Statement';
import { TokenKind } from '../lexer/TokenKind';
import { createVisitor, WalkMode, walkStatements } from './visitors';
import { isPrintStatement } from './reflection';
import { createCall, createToken, createVariableExpression } from './creators';
import { createVisitor, walkArray, WalkMode, walkStatements } from './visitors';
import { isBlock, isLiteralExpression, isPrintStatement } from './reflection';
import { createCall, createIntegerLiteral, createToken, createVariableExpression } from './creators';
import { createStackedVisitor } from './stackedVisitor';
import { AstEditor } from './AstEditor';
import { Parser } from '../parser/Parser';
import type { Statement, Expression, AstNode } from '../parser/AstNode';
import { expectZeroDiagnostics } from '../testHelpers.spec';
import type { LiteralExpression, VariableExpression } from '../parser/Expression';
import { BinaryExpression } from '../parser/Expression';

describe('astUtils visitors', () => {
const rootDir = process.cwd();
Expand Down Expand Up @@ -1068,6 +1070,88 @@ describe('astUtils visitors', () => {
).to.be.lengthOf(0);
});

it('walks everything when the first element is replaced', () => {
const { ast } = Parser.parse(`
sub main()
print 1
print 2
print 3
end sub
`);
const target = ast.findChild<Block>(isBlock).statements[0];
let callCount = 0;
ast.walk((astNode, parent, owner: Statement[], key) => {
if (isPrintStatement(astNode)) {
callCount++;
}
if (astNode === target) {
owner.splice(key, 1);
}
}, {
walkMode: WalkMode.visitAllRecursive
});
//the visitor should have been called for every statement
expect(callCount).to.eql(3);
expect(
(ast.statements[0] as FunctionStatement).func.body.statements
).not.to.include(target);
});

it('walks everything when the middle element is replaced', () => {
const { ast } = Parser.parse(`
sub main()
print 1
print 2
print 3
end sub
`);
const target = ast.findChild<Block>(isBlock).statements[1];
let callCount = 0;
ast.walk((astNode, parent, owner: Statement[], key) => {
if (isPrintStatement(astNode)) {
callCount++;
}
if (astNode === target) {
owner.splice(key, 1);
}
}, {
walkMode: WalkMode.visitAllRecursive
});
//the visitor should have been called for every statement
expect(callCount).to.eql(3);
expect(
(ast.statements[0] as FunctionStatement).func.body.statements
).not.to.include(target);
});

it('walks everything when the end element is replaced', () => {
const { ast } = Parser.parse(`
sub main()
print 1
print 2
print 3
end sub
`);
const target = ast.findChild<Block>(isBlock).statements[2];
let callCount = 0;
ast.walk((astNode, parent, owner: Statement[], key) => {
if (isPrintStatement(astNode)) {
callCount++;
}
if (astNode === target) {
owner.splice(key, 1);
}
}, {
walkMode: WalkMode.visitAllRecursive
});
//the visitor should have been called for every statement
expect(callCount).to.eql(3);
expect(
(ast.statements[0] as FunctionStatement).func.body.statements
).not.to.include(target);
});


it('can be used to insert statements', () => {
const { ast } = Parser.parse(`
sub main()
Expand All @@ -1085,23 +1169,142 @@ describe('astUtils visitors', () => {
//add another expression to the list every time. This should result in 1 the first time, 2 the second, 3 the third.
calls.push(new ExpressionStatement(
createCall(
createVariableExpression('doSomethingBeforePrint')
createVariableExpression('doSomethingBeforePrint'),
[
createIntegerLiteral(callExpressionCount.toString())
]
)
));
owner.splice(key, 0, ...calls);
owner.splice(key + 1, 0, ...calls.map(x => x.clone()));
},
CallExpression: () => {
CallExpression: (call) => {
callExpressionCount++;
console.log('call visitor for', (call.args[0] as LiteralExpression).token.text);
}
}), {
walkMode: WalkMode.visitAllRecursive
});
//the visitor should have been called for every statement
expect(printStatementCount).to.eql(3);
expect(callExpressionCount).to.eql(0);
//since the calls were injected after each print statement, we should have 1 call for the first print, 2 for the second, and 3 for the third
expect(callExpressionCount).to.eql(6);
expect(
(ast.statements[0] as FunctionStatement).func.body.statements
).to.be.lengthOf(9);
});

it('walks a new child when returned from a visitor', () => {
let walkedLiterals: string[] = [];

Parser.parse(`
sub main()
print 1 + 2
end sub
`).ast.walk(createVisitor({
BinaryExpression: (node, parent, owner: Statement[], key) => {
//replace the `1 + 2` binary expression with a new binary expression
if (isLiteralExpression(node.left) && node.left.token.text === '1') {
return new BinaryExpression(
createIntegerLiteral('3'),
createToken(TokenKind.Plus),
createIntegerLiteral('4')
);
}
},
LiteralExpression: (node) => {
walkedLiterals.push(node.token.text);
}
}), {
walkMode: WalkMode.visitAllRecursive
});

expect(walkedLiterals).to.eql(['3', '4']);
});

it('walks a new child when returned from a visitor and using an AstEditor', () => {
let walkedLiterals: string[] = [];

Parser.parse(`
sub main()
print 1 + 2
end sub
`).ast.walk(createVisitor({
BinaryExpression: (node, parent, owner: Statement[], key) => {
//replace the `1 + 2` binary expression with a new binary expression
if (isLiteralExpression(node.left) && node.left.token.text === '1') {
return new BinaryExpression(
createIntegerLiteral('3'),
createToken(TokenKind.Plus),
createIntegerLiteral('4')
);
}
},
LiteralExpression: (node) => {
walkedLiterals.push(node.token.text);
}
}), {
walkMode: WalkMode.visitAllRecursive,
editor: new AstEditor()
});

expect(walkedLiterals).to.eql(['3', '4']);
});
});

describe('walkArray', () => {
const one = createVariableExpression('one');
const two = createVariableExpression('two');
const three = createVariableExpression('three');
const four = createVariableExpression('four');
const five = createVariableExpression('five');

function doTest(startingArray: VariableExpression[], expected: VariableExpression[], visitor?: (item: AstNode, parent: AstNode, owner: any, key: number) => any) {
const visitedItems: VariableExpression[] = [];
walkArray(startingArray, (item, parent, owner, key) => {
visitedItems.push(item as any);
return visitor?.(item, parent, owner, key);
}, { walkMode: WalkMode.visitAllRecursive });
expect(
visitedItems.map(x => x.name.text)
).to.eql(
expected.map(x => x.name.text)
);
}

it('walks every element in the array', () => {
doTest(
[one, two, three, four, five],
[one, two, three, four, five]
);
});

it('walks new items added to the array', () => {
doTest(
[one, two],
[one, three, two, four],
(item, parent, owner, key) => {
//insert a value after one
if (item === one) {
owner.splice(key + 1, 0, three);
//insert a value after one
} else if (item === two) {
owner.splice(key + 1, 0, four);
}
}
);
});

it('triggers on nodes that were skiped due to insertions', () => {
doTest(
[one, two, three],
[one, two, four, three],
(item, parent, owner, key) => {
//insert a value after one
if (item === two) {
owner.splice(key, 0, four);
}
}
);
});
});
});
64 changes: 48 additions & 16 deletions src/astUtils/visitors.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,48 +25,65 @@ export type WalkVisitor = <T = AstNode>(node: AstNode, parent?: AstNode, owner?:

/**
* A helper function for Statement and Expression `walkAll` calls.
* @returns a new AstNode if it was changed by returning from the visitor, or undefined if not
*/
export function walk<T>(owner: T, key: keyof T, visitor: WalkVisitor, options: WalkOptions, parent?: AstNode) {
export function walk<T>(owner: T, key: keyof T, visitor: WalkVisitor, options: WalkOptions, parent?: AstNode): AstNode | void {
let returnValue: AstNode | void;

//stop processing if canceled
if (options.cancel?.isCancellationRequested) {
return;
return returnValue;
}

//the object we're visiting
let element = owner[key] as any as AstNode;
if (!element) {
return;
return returnValue;
}

//link this node to its parent
element.parent = parent ?? owner as unknown as AstNode;
parent = parent ?? owner as unknown as AstNode;
element.parent = parent;


//notify the visitor of this element
if (element.visitMode & options.walkMode) {
const result = visitor?.(element, element.parent as any, owner, key);
returnValue = visitor?.(element, element.parent as any, owner, key);

//replace the value on the parent if the visitor returned a Statement or Expression (this is how visitors can edit AST)
if (result && (isExpression(result) || isStatement(result))) {
if (returnValue && (isExpression(returnValue) || isStatement(returnValue))) {
//if we have an editor, use that to modify the AST
if (options.editor) {
options.editor.setProperty(owner, key, result as any);
options.editor.setProperty(owner, key, returnValue as any);

//we don't have an editor, modify the AST directly
} else {
(owner as any)[key] = result;
//don't walk the new element
return;
(owner as any)[key] = returnValue;
}
}
}

//stop processing if canceled
if (options.cancel?.isCancellationRequested) {
return;
return returnValue;
}

//get the element again in case it was replaced by the visitor
element = owner[key] as any as AstNode;
if (!element) {
return returnValue;
}

//set the parent of this new expression
element.parent = parent;

if (!element.walk) {
throw new Error(`${owner.constructor.name}["${String(key)}"]${parent ? ` for ${parent.constructor.name}` : ''} does not contain a "walk" method`);
}
//walk the child expressions
element.walk(visitor, options);

return returnValue;
}

/**
Expand All @@ -77,13 +94,28 @@ export function walk<T>(owner: T, key: keyof T, visitor: WalkVisitor, options: W
* @param parent the parent AstNode of each item in the array
* @param filter a function used to filter items from the array. return true if that item should be walked
*/
export function walkArray<T = AstNode>(array: Array<T>, visitor: WalkVisitor, options: WalkOptions, parent?: AstNode, filter?: <T>(element: T) => boolean) {
export function walkArray<T extends AstNode = AstNode>(array: Array<T>, visitor: WalkVisitor, options: WalkOptions, parent?: AstNode, filter?: <T>(element: T) => boolean) {
let processedNodes = new Set<AstNode>();

for (let i = 0; i < array?.length; i++) {
if (!filter || filter(array[i])) {
const startLength = array.length;
walk(array, i, visitor, options, parent);
//compensate for deleted or added items.
i += array.length - startLength;
let item = array[i];
//skip already processed nodes for this array walk
if (processedNodes.has(item)) {
continue;
}
processedNodes.add(item);

//if the walk produced a new node, we will assume the original node was handled, and the new node's children were walked, so we can skip it if we enter recovery mode
const newNode = walk(array, i, visitor, options, parent);
if (newNode) {
processedNodes.add(newNode);
}

//if the current item changed, restart the entire loop (we'll skip any already-processed items)
if (array[i] !== item) {
i = -1;
}
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions src/parser/Statement.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ import { DynamicType } from '../types/DynamicType';
import type { BscType } from '../types/BscType';
import type { TranspileState } from './TranspileState';
import { SymbolTable } from '../SymbolTable';
import type { Expression } from './AstNode';
import type { AstNode, Expression } from './AstNode';
import { Statement } from './AstNode';

export class EmptyStatement extends Statement {
Expand Down Expand Up @@ -699,7 +699,7 @@ export class PrintStatement extends Statement {
walk(visitor: WalkVisitor, options: WalkOptions) {
if (options.walkMode & InternalWalkMode.walkExpressions) {
//sometimes we have semicolon Tokens in the expressions list (should probably fix that...), so only walk the actual expressions
walkArray(this.expressions, visitor, options, this, (item) => isExpression(item as any));
walkArray(this.expressions as AstNode[], visitor, options, this, (item) => isExpression(item as any));
}
}

Expand Down

0 comments on commit 3d998ef

Please sign in to comment.