Better formatting in some cases

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/interfaces.kt

@@ -30,7 +30,7 @@ interface ChainedCallGenerator {
     /**
      * @param pattern pattern string for generating the part of the call to the right from the dot
      */
-    fun generate(pattern: String, vararg args: Any): KtExpression
+    fun generate(pattern: String, vararg args: Any, receiver: KtExpression = this.receiver, safeCall: Boolean = false): KtExpression
 }
 
 /**
@@ -42,7 +42,7 @@ interface Transformation {
 
     val presentation: String
 
-    open fun buildPresentation(prevTransformationsPresentation: String?): String {
+    fun buildPresentation(prevTransformationsPresentation: String?): String {
         return if (prevTransformationsPresentation != null)
             prevTransformationsPresentation + "." + presentation
         else
@@ -50,6 +50,10 @@ interface Transformation {
     }
 
     fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression
+
+
+    val chainCallCount: Int
+        get() = 1
 }
 
 /**
@@ -59,7 +63,6 @@ interface SequenceTransformation : Transformation {
     fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? = null
 
     val affectsIndex: Boolean
-
 }
 
 /**

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt

@@ -169,7 +169,8 @@ private fun ResultTransformationMatch.generateCallChain(loop: KtForExpression):
         sequenceTransformations = sequenceTransformations.dropLast(1)
     }
 
-    val lineBreak = if (sequenceTransformations.isNotEmpty()) "\n" else ""
+    val chainCallCount = sequenceTransformations.sumBy { it.chainCallCount } + resultTransformation.chainCallCount
+    val lineBreak = if (chainCallCount > 1) "\n" else ""
 
     var callChain = loop.loopRange!!
 
@@ -178,9 +179,10 @@ private fun ResultTransformationMatch.generateCallChain(loop: KtForExpression):
         override val receiver: KtExpression
             get() = callChain
 
-        override fun generate(pattern: String, vararg args: Any): KtExpression {
+        override fun generate(pattern: String, vararg args: Any, receiver: KtExpression, safeCall: Boolean): KtExpression {
-            val newPattern = "\$${args.size}$lineBreak.$pattern"
+            val dot = if (safeCall) "?." else "."
-            return psiFactory.createExpressionByPattern(newPattern, *args, callChain)
+            val newPattern = "$" + args.size + lineBreak + dot + pattern
+            return psiFactory.createExpressionByPattern(newPattern, *args, receiver)
         }
     }
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/AddToCollectionTransformation.kt

@@ -65,6 +65,9 @@ class AddToCollectionTransformation(
             "+="
     }
 
+    override val chainCallCount: Int
+        get() = 0
+
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
         return KtPsiFactory(loop).createExpressionByPattern("$0 += $1", targetCollection, chainedCallGenerator.receiver)
     }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndAssignTransformation.kt

@@ -43,6 +43,9 @@ class FindAndAssignTransformation(
     override val presentation: String
         get() = generator.functionName + (if (filter != null) "{}" else "()")
 
+    override val chainCallCount: Int
+        get() = generator.chainCallCount
+
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
         return generator.generate(chainedCallGenerator, filter)
     }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndReturnTransformation.kt

@@ -47,6 +47,9 @@ class FindAndReturnTransformation(
     override val presentation: String
         get() = generator.functionName + (if (filter != null) "{}" else "()")
 
+    override val chainCallCount: Int
+        get() = generator.chainCallCount
+
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
         return generator.generate(chainedCallGenerator, filter)
     }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/utils.kt

@@ -114,7 +114,11 @@ fun KtProperty.hasWriteUsages(): Boolean {
 
 interface FindOperatorGenerator {
     val functionName: String
+
     fun generate(chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?): KtExpression
+
+    val chainCallCount: Int
+        get() = 1
 }
 
 fun buildFindOperationGenerator(
@@ -150,10 +154,7 @@ fun buildFindOperationGenerator(
         // we cannot use ?: if found value can be null
         if (inputVariableCanHoldNull) return null
 
-        return object: FindOperatorGenerator {
+        return object: FindOperatorGenerator by this {
-            override val functionName: String
-                get() = this@useElvisOperatorIfNeeded.functionName
-
             override fun generate(chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?): KtExpression {
                 val generated = this@useElvisOperatorIfNeeded.generate(chainedCallGenerator, filter)
                 return KtPsiFactory(generated).createExpressionByPattern("$0 ?: $1", generated, valueIfNotFound)
@@ -184,9 +185,12 @@ fun buildFindOperationGenerator(
                         override val functionName: String
                             get() = "firstOrNull"
 
+                        override val chainCallCount: Int
+                            get() = 2
+
                         override fun generate(chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?): KtExpression {
                             val findFirstCall = generateChainedCall(functionName, chainedCallGenerator, filter)
-                            return KtPsiFactory(findFirstCall).createExpressionByPattern("$0?.$1", findFirstCall, selector)
+                            return chainedCallGenerator.generate("$0", selector, receiver = findFirstCall, safeCall = true)
                         }
                     }.useElvisOperatorIfNeeded()
                 }
@@ -197,12 +201,15 @@ fun buildFindOperationGenerator(
 
             return object: FindOperatorGenerator {
                 override val functionName: String
-                    get() = "firstOrNull" //TODO
+                    get() = "firstOrNull"
+
+                override val chainCallCount: Int
+                    get() = 2 // also includes "let"
 
                 override fun generate(chainedCallGenerator: ChainedCallGenerator, filter: KtExpression?): KtExpression {
                     val findFirstCall = generateChainedCall(functionName, chainedCallGenerator, filter)
                     val letBody = generateLambda(inputVariable, valueIfFound)
-                    return KtPsiFactory(findFirstCall).createExpressionByPattern("$0?.let $1:'{}'", findFirstCall, letBody)
+                    return chainedCallGenerator.generate("let $0:'{}'", letBody, receiver = findFirstCall, safeCall = true)
                 }
             }.useElvisOperatorIfNeeded()
         }

idea/testData/intentions/loopToCallChain/firstOrNull_let.kt.after

@@ -1,7 +1,9 @@
 // WITH_RUNTIME
 // INTENTION_TEXT: "Replace with 'firstOrNull{}'"
 fun foo(list: List<String>) {
-    <caret>val result: String? = list.firstOrNull { it.length > 0 }?.let { bar(it) }
+    <caret>val result: String? = list
+            .firstOrNull { it.length > 0 }
+            ?.let { bar(it) }
 }
 
 fun bar(s: String): String = s

idea/testData/intentions/loopToCallChain/firstOrNull_let2.kt.after

@@ -1,7 +1,9 @@
 // WITH_RUNTIME
 // INTENTION_TEXT: "Replace with 'firstOrNull{}'"
 fun foo(list: List<String>) {
-    <caret>val result: String? = list.firstOrNull { it.length > 0 }?.let { it.substring(0, it.length - 1) }
+    <caret>val result: String? = list
+            .firstOrNull { it.length > 0 }
+            ?.let { it.substring(0, it.length - 1) }
 }
 
 fun bar(s: String): String = s

idea/testData/intentions/loopToCallChain/firstOrNull_letOrNotNull.kt.after

@@ -1,7 +1,9 @@
 // WITH_RUNTIME
 // INTENTION_TEXT: "Replace with 'firstOrNull{}'"
 fun foo(list: List<String>) {
-    <caret>val result = list.firstOrNull { it.length > 0 }?.let { bar(it) } ?: ""
+    <caret>val result = list
+            .firstOrNull { it.length > 0 }
+            ?.let { bar(it) } ?: ""
 }
 
 fun bar(s: String): String = s

idea/testData/intentions/loopToCallChain/firstOrNull_returnExpression.kt.after

@@ -1,5 +1,7 @@
 // WITH_RUNTIME
 // INTENTION_TEXT: "Replace with 'firstOrNull{}'"
 fun foo(list: List<String>): Int? {
-    <caret>return list.firstOrNull { it.isNotEmpty() }?.length
+    <caret>return list
+            .firstOrNull { it.isNotEmpty() }
+            ?.length
 }

idea/testData/intentions/loopToCallChain/firstOrNull_returnExpressionOrNotNull.kt.after

@@ -1,5 +1,7 @@
 // WITH_RUNTIME
 // INTENTION_TEXT: "Replace with 'firstOrNull{}'"
 fun foo(list: List<String>): Int {
-    <caret>return list.firstOrNull { it.isNotEmpty() }?.length ?: -1
+    <caret>return list
+            .firstOrNull { it.isNotEmpty() }
+            ?.length ?: -1
 }

idea/testData/intentions/loopToCallChain/firstOrNull_safeDotExpression.kt.after

@@ -1,5 +1,7 @@
 // WITH_RUNTIME
 // INTENTION_TEXT: "Replace with 'firstOrNull{}'"
 fun foo(list: List<String?>) {
-    <caret>val result: String? = list.firstOrNull { it != "" }?.substring(1)
+    <caret>val result: String? = list
+            .firstOrNull { it != "" }
+            ?.substring(1)
 }