Don't insert blank lines between line comments at the end of files

There was already code/testing for this, but it wasn't covering a strange case where a comment was the first line of the file.
facebook · May 17, 2023 · 85ea42a · 85ea42a
1 parent 1936398
commit 85ea42a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 9 deletions.
diff --git a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -2390,16 +2390,19 @@ class KotlinInputAstVisitor(
 
   override fun visitKtFile(file: KtFile) {
     markForPartialFormat()
-    var importListEmpty = false
+    val importListEmpty = file.importList?.text?.isBlank() ?: true
+
     var isFirst = true
     for (child in file.children) {
-      if (child.text.isBlank()) {
-        importListEmpty = child is KtImportList
-        continue
-      }
-      if (!isFirst && child !is PsiComment && (child !is KtScript || !importListEmpty)) {
-        builder.blankLineWanted(OpsBuilder.BlankLineWanted.YES)
-      }
+      if (child.text.isBlank()) continue
+
+      builder.blankLineWanted(when {
+        isFirst -> OpsBuilder.BlankLineWanted.NO
+        child is PsiComment -> OpsBuilder.BlankLineWanted.NO
+        child is KtScript && importListEmpty -> OpsBuilder.BlankLineWanted.PRESERVE
+        else -> OpsBuilder.BlankLineWanted.YES
+      })
+
       visit(child)
       isFirst = false
     }

diff --git a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -440,7 +440,7 @@ class FormatterTest {
           deduceMaxWidth = true)
 
   @Test
-  fun `don't keep adding newlines between these two comments when they're at end of file`() =
+  fun `don't keep adding newlines between these two comments when they're at end of file`() {
       assertFormatted(
           """
      |package foo
@@ -450,6 +450,17 @@ class FormatterTest {
      |"""
               .trimMargin())
 
+    assertFormatted(
+          """
+     |// Comment as first element
+     |package foo
+     |// a
+     |
+     |/* Another comment */
+     |"""
+              .trimMargin())
+  }
+
   @Test
   fun `properties with line comment above initializer`() =
       assertFormatted(