From 342488e66e0d85b8a7c8b372dcd94b8740fa7471 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Mon, 6 Mar 2023 16:31:47 +0900
Subject: [PATCH] [LL API] Restrict accepted declarations for on-air resolution

After 1cfd2ae7ff398c4aab0d1c240f6b65a102b894bb, all non-local
declarations became targets for on-air resolution. However, some of
them, including secondary constructors, are not supposed to be used
in that way. The commit restores the previous behavior.

This commits fixes failing tests in the following groups:
- FirIdeDependentAnalysisSourceModuleIsUsedAsExpressionTestGenerated
- FirIdeDependentAnalysisSourceModuleReferenceResolveTestGenerated
- FirOnAirResolveTestGenerated
---
 .../LowLevelFirApiFacadeForResolveOnAir.kt    | 26 +++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
index bd82da9797f..dbdea70555a 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.api
 
+import com.intellij.psi.PsiElement
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.analysis.low.level.api.fir.DeclarationCopyBuilder.withBodyFrom
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionDepended
@@ -47,6 +48,27 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 
 object LowLevelFirApiFacadeForResolveOnAir {
+    private fun PsiElement.onAirGetNonLocalContainingOrThisDeclarationFor(): KtDeclaration? {
+        return getNonLocalContainingOrThisDeclaration { declaration ->
+            when (declaration) {
+                is KtNamedFunction,
+                is KtProperty,
+                is KtTypeAlias,
+                is KtScript -> {
+                    true
+                }
+
+                is KtClassOrObject -> {
+                    declaration !is KtEnumEntry
+                }
+
+                else -> {
+                    false
+                }
+            }
+        }
+    }
+
     private fun recordOriginalDeclaration(targetDeclaration: KtDeclaration, originalDeclaration: KtDeclaration) {
         require(originalDeclaration.containingKtFile !== targetDeclaration.containingKtFile)
         val originalDeclarationParents = originalDeclaration.parentsOfType<KtDeclaration>().toList()
@@ -136,7 +158,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         require(originalFirResolveSession is LLFirResolvableResolveSession)
         require(elementToAnalyze !is KtFile) { "KtFile for dependency element not supported" }
 
-        val dependencyNonLocalDeclaration = elementToAnalyze.getNonLocalContainingOrThisDeclaration() as? KtNamedDeclaration
+        val dependencyNonLocalDeclaration = elementToAnalyze.onAirGetNonLocalContainingOrThisDeclarationFor() as? KtNamedDeclaration
 
         if (dependencyNonLocalDeclaration == null) {
             val towerDataContext = onAirGetTowerContextForFile(originalFirResolveSession, originalKtFile)
@@ -205,7 +227,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         onAirCreatedDeclaration: Boolean,
         collector: FirTowerDataContextCollector? = null,
     ): FirElement {
-        val nonLocalDeclaration = replacement.from.getNonLocalContainingOrThisDeclaration()
+        val nonLocalDeclaration = replacement.from.onAirGetNonLocalContainingOrThisDeclarationFor()
         val originalFirFile = firResolveSession.getOrBuildFirFile(replacement.from.containingKtFile)
 
         if (nonLocalDeclaration == null) {