From dea01125d694971446aba0628a73784ae21a4a98 Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Wed, 16 Dec 2020 15:48:21 -0800
Subject: [PATCH] FIR deserializer: keep SourceElement for more precise
 Fir2IrLazyClass's source

---
 .../kotlin/fir/deserialization/ClassDeserialization.kt        | 2 ++
 .../src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt      | 2 +-
 .../jetbrains/kotlin/fir/declarations/FirDeclarationUtil.kt   | 4 ++++
 .../box/annotations/typeAnnotations/typeAnnotationTarget6.kt  | 1 -
 4 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/ClassDeserialization.kt b/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/ClassDeserialization.kt
index 38bbaa68b48..5c0b05b084b 100644
--- a/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/ClassDeserialization.kt
+++ b/compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/ClassDeserialization.kt
@@ -203,6 +203,8 @@ fun deserializeClassToSymbol(
             context.annotationDeserializer.loadClassAnnotations(classProto, context.nameResolver)
 
         it.versionRequirementsTable = context.versionRequirementTable
+
+        it.sourceElement = containerSource
     }
 }
 
diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt
index 7fcc8ad800e..8e88c2b00f1 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt
@@ -45,7 +45,7 @@ class Fir2IrLazyClass(
     override lateinit var parent: IrDeclarationParent
 
     override val source: SourceElement
-        get() = SourceElement.NO_SOURCE
+        get() = fir.sourceElement ?: SourceElement.NO_SOURCE
 
     @ObsoleteDescriptorBasedAPI
     override val descriptor: ClassDescriptor
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationUtil.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationUtil.kt
index abeeaddb8cd..b0802632424 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationUtil.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationUtil.kt
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.SourceElement
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.declarations.builder.FirRegularClassBuilder
 import org.jetbrains.kotlin.fir.declarations.builder.FirTypeParameterBuilder
@@ -128,6 +129,9 @@ fun FirRegularClass.addDeclaration(declaration: FirDeclaration) {
     }
 }
 
+private object SourceElementKey : FirDeclarationDataKey()
+var FirRegularClass.sourceElement: SourceElement? by FirDeclarationDataRegistry.data(SourceElementKey)
+
 private object IsFromVarargKey : FirDeclarationDataKey()
 var FirProperty.isFromVararg: Boolean? by FirDeclarationDataRegistry.data(IsFromVarargKey)
 private object IsReferredViaField : FirDeclarationDataKey()
diff --git a/compiler/testData/codegen/box/annotations/typeAnnotations/typeAnnotationTarget6.kt b/compiler/testData/codegen/box/annotations/typeAnnotations/typeAnnotationTarget6.kt
index 0fa0a911574..a849dfc7adb 100644
--- a/compiler/testData/codegen/box/annotations/typeAnnotations/typeAnnotationTarget6.kt
+++ b/compiler/testData/codegen/box/annotations/typeAnnotations/typeAnnotationTarget6.kt
@@ -1,6 +1,5 @@
 // KOTLIN_CONFIGURATION_FLAGS: +JVM.EMIT_JVM_TYPE_ANNOTATIONS
 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND_FIR: JVM_IR
 
 // No virtual method getAnnotatedReturnType()Ljava/lang/reflect/AnnotatedType
 // IGNORE_BACKEND: ANDROID