Remove usages of getPsi()

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

@@ -3,6 +3,8 @@ 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 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
@@ -14,6 +16,7 @@ public final class com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionKt
 	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 +59,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 +529,13 @@ 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 fun getEXPRESSIONS ()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

@@ -8,12 +8,12 @@ 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.com.intellij.lang.ASTNode
-import org.jetbrains.kotlin.com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 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.psi.psiUtil.leaves
 import org.jetbrains.kotlin.util.prefixIfNot
 import org.jetbrains.kotlin.utils.addToStdlib.applyIf
@@ -183,11 +183,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) {
@@ -223,10 +230,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 +570,10 @@ public fun ASTNode.replaceWith(node: ASTNode) {
 public fun ASTNode.remove() {
     treeParent.removeChild(this)
 }
+
+public fun ASTNode.findChildByTypeRecursively(
+    elementType: IElementType,
+    includeSelf: Boolean,
+): ASTNode? = recursiveChildren(includeSelf).firstOrNull { it.elementType == elementType }
+
+public fun ASTNode.endOffset(): Int = textRange.endOffset

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

@@ -0,0 +1,78 @@
+package com.pinterest.ktlint.rule.engine.core.api
+
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.ANNOTATED_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.ARRAY_ACCESS_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.BINARY_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.BINARY_WITH_TYPE
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.BLOCK_COMMENT
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.CALLABLE_REFERENCE_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.CALL_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.CLASS_LITERAL_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.COLLECTION_LITERAL_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.DOT_QUALIFIED_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.DO_KEYWORD
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.ENUM_ENTRY_SUPERCLASS_REFERENCE_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.EOL_COMMENT
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.EXPRESSION_CODE_FRAGMENT
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.FUNCTION_LITERAL
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.IF_KEYWORD
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.INDICES
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.IS_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.KDOC
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.LABEL
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.LABELED_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.LABEL_QUALIFIER
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.LAMBDA_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.OBJECT_LITERAL
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.OPERATION_REFERENCE
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.POSTFIX_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.PREFIX_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.REFERENCE_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.SAFE_ACCESS_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.SUPER_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.THIS_EXPRESSION
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.TYPE_CODE_FRAGMENT
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.WHEN
+import com.pinterest.ktlint.rule.engine.core.api.ElementType.WHILE_KEYWORD
+import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
+
+public object TokenSets {
+    public val COMMENTS: TokenSet = TokenSet.create(BLOCK_COMMENT, EOL_COMMENT, KDOC)
+    public val EXPRESSIONS: TokenSet =
+        TokenSet.create(
+            ANNOTATED_EXPRESSION,
+            ARRAY_ACCESS_EXPRESSION,
+            BINARY_EXPRESSION,
+            BINARY_WITH_TYPE,
+            CALLABLE_REFERENCE_EXPRESSION,
+            CALL_EXPRESSION,
+            CLASS_LITERAL_EXPRESSION,
+            COLLECTION_LITERAL_EXPRESSION,
+            DOT_QUALIFIED_EXPRESSION,
+            ENUM_ENTRY_SUPERCLASS_REFERENCE_EXPRESSION,
+            EXPRESSION_CODE_FRAGMENT,
+            FUNCTION_LITERAL,
+            INDICES,
+            IS_EXPRESSION,
+            LABEL,
+            LABELED_EXPRESSION,
+            LABEL_QUALIFIER,
+            LAMBDA_EXPRESSION,
+            OBJECT_LITERAL,
+            OPERATION_REFERENCE,
+            POSTFIX_EXPRESSION,
+            PREFIX_EXPRESSION,
+            REFERENCE_EXPRESSION,
+            SAFE_ACCESS_EXPRESSION,
+            SUPER_EXPRESSION,
+            THIS_EXPRESSION,
+            TYPE_CODE_FRAGMENT,
+            WHEN,
+        )
+    public val CONTROL_FLOW_KEYWORDS: TokenSet =
+        TokenSet.create(
+            DO_KEYWORD,
+            IF_KEYWORD,
+            WHILE_KEYWORD,
+        )
+}

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

@@ -1089,6 +1089,37 @@ 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()
+        }
+    }
+
     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()) {