Skip to content

Commit

Permalink
Improve findChild and findAncestor AST methods (#807)
Browse files Browse the repository at this point in the history
  • Loading branch information
TwitchBronBron committed May 11, 2023
1 parent 4a02190 commit 2c144e4
Show file tree
Hide file tree
Showing 2 changed files with 170 additions and 18 deletions.
164 changes: 153 additions & 11 deletions src/parser/AstNode.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,11 @@ import { expect } from '../chai-config.spec';
import type { DottedGetExpression } from './Expression';
import { expectZeroDiagnostics } from '../testHelpers.spec';
import { tempDir, rootDir, stagingDir } from '../testHelpers.spec';
import { isAssignmentStatement, isClassStatement, isDottedGetExpression, isPrintStatement } from '../astUtils/reflection';
import type { FunctionStatement } from './Statement';
import { AssignmentStatement } from './Statement';

describe('Program', () => {
describe('AstNode', () => {
let program: Program;

beforeEach(() => {
Expand All @@ -23,23 +26,162 @@ describe('Program', () => {
program.dispose();
});

describe('AstNode', () => {
describe('findNodeAtPosition', () => {
it('finds deepest AstNode that matches the position', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
describe('findChildAtPosition', () => {
it('finds deepest AstNode that matches the position', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
program.validate();
expectZeroDiagnostics(program);
const delta = file.ast.findChildAtPosition<DottedGetExpression>(util.createPosition(3, 52));
expect(delta.name.text).to.eql('delta');
program.validate();
expectZeroDiagnostics(program);
const delta = file.ast.findChildAtPosition<DottedGetExpression>(util.createPosition(3, 52));
expect(delta.name.text).to.eql('delta');

const foxtrot = file.ast.findChildAtPosition<DottedGetExpression>(util.createPosition(3, 71));
expect(foxtrot.name.text).to.eql('foxtrot');
const foxtrot = file.ast.findChildAtPosition<DottedGetExpression>(util.createPosition(3, 71));
expect(foxtrot.name.text).to.eql('foxtrot');
});
});

describe('findChild', () => {
it('finds a child that matches the matcher', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
expect(
file.ast.findChild((node) => {
return isAssignmentStatement(node) && node.name.text === 'alpha';
})
).instanceof(AssignmentStatement);
});

it('returns the exact node that matches', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha1 = invalid
alpha2 = invalid
end sub
`);
let count = 0;
const instance = file.ast.findChild((node) => {
if (isAssignmentStatement(node)) {
count++;
if (count === 2) {
return true;
}
}
});
const expected = (file.ast.statements[0] as FunctionStatement).func.body.statements[1];
expect(instance).to.equal(expected);
});

it('returns undefined when matcher never returned true', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
expect(
file.ast.findChild((node) => false)
).not.to.exist;
});

it('returns the value returned from the matcher', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
const secondStatement = (file.ast.statements[0] as FunctionStatement).func.body.statements[1];
expect(
file.ast.findChild((node) => secondStatement)
).to.equal(secondStatement);
});

it('cancels properly', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
let count = 0;
file.ast.findChild((node, cancelToken) => {
count++;
cancelToken.cancel();
});
expect(count).to.eql(1);
});
});

describe('findAncestor', () => {
it('returns node when matcher returns true', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
const secondStatement = (file.ast.statements[0] as FunctionStatement).func.body.statements[1];
const foxtrot = file.ast.findChild((node) => {
return isDottedGetExpression(node) && node.name?.text === 'foxtrot';
});
expect(
foxtrot.findAncestor(isPrintStatement)
).to.equal(secondStatement);
});

it('returns undefined when no match found', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
const foxtrot = file.ast.findChild((node) => {
return isDottedGetExpression(node) && node.name?.text === 'foxtrot';
});
expect(
foxtrot.findAncestor(isClassStatement)
).to.be.undefined;
});

it('returns overridden node when returned in matcher', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
const firstStatement = (file.ast.statements[0] as FunctionStatement).func.body.statements[0];
const foxtrot = file.ast.findChild((node) => {
return isDottedGetExpression(node) && node.name?.text === 'foxtrot';
});
expect(
foxtrot.findAncestor(node => firstStatement)
).to.equal(firstStatement);
});

it('returns overridden node when returned in matcher', () => {
const file = program.setFile<BrsFile>('source/main.brs', `
sub main()
alpha = invalid
print alpha.beta.charlie.delta(alpha.echo.foxtrot())
end sub
`);
let count = 0;
const firstStatement = (file.ast.statements[0] as FunctionStatement).func.body.statements[0];
firstStatement.findAncestor((node, cancel) => {
count++;
cancel.cancel();
});
expect(count).to.eql(1);
});
});
});
24 changes: 17 additions & 7 deletions src/parser/AstNode.ts
Original file line number Diff line number Diff line change
Expand Up @@ -52,13 +52,23 @@ export abstract class AstNode {
}

/**
* Walk upward and return the first node that results in `true` from the matcher
* Walk upward and return the first node that results in `true` from the matcher.
* @param matcher a function called for each node. If you return true, this function returns the specified node. If you return a node, that node is returned. all other return values continue the loop
* The function's second parameter is a cancellation token. If you'd like to short-circuit the walk, call `cancellationToken.cancel()`, then this function will return `undefined`
*/
public findAncestor<TNode extends AstNode = AstNode>(matcher: (node: AstNode) => boolean | undefined) {
public findAncestor<TNode extends AstNode = AstNode>(matcher: (node: AstNode, cancellationToken: CancellationTokenSource) => boolean | AstNode | undefined | void): TNode {
let node = this.parent;

const cancel = new CancellationTokenSource();
while (node) {
if (matcher(node)) {
return node as TNode;
let matcherValue = matcher(node, cancel);
if (cancel.token.isCancellationRequested) {
return;
}
if (matcherValue) {
cancel.cancel();
return (matcherValue === true ? node : matcherValue) as TNode;

}
node = node.parent;
}
Expand All @@ -68,11 +78,11 @@ export abstract class AstNode {
* Find the first child where the matcher evaluates to true.
* @param matcher a function called for each node. If you return true, this function returns the specified node. If you return a node, that node is returned. all other return values continue the loop
*/
public findChild<TNodeType extends AstNode = AstNode>(matcher: (node: AstNode) => boolean | AstNode, options?: WalkOptions) {
public findChild<TNode extends AstNode = AstNode>(matcher: (node: AstNode, cancellationSource) => boolean | AstNode | undefined | void, options?: WalkOptions) {
const cancel = new CancellationTokenSource();
let result: AstNode;
this.walk((node) => {
const matcherValue = matcher(node);
const matcherValue = matcher(node, cancel);
if (matcherValue) {
cancel.cancel();
result = matcherValue === true ? node : matcherValue;
Expand All @@ -82,7 +92,7 @@ export abstract class AstNode {
...options ?? {},
cancel: cancel.token
});
return result as TNodeType;
return result as TNode;
}

/**
Expand Down

0 comments on commit 2c144e4

Please sign in to comment.