1 files changed, 94 insertions, 0 deletions
diff --git a/ktlint-test/src/main/kotlin/com/github/shyiko/ktlint/test/DumpAST.kt b/ktlint-test/src/main/kotlin/com/github/shyiko/ktlint/test/DumpAST.kt
new file mode 100644
index 00000000..a22223ee
--- /dev/null
+++ b/ktlint-test/src/main/kotlin/com/github/shyiko/ktlint/test/DumpAST.kt
@@ -0,0 +1,94 @@
+package com.github.shyiko.ktlint.test
+
+import com.andreapivetta.kolor.Color
+import com.andreapivetta.kolor.Kolor
+import com.github.shyiko.ktlint.core.Rule
+import org.jetbrains.kotlin.com.intellij.lang.ASTNode
+import org.jetbrains.kotlin.com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.diagnostics.DiagnosticUtils
+import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
+import java.io.PrintStream
+
+val debugAST = {
+    (System.getProperty("ktlintDebug") ?: System.getenv("KTLINT_DEBUG") ?: "")
+        .toLowerCase().split(",").contains("ast")
+}
+
+class DumpAST @JvmOverloads constructor(
+    private val out: PrintStream = System.err,
+    private val color: Boolean = false
+) : Rule("dump") {
+
+    private var lineNumberColumnLength: Int = 0
+    private var lastNode: ASTNode? = null
+
+    override fun visit(
+        node: ASTNode,
+        autoCorrect: Boolean,
+        emit: (offset: Int, errorMessage: String, corrected: Boolean) -> Unit
+    ) {
+        if (node.elementType == KtStubElementTypes.FILE) {
+            lineNumberColumnLength = (location(PsiTreeUtil.getDeepestLast(node.psi).node)?.line ?: 1)
+                .let { var v = it; var c = 0; while (v > 0) { c++; v /= 10 }; c }
+            lastNode = lastChildNodeOf(node)
+        }
+        var level = -1
+        var parent: ASTNode? = node
+        do {
+            level++
+            parent = parent?.treeParent
+        } while (parent != null)
+        out.println((
+            location(node)
+                ?.let { String.format("%${lineNumberColumnLength}s: ", it.line).gray() }
+                // should only happen when autoCorrect=true and other rules mutate AST in a way that changes text length
+                ?: String.format("%${lineNumberColumnLength}s: ", "?").gray()
+            ) +
+            "  ".repeat(level).gray() +
+            colorClassName(node.psi.className) +
+            " (".gray() + colorClassName(node.elementType.className) + "." + node.elementType + ")".gray() +
+            if (node.getChildren(null).isEmpty()) " \"" + node.text.escape().yellow() + "\"" else "")
+        if (lastNode == node) {
+            out.println()
+            out.println(" ".repeat(lineNumberColumnLength) +
+                "  format: <line_number:> <node.psi::class> (<node.elementType>) \"<node.text>\"".gray())
+            out.println(" ".repeat(lineNumberColumnLength) +
+                "  legend: ~ = org.jetbrains.kotlin, c.i.p = com.intellij.psi".gray())
+            out.println()
+        }
+    }
+
+    private tailrec fun lastChildNodeOf(node: ASTNode): ASTNode? =
+        if (node.lastChildNode == null) node else lastChildNodeOf(node.lastChildNode)
+
+    private fun location(node: ASTNode) =
+        node.psi.containingFile?.let { psiFile ->
+            try {
+                DiagnosticUtils.getLineAndColumnInPsiFile(
+                    psiFile,
+                    TextRange(node.startOffset, node.startOffset)
+                )
+            } catch (e: Exception) {
+                null // DiagnosticUtils has no knowledge of mutated AST
+            }
+        }
+
+    private fun colorClassName(className: String): String {
+        val name = className.substringAfterLast(".")
+        return className.substring(0, className.length - name.length).gray() + name
+    }
+
+    private fun String.yellow() =
+        if (color) Kolor.foreground(this, Color.YELLOW) else this
+    private fun String.gray() =
+        if (color) Kolor.foreground(this, Color.DARK_GRAY) else this
+
+    private val Any.className
+        get() = this.javaClass.name
+            .replace("org.jetbrains.kotlin.", "~.")
+            .replace("com.intellij.psi.", "c.i.p.")
+
+    private fun String.escape() =
+        this.replace("\\", "\\\\").replace("\n", "\\n").replace("\t", "\\t").replace("\r", "\\r")
+}