From 8a15edf5e6493ae77ed0ef78b04a5bb73545dd37 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 19 May 2020 18:40:36 +0300
Subject: [PATCH] [IR] Extract IrDeclaration(Access).psiElement extension

---
 .../jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt   | 2 +-
 .../kotlin/backend/jvm/codegen/CoroutineCodegen.kt         | 4 ++--
 .../backend/jvm/codegen/IrSourceCompilerForInline.kt       | 2 +-
 .../jetbrains/kotlin/backend/jvm/codegen/irCodegenUtils.kt | 7 +++++++
 4 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt
index accfa3cf45b..3230c312f2a 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ClassCodegen.kt
@@ -83,7 +83,7 @@ abstract class ClassCodegen protected constructor(
             throw IllegalStateException("Generating class with invalid name '${type.className}': ${irClass.dump()}")
         }
         defineClass(
-            irClass.descriptor.psiElement,
+            irClass.psiElement,
             state.classFileVersion,
             irClass.flags,
             signature.name,
diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt
index 6f7ebc065d9..953b0adbc57 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt
@@ -44,9 +44,9 @@ internal fun MethodNode.acceptWithStateMachine(
     val languageVersionSettings = state.languageVersionSettings
     assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
     val element = if (irFunction.isSuspend)
-        irFunction.symbol.descriptor.psiElement ?: classCodegen.irClass.descriptor.psiElement
+        irFunction.psiElement ?: classCodegen.irClass.psiElement
     else
-        classCodegen.context.suspendLambdaToOriginalFunctionMap[classCodegen.irClass.attributeOwnerId]!!.symbol.descriptor.psiElement
+        classCodegen.context.suspendLambdaToOriginalFunctionMap[classCodegen.irClass.attributeOwnerId]!!.psiElement
     val visitor = CoroutineTransformerMethodVisitor(
         methodVisitor, access, name, desc, signature, exceptions.toTypedArray(),
         obtainClassBuilderForCoroutineState = obtainContinuationClassBuilder,
diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt
index 9c81a03ef46..1b0cf53cd68 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt
@@ -142,5 +142,5 @@ class IrSourceCompilerForInline(
             .report(SUSPENSION_POINT_INSIDE_MONITOR, stackTraceElement)
     }
 
-    private fun findElement() = (callElement.symbol.descriptor.original as? DeclarationDescriptorWithSource)?.source?.getPsi() as? KtElement
+    private fun findElement() = callElement.psiElement as? KtElement
 }
diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/irCodegenUtils.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/irCodegenUtils.kt
index 2820e7d6bcc..be47cfde0a5 100644
--- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/irCodegenUtils.kt
+++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/irCodegenUtils.kt
@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrGetEnumValue
+import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.types.IrSimpleType
@@ -425,3 +426,9 @@ val IrFunction.deprecationFlags: Int
         val propertyFlags = (this as? IrSimpleFunction)?.correspondingPropertySymbol?.owner?.deprecationFlags ?: 0
         return originFlags or propertyFlags or (this as IrAnnotationContainer).deprecationFlags
     }
+
+val IrDeclaration.psiElement: PsiElement?
+    get() = (descriptor as? DeclarationDescriptorWithSource)?.psiElement
+
+val IrMemberAccessExpression.psiElement: PsiElement?
+    get() = (symbol.descriptor.original as? DeclarationDescriptorWithSource)?.psiElement