Remove usages of getPsi()

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

@@ -6,6 +6,7 @@ public final class com/pinterest/ktlint/rule/engine/core/api/ASTNodeExtensionKt
 	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
+	public static final fun getPrevSiblingIgnoringWhitespaceAndComments (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;
 	public static final fun hasModifier (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/psi/tree/IElementType;)Z
 	public static final fun hasNewLineInClosedRange (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;)Z
 	public static final fun indent (Lorg/jetbrains/kotlin/com/intellij/lang/ASTNode;Z)Ljava/lang/String;
@@ -14,6 +15,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 +58,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 +528,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 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,13 +8,14 @@ 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.psi.psiUtil.siblings
 import org.jetbrains.kotlin.util.prefixIfNot
 import org.jetbrains.kotlin.utils.addToStdlib.applyIf
 import kotlin.reflect.KClass
@@ -183,11 +184,17 @@ public fun ASTNode.parent(
     return null
 }
 
+public fun ASTNode.isPartOf(tokenSet: TokenSet): Boolean = parent(predicate = { tokenSet.contains(it.elementType) }, strict = false) != null
+
 /**
  * @param elementType [ElementType].*
  */
 public fun ASTNode.isPartOf(elementType: IElementType): Boolean = parent(elementType, strict = false) != null
 
+@Deprecated(
+    "psi is a performance issue, see https://github.com/pinterest/ktlint/pull/2901",
+    replaceWith = ReplaceWith("ASTNode.isPartOf(elementType: IElementType) or ASTNode.isPartOf(tokenSet: TokenSet)"),
+)
 public fun ASTNode.isPartOf(klass: KClass<out PsiElement>): Boolean {
     var n: ASTNode? = this
     while (n != null) {
@@ -223,10 +230,21 @@ 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 {
+            yield(it)
+            yieldAll(it.recursiveChildren())
+        }
+    }
+
 /**
  * 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 +573,9 @@ public fun ASTNode.replaceWith(node: ASTNode) {
 public fun ASTNode.remove() {
     treeParent.removeChild(this)
 }
+
+public fun ASTNode.getPrevSiblingIgnoringWhitespaceAndComments(): ASTNode? =
+    siblings(forward = false)
+        .filter {
+            !it.isWhiteSpace() && !TokenSets.COMMENTS.contains(it.elementType)
+        }.firstOrNull()

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

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

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

@@ -17,6 +17,7 @@ import com.pinterest.ktlint.rule.engine.core.api.isPartOfComment
 import com.pinterest.ktlint.rule.engine.core.api.isRoot
 import com.pinterest.ktlint.rule.engine.core.api.isWhiteSpace
 import com.pinterest.ktlint.rule.engine.core.api.nextCodeSibling
+import com.pinterest.ktlint.rule.engine.core.api.recursiveChildren
 import com.pinterest.ktlint.rule.engine.core.api.replaceWith
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
@@ -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,15 @@ private fun ASTNode.existingSuppressions() =
     existingSuppressionsFromNamedArgumentOrNull()
         ?: getValueArguments()
 
-private fun ASTNode.existingSuppressionsFromNamedArgumentOrNull() =
-    psi
-        .findDescendantOfType<KtCollectionLiteralExpression>()
-        ?.children
-        ?.map { it.text }
-        ?.toSet()
+private fun ASTNode.existingSuppressionsFromNamedArgumentOrNull(): Set<String>? =
+    recursiveChildren()
+        .firstOrNull { it.elementType == ElementType.COLLECTION_LITERAL_EXPRESSION }
+        ?.run {
+            children()
+                .filter { it.elementType == ElementType.STRING_TEMPLATE }
+                .map { it.text }
+                .toSet()
+        }
 
 private fun ASTNode.findSuppressionAnnotations(): Map<SuppressAnnotationType, ASTNode> =
     if (this.isRoot()) {

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

@@ -1,20 +1,19 @@
 package com.pinterest.ktlint.rule.engine.internal
 
+import com.pinterest.ktlint.rule.engine.core.api.ElementType
 import com.pinterest.ktlint.rule.engine.core.api.ElementType.RBRACE
 import com.pinterest.ktlint.rule.engine.core.api.IgnoreKtlintSuppressions
 import com.pinterest.ktlint.rule.engine.core.api.Rule
 import com.pinterest.ktlint.rule.engine.core.api.RuleId
+import com.pinterest.ktlint.rule.engine.core.api.TokenSets
 import com.pinterest.ktlint.rule.engine.core.api.editorconfig.EditorConfig
 import com.pinterest.ktlint.rule.engine.core.api.nextSibling
+import com.pinterest.ktlint.rule.engine.core.api.parent
+import com.pinterest.ktlint.rule.engine.core.api.recursiveChildren
 import com.pinterest.ktlint.rule.engine.internal.SuppressionLocator.CommentSuppressionHint.Type.BLOCK_END
 import com.pinterest.ktlint.rule.engine.internal.SuppressionLocator.CommentSuppressionHint.Type.BLOCK_START
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
-import org.jetbrains.kotlin.com.intellij.psi.PsiComment
-import org.jetbrains.kotlin.psi.KtAnnotated
-import org.jetbrains.kotlin.psi.KtAnnotationEntry
-import org.jetbrains.kotlin.psi.ValueArgument
-import org.jetbrains.kotlin.psi.psiUtil.endOffset
-import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
 
 internal class SuppressionLocator(
     editorConfig: EditorConfig,
@@ -57,34 +56,27 @@ internal class SuppressionLocator(
     private fun findSuppressionHints(rootNode: ASTNode): List<SuppressionHint> {
         val suppressionHints = ArrayList<SuppressionHint>()
         val commentSuppressionsHints = mutableListOf<CommentSuppressionHint>()
-        rootNode.findSuppressionHints { node ->
-            when (val psi = node.psi) {
-                is PsiComment -> {
+        rootNode.recursiveChildren(includeSelf = true).forEach { node ->
+            val eType = node.elementType
+            when {
+                TokenSets.COMMENTS.contains(eType) -> {
                     node
                         .createSuppressionHintFromComment()
                         ?.let { commentSuppressionsHints.add(it) }
                 }
 
-                is KtAnnotated -> {
-                    psi
-                        .createSuppressionHintFromAnnotations()
-                        ?.let { suppressionHints.add(it) }
+                eType == ElementType.ANNOTATION_ENTRY -> {
+                    createSuppressionHintFromAnnotations(
+                        annotation = node,
+                    )?.let { suppressionHints.add(it) }
                 }
             }
         }
-
         return suppressionHints.plus(
             commentSuppressionsHints.toSuppressionHints(),
         )
     }
 
-    private fun ASTNode.findSuppressionHints(block: (node: ASTNode) -> Unit) {
-        block(this)
-        this
-            .getChildren(null)
-            .forEach { it.findSuppressionHints(block) }
-    }
-
     private fun ASTNode.createSuppressionHintFromComment(): CommentSuppressionHint? =
         text
             .removePrefix("//")
@@ -178,42 +170,41 @@ internal class SuppressionLocator(
 
     private fun <T> List<T>.tail() = this.subList(1, this.size)
 
-    private fun KtAnnotated.createSuppressionHintFromAnnotations(): SuppressionHint? =
-        annotationEntries
-            .filter {
-                it
-                    .calleeExpression
-                    ?.constructorReferenceExpression
-                    ?.getReferencedName() in SUPPRESS_ANNOTATIONS
-            }.flatMap(KtAnnotationEntry::getValueArguments)
-            .flatMap { it.findRuleSuppressionIds() }
-            .let { suppressedRuleIds ->
-                when {
-                    suppressedRuleIds.isEmpty() -> {
-                        null
-                    }
-
-                    suppressedRuleIds.contains(ALL_KTLINT_RULES_SUPPRESSION_ID) -> {
-                        SuppressionHint(
-                            IntRange(startOffset, endOffset - 1),
-                            emptySet(),
-                        )
-                    }
+    private fun createSuppressionHintFromAnnotations(annotation: ASTNode): SuppressionHint? {
+        if (annotation
+                .findChildByType(ElementType.CONSTRUCTOR_CALLEE)
+                ?.findChildByType(ElementType.TYPE_REFERENCE)
+                ?.text !in SUPPRESS_ANNOTATIONS
+        ) {
+            return null
+        }
 
-                    else -> {
-                        SuppressionHint(
-                            IntRange(startOffset, endOffset - 1),
-                            suppressedRuleIds.toSet(),
-                        )
-                    }
+        val suppressedRuleIds =
+            annotation
+                .recursiveChildren()
+                .filter { it.elementType == ElementType.VALUE_ARGUMENT }
+                .flatMapTo(mutableListOf()) {
+                    it.text.findRuleSuppressionIds()
                 }
-            }
 
-    private fun ValueArgument.findRuleSuppressionIds(): List<String> =
-        getArgumentExpression()
-            ?.text
-            ?.removeSurrounding("\"")
-            ?.let { argumentExpressionText ->
+        if (suppressedRuleIds.isEmpty()) return null
+
+        val owner =
+            annotation.parent {
+                ANNOTATED_ELEMENT_TYPES.contains(it.elementType)
+            } ?: return null
+
+        val tr = owner.textRange
+
+        return SuppressionHint(
+            IntRange(tr.startOffset, tr.endOffset - 1),
+            if (suppressedRuleIds.contains(ALL_KTLINT_RULES_SUPPRESSION_ID)) emptySet() else suppressedRuleIds.toSet(),
+        )
+    }
+
+    private fun String.findRuleSuppressionIds(): List<String> =
+        removeSurrounding("\"")
+            .let { argumentExpressionText ->
                 when {
                     argumentExpressionText == "ktlint" -> {
                         // Disable all rules
@@ -280,5 +271,21 @@ internal class SuppressionLocator(
             )
         val SUPPRESS_ANNOTATIONS = setOf("Suppress", "SuppressWarnings")
         const val ALL_KTLINT_RULES_SUPPRESSION_ID = "ktlint:suppress-all-rules"
+        val ANNOTATED_ELEMENT_TYPES =
+            TokenSet.create(
+                ElementType.CLASS,
+                ElementType.OBJECT_DECLARATION,
+                ElementType.ENUM_ENTRY,
+                ElementType.FUN,
+                ElementType.PROPERTY,
+                ElementType.PROPERTY_ACCESSOR,
+                ElementType.TYPE_PARAMETER,
+                ElementType.TYPEALIAS,
+                ElementType.PRIMARY_CONSTRUCTOR,
+                ElementType.SECONDARY_CONSTRUCTOR,
+                ElementType.ANNOTATED_EXPRESSION,
+                ElementType.EXPRESSION_CODE_FRAGMENT,
+                ElementType.FILE,
+            )
     }
 }

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

@@ -17,6 +17,7 @@ import com.pinterest.ktlint.rule.engine.core.api.nextLeaf
 import com.pinterest.ktlint.rule.engine.core.api.nextSibling
 import com.pinterest.ktlint.rule.engine.core.api.parent
 import com.pinterest.ktlint.rule.engine.core.api.prevLeaf
+import com.pinterest.ktlint.rule.engine.core.api.recursiveChildren
 import com.pinterest.ktlint.rule.engine.core.api.remove
 import com.pinterest.ktlint.rule.engine.core.api.replaceWith
 import com.pinterest.ktlint.rule.engine.internal.KTLINT_SUPPRESSION_ID_ALL_RULES
@@ -42,9 +43,9 @@ import org.jetbrains.kotlin.psi.KtScriptInitializer
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
 import org.jetbrains.kotlin.psi.KtValueArgument
 import org.jetbrains.kotlin.psi.KtValueArgumentList
-import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
 import org.jetbrains.kotlin.psi.psiUtil.siblings
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.utils.addToStdlib.applyIf
 
 /**
@@ -105,10 +106,10 @@ public class KtlintSuppressionRule(
         emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> AutocorrectDecision,
     ) {
         node
-            .psi
-            .findDescendantOfType<KtLiteralStringTemplateEntry>()
-            ?.node
-            ?.let { literalStringTemplateEntry ->
+            .recursiveChildren(includeSelf = true)
+            .firstOrNull {
+                it.elementType == ElementType.LITERAL_STRING_TEMPLATE_ENTRY
+            }?.let { literalStringTemplateEntry ->
                 val prefixedSuppression =
                     literalStringTemplateEntry
                         .text