From 38ae16e9d32241e64ace4a78a60c37cb23fc3291 Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Fri, 2 Jun 2023 14:27:26 +0000
Subject: [PATCH] [AA] explicitly search for sources of synthetic functions
 declarations

these synthetic classes are found by `FirBuiltinSyntheticFunctionInterfaceProvider`,
which isn't based on stubs and has no explicit source

This fixes tests in IJ repo, such as uast based tests:
`FirLegacyUastResolveEverythingTestGenerated#testLambdaReturn`

Merge-request: KT-MR-10406
Merged-by: Anna Kozlova <Anna.Kozlova@jetbrains.com>
---
 ...yntheticFunctionInterfaceSourceProvider.kt | 56 +++++++++++++++++++
 .../kotlin/analysis/api/fir/psiUtils.kt       |  4 +-
 2 files changed, 58 insertions(+), 2 deletions(-)
 create mode 100644 analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirSyntheticFunctionInterfaceSourceProvider.kt

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirSyntheticFunctionInterfaceSourceProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirSyntheticFunctionInterfaceSourceProvider.kt
new file mode 100644
index 00000000000..ec443ebbaca
--- /dev/null
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirSyntheticFunctionInterfaceSourceProvider.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2010-2023 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.analysis.api.fir
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.search.ProjectScope
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
+import org.jetbrains.kotlin.analysis.project.structure.KtBuiltinsModule
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
+import org.jetbrains.kotlin.fir.containingClassLookupTag
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlin.psi.KtClassOrObject
+
+//todo introduce LibraryModificationTracker based cache?
+internal object FirSyntheticFunctionInterfaceSourceProvider {
+    fun findPsi(fir: FirDeclaration): PsiElement? {
+        return when (fir) {
+            is FirSimpleFunction -> provideSourceForInvokeFunction(fir)
+            is FirClass -> provideSourceForFunctionClass(fir)
+            else -> null
+        }
+    }
+
+    private fun provideSourceForInvokeFunction(function: FirSimpleFunction): PsiElement? {
+        val classId = function.containingClassLookupTag()?.classId ?: return null
+        val classOrObject = classByClassId(classId, function.llFirSession.ktModule) ?: return null
+        return classOrObject.declarations.singleOrNull()
+    }
+
+    private fun provideSourceForFunctionClass(klass: FirClass): PsiElement? {
+        return classByClassId(klass.symbol.classId, klass.llFirSession.ktModule)
+    }
+
+    private fun classByClassId(classId: ClassId, ktModule: KtModule): KtClassOrObject? {
+        val project = ktModule.project
+        val correctedClassId = classIdMapping[classId] ?: return null
+        require(ktModule is KtBuiltinsModule) {
+            "Expected builtin module but found $ktModule"
+        }
+        return project.createDeclarationProvider(ProjectScope.getLibrariesScope(project), ktModule)
+            .getAllClassesByClassId(correctedClassId)
+            .firstOrNull { it.containingKtFile.isCompiled }
+    }
+
+    private val classIdMapping = (0..23).associate { i ->
+        StandardClassIds.FunctionN(i) to ClassId(FqName("kotlin.jvm.functions"), Name.identifier("Function$i"))
+    }
+}
diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/psiUtils.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/psiUtils.kt
index af582fe2768..23951ef9b1c 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/psiUtils.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/psiUtils.kt
@@ -34,7 +34,7 @@ fun FirElement.findPsi(): PsiElement? =
     getAllowedPsi()
 
 fun FirBasedSymbol<*>.findPsi(): PsiElement? =
-    fir.findPsi()
+    fir.findPsi() ?: FirSyntheticFunctionInterfaceSourceProvider.findPsi(fir)
 
 /**
  * Finds [PsiElement] which will be used as go-to referenced element for [KtPsiReference]
@@ -42,5 +42,5 @@ fun FirBasedSymbol<*>.findPsi(): PsiElement? =
  * Otherwise, behaves the same way as [findPsi] returns exact PSI declaration corresponding to passed [FirDeclaration]
  */
 fun FirDeclaration.findReferencePsi(): PsiElement? {
-    return psi
+    return psi ?: FirSyntheticFunctionInterfaceSourceProvider.findPsi(this)
 }