clean up alternative psi-type-checking strategy and remove getPsi() c…

…alls
pinterest · Jan 1, 2025 · 72fb3c0 · 72fb3c0
1 parent 2c7c9e8
commit 72fb3c0
Showing 1 changed file with 81 additions and 17 deletions.
diff --git a/...engine-core/src/main/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtension.kt b/...engine-core/src/main/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtension.kt
@@ -9,16 +9,29 @@ import com.pinterest.ktlint.rule.engine.core.api.ElementType.VAR_KEYWORD
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.WHITE_SPACE
 import org.jetbrains.kotlin.KtNodeType
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
+import org.jetbrains.kotlin.com.intellij.lang.FileASTNode
+import org.jetbrains.kotlin.com.intellij.lang.Language
+import org.jetbrains.kotlin.com.intellij.mock.MockProject
+import org.jetbrains.kotlin.com.intellij.openapi.Disposable
+import org.jetbrains.kotlin.com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.com.intellij.psi.AbstractFileViewProvider
+import org.jetbrains.kotlin.com.intellij.psi.FileViewProvider
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.com.intellij.psi.PsiFile
+import org.jetbrains.kotlin.com.intellij.psi.PsiReference
+import org.jetbrains.kotlin.com.intellij.psi.impl.PsiManagerImpl
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.CompositeElement
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.LeafElement
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.PsiWhiteSpaceImpl
 import org.jetbrains.kotlin.com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
+import org.jetbrains.kotlin.com.intellij.testFramework.LightVirtualFile
+import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.lexer.KtKeywordToken
 import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.leaves
 import org.jetbrains.kotlin.psi.stubs.elements.KtFileElementType
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementType
@@ -609,20 +622,7 @@ private val elementTypeCache = hashMapOf<IElementType, PsiElement>()
  */
 public fun ASTNode.isKtAnnotated(): Boolean = psiType { it is KtAnnotated }
 
-private inline fun ASTNode.psiType(predicate: (psiElement: PsiElement) -> Boolean): Boolean =
-    elementTypeCache
-        .getOrPut(elementType) {
-            // Create a dummy Psi element based on the current node, so that we can store the Psi Type for this ElementType.
-            // Creating this cache entry once per elementType is cheaper than accessing the psi for every node.
-            when (elementType) {
-                is KtFileElementType -> this.psi
-                is KtKeywordToken -> this.psi
-                is KtNodeType -> (elementType as KtNodeType).createPsi(this)
-                is KtStubElementType<*, *> -> (elementType as KtStubElementType<*, *>).createPsiFromAst(this)
-                is KtToken -> this.psi
-                else -> throw NotImplementedError("Cannot create dummy psi for $elementType (${elementType::class})")
-            }
-        }.let { predicate(it) }
+private inline fun ASTNode.psiType(predicate: (psiElement: PsiElement) -> Boolean): Boolean = predicate(dummyPsiElement())
 
 /**
  * Checks if the [AstNode] extends the [T] interface which implements [KtElement]. Call this function like:
@@ -644,11 +644,75 @@ public fun ASTNode.dummyPsiElement(): PsiElement =
             // Create a dummy Psi element based on the current node, so that we can store the Psi Type for this ElementType.
             // Creating this cache entry once per elementType is cheaper than accessing the psi for every node.
             when (elementType) {
-                is KtFileElementType -> this.psi
-                is KtKeywordToken -> this.psi
+                is KtFileElementType -> KtFile(dummyFileViewProvider, false)
+                is KtKeywordToken -> this as PsiElement
                 is KtNodeType -> (elementType as KtNodeType).createPsi(this)
                 is KtStubElementType<*, *> -> (elementType as KtStubElementType<*, *>).createPsiFromAst(this)
-                is KtToken -> this.psi
+                is KtToken -> this as PsiElement
                 else -> throw NotImplementedError("Cannot create dummy psi for $elementType (${elementType::class})")
             }
         }
+
+private val dummyFileViewProvider by lazy {
+    object : AbstractFileViewProvider(
+        PsiManagerImpl(
+            MockProject(
+                null,
+                object : Disposable {
+                    override fun dispose() {
+                        TODO("Not yet implemented")
+                    }
+                },
+            ),
+        ),
+        LightVirtualFile(),
+        false,
+    ) {
+        override fun getPsiInner(p0: Language?): PsiFile? {
+            TODO("Not yet implemented")
+        }
+
+        override fun getCachedPsi(p0: Language): PsiFile? {
+            TODO("Not yet implemented")
+        }
+
+        override fun getCachedPsiFiles(): List<PsiFile?> {
+            TODO("Not yet implemented")
+        }
+
+        override fun getKnownTreeRoots(): List<FileASTNode?> {
+            TODO("Not yet implemented")
+        }
+
+        override fun getBaseLanguage(): Language {
+            TODO("Not yet implemented")
+        }
+
+        private val languages = setOf(KotlinLanguage.INSTANCE)
+
+        override fun getLanguages(): Set<Language> = languages
+
+        override fun getAllFiles(): List<PsiFile?> {
+            TODO("Not yet implemented")
+        }
+
+        override fun findElementAt(p0: Int): PsiElement? {
+            TODO("Not yet implemented")
+        }
+
+        override fun findReferenceAt(p0: Int): PsiReference? {
+            TODO("Not yet implemented")
+        }
+
+        override fun findElementAt(
+            p0: Int,
+            p1: Class<out Language?>,
+        ): PsiElement? {
+            TODO("Not yet implemented")
+        }
+
+        override fun createCopy(p0: VirtualFile): FileViewProvider {
+            TODO("Not yet implemented")
+        }
+    }
+}