From acc15e5fad579837a9700c6775a3debb4f99b5d4 Mon Sep 17 00:00:00 2001
From: Dmitry Gridin <dmitry.gridin@jetbrains.com>
Date: Thu, 4 Jun 2020 13:26:21 +0700
Subject: [PATCH] LineIndentProvider: fix for do-while

Part of #KT-22211
---
 .../KotlinLikeLangLineIndentProvider.kt       | 57 +++++++++++++------
 ...ormanceTypingIndentationTestGenerated.java | 35 ++++++++++++
 .../indentationOnNewline/DoWhile.after.kt     |  4 ++
 idea/testData/indentationOnNewline/DoWhile.kt |  3 +
 .../indentationOnNewline/DoWhile2.after.kt    |  4 ++
 .../testData/indentationOnNewline/DoWhile2.kt |  3 +
 .../indentationOnNewline/DoWhile3.after.kt    |  6 ++
 .../testData/indentationOnNewline/DoWhile3.kt |  5 ++
 .../indentationOnNewline/DoWhile4.after.kt    |  4 ++
 .../testData/indentationOnNewline/DoWhile4.kt |  3 +
 .../indentationOnNewline/DoWhile5.after.kt    |  5 ++
 .../testData/indentationOnNewline/DoWhile5.kt |  4 ++
 .../indentationOnNewline/DoWhile6.after.kt    |  8 +++
 .../testData/indentationOnNewline/DoWhile6.kt |  7 +++
 .../indentationOnNewline/DoWhile7.after.kt    |  6 ++
 .../testData/indentationOnNewline/DoWhile7.kt |  5 ++
 .../TypingIndentationTestBaseGenerated.java   | 35 ++++++++++++
 17 files changed, 178 insertions(+), 16 deletions(-)
 create mode 100644 idea/testData/indentationOnNewline/DoWhile.after.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile2.after.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile2.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile3.after.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile3.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile4.after.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile4.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile5.after.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile5.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile6.after.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile6.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile7.after.kt
 create mode 100644 idea/testData/indentationOnNewline/DoWhile7.kt

diff --git a/idea/line-indent-provider/src/org/jetbrains/kotlin/idea/formatter/lineIndent/KotlinLikeLangLineIndentProvider.kt b/idea/line-indent-provider/src/org/jetbrains/kotlin/idea/formatter/lineIndent/KotlinLikeLangLineIndentProvider.kt
index dcffbb51cff..14825e26412 100644
--- a/idea/line-indent-provider/src/org/jetbrains/kotlin/idea/formatter/lineIndent/KotlinLikeLangLineIndentProvider.kt
+++ b/idea/line-indent-provider/src/org/jetbrains/kotlin/idea/formatter/lineIndent/KotlinLikeLangLineIndentProvider.kt
@@ -80,43 +80,68 @@ abstract class KotlinLikeLangLineIndentProvider : JavaLikeLangLineIndentProvider
                 Indent.getNoneIndent(),
                 IndentCalculator.LINE_BEFORE,
             )
+        }
 
-            before.isInControlFlowStatement() -> return factory.createIndentCalculator(
+        before.controlFlowStatementBefore()?.let { controlFlowKeywordPosition ->
+            return factory.createIndentCalculator(
                 when {
                     after.isAt(LeftParenthesis) -> Indent.getContinuationIndent()
-                    after.isAtAnyOf(BlockOpeningBrace, Arrow) -> Indent.getNoneIndent()
+                    after.isAtAnyOf(BlockOpeningBrace, Arrow) || controlFlowKeywordPosition.isWhileInsideDoWhile() -> Indent.getNoneIndent()
                     else -> Indent.getNormalIndent()
                 },
                 IndentCalculator.LINE_BEFORE,
             )
         }
 
-        after
-            .takeIf { it.isAt(TemplateEntryClose) }
-            ?.let { templateEntryPosition ->
-                val indent = if (currentPosition.hasEmptyLineAfter(offset)) Indent.getNormalIndent() else Indent.getNoneIndent()
-                templateEntryPosition.moveBeforeParentheses(TemplateEntryOpen, TemplateEntryClose)
-                val baseLineOffset = templateEntryPosition.startOffset
-                return factory.createIndentCalculator(indent) { baseLineOffset }
-            }
+        after.takeIf { it.isAt(TemplateEntryClose) }?.let { templateEntryPosition ->
+            val indent = if (currentPosition.hasEmptyLineAfter(offset)) Indent.getNormalIndent() else Indent.getNoneIndent()
+            templateEntryPosition.moveBeforeParentheses(TemplateEntryOpen, TemplateEntryClose)
+            val baseLineOffset = templateEntryPosition.startOffset
+            return factory.createIndentCalculator(indent) { baseLineOffset }
+        }
 
         return null
     }
 
