From b3cf13a7307fe568b45f0dca39b7ffe8b600de88 Mon Sep 17 00:00:00 2001
From: Beth Cutler <bcutler@google.com>
Date: Wed, 23 Feb 2022 11:22:38 -0800
Subject: [PATCH] #291: move comments between imports above import list (#292)

Summary:
Fix for https://github.com/facebookincubator/ktfmt/issues/291

Pull Request resolved: https://github.com/facebookincubator/ktfmt/pull/292

Reviewed By: cgrushko

Differential Revision: D34424976

Pulled By: strulovich

fbshipit-source-id: 199b6d6644c4e51b08ed06fd3ab052f3b9339ba5
---
 .../com/facebook/ktfmt/format/Formatter.kt    | 31 +++++++--------
 .../facebook/ktfmt/format/FormatterTest.kt    | 39 ++++++++++++-------
 2 files changed, 41 insertions(+), 29 deletions(-)

diff --git a/core/src/main/java/com/facebook/ktfmt/format/Formatter.kt b/core/src/main/java/com/facebook/ktfmt/format/Formatter.kt
index d50c7846..3bf2e2aa 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/Formatter.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/Formatter.kt
@@ -33,6 +33,7 @@ import com.google.googlejavaformat.java.FormatterException
 import com.google.googlejavaformat.java.JavaOutput
 import org.jetbrains.kotlin.com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.com.intellij.openapi.util.text.StringUtilRt
+import org.jetbrains.kotlin.com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
@@ -144,22 +145,19 @@ object Formatter {
       return code
     }
 
-    fun findNonImportElement(): PsiElement? {
-      var element = importList.firstChild
-      while (element != null) {
-        if (element !is KtImportDirective && element !is PsiWhiteSpace) {
-          return element
-        }
-        element = element.nextSibling
+    val commentList = mutableListOf<PsiElement>()
+    // Find non-import elements; comments are moved, in order, to the top of the import list. Other
+    // non-import elements throw a ParseError.
+    var element = importList.firstChild
+    while (element != null) {
+      if (element is PsiComment) {
+        commentList.add(element)
+      } else if (element !is KtImportDirective && element !is PsiWhiteSpace) {
+        throw ParseError(
+            "Imports not contiguous: " + element.text,
+            StringUtil.offsetToLineColumn(code, element.startOffset))
       }
-      return null
-    }
-
-    val nonImportElement = findNonImportElement()
-    if (nonImportElement != null) {
-      throw ParseError(
-          "Imports not contiguous (perhaps a comment separates them?): " + nonImportElement.text,
-          StringUtil.offsetToLineColumn(code, nonImportElement.startOffset))
+      element = element.nextSibling
     }
     fun canonicalText(importDirective: KtImportDirective) =
         importDirective.importedFqName?.asString() +
@@ -169,10 +167,11 @@ object Formatter {
             if (importDirective.isAllUnder) "*" else ""
 
     val sortedImports = importList.imports.sortedBy(::canonicalText).distinctBy(::canonicalText)
+    val importsWithComments = commentList + sortedImports
 
     return code.replaceRange(
         importList.startOffset,
         importList.endOffset,
-        sortedImports.joinToString(separator = "\n") { imprt -> imprt.text } + "\n")
+        importsWithComments.joinToString(separator = "\n") { imprt -> imprt.text } + "\n")
   }
 }
diff --git a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
index 3527d981..b96ac267 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -1031,24 +1031,37 @@ class FormatterTest {
   }
 
   @Test
-  fun `comments between imports are not allowed`() {
+  fun `comments between imports are moved above import list`() {
     val code =
         """
         |package com.facebook.ktfmt
         |
-        |import com.google.common.truth.Subject
-        |/* add */
-        |import com.google.common.truth.FailureMetadata as Foo
+        |/* leading comment */
+        |import com.example.abc
+        |/* internal comment 1 */
+        |import com.example.bcd
+        |// internal comment 2
+        |import com.example.Sample
+        |// trailing comment
+        |
+        |val x = Sample(abc, bcd)
         |""".trimMargin()
-
-    try {
-      Formatter.format(code)
-      fail()
-    } catch (e: ParseError) {
-      assertThat(e.errorDescription).contains("Imports not contiguous")
-      assertThat(e.lineColumn.line).isEqualTo(3)
-      assertThat(e.lineColumn.column).isEqualTo(0)
-    }
+    val expected =
+        """
+        |package com.facebook.ktfmt
+        |
+        |/* leading comment */
+        |/* internal comment 1 */
+        |// internal comment 2
+        |import com.example.Sample
+        |import com.example.abc
+        |import com.example.bcd
+        |
+        |// trailing comment
+        |
+        |val x = Sample(abc, bcd)
+        |""".trimMargin()
+    assertThatFormatting(code).isEqualTo(expected)
   }
 
   @Test