FIR IDE: Implement search of psi elements for deserialized Kotlin declarations

Fix "FIR IDE: Implement search of psi elements for deserialized Kotlin declarations"

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

@@ -55,6 +55,10 @@ class RawFirBuilder(
         return file.accept(Visitor(), Unit) as FirFile
     }
 
+    fun buildTypeReference(reference: KtTypeReference): FirTypeRef {
+        return reference.accept(Visitor(), Unit) as FirTypeRef
+    }
+
     override fun PsiElement.toFirSourceElement(): FirPsiSourceElement<*> {
         return this.toFirPsiSourceElement()
     }

generators/build.gradle.kts

@@ -38,6 +38,7 @@ dependencies {
     testCompile(projectTests(":compiler:cli"))
     testCompile(projectTests(":idea:idea-maven"))
     testCompile(projectTests(":idea:idea-fir"))
+    testCompile(projectTests(":idea:idea-frontend-fir"))
     testCompile(projectTests(":j2k"))
     testCompile(projectTests(":nj2k"))
     if (Ide.IJ()) {

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt

@@ -85,6 +85,7 @@ import org.jetbrains.kotlin.idea.editor.quickDoc.AbstractQuickDocProviderTest
 import org.jetbrains.kotlin.idea.filters.AbstractKotlinExceptionFilterTest
 import org.jetbrains.kotlin.idea.fir.AbstractFirLazyResolveTest
 import org.jetbrains.kotlin.idea.fir.AbstractFirMultiModuleResolveTest
+import org.jetbrains.kotlin.idea.fir.AbstractKtDeclarationAndFirDeclarationEqualityChecker
 import org.jetbrains.kotlin.idea.folding.AbstractKotlinFoldingTest
 import org.jetbrains.kotlin.idea.hierarchy.AbstractHierarchyTest
 import org.jetbrains.kotlin.idea.hierarchy.AbstractHierarchyWithLibTest
@@ -918,10 +919,16 @@ fun main(args: Array<String>) {
             }
         }
 
-        testGroup("idea/idea-fir/tests", "idea/testData") {
-            testClass<AbstractFirMultiModuleResolveTest> {
-                model("fir/multiModule", recursive = false, extension = null)
-            }
+        testGroup("idea/idea-frontend-fir/tests", "idea/idea-frontend-fir/testData") {
+        testClass<AbstractKtDeclarationAndFirDeclarationEqualityChecker> {
+            model("ktDeclarationAndFirDeclarationEqualityChecker")
+        }
+    }
+
+    testGroup("idea/idea-fir/tests", "idea/testData") {
+        testClass<AbstractFirMultiModuleResolveTest> {
+            model("fir/multiModule", recursive = false, extension = null)
+        }
 
             testClass<AbstractFirHighlightingTest> {
             model("highlighter")

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

@@ -0,0 +1,145 @@
+/*
+ * 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.fir
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.idea.stubindex.KotlinFullClassNameIndex
+import org.jetbrains.kotlin.idea.stubindex.KotlinTopLevelFunctionFqnNameIndex
+import org.jetbrains.kotlin.idea.stubindex.KotlinTopLevelTypeAliasFqNameIndex
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtNamedFunction
+
+
+//todo introduce LibraryModificationTracker based cache?
+object FirIdeDeserializedDeclarationSourceProvider {
+    fun findPsi(fir: FirElement, project: Project): PsiElement? {
+        return when (fir) {
+            is FirSimpleFunction -> provideSourceForFunction(fir, project)
+            is FirProperty -> provideSourceForProperty(fir, project)
+            is FirClass<*> -> provideSourceForClass(fir, project)
+            is FirTypeAlias -> provideSourceForTypeAlias(fir, project)
+            is FirConstructor -> provideSourceForConstructor(fir, project)
+            else -> null
+        }
+    }
+
+    private fun provideSourceForFunction(
+        function: FirSimpleFunction,
+        project: Project
+    ): PsiElement? {
+        val candidates = if (function.isTopLevel) {
+            KotlinTopLevelFunctionFqnNameIndex.getInstance().get(
+                function.symbol.callableId.asFqNameForDebugInfo().asString(),
+                project,
+                function.scope(project)
+            ).filter(KtNamedFunction::isCompiled)
+        } else {
+            function.containingKtClass(project)?.body?.functions
+                ?.filter { it.name == function.name.asString() && it.isCompiled() }
+                .orEmpty()
+        }
+
+        return function.chooseCorrespondingPsi(candidates)
+    }
+
+    private fun provideSourceForProperty(property: FirProperty, project: Project): PsiElement? {
+        val candidates = if (property.isTopLevel) {
+            KotlinTopLevelFunctionFqnNameIndex.getInstance().get(
+                property.symbol.callableId.asFqNameForDebugInfo().asString(),
+                project,
+                property.scope(project)
+            )
+        } else {
+            property.containingKtClass(project)?.declarations
+                ?.filter { it.name == property.name.asString() }
+                .orEmpty()
+        }
+
+        return candidates.firstOrNull(KtElement::isCompiled)
+    }
+
+    private fun provideSourceForClass(klass: FirClass<*>, project: Project): PsiElement? =
+        classByClassId(klass.symbol.classId, klass.scope(project), project)
+
+    private fun provideSourceForTypeAlias(alias: FirTypeAlias, project: Project): PsiElement? {
+        val candidates = KotlinTopLevelTypeAliasFqNameIndex.getInstance().get(
+            alias.symbol.classId.asStringForUsingInIndexes(),
+            project,
+            alias.scope(project)
+        )
+        return candidates.firstOrNull(KtElement::isCompiled)
+    }
+
+    private fun provideSourceForConstructor(
+        constructor: FirConstructor,
+        project: Project
+    ): PsiElement? {
+        val containingKtClass = constructor.containingKtClass(project) ?: return null
+        if (constructor.isPrimary) return containingKtClass.primaryConstructor
+
+        return constructor.chooseCorrespondingPsi(containingKtClass.secondaryConstructors)
+    }
+
+    private fun FirFunction<*>.chooseCorrespondingPsi(
+        candidates: Collection<KtFunction>
+    ): KtFunction? {
+        if (candidates.isEmpty()) return null
+        for (candidate in candidates) {
+            assert(candidate.isCompiled()) {
+                "Candidate should be decompiled from metadata because it should have fqName types as we don't use resolve here"
+            }
+            if (KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(candidate, this)) {
+                return candidate
+            }
+        }
+        return null
+    }
+
+    private fun FirDeclaration.scope(project: Project): GlobalSearchScope {
+        return GlobalSearchScope.allScope(project)
+        /* TODO:
+         val session = session as? FirLibrarySession
+         return session?.scope ?: GlobalSearchScope.allScope(project)*/
+    }
+
+    private fun FirCallableDeclaration<*>.containingKtClass(project: Project): KtClassOrObject? =
+        symbol.callableId.classId?.let { classByClassId(it, scope(project), project) }
+
+    private fun classByClassId(classId: ClassId, scope: GlobalSearchScope, project: Project): KtClassOrObject? {
+        val fqName = classId.asStringForUsingInIndexes().let { classIdMapping[it] ?: it }
+        return KotlinFullClassNameIndex.getInstance().get(
+            fqName,
+            project,
+            scope
+        ).firstOrNull(KtElement::isCompiled)
+    }
+
+    private val FirCallableDeclaration<*>.isTopLevel
+        get() = symbol.callableId.className == null
+
+    private fun ClassId.asStringForUsingInIndexes() = asString().replace('/', '.')
+
+    private val classIdMapping = (0..23).associate { i ->
+        "kotlin.Function$i" to "kotlin.jvm.functions.Function$i"
+    }
+}
+
+private fun KtElement.isCompiled(): Boolean = containingKtFile.isCompiled
+
+fun FirElement.findPsi(project: Project): PsiElement? =
+    psi ?: FirIdeDeserializedDeclarationSourceProvider.findPsi(this, project)
+
+fun FirElement.findPsi(session: FirSession): PsiElement? =
+    findPsi(session.sessionProvider!!.project)

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

@@ -0,0 +1,185 @@
+/*
+ * 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.fir
+
+import org.jetbrains.annotations.TestOnly
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.builder.RawFirBuilder
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.KtTypeReference
+import org.jetbrains.kotlin.types.Variance
+
+// TODO replace with structural type comparision?
+object KtDeclarationAndFirDeclarationEqualityChecker {
+    fun representsTheSameDeclaration(psi: KtFunction, fir: FirFunction<*>): Boolean {
+        if ((fir.receiverTypeRef != null) != (psi.receiverTypeReference != null)) return false
+        if (fir.receiverTypeRef != null
+            && !isTheSameTypes(
+                psi.receiverTypeReference!!,
+                fir.receiverTypeRef!!,
+                fir.session,
+                isVararg = false
+            )
+        ) {
+            return false
+        }
+        if (fir.valueParameters.size != psi.valueParameters.size) return false
+        fir.valueParameters.zip(psi.valueParameters) { expectedParameter, candidateParameter ->
+            if (expectedParameter.name.toString() != candidateParameter.name) return false
+            if (expectedParameter.isVararg != candidateParameter.isVarArg) return false
+            val candidateParameterType = candidateParameter.typeReference ?: return false
+            if (!isTheSameTypes(
+                    candidateParameterType,
+                    expectedParameter.returnTypeRef,
+                    fir.session,
+                    isVararg = expectedParameter.isVararg
+                )
+            ) {
+                return false
+            }
+        }
+        return true
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun FirTypeRef.renderTypeAsKotlinType(isVararg: Boolean = false): String {
+        val rendered = when (this) {
+            is FirResolvedTypeRef -> type.renderTypeAsKotlinType()
+            is FirUserTypeRef -> {
+                val renderedQualifier = qualifier.joinToString(separator = ".") { part ->
+                    buildString {
+                        append(part.name)
+                        if (part.typeArguments.isNotEmpty()) {
+                            part.typeArguments.joinTo(this, prefix = "<", postfix = ">") { it.renderTypeAsKotlinType() }
+                        }
+                    }
+                }
+                if (isMarkedNullable) "$renderedQualifier?" else renderedQualifier
+            }
+            is FirFunctionTypeRef -> {
+                val classId = if (isSuspend) {
+                    KotlinBuiltIns.getSuspendFunctionClassId(parametersCount)
+                } else {
+                    KotlinBuiltIns.getFunctionClassId(parametersCount)
+                }
+                buildString {
+                    append(classId.asSingleFqName().toString())
+                    val parameters = buildList {
+                        receiverTypeRef?.let(::add)
+                        valueParameters.mapTo(this) { it.returnTypeRef }
+                        returnTypeRef.let(::add)
+                    }
+                    if (parameters.isNotEmpty()) {
+                        append(parameters.joinToString(prefix = "<", postfix = ">") { it.renderTypeAsKotlinType() })
+                    }
+                }
+            }
+            else -> error("Invalid type reference $this")
+        }
+        return if (isVararg) {
+            "kotlin.Array<out $rendered>"
+        } else {
+            rendered
+        }
+    }
+
+    private fun FirTypeProjection.renderTypeAsKotlinType() = when (this) {
+        is FirStarProjection -> "*"
+        is FirTypeProjectionWithVariance -> buildString {
+            append(variance.label)
+            if (variance != Variance.INVARIANT) {
+                append(" ")
+            }
+            append(typeRef.renderTypeAsKotlinType())
+        }
+        else -> error("Invalid type projection $this")
+    }
+
+    private fun isTheSameTypes(
+        psiTypeReference: KtTypeReference,
+        coneTypeReference: FirTypeRef,
+        session: FirSession,
+        isVararg: Boolean
+    ): Boolean =
+        psiTypeReference.toKotlinTypReference(session).renderTypeAsKotlinType(isVararg) == coneTypeReference.renderTypeAsKotlinType()
+
+    private fun KtTypeReference.toKotlinTypReference(session: FirSession): FirTypeRef {
+        // Maybe resolve all types here to not to work with FirTypeRef directly
+        return RawFirBuilder(session, DummyScopeProvider, stubMode = true).buildTypeReference(this)
+    }
+
+    private fun ConeKotlinType.renderTypeAsKotlinType(): String {
+        val rendered = when (this) {
+            is ConeClassLikeType -> buildString {
+                append(lookupTag.classId.asString())
+                if (typeArguments.isNotEmpty()) {
+                    append(typeArguments.joinToString(prefix = "<", postfix = ">", separator = ", ") { it.renderTypeAsKotlinType() })
+                }
+            }
+            is ConeTypeVariableType -> lookupTag.name.asString()
+            is ConeLookupTagBasedType -> lookupTag.name.asString()
+            is ConeFlexibleType -> {
+                // Can be present as return type
+                "${lowerBound.renderTypeAsKotlinType()}..${upperBound.renderTypeAsKotlinType()}"
+            }
+            else -> error("Type $this should not be present in Kotlin declaration")
+        }.replace('/', '.')
+        return rendered + nullability.suffix
+    }
+
+    private fun ConeTypeProjection.renderTypeAsKotlinType(): String = when (this) {
+        ConeStarProjection -> "*"
+        is ConeKotlinTypeProjectionIn -> "in ${type.renderTypeAsKotlinType()}"
+        is ConeKotlinTypeProjectionOut -> "out ${type.renderTypeAsKotlinType()}"
+        is ConeKotlinType -> renderTypeAsKotlinType()
+    }
+
+    @TestOnly
+    internal fun renderPsi(ktFunction: KtFunction, session: FirSession): String = buildString {
+        appendLine("receiver: ${ktFunction.receiverTypeReference?.toKotlinTypReference(session)?.renderTypeAsKotlinType()}")
+        ktFunction.valueParameters.forEach { parameter ->
+            appendLine("${parameter.name}: ${parameter.typeReference?.toKotlinTypReference(session)?.renderTypeAsKotlinType()}")
+        }
+        appendLine("return: ${ktFunction.typeReference?.toKotlinTypReference(session)?.renderTypeAsKotlinType()}")
+    }
+
+    @TestOnly
+    internal fun renderFir(firFunction: FirFunction<*>): String = buildString {
+        appendLine("receiver: ${firFunction.receiverTypeRef?.renderTypeAsKotlinType()}")
+        firFunction.valueParameters.forEach { parameter ->
+            appendLine("${parameter.name}: ${parameter.returnTypeRef.renderTypeAsKotlinType()}")
+        }
+        appendLine("return: ${firFunction.returnTypeRef.renderTypeAsKotlinType()}")
+    }
+
+    private object DummyScopeProvider : FirScopeProvider() {
+        override fun getUseSiteMemberScope(klass: FirClass<*>, useSiteSession: FirSession, scopeSession: ScopeSession): FirScope {
+            shouldNotBeCalled()
+        }
+
+        override fun getStaticMemberScopeForCallables(
+            klass: FirClass<*>,
+            useSiteSession: FirSession,
+            scopeSession: ScopeSession
+        ): FirScope? {
+            shouldNotBeCalled()
+        }
+
+        override fun getNestedClassifierScope(klass: FirClass<*>, useSiteSession: FirSession, scopeSession: ScopeSession): FirScope? {
+            shouldNotBeCalled()
+        }
+
+        private fun shouldNotBeCalled(): Nothing =
+            error("Should not be called in RawFirBuilder while converting KtTypeReference")
+    }
+}

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

@@ -17,18 +17,15 @@ import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.idea.fir.firResolveState
-import org.jetbrains.kotlin.idea.fir.getOrBuildFir
+import org.jetbrains.kotlin.idea.fir.*
 import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtPackageDirective
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 
 object FirReferenceResolveHelper {
     fun FirResolvedTypeRef.toTargetPsi(session: FirSession): PsiElement? {
         val type = type as? ConeLookupTagBasedType ?: return null
-        return (type.lookupTag.toSymbol(session) as? AbstractFirBasedSymbol<*>)?.fir?.psi
+        return (type.lookupTag.toSymbol(session) as? AbstractFirBasedSymbol<*>)?.fir?.findPsi(session)
     }
 
     fun ClassId.toTargetPsi(session: FirSession, calleeReference: FirReference? = null): PsiElement? {
@@ -38,11 +35,11 @@ object FirReferenceResolveHelper {
                 val callee = calleeReference.resolvedSymbol.fir as? FirCallableMemberDeclaration
                 // TODO: check callee owner directly?
                 if (callee !is FirConstructor && callee?.isStatic != true) {
-                    classLikeDeclaration.companionObject?.let { return it.psi }
+                    classLikeDeclaration.companionObject?.let { return it.findPsi(session) }
                 }
             }
         }
-        return classLikeDeclaration?.psi
+        return classLikeDeclaration?.findPsi(session)
     }
 
     fun FirReference.toTargetPsi(session: FirSession): PsiElement? {
@@ -51,10 +48,10 @@ object FirReferenceResolveHelper {
                 resolvedSymbol.fir.psi
             }
             is FirResolvedCallableReference -> {
-                resolvedSymbol.fir.psi
+                resolvedSymbol.fir.findPsi(session)
             }
             is FirThisReference -> {
-                boundSymbol?.fir?.psi
+                boundSymbol?.fir?.findPsi(session)
             }
             is FirSuperReference -> {
                 (superTypeRef as? FirResolvedTypeRef)?.toTargetPsi(session)
@@ -116,15 +113,20 @@ object FirReferenceResolveHelper {
                 }
                 val name = fir.importedName ?: return emptyList()
                 val symbolProvider = session.firSymbolProvider
-                return symbolProvider.getTopLevelCallableSymbols(fir.packageFqName, name).mapNotNull { it.fir.psi } +
-                        listOfNotNull(symbolProvider.getClassLikeSymbolByFqName(ClassId(fir.packageFqName, name))?.fir?.psi)
+                return symbolProvider.getTopLevelCallableSymbols(fir.packageFqName, name)
+                    .mapNotNull { it.fir.findPsi(session) } +
+                        listOfNotNull(
+                            symbolProvider
+                                .getClassLikeSymbolByFqName(ClassId(fir.packageFqName, name))
+                                ?.fir?.findPsi(session)
+                        )
             }
             is FirFile -> {
                 if (expression.getNonStrictParentOfType<KtPackageDirective>() != null) {
                     // Special: package reference in the middle of package directive
                     return listOf(expression)
                 }
-                return listOfNotNull(fir.psi)
+                return listOfNotNull(fir.findPsi(session))
             }
             is FirArrayOfCall -> {
                 // We can't yet find PsiElement for arrayOf, intArrayOf, etc.

idea/idea-frontend-fir/testData/ktDeclarationAndFirDeclarationEqualityChecker/extensionMethods.kt

@@ -0,0 +1,7 @@
+fun kotlin.String.fromString(what: kotlin.Int, oops: kotlin.collections.Map<MyType<kotlin.Boolean>, kotlin.collections.MutableList<kotlin.Boolean>>): kotlin.Unit {}
+
+fun <T : kotlin.Number >kotlin.String.foo(what: kotlin.Int): kotlin.String { TODO() }
+
+fun <T> kotlin.String.withWhere(what: kotlin.Int): T where T : kotlin.Number, T: kotlin.Comparable<kotlin.Number> { TODO() }
+
+class MyType<in F>

idea/idea-frontend-fir/testData/ktDeclarationAndFirDeclarationEqualityChecker/functionalTypes.kt

@@ -0,0 +1 @@
+fun ((kotlin.String) -> kotlin.Int?).foo(other: kotlin.String.() -> kotlin.Unit?): kotlin.Function2<kotlin.Int, kotlin.String, kotlin.Unit> { TODO() }

idea/idea-frontend-fir/testData/ktDeclarationAndFirDeclarationEqualityChecker/inClass.kt

@@ -0,0 +1,11 @@
+class A<S> {
+    fun x(a: kotlin.Int, b: kotlin.String): Unit {}
+
+    fun <T> x(a: kotlin.Int, b: java.lang.String): kotlin.Int { TODO() }
+
+    fun <T, Q> x(a: kotlin.collections.Map<T, kotlin.collections.Map<in S, out Q>>, b: kotlin.collections.List<kotlin.Byte>): kotlin.Int { TODO() }
+
+    fun <T, Q> kotlin.collections.List<kotlin.collections.Map<T, Q>>.x(f: MyType<S>): kotlin.Int { TODO() }
+}
+
+class MyType<in F>

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/AbstractKtDeclarationAndFirDeclarationEqualityChecker.kt

@@ -0,0 +1,37 @@
+/*
+ * 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.fir
+
+import com.intellij.openapi.util.io.FileUtil
+import com.intellij.rt.execution.junit.FileComparisonFailure
+import junit.framework.Assert
+import junit.framework.ComparisonFailure
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtFunction
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
+import java.io.File
+
+abstract class AbstractKtDeclarationAndFirDeclarationEqualityChecker : KotlinLightCodeInsightFixtureTestCase() {
+    protected fun doTest(path: String) {
+        val file = File(path)
+        val ktFile = myFixture.configureByText(file.name, FileUtil.loadFile(file)) as KtFile
+
+        val resolveState = ktFile.firResolveState()
+        ktFile.forEachDescendantOfType<KtFunction> { ktFunction ->
+            val firFunction = ktFunction.getOrBuildFir(resolveState) as FirFunction<*>
+            if (!KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(ktFunction, firFunction)) {
+                throw FileComparisonFailure(
+                    /* message= */          null,
+                                            KtDeclarationAndFirDeclarationEqualityChecker.renderPsi(ktFunction, firFunction.session),
+                                            KtDeclarationAndFirDeclarationEqualityChecker.renderFir(firFunction),
+                    /* expectedFilePath= */ null
+                )
+            }
+        }
+    }
+}

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/KtDeclarationAndFirDeclarationEqualityCheckerGenerated.java

