[FIR] Introduce extension unwrapFakeOverridesOrDelegated()

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirDeserializedDeclarationSourceProvider.kt

@@ -19,7 +19,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isData
 import org.jetbrains.kotlin.fir.declarations.utils.isEnumClass
 import org.jetbrains.kotlin.fir.resolve.getContainingClass
-import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
+import org.jetbrains.kotlin.fir.delegatedWrapperData
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.unwrapFakeOverrides
 import org.jetbrains.kotlin.name.ClassId

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirOverrideInfoProvider.kt

@@ -22,7 +22,7 @@ import org.jetbrains.kotlin.fir.originalForIntersectionOverrideAttr
 import org.jetbrains.kotlin.fir.originalForSubstitutionOverride
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.fir.resolve.toSymbol
-import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
+import org.jetbrains.kotlin.fir.delegatedWrapperData
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.util.ImplementationStatus
 

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirImplementationMismatchChecker.kt

@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
 import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.getDirectOverriddenMembers
 import org.jetbrains.kotlin.fir.scopes.getDirectOverriddenProperties
-import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
+import org.jetbrains.kotlin.fir.delegatedWrapperData
 import org.jetbrains.kotlin.fir.scopes.impl.toConeType
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeErrorType

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirNotImplementedOverrideChecker.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.fir.containingClassLookupTag
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.utils.*
+import org.jetbrains.kotlin.fir.delegatedWrapperData
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.scopes.MemberWithBaseScope
 import org.jetbrains.kotlin.fir.scopes.getDirectOverriddenMembersWithBaseScope

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/DelegatedMemberGenerator.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.modality
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.scopes.*
-import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
+import org.jetbrains.kotlin.fir.delegatedWrapperData
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirDelegatedMemberScope.kt

@@ -7,16 +7,17 @@ package org.jetbrains.kotlin.fir.scopes.impl
 
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.DelegatedWrapperData
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.modality
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
+import org.jetbrains.kotlin.fir.delegatedWrapperData
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.scope
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.scopes.*
-import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.ConeFlexibleType
@@ -206,18 +207,6 @@ var FirCallableDeclaration.multipleDelegatesWithTheSameSignature: Boolean? by Fi
 val FirCallableSymbol<*>.multipleDelegatesWithTheSameSignature: Boolean?
     get() = fir.multipleDelegatesWithTheSameSignature
 
