FIR IDE: Implement FIR reference resolve for some kinds of references

* implement for:
   * KtForLoopInReference
   * KtInvokeFunctionReference
   * KtPropertyDelegationMethodsReference
   * KtDestructuringDeclarationReference

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirModuleResolveState.kt

@@ -18,6 +18,8 @@ import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
+import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.idea.caches.project.LibrarySourceInfo
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.psi.KtElement
@@ -27,13 +29,18 @@ interface FirModuleResolveState {
     val sessionProvider: FirProjectSessionProvider
 
     fun getSession(psi: KtElement): FirSession {
-        val moduleInfo = psi.getModuleInfo() as ModuleSourceInfo
+        val moduleInfo = psi.getModuleInfo()
         return getSession(psi.project, moduleInfo)
     }
 
-    fun getSession(project: Project, moduleInfo: ModuleSourceInfo): FirSession {
+    fun getSession(project: Project, moduleInfo: IdeaModuleInfo): FirSession {
         sessionProvider.getSession(moduleInfo)?.let { return it }
-        return synchronized(moduleInfo.module) {
+        val lock = when (moduleInfo) {
+            is ModuleSourceInfo -> moduleInfo.module
+            is LibrarySourceInfo -> moduleInfo.library
+            else -> TODO(moduleInfo.toString())
+        }
+        return synchronized(lock) {
             val session = sessionProvider.getSession(moduleInfo) ?: FirIdeJavaModuleBasedSession.create(
                 project, moduleInfo, sessionProvider, moduleInfo.contentScope()
             ).also { moduleBasedSession ->
@@ -89,7 +96,7 @@ class FirModuleResolveStateImpl(override val sessionProvider: FirProjectSessionP
     }
 
     override fun record(psi: KtElement, fir: FirElement) {
-        cache[psi] = fir
+        cache.putIfAbsent(psi, fir)
     }
 
     override fun record(psi: KtElement, diagnostic: Diagnostic) {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt

@@ -223,6 +223,11 @@ private class FirDesignatedBodyResolveTransformerForIDE(
     }
 }
 
+inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
+    state: FirModuleResolveState,
+    phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
+) = getOrBuildFir(state, phase) as? E
+
 fun KtElement.getOrBuildFir(
     state: FirModuleResolveState,
     phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt

@@ -7,11 +7,21 @@ package org.jetbrains.kotlin.idea.fir
 
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.idea.references.KtReference
+import org.jetbrains.kotlin.util.OperatorNameConventions
 
 fun FirFunctionCall.isImplicitFunctionCall(): Boolean {
     if (dispatchReceiver !is FirQualifiedAccessExpression) return false
     val resolvedCalleeSymbol = (calleeReference as? FirResolvedNamedReference)?.resolvedSymbol
-    return (resolvedCalleeSymbol as? FirNamedFunctionSymbol)?.fir?.name?.asString() == "invoke"
-}
+    return (resolvedCalleeSymbol as? FirNamedFunctionSymbol)?.fir?.name == OperatorNameConventions.INVOKE
+}
+
+fun FirFunctionCall.getCalleeSymbol(): FirBasedSymbol<*>? =
+    calleeReference.getResolvedSymbolOfNameReference()
+
+fun FirReference.getResolvedSymbolOfNameReference(): FirBasedSymbol<*>? =
+    (this as? FirResolvedNamedReference)?.resolvedSymbol

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirKtReference.kt

@@ -5,8 +5,17 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import com.intellij.psi.impl.source.resolve.ResolveCache
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.idea.fir.FirModuleResolveState
+import org.jetbrains.kotlin.idea.frontend.api.fir.AnalysisSessionFirImpl
 
 interface FirKtReference : KtReference {
+    fun getResolvedToPsi(
+        analysisSession: AnalysisSessionFirImpl,
+        session: FirSession,
+        state: FirModuleResolveState
+    ): Collection<PsiElement>
+
     override val resolver get() = KtFirReferenceResolver
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt

@@ -8,9 +8,10 @@ package org.jetbrains.kotlin.idea.references
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.*
+import org.jetbrains.kotlin.fir.resolve.calls.SyntheticPropertySymbol
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
@@ -45,7 +46,18 @@ object FirReferenceResolveHelper {
     fun FirReference.toTargetPsi(session: FirSession): PsiElement? {
         return when (this) {
             is FirResolvedNamedReference -> {
-                resolvedSymbol.fir.psi
+                val targetFir = when (val symbol = resolvedSymbol) {
+                    is SyntheticPropertySymbol -> {
+                        val syntheticProperty = symbol.fir as FirSyntheticProperty
+                        if (syntheticProperty.getter.delegate.symbol.callableId == symbol.accessorId) {
+                            syntheticProperty.getter.delegate
+                        } else {
+                            syntheticProperty.setter!!.delegate
+                        }
+                    }
+                    else -> symbol.fir
+                }
+                targetFir.findPsi(session)
             }
             is FirResolvedCallableReference -> {
                 resolvedSymbol.fir.findPsi(session)
@@ -62,13 +74,26 @@ object FirReferenceResolveHelper {
         }
     }
 
-    fun resolveToPsiElements(ref: AbstractKtReference<*>): Collection<PsiElement> {
+    internal fun resolveSimpleNameReference(
+        ref: KtSimpleNameReferenceFirImpl,
+        session: FirSession,
+        state: FirModuleResolveState
+    ): Collection<PsiElement> {
         val expression = ref.expression
-        val state = expression.firResolveState()
-        val session = state.getSession(expression)
         when (val fir = expression.getOrBuildFir(state)) {
             is FirResolvable -> {
-                return listOfNotNull(fir.calleeReference.toTargetPsi(session))
+                val calleeReference =
+                    if (fir is FirFunctionCall
+                        && fir.isImplicitFunctionCall()
+                        && expression is KtNameReferenceExpression
+                    ) {
+                        // we are resolving implicit invoke call, like
+                        // fun foo(a: () -> Unit) {
+                        //     <expression>a</expression>()
+                        // }
+                        (fir.dispatchReceiver as FirQualifiedAccessExpression).calleeReference
+                    } else fir.calleeReference
+                return listOfNotNull(calleeReference.toTargetPsi(session))
             }
             is FirResolvedTypeRef -> {
                 return listOfNotNull(fir.toTargetPsi(session))
@@ -132,6 +157,11 @@ object FirReferenceResolveHelper {
                 // We can't yet find PsiElement for arrayOf, intArrayOf, etc.
                 return emptyList()
             }
+            is FirReturnExpression -> {
+                return if (expression is KtLabelReferenceExpression) {
+                    listOfNotNull(fir.target.labeledElement.findPsi(session))
+                } else emptyList()
+            }
             is FirErrorNamedReference -> {
                 return emptyList()
             }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt

@@ -6,7 +6,9 @@
 package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.PsiElement
-import com.intellij.psi.impl.source.resolve.ResolveCache
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.idea.fir.FirModuleResolveState
+import org.jetbrains.kotlin.idea.frontend.api.fir.AnalysisSessionFirImpl
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtImportAlias
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
@@ -14,12 +16,18 @@ import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 internal class KtSimpleNameReferenceFirImpl(
     expression: KtSimpleNameExpression
 ) : KtSimpleNameReference(expression), FirKtReference {
+    override fun getResolvedToPsi(
+        analysisSession: AnalysisSessionFirImpl,
+        session: FirSession,
+        state: FirModuleResolveState
+    ) = FirReferenceResolveHelper.resolveSimpleNameReference(this, session, state)
+
     override fun doCanBeReferenceTo(candidateTarget: PsiElement): Boolean {
         return true // TODO
     }
 
     override fun isReferenceToWithoutExtensionChecking(candidateTarget: PsiElement): Boolean {
-        return resolve() == candidateTarget
+        return resolve() == candidateTarget //todo
     }
 
     override fun handleElementRename(newElementName: String): PsiElement? {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KotlinFirReferenceContributor.kt

@@ -9,6 +9,10 @@ class KotlinFirReferenceContributor : KotlinReferenceProviderContributor {
     override fun registerReferenceProviders(registrar: KotlinPsiReferenceRegistrar) {
         with(registrar) {
             registerProvider(factory = ::KtSimpleNameReferenceFirImpl)
+            registerProvider(factory = ::KtForLoopInReferenceFirImpl)
+            registerProvider(factory = ::KtInvokeFunctionReferenceFirImpl)
+            registerProvider(factory = ::KtPropertyDelegationMethodsReferenceFirImpl)
+            registerProvider(factory = ::KtDestructuringDeclarationReferenceFirImpl)
         }
     }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtDestructuringDeclarationReferenceFirImpl.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.expressions.FirComponentCall
+import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.frontend.api.fir.AnalysisSessionFirImpl
+import org.jetbrains.kotlin.psi.KtDestructuringDeclarationEntry
+
+class KtDestructuringDeclarationReferenceFirImpl(
+    element: KtDestructuringDeclarationEntry
+) : KtDestructuringDeclarationReference(element), FirKtReference {
+    override fun canRename(): Boolean = false //todo
+
+    override fun getResolvedToPsi(
+        analysisSession: AnalysisSessionFirImpl,
+        session: FirSession,
+        state: FirModuleResolveState
+    ): Collection<PsiElement> {
+        val fir = expression.getOrBuildFirSafe<FirProperty>(state) ?: return emptyList()
+        val componentFunctionSymbol = (fir.initializer as? FirComponentCall)?.getCalleeSymbol() ?: return emptyList()
+        return listOfNotNull(componentFunctionSymbol.fir.findPsi(session))
+    }
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt

@@ -9,13 +9,20 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementResolveResult
 import com.intellij.psi.ResolveResult
 import com.intellij.psi.impl.source.resolve.ResolveCache
+import org.jetbrains.kotlin.idea.fir.firResolveState
+import org.jetbrains.kotlin.idea.frontend.api.fir.AnalysisSessionFirImpl
 
 object KtFirReferenceResolver : ResolveCache.PolyVariantResolver<KtReference> {
     class KotlinResolveResult(element: PsiElement) : PsiElementResolveResult(element)
 
     override fun resolve(ref: KtReference, incompleteCode: Boolean): Array<ResolveResult> {
+        check(ref is FirKtReference) { "reference should be FirKtReference, but was ${ref::class}" }
         check(ref is AbstractKtReference<*>) { "reference should be AbstractKtReference, but was ${ref::class}" }
-        val resolveToPsiElements = FirReferenceResolveHelper.resolveToPsiElements(ref)
+        val expression = ref.expression
+        val state = expression.firResolveState()
+        val session = state.getSession(expression)
+        val analysisSession = AnalysisSessionFirImpl(state)
+        val resolveToPsiElements = ref.getResolvedToPsi(analysisSession, session, state)
         return resolveToPsiElements.map { KotlinResolveResult(it) }.toTypedArray()
     }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtForLoopInReferenceFirImpl.kt

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.expressions.FirWhileLoop
+import org.jetbrains.kotlin.idea.fir.FirModuleResolveState
+import org.jetbrains.kotlin.idea.fir.findPsi
+import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
+import org.jetbrains.kotlin.idea.fir.getResolvedSymbolOfNameReference
+import org.jetbrains.kotlin.idea.frontend.api.fir.AnalysisSessionFirImpl
+import org.jetbrains.kotlin.psi.KtForExpression
+
+open class KtForLoopInReferenceFirImpl(expression: KtForExpression) : KtForLoopInReference(expression), FirKtReference {
+    override fun getResolvedToPsi(
+        analysisSession: AnalysisSessionFirImpl,
+        session: FirSession,
+        state: FirModuleResolveState
+    ): Collection<PsiElement> {
+        val firLoop = expression.getOrBuildFirSafe<FirWhileLoop>(state) ?: return emptyList()
+        val condition = firLoop.condition as? FirFunctionCall
+        val iterator = run {
+            val callee = (condition?.explicitReceiver as? FirQualifiedAccessExpression)?.calleeReference
+            (callee?.getResolvedSymbolOfNameReference()?.fir as? FirProperty)?.getInitializerFunctionCall()
+        }
+        val hasNext = condition?.calleeReference?.getResolvedSymbolOfNameReference()
+        val next = (firLoop.block.statements.firstOrNull() as? FirProperty?)?.getInitializerFunctionCall()
+        return listOfNotNull(
+            iterator?.fir?.findPsi(session),
+            hasNext?.fir?.findPsi(session),
+            next?.fir?.findPsi(session),
+        )
+    }
+
+    private fun FirProperty.getInitializerFunctionCall() =
+        (initializer as? FirFunctionCall)?.calleeReference?.getResolvedSymbolOfNameReference()
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtInvokeFunctionReferenceFirImpl.kt

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.idea.fir.FirModuleResolveState
+import org.jetbrains.kotlin.idea.frontend.api.VariableAsFunctionLikeCallInfo
+import org.jetbrains.kotlin.idea.frontend.api.fir.AnalysisSessionFirImpl
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtExpression
+
+class KtInvokeFunctionReferenceFirImpl(expression: KtCallExpression) : KtInvokeFunctionReference(expression), FirKtReference {
+    override fun doRenameImplicitConventionalCall(newName: String?): KtExpression {
+        TODO("Not yet implemented")
+    }
+
+    override fun getResolvedToPsi(
+        analysisSession: AnalysisSessionFirImpl,
+        session: FirSession,
+        state: FirModuleResolveState
+    ): Collection<PsiElement> {
+        val call = analysisSession.resolveCall(expression) ?: return emptyList()
+        if (call is VariableAsFunctionLikeCallInfo) {
+            return listOf(call.invokeFunction)
+        }
+        return emptyList()
+    }
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtPropertyDelegationMethodsReferenceFirImpl.kt

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.frontend.api.fir.AnalysisSessionFirImpl
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtPropertyDelegate
+
+class KtPropertyDelegationMethodsReferenceFirImpl(
+    element: KtPropertyDelegate
+) : KtPropertyDelegationMethodsReference(element), FirKtReference {
+    override fun getResolvedToPsi(
+        analysisSession: AnalysisSessionFirImpl,
+        session: FirSession,
+        state: FirModuleResolveState
+    ): Collection<PsiElement> {
+        val property = (expression.parent as? KtElement)?.getOrBuildFirSafe<FirProperty>(state) ?: return emptyList()
+        if (property.delegate == null) return emptyList()
+        val getValueSymbol = (property.getter?.singleStatementOfType<FirReturnExpression>()?.result as? FirFunctionCall)?.getCalleeSymbol()
+        val setValueSymbol = property.setter?.singleStatementOfType<FirFunctionCall>()?.getCalleeSymbol()
+        return listOfNotNull(
+            getValueSymbol?.fir?.findPsi(session),
+            setValueSymbol?.fir?.findPsi(session)
+        )
+    }
+
+    private inline fun <reified S : FirStatement> FirPropertyAccessor.singleStatementOfType(): S? =
+        body?.statements?.singleOrNull() as? S
+}