@@ -0,0 +1,45 @@
+/*
+ * 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.fir;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("idea/idea-frontend-fir/testData/ktDeclarationAndFirDeclarationEqualityChecker")
+@TestDataPath("$PROJECT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+public class KtDeclarationAndFirDeclarationEqualityCheckerGenerated extends AbstractKtDeclarationAndFirDeclarationEqualityChecker {
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+    }
+
+    public void testAllFilesPresentInKtDeclarationAndFirDeclarationEqualityChecker() throws Exception {
+        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/ktDeclarationAndFirDeclarationEqualityChecker"), Pattern.compile("^(.+)\\.kt$"), null, true);
+    }
+
+    @TestMetadata("extensionMethods.kt")
+    public void testExtensionMethods() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/ktDeclarationAndFirDeclarationEqualityChecker/extensionMethods.kt");
+    }
+
+    @TestMetadata("functionalTypes.kt")
+    public void testFunctionalTypes() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/ktDeclarationAndFirDeclarationEqualityChecker/functionalTypes.kt");
+    }
+
+    @TestMetadata("inClass.kt")
+    public void testInClass() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/ktDeclarationAndFirDeclarationEqualityChecker/inClass.kt");
+    }
+}

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/highlighter/KotlinBeforeResolveHighlightingPass.kt

@@ -9,6 +9,7 @@ import com.intellij.codeHighlighting.*
 import com.intellij.codeInsight.daemon.impl.AnnotationHolderImpl
 import com.intellij.codeInsight.daemon.impl.HighlightInfo
 import com.intellij.codeInsight.daemon.impl.UpdateHighlightersUtil
+import com.intellij.lang.annotation.AnnotationHolder
 import com.intellij.lang.annotation.AnnotationSession
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
@@ -32,7 +33,7 @@ class KotlinBeforeResolveHighlightingPass(
     override fun doCollectInformation(progress: ProgressIndicator) {
         val annotationHolder = AnnotationHolderImpl(AnnotationSession(file))
         val visitor = BeforeResolveHighlightingVisitor(annotationHolder)
-        val extensions = EP_NAME.extensionList
+        val extensions = EP_NAME.extensionList.map { it.createVisitor(annotationHolder) }
         file.accept(object : PsiRecursiveElementVisitor() {
             override fun visitElement(element: PsiElement) {
                 super.visitElement(element)
@@ -72,6 +73,10 @@ class KotlinBeforeResolveHighlightingPass(
     }
 
     companion object {
-        val EP_NAME = ExtensionPointName.create<HighlightingVisitor>("org.jetbrains.kotlin.beforeResolveHighlightingVisitor")
+        val EP_NAME = ExtensionPointName.create<BeforeResolveHighlightingExtension>("org.jetbrains.kotlin.beforeResolveHighlightingVisitor")
     }
 }
+
+interface BeforeResolveHighlightingExtension {
+    fun createVisitor(holder: AnnotationHolder): HighlightingVisitor
+}

idea/idea-frontend-independent/src/org/jetbrains/kotlin/idea/highlighter/KotlinBeforeResolveHighlightingPass.kt.192

@@ -17,6 +17,7 @@ import com.intellij.lang.annotation.AnnotationHolder
 import com.intellij.openapi.components.ProjectComponent
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.project.DumbAware
 import com.intellij.psi.PsiElement
@@ -35,7 +36,7 @@ class KotlinBeforeResolveHighlightingPass(
     override fun doCollectInformation(progress: ProgressIndicator) {
         val annotationHolder = AnnotationHolderImpl(AnnotationSession(file))
         val visitor = BeforeResolveHighlightingVisitor(annotationHolder)
-        val extensions = EP_NAME.extensionList
+        val extensions = EP_NAME.extensionList.map { it.createVisitor(annotationHolder) }
         file.accept(object : PsiRecursiveElementVisitor() {
             override fun visitElement(element: PsiElement) {
                 super.visitElement(element)
@@ -72,7 +73,11 @@ class KotlinBeforeResolveHighlightingPass(
         }
     }
 
-     companion object {
-        val EP_NAME = ExtensionPointName.create<HighlightingVisitor>("org.jetbrains.kotlin.beforeResolveHighlightingVisitor")
+    companion object {
+        val EP_NAME = ExtensionPointName.create<BeforeResolveHighlightingExtension>("org.jetbrains.kotlin.beforeResolveHighlightingVisitor")
     }
 }
+
+interface BeforeResolveHighlightingExtension {
+    fun createVisitor(holder: AnnotationHolder): HighlightingVisitor
+}