Support labeled lambdas in block-like scoping functions (#403)

Summary: Pull Request resolved: #403 Reviewed By: strulovich Differential Revision: D46808245 Pulled By: cgrushko fbshipit-source-id: a6e292cb37c0dc822daf0a1e5ebb6450fcf0f0dd
facebook · Jun 18, 2023 · 662a3a5 · 662a3a5
1 parent b6e5f35
commit 662a3a5
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 18 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
@@ -1415,16 +1415,24 @@ class KotlinInputAstVisitor(
     if (expression.getPrevSiblingIgnoringWhitespace() is PsiComment) {
       return false // Leading comments cause weird indentation.
     }
-    if (expression is KtLambdaExpression) {
-      return true
+
+    var carry = expression
+    if (carry is KtCallExpression) {
+      if (carry.valueArgumentList?.leftParenthesis == null &&
+          carry.lambdaArguments.isNotEmpty() &&
+          carry.typeArgumentList?.arguments.isNullOrEmpty()) {
+        carry = carry.lambdaArguments[0].getArgumentExpression()
+      } else {
+        return false
+      }
     }
-    if (expression is KtCallExpression &&
-        expression.valueArgumentList?.leftParenthesis == null &&
-        expression.lambdaArguments.isNotEmpty() &&
-        expression.typeArgumentList?.arguments.isNullOrEmpty() &&
-        expression.lambdaArguments.first().getArgumentExpression() is KtLambdaExpression) {
+    if (carry is KtLabeledExpression) {
+      carry = carry.baseExpression
+    }
+    if (carry is KtLambdaExpression) {
       return true
     }
+
     return false
   }
 
@@ -1433,18 +1441,22 @@ class KotlinInputAstVisitor(
     val breakToExpr = genSym()
     builder.breakOp(Doc.FillMode.INDEPENDENT, " ", expressionBreakIndent, Optional.of(breakToExpr))
 
-    val lambdaExpression =
-        when (expr) {
-          is KtLambdaExpression -> expr
-          is KtCallExpression -> {
-            visit(expr.calleeExpression)
-            builder.space()
-            expr.lambdaArguments[0].getLambdaExpression() ?: fail()
-          }
-          else -> throw AssertionError(expr)
-        }
+    var carry = expr
+    if (carry is KtCallExpression) {
+      visit(carry.calleeExpression)
+      builder.space()
+      carry = carry.lambdaArguments[0].getArgumentExpression()
+    }
+    if (carry is KtLabeledExpression) {
+      visit(carry.labelQualifier)
+      carry = carry.baseExpression ?: fail()
+    }
+    if (carry is KtLambdaExpression) {
+      visitLambdaExpressionInternal(carry, brokeBeforeBrace = breakToExpr)
+      return
+    }
 
-    visitLambdaExpressionInternal(lambdaExpression, brokeBeforeBrace = breakToExpr)
+    throw AssertionError(carry)
   }
 
   override fun visitClassOrObject(classOrObject: KtClassOrObject) {

diff --git a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -5341,13 +5341,24 @@ class FormatterTest {
       |  //
       |}
       |
+      |fun foo() = scope label@{
+      |  foo()
+      |  //
+      |}
+      |
       |fun foo() =
       |    coroutineScope { x ->
       |      foo()
       |      //
       |    }
       |
       |fun foo() =
+      |    coroutineScope label@{
+      |      foo()
+      |      //
+      |    }
+      |
+      |fun foo() =
       |    Runnable @Px {
       |      foo()
       |      //