[LL FIR] rebind returnTarget for delegate accessors during psi2fir

It is safe to rebind returnTarget to the original declaration during
psi2fir as it doesn't affect the resolution. We already have the same
logic for regular functions/property accessors.
This change is crucial in the context of parallel implicit type
resolution because it affects CFG builder in a bad way in the case
of on-air resolution

^KT-56551

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -292,15 +292,6 @@ private fun rebindReturnExpression(returnExpression: FirStatement, newTarget: Fi
         withFirEntry("expression", returnExpression)
     }
 
-    val target = returnExpression.target
-    requireWithAttachment(target.labeledElement == oldTarget, { "Unexpected return target" }) {
-        withFirSymbolEntry("newTarget", newTarget.propertySymbol)
-        withFirSymbolEntry("oldTarget", oldTarget.propertySymbol)
-        withFirEntry("target", target.labeledElement)
-    }
-
-    target.bind(newTarget)
-
     val functionCall = returnExpression.result
     rebindFunctionCall(functionCall, newTarget, oldTarget)
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
-import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
 import org.jetbrains.kotlin.analysis.utils.errors.withPsiEntry
 import org.jetbrains.kotlin.fir.*
@@ -32,6 +31,7 @@ import org.jetbrains.kotlin.name.NameUtils
 import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.hasExpectModifier
+import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 
 internal class RawFirNonLocalDeclarationBuilder private constructor(
     session: FirSession,
@@ -103,8 +103,26 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
     }
 
     override fun bindFunctionTarget(target: FirFunctionTarget, function: FirFunction) {
-        val rewrittenTarget = functionsToRebind?.firstOrNull { it.realPsi == function.realPsi } ?: function
+        super.bindFunctionTarget(target, computeRebindTarget(function) ?: function)
-        super.bindFunctionTarget(target, rewrittenTarget)
+    }
+
+    /**
+     * @return [FirFunction] if another function should be used instead of [function] for [FirFunctionTarget]
+     *
+     * @see bindFunctionTarget
+     * @see functionsToRebind
+     */
+    private fun computeRebindTarget(function: FirFunction): FirFunction? {
+        if (functionsToRebind.isNullOrEmpty()) return null
+        val realPsi = function.realPsi
+        if (realPsi != null) {
+            return functionsToRebind.firstOrNull { it.realPsi == realPsi }
+        }
+
+        val accessor = function as? FirPropertyAccessor ?: return null
+        val accessorPsi = accessor.psi ?: return null
+
+        return functionsToRebind.firstOrNull { it is FirPropertyAccessor && it.isGetter == accessor.isGetter && it.psi == accessorPsi }
     }
 
     override fun addCapturedTypeParameters(

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt

@@ -59,7 +59,10 @@ open class PsiRawFirBuilder(
     val baseScopeProvider: FirScopeProvider,
     bodyBuildingMode: BodyBuildingMode = BodyBuildingMode.NORMAL,
 ) : AbstractRawFirBuilder<PsiElement>(session) {
-    protected open fun bindFunctionTarget(target: FirFunctionTarget, function: FirFunction) = target.bind(function)
+    protected open fun bindFunctionTarget(target: FirFunctionTarget, function: FirFunction) {
+        target.bind(function)
+    }
+
     protected open fun FirFunctionBuilder.additionalFunctionInit() {}
     protected open fun FirPropertyBuilder.additionalPropertyInit() {}
     protected open fun FirPropertyAccessorBuilder.additionalPropertyAccessorInit() {}
@@ -2080,6 +2083,7 @@ open class PsiRawFirBuilder(
                                 isExtension = receiverTypeReference != null,
                                 lazyDelegateExpression = lazyDelegateExpression,
                                 lazyBodyForGeneratedAccessors = lazyBody,
+                                ::bindFunctionTarget,
                             )
                         }
                     }

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/ConversionUtils.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.fir.contracts.FirLegacyRawContractDescription
 import org.jetbrains.kotlin.fir.contracts.builder.buildLegacyRawContractDescription
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirReceiverParameter
 import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.declarations.builder.*
@@ -334,6 +335,7 @@ fun <T> FirPropertyBuilder.generateAccessorsByDelegate(
     isExtension: Boolean,
     lazyDelegateExpression: FirLazyExpression? = null,
     lazyBodyForGeneratedAccessors: FirLazyBlock? = null,
+    bindFunction: (target: FirFunctionTarget, function: FirFunction) -> Unit = FirFunctionTarget::bind,
 ) {
     if (delegateBuilder == null) return
     val delegateFieldSymbol = FirDelegateFieldSymbol(symbol.callableId).also {
@@ -481,10 +483,11 @@ fun <T> FirPropertyBuilder.generateAccessorsByDelegate(
             }
             propertySymbol = this@generateAccessorsByDelegate.symbol
         }.also {
-            returnTarget.bind(it)
+            bindFunction(returnTarget, it)
             it.initContainingClassAttr(context)
         }
     }
+
     if (isVar && (setter == null || setter is FirDefaultPropertyAccessor)) {
         val annotations = setter?.annotations
         val returnTarget = FirFunctionTarget(null, isLambda = false)
@@ -549,7 +552,7 @@ fun <T> FirPropertyBuilder.generateAccessorsByDelegate(
             }
             propertySymbol = this@generateAccessorsByDelegate.symbol
         }.also {
-            returnTarget.bind(it)
+            bindFunction(returnTarget, it)
             it.initContainingClassAttr(context)
         }
     }