feat(codewhisperer): crossfile js, ts jsx, tsx support #3668

* restructure crossfile/utg code path
* add ts js tsx jsx crossfile support
* cleanup test states in both beforeEach and afterEach in editorUtilities.test.ts

Author: Will-ShaoHua

.changes/next-release/Feature-5baff9c0-54b5-4ed3-8a3e-a74800cc5af5.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "CodeWhisperer: Improve Java suggestion quality with enhanced file context fetching"
+}

src/codewhisperer/models/constants.ts

@@ -76,6 +76,10 @@ export const javascript = 'javascript'
 
 export const typescript = 'typescript'
 
+export const jsx = 'javascriptreact'
+
+export const tsx = 'typescriptreact'
+
 export const plaintext = 'plaintext'
 
 // use vscode languageId here

src/codewhisperer/util/supplementalContext/codeParsingUtil.ts

@@ -3,9 +3,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import * as vscode from 'vscode'
 import path = require('path')
 import { DependencyGraph } from '../dependencyGraph/dependencyGraph'
-import * as vscode from 'vscode'
 
 export interface utgLanguageConfig {
     extension: string
@@ -68,29 +68,30 @@ export function countSubstringMatches(arr1: string[], arr2: string[]): number {
     return count
 }
 
-export async function isTestFile(editor: vscode.TextEditor, dependencyGraph: DependencyGraph): Promise<boolean> {
-    const languageConfig = utgLanguageConfigs[editor.document.languageId]
-    if (!languageConfig) {
-        // We have enabled the support only for python and Java for this check
-        // as we depend on Regex for this validation.
-        return false
-    }
-
-    // TODO (Metrics): Add total number of calls to isTestFile
-    if (isTestFileByName(editor.document.uri.fsPath, editor.document.languageId)) {
-        return true
+export async function isTestFile(
+    filePath: string,
+    languageConfig: {
+        languageId: vscode.TextDocument['languageId']
+        dependencyGraph?: DependencyGraph
+        fileContent?: string
     }
+): Promise<boolean> {
+    const pathContainsTest = filePath.includes(`tests/`) || filePath.includes('test/') || filePath.includes('tst/')
+    const fileNameMatchTestPatterns = isTestFileByName(filePath, languageConfig.languageId)
 
-    // TODO (Metrics): Add metrics for isTestFileByName Failure
-    // (to help us determine if people follow naming conventions)
-    if (await dependencyGraph.isTestFile(editor.document.getText())) {
+    if (pathContainsTest || fileNameMatchTestPatterns) {
         return true
     }
 
-    return false
+    // This run slowly thus lazily execute
+    const fileHasTestDependency =
+        languageConfig.dependencyGraph && languageConfig.fileContent
+            ? await languageConfig.dependencyGraph.isTestFile(languageConfig.fileContent)
+            : false
+    return fileHasTestDependency
 }
 
-export function isTestFileByName(filePath: string, language: string): boolean {
+function isTestFileByName(filePath: string, language: vscode.TextDocument['languageId']): boolean {
     const languageConfig = utgLanguageConfigs[language]
     if (!languageConfig) {
         // We have enabled the support only for python and Java for this check
@@ -100,8 +101,6 @@ export function isTestFileByName(filePath: string, language: string): boolean {
     const testFilenamePattern = languageConfig.testFilenamePattern
 
     const filename = path.basename(filePath)
-    if (testFilenamePattern.test(filename)) {
-        return true
-    }
-    return false
+
+    return testFilenamePattern.test(filename)
 }

src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts

@@ -5,16 +5,40 @@
 
 import * as vscode from 'vscode'
 import * as fs from 'fs-extra'
-import { DependencyGraph } from '../dependencyGraph/dependencyGraph'
+import path = require('path')
 import { BM25Document, BM25Okapi } from './rankBm25'
-import { isRelevant } from './editorFilesUtil'
 import { ToolkitError } from '../../../shared/errors'
 import { UserGroup, crossFileContextConfig, supplemetalContextFetchingTimeoutMsg } from '../../models/constants'
 import { CancellationError } from '../../../shared/utilities/timeoutUtils'
 import { CodeWhispererSupplementalContextItem } from './supplementalContextUtil'
 import { CodeWhispererUserGroupSettings } from '../userGroupUtil'
+import { isTestFile } from './codeParsingUtil'
+import { getOpenFilesInWindow } from '../../../shared/utilities/editorUtilities'
+
+type CrossFileSupportedLanguage =
+    | 'java'
+    | 'python'
+    | 'javascript'
+    | 'typescript'
+    | 'javascriptreact'
+    | 'typescriptreact'
+
+// TODO: ugly, can we make it prettier? like we have to manually type 'java', 'javascriptreact' which is error prone
+// TODO: Move to another config file or constants file
+// Supported language to its corresponding file ext
+const supportedLanguageToDialects: Readonly<Record<CrossFileSupportedLanguage, Set<string>>> = {
+    java: new Set<string>(['.java']),
+    python: new Set<string>(['.py']),
+    javascript: new Set<string>(['.js', '.jsx']),
+    javascriptreact: new Set<string>(['.js', '.jsx']),
+    typescript: new Set<string>(['.ts', '.tsx']),
+    typescriptreact: new Set<string>(['.ts', '.tsx']),
+}
+
+function isCrossFileSupported(languageId: string): languageId is CrossFileSupportedLanguage {
+    return Object.keys(supportedLanguageToDialects).includes(languageId)
+}
 
-const crossFileLanguageConfigs = ['java']
 interface Chunk {
     fileName: string
     content: string
@@ -24,20 +48,21 @@ interface Chunk {
 
 export async function fetchSupplementalContextForSrc(
     editor: vscode.TextEditor,
-    dependencyGraph: DependencyGraph,
     cancellationToken: vscode.CancellationToken
 ): Promise<CodeWhispererSupplementalContextItem[] | undefined> {
-    if (crossFileLanguageConfigs.includes(editor.document.languageId) === false) {
-        return undefined
-    }
+    const shouldProceed = shouldFetchCrossFileContext(
+        editor.document.languageId,
+        CodeWhispererUserGroupSettings.instance.userGroup
+    )
 
-    if (CodeWhispererUserGroupSettings.instance.userGroup !== UserGroup.CrossFile) {
-        return []
+    if (!shouldProceed) {
+        return shouldProceed === undefined ? undefined : []
     }
 
     // Step 1: Get relevant cross files to refer
-    const relevantCrossFilePaths = await getRelevantCrossFiles(editor, dependencyGraph)
+    const relevantCrossFilePaths = await getCrossFileCandidates(editor)
     throwIfCancelled(cancellationToken)
+
     // Step 2: Split files to chunks with upper bound on chunkCount
     // We restrict the total number of chunks to improve on latency.
     // Chunk linking is required as we want to pass the next chunk value for matched chunk.
@@ -110,6 +135,27 @@ function getInputChunk(editor: vscode.TextEditor, chunkSize: number) {
     return inputChunk
 }
 
+/**
+ * Util to decide if we need to fetch crossfile context since CodeWhisperer CrossFile Context feature is gated by userGroup and language level
+ * @param languageId: VSCode language Identifier
+ * @param userGroup: CodeWhisperer user group settings, refer to userGroupUtil.ts
+ * @returns specifically returning undefined if the langueage is not supported,
+ * otherwise true/false depending on if the language is fully supported or not belonging to the user group
+ */
+function shouldFetchCrossFileContext(languageId: string, userGroup: UserGroup): boolean | undefined {
+    if (!isCrossFileSupported(languageId)) {
+        return undefined
+    }
+
+    if (languageId === 'java') {
+        return true
+    } else if (supportedLanguageToDialects[languageId] && userGroup === UserGroup.CrossFile) {
+        return true
+    } else {
+        return false
+    }
+}
+
 /**
  * This linking is required from science experimentations to pass the next contnet chunk
  * when a given chunk context passes the match in BM25.
@@ -158,29 +204,27 @@ function splitFileToChunks(filePath: string, chunkSize: number): Chunk[] {
  * This function will return relevant cross files for the given editor file
  * by referencing open files, imported files and same package files.
  */
-async function getRelevantCrossFiles(editor: vscode.TextEditor, dependencyGraph: DependencyGraph): Promise<string[]> {
-    return getOpenFilesInWindow().filter(file => {
-        return isRelevant(editor.document.fileName, file, editor.document.languageId)
+async function getCrossFileCandidates(editor: vscode.TextEditor): Promise<string[]> {
+    const targetFile = editor.document.uri.fsPath
+    const language = editor.document.languageId as CrossFileSupportedLanguage
+    const dialects = supportedLanguageToDialects[language]
+
+    /**
+     * Consider a file which
+     * 1. is different from the target
+     * 2. has the same file extension or it's one of the dialect of target file (e.g .js vs. .jsx)
+     * 3. is not a test file
+     */
+    return await getOpenFilesInWindow(async candidateFile => {
+        return (
+            targetFile !== candidateFile &&
+            (path.extname(targetFile) === path.extname(candidateFile) ||
+                (dialects && dialects.has(path.extname(candidateFile)))) &&
+            !(await isTestFile(candidateFile, { languageId: language }))
+        )
     })
 }
 
-function getOpenFilesInWindow(): string[] {
-    const filesOpenedInEditor: string[] = []
-
-    try {
-        const tabArrays = vscode.window.tabGroups.all
-        tabArrays.forEach(tabArray => {
-            tabArray.tabs.forEach(tab => {
-                filesOpenedInEditor.push((tab.input as any).uri.path)
-            })
-        })
-    } catch (e) {
-        // Older versions of VSC do not have the tab API
-    }
-
-    return filesOpenedInEditor
-}
-
 function throwIfCancelled(token: vscode.CancellationToken): void | never {
     if (token.isCancellationRequested) {
         throw new ToolkitError(supplemetalContextFetchingTimeoutMsg, { cause: new CancellationError('timeout') })

src/codewhisperer/util/supplementalContext/editorFilesUtil.ts

@@ -35,19 +35,18 @@ export async function fetchSupplementalContext(
     const timesBeforeFetching = performance.now()
     const dependencyGraph = DependencyGraphFactory.getDependencyGraph(editor.document.languageId)
 
-    if (dependencyGraph === undefined) {
-        // This is a general check for language support of CW.
-        // We perform feature level language filtering later.
-        return undefined
-    }
+    const isUtg = await isTestFile(editor.document.uri.fsPath, {
+        languageId: editor.document.languageId,
+        dependencyGraph: dependencyGraph,
+        fileContent: editor.document.getText(),
+    })
 
-    const isUtg = await isTestFile(editor, dependencyGraph)
     let supplementalContextPromise: Promise<CodeWhispererSupplementalContextItem[] | undefined>
 
     if (isUtg) {
-        supplementalContextPromise = fetchSupplementalContextForTest(editor, dependencyGraph, cancellationToken)
+        supplementalContextPromise = fetchSupplementalContextForTest(editor, cancellationToken)
     } else {
-        supplementalContextPromise = fetchSupplementalContextForSrc(editor, dependencyGraph, cancellationToken)
+        supplementalContextPromise = fetchSupplementalContextForSrc(editor, cancellationToken)
     }
 
     return supplementalContextPromise