From 9ac639e0d3d23b2b289da59ae122f1f9f9ec3a43 Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Fri, 22 Sep 2023 15:12:30 +0200
Subject: [PATCH] [PSI] Add `forEachDescendantOfTypeInPreorder`

- `forEachDescendantOfTypeInPreorder` is a preorder variant of
  `forEachDescendantOfType`, which operates in postorder.
---
 .../jetbrains/kotlin/psi/psiUtil/psiUtils.kt  | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt
index 188e8cd0fca..e0cbccc4366 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt
@@ -239,6 +239,28 @@ inline fun <reified T : PsiElement> PsiElement.forEachDescendantOfType(
     })
 }
 
+inline fun <reified T : PsiElement> PsiElement.forEachDescendantOfTypeInPreorder(noinline action: (T) -> Unit) {
+    forEachDescendantOfTypeInPreorder({ true }, action)
+}
+
+inline fun <reified T : PsiElement> PsiElement.forEachDescendantOfTypeInPreorder(
+    crossinline canGoInside: (PsiElement) -> Boolean,
+    noinline action: (T) -> Unit,
+) {
+    checkDecompiledText()
+    this.accept(object : PsiRecursiveElementVisitor() {
+        override fun visitElement(element: PsiElement) {
+            if (element is T) {
+                action(element)
+            }
+
+            if (canGoInside(element)) {
+                super.visitElement(element)
+            }
+        }
+    })
+}
+
 inline fun <reified T : PsiElement> PsiElement.anyDescendantOfType(noinline predicate: (T) -> Boolean = { true }): Boolean {
     return findDescendantOfType(predicate) != null
 }