Option to use continuation indent in expression bodies (KT-9818)

idea/formatter/src/org/jetbrains/kotlin/idea/core/formatter/KotlinCodeStyleSettings.java

@@ -41,6 +41,7 @@ public class KotlinCodeStyleSettings extends CustomCodeStyleSettings {
     public int NAME_COUNT_TO_USE_STAR_IMPORT_FOR_MEMBERS = ApplicationManager.getApplication().isUnitTestMode() ? Integer.MAX_VALUE : 3;
     public boolean IMPORT_NESTED_CLASSES = false;
     public boolean CONTINUATION_INDENT_IN_PARAMETER_LISTS = true;
+    public boolean CONTINUATION_INDENT_FOR_EXPRESSION_BODIES = true;
 
     public KotlinCodeStyleSettings(CodeStyleSettings container) {
         super("JetCodeStyleSettings", container);

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

@@ -30,7 +30,9 @@ import org.jetbrains.kotlin.kdoc.lexer.KDocTokens
 import org.jetbrains.kotlin.kdoc.parser.KDocElementTypes
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.lexer.KtTokens.*
+import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtExpression
 import java.util.*
 
 private val QUALIFIED_OPERATION = TokenSet.create(DOT, SAFE_ACCESS)
@@ -361,6 +363,18 @@ private val INDENT_RULES = arrayOf<NodeIndentStrategy>(
                 .within(KtNodeTypes.THEN, KtNodeTypes.ELSE).notForType(KtNodeTypes.BLOCK)
                 .set(Indent.getNormalIndent()),
 
+        strategy("Expression body")
+                .within(KtNodeTypes.FUN)
+                .forElement {
+                    it.psi is KtExpression && it.psi !is KtBlockExpression
+                }
+                .set { settings ->
+                    if (settings.kotlinSettings.CONTINUATION_INDENT_FOR_EXPRESSION_BODIES)
+                        Indent.getContinuationIndent()
+                    else
+                        Indent.getNormalIndent()
+                },
+
         strategy("Indent for parts")
                 .within(KtNodeTypes.PROPERTY, KtNodeTypes.FUN, KtNodeTypes.DESTRUCTURING_DECLARATION, KtNodeTypes.SECONDARY_CONSTRUCTOR)
                 .notForType(KtNodeTypes.BLOCK, FUN_KEYWORD, VAL_KEYWORD, VAR_KEYWORD, CONSTRUCTOR_KEYWORD)
@@ -406,7 +420,7 @@ private val INDENT_RULES = arrayOf<NodeIndentStrategy>(
                 .notForType(KtNodeTypes.BLOCK, KtNodeTypes.WHEN_CONDITION_EXPRESSION, KtNodeTypes.WHEN_CONDITION_IN_RANGE, KtNodeTypes.WHEN_CONDITION_IS_PATTERN, ELSE_KEYWORD, ARROW)
                 .set(Indent.getNormalIndent()),
 
-         strategy("Parameter list")
+        strategy("Parameter list")
                 .within(KtNodeTypes.VALUE_PARAMETER_LIST)
                 .forElement { it.elementType == KtNodeTypes.VALUE_PARAMETER && it.psi.prevSibling != null }
                 .set { settings ->
@@ -416,6 +430,7 @@ private val INDENT_RULES = arrayOf<NodeIndentStrategy>(
                         Indent.getNormalIndent()
                 })
 
+
 private fun getOperationType(node: ASTNode): IElementType? = node.findChildByType(KtNodeTypes.OPERATION_REFERENCE)?.firstChildNode?.elementType
 
 private fun getAlignmentForChildInParenthesis(

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

@@ -41,6 +41,7 @@ abstract class NodeIndentStrategy {
         private val notIn = ArrayList<IElementType>()
         private val forElement = ArrayList<IElementType>()
         private val notForElement = ArrayList<IElementType>()
+        private var forElementCallback: ((ASTNode) -> Boolean)? = null
 
         override fun toString(): String {
             return "PositionStrategy " + (debugInfo ?: "No debug info")
@@ -98,17 +99,25 @@ abstract class NodeIndentStrategy {
             return this
         }
 
+        fun forElement(callback: (ASTNode) -> Boolean): PositionStrategy {
+            forElementCallback = callback
+            return this
+        }
+
         override fun getIndent(node: ASTNode, settings: CodeStyleSettings): Indent? {
             if (!forElement.isEmpty()) {
                 if (!forElement.contains(node.elementType)) {
                     return null
                 }
             }
-
             if (notForElement.contains(node.elementType)) {
                 return null
             }
 
+            if (forElementCallback?.invoke(node) == false) {
+                return null
+            }
+
             val parent = node.treeParent
             if (parent != null) {
                 if (!within.isEmpty()) {