[FIR] Extract symbol from candidates for function contract DFA

^KT-61055 Fixed

analysis/low-level-api-fir/testData/lazyResolve/lazyProperty.txt

@@ -316,7 +316,9 @@ FILE: [ResolvedTo(IMPORTS)] lazyProperty.kt
         public final [ResolvedTo(STATUS)] [IsFromPrimaryConstructor=true] val value: R|T| = R|<local>/value|
             public [ResolvedTo(STATUS)] [ContainingClassKey=LazyDelegate] get(): R|T|
 
-        public final operator [ResolvedTo(STATUS)] fun getValue([ResolvedTo(STATUS)] thisRef: R|kotlin/Any?|, [ResolvedTo(STATUS)] property: R|kotlin/reflect/KProperty<*>|): R|T| { LAZY_BLOCK }
+        public final operator [ResolvedTo(CONTRACTS)] fun getValue([ResolvedTo(CONTRACTS)] thisRef: R|kotlin/Any?|, [ResolvedTo(CONTRACTS)] property: R|kotlin/reflect/KProperty<*>|): R|T| {
+            ^getValue value#
+        }
 
     }
     public final [ResolvedTo(CONTRACTS)] fun <[ResolvedTo(CONTRACTS)] T> lazy([ResolvedTo(CONTRACTS)] block: R|() -> T|): R|LazyDelegate<T>| {

analysis/low-level-api-fir/testData/lazyResolve/lazyPropertyScript.txt

@@ -419,7 +419,9 @@ FILE: [ResolvedTo(IMPORTS)] lazyPropertyScript.kts
             public final [ResolvedTo(STATUS)] [IsFromPrimaryConstructor=true] val value: R|T| = R|<local>/value|
                 public [ResolvedTo(STATUS)] [ContainingClassKey=LazyDelegate] get(): R|T|
 
-            public final operator [ResolvedTo(STATUS)] fun getValue([ResolvedTo(STATUS)] thisRef: R|kotlin/Any?|, [ResolvedTo(STATUS)] property: R|kotlin/reflect/KProperty<*>|): R|T| { LAZY_BLOCK }
+            public final operator [ResolvedTo(CONTRACTS)] fun getValue([ResolvedTo(CONTRACTS)] thisRef: R|kotlin/Any?|, [ResolvedTo(CONTRACTS)] property: R|kotlin/reflect/KProperty<*>|): R|T| {
+                ^getValue value#
+            }
 
         }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
+import org.jetbrains.kotlin.fir.references.symbol
 import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
@@ -29,10 +30,7 @@ import org.jetbrains.kotlin.fir.scopes.getFunctions
 import org.jetbrains.kotlin.fir.scopes.impl.toConeType
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.StandardClassIds
@@ -924,12 +922,10 @@ abstract class FirDataFlowAnalyzer(
         // contracts has no effect on non-body resolve stages
         if (!components.transformer.baseTransformerPhase.isBodyResolve) return
 
-        // TODO: Consider using something besides `toResolvedCallableSymbol` as the latter only works
-        //  for completed calls without candidates (KT-61055 for tracking)
         val callee = when (qualifiedAccess) {
-            is FirFunctionCall -> qualifiedAccess.toResolvedCallableSymbol()?.fir as? FirSimpleFunction
-            is FirQualifiedAccessExpression -> qualifiedAccess.calleeReference.toResolvedPropertySymbol()?.fir?.getter
-            is FirVariableAssignment -> qualifiedAccess.calleeReference?.toResolvedPropertySymbol()?.fir?.setter
+            is FirFunctionCall -> qualifiedAccess.calleeReference.symbol?.fir as? FirSimpleFunction
+            is FirQualifiedAccessExpression -> qualifiedAccess.calleeReference.symbol?.let { it.fir as? FirProperty }?.getter
+            is FirVariableAssignment -> qualifiedAccess.calleeReference?.symbol?.let { it.fir as? FirProperty }?.setter
             else -> null
         } ?: return
 

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt

@@ -8,16 +8,13 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.references.FirNamedReferenceWithCandidateBase
-import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
-import org.jetbrains.kotlin.fir.references.FirThisReference
+import org.jetbrains.kotlin.fir.references.symbol
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirSyntheticPropertySymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeTypeContext
-import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.unwrapFakeOverrides
 
 fun TypeStatement?.smartCastedType(context: ConeTypeContext, originalType: ConeKotlinType): ConeKotlinType =
@@ -39,8 +36,8 @@ fun FirOperation.isEq(): Boolean {
 @DfaInternals
 val FirElement.symbol: FirBasedSymbol<*>?
     get() = when (this) {
-        is FirResolvable -> symbol.unwrapFakeOverridesIfNecessary()
-        is FirVariableAssignment -> unwrapLValue()?.symbol
+        is FirResolvable -> calleeReference.symbol.unwrapFakeOverridesIfNecessary()
+        is FirVariableAssignment -> unwrapLValue()?.calleeReference?.symbol
         is FirDeclaration -> symbol.unwrapFakeOverridesIfNecessary()
         is FirWhenSubjectExpression -> whenRef.value.subject?.symbol
         is FirSafeCallExpression -> selector.symbol
@@ -62,15 +59,6 @@ private fun FirBasedSymbol<*>?.unwrapFakeOverridesIfNecessary(): FirBasedSymbol<
     return this.unwrapFakeOverrides()
 }
 
-@DfaInternals
-internal val FirResolvable.symbol: FirBasedSymbol<*>?
-    get() = when (val reference = calleeReference) {
-        is FirThisReference -> reference.boundSymbol
-        is FirResolvedNamedReference -> reference.resolvedSymbol
-        is FirNamedReferenceWithCandidateBase -> reference.candidateSymbol
-        else -> null
-    }
-
 @DfaInternals
 fun FirElement.unwrapElement(): FirElement = when (this) {
     is FirWhenSubjectExpression -> whenRef.value.let { it.subjectVariable ?: it.subject }?.unwrapElement() ?: this

compiler/fir/tree/src/org/jetbrains/kotlin/fir/references/FirReferenceUtils.kt

@@ -11,17 +11,24 @@ import org.jetbrains.kotlin.fir.symbols.impl.*
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
 
+val FirReference.symbol: FirBasedSymbol<*>?
+    get() = when (this) {
+        is FirThisReference -> boundSymbol
+        is FirResolvedNamedReference -> resolvedSymbol
+        is FirNamedReferenceWithCandidateBase -> candidateSymbol
+        else -> null
+    }
+
 val FirReference.resolved: FirResolvedNamedReference? get() = this as? FirResolvedNamedReference
 
 @Suppress("INVISIBLE_MEMBER", "INVISIBLE_REFERENCE",)
 inline fun <reified T : FirBasedSymbol<*>> FirReference.toResolvedSymbol(
     discardErrorReference: Boolean = false
 ): @kotlin.internal.NoInfer T? {
-    val resolvedReference = resolved ?: return null
-    if (discardErrorReference && resolvedReference is FirResolvedErrorReference) {
+    if (discardErrorReference && this is FirResolvedErrorReference) {
         return null
     }
-    return resolvedReference.resolvedSymbol as? T
+    return resolved?.resolvedSymbol as? T
 }
 
 fun FirReference.toResolvedBaseSymbol(discardErrorReference: Boolean = false): FirBasedSymbol<*>? {