Remove usages of getPsi()

gradle/libs.versions.toml

@@ -3,13 +3,13 @@
 java-compilation = "21"
 # The java-target version is the lowest supported LTS version of Java. Jar's produced are bytecode compatible with this version.
 java-target = "8"
-kotlin = "2.1.0"
-kotlinDev = "2.1.0"
+kotlin = "2.1.10"
+kotlinDev = "2.1.10"
 
 [plugins]
 kotlin-jvm = { id = "org.jetbrains.kotlin.jvm", version.ref = "kotlin" }
 checksum = "org.gradle.crypto.checksum:1.4.0"
-shadow = "com.gradleup.shadow:8.3.5"
+shadow = "com.gradleup.shadow:8.3.6"
 kotlinx-binary-compatibiltiy-validator = "org.jetbrains.kotlinx.binary-compatibility-validator:0.17.0"
 
 [libraries]
@@ -19,18 +19,18 @@ kotlin-plugin-dev = { module = "org.jetbrains.kotlin:kotlin-gradle-plugin", vers
 clikt = "com.github.ajalt.clikt:clikt:5.0.2"
 dokka = "org.jetbrains.dokka:dokka-gradle-plugin:2.0.0"
 ec4j = "org.ec4j.core:ec4j-core:1.1.0"
-logging = "io.github.oshai:kotlin-logging-jvm:7.0.3"
+logging = "io.github.oshai:kotlin-logging-jvm:7.0.4"
 slf4j = "org.slf4j:slf4j-simple:2.0.16"
 poko = "dev.drewhamilton.poko:poko-gradle-plugin:0.18.2"
 # Use logback-classic as the logger for kotlin-logging / slf4j as it allow changing the log level at runtime.
 # TODO: Update "renovate.json" once logback-classic is updated to 1.4 (once java8 support for ktlint is dropped)
 logback = "ch.qos.logback:logback-classic:1.3.15"
-logcaptor = "io.github.hakky54:logcaptor:2.10.0"
+logcaptor = "io.github.hakky54:logcaptor:2.10.1"
 # Required for logback-test.xml conditional configuration so that trace-logging in unit test can be automatically enabled using an
 # environment variable
 janino = "org.codehaus.janino:janino:3.1.12"
 # Testing libraries
 junit5 = "org.junit.jupiter:junit-jupiter:5.11.4"
-assertj = "org.assertj:assertj-core:3.27.2"
+assertj = "org.assertj:assertj-core:3.27.3"
 sarif4k = "io.github.detekt.sarif4k:sarif4k:0.6.0"
 jimfs = "com.google.jimfs:jimfs:1.3.0"

gradle/wrapper/gradle-wrapper.properties

@@ -1,7 +1,7 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionSha256Sum=7a00d51fb93147819aab76024feece20b6b84e420694101f276be952e08bef03
-distributionUrl=https\://services.gradle.org/distributions/gradle-8.12-bin.zip
+distributionSha256Sum=8d97a97984f6cbd2b85fe4c60a743440a347544bf18818048e611f5288d46c94
+distributionUrl=https\://services.gradle.org/distributions/gradle-8.12.1-bin.zip
 networkTimeout=10000
 validateDistributionUrl=true
 zipStoreBase=GRADLE_USER_HOME

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)