From 749ea40ff4e9cf0e921241de9dcd1efcc9f52381 Mon Sep 17 00:00:00 2001
From: Bronley Plumb <bronley@gmail.com>
Date: Fri, 23 Sep 2022 20:31:07 -0400
Subject: [PATCH] Add AST child searching functionality. (#695)

---
 src/parser/AstNode.spec.ts | 49 ++++++++++++++++++++++++++++++++++++++
 src/parser/AstNode.ts      | 38 ++++++++++++++++++++++++++++-
 2 files changed, 86 insertions(+), 1 deletion(-)
 create mode 100644 src/parser/AstNode.spec.ts

diff --git a/src/parser/AstNode.spec.ts b/src/parser/AstNode.spec.ts
new file mode 100644
index 000000000..ba32ef430
--- /dev/null
+++ b/src/parser/AstNode.spec.ts
@@ -0,0 +1,49 @@
+import { standardizePath as s, util } from '../util';
+import * as fsExtra from 'fs-extra';
+import { Program } from '../Program';
+import type { BrsFile } from '../files/BrsFile';
+import { expect } from 'chai';
+import type { DottedGetExpression } from './Expression';
+import { expectZeroDiagnostics } from '../testHelpers.spec';
+
+let tempDir = s`${process.cwd()}/.tmp`;
+let rootDir = s`${tempDir}/rootDir`;
+let stagingDir = s`${tempDir}/staging`;
+
+describe('Program', () => {
+    let program: Program;
+
+    beforeEach(() => {
+        fsExtra.emptyDirSync(tempDir);
+        program = new Program({
+            rootDir: rootDir,
+            stagingFolderPath: stagingDir
+        });
+        program.createSourceScope(); //ensure source scope is created
+    });
+    afterEach(() => {
+        fsExtra.emptyDirSync(tempDir);
+        program.dispose();
+    });
+
+    describe('AstNode', () => {
+        describe('findNodeAtPosition', () => {
+            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');
+
+                const foxtrot = file.ast.findChildAtPosition<DottedGetExpression>(util.createPosition(3, 71));
+                expect(foxtrot.name.text).to.eql('foxtrot');
+            });
+        });
+    });
+});
+
diff --git a/src/parser/AstNode.ts b/src/parser/AstNode.ts
index 31f274c6f..d4be057b7 100644
--- a/src/parser/AstNode.ts
+++ b/src/parser/AstNode.ts
@@ -1,10 +1,13 @@
 import type { WalkVisitor, WalkOptions } from '../astUtils/visitors';
-import type { Range } from 'vscode-languageserver';
+import { WalkMode } from '../astUtils/visitors';
+import type { Position, Range } from 'vscode-languageserver';
+import { CancellationTokenSource } from 'vscode-languageserver';
 import { InternalWalkMode } from '../astUtils/visitors';
 import type { SymbolTable } from '../SymbolTable';
 import type { BrsTranspileState } from './BrsTranspileState';
 import type { TranspileResult } from '../interfaces';
 import type { AnnotationExpression } from './Expression';
+import util from '../util';
 
 /**
  * A BrightScript AST node
@@ -58,6 +61,39 @@ export abstract class AstNode {
             node = node.parent;
         }
     }
+
+    /**
+     * 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) {
+        const cancel = new CancellationTokenSource();
+        let result: AstNode;
+        this.walk((node) => {
+            const matcherValue = matcher(node);
+            if (matcherValue) {
+                cancel.cancel();
+                result = matcherValue === true ? node : matcherValue;
+            }
+        }, {
+            walkMode: WalkMode.visitAllRecursive,
+            ...options ?? {},
+            cancel: cancel.token
+        });
+        return result as TNodeType;
+    }
+
+    /**
+     * FInd the deepest child that includes the given position
+     */
+    public findChildAtPosition<TNodeType extends AstNode = AstNode>(position: Position, options?: WalkOptions): TNodeType {
+        return this.findChild<TNodeType>((node) => {
+            //if the current node includes this range, keep that node
+            if (util.rangeContains(node.range, position)) {
+                return node.findChildAtPosition(position, options) ?? node;
+            }
+        }, options);
+    }
 }
 
 export abstract class Statement extends AstNode {