[LL] use green stub or load stubs from disk when no stub is available

e.g. if ast is loaded then no stub is attached to a file
but building deserialized fir can require information from enriched cls stub,
otherwise part of the information would be missed, e.g. contracts.

presumably, this can happen only after decompiled code is opened in the editor
(also meaning no source code is attached to the library),
green stub should exist until GCed,
thus performance degradation in the common case is not expected

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedAnnotationDeserializer.kt

@@ -52,7 +52,7 @@ class StubBasedAnnotationDeserializer(
     fun loadConstant(property: KtProperty, callableId: CallableId): FirExpression? {
         if (!property.hasModifier(KtTokens.CONST_KEYWORD)) return null
         constantCache[callableId]?.let { return it }
-        val propertyStub = property.stub as? KotlinPropertyStubImpl ?: return null
+        val propertyStub = (property.stub ?: loadStubByElement(property)) as? KotlinPropertyStubImpl ?: return null
         val constantValue = propertyStub.constantInitializer ?: return null
         return resolveValue(property, constantValue)
     }
@@ -66,7 +66,7 @@ class StubBasedAnnotationDeserializer(
         return deserializeAnnotation(
             ktAnnotation,
             userType.classId(),
-            (ktAnnotation.stub as? KotlinAnnotationEntryStubImpl)?.valueArguments,
+            ((ktAnnotation.stub ?: loadStubByElement(ktAnnotation)) as? KotlinAnnotationEntryStubImpl)?.valueArguments,
             ktAnnotation.useSiteTarget?.getAnnotationUseSiteTarget()
         )
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedClassDeserialization.kt

@@ -5,7 +5,13 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.stubBased.deserialization
 
+import com.intellij.extapi.psi.StubBasedPsiElementBase
+import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiElement
+import com.intellij.psi.impl.source.tree.FileElement
+import com.intellij.psi.stubs.Stub
+import com.intellij.psi.stubs.StubTreeLoader
+import com.intellij.psi.util.PsiUtilCore
 import org.jetbrains.kotlin.KtFakeSourceElement
 import org.jetbrains.kotlin.KtRealPsiSourceElement
 import org.jetbrains.kotlin.descriptors.*
@@ -23,6 +29,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
+import java.lang.ref.WeakReference
 
 internal val KtModifierListOwner.visibility: Visibility
     get() = with(modifierList) {
@@ -45,6 +52,29 @@ internal val KtDeclaration.modality: Modality
         }
     }
 
+private val STUBS_KEY = Key.create<WeakReference<List<Stub>?>>("STUBS")
+internal fun <S, T> loadStubByElement(ktElement: T): S? where T : StubBasedPsiElementBase<*>, T : KtElement {
+    ktElement.greenStub?.let {
+        @Suppress("UNCHECKED_CAST")
+        return it as S
+    }
+    val ktFile = ktElement.containingKtFile
+    require(ktFile.isCompiled) {
+        "Expected compiled file $ktFile"
+    }
+    val virtualFile = PsiUtilCore.getVirtualFile(ktFile) ?: return null
+    var stubList = ktFile.getUserData(STUBS_KEY)?.get()
+    if (stubList == null) {
+        val stubTree = StubTreeLoader.getInstance().readFromVFile(ktElement.project, virtualFile)
+        stubList = stubTree?.plainList ?: emptyList()
+        ktFile.putUserData(STUBS_KEY, WeakReference(stubList))
+    }
+    val nodeList = (ktFile.node as FileElement).stubbedSpine.spineNodes
+    if (stubList.size != nodeList.size) return null
+    @Suppress("UNCHECKED_CAST")
+    return stubList[nodeList.indexOf(ktElement.node)] as S
+}
+
 internal fun deserializeClassToSymbol(
     classId: ClassId,
     classOrObject: KtClassOrObject,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirContractDeserializer.kt

@@ -14,14 +14,16 @@ import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.stubs.impl.*
+import org.jetbrains.kotlin.psi.stubs.impl.KotlinContractConstantValues
+import org.jetbrains.kotlin.psi.stubs.impl.KotlinFunctionStubImpl
+import org.jetbrains.kotlin.psi.stubs.impl.KotlinTypeBean
 
 internal class StubBasedFirContractDeserializer(
     private val simpleFunction: FirSimpleFunction,
-    private val typeDeserializer: StubBasedFirTypeDeserializer
+    private val typeDeserializer: StubBasedFirTypeDeserializer,
 ) {
     fun loadContract(function: KtNamedFunction): FirContractDescription? {
-        val functionStub = function.stub as? KotlinFunctionStubImpl ?: return null
+        val functionStub = function.stub as? KotlinFunctionStubImpl ?: loadStubByElement(function) ?: return null
         val effects = functionStub.contract?.map {
             it.accept(ContractDescriptionConvertingVisitor(), null)
         }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirTypeDeserializer.kt

@@ -103,7 +103,7 @@ internal class StubBasedFirTypeDeserializer(
 
     fun type(typeReference: KtTypeReference): ConeKotlinType {
         val annotations = annotationDeserializer.loadAnnotations(typeReference).toMutableList()
-        val parent = typeReference.stub?.parentStub
+        val parent = (typeReference.stub ?: loadStubByElement(typeReference))?.parentStub
         if (parent is KotlinParameterStubImpl) {
             (parent as? KotlinParameterStubImpl)?.functionTypeParameterName?.let { paramName ->
                 annotations += buildAnnotation {
@@ -163,7 +163,8 @@ internal class StubBasedFirTypeDeserializer(
     }
 
     private fun type(typeReference: KtTypeReference, attributes: ConeAttributes): ConeKotlinType {
-        val upperBoundType = ((typeReference.typeElement as? KtUserType)?.stub as? KotlinUserTypeStubImpl)?.upperBound
+        val userType = typeReference.typeElement as? KtUserType
+        val upperBoundType = (userType?.let { it.stub ?: loadStubByElement(it) } as? KotlinUserTypeStubImpl)?.upperBound
         if (upperBoundType != null) {
             val lowerBound = simpleType(typeReference, attributes)
             val upperBound = type(upperBoundType)
@@ -281,7 +282,7 @@ internal fun KtUserType.classId(): ClassId {
         val referenceExpression = type.referenceExpression as? KtNameReferenceExpression
         if (referenceExpression != null) {
             val referencedName = referenceExpression.getReferencedName()
-            val stub = referenceExpression.stub
+            val stub = referenceExpression.stub ?: loadStubByElement(referenceExpression)
             if (stub is KotlinNameReferenceExpressionStubImpl && stub.isClassRef) {
                 classFragments.add(referencedName)
             } else {
@@ -290,14 +291,6 @@ internal fun KtUserType.classId(): ClassId {
         }
     }
     collectFragments(this)
-    if (classFragments.isEmpty()) {
-        //stub is re-built from decompiled text and additional information is already missed
-        return ClassId(
-            FqName.fromSegments(packageFragments).parent(),
-            FqName(packageFragments.last()),
-            /* local = */ false
-        )
-    }
     return ClassId(
         FqName.fromSegments(packageFragments),
         FqName.fromSegments(classFragments),