From 11c4fcda5fa90a71fec9483db279fa927549c82f Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Wed, 28 Jun 2023 13:11:24 +0200
Subject: [PATCH] [psi][AA] don't decompile code when stubs have enough
 information

^ KTIJ-25979
+ do not assert when there are no stubs (text is already computed)
---
 .../references/FirReferenceResolveHelper.kt   | 24 +++++++++++++++----
 .../jetbrains/kotlin/psi/psiUtil/psiUtils.kt  |  2 +-
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirReferenceResolveHelper.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirReferenceResolveHelper.kt
index c2f7f137b28..1fc93beebe0 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirReferenceResolveHelper.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirReferenceResolveHelper.kt
@@ -50,6 +50,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
+import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
@@ -142,15 +143,30 @@ internal object FirReferenceResolveHelper {
         return when (qualified) {
             null -> FqName(expression.getReferencedName())
             else -> {
-                qualified
-                    .collectDescendantsOfType<KtSimpleNameExpression>()
-                    .dropWhile { it.getReferencedName() == ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE }
-                    .joinToString(separator = ".") { it.getReferencedName() }
+                val refs =
+                    if (qualified is KtUserType && qualified.stub != null && (qualified.containingFile as? KtFile)?.isCompiled == true) {
+                        collectTypeReferences(qualified)
+                    } else {
+                        qualified.collectDescendantsOfType<KtSimpleNameExpression>()
+                    }
+                refs.map { it.getReferencedName() }
+                    .dropWhile { it == ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE }
+                    .joinToString(separator = ".")
                     .let(::FqName)
             }
         }
     }
 
+    private fun collectTypeReferences(qualified: KtUserType): MutableList<KtNameReferenceExpression> {
+        val refs = mutableListOf<KtNameReferenceExpression>()
+        fun collectFragments(type: KtUserType) {
+            type.getStubOrPsiChild(KtStubElementTypes.USER_TYPE)?.let { collectFragments(it) }
+            refs.add(type.referenceExpression as? KtNameReferenceExpression ?: return)
+        }
+        collectFragments(qualified)
+        return refs
+    }
+
     private fun KtSimpleNameExpression.isPartOfQualifiedExpression(): Boolean {
         var parent = parent
         while (parent is KtDotQualifiedExpression) {
diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt
index 6d65f5fb034..188e8cd0fca 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt
@@ -278,7 +278,7 @@ inline fun <reified T : PsiElement> PsiElement.findDescendantOfType(
 
 fun PsiElement.checkDecompiledText() {
     val file = containingFile
-    if (file is KtFile && file.isCompiled) {
+    if (file is KtFile && file.isCompiled && file.stub != null) {
         error("Attempt to load decompiled text, please use stubs instead. Decompile process might be slow and should be avoided")
     }
 }