-    private fun SemanticEditorPosition.isInControlFlowStatement(): Boolean = with(copy()) {
+    private fun SemanticEditorPosition.moveBeforeIfThisIsWhiteSpaceOrComment() = moveBeforeOptionalMix(*WHITE_SPACE_OR_COMMENT_BIT_SET)
+
+    private fun SemanticEditorPosition.isWhileInsideDoWhile(): Boolean {
+        if (!isAt(WhileKeyword)) return false
+        with(copy()) {
+            moveBefore()
+            var whileKeywordLevel = 1
+            while (!isAtEnd) when {
+                isAt(BlockOpeningBrace) -> return false
+                isAt(DoKeyword) -> {
+                    if (--whileKeywordLevel == 0) return true
+                    moveBefore()
+                }
+                isAt(WhileKeyword) -> {
+                    ++whileKeywordLevel
+                    moveBefore()
+                }
+                isAt(BlockClosingBrace) -> moveBeforeParentheses(BlockOpeningBrace, BlockClosingBrace)
+                else -> moveBefore()
+            }
+        }
+
+        return false
+    }
+
+    private fun SemanticEditorPosition.controlFlowStatementBefore(): SemanticEditorPosition? = with(copy()) {
         if (isAt(BlockOpeningBrace)) {
             moveBefore()
             moveBeforeParentheses(LeftParenthesis, RightParenthesis)
-            moveBeforeOptionalMix(*WHITE_SPACE_OR_COMMENT_BIT_SET)
+            moveBeforeIfThisIsWhiteSpaceOrComment()
         }
 
-        if (currElement in CONTROL_FLOW_CONSTRUCTIONS) return true
-        if (!isAt(RightParenthesis)) return false
+        if (currElement in CONTROL_FLOW_CONSTRUCTIONS) return this
+        if (!isAt(RightParenthesis)) return null
 
         moveBeforeParentheses(LeftParenthesis, RightParenthesis)
-        moveBeforeOptionalMix(*WHITE_SPACE_OR_COMMENT_BIT_SET)
+        moveBeforeIfThisIsWhiteSpaceOrComment()
 
-        return currElement in CONTROL_FLOW_CONSTRUCTIONS
+        return takeIf { currElement in CONTROL_FLOW_CONSTRUCTIONS }
     }
 
     enum class KotlinElement : SemanticEditorPosition.SyntaxElement {
diff --git a/idea/performanceTests/test/org/jetbrains/kotlin/idea/perf/PerformanceTypingIndentationTestGenerated.java b/idea/performanceTests/test/org/jetbrains/kotlin/idea/perf/PerformanceTypingIndentationTestGenerated.java
index 9bb082049b4..2c651a520b0 100644
--- a/idea/performanceTests/test/org/jetbrains/kotlin/idea/perf/PerformanceTypingIndentationTestGenerated.java
+++ b/idea/performanceTests/test/org/jetbrains/kotlin/idea/perf/PerformanceTypingIndentationTestGenerated.java
@@ -118,6 +118,41 @@ public class PerformanceTypingIndentationTestGenerated extends AbstractPerforman
         runTest("idea/testData/indentationOnNewline/DoInFun.kt");
     }
 
+    @TestMetadata("DoWhile.kt")
+    public void testDoWhile() throws Exception {
+        runTest("idea/testData/indentationOnNewline/DoWhile.kt");
+    }
+
+    @TestMetadata("DoWhile2.kt")
+    public void testDoWhile2() throws Exception {
+        runTest("idea/testData/indentationOnNewline/DoWhile2.kt");
+    }
+
+    @TestMetadata("DoWhile3.kt")
+    public void testDoWhile3() throws Exception {
+        runTest("idea/testData/indentationOnNewline/DoWhile3.kt");
+    }
+
+    @TestMetadata("DoWhile4.kt")
+    public void testDoWhile4() throws Exception {
+        runTest("idea/testData/indentationOnNewline/DoWhile4.kt");
+    }
+
+    @TestMetadata("DoWhile5.kt")
+    public void testDoWhile5() throws Exception {
+        runTest("idea/testData/indentationOnNewline/DoWhile5.kt");
+    }
+
+    @TestMetadata("DoWhile6.kt")
+    public void testDoWhile6() throws Exception {
+        runTest("idea/testData/indentationOnNewline/DoWhile6.kt");
+    }
+
+    @TestMetadata("DoWhile7.kt")
+    public void testDoWhile7() throws Exception {
+        runTest("idea/testData/indentationOnNewline/DoWhile7.kt");
+    }
+
     @TestMetadata("DoWithBraces.kt")
     public void testDoWithBraces() throws Exception {
         runTest("idea/testData/indentationOnNewline/DoWithBraces.kt");
diff --git a/idea/testData/indentationOnNewline/DoWhile.after.kt b/idea/testData/indentationOnNewline/DoWhile.after.kt
new file mode 100644
index 00000000000..2ee1b28dc27
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile.after.kt
@@ -0,0 +1,4 @@
+fun some() {
+    do println() while (true)
+    <caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile.kt b/idea/testData/indentationOnNewline/DoWhile.kt
new file mode 100644
index 00000000000..3aac6d32ac5
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile.kt
@@ -0,0 +1,3 @@
+fun some() {
+    do println() while (true)<caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile2.after.kt b/idea/testData/indentationOnNewline/DoWhile2.after.kt
new file mode 100644
index 00000000000..06204d29dbc
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile2.after.kt
@@ -0,0 +1,4 @@
+fun some() {
+    do do println() while (true) while (true)
+    <caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile2.kt b/idea/testData/indentationOnNewline/DoWhile2.kt
new file mode 100644
index 00000000000..4e00285bfdc
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile2.kt
@@ -0,0 +1,3 @@
+fun some() {
+    do do println() while (true) while (true)<caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile3.after.kt b/idea/testData/indentationOnNewline/DoWhile3.after.kt
new file mode 100644
index 00000000000..60e6824e331
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile3.after.kt
@@ -0,0 +1,6 @@
+fun some() {
+    do {
+        while (true) println()
+    } while (true)
+    <caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile3.kt b/idea/testData/indentationOnNewline/DoWhile3.kt
new file mode 100644
index 00000000000..08797f4ca0f
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile3.kt
@@ -0,0 +1,5 @@
+fun some() {
+    do {
+        while (true) println()
+    } while (true)<caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile4.after.kt b/idea/testData/indentationOnNewline/DoWhile4.after.kt
new file mode 100644
index 00000000000..70065b8444b
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile4.after.kt
@@ -0,0 +1,4 @@
+fun some() {
+    while (true)
+        <caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile4.kt b/idea/testData/indentationOnNewline/DoWhile4.kt
new file mode 100644
index 00000000000..80f3a4c9d2a
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile4.kt
@@ -0,0 +1,3 @@
+fun some() {
+    while (true)<caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile5.after.kt b/idea/testData/indentationOnNewline/DoWhile5.after.kt
new file mode 100644
index 00000000000..a983bd42f4b
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile5.after.kt
@@ -0,0 +1,5 @@
+fun some() {
+    while (true)
+        while (true)
+            <caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile5.kt b/idea/testData/indentationOnNewline/DoWhile5.kt
new file mode 100644
index 00000000000..7f729569020
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile5.kt
@@ -0,0 +1,4 @@
+fun some() {
+    while (true)
+        while (true)<caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile6.after.kt b/idea/testData/indentationOnNewline/DoWhile6.after.kt
new file mode 100644
index 00000000000..29b2a750d02
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile6.after.kt
@@ -0,0 +1,8 @@
+fun some() {
+    while (true) {
+        println()
+    }
+
+    while (true)
+        <caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile6.kt b/idea/testData/indentationOnNewline/DoWhile6.kt
new file mode 100644
index 00000000000..0a186df6070
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile6.kt
@@ -0,0 +1,7 @@
+fun some() {
+    while (true) {
+        println()
+    }
+
+    while (true)<caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile7.after.kt b/idea/testData/indentationOnNewline/DoWhile7.after.kt
new file mode 100644
index 00000000000..bc9aab32651
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile7.after.kt
@@ -0,0 +1,6 @@
+fun some() {
+    do println() while (true)
+    println()
+    while (true)
+        <caret>
+}
diff --git a/idea/testData/indentationOnNewline/DoWhile7.kt b/idea/testData/indentationOnNewline/DoWhile7.kt
new file mode 100644
index 00000000000..8c8ca9f0f7f
--- /dev/null
+++ b/idea/testData/indentationOnNewline/DoWhile7.kt
@@ -0,0 +1,5 @@
+fun some() {
+    do println() while (true)
+    println()
+    while (true)<caret>
+}
diff --git a/idea/tests/org/jetbrains/kotlin/formatter/TypingIndentationTestBaseGenerated.java b/idea/tests/org/jetbrains/kotlin/formatter/TypingIndentationTestBaseGenerated.java
index 652f0b4d47a..efbb74b30d7 100644
--- a/idea/tests/org/jetbrains/kotlin/formatter/TypingIndentationTestBaseGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/formatter/TypingIndentationTestBaseGenerated.java
@@ -120,6 +120,41 @@ public class TypingIndentationTestBaseGenerated extends AbstractTypingIndentatio
             runTest("idea/testData/indentationOnNewline/DoInFun.after.kt");
         }
 
+        @TestMetadata("DoWhile.after.kt")
+        public void testDoWhile() throws Exception {
+            runTest("idea/testData/indentationOnNewline/DoWhile.after.kt");
+        }
+
+        @TestMetadata("DoWhile2.after.kt")
+        public void testDoWhile2() throws Exception {
+            runTest("idea/testData/indentationOnNewline/DoWhile2.after.kt");
+        }
+
+        @TestMetadata("DoWhile3.after.kt")
+        public void testDoWhile3() throws Exception {
+            runTest("idea/testData/indentationOnNewline/DoWhile3.after.kt");
+        }
+
+        @TestMetadata("DoWhile4.after.kt")
+        public void testDoWhile4() throws Exception {
+            runTest("idea/testData/indentationOnNewline/DoWhile4.after.kt");
+        }
+
+        @TestMetadata("DoWhile5.after.kt")
+        public void testDoWhile5() throws Exception {
+            runTest("idea/testData/indentationOnNewline/DoWhile5.after.kt");
+        }
+
+        @TestMetadata("DoWhile6.after.kt")
+        public void testDoWhile6() throws Exception {
+            runTest("idea/testData/indentationOnNewline/DoWhile6.after.kt");
+        }
+
+        @TestMetadata("DoWhile7.after.kt")
+        public void testDoWhile7() throws Exception {
+            runTest("idea/testData/indentationOnNewline/DoWhile7.after.kt");
+        }
+
         @TestMetadata("DoWithBraces.after.kt")
         public void testDoWithBraces() throws Exception {
             runTest("idea/testData/indentationOnNewline/DoWithBraces.after.kt");