Merge pull request #19 from orangain/heuristic-whitespaces

Heuristic whitespaces and semicolons

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

@@ -30,9 +30,12 @@ open class Converter {
         names = v.packageNames.map(::convertName),
     ).map(v)
 
-    open fun convertImports(v: KtImportList) = Node.Imports(
-        elements = v.imports.map(::convertImport),
-    ).map(v)
+    open fun convertImports(v: KtImportList): Node.Imports? = if (v.imports.isEmpty())
+        null // Explicitly returns null here. This is because, unlike other PsiElements, KtImportList does exist even when there is no import statement.
+    else
+        Node.Imports(
+            elements = v.imports.map(::convertImport),
+        ).map(v)
 
     open fun convertImport(v: KtImportDirective) = Node.Import(
         importKeyword = convertKeyword(v.importKeyword, Node.Keyword::Import),
@@ -55,12 +58,8 @@ open class Converter {
         name = convertName(v.nameIdentifier ?: error("No name identifier for $v")),
     ).map(v)
 
-    open fun convertDecls(v: KtClassBody) = Node.Decls(
-        elements = v.declarations.map(::convertDecl),
-    ).map(v)
-
     open fun convertDecl(v: KtDeclaration): Node.Decl = when (v) {
-        is KtEnumEntry -> convertEnumEntry(v)
+        is KtEnumEntry -> error("KtEnumEntry is handled in convertEnumEntry")
         is KtClassOrObject -> convertStructured(v)
         is KtAnonymousInitializer -> convertInit(v)
         is KtNamedFunction -> convertFunc(v)
@@ -85,7 +84,7 @@ open class Converter {
                 constraints = convertTypeConstraints(typeConstraintList),
             ).mapNotCorrespondsPsiElement(v)
         },
-        body = v.body?.let { convertDecls(it) },
+        body = v.body?.let(::convertClassBody),
     ).map(v)
 
     open fun convertParents(v: KtSuperTypeList) = Node.Decl.Structured.Parents(
@@ -266,17 +265,22 @@ open class Converter {
             args = v.valueArgumentList?.let(::convertValueArgs)
         ).map(v)
 
-    open fun convertEnumEntry(v: KtEnumEntry) = Node.Decl.EnumEntry(
+    open fun convertEnumEntry(v: KtEnumEntry): Node.EnumEntry = Node.EnumEntry(
         mods = v.modifierList?.let(::convertModifiers),
         name = v.nameIdentifier?.let(::convertName) ?: error("Unnamed enum"),
         args = v.initializerList?.let(::convertValueArgs),
         body = v.body?.let(::convertClassBody),
-        hasComma = v.comma != null,
     ).map(v)
 
-    open fun convertClassBody(v: KtClassBody) = Node.Decl.Structured.Body(
-        decls = v.declarations.map(::convertDecl),
-    ).map(v)
+    open fun convertClassBody(v: KtClassBody): Node.Decl.Structured.Body {
+        val ktEnumEntries = v.declarations.filterIsInstance<KtEnumEntry>()
+        val declarationsExcludingKtEnumEntry = v.declarations.filter { it !is KtEnumEntry }
+        return Node.Decl.Structured.Body(
+            enumEntries = ktEnumEntries.map(::convertEnumEntry),
+            hasTrailingCommaInEnumEntries = ktEnumEntries.lastOrNull()?.comma != null,
+            decls = declarationsExcludingKtEnumEntry.map(::convertDecl),
+        ).map(v)
+    }
 
     open fun convertInitializer(equalsToken: PsiElement, expr: KtExpression, parent: PsiElement) = Node.Initializer(
         equals = convertKeyword(equalsToken, Node.Keyword::Equal),

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

@@ -6,6 +6,7 @@ import org.jetbrains.kotlin.com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.KtImportList
 import org.jetbrains.kotlin.psi.psiUtil.allChildren
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import java.util.*
@@ -48,22 +49,31 @@ open class ConverterWithExtras : Converter(), ExtrasMap {
 
     protected open fun nodeExtraElems(elem: PsiElement): Triple<List<PsiElement>, List<PsiElement>, List<PsiElement>> {
         // Before starts with all directly above ws/comments (reversed to be top-down)
-        val before = elem.siblings(forward = false, withItself = false).takeWhile {
-            it is PsiWhiteSpace || it is PsiComment || it.node.elementType == KtTokens.SEMICOLON
-        }.toList().reversed()
+        val before = elem.siblings(forward = false, withItself = false)
+            .filterNot(::shouldSkip)
+            .takeWhile(::isExtra)
+            .toList().reversed()
 
         // Go over every child...
-        val within = elem.allChildren.filter {
-            it is PsiWhiteSpace || it is PsiComment || it.node.elementType == KtTokens.SEMICOLON
-        }.toList()
+        val within = elem.allChildren
+            .filterNot(::shouldSkip)
+            .filter(::isExtra)
+            .toList()
 
-        val after = elem.siblings(forward = true, withItself = false).takeWhile {
-            it is PsiWhiteSpace || it is PsiComment || it.node.elementType == KtTokens.SEMICOLON
-        }.toList()
+        val after = elem.siblings(forward = true, withItself = false)
+            .filterNot(::shouldSkip)
+            .takeWhile(::isExtra)
+            .toList()
 
         return Triple(before, within, after)
     }
 
+    protected open fun shouldSkip(e: PsiElement) =
+        e is KtImportList && e.imports.isEmpty()
+
+    protected open fun isExtra(e: PsiElement) =
+        e is PsiWhiteSpace || e is PsiComment || e.node.elementType == KtTokens.SEMICOLON
+
     protected open fun convertExtras(elems: List<PsiElement>): List<Node.Extra> = elems.mapNotNull { elem ->
         // Ignore elems we've done before
         val elemId = System.identityHashCode(elem)

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

@@ -15,7 +15,6 @@ class ConverterTest {
             """.trimIndent(),
             """
                 Node.File
-                  Node.Imports
                   Node.Decl.Property
                     Node.Keyword.ValOrVar
                     AFTER: Node.Extra.Whitespace
@@ -38,7 +37,6 @@ class ConverterTest {
             """.trimIndent(),
             """
                 Node.File
-                  Node.Imports
                   Node.Decl.Property
                     Node.Keyword.ValOrVar
                     AFTER: Node.Extra.Whitespace

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

@@ -0,0 +1,57 @@
+package ktast.ast.psi
+
+import ktast.ast.Dumper
+import ktast.ast.Writer
+import org.jetbrains.kotlin.com.intellij.openapi.util.text.StringUtilRt
+import org.junit.Assume
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import kotlin.test.assertEquals
+
+@RunWith(Parameterized::class)
+class CorpusParseAndWriteHeuristicTest(private val unit: Corpus.Unit) {
+
+    @Test
+    fun testParseAndWriteHeuristic() {
+        // In order to test, we parse the test code (failing and validating errors if present),
+        // convert to our AST, write out our AST, re-parse what we wrote, re-convert, and compare
+        try {
+            val origCode = StringUtilRt.convertLineSeparators(unit.read())
+            println("----ORIG CODE----\n$origCode\n------------")
+            val origFile = Parser.parseFile(origCode)
+            val origDump = Dumper.dump(origFile)
+            println("----ORIG AST----\n$origDump\n------------")
+
+            val newCode = Writer.write(origFile)
+            println("----NEW CODE----\n$newCode\n-----------")
+            val newFile = Parser.parseFile(newCode)
+            val newDump = Dumper.dump(newFile)
+
+            assertEquals(origDump, newDump)
+            assertEquals(origFile, newFile)
+        } catch (e: Converter.Unsupported) {
+            Assume.assumeNoException(e.message, e)
+        } catch (e: Parser.ParseError) {
+            if (unit.errorMessages.isEmpty()) throw e
+            assertEquals(unit.errorMessages.toSet(), e.errors.map { it.errorDescription }.toSet())
+            Assume.assumeTrue("Partial parsing not supported (expected parse errors: ${unit.errorMessages})", false)
+        }
+    }
+
+    companion object {
+        const val debug = false
+
+        @JvmStatic
+        @Parameterized.Parameters(name = "{0}")
+        fun data() = Corpus.default
+        // Uncomment to test a specific file
+        //.filter { it.relativePath.toString().endsWith("list\\basic.kt") }
+
+        // Good for quick testing
+//        @JvmStatic @Parameterized.Parameters(name = "{0}")
+//        fun data() = listOf(Corpus.Unit.FromString("temp", """
+//            val lambdaType: (@A() (() -> C))
+//        """))
+    }
+}

ast-psi/src/test/kotlin/ktast/ast/psi/CorpusTest.kt → ast-psi/src/test/kotlin/ktast/ast/psi/CorpusParseAndWriteWithExtrasTest.kt

@@ -11,12 +11,12 @@ import org.junit.runners.Parameterized
 import kotlin.test.assertEquals
 
 @RunWith(Parameterized::class)
-class CorpusTest(private val unit: Corpus.Unit) {
+class CorpusParseAndWriteWithExtrasTest(private val unit: Corpus.Unit) {
 
     @Test
-    fun testParseAndConvert() {
+    fun testParseAndWriteWithExtras() {
         // In order to test, we parse the test code (failing and validating errors if present),
-        // convert to our AST, write out our AST, re-parse what we wrote, re-convert, and compare
+        // convert to our AST, write out our AST, and compare
         try {
             val origExtrasConv = ConverterWithExtras()
             val origCode = StringUtilRt.convertLineSeparators(unit.read())

ast-psi/src/test/resources/localTestData/enumSemicolon.kt

@@ -0,0 +1,7 @@
+enum class Color(val rgb : Int) {
+    RED(0xFF000),
+    GREEN(0x00FF00),
+    BLUE(0x0000FF),
+    ;
+    fun draw() { }
+}

ast/src/commonMain/kotlin/ktast/ast/Dumper.kt

@@ -73,11 +73,18 @@ class Dumper(
         app.append(this::class.qualifiedName?.substring(10)) // 10 means length of "ktast.ast."
         if (verbose) {
             when (this) {
+                is Node.Decl.SecondaryConstructor.DelegationCall -> mapOf("target" to target)
+                is Node.Expr.BinaryOp.Oper.Infix -> mapOf("str" to str)
+                is Node.Expr.BinaryOp.Oper.Token -> mapOf("token" to token)
+                is Node.Expr.UnaryOp -> mapOf("prefix" to prefix)
+                is Node.Expr.UnaryOp.Oper -> mapOf("token" to token)
+                is Node.Expr.TypeOp.Oper -> mapOf("token" to token)
+                is Node.Modifier.AnnotationSet -> mapOf("target" to target)
                 is Node.Modifier.Lit -> mapOf("keyword" to keyword)
-                is Node.Decl.Func -> mapOf("name" to name)
-                is Node.Decl.Property.Variable.Single -> mapOf("name" to name)
                 is Node.Expr.Name -> mapOf("name" to name)
-                is Node.Expr.Const -> mapOf("value" to value)
+                is Node.Expr.Const -> mapOf("value" to value, "form" to form)
+                is Node.Keyword.ValOrVar -> mapOf("token" to token)
+                is Node.Keyword.Declaration -> mapOf("token" to token)
                 is Node.Extra.Comment -> mapOf("text" to text)
                 else -> null
             }?.let {

ast/src/commonMain/kotlin/ktast/ast/MutableVisitor.kt

@@ -39,9 +39,6 @@ open class MutableVisitor(
                     is Node.Import.Alias -> copy(
                         name = visitChildren(name, newCh),
                     )
-                    is Node.Decls -> copy(
-                        elements = visitChildren(elements, newCh),
-                    )
                     is Node.Decl.Structured -> copy(
                         mods = visitChildren(mods, newCh),
                         declarationKeyword = visitChildren(declarationKeyword, newCh),
@@ -75,6 +72,7 @@ open class MutableVisitor(
                         params = visitChildren(params, newCh)
                     )
                     is Node.Decl.Structured.Body -> copy(
+                        enumEntries = visitChildren(enumEntries, newCh),
                         decls = visitChildren(decls, newCh),
                     )
                     is Node.Decl.Init -> copy(
@@ -168,7 +166,7 @@ open class MutableVisitor(
                     is Node.Decl.SecondaryConstructor.DelegationCall -> copy(
                         args = visitChildren(args, newCh)
                     )
-                    is Node.Decl.EnumEntry -> copy(
+                    is Node.EnumEntry -> copy(
                         mods = visitChildren(mods, newCh),
                         name = visitChildren(name, newCh),
                         args = visitChildren(args, newCh),

ast/src/commonMain/kotlin/ktast/ast/Node.kt

@@ -39,12 +39,20 @@ sealed class Node {
         val postMods: List<PostModifier>
     }
 
+    interface StatementsContainer {
+        val statements: List<Statement>
+    }
+
+    interface DeclsContainer {
+        val decls: List<Decl>
+    }
+
     data class File(
         override val anns: List<Modifier.AnnotationSet>,
         override val pkg: Package?,
         override val imports: Imports?,
-        val decls: List<Decl>
-    ) : Node(), Entry
+        override val decls: List<Decl>
+    ) : Node(), Entry, DeclsContainer
 
     data class Script(
         override val anns: List<Modifier.AnnotationSet>,
@@ -91,13 +99,6 @@ sealed class Node {
      */
     sealed class Statement : Node()
 
-    /**
-     * AST node corresponds to KtClassBody.
-     */
-    data class Decls(
-        override val elements: List<Decl>,
-    ) : NodeList<Decl>("{", "}")
-
     sealed class Decl : Statement() {
         /**
          * AST node corresponds to KtClassOrObject.
@@ -110,7 +111,7 @@ sealed class Node {
             val primaryConstructor: PrimaryConstructor?,
             val parents: Parents?,
             val typeConstraints: PostModifier.TypeConstraints?,
-            val body: Decls?,
+            val body: Body?,
         ) : Decl(), WithModifiers {
 
             val isClass = declarationKeyword.token == Keyword.DeclarationToken.CLASS
@@ -168,8 +169,10 @@ sealed class Node {
              * AST node corresponds to KtClassBody.
              */
             data class Body(
-                val decls: List<Decl>
-            ) : Node()
+                val enumEntries: List<EnumEntry>,
+                val hasTrailingCommaInEnumEntries: Boolean,
+                override val decls: List<Decl>,
+            ) : Node(), DeclsContainer
         }
 
         /**
@@ -279,20 +282,22 @@ sealed class Node {
              * AST node corresponds to KtPropertyAccessor.
              */
             sealed class Accessor : Node(), WithModifiers, WithPostModifiers {
+                abstract val body: Func.Body?
+
                 data class Get(
                     override val mods: Modifiers?,
                     val getKeyword: Keyword.Get,
                     val typeRef: TypeRef?,
                     override val postMods: List<PostModifier>,
-                    val body: Func.Body?
+                    override val body: Func.Body?
                 ) : Accessor()
 
                 data class Set(
                     override val mods: Modifiers?,
                     val setKeyword: Keyword.Set,
                     val params: Params?,
                     override val postMods: List<PostModifier>,
-                    val body: Func.Body?
+                    override val body: Func.Body?
                 ) : Accessor()
 
                 /**
@@ -337,22 +342,18 @@ sealed class Node {
             enum class DelegationTarget { THIS, SUPER }
         }
 
-        /**
-         * AST node corresponds to KtEnumEntry.
-         */
-        data class EnumEntry(
-            override val mods: Modifiers?,
-            val name: Expr.Name,
-            val args: ValueArgs?,
-            val body: Structured.Body?,
-            /**
-             * Whether this enum entry has comma or not. All entries excluding the last one should have value true.
-             * The last entry can have both true or false.
-             */
-            val hasComma: Boolean,
-        ) : Decl(), WithModifiers
     }
 
+    /**
+     * AST node corresponds to KtEnumEntry.
+     */
+    data class EnumEntry(
+        override val mods: Modifiers?,
+        val name: Expr.Name,
+        val args: ValueArgs?,
+        val body: Decl.Structured.Body?,
+    ) : Node(), WithModifiers
+
     /**
      * Virtual AST node corresponds to a pair of equals and expression.
      */
@@ -737,7 +738,7 @@ sealed class Node {
              *
              * <Lambda> = { <Param>, <Param> -> <Body> }
              */
-            data class Body(val statements: List<Statement>) : Expr()
+            data class Body(override val statements: List<Statement>) : Expr(), StatementsContainer
         }
 
         /**
@@ -921,7 +922,7 @@ sealed class Node {
         /**
          * AST node corresponds to KtBlockExpression.
          */
-        data class Block(val statements: List<Statement>) : Expr()
+        data class Block(override val statements: List<Statement>) : Expr(), StatementsContainer
     }
 
     /**