Merge pull request #103 from orangain/node-extra

Now each node has extra information

README.md

@@ -71,7 +71,7 @@ Examples below are simple Kotlin scripts.
 
 #### Parsing code
 
-In this example, we use the wrapper around the Kotlin compiler's parser:
+In this example, we use the Parser, which is a wrapper around the Kotlin compiler's parser:
 
 ```kotlin
 import ktast.ast.psi.Parser
@@ -88,25 +88,22 @@ val code = """
 """.trimIndent()
 // Call the parser with the code
 val file = Parser.parseFile(code)
-// The file var is now a ktast.ast.Node.KotlinFile that is used in future examples...
 ```
 
-The `file` variable has the full AST. Note, if you want to parse with blank line and comment information, you can create
-a converter that holds the extras:
+The `file` variable is now a `ktast.ast.Node.KotlinFile`. Each AST nodes have blank line and comment information. If you
+don't need them, you can pass a `Converter` instance to the constructor argument of the `Parser` instead:
 
 ```kotlin
-import ktast.ast.psi.ConverterWithExtras
+import ktast.ast.psi.Parser
+import ktast.ast.psi.Converter
 
 // ...
-
-val extrasMap = ConverterWithExtras()
-val file = Parser(extrasMap).parseFile(code)
-// extrasMap is an instance of ktast.ast.ExtrasMap
+val fileWithoutExtras = Parser(Converter()).parseFile(code)
 ```
 
 #### Writing code
 
-To write the code created above, simply use the writer
+To write the code from the node created above, simply use the Writer:
 
 ```kotlin
 import ktast.ast.Writer
@@ -116,20 +113,23 @@ import ktast.ast.Writer
 println(Writer.write(file))
 ```
 
-This outputs a string of the written code from the AST `file` object. To include the extra blank line and comment info
-from the previous parse example, pass in the extras map:
+This outputs the following code, which is exactly the same code as the input. This is because the AST nodes have blank
+line and comment information.
 
 ```kotlin
-// ...
+package foo
 
-println(Writer.write(file, extrasMap))
-```
+fun bar() {
+    // Print hello
+    println("Hello, World!")
+}
 
-This outputs the code with the comments.
+fun baz() = println("Hello, again!")
+```
 
 #### Visiting nodes
 
-This will get all strings:
+To get all strings from the file, we can use the Visitor:
 
 ```kotlin
 import ktast.ast.Node
