Simplify invocation counting in emitQualifiedExpression

Summary: The current logic for figuring out `invocationCount` (as well as some other flags) is a little hard to follow, so rewrite it to be clearer. This has the downside of going through the list twice (rather than once).

Reviewed By: strulovich

Differential Revision: D33372295

fbshipit-source-id: be3b2e0be034ef02baa2df0e2de1eaaa5c8f3e42

1 files changed, 23 insertions, 27 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 0c59b7b..10187c8 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -472,31 +472,19 @@ class KotlinInputAstVisitor(
     // treat the type name-shaped part as a single syntactic unit.
     TypeNameClassifier.typePrefixLength(simpleNames(parts)).ifPresent { prefixes.add(it) }
 
-    var invocationCount = 0
-    var firstInvocationIndex = -1
-    var isFirstInvocationLambda = false
-    for ((i, part) in parts.withIndex()) {
-      val callExpression = extractCallExpression(part)
-      if (callExpression != null) {
-        // Don't count trailing lambdas as call expressions so they look like
-        // ```
-        // blah.foo().bar().map {
-        //   // blah
-        // }
-        // ```
-        if (invocationCount > 0 &&
-            i == parts.size - 1 &&
-            callExpression.lambdaArguments.isNotEmpty()) {
-          continue
-        }
-        invocationCount++
-        if (firstInvocationIndex < 0) {
-          firstInvocationIndex = i
-          if (callExpression.lambdaArguments.isNotEmpty()) {
-            isFirstInvocationLambda = true
-          }
-        }
-      }
+    var invocationCount = parts.count { it.isCallExpression() }
+    val firstInvocationIndex = parts.indexOfFirst { it.isCallExpression() }
+    val isFirstInvocationLambda = parts.getOrNull(firstInvocationIndex)?.isLambda() ?: false
+    val hasTrailingLambda = parts.last().isLambda()
+
+    // Don't count trailing lambdas as call expressions so they look like
+    // ```
+    // blah.foo().bar().map {
+    //   // blah
+    // }
+    // ```
+    if (invocationCount > 1 && hasTrailingLambda) {
+      invocationCount--
     }
 
     // If there's only one invocation, treat leading field accesses as a single
@@ -516,8 +504,6 @@ class KotlinInputAstVisitor(
     // myField
     //     .foo();
     //
-    val hasTrailingLambda =
-        extractCallExpression(parts.last())?.lambdaArguments?.isNotEmpty() == true
     if (invocationCount == 1 && firstInvocationIndex > 0) {
       if (firstInvocationIndex != parts.size - 1 && isFirstInvocationLambda) {
         prefixes.add(firstInvocationIndex - 1)
@@ -557,6 +543,16 @@ class KotlinInputAstVisitor(
     return parts.toList()
   }
 
+  /** Returns true if the expression represents an invocation */
+  private fun KtExpression.isCallExpression(): Boolean {
+    return extractCallExpression(this) != null
+  }
+
+  /** Returns true if the expression represents an invocation that is also a lambda */
+  private fun KtExpression.isLambda(): Boolean {
+    return extractCallExpression(this)?.lambdaArguments?.isNotEmpty() ?: false
+  }
+
   /**
    * emitQualifiedExpression formats call expressions that are either part of a qualified
    * expression, or standing alone. This method makes it easier to handle both cases uniformly.