#291: move comments between imports above import list (#292)

Summary:
Fix for #291

Pull Request resolved: #292

Reviewed By: cgrushko

Differential Revision: D34424976

Pulled By: strulovich

fbshipit-source-id: 199b6d6644c4e51b08ed06fd3ab052f3b9339ba5

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")
  }
 }

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