From 7d6275e22840ef33b63b2fdcd4991697e134e7e9 Mon Sep 17 00:00:00 2001 From: Anna Kozlova Date: Thu, 1 Jun 2023 16:07:17 +0200 Subject: [PATCH] [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 --- .../StubBasedAnnotationDeserializer.kt | 4 +-- .../StubBasedClassDeserialization.kt | 30 +++++++++++++++++++ .../StubBasedFirContractDeserializer.kt | 8 +++-- .../StubBasedFirTypeDeserializer.kt | 15 +++------- 4 files changed, 41 insertions(+), 16 deletions(-) diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedAnnotationDeserializer.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedAnnotationDeserializer.kt index f3ef49e6a50..4f5b4a029fb 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedAnnotationDeserializer.kt +++ b/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() ) } diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedClassDeserialization.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedClassDeserialization.kt index 2fa283ac359..37708995b54 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedClassDeserialization.kt +++ b/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?>>("STUBS") +internal fun 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, diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirContractDeserializer.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirContractDeserializer.kt index 0c3ab7b715e..35bd428ec28 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirContractDeserializer.kt +++ b/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) } diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirTypeDeserializer.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirTypeDeserializer.kt index 6925458d976..d9377085d80 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirTypeDeserializer.kt +++ b/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),