From d46088548d3c9b54e6a20a4bd6e0eab6a6aea35a Mon Sep 17 00:00:00 2001 From: Ilya Kirillov Date: Wed, 20 May 2020 21:48:30 +0300 Subject: [PATCH] FIR IDE: Implement FIR reference resolve for some kinds of references * implement for: * KtForLoopInReference * KtInvokeFunctionReference * KtPropertyDelegationMethodsReference * KtDestructuringDeclarationReference --- .../kotlin/idea/fir/FirModuleResolveState.kt | 15 +++++-- .../kotlin/idea/fir/FirResolutionApi.kt | 5 +++ .../org/jetbrains/kotlin/idea/fir/FirUtils.kt | 14 +++++- .../kotlin/idea/references/FirKtReference.kt | 11 ++++- .../references/FirReferenceResolveHelper.kt | 42 +++++++++++++++--- .../idea/references/FirSimpleNameReference.kt | 12 ++++- .../KotlinFirReferenceContributor.kt | 4 ++ ...estructuringDeclarationReferenceFirImpl.kt | 30 +++++++++++++ .../idea/references/KtFirReferenceResolver.kt | 9 +++- .../references/KtForLoopInReferenceFirImpl.kt | 44 +++++++++++++++++++ .../KtInvokeFunctionReferenceFirImpl.kt | 32 ++++++++++++++ ...opertyDelegationMethodsReferenceFirImpl.kt | 40 +++++++++++++++++ 12 files changed, 242 insertions(+), 16 deletions(-) create mode 100644 idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtDestructuringDeclarationReferenceFirImpl.kt create mode 100644 idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtForLoopInReferenceFirImpl.kt create mode 100644 idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtInvokeFunctionReferenceFirImpl.kt create mode 100644 idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtPropertyDelegationMethodsReferenceFirImpl.kt diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirModuleResolveState.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirModuleResolveState.kt index ed1122e83ec..134a7770664 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirModuleResolveState.kt +++ b/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) { diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt index 8d25a071913..83d5c307d9e 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt +++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirResolutionApi.kt @@ -223,6 +223,11 @@ private class FirDesignatedBodyResolveTransformerForIDE( } } +inline fun KtElement.getOrBuildFirSafe( + state: FirModuleResolveState, + phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE +) = getOrBuildFir(state, phase) as? E + fun KtElement.getOrBuildFir( state: FirModuleResolveState, phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt index 38b7bb8edf3..bfbc5fa2a34 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt +++ b/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" -} \ No newline at end of file + return (resolvedCalleeSymbol as? FirNamedFunctionSymbol)?.fir?.name == OperatorNameConventions.INVOKE +} + +fun FirFunctionCall.getCalleeSymbol(): FirBasedSymbol<*>? = + calleeReference.getResolvedSymbolOfNameReference() + +fun FirReference.getResolvedSymbolOfNameReference(): FirBasedSymbol<*>? = + (this as? FirResolvedNamedReference)?.resolvedSymbol \ No newline at end of file diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirKtReference.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirKtReference.kt index e5056e0f86a..e48ffc119b2 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirKtReference.kt +++ b/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 + override val resolver get() = KtFirReferenceResolver } \ No newline at end of file diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt index e2f06b0bf28..4c8ed35dcd9 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt +++ b/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 { + internal fun resolveSimpleNameReference( + ref: KtSimpleNameReferenceFirImpl, + session: FirSession, + state: FirModuleResolveState + ): Collection { 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) { + // a() + // } + (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() } diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt index 247010308cf..d763ddb1d68 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt +++ b/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? { diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KotlinFirReferenceContributor.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KotlinFirReferenceContributor.kt index 612b3b52405..19233987cc9 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KotlinFirReferenceContributor.kt +++ b/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) } } } diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtDestructuringDeclarationReferenceFirImpl.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtDestructuringDeclarationReferenceFirImpl.kt new file mode 100644 index 00000000000..83274d0802b --- /dev/null +++ b/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 { + val fir = expression.getOrBuildFirSafe(state) ?: return emptyList() + val componentFunctionSymbol = (fir.initializer as? FirComponentCall)?.getCalleeSymbol() ?: return emptyList() + return listOfNotNull(componentFunctionSymbol.fir.findPsi(session)) + } +} diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt index 18f204e80f7..7b10e4f5688 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirReferenceResolver.kt +++ b/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 { class KotlinResolveResult(element: PsiElement) : PsiElementResolveResult(element) override fun resolve(ref: KtReference, incompleteCode: Boolean): Array { + 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() } } \ No newline at end of file diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtForLoopInReferenceFirImpl.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtForLoopInReferenceFirImpl.kt new file mode 100644 index 00000000000..4303d1de6bc --- /dev/null +++ b/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 { + val firLoop = expression.getOrBuildFirSafe(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() +} \ No newline at end of file diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtInvokeFunctionReferenceFirImpl.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtInvokeFunctionReferenceFirImpl.kt new file mode 100644 index 00000000000..d0b5b7e85bf --- /dev/null +++ b/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 { + val call = analysisSession.resolveCall(expression) ?: return emptyList() + if (call is VariableAsFunctionLikeCallInfo) { + return listOf(call.invokeFunction) + } + return emptyList() + } +} \ No newline at end of file diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtPropertyDelegationMethodsReferenceFirImpl.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtPropertyDelegationMethodsReferenceFirImpl.kt new file mode 100644 index 00000000000..656cee68a2a --- /dev/null +++ b/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 { + val property = (expression.parent as? KtElement)?.getOrBuildFirSafe(state) ?: return emptyList() + if (property.delegate == null) return emptyList() + val getValueSymbol = (property.getter?.singleStatementOfType()?.result as? FirFunctionCall)?.getCalleeSymbol() + val setValueSymbol = property.setter?.singleStatementOfType()?.getCalleeSymbol() + return listOfNotNull( + getValueSymbol?.fir?.findPsi(session), + setValueSymbol?.fir?.findPsi(session) + ) + } + + private inline fun FirPropertyAccessor.singleStatementOfType(): S? = + body?.statements?.singleOrNull() as? S +} \ No newline at end of file