[Java Resolution] use a smart psi pointer to store PSI in JavaElement for IDE

`JavaElement`s are reused between read actions, so underlying PSI elements might be invalidated when using hard PSI references ^KT-58194 fixed
2023-05-08 17:51:39 +02:00
parent c114cb67cb
commit 2d08d29dac
6 changed files with 124 additions and 6 deletions
@@ -8,6 +8,7 @@ dependencies {
    api(project(":analysis:analysis-api"))
    api(project(":analysis:analysis-api-impl-barebone"))
    api(project(":analysis:kt-references"))
+    api(project(":compiler:resolution.common.jvm"))
    api(intellijCore())
    implementation(project(":analysis:analysis-internal-utils"))

@@ -0,0 +1,28 @@
+/*
+ * 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.impl.base.java.source
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.SmartPsiElementPointer
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
+
+internal class JavaElementPsiSourceWithSmartPointer<PSI : PsiElement>(
+    private val pointer: SmartPsiElementPointer<PSI>,
+    override val factory: JavaElementSourceFactory,
+) : JavaElementPsiSource<PSI>() {
+
+    override val psi: PSI
+        get() {
+            return pointer.element
+                ?: error("Cannot restore a PsiElement from $pointer")
+        }
+
+    override fun toString(): String {
+        return pointer.element?.toString() ?: "Cannot restore a PsiElement from $pointer"
+    }
+}
+
@@ -0,0 +1,28 @@
+/*
+ * 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.impl.base.java.source
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiType
+import com.intellij.psi.SmartPointerManager
+import com.intellij.psi.SmartTypePointerManager
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementPsiSource
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource
+
+ class JavaElementSourceWithSmartPointerFactory(project: Project) : JavaElementSourceFactory() {
+    private val smartTypePointerManager = SmartTypePointerManager.getInstance(project)
+    private val smartPsiPointerManager = SmartPointerManager.getInstance(project)
+
+    override fun <PSI : PsiElement> createPsiSource(psi: PSI): JavaElementPsiSource<PSI> {
+        return JavaElementPsiSourceWithSmartPointer(smartPsiPointerManager.createSmartPsiElementPointer(psi), this)
+    }
+
+    override fun <TYPE : PsiType> createTypeSource(type: TYPE): JavaElementTypeSource<TYPE> {
+        return JavaElementTypeSourceWithSmartPointer(smartTypePointerManager.createSmartTypePointer(type), this)
+    }
+}
@@ -0,0 +1,29 @@
+/*
+ * 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.impl.base.java.source
+
+import com.intellij.psi.PsiType
+import com.intellij.psi.SmartTypePointer
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementTypeSource
+
+internal class JavaElementTypeSourceWithSmartPointer<TYPE : PsiType>(
+    private val pointer: SmartTypePointer,
+    override val factory: JavaElementSourceFactory,
+) : JavaElementTypeSource<TYPE>() {
+    override val type: TYPE
+        get() {
+            val type = pointer.type
+                ?: error("Cannot restore a PsiType from $pointer")
+
+            @Suppress("UNCHECKED_CAST")
+            return type as TYPE
+        }
+
+    override fun toString(): String {
+        return pointer.type?.toString() ?: "Cannot restore a PsiType from $pointer"
+    }
+}
@@ -7,6 +7,7 @@ dependencies {
    implementation(intellijCore())
    implementation(kotlinStdlib())
    implementation(project(":compiler:psi"))
+    implementation(project(":analysis:analysis-api-impl-base"))
    api(project(":compiler:cli-base"))
    api(project(":analysis:analysis-api"))
    api(project(":analysis:analysis-api-providers"))
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.api.standalone.base.project.structure

 import com.intellij.codeInsight.ExternalAnnotationsManager
 import com.intellij.codeInsight.InferredAnnotationsManager
+import com.intellij.core.CoreApplicationEnvironment
 import com.intellij.core.CoreJavaFileManager
 import com.intellij.core.CorePackageIndex
 import com.intellij.ide.highlighter.JavaFileType
@@ -15,14 +16,16 @@ import com.intellij.openapi.Disposable
 import com.intellij.openapi.roots.PackageIndex
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
-import com.intellij.psi.PsiDirectory
-import com.intellij.psi.PsiFileSystemItem
-import com.intellij.psi.PsiJavaFile
+import com.intellij.psi.*
 import com.intellij.psi.impl.file.impl.JavaFileManager
+import com.intellij.psi.impl.smartPointers.PsiClassReferenceTypePointerFactory
+import com.intellij.psi.impl.smartPointers.SmartPointerManagerImpl
+import com.intellij.psi.impl.smartPointers.SmartTypePointerManagerImpl
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.ProjectScope
 import com.intellij.util.io.URLUtil.JAR_PROTOCOL
 import com.intellij.util.io.URLUtil.JAR_SEPARATOR
+import org.jetbrains.kotlin.analysis.api.impl.base.java.source.JavaElementSourceWithSmartPointerFactory
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.cli.jvm.compiler.*
@@ -35,9 +38,11 @@ import org.jetbrains.kotlin.cli.jvm.modules.CliJavaModuleResolver
 import org.jetbrains.kotlin.cli.jvm.modules.CoreJrtFileSystem
 import org.jetbrains.kotlin.cli.jvm.modules.JavaModuleGraph
 import org.jetbrains.kotlin.config.*
+import org.jetbrains.kotlin.load.java.structure.impl.source.JavaElementSourceFactory
 import org.jetbrains.kotlin.load.kotlin.MetadataFinderFactory
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinderFactory
 import org.jetbrains.kotlin.resolve.ModuleAnnotationsResolver
+import org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.popLast
@@ -51,12 +56,31 @@ object StandaloneProjectFactory {
    ): KotlinCoreProjectEnvironment {
        val applicationEnvironment =
            KotlinCoreEnvironment.getOrCreateApplicationEnvironmentForTests(applicationDisposable, compilerConfiguration)
+        registerApplicationExtensionPoints(applicationEnvironment, applicationDisposable)

        return KotlinCoreProjectEnvironment(projectDisposable, applicationEnvironment).apply {
            registerJavaPsiFacade(project)
        }
    }

+    private fun registerApplicationExtensionPoints(
+        applicationEnvironment: KotlinCoreApplicationEnvironment,
+        applicationDisposable: Disposable,
+    ) {
+        val applicationArea = applicationEnvironment.application.extensionArea
+
+        if (applicationArea.hasExtensionPoint(ClassTypePointerFactory.EP_NAME)) return
+        KotlinCoreEnvironment.underApplicationLock {
+            if (applicationArea.hasExtensionPoint(ClassTypePointerFactory.EP_NAME)) return@underApplicationLock
+            CoreApplicationEnvironment.registerApplicationExtensionPoint(
+                ClassTypePointerFactory.EP_NAME,
+                ClassTypePointerFactory::class.java
+            )
+            applicationArea.getExtensionPoint(ClassTypePointerFactory.EP_NAME)
+                .registerExtension(PsiClassReferenceTypePointerFactory(), applicationDisposable)
+        }
+    }
+
    private fun registerJavaPsiFacade(project: MockProject) {
        with(project) {
            registerService(
@@ -83,7 +107,14 @@ object StandaloneProjectFactory {
        val project = environment.project

        KotlinCoreEnvironment.registerProjectExtensionPoints(project.extensionArea)
-        KotlinCoreEnvironment.registerProjectServices(project)
+        with(project) {
+            registerService(SmartTypePointerManager::class.java, SmartTypePointerManagerImpl::class.java)
+            registerService(SmartPointerManager::class.java, SmartPointerManagerImpl::class.java)
+            registerService(JavaElementSourceFactory::class.java, JavaElementSourceWithSmartPointerFactory::class.java)
+
+            registerService(KotlinJavaPsiFacade::class.java, KotlinJavaPsiFacade(this))
+            registerService(ModuleAnnotationsResolver::class.java, CliModuleAnnotationsResolver())
+        }

        project.registerService(ProjectStructureProvider::class.java, projectStructureProvider)
        initialiseVirtualFileFinderServices(environment, modules, sourceFiles, languageVersionSettings, jdkHome)
@@ -202,14 +233,14 @@ object StandaloneProjectFactory {

    fun getAllBinaryRoots(
        modules: List<KtModule>,
-        environment: KotlinCoreProjectEnvironment
+        environment: KotlinCoreProjectEnvironment,
    ): List<JavaRoot> = withAllTransitiveDependencies(modules)
        .filterIsInstance<KtBinaryModule>()
        .flatMap { it.getJavaRoots(environment) }

    fun getVirtualFilesForLibraryRoots(
        roots: Collection<Path>,
-        environment: KotlinCoreProjectEnvironment
+        environment: KotlinCoreProjectEnvironment,
    ): List<VirtualFile> {
        return roots.mapNotNull { path ->
            val pathString = path.toAbsolutePath().toString()