From ab99bf843a8cea19bdd7c136e38133f5a4cebb3f Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Fri, 8 Dec 2017 14:07:21 +0100
Subject: [PATCH] Implement option for wrapping assignment statements

 #KT-21718 Fixed
---
 .../idea/formatter/KotlinCommonBlock.kt       | 30 ++++++++++++++++---
 ...KotlinLanguageCodeStyleSettingsProvider.kt |  3 +-
 .../formatter/AssignmentWrap.after.kt         | 11 +++++++
 idea/testData/formatter/AssignmentWrap.kt     |  8 +++++
 .../formatter/FormatterTestGenerated.java     |  6 ++++
 5 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 idea/testData/formatter/AssignmentWrap.after.kt
 create mode 100644 idea/testData/formatter/AssignmentWrap.kt

diff --git a/idea/formatter/src/org/jetbrains/kotlin/idea/formatter/KotlinCommonBlock.kt b/idea/formatter/src/org/jetbrains/kotlin/idea/formatter/KotlinCommonBlock.kt
index ac86d69dc6d..8e52e98483b 100644
--- a/idea/formatter/src/org/jetbrains/kotlin/idea/formatter/KotlinCommonBlock.kt
+++ b/idea/formatter/src/org/jetbrains/kotlin/idea/formatter/KotlinCommonBlock.kt
@@ -415,10 +415,32 @@ abstract class KotlinCommonBlock(
                 }
 
             nodePsi is KtProperty ->
-                return wrapAfterAnnotation(if (nodePsi.isLocal)
-                                               commonSettings.VARIABLE_ANNOTATION_WRAP
-                                           else
-                                               commonSettings.FIELD_ANNOTATION_WRAP)
+                return object : WrappingStrategy {
+                    override fun getWrap(childElement: ASTNode): Wrap? {
+                        val wrapSetting = if (nodePsi.isLocal) commonSettings.VARIABLE_ANNOTATION_WRAP else commonSettings.FIELD_ANNOTATION_WRAP
+                        getWrapAfterAnnotation(childElement, wrapSetting)?.let {
+                            return it
+                        }
+                        if (getPrevWithoutWhitespaceAndComments(childElement)?.elementType == KtTokens.EQ) {
+                            return Wrap.createWrap(settings.kotlinCommonSettings.ASSIGNMENT_WRAP, true)
+                        }
+                        return null
+                    }
+                }
+
+            nodePsi is KtBinaryExpression -> {
+                if (nodePsi.operationToken == KtTokens.EQ) {
+                    return object : WrappingStrategy {
+                        override fun getWrap(childElement: ASTNode): Wrap? {
+                            if (getPrevWithoutWhitespaceAndComments(childElement)?.elementType == KtNodeTypes.OPERATION_REFERENCE) {
+                                return Wrap.createWrap(settings.kotlinCommonSettings.ASSIGNMENT_WRAP, true)
+                            }
+                            return null
+                        }
+                    }
+                }
+                return WrappingStrategy.NoWrapping
+            }
         }
 
         return WrappingStrategy.NoWrapping
diff --git a/idea/src/org/jetbrains/kotlin/idea/formatter/KotlinLanguageCodeStyleSettingsProvider.kt b/idea/src/org/jetbrains/kotlin/idea/formatter/KotlinLanguageCodeStyleSettingsProvider.kt
index 3149f99bf6b..9be427230aa 100644
--- a/idea/src/org/jetbrains/kotlin/idea/formatter/KotlinLanguageCodeStyleSettingsProvider.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/formatter/KotlinLanguageCodeStyleSettingsProvider.kt
@@ -260,7 +260,8 @@ class KotlinLanguageCodeStyleSettingsProvider : LanguageCodeStyleSettingsProvide
                         "CALL_PARAMETERS_RPAREN_ON_NEXT_LINE",
                         "ENUM_CONSTANTS_WRAP",
                         "METHOD_CALL_CHAIN_WRAP",
-                        "WRAP_FIRST_METHOD_IN_CALL_CHAIN"
+                        "WRAP_FIRST_METHOD_IN_CALL_CHAIN",
+                        "ASSIGNMENT_WRAP"
                 )
                 consumer.renameStandardOption(CodeStyleSettingsCustomizable.WRAPPING_SWITCH_STATEMENT, "'when' statements")
                 consumer.renameStandardOption("FIELD_ANNOTATION_WRAP", "Property annotations")
diff --git a/idea/testData/formatter/AssignmentWrap.after.kt b/idea/testData/formatter/AssignmentWrap.after.kt
new file mode 100644
index 00000000000..3d80d0534b2
--- /dev/null
+++ b/idea/testData/formatter/AssignmentWrap.after.kt
@@ -0,0 +1,11 @@
+val x =
+        1
+
+fun foo() {
+    val y =
+            2
+    y =
+            3
+}
+
+// SET_INT: ASSIGNMENT_WRAP = 2
diff --git a/idea/testData/formatter/AssignmentWrap.kt b/idea/testData/formatter/AssignmentWrap.kt
new file mode 100644
index 00000000000..4762bef2331
--- /dev/null
+++ b/idea/testData/formatter/AssignmentWrap.kt
@@ -0,0 +1,8 @@
+val x = 1
+
+fun foo() {
+    val y = 2
+    y = 3
+}
+
+// SET_INT: ASSIGNMENT_WRAP = 2
diff --git a/idea/tests/org/jetbrains/kotlin/formatter/FormatterTestGenerated.java b/idea/tests/org/jetbrains/kotlin/formatter/FormatterTestGenerated.java
index dcfe9ac089a..997df135094 100644
--- a/idea/tests/org/jetbrains/kotlin/formatter/FormatterTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/formatter/FormatterTestGenerated.java
@@ -74,6 +74,12 @@ public class FormatterTestGenerated extends AbstractFormatterTest {
             doTest(fileName);
         }
 
+        @TestMetadata("AssignmentWrap.after.kt")
+        public void testAssignmentWrap() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/AssignmentWrap.after.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("BinaryExpressionAlignmentSpread.after.kt")
         public void testBinaryExpressionAlignmentSpread() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/formatter/BinaryExpressionAlignmentSpread.after.kt");