From 30c35d79f4a39495752d3aca6169e830c6e2a5c0 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Fri, 20 Oct 2023 13:53:00 +0200
Subject: [PATCH] LL API: unwrap callable copies while determining their
 sources

In previous commit (K2: prefer derived class sources...)
we changed default source for callable copy (e.g. substitution override)
to its owner class. This can create an unexpected situation in LL API,
because normally KtCallableDeclaration is expected to be a source.
This commit introduces "compensation" on analysis API side.
---
 .../api/fir/lazy/resolve/FirLazyBodiesCalculator.kt    | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt
index 3a5af982c21..b04138916a9 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt
@@ -78,7 +78,7 @@ internal object FirLazyBodiesCalculator {
 
 private inline fun <reified T : FirDeclaration> revive(
     designation: FirDesignation,
-    psiFactory: (FirDesignation) -> PsiElement? = { it.target.psi }
+    psi: PsiElement? = designation.target.psi
 ): T {
     val session = designation.target.moduleData.session
 
@@ -86,7 +86,7 @@ private inline fun <reified T : FirDeclaration> revive(
         session = session,
         scopeProvider = session.kotlinScopeProvider,
         designation = designation,
-        rootNonLocalDeclaration = psiFactory(designation) as KtAnnotated,
+        rootNonLocalDeclaration = psi as KtAnnotated,
     ) as T
 }
 
@@ -157,7 +157,7 @@ private fun calculateLazyBodiesForFunction(designation: FirDesignation) {
     val simpleFunction = designation.target as FirSimpleFunction
     require(needCalculatingLazyBodyForFunction(simpleFunction))
 
-    val newSimpleFunction = revive<FirSimpleFunction>(designation)
+    val newSimpleFunction = revive<FirSimpleFunction>(designation, simpleFunction.unwrapFakeOverridesOrDelegated().psi)
 
     replaceLazyContractDescription(simpleFunction, newSimpleFunction)
     replaceLazyBody(simpleFunction, newSimpleFunction)
@@ -180,7 +180,7 @@ private fun calculateLazyBodyForProperty(designation: FirDesignation) {
     val firProperty = designation.target as FirProperty
     if (!needCalculatingLazyBodyForProperty(firProperty)) return
 
-    val newProperty = revive<FirProperty>(designation)
+    val newProperty = revive<FirProperty>(designation, firProperty.unwrapFakeOverridesOrDelegated().psi)
 
     firProperty.getter?.let { getter ->
         val newGetter = newProperty.getter!!
@@ -238,7 +238,7 @@ private fun calculateLazyBodiesForField(designation: FirDesignation) {
     val field = designation.target as FirField
     require(field.initializer is FirLazyExpression)
 
-    val newField = revive<FirField>(designation) { it.path.last().psi }
+    val newField = revive<FirField>(designation, designation.path.last().psi)
     field.replaceInitializer(newField.initializer)
 }