From 233a979932d7157770c7ab709fd4b87378a58f85 Mon Sep 17 00:00:00 2001
From: Dmitry Petrov <Dmitry.Petrov@jetbrains.com>
Date: Mon, 19 Sep 2016 14:07:21 +0300
Subject: [PATCH] Source element for file classes.

---
 .../backend/jvm/lower/JvmFileClassProvider.kt | 20 ++++++++------
 .../kotlin/psi2ir/PsiSourceManager.kt         | 27 ++++++++++---------
 2 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/JvmFileClassProvider.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/JvmFileClassProvider.kt
index 8eadc0c8f06..8f7ca01d182 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/JvmFileClassProvider.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/JvmFileClassProvider.kt
@@ -24,10 +24,10 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.SourceManager
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi2ir.Psi2IrTranslator
 import org.jetbrains.kotlin.psi2ir.PsiSourceManager
 import org.jetbrains.kotlin.psi2ir.generators.GeneratorContext
+import org.jetbrains.kotlin.resolve.source.KotlinSourceElement
 import java.lang.AssertionError
 import java.util.*
 
@@ -42,20 +42,24 @@ class JvmFileClassProvider : Psi2IrTranslator.PostprocessingStep {
     }
 
     private val fileClassInfoByFileEntry = HashMap<SourceManager.FileEntry, JvmFileClassInfo>()
+    private val sourceElementByFileEntry = HashMap<SourceManager.FileEntry, SourceElement>()
 
     private fun recordFileClassInfo(context: GeneratorContext, irFile: IrFile) {
-        val jvmFileClassInfo = context.sourceManager.getFileClassInfo(irFile.fileEntry)
+        context.sourceManager.getKtFile(irFile)?.let { ktFile ->
+            sourceElementByFileEntry[irFile.fileEntry] = KotlinSourceElement(ktFile)
+        }
+        val jvmFileClassInfo = context.sourceManager.getFileClassInfo(irFile) ?: return
         fileClassInfoByFileEntry[irFile.fileEntry] = jvmFileClassInfo
     }
 
-    private fun PsiSourceManager.getKtFile(fileEntry: SourceManager.FileEntry) =
-            getPsiFile(fileEntry as PsiSourceManager.PsiFileEntry) as KtFile
-
-    private fun PsiSourceManager.getFileClassInfo(fileEntry: SourceManager.FileEntry): JvmFileClassInfo =
-            JvmFileClassUtil.getFileClassInfoNoResolve(getKtFile(fileEntry))
+    private fun PsiSourceManager.getFileClassInfo(irFile: IrFile): JvmFileClassInfo? {
+        val file = getKtFile(irFile) ?: return null
+        return JvmFileClassUtil.getFileClassInfoNoResolve(file)
+    }
 
     fun createFileClassDescriptor(fileEntry: SourceManager.FileEntry, packageFragment: PackageFragmentDescriptor): FileClassDescriptor {
         val fileClassInfo = fileClassInfoByFileEntry[fileEntry] ?: throw AssertionError("No file class info for ${fileEntry.name})")
-        return FileClassDescriptorImpl(fileClassInfo.fileClassFqName.shortName(), packageFragment, SourceElement.NO_SOURCE)
+        val sourceElement = sourceElementByFileEntry[fileEntry] ?: SourceElement.NO_SOURCE
+        return FileClassDescriptorImpl(fileClassInfo.fileClassFqName.shortName(), packageFragment, sourceElement)
     }
 }
diff --git a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/PsiSourceManager.kt b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/PsiSourceManager.kt
index a3bf95a0a4a..09cc8510707 100644
--- a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/PsiSourceManager.kt
+++ b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/PsiSourceManager.kt
@@ -20,6 +20,7 @@ import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.ir.SourceManager
 import org.jetbrains.kotlin.ir.SourceRangeInfo
 import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.psi.KtFile
 import java.util.*
 
 class PsiSourceManager : SourceManager {
@@ -67,15 +68,15 @@ class PsiSourceManager : SourceManager {
         override fun toString(): String = getRecognizableName()
     }
 
-    private val fileEntriesByPsiFile = HashMap<PsiFile, PsiFileEntry>()
+    private val fileEntriesByKtFile = HashMap<KtFile, PsiFileEntry>()
     private val fileEntriesByIrFile = HashMap<IrFile, PsiFileEntry>()
-    private val psiFileByFileEntry = HashMap<PsiFileEntry, PsiFile>()
+    private val ktFileByFileEntry = HashMap<PsiFileEntry, KtFile>()
 
-    fun createFileEntry(psiFile: PsiFile): PsiFileEntry {
-        if (psiFile in fileEntriesByPsiFile) error("PsiFileEntry is already created for $psiFile")
-        val newEntry = PsiFileEntry(psiFile)
-        fileEntriesByPsiFile[psiFile] = newEntry
-        psiFileByFileEntry[newEntry] = psiFile
+    fun createFileEntry(ktFile: KtFile): PsiFileEntry {
+        if (ktFile in fileEntriesByKtFile) error("PsiFileEntry is already created for $ktFile")
+        val newEntry = PsiFileEntry(ktFile)
+        fileEntriesByKtFile[ktFile] = newEntry
+        ktFileByFileEntry[newEntry] = ktFile
         return newEntry
     }
 
@@ -83,14 +84,14 @@ class PsiSourceManager : SourceManager {
         fileEntriesByIrFile[irFile] = fileEntry
     }
 
-    fun getOrCreateFileEntry(psiFile: PsiFile): PsiFileEntry =
-            fileEntriesByPsiFile.getOrElse(psiFile) { createFileEntry(psiFile) }
+    fun getOrCreateFileEntry(ktFile: KtFile): PsiFileEntry =
+            fileEntriesByKtFile.getOrElse(ktFile) { createFileEntry(ktFile) }
 
-    fun getFileEntry(psiFile: PsiFile): PsiFileEntry? =
-            fileEntriesByPsiFile[psiFile]
+    fun getKtFile(fileEntry: PsiFileEntry): KtFile? =
+            ktFileByFileEntry[fileEntry]
 
-    fun getPsiFile(fileEntry: PsiFileEntry) =
-            psiFileByFileEntry[fileEntry]
+    fun getKtFile(irFile: IrFile): KtFile? =
+            (irFile.fileEntry as? PsiFileEntry)?.let { ktFileByFileEntry[it] }
 
     override fun getFileEntry(irFile: IrFile): SourceManager.FileEntry =
             fileEntriesByIrFile[irFile]!!