Reduce usages of psi functions (#2901)

Using psi functions from the kotlin embedded compiler will incurr performance penalties in certain use case as some psi functionality checks for cancellations. See #2901 for more background.

Closes #2901

Co-authored-by: Paul Dingemans <[email protected]>

Author: mgroth0

ktlint-rule-engine-core/api/ktlint-rule-engine-core.api

@@ -3,6 +3,9 @@ public final class com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionKt
 	public static final fun beforeCodeSibling (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/psi/tree/IElementType;)Z
 	public static final fun betweenCodeSiblings (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/psi/tree/IElementType;Lorg/jetbrains/kotlin/com/intellij/psi/tree/IElementType;)Z
 	public static final fun children (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Lkotlin/sequences/Sequence;
+	public static final fun dummyPsiElement (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Lorg/jetbrains/kotlin/com/intellij/psi/PsiElement;
+	public static final fun endOffset (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)I
+	public static final fun findChildByTypeRecursively (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/psi/tree/IElementType;Z)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
 	public static final fun findCompositeParentElementOfType (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/psi/tree/IElementType;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
 	public static final fun firstChildLeafOrSelf (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
 	public static final fun getColumn (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)I
@@ -11,9 +14,12 @@ public final class com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionKt
 	public static final fun indent (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Z)Ljava/lang/String;
 	public static synthetic fun indent$default (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;ZILjava/lang/Object;)Ljava/lang/String;
 	public static final fun isCodeLeaf (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Z
+	public static final fun isDeclaration (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Z
+	public static final fun isKtAnnotated (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Z
 	public static final fun isLeaf (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Z
 	public static final fun isPartOf (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lkotlin/reflect/KClass;)Z
 	public static final fun isPartOf (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/psi/tree/IElementType;)Z
+	public static final fun isPartOf (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/psi/tree/TokenSet;)Z
 	public static final fun isPartOfComment (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Z
 	public static final fun isPartOfCompositeElementOfType (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/psi/tree/IElementType;)Z
 	public static final fun isPartOfString (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Z
@@ -56,6 +62,8 @@ public final class com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionKt
 	public static synthetic fun prevLeaf$default (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;ZILjava/lang/Object;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
 	public static final fun prevSibling (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lkotlin/jvm/functions/Function1;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
 	public static synthetic fun prevSibling$default (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
+	public static final fun recursiveChildren (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Z)Lkotlin/sequences/Sequence;
+	public static synthetic fun recursiveChildren$default (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;ZILjava/lang/Object;)Lkotlin/sequences/Sequence;
 	public static final fun remove (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)V
 	public static final fun replaceWith (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)V
 	public static final fun upsertWhitespaceAfterMe (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Ljava/lang/String;)V
@@ -524,6 +532,12 @@ public final class com/pinterest/ktlint/rule/engine/core/api/SinceKtlint$Status
 	public static fun values ()[Lcom/pinterest/ktlint/rule/engine/core/api/SinceKtlint$Status;
 }
 
+public final class com/pinterest/ktlint/rule/engine/core/api/TokenSets {
+	public static final field INSTANCE Lcom/pinterest/ktlint/rule/engine/core/api/TokenSets;
+	public final fun getCOMMENTS ()Lorg/jetbrains/kotlin/com/intellij/psi/tree/TokenSet;
+	public final fun getCONTROL_FLOW_KEYWORDS ()Lorg/jetbrains/kotlin/com/intellij/psi/tree/TokenSet;
+}
+
 public final class com/pinterest/ktlint/rule/engine/core/api/editorconfig/CodeStyleEditorConfigPropertyKt {
 	public static final fun getCODE_STYLE_PROPERTY ()Lcom/pinterest/ktlint/rule/engine/core/api/editorconfig/EditorConfigProperty;
 	public static final fun getCODE_STYLE_PROPERTY_TYPE ()Lorg/ec4j/core/model/PropertyType$LowerCasingPropertyType;

ktlint-rule-engine-core/src/main/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtension.kt

@@ -7,16 +7,34 @@ import com.pinterest.ktlint.rule.engine.core.api.ElementType.VAL_KEYWORD
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.VARARG_KEYWORD
 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.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
-import org.jetbrains.kotlin.com.intellij.psi.PsiComment
+import org.jetbrains.kotlin.com.intellij.mock.MockProject
+import org.jetbrains.kotlin.com.intellij.openapi.util.Disposer
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.com.intellij.psi.PsiFileFactory
 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.config.CompilerConfiguration
+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
+import org.jetbrains.kotlin.psi.stubs.elements.KtTokenSets
 import org.jetbrains.kotlin.util.prefixIfNot
 import org.jetbrains.kotlin.utils.addToStdlib.applyIf
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
 import kotlin.reflect.KClass
 
 public fun ASTNode.nextLeaf(
@@ -183,11 +201,18 @@ public fun ASTNode.parent(
     return null
 }
 
+public fun ASTNode.isPartOf(tokenSet: TokenSet): Boolean = parent(strict = false) { tokenSet.contains(it.elementType) } != null
+
 /**
  * @param elementType [ElementType].*
  */
 public fun ASTNode.isPartOf(elementType: IElementType): Boolean = parent(elementType, strict = false) != null
 
+@Deprecated(
+    "Marked for removal in Ktlint 2.x. Replace with ASTNode.isPartOf(elementType: IElementType) or ASTNode.isPartOf(tokenSet: TokenSet). " +
+        "This method might cause performance issues, see https://github.com/pinterest/ktlint/pull/2901",
+    replaceWith = ReplaceWith("this.isPartOf(elementTypeOrTokenSet)"),
+)
 public fun ASTNode.isPartOf(klass: KClass<out PsiElement>): Boolean {
     var n: ASTNode? = this
     while (n != null) {
@@ -207,7 +232,13 @@ public fun ASTNode.findCompositeParentElementOfType(iElementType: IElementType):
 
 public fun ASTNode.isPartOfString(): Boolean = parent(STRING_TEMPLATE, strict = false) != null
 
-public fun ASTNode?.isWhiteSpace(): Boolean = this != null && elementType == WHITE_SPACE
+@OptIn(ExperimentalContracts::class)
+public fun ASTNode?.isWhiteSpace(): Boolean {
+    contract {
+        returns(true) implies (this@isWhiteSpace != null)
+    }
+    return this != null && elementType == WHITE_SPACE
+}
 
 public fun ASTNode?.isWhiteSpaceWithNewline(): Boolean = this != null && elementType == WHITE_SPACE && textContains('\n')
 
@@ -223,10 +254,18 @@ public fun ASTNode.isLeaf(): Boolean = firstChildNode == null
  */
 public fun ASTNode.isCodeLeaf(): Boolean = isLeaf() && !isWhiteSpace() && !isPartOfComment()
 
-public fun ASTNode.isPartOfComment(): Boolean = parent(strict = false) { it.psi is PsiComment } != null
+public fun ASTNode.isPartOfComment(): Boolean = isPartOf(TokenSets.COMMENTS)
 
 public fun ASTNode.children(): Sequence<ASTNode> = generateSequence(firstChildNode) { node -> node.treeNext }
 
+public fun ASTNode.recursiveChildren(includeSelf: Boolean = false): Sequence<ASTNode> =
+    sequence {
+        if (includeSelf) {
+            yield(this@recursiveChildren)
+        }
+        children().forEach { yieldAll(it.recursiveChildren(includeSelf = true)) }
+    }
+
 /**
  * Updates or inserts a new whitespace element with [text] before the given node. If the node itself is a whitespace
  * then its contents is replaced with [text]. If the node is a (nested) composite element, the whitespace element is
@@ -555,3 +594,84 @@ public fun ASTNode.replaceWith(node: ASTNode) {
 public fun ASTNode.remove() {
     treeParent.removeChild(this)
 }
+
+/**
+ * Searches the receiver [ASTNode] recursively, returning the first child with type [elementType]. If none are found, returns `null`.
+ * If [includeSelf] is `true`, includes the receiver in the search. The receiver would then be the first element searched, so it is
+ * guaranteed to be returned if it has type [elementType].
+ */
+public fun ASTNode.findChildByTypeRecursively(
+    elementType: IElementType,
+    includeSelf: Boolean,
+): ASTNode? = recursiveChildren(includeSelf).firstOrNull { it.elementType == elementType }
+
+/**
+ * Returns the end offset of the text of this [ASTNode]
+ */
+public fun ASTNode.endOffset(): Int = textRange.endOffset
+
+private val elementTypeCache = hashMapOf<IElementType, PsiElement>()
+
+/**
+ * Checks if the [AstNode] extends the [KtAnnotated] interface. Using this function to check the interface is more performant than checking
+ * whether `psi is KtAnnotated` as the psi does not need to be derived from [ASTNode].
+ */
+public fun ASTNode.isKtAnnotated(): Boolean = psiType { it is KtAnnotated }
+
+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:
+ * ```
+ * astNode.isPsiType<KtAnnotated>()
+ * ```
+ * Using this function to check the [PsiElement] type of the [ASTNode] is more performant than checking whether `astNode.psi is KtAnnotated`
+ * as the psi does not need to be derived from [ASTNode].
+ */
+public inline fun <reified T : KtElement> ASTNode.isPsiType(): Boolean = this.dummyPsiElement() is T
+
+/**
+ * FOR INTERNAL USE ONLY. The returned element is a stub version of a [PsiElement] of the same type as the given [ASTNode]. The returned
+ * result may only be used to validate the type of the [PsiElement].
+ */
+public fun ASTNode.dummyPsiElement(): PsiElement =
+    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 -> createDummyKtFile()
+                is KtKeywordToken -> this as PsiElement
+                is KtNodeType -> (elementType as KtNodeType).createPsi(this)
+                is KtStubElementType<*, *> -> (elementType as KtStubElementType<*, *>).createPsiFromAst(this)
+                is KtToken -> this as PsiElement
+                else -> throw NotImplementedError("Cannot create dummy psi for $elementType (${elementType::class})")
+            }
+        }
+
+private fun createDummyKtFile(): KtFile {
+    val disposable = Disposer.newDisposable()
+    try {
+        val project =
+            KotlinCoreEnvironment
+                .createForProduction(
+                    disposable,
+                    CompilerConfiguration(),
+                    EnvironmentConfigFiles.JVM_CONFIG_FILES,
+                ).project as MockProject
+
+        return PsiFileFactory
+            .getInstance(project)
+            .createFileFromText("dummy-file.kt", KotlinLanguage.INSTANCE, "") as KtFile
+    } finally {
+        // Dispose explicitly to (possibly) prevent memory leak
+        // https://discuss.kotlinlang.org/t/memory-leak-in-kotlincoreenvironment-and-kotlintojvmbytecodecompiler/21950
+        // https://youtrack.jetbrains.com/issue/KT-47044
+        disposable.dispose()
+    }
+}
+
+/**
+ * Returns true if the receiver is not null, and it represents a declaration
+ */
+public fun ASTNode?.isDeclaration(): Boolean = this != null && elementType in KtTokenSets.DECLARATION_TYPES

ktlint-rule-engine-core/src/main/kotlin/com/pinterest/ktlint/rule/engine/core/api/TokenSets.kt

@@ -0,0 +1,31 @@
+package com.pinterest.ktlint.rule.engine.core.api
+
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.DO_KEYWORD
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.FOR_KEYWORD
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.IF_KEYWORD
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.OBJECT_KEYWORD
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.TRY_KEYWORD
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.WHEN_KEYWORD
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.WHILE_KEYWORD
+import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
+import org.jetbrains.kotlin.lexer.KtTokens
+
+public object TokenSets {
+    public val COMMENTS: TokenSet = KtTokens.COMMENTS
+
+    /**
+     *
+     * Reference: This is a subset of [KotlinExpressionParsing.EXPRESSION_FIRST]
+     */
+    public val CONTROL_FLOW_KEYWORDS: TokenSet =
+        TokenSet.create(
+            IF_KEYWORD, // if
+            WHEN_KEYWORD, // when
+            TRY_KEYWORD, // try
+            OBJECT_KEYWORD, // object
+            // loop
+            FOR_KEYWORD,
+            WHILE_KEYWORD,
+            DO_KEYWORD,
+        )
+}

ktlint-rule-engine-core/src/test/kotlin/com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionTest.kt

@@ -5,6 +5,7 @@ import com.pinterest.ktlint.rule.engine.api.KtLintRuleEngine
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.ANNOTATION_ENTRY
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.CLASS
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.CLASS_BODY
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.CLASS_KEYWORD
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.ENUM_ENTRY
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.FILE
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.FUN
@@ -25,6 +26,7 @@ import org.assertj.core.api.Assertions.assertThatNoException
 import org.assertj.core.api.Assertions.entry
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.lang.FileASTNode
+import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.psiUtil.leaves
 import org.junit.jupiter.api.Nested
 import org.junit.jupiter.api.Test
@@ -1089,6 +1091,55 @@ class ASTNodeExtensionTest {
         }
     }
 
+    @Nested
+    inner class FindChildByTypeRecursively {
+        @Test
+        fun `Given a node with a target type return non-null`() {
+            val code =
+                """
+                class MyClass {
+                    fun foo() = 42
+                }
+                """.trimIndent()
+            val result =
+                transformCodeToAST(code)
+                    .findChildByTypeRecursively(FUN, includeSelf = false)
+            assertThat(result).isNotNull()
+        }
+
+        @Test
+        fun `Given a node without a target type return null`() {
+            val code =
+                """
+                class MyClass {
+
+                }
+                """.trimIndent()
+            val result =
+                transformCodeToAST(code)
+                    .findChildByTypeRecursively(FUN, includeSelf = false)
+            assertThat(result).isNull()
+        }
+    }
+
+    @Test
+    fun `Given a simple class declaration without body then the declaration itself is derived from KtAnnotated while its child elements are not derived from KtAnnotated`() {
+        val code =
+            """
+            class Foo
+            """.trimIndent()
+
+        val actual = transformCodeToAST(code).findChildByType(CLASS)!!
+
+        assertThat(actual.isKtAnnotated()).isTrue()
+        assertThat(actual.findChildByType(CLASS_KEYWORD)!!.isKtAnnotated()).isFalse()
+        assertThat(actual.findChildByType(IDENTIFIER)!!.isKtAnnotated()).isFalse()
+
+        assertThat(actual.isPsiType<KtAnnotated>()).isTrue()
+        assertThat(actual.findChildByType(CLASS_KEYWORD)!!.isPsiType<KtAnnotated>()).isFalse()
+        assertThat(actual.findChildByType(IDENTIFIER)!!.isPsiType<KtAnnotated>()).isFalse()
+    }
+
     private inline fun String.transformAst(block: FileASTNode.() -> Unit): FileASTNode =
         transformCodeToAST(this)
             .apply(block)

ktlint-rule-engine/src/main/kotlin/com/pinterest/ktlint/rule/engine/internal/KtlintSuppression.kt

@@ -11,6 +11,7 @@ import com.pinterest.ktlint.rule.engine.core.api.ElementType.VALUE_ARGUMENT
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.VALUE_ARGUMENT_LIST
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.VALUE_PARAMETER
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.VALUE_PARAMETER_LIST
+import com.pinterest.ktlint.rule.engine.core.api.findChildByTypeRecursively
 import com.pinterest.ktlint.rule.engine.core.api.firstChildLeafOrSelf
 import com.pinterest.ktlint.rule.engine.core.api.indent
 import com.pinterest.ktlint.rule.engine.core.api.isPartOfComment
@@ -30,7 +31,6 @@ import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtClassInitializer
-import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtDeclarationModifierList
 import org.jetbrains.kotlin.psi.KtExpression
@@ -47,7 +47,6 @@ import org.jetbrains.kotlin.psi.KtScript
 import org.jetbrains.kotlin.psi.KtScriptInitializer
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
 import org.jetbrains.kotlin.psi.psiUtil.children
-import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
 import org.jetbrains.kotlin.util.prefixIfNot
 
@@ -202,12 +201,14 @@ private fun ASTNode.existingSuppressions() =
     existingSuppressionsFromNamedArgumentOrNull()
         ?: getValueArguments()
 
-private fun ASTNode.existingSuppressionsFromNamedArgumentOrNull() =
-    psi
-        .findDescendantOfType<KtCollectionLiteralExpression>()
-        ?.children
-        ?.map { it.text }
-        ?.toSet()
+private fun ASTNode.existingSuppressionsFromNamedArgumentOrNull(): Set<String>? =
+    findChildByTypeRecursively(ElementType.COLLECTION_LITERAL_EXPRESSION, includeSelf = false)
+        ?.run {
+            children()
+                .filter { it.elementType == ElementType.STRING_TEMPLATE }
+                .map { it.text }
+                .toSet()
+        }
 
 private fun ASTNode.findSuppressionAnnotations(): Map<SuppressAnnotationType, ASTNode> =
     if (this.isRoot()) {