diff --git a/src/astUtils/visitors.spec.ts b/src/astUtils/visitors.spec.ts index d3a583107..7355ed693 100644 --- a/src/astUtils/visitors.spec.ts +++ b/src/astUtils/visitors.spec.ts @@ -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(); @@ -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(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(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(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() @@ -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); + } + } + ); + }); }); }); diff --git a/src/astUtils/visitors.ts b/src/astUtils/visitors.ts index 254b9ddb2..24303c883 100644 --- a/src/astUtils/visitors.ts +++ b/src/astUtils/visitors.ts @@ -25,48 +25,65 @@ export type WalkVisitor = (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(owner: T, key: keyof T, visitor: WalkVisitor, options: WalkOptions, parent?: AstNode) { +export function walk(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; } /** @@ -77,13 +94,28 @@ export function walk(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(array: Array, visitor: WalkVisitor, options: WalkOptions, parent?: AstNode, filter?: (element: T) => boolean) { +export function walkArray(array: Array, visitor: WalkVisitor, options: WalkOptions, parent?: AstNode, filter?: (element: T) => boolean) { + let processedNodes = new Set(); + 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; + } } } } diff --git a/src/parser/Statement.ts b/src/parser/Statement.ts index 6af057376..3b9baf8bb 100644 --- a/src/parser/Statement.ts +++ b/src/parser/Statement.ts @@ -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 { @@ -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)); } }