fix: fix trailing newline issue for kotlinFile. Add test.

* Add trailing newline tests

* fix: Whitespace.countTerminalNewlines is subtly different for kotlinFile vs script.

* Make trim consistent with other tests

---------

Co-authored-by: Tony Robalik <[email protected]>

Author: yissachar

core/src/main/kotlin/cash/grammar/kotlindsl/utils/Whitespace.kt

@@ -4,9 +4,11 @@ import com.squareup.cash.grammar.KotlinLexer
 import com.squareup.cash.grammar.KotlinParser.KotlinFileContext
 import com.squareup.cash.grammar.KotlinParser.ScriptContext
 import com.squareup.cash.grammar.KotlinParser.SemiContext
+import com.squareup.cash.grammar.KotlinParser.SemisContext
 import org.antlr.v4.runtime.CommonTokenStream
 import org.antlr.v4.runtime.ParserRuleContext
 import org.antlr.v4.runtime.Token
+import org.antlr.v4.runtime.tree.ParseTree
 
 /**
  * Utilities for working with whitespace, including newlines, carriage returns, etc.
@@ -119,49 +121,67 @@ public object Whitespace {
    * Use this in conjunction with [trimGently] to maintain original end-of-file formatting.
    */
   public fun countTerminalNewlines(ctx: ScriptContext, tokens: CommonTokenStream): Int {
-    return countTerminalNewlines(ctx as ParserRuleContext, tokens)
+    return ctx.children
+      // Start iterating from EOF
+      .reversed()
+      .asSequence()
+      // Drop `EOF` (every file must have this)
+      .drop(1)
+      // Take only the "semis", which is a semi-colon or a newline, followed by 0 or more newlines
+      .takeWhile { parseTree ->
+        parseTree.javaClass == SemiContext::class.java && hasNoComments(parseTree, tokens)
+      }
+      // Note this is a `SemiContext` (singular!)
+      .filterIsInstance<SemiContext>()
+      // This is the "a newline, followed by 0 or more newlines"
+      .flatMap { it.NL() }
+      .count()
   }
 
   /**
    * Use this in conjunction with [trimGently] to maintain original end-of-file formatting.
    */
   public fun countTerminalNewlines(ctx: KotlinFileContext, tokens: CommonTokenStream): Int {
-    return countTerminalNewlines(ctx as ParserRuleContext, tokens)
-  }
-
-  private fun countTerminalNewlines(ctx: ParserRuleContext, tokens: CommonTokenStream): Int {
     return ctx.children
       // Start iterating from EOF
       .reversed()
       .asSequence()
       // Drop `EOF` (every file must have this)
       .drop(1)
+      .filterIsInstance<ParserRuleContext>()
+      // We need to reverse the order of the children too. If a node doesn't have children, use it.
+      .flatMap { it.children?.reversed() ?: listOf(it) }
       // Take only the "semis", which is a semi-colon or a newline, followed by 0 or more newlines
       .takeWhile { parseTree ->
-        // Because comments are not part of the parse tree (they are shunted to a "hidden" channel),
-        // we need to check for them. Otherwise, we'll "detect" too many newlines at the end of a
-        // file, when that file has only comments and newlines at the end.
-        val hasNoComments = if (parseTree is ParserRuleContext) {
-          val toLeft = parseTree.stop?.let { stop ->
-            tokens.getHiddenTokensToLeft(stop.tokenIndex).orEmpty().isEmpty()
-          } ?: true
-          val toRight = parseTree.stop?.let { stop ->
-            tokens.getHiddenTokensToRight(stop.tokenIndex).orEmpty().isEmpty()
-          } ?: true
-
-          toLeft && toRight
-        } else {
-          true
-        }
-
-        parseTree.javaClass == SemiContext::class.java && hasNoComments
+        parseTree.javaClass == SemisContext::class.java && hasNoComments(parseTree, tokens)
       }
-      .filterIsInstance<SemiContext>()
+      // Note this is a `SemisContext` (plural!)
+      .filterIsInstance<SemisContext>()
       // This is the "a newline, followed by 0 or more newlines"
       .flatMap { it.NL() }
       .count()
   }
 
