Improved logic for chained lambda indentation

#KT-22071 Fixed

idea/formatter/src/org/jetbrains/kotlin/idea/formatter/KotlinCommonBlock.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.formatter
 
 import com.intellij.formatting.*
 import com.intellij.lang.ASTNode
+import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.TokenType
 import com.intellij.psi.codeStyle.CodeStyleSettings
@@ -27,7 +28,6 @@ import org.jetbrains.kotlin.psi.psiUtil.children
 import org.jetbrains.kotlin.psi.psiUtil.leaves
 import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.psi.psiUtil.siblings
-import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 
 private val QUALIFIED_OPERATION = TokenSet.create(DOT, SAFE_ACCESS)
 private val QUALIFIED_EXPRESSIONS = TokenSet.create(KtNodeTypes.DOT_QUALIFIED_EXPRESSION, KtNodeTypes.SAFE_ACCESS_EXPRESSION)
@@ -112,14 +112,16 @@ abstract class KotlinCommonBlock(
             // Create fake ".something" or "?.something" block here, so child indentation will be
             // relative to it when it starts from new line (see Indent javadoc).
 
-            val indentType = if (settings.kotlinCustomSettings.CONTINUATION_INDENT_FOR_CHAINED_CALLS)
-                Indent.Type.CONTINUATION
-            else
-                Indent.Type.NORMAL
             val isNonFirstChainedCall = operationBlockIndex > 0 && isCallBlock(nodeSubBlocks[operationBlockIndex - 1])
+            val enforceIndentToChildren = isNonFirstChainedCall && hasLineBreakBefore(nodeSubBlocks[operationBlockIndex])
+            val indentType = if (settings.kotlinCustomSettings.CONTINUATION_INDENT_FOR_CHAINED_CALLS) {
+                if (enforceIndentToChildren) Indent.Type.CONTINUATION else Indent.Type.CONTINUATION_WITHOUT_FIRST
+            } else {
+                Indent.Type.NORMAL
+            }
             val indent = Indent.getIndent(
                 indentType, false,
-                isNonFirstChainedCall && hasLineBreakBefore(nodeSubBlocks[operationBlockIndex - 1])
+                enforceIndentToChildren
             )
             val wrap = if ((settings.kotlinCommonSettings.WRAP_FIRST_METHOD_IN_CALL_CHAIN || isNonFirstChainedCall) &&
                 canWrapCallChain(node))
@@ -566,18 +568,10 @@ fun needWrapArgumentList(psi: PsiElement): Boolean {
 }
 
 private fun hasLineBreakBefore(block: ASTBlock): Boolean {
-    val topLevelParent = block.node.parents().firstOrNull {
-        it.treeParent.elementType == KtNodeTypes.BLOCK || it.treeParent.elementType == KtStubElementTypes.FILE
-    } ?: return false
-    for (leaf in block.node.leaves(forward = false)) {
-        if (leaf.textContains('\n')) {
-            return true
-        }
-        if (leaf.textRange.startOffset == topLevelParent.startOffset) {
-            break
-        }
-    }
-    return false
+    val prevSibling = block.node.leaves(false)
+        .dropWhile { it.psi is PsiComment || it.elementType == KtTokens.RBRACE }
+        .firstOrNull()
+    return prevSibling?.elementType == TokenType.WHITE_SPACE && prevSibling?.textContains('\n') == true
 }
 
 fun NodeIndentStrategy.PositionStrategy.continuationIf(

idea/testData/formatter/KT22071.after.kt

@@ -0,0 +1,51 @@
+package templates
+
+import kotlin.coroutines.experimental.buildSequence
+import kotlin.reflect.KTypeProjection
+import kotlin.reflect.full.createType
+import kotlin.reflect.full.isSubtypeOf
+
+
+typealias TemplateGroup = () -> Sequence<MemberTemplate>
+
+fun templateGroupOf(vararg templates: MemberTemplate): TemplateGroup = { templates.asSequence() }
+
+abstract class TemplateGroupBase : TemplateGroup {
+
+    override fun invoke(): Sequence<MemberTemplate> = buildSequence {
+        with(this@TemplateGroupBase) {
+            this::class.members.filter { it.name.startsWith("f_") }.forEach {
+                require(it.parameters.size == 1) { "Member $it violates naming convention" }
+                when {
+                    it.returnType.isSubtypeOf(typeMemberTemplate) ->
+                        yield(it.call(this) as MemberTemplate)
+                    it.returnType.isSubtypeOf(typeIterableOfMemberTemplates) ->
+                        @Suppress("UNCHECKED_CAST")
+                        yieldAll(it.call(this) as Iterable<MemberTemplate>)
+                    else ->
+                        error("Member $it violates naming convention")
+                }
+            }
+        }
+    }.run {
+        if (defaultActions.isEmpty()) this else onEach { t -> defaultActions.forEach(t::builder) }
+    }
+
+    private val defaultActions = mutableListOf<MemberBuildAction>()
+
+    fun defaultBuilder(builderAction: MemberBuildAction) {
+        defaultActions += builderAction
+    }
+
+    companion object {
+        private val typeMemberTemplate = MemberTemplate::class.createType()
+        private val typeIterableOfMemberTemplates = Iterable::class.createType(arguments = listOf(KTypeProjection.invariant(typeMemberTemplate)))
+    }
+
+}
+
+fun foo() {
+    cacheFile.objectInputStream().use {
+        // one indent here
+    }
+}

idea/testData/formatter/KT22071.kt

@@ -0,0 +1,51 @@
+package templates
+
+import kotlin.coroutines.experimental.buildSequence
+import kotlin.reflect.KTypeProjection
+import kotlin.reflect.full.createType
+import kotlin.reflect.full.isSubtypeOf
+
+
+typealias TemplateGroup = () -> Sequence<MemberTemplate>
+
+fun templateGroupOf(vararg templates: MemberTemplate): TemplateGroup = { templates.asSequence() }
+
+abstract class TemplateGroupBase : TemplateGroup {
+
+    override fun invoke(): Sequence<MemberTemplate> = buildSequence {
+        with(this@TemplateGroupBase) {
+            this::class.members.filter { it.name.startsWith("f_") }.forEach {
+                    require(it.parameters.size == 1) { "Member $it violates naming convention" }
+                    when {
+                        it.returnType.isSubtypeOf(typeMemberTemplate) ->
+                            yield(it.call(this) as MemberTemplate)
+                        it.returnType.isSubtypeOf(typeIterableOfMemberTemplates) ->
+                            @Suppress("UNCHECKED_CAST")
+                            yieldAll(it.call(this) as Iterable<MemberTemplate>)
+                        else ->
+                            error("Member $it violates naming convention")
+                    }
+                }
+        }
+    }.run {
+        if (defaultActions.isEmpty()) this else onEach { t -> defaultActions.forEach(t::builder) }
+    }
+
+    private val defaultActions = mutableListOf<MemberBuildAction>()
+
+    fun defaultBuilder(builderAction: MemberBuildAction) {
+        defaultActions += builderAction
+    }
+
+    companion object {
+        private val typeMemberTemplate = MemberTemplate::class.createType()
+        private val typeIterableOfMemberTemplates = Iterable::class.createType(arguments = listOf(KTypeProjection.invariant(typeMemberTemplate)))
+    }
+
+}
+
+fun foo() {
+    cacheFile.objectInputStream().use {
+        // one indent here
+    }
+}

idea/testData/formatter/KT22115.after.kt

@@ -0,0 +1,9 @@
+private fun String.foo(): List<String> = emptyList()
+
+fun test(s: String) {
+    s.foo().flatMap { m ->
+        m.foo().map { e ->
+            e to m
+        }
+    }
+}

idea/testData/formatter/KT22115.kt

@@ -0,0 +1,9 @@
+private fun String.foo(): List<String> = emptyList()
+
+fun test(s: String) {
+    s.foo().flatMap { m ->
+        m.foo().map { e ->
+            e to m
+        }
+    }
+}

idea/testData/formatter/KT22148.after.kt

@@ -0,0 +1,14 @@
+fun foo() {
+    async {
+        // ...
+    }.logFailures().invokeOnCompletion {
+        if (it != null) showErrorAndContinue()
+        finish()
+    }
+
+    FirebaseAuth.getInstance().addAuthStateListener(object : FirebaseAuth.AuthStateListener {
+        override fun onAuthStateChanged(auth: FirebaseAuth) {
+            // ...
+        }
+    })
+}

idea/testData/formatter/KT22148.kt

@@ -0,0 +1,14 @@
+fun foo() {
+    async {
+        // ...
+    }.logFailures().invokeOnCompletion {
+        if (it != null) showErrorAndContinue()
+        finish()
+    }
+
+    FirebaseAuth.getInstance().addAuthStateListener(object : FirebaseAuth.AuthStateListener {
+        override fun onAuthStateChanged(auth: FirebaseAuth) {
+            // ...
+        }
+    })
+}

idea/tests/org/jetbrains/kotlin/formatter/FormatterTestGenerated.java

@@ -1,17 +1,6 @@
 /*
- * Copyright 2010-2017 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
  */
 
 package org.jetbrains.kotlin.formatter;
@@ -542,6 +531,24 @@ public class FormatterTestGenerated extends AbstractFormatterTest {
             doTest(fileName);
         }
 
+        @TestMetadata("KT22071.after.kt")
+        public void testKT22071() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/KT22071.after.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("KT22115.after.kt")
+        public void testKT22115() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/KT22115.after.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("KT22148.after.kt")
+        public void testKT22148() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/KT22148.after.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("KeepLineBreak.after.kt")
         public void testKeepLineBreak() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/KeepLineBreak.after.kt");