-private object DelegatedWrapperDataKey : FirDeclarationDataKey()
-class DelegatedWrapperData<D : FirCallableDeclaration>(
-    val wrapped: D,
-    val containingClass: ConeClassLikeLookupTag,
-    val delegateField: FirField,
-)
-var <D : FirCallableDeclaration>
-        D.delegatedWrapperData: DelegatedWrapperData<D>? by FirDeclarationDataRegistry.data(DelegatedWrapperDataKey)
-
-val <D : FirCallableDeclaration> FirCallableSymbol<out D>.delegatedWrapperData: DelegatedWrapperData<D>?
-    get() = fir.delegatedWrapperData
-
 
 // From the definition of function interfaces in the Java specification (pt. 9.8):
 // "methods that are members of I that do not have the same signature as any public instance method of the class Object"

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -29,7 +29,6 @@ import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
-import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
 import org.jetbrains.kotlin.fir.scopes.impl.importedFromObjectOrStaticData
 import org.jetbrains.kotlin.fir.scopes.processOverriddenFunctions
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -556,14 +555,12 @@ fun FirFunction.getAsForbiddenNamedArgumentsTarget(
     if (hasStableParameterNames) return null
 
     return when (origin) {
-        FirDeclarationOrigin.Delegated -> delegatedWrapperData?.wrapped?.getAsForbiddenNamedArgumentsTarget(session)
-
         FirDeclarationOrigin.ImportedFromObjectOrStatic ->
             importedFromObjectOrStaticData?.original?.getAsForbiddenNamedArgumentsTarget(session)
 
-        FirDeclarationOrigin.IntersectionOverride, FirDeclarationOrigin.SubstitutionOverride -> {
+        FirDeclarationOrigin.IntersectionOverride, FirDeclarationOrigin.SubstitutionOverride, FirDeclarationOrigin.Delegated -> {
             var result: ForbiddenNamedArgumentsTarget? =
-                originalIfFakeOverride()?.getAsForbiddenNamedArgumentsTarget(session) ?: return null
+                unwrapFakeOverridesOrDelegated().getAsForbiddenNamedArgumentsTarget(session) ?: return null
             originScope?.processOverriddenFunctions(symbol as FirNamedFunctionSymbol) {
                 if (it.fir.getAsForbiddenNamedArgumentsTarget(session) == null) {
                     result = null

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFull
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.runContractResolveForLocalClass
 import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
 import org.jetbrains.kotlin.fir.scopes.fakeOverrideSubstitution
-import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
+import org.jetbrains.kotlin.fir.delegatedWrapperData
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirSyntheticPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase

compiler/fir/tree/src/org/jetbrains/kotlin/fir/ClassMembers.kt

@@ -78,6 +78,9 @@ inline val <reified S : FirCallableSymbol<*>> S.baseForIntersectionOverride: S?
 inline fun <reified D : FirCallableDeclaration> D.originalIfFakeOverride(): D? =
     originalForSubstitutionOverride ?: baseForIntersectionOverride
 
+inline fun <reified D : FirCallableDeclaration> D.originalIfFakeOverrideOrDelegated(): D? =
+    originalForSubstitutionOverride ?: baseForIntersectionOverride ?: delegatedWrapperData?.wrapped
+
 inline fun <reified S : FirCallableSymbol<*>> S.originalIfFakeOverride(): S? =
     fir.originalIfFakeOverride()?.symbol as S?
 
@@ -100,6 +103,15 @@ inline fun <reified D : FirCallableDeclaration> D.unwrapFakeOverrides(): D {
     } while (true)
 }
 
+inline fun <reified D : FirCallableDeclaration> D.unwrapFakeOverridesOrDelegated(): D {
+    var current = this
+
+    do {
+        val next = current.originalIfFakeOverrideOrDelegated() ?: return current
+        current = next
+    } while (true)
+}
+
 inline fun <reified D : FirCallableDeclaration> D.unwrapSubstitutionOverrides(): D {
     var current = this
 
@@ -168,3 +180,16 @@ private object IsCatchParameterProperty : FirDeclarationDataKey()
 
 var FirProperty.isCatchParameter: Boolean? by FirDeclarationDataRegistry.data(IsCatchParameterProperty)
 
+private object DelegatedWrapperDataKey : FirDeclarationDataKey()
+
+class DelegatedWrapperData<D : FirCallableDeclaration>(
+    val wrapped: D,
+    val containingClass: ConeClassLikeLookupTag,
+    val delegateField: FirField,
+)
+
+var <D : FirCallableDeclaration>
+        D.delegatedWrapperData: DelegatedWrapperData<D>? by FirDeclarationDataRegistry.data(DelegatedWrapperDataKey)
+
+val <D : FirCallableDeclaration> FirCallableSymbol<out D>.delegatedWrapperData: DelegatedWrapperData<D>?
+    get() = fir.delegatedWrapperData

compiler/visualizer/render-fir/src/org/jetbrains/kotlin/compiler/visualizer/FirVisualizer.kt

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.renderer.ConeTypeRendererForDebugging
 import org.jetbrains.kotlin.fir.resolve.defaultType
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.toFirRegularClassSymbol
-import org.jetbrains.kotlin.fir.scopes.impl.delegatedWrapperData
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
@@ -462,24 +461,6 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
             }
         }
 
-        private inline fun <reified D : FirCallableDeclaration> D.originalIfFakeOverrideOrDelegated(): D? {
-            return when (origin) {
-                FirDeclarationOrigin.SubstitutionOverride, FirDeclarationOrigin.IntersectionOverride -> originalIfFakeOverride()
-                FirDeclarationOrigin.Delegated -> delegatedWrapperData!!.wrapped
-                else -> this
-            }
-        }
-
-        private inline fun <reified D : FirCallableDeclaration> D.getOriginal(): D {
-            var current = this
-            var next = this.originalIfFakeOverrideOrDelegated() ?: this
-            while (current != next) {
-                current = next
-                next = current.originalIfFakeOverrideOrDelegated() ?: current
-            }
-            return current
-        }
-
         private fun renderImplicitReceiver(symbol: FirBasedSymbol<*>, psi: PsiElement?) {
             val receiverType = (symbol.fir as? FirCallableDeclaration)?.dispatchReceiverType ?: return
             val implicitReceiverIndex = implicitReceivers.indexOf(receiverType)
@@ -493,7 +474,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
         }
 
         private fun renderField(field: FirField, data: StringBuilder) {
-            val original = field.getOriginal()
+            val original = field.unwrapFakeOverridesOrDelegated()
             if (original.isVal) data.append("val ") else if (original.isVar) data.append("var ")
 
             data.append("(")
@@ -537,7 +518,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
         }
 
         private fun renderPropertySymbol(symbol: FirPropertySymbol, data: StringBuilder) {
-            val fir = symbol.fir.getOriginal()
+            val fir = symbol.fir.unwrapFakeOverridesOrDelegated()
             data.append(if (fir.isVar) "var" else "val").append(" ")
             renderListInTriangles(fir.typeParameters, data, withSpace = true)
 
@@ -559,7 +540,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
 
         private fun renderFunctionSymbol(symbol: FirNamedFunctionSymbol, data: StringBuilder, call: FirFunctionCall? = null) {
             data.append("fun ")
-            val fir = symbol.fir.getOriginal()
+            val fir = symbol.fir.unwrapFakeOverridesOrDelegated()
             renderListInTriangles(fir.typeParameters, data, true)
 
             val id = getSymbolId(symbol)
@@ -625,7 +606,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
 
         override fun visitConstructor(constructor: FirConstructor, data: StringBuilder) {
             renderConstructorSymbol(constructor.symbol, data)
-            visitValueParameters(constructor.getOriginal().valueParameters, data)
+            visitValueParameters(constructor.unwrapFakeOverridesOrDelegated().valueParameters, data)
         }
 
         override fun visitTypeParameterRef(typeParameterRef: FirTypeParameterRef, data: StringBuilder) {