[Test] Create KtFiles with proper virtual files

This is needed to properly report diagnostics in KAPT3 tests, which
  uses psiElement.containingFile as origin

compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/SourceFileProvider.kt

@@ -6,8 +6,14 @@
 package org.jetbrains.kotlin.test.services
 
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.text.StringUtilRt
+import com.intellij.openapi.vfs.StandardFileSystems
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiManager
+import com.intellij.testFramework.LightVirtualFile
 import org.jetbrains.kotlin.KtInMemoryTextSourceFile
 import org.jetbrains.kotlin.fir.lightTree.LightTree2Fir
+import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.sourceFiles.LightTreeFile
 import org.jetbrains.kotlin.test.model.TestFile
@@ -86,9 +92,13 @@ class SourceFileProviderImpl(val testServices: TestServices, override val prepro
     }
 }
 
-fun SourceFileProvider.getKtFileForSourceFile(testFile: TestFile, project: Project): KtFile {
+fun SourceFileProvider.getKtFileForSourceFile(testFile: TestFile, project: Project, findViaVfs: Boolean = false): KtFile {
-// TODO
+    if (findViaVfs) {
-//    return TestCheckerUtil.createCheckAndReturnPsiFile(
+        val realFile = getRealFileForSourceFile(testFile)
+        StandardFileSystems.local().findFileByPath(realFile.path)
+            ?.let { PsiManager.getInstance(project).findFile(it) as? KtFile }
+            ?.let { return it }
+    }
     return KtTestUtil.createFile(
         testFile.name,
         getContentOfSourceFile(testFile),
@@ -96,10 +106,10 @@ fun SourceFileProvider.getKtFileForSourceFile(testFile: TestFile, project: Proje
     )
 }
 
-fun SourceFileProvider.getKtFilesForSourceFiles(testFiles: Collection<TestFile>, project: Project): Map<TestFile, KtFile> {
+fun SourceFileProvider.getKtFilesForSourceFiles(testFiles: Collection<TestFile>, project: Project, findViaVfs: Boolean = false): Map<TestFile, KtFile> {
     return testFiles.mapNotNull {
         if (!it.isKtFile) return@mapNotNull null
-        it to getKtFileForSourceFile(it, project)
+        it to getKtFileForSourceFile(it, project, findViaVfs)
     }.toMap()
 }