From cd5fdc76650f8a6930a1c80b3e9f954b7a18c2a3 Mon Sep 17 00:00:00 2001
From: nickreid <nickreid@google.com>
Date: Thu, 2 Feb 2023 09:28:43 -0800
Subject: Fix indentation of initializer-like expressions with leading comments
 (#382)

Summary: Pull Request resolved: https://github.com/facebook/ktfmt/pull/382

Reviewed By: hick209

Differential Revision: D42801455

Pulled By: cgrushko

fbshipit-source-id: c53eb78591d5c23301bb2bafd84f89c5bb7a7a33
---
 .../facebook/ktfmt/format/KotlinInputAstVisitor.kt | 15 ++++++-
 .../com/facebook/ktfmt/format/FormatterTest.kt     | 52 ++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 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 86d5205..433ae1a 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -121,6 +121,7 @@ import org.jetbrains.kotlin.psi.KtWhenConditionWithExpression
 import org.jetbrains.kotlin.psi.KtWhenExpression
 import org.jetbrains.kotlin.psi.KtWhileExpression
 import org.jetbrains.kotlin.psi.psiUtil.children
+import org.jetbrains.kotlin.psi.psiUtil.getPrevSiblingIgnoringWhitespace
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.psi.psiUtil.startsWithComment
 
@@ -1346,7 +1347,10 @@ class KotlinInputAstVisitor(
           visit(delegate)
         } else {
           builder.breakOp(Doc.FillMode.UNIFIED, " ", expressionBreakIndent)
-          builder.block(expressionBreakIndent) { visit(delegate) }
+          builder.block(expressionBreakIndent) {
+            builder.fenceComments()
+            visit(delegate)
+          }
         }
       } else if (initializer != null) {
         builder.space()
@@ -1355,7 +1359,10 @@ class KotlinInputAstVisitor(
           visitLambdaOrScopingFunction(initializer)
         } else {
           builder.breakOp(Doc.FillMode.UNIFIED, " ", expressionBreakIndent)
-          builder.block(expressionBreakIndent) { visit(initializer) }
+          builder.block(expressionBreakIndent) {
+            builder.fenceComments()
+            visit(initializer)
+          }
         }
       }
     }
@@ -1409,6 +1416,10 @@ class KotlinInputAstVisitor(
    * 2. '... = Runnable @Annotation { ... }' due to the annotation
    */
   private fun isLambdaOrScopingFunction(expression: KtExpression?): Boolean {
+    if (expression == null) return false
+    if (expression.getPrevSiblingIgnoringWhitespace() is PsiComment) {
+      return false // Leading comments cause weird indentation.
+    }
     if (expression is KtLambdaExpression) {
       return true
     }
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 88b16fe..ae9d05e 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -445,6 +445,58 @@ class FormatterTest {
      |"""
               .trimMargin())
 
+  @Test
+  fun `properties with line comment above initializer`() =
+      assertFormatted(
+          """
+      |class Foo {
+      |  var x: Int =
+      |      // Comment
+      |      0
+      |
+      |  var y: Int =
+      |      // Comment
+      |      scope {
+      |        0 //
+      |      }
+      |
+      |  var z: Int =
+      |      // Comment
+      |      if (cond) {
+      |        0
+      |      } else {
+      |        1
+      |      }
+      |}
+      |"""
+              .trimMargin())
+
+  @Test
+  fun `properties with line comment above delegate`() =
+      assertFormatted(
+          """
+      |class Foo {
+      |  var x: Int by
+      |      // Comment
+      |      0
+      |
+      |  var y: Int by
+      |      // Comment
+      |      scope {
+      |        0 //
+      |      }
+      |
+      |  var z: Int by
+      |      // Comment
+      |      if (cond) {
+      |        0
+      |      } else {
+      |        1
+      |      }
+      |}
+      |"""
+              .trimMargin())
+
   @Test
   fun `properties with accessors`() =
       assertFormatted(
-- 
cgit v1.2.3