Add wrap option for expression body functions

#KT-21470 Fixed

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

@@ -46,6 +46,7 @@ public class KotlinCodeStyleSettings extends CustomCodeStyleSettings {
     public boolean CONTINUATION_INDENT_FOR_CHAINED_CALLS = true;
     public boolean CONTINUATION_INDENT_IN_SUPERTYPE_LISTS = true;
     public int BLANK_LINES_AROUND_BLOCK_WHEN_BRANCHES = 0;
+    public int WRAP_EXPRESSION_BODY_FUNCTIONS = 0;
 
     public KotlinCodeStyleSettings(CodeStyleSettings container) {
         super("JetCodeStyleSettings", container);

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

@@ -337,6 +337,7 @@ abstract class KotlinCommonBlock(
     private fun getWrappingStrategy(): WrappingStrategy {
         val commonSettings = settings.kotlinCommonSettings
         val elementType = node.elementType
+        val parentElementType = node.treeParent?.elementType
         val nodePsi = node.psi
 
         when {
@@ -344,7 +345,6 @@ abstract class KotlinCommonBlock(
                 return getWrappingStrategyForItemList(commonSettings.CALL_PARAMETERS_WRAP, KtNodeTypes.VALUE_ARGUMENT)
 
             elementType === KtNodeTypes.VALUE_PARAMETER_LIST -> {
-                val parentElementType = node.treeParent.elementType
                 if (parentElementType === KtNodeTypes.FUN ||
                     parentElementType === KtNodeTypes.PRIMARY_CONSTRUCTOR ||
                     parentElementType === KtNodeTypes.SECONDARY_CONSTRUCTOR) {
@@ -402,7 +402,17 @@ abstract class KotlinCommonBlock(
                 return wrapAfterAnnotation(commonSettings.CLASS_ANNOTATION_WRAP)
 
             nodePsi is KtNamedFunction ->
-                return wrapAfterAnnotation(commonSettings.METHOD_ANNOTATION_WRAP)
+                return object : WrappingStrategy {
+                    override fun getWrap(childElement: ASTNode): Wrap? {
+                        getWrapAfterAnnotation(childElement, commonSettings.METHOD_ANNOTATION_WRAP)?.let {
+                            return it
+                        }
+                        if (getPrevWithoutWhitespaceAndComments(childElement)?.elementType == KtTokens.EQ) {
+                            return Wrap.createWrap(settings.kotlinCustomSettings.WRAP_EXPRESSION_BODY_FUNCTIONS, true)
+                        }
+                        return null
+                    }
+                }
 
             nodePsi is KtProperty ->
                 return wrapAfterAnnotation(if (nodePsi.isLocal)
@@ -422,21 +432,25 @@ private fun ASTNode.isFirstParameter(): Boolean = treePrev?.elementType == KtTok
 private fun wrapAfterAnnotation(wrapType: Int): WrappingStrategy {
     return object : WrappingStrategy {
         override fun getWrap(childElement: ASTNode): Wrap? {
-            if (childElement.elementType in KtTokens.COMMENTS) return null
+            return getWrapAfterAnnotation(childElement, wrapType)
-            var prevLeaf = childElement.treePrev
-            while (prevLeaf?.elementType == TokenType.WHITE_SPACE) {
-                prevLeaf = prevLeaf.treePrev
-            }
-            if (prevLeaf?.elementType == KtNodeTypes.MODIFIER_LIST) {
-                if (prevLeaf?.lastChildNode?.elementType in ANNOTATIONS) {
-                    return Wrap.createWrap(wrapType, true)
-                }
-            }
-            return null
         }
     }
 }
 
+private fun getWrapAfterAnnotation(childElement: ASTNode, wrapType: Int): Wrap? {
+    if (childElement.elementType in COMMENTS) return null
+    var prevLeaf = childElement.treePrev
+    while (prevLeaf?.elementType == TokenType.WHITE_SPACE) {
+        prevLeaf = prevLeaf.treePrev
+    }
+    if (prevLeaf?.elementType == KtNodeTypes.MODIFIER_LIST) {
+        if (prevLeaf?.lastChildNode?.elementType in ANNOTATIONS) {
+            return Wrap.createWrap(wrapType, true)
+        }
+    }
+    return null
+}
+
 private val INDENT_RULES = arrayOf<NodeIndentStrategy>(
         strategy("No indent for braces in blocks")
                 .within(KtNodeTypes.BLOCK, KtNodeTypes.CLASS_BODY, KtNodeTypes.FUNCTION_LITERAL)

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

@@ -1,58 +0,0 @@
-/*
- * 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.
- */
-
-package org.jetbrains.kotlin.idea.formatter
-
-import com.intellij.application.options.IndentOptionsEditor
-import com.intellij.application.options.SmartIndentOptionsEditor
-import com.intellij.psi.codeStyle.CodeStyleSettings
-import com.intellij.psi.codeStyle.CommonCodeStyleSettings
-import org.jetbrains.kotlin.idea.core.formatter.KotlinCodeStyleSettings
-import javax.swing.JCheckBox
-
-class KotlinIndentOptionsEditor : SmartIndentOptionsEditor() {
-    private val useContinuationIndentForExpressionBodies = JCheckBox("Use continuation indent for expression bodies")
-
-    override fun addComponents() {
-        super.addComponents()
-        add(useContinuationIndentForExpressionBodies)
-    }
-
-    override fun isModified(settings: CodeStyleSettings, options: CommonCodeStyleSettings.IndentOptions): Boolean {
-        var isModified = super.isModified(settings, options)
-        val kotlinSettings = settings.kotlinCustomSettings
-        isModified = isModified || IndentOptionsEditor.isFieldModified(useContinuationIndentForExpressionBodies,
-                                                                       kotlinSettings.CONTINUATION_INDENT_FOR_EXPRESSION_BODIES)
-        return isModified
-    }
-
-    override fun apply(settings: CodeStyleSettings, options: CommonCodeStyleSettings.IndentOptions) {
-        super.apply(settings, options)
-        val kotlinSettings = settings.getCustomSettings(KotlinCodeStyleSettings::class.java)
-        kotlinSettings.CONTINUATION_INDENT_FOR_EXPRESSION_BODIES = useContinuationIndentForExpressionBodies.isSelected
-    }
-
-    override fun reset(settings: CodeStyleSettings, options: CommonCodeStyleSettings.IndentOptions) {
-        super.reset(settings, options)
-        val kotlinSettings = settings.getCustomSettings(KotlinCodeStyleSettings::class.java)
-        useContinuationIndentForExpressionBodies.isSelected = kotlinSettings.CONTINUATION_INDENT_FOR_EXPRESSION_BODIES
-    }
-
-    override fun setEnabled(enabled: Boolean) {
-        super.setEnabled(enabled)
-        useContinuationIndentForExpressionBodies.isEnabled = enabled
-    }
-}

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

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.formatter
 
 import com.intellij.application.options.IndentOptionsEditor
+import com.intellij.application.options.SmartIndentOptionsEditor
 import com.intellij.psi.codeStyle.CodeStyleSettingsCustomizable
 import com.intellij.psi.codeStyle.CommonCodeStyleSettings
 import com.intellij.psi.codeStyle.LanguageCodeStyleSettingsProvider
@@ -66,6 +67,8 @@ class KotlinLanguageCodeStyleSettingsProvider : LanguageCodeStyleSettingsProvide
                enum class Enumeration {
                    A, B
                }
+
+               fun veryLongExpressionBodyMethod() = "abc"
             """.trimIndent()
 
         LanguageCodeStyleSettingsProvider.SettingsType.BLANK_LINES_SETTINGS ->
@@ -167,8 +170,8 @@ class KotlinLanguageCodeStyleSettingsProvider : LanguageCodeStyleSettingsProvide
     override fun getLanguageName(): String = KotlinLanguage.NAME
 
     override fun customizeSettings(consumer: CodeStyleSettingsCustomizable, settingsType: LanguageCodeStyleSettingsProvider.SettingsType) {
-        fun showCustomOption(field: KProperty<*>, title: String, groupName: String? = null) {
+        fun showCustomOption(field: KProperty<*>, title: String, groupName: String? = null, vararg options: Any) {
-            consumer.showCustomOption(KotlinCodeStyleSettings::class.java, field.name, title, groupName)
+            consumer.showCustomOption(KotlinCodeStyleSettings::class.java, field.name, title, groupName, *options)
         }
 
         when (settingsType) {
@@ -285,6 +288,17 @@ class KotlinLanguageCodeStyleSettingsProvider : LanguageCodeStyleSettingsProvide
                         KotlinCodeStyleSettings::CONTINUATION_INDENT_IN_SUPERTYPE_LISTS,
                         "Use continuation indent",
                         CodeStyleSettingsCustomizable.WRAPPING_EXTENDS_LIST)
+
+                showCustomOption(
+                        KotlinCodeStyleSettings::WRAP_EXPRESSION_BODY_FUNCTIONS,
+                        "Expression body functions",
+                        options = *arrayOf(CodeStyleSettingsCustomizable.WRAP_OPTIONS_FOR_SINGLETON, CodeStyleSettingsCustomizable.WRAP_VALUES_FOR_SINGLETON)
+                )
+                showCustomOption(
+                        KotlinCodeStyleSettings::CONTINUATION_INDENT_FOR_EXPRESSION_BODIES,
+                        "Use continuation indent",
+                        "Expression body functions"
+                )
             }
             LanguageCodeStyleSettingsProvider.SettingsType.BLANK_LINES_SETTINGS -> {
                 consumer.showStandardOptions(
@@ -301,7 +315,7 @@ class KotlinLanguageCodeStyleSettingsProvider : LanguageCodeStyleSettingsProvide
         }
     }
 
-    override fun getIndentOptionsEditor(): IndentOptionsEditor = KotlinIndentOptionsEditor()
+    override fun getIndentOptionsEditor(): IndentOptionsEditor = SmartIndentOptionsEditor()
 
     override fun getDefaultCommonSettings(): CommonCodeStyleSettings =
         CommonCodeStyleSettings(language).apply {

idea/testData/formatter/ExpressionBodyWrap.after.kt

@@ -0,0 +1,4 @@
+fun foo() =
+        "abc"
+
+// SET_INT: WRAP_EXPRESSION_BODY_FUNCTIONS = 2

idea/testData/formatter/ExpressionBodyWrap.kt

@@ -0,0 +1,3 @@
+fun foo() = "abc"
+
+// SET_INT: WRAP_EXPRESSION_BODY_FUNCTIONS = 2

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

@@ -344,6 +344,12 @@ public class FormatterTestGenerated extends AbstractFormatterTest {
             doTest(fileName);
         }
 
+        @TestMetadata("ExpressionBodyWrap.after.kt")
+        public void testExpressionBodyWrap() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/ExpressionBodyWrap.after.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("ExtendsListWrap.after.kt")
         public void testExtendsListWrap() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/ExtendsListWrap.after.kt");