feat(codewhisperer): consider filedistance for crossfile context #3679

Before:
Crossfile file candidates are fetched by opened tabs without any sorting, e.g. files will be returned based on theirs orders in the IDE tabs (from the left to the right)

After:
Crossfile file candidates fetched will be sorted by their file distance ascendingly. Files with smaller file distance will be prioritized than others.

Author: Will-ShaoHua

src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts

@@ -13,6 +13,7 @@ import { CancellationError } from '../../../shared/utilities/timeoutUtils'
 import { CodeWhispererSupplementalContextItem } from './supplementalContextUtil'
 import { CodeWhispererUserGroupSettings } from '../userGroupUtil'
 import { isTestFile } from './codeParsingUtil'
+import { getFileDistance } from '../../../shared/filesystemUtilities'
 import { getOpenFilesInWindow } from '../../../shared/utilities/editorUtilities'
 
 type CrossFileSupportedLanguage =
@@ -201,10 +202,10 @@ function splitFileToChunks(filePath: string, chunkSize: number): Chunk[] {
 }
 
 /**
- * This function will return relevant cross files for the given editor file
+ * This function will return relevant cross files sorted by file distance for the given editor file
  * by referencing open files, imported files and same package files.
  */
-async function getCrossFileCandidates(editor: vscode.TextEditor): Promise<string[]> {
+export 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]
@@ -215,14 +216,28 @@ async function getCrossFileCandidates(editor: vscode.TextEditor): Promise<string
      * 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 => {
+    const unsortedCandidates = await getOpenFilesInWindow(async candidateFile => {
         return (
             targetFile !== candidateFile &&
             (path.extname(targetFile) === path.extname(candidateFile) ||
                 (dialects && dialects.has(path.extname(candidateFile)))) &&
             !(await isTestFile(candidateFile, { languageId: language }))
         )
     })
+
+    return unsortedCandidates
+        .map(candidate => {
+            return {
+                file: candidate,
+                fileDistance: getFileDistance(targetFile, candidate),
+            }
+        })
+        .sort((file1, file2) => {
+            return file1.fileDistance - file2.fileDistance
+        })
+        .map(fileToDistance => {
+            return fileToDistance.file
+        })
 }
 
 function throwIfCancelled(token: vscode.CancellationToken): void | never {

src/test/codewhisperer/util/crossFileContextUtil.test.ts

@@ -11,7 +11,15 @@ import * as crossFile from '../../../codewhisperer/util/supplementalContext/cros
 import { createMockTextEditor } from '../testUtil'
 import { CodeWhispererUserGroupSettings } from '../../../codewhisperer/util/userGroupUtil'
 import { UserGroup } from '../../../codewhisperer/models/constants'
-import { assertTabCount, closeAllEditors, createTestWorkspaceFolder, openATextEditorWithText } from '../../testUtil'
+import {
+    assertTabCount,
+    closeAllEditors,
+    createTestWorkspaceFolder,
+    openATextEditorWithText,
+    shuffleList,
+} from '../../testUtil'
+import { areEqual, normalize } from '../../../shared/utilities/pathUtils'
+import * as path from 'path'
 import { getMinVscodeVersion } from '../../../shared/vscode/env'
 
 const userGroupSettings = CodeWhispererUserGroupSettings.instance
@@ -49,6 +57,71 @@ describe('crossFileContextUtil', function () {
         })
     })
 
+    describe('getCrossFileCandidate', function () {
+        before(async function () {
+            this.timeout(60000)
+        })
+
+        beforeEach(async function () {
+            tempFolder = (await createTestWorkspaceFolder()).uri.fsPath
+        })
+
+        afterEach(async function () {
+            await closeAllEditors()
+        })
+
+        it('should return opened files, exclude test files and sorted ascendingly by file distance', async function () {
+            if (!shouldRunTheTest()) {
+                this.skip()
+            }
+
+            const targetFile = path.join('src', 'service', 'microService', 'CodeWhispererFileContextProvider.java')
+            const fileWithDistance3 = path.join('src', 'service', 'CodewhispererRecommendationService.java')
+            const fileWithDistance5 = path.join('src', 'util', 'CodeWhispererConstants.java')
+            const fileWithDistance6 = path.join('src', 'ui', 'popup', 'CodeWhispererPopupManager.java')
+            const fileWithDistance7 = path.join('src', 'ui', 'popup', 'components', 'CodeWhispererPopup.java')
+            const fileWithDistance8 = path.join(
+                'src',
+                'ui',
+                'popup',
+                'components',
+                'actions',
+                'AcceptRecommendationAction.java'
+            )
+            const testFile1 = path.join('test', 'service', 'CodeWhispererFileContextProviderTest.java')
+            const testFile2 = path.join('test', 'ui', 'CodeWhispererPopupManagerTest.java')
+
+            const expectedFilePaths = [
+                fileWithDistance3,
+                fileWithDistance5,
+                fileWithDistance6,
+                fileWithDistance7,
+                fileWithDistance8,
+            ]
+
+            const shuffledFilePaths = shuffleList(expectedFilePaths)
+
+            for (const filePath of shuffledFilePaths) {
+                await openATextEditorWithText('', filePath, tempFolder, { preview: false })
+            }
+
+            await openATextEditorWithText('', testFile1, tempFolder, { preview: false })
+            await openATextEditorWithText('', testFile2, tempFolder, { preview: false })
+            const editor = await openATextEditorWithText('', targetFile, tempFolder, { preview: false })
+
+            await assertTabCount(shuffledFilePaths.length + 3)
+
+            const actual = await crossFile.getCrossFileCandidates(editor)
+
+            assert.ok(actual.length === 5)
+            actual.forEach((actualFile, index) => {
+                const expectedFile = path.join(tempFolder, expectedFilePaths[index])
+                assert.strictEqual(normalize(expectedFile), normalize(actualFile))
+                assert.ok(areEqual(tempFolder, actualFile, expectedFile))
+            })
+        })
+    })
+
     describe('partial support - control group', function () {
         before(async function () {
             this.timeout(60000)

src/test/testUtil.ts

@@ -388,3 +388,17 @@ export function captureEventOnce<T>(event: vscode.Event<T>, timeout?: number): P
         }
     })
 }
+
+/**
+ * Shuffle a list, Fisher-Yates Sorting Algorithm
+ */
+export function shuffleList<T>(list: T[]): T[] {
+    const shuffledList = [...list]
+
+    for (let i = shuffledList.length - 1; i > 0; i--) {
+        const j = Math.floor(Math.random() * (i + 1))
+        ;[shuffledList[i], shuffledList[j]] = [shuffledList[j], shuffledList[i]]
+    }
+
+    return shuffledList
+}