From 8118a02f42bfdac0f4bc75277a1ab38ae79dafea Mon Sep 17 00:00:00 2001
From: nickreid <nickreid@google.com>
Date: Tue, 22 Mar 2022 15:13:46 -0700
Subject: [PATCH] Preserve line breaks associated with expression annotations.

These breaks are meaningful, since adding or removing them can change which parts of the expression the annotation applies to.
---
 .../ktfmt/format/KotlinInputAstVisitor.kt     | 22 ++++---
 .../facebook/ktfmt/format/FormatterTest.kt    | 60 ++++++++++++++++++-
 2 files changed, 67 insertions(+), 15 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 217c82ea..55bc083f 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -1656,20 +1656,18 @@ 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)
-            }
-          }
-          else -> visit(psi)
+      val hasBreak =
+          expression.node.children().filter { it is PsiWhiteSpace }.any { it.text.contains('\n') }
+
+      expression.annotationEntries.forEach {
+        visit(it)
+        if (hasBreak) {
+          builder.forcedBreak()
+        } else {
+          builder.space()
         }
       }
+      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 60bad094..1bf9087b 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -3824,16 +3824,70 @@ 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)
       |  @Anno3
       |  @Anno4(param = Param2::class)
-      |  add(10)
+      |  add(10) && add(20)
       |}
-      |""".trimMargin())
+      |""".trimMargin(),
+          deduceMaxWidth = true)
+
+  @Test
+  fun `annotated expressions scope is preserverd when breaks are added`() {
+    val code =
+        """
+      |/////////////////////////////////////////////
+      |fun f() {
+      |  @EntireExpression
+      |  @AlsoEntireExpression add(10) && add(20)
+      |
+      |  @Just10 @AlsoJust10 @StillJust10 add(10) && add(20)
+      |
+      |  @Just10 @AlsoJust10 add(10) && @Just20 add(20)
+      |
+      |  @Anno1 @Anno2(param = Param1::class)
+      |  @Anno3 @Anno4(param = Param2::class)
+      |  add(10) && add(20)
+      |
+      |  @EntireExpression(breakInside = "afterThis") add(10) && add(20)
+      |}
+      |""".trimMargin()
+    val expected =
+        """
+      |/////////////////////////////////////////////
+      |fun f() {
+      |  @EntireExpression
+      |  @AlsoEntireExpression
+      |  add(10) && add(20)
+      |
+      |  @Just10 @AlsoJust10 @StillJust10 add(10) &&
+      |      add(20)
+      |
+      |  @Just10 @AlsoJust10 add(10) &&
+      |      @Just20 add(20)
+      |
+      |  @Anno1
+      |  @Anno2(param = Param1::class)
+      |  @Anno3
+      |  @Anno4(param = Param2::class)
+      |  add(10) && add(20)
+      |
+      |  @EntireExpression(
+      |      breakInside = "afterThis") add(10) &&
+      |      add(20)
+      |}
+      |""".trimMargin()
+
+    assertThatFormatting(code).withOptions(FormattingOptions(maxWidth = 45)).isEqualTo(expected)
+  }
 
   @Test
   fun `annotated function declarations`() =