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

1 files changed, 13 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
     }