@@ -150,11 +150,12 @@ println(strings)
 
 The parameter of the lambda is
 a [NodePath](https://orangain.github.io/ktast/latest/api/ast/ktast.ast/-node-path/index.html) object that has `node`
-and `parent` NodePath. There is a `tag` var on each node that can be used to store per-node state if desired.
+and `parent` NodePath.
 
 #### Modifying nodes
 
-This will change "Hello, World!" and "Hello, again!" to "Howdy, World!" and "Howdy, again":
+To modify the file, we can use the MutableVisitor. The following code will change "Hello, World!" and "Hello, again!"
+to "Howdy, World!" and "Howdy, again":
 
 ```kotlin
 import ktast.ast.MutableVisitor
@@ -166,16 +167,25 @@ val newFile = MutableVisitor.traverse(file) { path ->
     if (node !is Node.Expression.StringLiteralExpression.LiteralStringEntry) node
     else node.copy(text = node.text.replace("Hello", "Howdy"))
 }
+
+Writer.write(newFile)
 ```
 
-Now `newFile` is a transformed version of `file`. As before, the parameter of the lambda is a NodePath object.
+Now `newFile` is a transformed version of `file`. As before, the parameter of the lambda is a NodePath object. The
+output will be the following code:
+
+```kotlin
+package foo
+
+fun bar() {
+    // Print hello
+    println("Howdy, World!")
+}
+
+fun baz() = println("Howdy, again!")
+```
 
-Note, since `extraMap` support relies on object identities and this creates entirely new objects in the immutable tree,
-the extra map becomes invalid on this step. When you mutate an AST tree, it is recommended to
-use [ConverterWithMutableExtras](https://orangain.github.io/ktast/latest/api/ast-psi/ktast.ast.psi/-converter-with-mutable-extras/index.html)
-that
-implements [MutableExtrasMap](https://orangain.github.io/ktast/latest/api/ast/ktast.ast/-mutable-extras-map/index.html)
-interface.
+Note that the comments and blank lines are preserved.
 
 ## Running tests
 

ast-psi/src/main/kotlin/ktast/ast/psi/ConverterWithExtras.kt

@@ -1,6 +1,5 @@
 package ktast.ast.psi
 
-import ktast.ast.ExtrasMap
 import ktast.ast.Node
 import org.jetbrains.kotlin.com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
@@ -11,30 +10,21 @@ import org.jetbrains.kotlin.psi.KtEnumEntry
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.allChildren
 import org.jetbrains.kotlin.psi.psiUtil.siblings
-import java.util.*
-import kotlin.collections.ArrayDeque
 
 /**
  * Converts PSI elements to AST nodes and keeps track of extras.
  */
-open class ConverterWithExtras : Converter(), ExtrasMap {
+open class ConverterWithExtras : Converter() {
     // Sometimes many nodes are created from the same element, but we only want the last node, i.e. the most ancestor node.
     // We remove the previous nodes we've found for the same identity when we see a new one. So we don't have to
     // keep PSI elements around, we hold a map to the element's identity hash code. Then we use that number to tie
     // to the extras to keep duplicates out. Usually using identity hash codes would be problematic due to
     // potential reuse, we know the PSI objects are all around at the same time, so it's good enough.
     protected val psiIdentitiesToNodes = mutableMapOf<Int, Node>()
-    protected val extrasBefore = IdentityHashMap<Node, List<Node.Extra>>()
-    protected val extrasWithin = IdentityHashMap<Node, List<Node.Extra>>()
-    protected val extrasAfter = IdentityHashMap<Node, List<Node.Extra>>()
 
     // This keeps track of ws nodes we've seen before so we don't duplicate them
     private val seenExtraPsiIdentities = mutableSetOf<Int>()
 
-    override fun extrasBefore(node: Node) = extrasBefore[node] ?: emptyList()
-    override fun extrasWithin(node: Node) = extrasWithin[node] ?: emptyList()
-    override fun extrasAfter(node: Node) = extrasAfter[node] ?: emptyList()
-
     override fun onNode(node: Node, element: PsiElement) {
         // We ignore whitespace and comments here to prevent recursion
         if (element is PsiWhiteSpace || element is PsiComment) return
@@ -44,6 +34,8 @@ open class ConverterWithExtras : Converter(), ExtrasMap {
     }
 
     override fun convert(v: KtFile): Node.KotlinFile {
+        psiIdentitiesToNodes.clear()
+        seenExtraPsiIdentities.clear()
         return super.convert(v).also {
             fillWholeExtras(it, v)
         }
@@ -106,21 +98,27 @@ open class ConverterWithExtras : Converter(), ExtrasMap {
 
             private fun fillExtrasBefore(node: Node) {
                 convertExtras(extraElementsSinceLastNode).also {
-                    if (it.isNotEmpty()) extrasBefore[node] = (extrasBefore[node] ?: listOf()) + it
+                    if (it.isNotEmpty()) {
+                        node.supplement.extrasBefore += it
+                    }
                 }
                 extraElementsSinceLastNode.clear()
             }
 
             private fun fillExtrasAfter(node: Node) {
                 convertExtras(extraElementsSinceLastNode).also {
-                    if (it.isNotEmpty()) extrasAfter[node] = (extrasAfter[node] ?: listOf()) + it
+                    if (it.isNotEmpty()) {
+                        node.supplement.extrasAfter += it
+                    }
                 }
                 extraElementsSinceLastNode.clear()
             }
 
             private fun fillExtrasWithin(node: Node) {
                 convertExtras(extraElementsSinceLastNode).also {
-                    if (it.isNotEmpty()) extrasWithin[node] = (extrasWithin[node] ?: listOf()) + it
+                    if (it.isNotEmpty()) {
+                        node.supplement.extrasWithin += it
+                    }
                 }
                 extraElementsSinceLastNode.clear()
             }

ast-psi/src/main/kotlin/ktast/ast/psi/Parser.kt

@@ -17,7 +17,7 @@ import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 /**
  * Parses Kotlin source codes into PSI and converts it to AST.
  */
-open class Parser(protected val converter: Converter = Converter()) {
+open class Parser(protected val converter: Converter = ConverterWithExtras()) {
     companion object : Parser() {
         init {
             // To hide annoying warning on Windows

ast-psi/src/test/kotlin/ktast/ast/DumperTest.kt

@@ -0,0 +1,101 @@
+package ktast.ast
+
+import ktast.ast.psi.Parser
+import org.junit.Test
+import kotlin.test.assertEquals
+
+class DumperTest {
+
+    private val code = """
+        val x = {
+            // x is empty
+        }
+    """.trimIndent()
+
+    @Test
+    fun testWithExtrasAndProperties() {
+        assertDumpedAs(
+            code,
+            """
+                Node.KotlinFile
+                  Node.Declaration.PropertyDeclaration
+                    Node.Keyword.Val{text="val"}
+                    Node.Variable
+                      BEFORE: Node.Extra.Whitespace{text=" "}
+                      Node.Expression.NameExpression{text="x"}
+                      AFTER: Node.Extra.Whitespace{text=" "}
+                    BEFORE: Node.Extra.Whitespace{text=" "}
+                    Node.Expression.LambdaExpression
+                      WITHIN: Node.Extra.Whitespace{text="\n    "}
+                      WITHIN: Node.Extra.Comment{text="// x is empty"}
+                      WITHIN: Node.Extra.Whitespace{text="\n"}
+            """.trimIndent(),
+            withExtras = true,
+            withProperties = true,
+        )
+    }
+
+    @Test
+    fun testWithExtrasButWithoutProperties() {
+        assertDumpedAs(
+            code,
+            """
+                Node.KotlinFile
+                  Node.Declaration.PropertyDeclaration
+                    Node.Keyword.Val
+                    Node.Variable
+                      BEFORE: Node.Extra.Whitespace
+                      Node.Expression.NameExpression
+                      AFTER: Node.Extra.Whitespace
+                    BEFORE: Node.Extra.Whitespace
+                    Node.Expression.LambdaExpression
+                      WITHIN: Node.Extra.Whitespace
+                      WITHIN: Node.Extra.Comment
+                      WITHIN: Node.Extra.Whitespace
+            """.trimIndent(),
+            withExtras = true,
+            withProperties = false,
+        )
+    }
+
+    @Test
+    fun testWithoutExtrasButWithProperties() {
+        assertDumpedAs(
+            code,
+            """
+                Node.KotlinFile
+                  Node.Declaration.PropertyDeclaration
+                    Node.Keyword.Val{text="val"}
+                    Node.Variable
+                      Node.Expression.NameExpression{text="x"}
+                    Node.Expression.LambdaExpression
+            """.trimIndent(),
+            withExtras = false,
+            withProperties = true,
+        )
+    }
+
+    @Test
+    fun testWithoutExtrasOrProperties() {
+        assertDumpedAs(
+            code,
+            """
+                Node.KotlinFile
+                  Node.Declaration.PropertyDeclaration
+                    Node.Keyword.Val
+                    Node.Variable
+                      Node.Expression.NameExpression
+                    Node.Expression.LambdaExpression
+            """.trimIndent(),
+            withExtras = false,
+            withProperties = false,
+        )
+    }
+
+    private fun assertDumpedAs(code: String, expectedDump: String, withExtras: Boolean, withProperties: Boolean) {
+        val node = Parser.parseFile(code)
+        val actualDump = Dumper.dump(node, withExtras = withExtras, withProperties = withProperties)
+        assertEquals(expectedDump.trim(), actualDump.trim())
+    }
+
+}

ast-psi/src/test/kotlin/ktast/ast/MutableVisitorTest.kt

@@ -1,6 +1,5 @@
 package ktast.ast
 
-import ktast.ast.psi.ConverterWithMutableExtras
 import ktast.ast.psi.Parser
 import org.junit.Test
 
@@ -53,12 +52,11 @@ private fun assertMutateAndWriteExact(
     fn: (path: NodePath<*>) -> Node,
     expectedCode: String
 ) {
-    val origExtrasConv = ConverterWithMutableExtras()
-    val origFile = Parser(origExtrasConv).parseFile(origCode)
-    println("----ORIG AST----\n${Dumper.dump(origFile, origExtrasConv)}\n------------")
+    val origFile = Parser.parseFile(origCode)
+    println("----ORIG AST----\n${Dumper.dump(origFile)}\n------------")
 
-    val newFile = MutableVisitor.traverse(origFile, origExtrasConv, fn)
-    val newCode = Writer.write(newFile, origExtrasConv)
+    val newFile = MutableVisitor.traverse(origFile, fn)
+    val newCode = Writer.write(newFile)
     kotlin.test.assertEquals(
         expectedCode.trim(),
         newCode.trim(),

ast-psi/src/test/kotlin/ktast/ast/NodeTest.kt

@@ -1,5 +1,6 @@
 package ktast.ast
 
+import ktast.ast.psi.Converter
 import ktast.ast.psi.Parser
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -11,7 +12,7 @@ import kotlin.test.assertNull
 class DoubleColonExpressionTypeTest(private val code: String) {
     @Test
     fun testType() {
-        val node = Parser.parseFile(code)
+        val node = Parser(Converter()).parseFile(code)
         val properties = node.declarations.filterIsInstance<Node.Declaration.PropertyDeclaration>()
         assertEquals(properties.size, 1)
         properties.forEach { property ->
@@ -51,7 +52,7 @@ class DoubleColonExpressionTypeTest(private val code: String) {
 class DoubleColonExpressionExpressionTest(private val code: String) {
     @Test
     fun testExpression() {
-        val node = Parser.parseFile(code)
+        val node = Parser(Converter()).parseFile(code)
         val properties = node.declarations.filterIsInstance<Node.Declaration.PropertyDeclaration>()
         assertEquals(properties.size, 1)
         properties.forEach { property ->
@@ -80,7 +81,7 @@ class DoubleColonExpressionExpressionTest(private val code: String) {
 class LambdaArgLambdaExpressionTest(private val code: String) {
     @Test
     fun testLambdaExpression() {
-        val node = Parser.parseFile(code)
+        val node = Parser(Converter()).parseFile(code)
         val functionDeclaration = node.declarations.filterIsInstance<Node.Declaration.FunctionDeclaration>().first()
         val callExpressions =
             (functionDeclaration.body as Node.Expression.BlockExpression).statements.filterIsInstance<Node.Expression.CallExpression>()