Break after annotations iff it is a block-level expression (#302)

Summary:
These breaks are meaningful, since adding or removing them can change which parts of the expression the annotation applies to.

Fixes https://github.com/facebookincubator/ktfmt/issues/247

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

Test Plan:
Imported from GitHub, without a `Test Plan:` line.

See for changes in codebase: D36070191 (none for us, just the side effect of dealing nicely with annotated lambdas)

Reviewed By: cgrushko

Differential Revision: D35072530

Pulled By: strulovich

fbshipit-source-id: 2660cf57c1f06e3425cc4928483250441c18f03d

2 files changed, 91 insertions, 16 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
index ecf490b..1f56f33 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -1557,20 +1557,30 @@ class KotlinInputAstVisitor(
   override fun visitAnnotatedExpression(expression: KtAnnotatedExpression) {
     builder.sync(expression)
     builder.block(ZERO) {
-      loop@ for (child in expression.node.children()) {
-        when (val psi = child.psi) {
-          is PsiWhiteSpace -> continue@loop
-          is KtAnnotationEntry -> {
-            visit(psi)
-            if (expression.annotationEntries.size != 1) {
-              builder.forcedBreak()
-            } else {
-              builder.breakOp(Doc.FillMode.UNIFIED, " ", ZERO)
-            }
+      val baseExpression = expression.baseExpression
+
+      builder.block(ZERO) {
+        val annotationEntries = expression.annotationEntries
+        for (annotationEntry in annotationEntries) {
+          if (annotationEntry !== annotationEntries.first()) {
+            builder.breakOp(Doc.FillMode.UNIFIED, " ", ZERO)
           }
-          else -> visit(psi)
+          visit(annotationEntry)
         }
       }
+
+      // Binary expressions in a block have a different meaning according to their formatting.
+      // If there in the line above, they refer to the entire expression, if they're in the same
+      // line then only to the first operand of the operator.
+      // We force a break to avoid such semantic changes
+      when {
+        (baseExpression is KtBinaryExpression || baseExpression is KtBinaryExpressionWithTypeRHS) &&
+            expression.parent is KtBlockExpression -> builder.forcedBreak()
+        baseExpression is KtLambdaExpression -> builder.space()
+        else -> builder.breakOp(Doc.FillMode.UNIFIED, " ", ZERO)
+      }
+
+      visit(expression.baseExpression)
     }
   }
 
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 33b5b19..ff83cb9 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -4034,16 +4034,38 @@ class FormatterTest {
   fun `annotated expressions`() =
       assertFormatted(
           """
+      |------------------------------------------------
       |fun f() {
-      |  @Suppress("MagicNumber") add(10)
+      |  @Suppress("MagicNumber") add(10) && add(20)
+      |
+      |  @Suppress("MagicNumber")
+      |  add(10) && add(20)
+      |
+      |  @Anno1 @Anno2(param = Param1::class)
+      |  add(10) && add(20)
       |
       |  @Anno1
       |  @Anno2(param = Param1::class)
       |  @Anno3
       |  @Anno4(param = Param2::class)
-      |  add(10)
+      |  add(10) && add(20)
+      |
+      |  @Anno1
+      |  @Anno2(param = Param1::class)
+      |  @Anno3
+      |  @Anno4(param = Param2::class)
+      |  add(10) && add(20)
+      |
+      |  @Suppress("MagicNumber") add(10) &&
+      |      add(20) &&
+      |      add(30)
+      |
+      |  add(@Suppress("MagicNumber") 10) &&
+      |      add(20) &&
+      |      add(30)
       |}
-      |""".trimMargin())
+      |""".trimMargin(),
+          deduceMaxWidth = true)
 
   @Test
   fun `annotated function declarations`() =
@@ -4688,8 +4710,7 @@ class FormatterTest {
       |    }
       |
       |fun foo() =
-      |    Runnable @Px
-      |    {
+      |    Runnable @Px {
       |      foo()
       |      //
       |    }
@@ -5687,4 +5708,48 @@ class FormatterTest {
       |    .z { it }
       |""".trimMargin(),
           deduceMaxWidth = true)
+
+  @Test
+  fun `annotations for expressions`() =
+      assertFormatted(
+          """
+      |fun f() {
+      |  var b
+      |  @Suppress("UNCHECKED_CAST") b = f(1) as Int
+      |  @Suppress("UNCHECKED_CAST")
+      |  b = f(1) as Int
+      |
+      |  @Suppress("UNCHECKED_CAST") b = f(1) to 5
+      |  @Suppress("UNCHECKED_CAST")
+      |  b = f(1) to 5
+      |
+      |  @Suppress("UNCHECKED_CAST") f(1) as Int + 5
+      |  @Suppress("UNCHECKED_CAST")
+      |  f(1) as Int + 5
+      |
+      |  @Anno1 /* comment */ @Anno2 f(1) as Int
+      |}
+      |""".trimMargin())
+
+  @Test
+  fun `annotations for expressions 2`() {
+    val code =
+        """
+      |fun f() {
+      |  @Suppress("UNCHECKED_CAST") f(1 + f(1) as Int)
+      |  @Suppress("UNCHECKED_CAST")
+      |  f(1 + f(1) as Int)
+      |}
+      |""".trimMargin()
+
+    val expected =
+        """
+      |fun f() {
+      |  @Suppress("UNCHECKED_CAST") f(1 + f(1) as Int)
+      |  @Suppress("UNCHECKED_CAST") f(1 + f(1) as Int)
+      |}
+      |""".trimMargin()
+
+    assertThatFormatting(code).isEqualTo(expected)
+  }
 }