Repository: cqfn/diKTat (kotlin)

Issue:
For if-statement rule need to properly handle comments
Continuation of #801. In case we have comments in the if-block - they should be properly managed, 
For example:
```
if (statement1) {
    // comments 
    if (statement2) {
     }
}
```

should be converted to: 
```
if (statement1 &&
    // comments 
     statement2) {

```

Pull request:
### What's done:
* Change collapseThenBlocks function: KtBlockExpression.statements didn't hold comments, now it's fixed
* Monitor, which types of nodes are followed after nested if (more proper way), instead of check, that this if is the last child
* Fix indentions in tests

Patch:
diff --git a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/DiktatRuleSetProvider.kt b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/DiktatRuleSetProvider.kt
index dd1d3110cd..9322a7e9fe 100644
--- a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/DiktatRuleSetProvider.kt
+++ b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/DiktatRuleSetProvider.kt
@@ -190,6 +190,7 @@ class DiktatRuleSetProvider(private var diktatConfigFile: String = DIKTAT_ANALYS
             ::StringConcatenationRule,
             ::StringTemplateFormatRule,
             ::AccurateCalculationsRule,
+            ::CollapseIfStatementsRule,
             ::LineLength,
             ::RunInScript,
             ::TypeAliasRule,
@@ -207,7 +208,6 @@ class DiktatRuleSetProvider(private var diktatConfigFile: String = DIKTAT_ANALYS
             ::AvoidNestedFunctionsRule,
             ::ExtensionFunctionsSameNameRule,
             ::LambdaLengthRule,
-            ::CollapseIfStatementsRule,
             // formatting: moving blocks, adding line breaks, indentations etc.
             ::BlockStructureBraces,
             ::ConsecutiveSpacesRule,
diff --git a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/CollapseIfStatementsRule.kt b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/CollapseIfStatementsRule.kt
index 690b71e125..8ed3e28332 100644
--- a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/CollapseIfStatementsRule.kt
+++ b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter3/CollapseIfStatementsRule.kt
@@ -11,10 +11,11 @@ import com.pinterest.ktlint.core.ast.ElementType.BINARY_EXPRESSION
 import com.pinterest.ktlint.core.ast.ElementType.BLOCK_COMMENT
 import com.pinterest.ktlint.core.ast.ElementType.EOL_COMMENT
 import com.pinterest.ktlint.core.ast.ElementType.IF
-import com.pinterest.ktlint.core.ast.ElementType.KDOC
 import com.pinterest.ktlint.core.ast.ElementType.LBRACE
+import com.pinterest.ktlint.core.ast.ElementType.LPAR
 import com.pinterest.ktlint.core.ast.ElementType.OPERATION_REFERENCE
 import com.pinterest.ktlint.core.ast.ElementType.RBRACE
+import com.pinterest.ktlint.core.ast.ElementType.RPAR
 import com.pinterest.ktlint.core.ast.ElementType.THEN
 import com.pinterest.ktlint.core.ast.ElementType.WHITE_SPACE
 import com.pinterest.ktlint.core.ast.children
@@ -85,47 +86,63 @@ class CollapseIfStatementsRule(configRules: List<RulesConfig>) : DiktatRule(
     private fun findNestedIf(parentNode: ASTNode): ASTNode? {
         val parentThenNode = (parentNode.psi as KtIfExpression).then?.node ?: return null
         val nestedIfNode = parentThenNode.findChildByType(IF) ?: return null
-        // Nested `if` node should be the last child, but actually,
-        // the last children are WHITESPACE and `}`, so take treePrev
-        if ((nestedIfNode.psi as KtIfExpression).node != parentThenNode.lastChildNode.treePrev.treePrev) {
-            return null
-        }
-        // We won't collapse, if statements some of them have `else` node
+        // We won't collapse if-statements, if some of them have `else` node
         if ((parentNode.psi as KtIfExpression).`else` != null ||
                 (nestedIfNode.psi as KtIfExpression).`else` != null) {
             return null
         }
-        // We monitor which types of nodes are followed before nested `if`
+        // We monitor which types of nodes are followed before and after nested `if`
         // and we allow only a limited number of types to pass through.
         // Otherwise discovered `if` is not nested
+        // We don't expect KDOC in `if-statements`, since it's a bad practise, and such code by meaning of our
+        // code analyzer is invalid
+        // However, if in some case we will hit the KDOC, than we won't collapse statements
         val listOfNodesBeforeNestedIf = parentThenNode.getChildren(null).takeWhile { it.elementType != IF }
-        val allowedTypes = listOf(LBRACE, WHITE_SPACE, RBRACE, KDOC, BLOCK_COMMENT, EOL_COMMENT)
-        if (listOfNodesBeforeNestedIf.any { it.elementType !in allowedTypes }) {
+        val listOfNodesAfterNestedIf = parentThenNode.getChildren(null).takeLastWhile { it != parentThenNode.findChildByType(IF) }
+        val allowedTypes = listOf(LBRACE, WHITE_SPACE, RBRACE, BLOCK_COMMENT, EOL_COMMENT)
+        if (listOfNodesBeforeNestedIf.any { it.elementType !in allowedTypes } ||
+                listOfNodesAfterNestedIf.any { it.elementType !in allowedTypes }) {
             return null
         }
         return nestedIfNode
     }
 
+    private fun takeCommentsBeforeNestedIf(node: ASTNode): List<ASTNode> {
+        val thenNode = (node.psi as KtIfExpression).then?.node
+        return thenNode?.children()?.takeWhile { it.elementType != IF }?.filter {
+            it.elementType == EOL_COMMENT ||
+                    it.elementType == BLOCK_COMMENT
+        }
+            ?.toList() ?: emptyList()
+    }
+
     private fun collapse(parentNode: ASTNode, nestedNode: ASTNode) {
         collapseConditions(parentNode, nestedNode)
         collapseThenBlocks(parentNode, nestedNode)
     }
 
     private fun collapseConditions(parentNode: ASTNode, nestedNode: ASTNode) {
+        // If there are comments before nested if, we will move them into parent condition
+        val comments = takeCommentsBeforeNestedIf(parentNode)
+        val commentsText = if (comments.isNotEmpty()) {
+            comments.joinToString(prefix = "\n", postfix = "\n", separator = "\n") { it.text }
+        } else {
+            " "
+        }
         // Merge parent and nested conditions
-        val parentCondition = (parentNode.psi as KtIfExpression).condition?.text
+        val parentConditionText = extractConditions(parentNode)
         val nestedCondition = (nestedNode.psi as KtIfExpression).condition
+        val nestedConditionText = extractConditions(nestedNode)
         // If the nested condition is compound,
         // we need to put it to the brackets, according algebra of logic
         val mergeCondition =
                 if (nestedCondition?.node?.elementType == BINARY_EXPRESSION &&
                         nestedCondition.node?.findChildByType(OPERATION_REFERENCE)?.text == "||"
                 ) {
-                    "if ($parentCondition && (${nestedCondition.text})) {}"
+                    "if ($parentConditionText &&$commentsText($nestedConditionText)) {}"
                 } else {
-                    "if ($parentCondition && ${nestedCondition?.text}) {}"
+                    "if ($parentConditionText &&$commentsText$nestedConditionText) {}"
                 }
-
         val newParentIfNode = KotlinParser().createNode(mergeCondition)
         // Remove THEN block
         newParentIfNode.removeChild(newParentIfNode.lastChildNode)
@@ -139,13 +156,48 @@ class CollapseIfStatementsRule(configRules: List<RulesConfig>) : DiktatRule(
         }
     }
 
+    // If condition contains comments, we need additional actions
+    // Because of `node.condition` will ignore comments
+    private fun extractConditions(node: ASTNode): String {
+        val condition = node.getChildren(null)
+            .takeLastWhile { it != node.findChildByType(LPAR) }
+            .takeWhile { it != node.findChildByType(RPAR) }
+        return condition.joinToString("") { it.text }
+    }
+
     private fun collapseThenBlocks(parentNode: ASTNode, nestedNode: ASTNode) {
+        // Remove comments from parent node, since we already moved them into parent condition
+        val comments = takeCommentsBeforeNestedIf(parentNode)
+        comments.forEach {
+            if (it.treeNext.elementType == WHITE_SPACE &&
+                    it.treePrev.elementType == WHITE_SPACE) {
+                parentNode.removeChild(it.treePrev)
+            }
+            parentNode.removeChild(it)
+        }
         // Merge parent and nested `THEN` blocks
         val nestedThenNode = (nestedNode.psi as KtIfExpression).then
-        val nestedThenText = (nestedThenNode as KtBlockExpression).statements.joinToString("\n") { it.text }
-        val newNestedNode = KotlinParser().createNode(nestedThenText)
+        val nestedContent = (nestedThenNode as KtBlockExpression).children().toMutableList()
+        // Remove {, }, and white spaces
+        repeat(2) {
+            val firstElType = nestedContent.first().elementType
+            if (firstElType == WHITE_SPACE ||
+                    firstElType == LBRACE) {
+                nestedContent.removeFirst()
+            }
+            val lastElType = nestedContent.last().elementType
+            if (lastElType == WHITE_SPACE ||
+                    lastElType == RBRACE) {
+                nestedContent.removeLast()
+            }
+        }
+        val nestedThenText = nestedContent.joinToString("") { it.text }
+        val newNestedNode = KotlinParser().createNode(nestedThenText).treeParent
         val parentThenNode = (parentNode.psi as KtIfExpression).then?.node
-        parentThenNode?.replaceChild(nestedNode, newNestedNode)
+        newNestedNode.getChildren(null).forEach {
+            parentThenNode?.addChild(it, nestedNode)
+        }
+        parentThenNode?.removeChild(nestedNode)
     }
 
     /**


Test patch:
diff --git a/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/CollapseIfStatementsRuleWarnTest.kt b/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/CollapseIfStatementsRuleWarnTest.kt
index 2eda46028a..9ea831b2b5 100644
--- a/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/CollapseIfStatementsRuleWarnTest.kt
+++ b/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter3/CollapseIfStatementsRuleWarnTest.kt
@@ -77,25 +77,198 @@ class CollapseIfStatementsRuleWarnTest : LintTestBase(::CollapseIfStatementsRule
 
     @Test
     @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
-    fun `comments check`() {
+    fun `nested if with incorrect indention`() {
+        lintMethod(
+            """
+            |fun foo () {
+            |     if (true) {
+            |         if (true) {doSomething()}
+            |     }
+            |}
+            """.trimMargin(),
+            LintError(3, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `eol comment`() {
+        lintMethod(
+            """
+            |fun foo () {
+            |     if (true) {
+            |         // Some
+            |         // comments
+            |         if (true) {
+            |             doSomething()
+            |         }
+            |     }
+            |}
+            """.trimMargin(),
+            LintError(5, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `block comment`() {
+        lintMethod(
+            """
+            |fun foo () {
+            |     if (true) {
+            |         /*
+            |          Some comments
+            |         */
+            |         if (true) {
+            |             doSomething()
+            |         }
+            |     }
+            |}
+            """.trimMargin(),
+            LintError(6, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `kdoc comment`() {
         lintMethod(
             """
             |fun foo () {
             |     if (true) {
             |         /**
-            |          * Some Comments
-            |          */
+            |          * Some comments
+            |         */
+            |         if (true) {
+            |             doSomething()
+            |         }
+            |     }
+            |}
+            """.trimMargin()
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `combine comments`() {
+        lintMethod(
+            """
+            |fun foo () {
+            |     if (true) {
             |         /*
-            |          More comments
+            |          Some Comments
             |         */
-            |         // Even more comments
+            |         // More comments
             |         if (true) {
+            |             // comment 1
+            |             val a = 5
+            |             // comment 2
             |             doSomething()
             |         }
+            |         // comment 3
             |     }
             |} 
             """.trimMargin(),
-            LintError(10, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+            LintError(7, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `comments in compound cond`() {
+        lintMethod(
+            """
+            |fun foo() {
+            |     // comment
+            |     if (cond1) {
+            |         /*
+            |          Some comments
+            |         */
+            |         // More comments
+            |         if (cond2 || cond3) {
+            |             doSomething()
+            |         }
+            |     }
+            |}
+            """.trimMargin(),
+            LintError(8, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `comments already in cond`() {
+        lintMethod(
+            """
+            |fun foo () {
+            |     if (/*comment*/ true) {
+            |         if (true) {
+            |             doSomething()
+            |         }
+            |     }
+            |}
+            """.trimMargin(),
+            LintError(3, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `comments already in complex cond`() {
+        lintMethod(
+            """
+            |fun foo () {
+            |     if (true && (true || false) /*comment*/) {
+            |         if (true /*comment*/) {
+            |             doSomething()
+            |         }
+            |     }
+            |}
+            """.trimMargin(),
+            LintError(3, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `multiline comments already in cond`() {
+        lintMethod(
+            """
+            |fun foo () {
+            |     if (true
+            |     /*comment
+            |     * more comments
+            |     */
+            |     ) {
+            |         if (true /*comment 2*/) {
+            |             doSomething()
+            |         }
+            |     }
+            |}
+            """.trimMargin(),
+            LintError(7, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
+        )
+    }
+
+    @Test
+    @Tag(WarningNames.COLLAPSE_IF_STATEMENTS)
+    fun `comments in multiple if-statements`() {
+        lintMethod(
+            """
+            |fun foo() {
+            |     if (cond1) {
+            |         // comment
+            |         if (cond2) {
+            |             // comment 2
+            |             if (cond3) {
+            |                 doSomething()
+            |             }
+            |         }
+            |     }
+            |}
+            """.trimMargin(),
+            LintError(4, 10, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true),
+            LintError(6, 14, ruleId, "${COLLAPSE_IF_STATEMENTS.warnText()} avoid using redundant nested if-statements", true)
         )
     }
 
diff --git a/diktat-rules/src/test/resources/test/paragraph3/collapse_if/CollapseIfStatementsExpected.kt b/diktat-rules/src/test/resources/test/paragraph3/collapse_if/CollapseIfStatementsExpected.kt
index 851f7778f9..1ce1065ab1 100644
--- a/diktat-rules/src/test/resources/test/paragraph3/collapse_if/CollapseIfStatementsExpected.kt
+++ b/diktat-rules/src/test/resources/test/paragraph3/collapse_if/CollapseIfStatementsExpected.kt
@@ -89,7 +89,7 @@ fun foo() {
 fun foo() {
     if (cond1 && (cond2 && cond3 || cond4)) {
         firstAction()
-secondAction()
+            secondAction()
     }
 }
 
@@ -130,7 +130,7 @@ fun foo () {
 fun foo() {
     if (cond1 && cond2) {
         firstAction()
-secondAction()
+            secondAction()
     }
     if (cond3) {
         secondAction()
@@ -140,7 +140,7 @@ secondAction()
 fun foo() {
     if (cond1 && (cond2 || cond3)) {
         firstAction()
-secondAction()
+            secondAction()
     }
     if (cond4) {
         secondAction()
@@ -167,3 +167,89 @@ fun foo() {
         val a = 5
     }
 }
+
+fun foo() {
+    if (true &&
+/*
+         Some Comments
+        */
+// More comments
+true) {
+        // comment 1
+            val a = 5
+            // comment 2
+            doSomething()
+        // comment 3
+    }
+}
+
+fun foo() {
+    if (true &&
+// Some
+// comments
+true) {
+        doSomething()
+    }
+}
+
+fun foo() {
+    // comment
+    if (cond1 &&
+/*
+         Some comments
+        */
+// More comments
+(cond2 || cond3)) {
+        doSomething()
+    }
+}
+
+fun foo() {
+    if (cond1 &&
+// comment
+cond2 &&
+// comment 2
+cond3) {
+        doSomething()
+    }
+}
+
+fun foo () {
+    if (true && true) {
+        doSomething()
+    }
+}
+
+fun foo () {
+     if (/*comment*/ true && true) {
+         doSomething()
+     }
+}
+
+fun foo () {
+    if (true /*comment*/ && true) {
+        doSomething()
+    }
+}
+
+fun foo () {
+    if (true && true /*comment*/) {
+        doSomething()
+    }
+}
+
+fun foo () {
+    if (true && (true || false) /*comment*/ && true /*comment*/) {
+        doSomething()
+    }
+}
+
+fun foo () {
+     if (true
+     /*comment
+     * more comments
+     */
+      && true /*comment 2*/) {
+         doSomething()
+     }
+}
diff --git a/diktat-rules/src/test/resources/test/paragraph3/collapse_if/CollapseIfStatementsTest.kt b/diktat-rules/src/test/resources/test/paragraph3/collapse_if/CollapseIfStatementsTest.kt
index a72b9390ce..136e73ae2e 100644
--- a/diktat-rules/src/test/resources/test/paragraph3/collapse_if/CollapseIfStatementsTest.kt
+++ b/diktat-rules/src/test/resources/test/paragraph3/collapse_if/CollapseIfStatementsTest.kt
@@ -199,3 +199,104 @@ fun foo() {
         val a = 5
     }
 }
+
+fun foo() {
+    if (true) {
+        /*
+         Some Comments
+        */
+        // More comments
+        if (true) {
+            // comment 1
+            val a = 5
+            // comment 2
+            doSomething()
+        }
+        // comment 3
+    }
+}
+
+fun foo() {
+    if (true) {
+        // Some
+        // comments
+        if (true) {
+            doSomething()
+        }
+    }
+}
+
+fun foo() {
+    // comment
+    if (cond1) {
+        /*
+         Some comments
+        */
+        // More comments
+        if (cond2 || cond3) {
+            doSomething()
+        }
+    }
+}
+
+fun foo() {
+    if (cond1) {
+        // comment
+        if (cond2) {
+            // comment 2
+            if (cond3) {
+                doSomething()
+            }
+        }
+    }
+}
+
+fun foo () {
+    if (true) {
+        if (true) {doSomething()}
+    }
+}
+
+fun foo () {
+     if (/*comment*/ true) {
+         if (true) {
+             doSomething()
+         }
+     }
+}
+
+fun foo () {
+    if (true /*comment*/) {
+        if (true) {
+            doSomething()
+        }
+    }
+}
+
+fun foo () {
+    if (true) {
+        if (true /*comment*/) {
+            doSomething()
+        }
+    }
+}
+
+fun foo () {
+    if (true && (true || false) /*comment*/) {
+        if (true /*comment*/) {
+            doSomething()
+        }
+    }
+}
+
+fun foo () {
+     if (true
+     /*comment
+     * more comments
+     */
+     ) {
+         if (true /*comment 2*/) {
+             doSomething()
+         }
+     }
+}