+  /**
+   * Because comments are not part of the parse tree (they are shunted to a "hidden" channel),
+   * we need to check for them. Otherwise, we'll "detect" too many newlines at the end of a
+   * file, when that file has only comments and newlines at the end.
+   */
+  private fun hasNoComments(parseTree: ParseTree, tokens: CommonTokenStream): Boolean {
+    return if (parseTree is ParserRuleContext) {
+      val toLeft = parseTree.stop?.let { stop ->
+        tokens.getHiddenTokensToLeft(stop.tokenIndex).orEmpty().isEmpty()
+      } ?: true
+      val toRight = parseTree.stop?.let { stop ->
+        tokens.getHiddenTokensToRight(stop.tokenIndex).orEmpty().isEmpty()
+      } ?: true
+
+      toLeft && toRight
+    } else {
+      true
+    }
+  }
+
   /**
    * Use this in conjunction with [countTerminalNewlines] to maintain original end-of-file
    * formatting.

core/src/test/kotlin/cash/grammar/kotlindsl/utils/WhitespaceTest.kt

@@ -4,6 +4,7 @@ import cash.grammar.kotlindsl.parse.Parser
 import cash.grammar.kotlindsl.utils.Blocks.isBuildscript
 import cash.grammar.kotlindsl.utils.Blocks.isSubprojects
 import cash.grammar.kotlindsl.utils.test.TestErrorListener
+import com.squareup.cash.grammar.KotlinParser
 import com.squareup.cash.grammar.KotlinParser.NamedBlockContext
 import com.squareup.cash.grammar.KotlinParserBaseListener
 import org.antlr.v4.runtime.CommonTokenStream
@@ -74,12 +75,68 @@ internal class WhitespaceTest {
       .allSatisfy { assertThat(it).isEqualTo(tuple(" ")) }
   }
 
+  @Test fun `gets trailing newlines for buildscript`() {
+    val buildScript =
+      """
+        plugins {
+          id("kotlin")
+        }
+        
+        subprojects {
+          buildscript {
+            repositories {}
+          }
+        }
+        
+      """.trimIndent()
+
+    val scriptListener = Parser(
+      file = buildScript,
+      errorListener = TestErrorListener {
+        throw RuntimeException("Syntax error: ${it?.message}", it)
+      },
+      listenerFactory = { _, tokens, _ -> TestListener(tokens) }
+    ).listener()
+
+    assertThat(scriptListener.trailingBuildscriptNewlines).isEqualTo(1)
+  }
+
+  @Test fun `gets trailing newlines for kotlin file`() {
+    val file =
+      """
+        class Foo {
+        }
+        
+      """.trimIndent()
+
+    val scriptListener = Parser(
+      file = file.byteInputStream(),
+      errorListener = TestErrorListener {
+        throw RuntimeException("Syntax error: ${it?.message}", it)
+      },
+      startRule = { parser -> parser.kotlinFile() },
+      listenerFactory = { _, tokens, _ -> TestListener(tokens) }
+    ).listener()
+
+    assertThat(scriptListener.trailingKotlinFileNewlines).isEqualTo(1)
+  }
+
   private class TestListener(
     private val tokens: CommonTokenStream
   ) : KotlinParserBaseListener() {
 
     var newlines: List<Token>? = null
     var whitespace: List<Token>? = null
+    var trailingBuildscriptNewlines = 0
+    var trailingKotlinFileNewlines = 0
+
+    override fun enterScript(ctx: KotlinParser.ScriptContext) {
+      trailingBuildscriptNewlines = Whitespace.countTerminalNewlines(ctx, tokens)
+    }
+
+    override fun enterKotlinFile(ctx: KotlinParser.KotlinFileContext) {
+      trailingKotlinFileNewlines = Whitespace.countTerminalNewlines(ctx, tokens)
+    }
 
     override fun exitNamedBlock(ctx: NamedBlockContext) {
       if (ctx.isSubprojects) {