Fix indent of expressions following elvis operator

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

@@ -39,6 +39,8 @@ import org.jetbrains.kotlin.psi.psiUtil.siblings
 
 private val QUALIFIED_OPERATION = TokenSet.create(DOT, SAFE_ACCESS)
 private val QUALIFIED_EXPRESSIONS = TokenSet.create(KtNodeTypes.DOT_QUALIFIED_EXPRESSION, KtNodeTypes.SAFE_ACCESS_EXPRESSION)
+private val ELVIS_SET = TokenSet.create(KtTokens.ELVIS)
+
 private val KDOC_COMMENT_INDENT = 1
 
 private val BINARY_EXPRESSIONS = TokenSet.create(KtNodeTypes.BINARY_EXPRESSION, KtNodeTypes.BINARY_WITH_TYPE, KtNodeTypes.IS_EXPRESSION)
@@ -95,29 +97,12 @@ abstract class KotlinCommonBlock(
         var nodeSubBlocks = buildSubBlocks()
 
         if (node.elementType in QUALIFIED_EXPRESSIONS) {
-            val operationBlockIndex = findNodeBlockIndex(nodeSubBlocks, QUALIFIED_OPERATION)
+            nodeSubBlocks = splitSubBlocksOnDot(nodeSubBlocks)
-            if (operationBlockIndex != -1) {
+        }
-                // Create fake ".something" or "?.something" block here, so child indentation will be
+        else {
-                // relative to it when it starts from new line (see Indent javadoc).
+            val psi = node.psi
-
+            if (psi is KtBinaryExpression && psi.operationToken == KtTokens.ELVIS) {
-                val operationBlock = nodeSubBlocks[operationBlockIndex]
+                nodeSubBlocks = splitSubBlocksOnElvis(nodeSubBlocks)
-                val indent = if (settings.kotlinCustomSettings.CONTINUATION_INDENT_FOR_CHAINED_CALLS)
-                    Indent.getContinuationWithoutFirstIndent()
-                else
-                    Indent.getNormalIndent()
-                val isNonFirstChainedCall = operationBlockIndex > 0 && isCallBlock(nodeSubBlocks[operationBlockIndex - 1])
-                val wrap = if ((settings.kotlinCommonSettings.WRAP_FIRST_METHOD_IN_CALL_CHAIN || isNonFirstChainedCall) &&
-                               canWrapCallChain(node))
-                    Wrap.createWrap(settings.kotlinCommonSettings.METHOD_CALL_CHAIN_WRAP, true)
-                else
-                    null
-
-                val operationSyntheticBlock = SyntheticKotlinBlock(
-                        operationBlock.node,
-                        nodeSubBlocks.subList(operationBlockIndex, nodeSubBlocks.size),
-                        null, indent, wrap, spacingBuilder) { createSyntheticSpacingNodeBlock(it) }
-
-                nodeSubBlocks = nodeSubBlocks.subList(0, operationBlockIndex) + operationSyntheticBlock
             }
         }
 
@@ -126,6 +111,38 @@ abstract class KotlinCommonBlock(
         return nodeSubBlocks
     }
 
+    private fun splitSubBlocksOnDot(nodeSubBlocks: List<ASTBlock>): List<ASTBlock> {
+        val operationBlockIndex = nodeSubBlocks.indexOfBlockWithType(QUALIFIED_OPERATION)
+        if (operationBlockIndex != -1) {
+            // 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 indent = if (settings.kotlinCustomSettings.CONTINUATION_INDENT_FOR_CHAINED_CALLS)
+                Indent.getContinuationWithoutFirstIndent()
+            else
+                Indent.getNormalIndent()
+            val isNonFirstChainedCall = operationBlockIndex > 0 && isCallBlock(nodeSubBlocks[operationBlockIndex - 1])
+            val wrap = if ((settings.kotlinCommonSettings.WRAP_FIRST_METHOD_IN_CALL_CHAIN || isNonFirstChainedCall) &&
+                           canWrapCallChain(node))
+                Wrap.createWrap(settings.kotlinCommonSettings.METHOD_CALL_CHAIN_WRAP, true)
+            else
+                null
+
+            return nodeSubBlocks.splitAtIndex(operationBlockIndex, indent, wrap)
+        }
+        return nodeSubBlocks
+    }
+
+    private fun List<ASTBlock>.splitAtIndex(index: Int, indent: Indent?, wrap: Wrap?): List<ASTBlock> {
+        val operationBlock = this[index]
+        val operationSyntheticBlock = SyntheticKotlinBlock(
+                operationBlock.node,
+                subList(index, size),
+                null, indent, wrap, spacingBuilder) { createSyntheticSpacingNodeBlock(it) }
+
+        return subList(0, index) + operationSyntheticBlock
+    }
+
     private fun isCallBlock(astBlock: ASTBlock): Boolean {
         val node = astBlock.node
         return node.elementType in QUALIFIED_EXPRESSIONS && node.lastChildNode?.elementType == KtNodeTypes.CALL_EXPRESSION
@@ -140,6 +157,18 @@ abstract class KotlinCommonBlock(
                callChainParent.elementType == KtNodeTypes.RETURN
     }
 
+    private fun splitSubBlocksOnElvis(nodeSubBlocks: List<ASTBlock>): List<ASTBlock> {
+        val elvisIndex = nodeSubBlocks.indexOfBlockWithType(ELVIS_SET)
+        if (elvisIndex >= 0) {
+            return nodeSubBlocks.splitAtIndex(
+                    elvisIndex,
+                    Indent.getContinuationIndent(),
+                    null
+            )
+        }
+        return nodeSubBlocks
+    }
+
     fun createChildIndent(child: ASTNode): Indent? {
         val childParent = child.treeParent
         val childType = child.elementType
@@ -703,6 +732,6 @@ private fun getWrappingStrategyForItemList(wrapType: Int, itemTypes: TokenSet, w
     }
 }
 
-private fun findNodeBlockIndex(blocks: List<ASTBlock>, tokenSet: TokenSet): Int {
+private fun List<ASTBlock>.indexOfBlockWithType(tokenSet: TokenSet): Int {
-    return blocks.indexOfFirst { block -> block.node?.elementType in tokenSet }
+    return indexOfFirst { block -> block.node?.elementType in tokenSet }
 }

idea/testData/formatter/ElvisIndent.after.kt

@@ -0,0 +1,14 @@
+fun foo() {
+    val topLevelDeclaration =
+        DescriptorUtils.getParentOfType(referencedDescriptor, PropertyDescriptor::class.java, false) as DeclarationDescriptor?
+                ?: DescriptorUtils.getParentOfType(
+                    referencedDescriptor,
+                    TypeAliasConstructorDescriptor::class.java,
+                    false
+                )?.typeAliasDescriptor
+                ?: DescriptorUtils.getParentOfType(referencedDescriptor, FunctionDescriptor::class.java, false)
+                ?: return emptyList()
+}
+
+// SET_FALSE: CONTINUATION_INDENT_IN_ARGUMENT_LISTS
+// SET_FALSE: CONTINUATION_INDENT_FOR_EXPRESSION_BODIES

idea/testData/formatter/ElvisIndent.kt

@@ -0,0 +1,14 @@
+fun foo() {
+    val topLevelDeclaration =
+        DescriptorUtils.getParentOfType(referencedDescriptor, PropertyDescriptor::class.java, false) as DeclarationDescriptor?
+        ?: DescriptorUtils.getParentOfType(
+                referencedDescriptor,
+                TypeAliasConstructorDescriptor::class.java,
+                false
+        )?.typeAliasDescriptor
+        ?: DescriptorUtils.getParentOfType(referencedDescriptor, FunctionDescriptor::class.java, false)
+        ?: return emptyList()
+}
+
+// SET_FALSE: CONTINUATION_INDENT_IN_ARGUMENT_LISTS
+// SET_FALSE: CONTINUATION_INDENT_FOR_EXPRESSION_BODIES

idea/testData/formatter/FunctionLiteralsInChainCalls.after.inv.kt

@@ -33,6 +33,14 @@ fun test3() {
     }
 }
 
+fun test4() {
+    val abc = ArrayList<Int>().mapTo(
+            LinkedHashSet()
+    ) {
+        it * 2
+    }
+}
+
 fun testWithComments() {
     val abc = ArrayList<Int>()
 //  .map {

idea/testData/formatter/FunctionLiteralsInChainCalls.after.kt

@@ -33,6 +33,14 @@ fun test3() {
     }
 }
 
+fun test4() {
+    val abc = ArrayList<Int>().mapTo(
+            LinkedHashSet()
+    ) {
+        it * 2
+    }
+}
+
 fun testWithComments() {
     val abc = ArrayList<Int>()
 //  .map {

idea/testData/formatter/FunctionLiteralsInChainCalls.kt

@@ -33,6 +33,14 @@ val abc = ArrayList<Int>().map {
 }
 }
 
+fun test4() {
+    val abc = ArrayList<Int>().mapTo(
+        LinkedHashSet()
+    ) {
+        it * 2
+    }
+}
+
 fun testWithComments() {
     val abc = ArrayList<Int>()
 //  .map {

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

@@ -284,6 +284,12 @@ public class FormatterTestGenerated extends AbstractFormatterTest {
             doTest(fileName);
         }
 
+        @TestMetadata("ElvisIndent.after.kt")
+        public void testElvisIndent() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/ElvisIndent.after.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("ElvisWrap.after.kt")
         public void testElvisWrap() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/ElvisWrap.after.kt");