TSC Graph Context: fix performance

Fixes CODY-1434

Previously, the `tsc-mixed` context retriever made  the UI freeze
because it was doing super inefficient things with context expansion.
This PR fixes the problem so that the performance is back to normal.

vscode/src/completions/context/retrievers/section-history/nextTick.ts

@@ -0,0 +1,3 @@
+export function nextTick() {
+    return new Promise(resolve => process.nextTick(resolve))
+}

vscode/src/completions/context/retrievers/section-history/section-history-retriever.test.ts

@@ -6,6 +6,7 @@ import { testFileUri } from '@sourcegraph/cody-shared'
 import { range, withPosixPathsInString } from '../../../../testutils/textDocument'
 import * as docContextGetters from '../../../doc-context-getters'
 
+import { nextTick } from './nextTick'
 import { SectionHistoryRetriever } from './section-history-retriever'
 
 const document1Uri = testFileUri('document1.ts')
@@ -274,7 +275,3 @@ describe('GraphSectionObserver', () => {
         })
     })
 })
-
-function nextTick() {
-    return new Promise(resolve => process.nextTick(resolve))
-}

vscode/src/completions/context/retrievers/tsc/SymbolFormatter.test.ts

@@ -19,7 +19,7 @@ describe('SymbolFormatter', () => {
         // quickly adds up, but it's good enough for now.
         const program = ts.createProgram(['test.ts'], {}, host)
         const checker = program.getTypeChecker()
-        const formatter = new SymbolFormatter(checker)
+        const formatter = new SymbolFormatter(checker, 10)
         const sourceFile = program.getSourceFile('test.ts')
         if (!sourceFile) {
             return []
@@ -34,7 +34,7 @@ describe('SymbolFormatter', () => {
             if (!symbol) {
                 continue
             }
-            result.push(formatter.formatSymbol(symbol))
+            result.push(formatter.formatSymbol(statement, symbol, 0))
         }
         return result
     }

vscode/src/completions/context/retrievers/tsc/SymbolFormatter.ts

@@ -26,24 +26,24 @@ const path = defaultPathFunctions()
  * `ts.{Symbol,Type}`).
  */
 export class SymbolFormatter {
-    private queue = new Set<ts.Symbol>()
+    public queue = new Map<ts.Symbol, number>()
     public isRendered = new Set<ts.Symbol>()
-    constructor(private checker: ts.TypeChecker) {}
+    private depth = 0
+    constructor(
+        private checker: ts.TypeChecker,
+        private maxDepth: number
+    ) {}
 
-    public formatSymbolWithQueue(sym: ts.Symbol): { formatted: string; queue: ts.Symbol[] } {
-        const formatted = this.formatSymbol(sym)
+    public formatSymbol(
+        declaration: ts.Node,
+        sym: ts.Symbol,
+        depth: number,
+        params?: { stripEnclosingInformation?: boolean }
+    ): string {
+        const oldDepth = this.depth
+        this.depth = depth
         this.queueRelatedSymbols(sym)
-
-        // TODO: return a queue of symbols that got added in this function call,
-        // not accumulated list
-        return { formatted, queue: [...this.queue] }
-    }
-
-    public formatSymbol(sym: ts.Symbol, params?: { stripEnclosingInformation?: boolean }): string {
-        const declaration = sym.declarations?.[0]
-        if (!declaration) {
-            return ''
-        }
+        this.depth = oldDepth
 
         if (ts.isClassLike(declaration) || ts.isInterfaceDeclaration(declaration)) {
             return this.formatClassOrInterface(declaration, sym)
@@ -65,6 +65,9 @@ export class SymbolFormatter {
     }
 
     private queueRelatedSymbols(sym: ts.Symbol): void {
+        if (this.depth > this.maxDepth) {
+            return
+        }
         const walkNode = (node: ts.Node | undefined): void => {
             if (!node) {
                 return
@@ -108,7 +111,7 @@ export class SymbolFormatter {
         if (isStdLibSymbol(s)) {
             return
         }
-        this.queue.add(s)
+        this.queue.set(s, this.depth + 1)
     }
 
     private registerRenderedNode(node: ts.Node): void {
@@ -184,9 +187,17 @@ export class SymbolFormatter {
                 }
                 const name = declarationName(decl)
                 if (name) {
-                    p.line(this.formatSymbol(member, { stripEnclosingInformation: true }))
+                    p.line(
+                        this.formatSymbol(decl, member, this.depth + 1, {
+                            stripEnclosingInformation: true,
+                        })
+                    )
                 } else if (memberName === ts.InternalSymbolName.Constructor) {
-                    p.line(this.formatSymbol(member, { stripEnclosingInformation: true }))
+                    p.line(
+                        this.formatSymbol(decl, member, this.depth + 1, {
+                            stripEnclosingInformation: true,
+                        })
+                    )
                 }
             }
         })

vscode/src/completions/context/retrievers/tsc/relevantTypeIdentifiers.ts

@@ -1,26 +1,43 @@
 import ts from 'typescript'
 import { declarationName } from './SymbolFormatter'
 
+export type NodeMatchKind =
+    | 'imports'
+    | 'call-expression'
+    | 'property-access'
+    | 'function-declaration'
+    | 'declaration'
+    | 'none'
+
 /**
  * Returns a list of identifier nodes that should be added to the Cody context.
  *
  * The logic for this function is going to be evolving as we add support for
  * more syntax constructs where we want to inject graph context.
  */
-export function relevantTypeIdentifiers(checker: ts.TypeChecker, node: ts.Node): ts.Node[] {
-    const result: ts.Node[] = []
-    pushTypeIdentifiers(result, checker, node)
-    return result
+export function relevantTypeIdentifiers(
+    checker: ts.TypeChecker,
+    node: ts.Node
+): { kind: NodeMatchKind; nodes: ts.Node[] } {
+    const nodes: ts.Node[] = []
+    const kind = pushTypeIdentifiers(nodes, checker, node)
+    return { kind, nodes }
 }
 
-export function pushTypeIdentifiers(result: ts.Node[], checker: ts.TypeChecker, node: ts.Node): void {
+export function pushTypeIdentifiers(
+    result: ts.Node[],
+    checker: ts.TypeChecker,
+    node: ts.Node
+): NodeMatchKind {
     if (ts.isSourceFile(node)) {
         ts.forEachChild(node, child => {
             if (ts.isImportDeclaration(child)) {
                 pushDescendentIdentifiers(result, child)
             }
         })
-    } else if (
+        return 'imports'
+    }
+    if (
         ts.isSetAccessorDeclaration(node) ||
         ts.isGetAccessorDeclaration(node) ||
         ts.isConstructorDeclaration(node) ||
@@ -36,19 +53,24 @@ export function pushTypeIdentifiers(result: ts.Node[], checker: ts.TypeChecker,
         if (node.type) {
             pushDescendentIdentifiers(result, node.type)
         }
-    } else if (ts.isCallExpression(node)) {
+        return 'function-declaration'
+    }
+    if (ts.isCallExpression(node)) {
         result.push(...rightmostIdentifier(node.expression))
-    } else if (ts.isPropertyAccessExpression(node)) {
+        return 'call-expression'
+    }
+    if (ts.isPropertyAccessExpression(node)) {
         result.push(...rightmostIdentifier(node.expression))
-    } else {
-        const name = declarationName(node)
-        if (name) {
-            result.push(name)
-        } else {
-            // Uncomment below to debug what kind of if (ts.isX) case to handle
-            // console.log({ text: node.getText(), kindString: ts.SyntaxKind[node.kind] })
-        }
+        return 'property-access'
+    }
+    const name = declarationName(node)
+    if (name) {
+        result.push(name)
+        return 'declaration'
     }
+    // Uncomment below to debug what kind of if (ts.isX) case to handle
+    // console.log({ text: node.getText(), kindString: ts.SyntaxKind[node.kind] })
+    return 'none'
 }
 
 // A hacky way to get the `ts.Identifier` node furthest to the right.  Ideally,

vscode/src/completions/context/retrievers/tsc/tsc-retriever.test.ts

@@ -47,7 +47,7 @@ describe.skipIf(isWindows())('TscRetriever', () => {
             docContext,
             document,
             position,
-            hints: { maxChars: 1000, maxMs: 100 },
+            hints: { maxChars: 10_000, maxMs: 100 },
         })
         return { snippets: result, moduleName, namespaceName }
     }
@@ -57,6 +57,28 @@ describe.skipIf(isWindows())('TscRetriever', () => {
     }
 
     it('imports', async () => {
+        expect(
+            await retrieveText(dedent`
+            import { execFileSync } from 'child_process'
+            const a = █
+    `)
+            // TODO: drill into Holder
+        ).toMatchInlineSnapshot(`
+          [
+            "function execFileSync(file: string): Buffer",
+            "function execFileSync(file: string, options: ExecFileSyncOptionsWithStringEncoding): string",
+            "function execFileSync(file: string, options: ExecFileSyncOptionsWithBufferEncoding): Buffer",
+            "function execFileSync(file: string, options?: ExecFileSyncOptions | undefined): string | Buffer",
+            "function execFileSync(file: string, args: readonly string[]): Buffer",
+            "function execFileSync(file: string, args: readonly string[], options: ExecFileSyncOptionsWithStringEncoding): string",
+            "function execFileSync(file: string, args: readonly string[], options: ExecFileSyncOptionsWithBufferEncoding): Buffer",
+            "function execFileSync(file: string, args?: readonly string[] | undefined, options?: ExecFileSyncOptions | undefined): string | Buffer",
+            "const a: any",
+          ]
+        `)
+    })
+
+    it('imports2', async () => {
         const { moduleName } = await retrieve(
             dedent`
             export interface Holder { bananas: number }
@@ -254,7 +276,9 @@ describe.skipIf(isWindows())('TscRetriever', () => {
             interface A { value: number }
             interface B { a(): A }
             const b: B = {}
-            b.a().█
+            function foo() {
+              b.a().█
+            }
     `)
         ).toMatchInlineSnapshot(`
           [