Further simplify invocation counting

Summary: Simply figuring out `invocationCount` even more. Rather than "not counting" the trailing lambda, simply update how we use the value.

Reviewed By: strulovich, hick209

Differential Revision: D33373511

fbshipit-source-id: bae297eaaee572ee922eb5ab931f579a80fceae2

1 files changed, 20 insertions, 17 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 10187c8..5da59d0 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -472,21 +472,11 @@ class KotlinInputAstVisitor(
     // treat the type name-shaped part as a single syntactic unit.
     TypeNameClassifier.typePrefixLength(simpleNames(parts)).ifPresent { prefixes.add(it) }
 
-    var invocationCount = parts.count { it.isCallExpression() }
+    val 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
     // unit. In the normal case we want to preserve the alignment of subsequent
     // method calls, and would emit e.g.:
@@ -504,14 +494,27 @@ class KotlinInputAstVisitor(
     // myField
     //     .foo();
     //
-    if (invocationCount == 1 && firstInvocationIndex > 0) {
-      if (firstInvocationIndex != parts.size - 1 && isFirstInvocationLambda) {
-        prefixes.add(firstInvocationIndex - 1)
-      } else {
-        prefixes.add(firstInvocationIndex)
-      }
+    val singleInvocation = invocationCount == 1
+
+    // For this case, don't count trailing lambdas as call expressions so they look like:
+    // ```
+    // blah.foo().bar().map {
+    //   // blah
+    // }
+    // ```
+    //
+    val singleInvocationWithTrailingLambda = invocationCount == 2 && hasTrailingLambda
+
+    if ((singleInvocation || singleInvocationWithTrailingLambda) && firstInvocationIndex > 0) {
+      prefixes.add(
+          if (firstInvocationIndex != parts.size - 1 && isFirstInvocationLambda) {
+            firstInvocationIndex - 1
+          } else {
+            firstInvocationIndex
+          })
     }
 
+    // keep `super` and `this` attached to the first dereference
     if (prefixes.isEmpty() &&
         (parts.first() is KtSuperExpression || parts.first() is KtThisExpression)) {
       prefixes.add(1)