Fixed NO_BRACES_IN_CONDITIONALS_AND_LOOPS fail inside scope functio…

…ns (#1739)

Fixed `NO_BRACES_IN_CONDITIONALS_AND_LOOPS` fail inside scope functions

### What's done:
* Reworked indent handling before `if` using `findIndentBeforeNode()` function.
* Added tests for cases when conditionals and loops are inside scope functions.
* Added `with` to scope functions list.

Closes #1270

diktat-rules/src/main/kotlin/com/saveourtool/diktat/ruleset/rules/chapter3/BracesInConditionalsAndLoopsRule.kt

@@ -12,7 +12,9 @@ import com.saveourtool.diktat.ruleset.utils.prevSibling
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.KtNodeTypes.BLOCK
 import org.jetbrains.kotlin.KtNodeTypes.CALL_EXPRESSION
+import org.jetbrains.kotlin.KtNodeTypes.ELSE
 import org.jetbrains.kotlin.KtNodeTypes.IF
+import org.jetbrains.kotlin.KtNodeTypes.LAMBDA_EXPRESSION
 import org.jetbrains.kotlin.KtNodeTypes.REFERENCE_EXPRESSION
 import org.jetbrains.kotlin.KtNodeTypes.SAFE_ACCESS_EXPRESSION
 import org.jetbrains.kotlin.KtNodeTypes.WHEN
@@ -66,12 +68,7 @@ class BracesInConditionalsAndLoopsRule(configRules: List<RulesConfig>) : DiktatR
         val thenNode = ifPsi.then?.node
         val elseKeyword = ifPsi.elseKeyword
         val elseNode = ifPsi.`else`?.node
-        val indent = node
-            .prevSibling { it.elementType == WHITE_SPACE }
-            ?.text
-            ?.lines()
-            ?.last()
-            ?.count { it == ' ' } ?: 0
+        val indent = node.findIndentBeforeNode()
 
         if (node.isSingleLineIfElse()) {
             return
@@ -132,11 +129,8 @@ class BracesInConditionalsAndLoopsRule(configRules: List<RulesConfig>) : DiktatR
         if (loopBodyNode == null || loopBodyNode.elementType != BLOCK) {
             NO_BRACES_IN_CONDITIONALS_AND_LOOPS.warnAndFix(configRules, emitWarn, isFixMode, node.elementType.toString(), node.startOffset, node) {
                 // fixme proper way to calculate indent? or get step size (instead of hardcoded 4)
-                val indent = node.prevSibling { it.elementType == WHITE_SPACE }!!
-                    .text
-                    .lines()
-                    .last()
-                    .count { it == ' ' }
+                val indent = node.findIndentBeforeNode()
+
                 loopBody?.run {
                     replaceWithBlock(indent)
                 }
@@ -152,6 +146,23 @@ class BracesInConditionalsAndLoopsRule(configRules: List<RulesConfig>) : DiktatR
         }
     }
 
+    private fun ASTNode.findIndentBeforeNode(): Int {
+        val isElseIfStatement = treeParent.elementType == ELSE
+        val primaryIfNode = if (isElseIfStatement) treeParent.treeParent else this
+
+        val indentNode = if (primaryIfNode.treeParent?.treeParent?.treeParent?.elementType == LAMBDA_EXPRESSION) {
+            primaryIfNode.treeParent.prevSibling { it.elementType == WHITE_SPACE }
+        } else {
+            primaryIfNode.prevSibling { it.elementType == WHITE_SPACE }
+        }
+
+        return indentNode
+            ?.text
+            ?.lines()
+            ?.last()
+            ?.count { it == ' ' } ?: 0
+    }
+
     @Suppress("UnsafeCallOnNullableType")
     private fun checkWhenBranches(node: ASTNode) {
         (node.psi as KtWhenExpression)
@@ -194,6 +205,6 @@ class BracesInConditionalsAndLoopsRule(configRules: List<RulesConfig>) : DiktatR
     companion object {
         private const val INDENT_STEP = 4
         const val NAME_ID = "races-rule"
-        private val scopeFunctions = listOf("let", "run", "apply", "also")
+        private val scopeFunctions = listOf("let", "run", "with", "apply", "also")
     }
 }

diktat-rules/src/test/kotlin/com/saveourtool/diktat/ruleset/chapter3/BracesRuleFixTest.kt

@@ -22,6 +22,18 @@ class BracesRuleFixTest : FixTestBase("test/paragraph3/braces", ::BracesInCondit
         fixAndCompare("LoopsBracesExpected.kt", "LoopsBracesTest.kt")
     }
 
+    @Test
+    @Tag(WarningNames.NO_BRACES_IN_CONDITIONALS_AND_LOOPS)
+    fun `should add braces to if-else statements inside scope functions`() {
+        fixAndCompare("IfElseBracesInsideScopeFunctionsExpected.kt", "IfElseBracesInsideScopeFunctionsTest.kt")
+    }
+
+    @Test
+    @Tag(WarningNames.NO_BRACES_IN_CONDITIONALS_AND_LOOPS)
+    fun `should add braces to loops inside scope functions`() {
+        fixAndCompare("LoopsBracesInsideScopeFunctionsExpected.kt", "LoopsBracesInsideScopeFunctionsTest.kt")
+    }
+
     @Test
     @Tag(WarningNames.NO_BRACES_IN_CONDITIONALS_AND_LOOPS)
     @Disabled("https://github.com/saveourtool/diktat/issues/1737")

diktat-rules/src/test/resources/test/paragraph3/braces/IfElseBracesInsideScopeFunctionsExpected.kt

@@ -0,0 +1,63 @@
+package test.paragraph3.braces
+
+fun foo1() {
+    str.apply {
+        if (x > 0) {
+            foo()
+        } else {
+            bar()
+        }
+    }
+}
+
+fun foo2() {
+    str.let { if (x > 0) { foo() }
+    else {
+     bar()
+ }
+    }
+}
+
+fun foo3() {
+    str.run {
+        while (x > 0) {
+            if (x > 0) {
+                foo()
+            } else {
+                bar()
+            }
+        }
+    }
+}
+
+fun foo4() {
+    str.with { while (x > 0) {
+        if (x > 0) {
+            foo()
+        } else { bar() }
+    }
+    }
+}
+
+fun foo5() {
+    str.also {
+        while (x > 0) { if (x > 0) {
+     foo()
+ } else {
+     bar()
+ }
+        }
+    }
+}
+
+fun foo6() {
+    str.apply {
+        if (x > 0) {
+            foo()
+        } else if (y > 0) {
+            abc()
+        } else {
+            bar()
+        }
+    }
+}

diktat-rules/src/test/resources/test/paragraph3/braces/IfElseBracesInsideScopeFunctionsTest.kt

@@ -0,0 +1,47 @@
+package test.paragraph3.braces
+
+fun foo1() {
+    str.apply {
+        if (x > 0) foo()
+        else bar()
+    }
+}
+
+fun foo2() {
+    str.let { if (x > 0) { foo() }
+    else bar()
+    }
+}
+
+fun foo3() {
+    str.run {
+        while (x > 0) {
+            if (x > 0) foo()
+            else bar()
+        }
+    }
+}
+
+fun foo4() {
+    str.with { while (x > 0) {
+        if (x > 0) foo()
+        else { bar() }
+    }
+    }
+}
+
+fun foo5() {
+    str.also {
+        while (x > 0) { if (x > 0) foo()
+            else bar()
+        }
+    }
+}
+
+fun foo6() {
+    str.apply {
+        if (x > 0) foo()
+        else if (y > 0) abc()
+        else bar()
+    }
+}

diktat-rules/src/test/resources/test/paragraph3/braces/LoopsBracesInsideScopeFunctionsExpected.kt

@@ -0,0 +1,44 @@
+package test.paragraph3.braces
+
+fun foo1() {
+    str.apply {
+        for (i in 1..100) {
+            println(i)
+        }
+    }
+}
+
+fun foo2() {
+    str.let { while (x > 0) {
+     println(i)
+ }
+    }
+}
+
+fun foo3() {
+    str.run {
+        do {
+            println(i)
+        }
+        while (x > 0)
+    }
+}
+
+fun foo4() {
+    str.with { do {
+     println(i)
+ }
+    while (x > 0)
+    }
+}
+
+fun foo5() {
+    str.also {
+        for (i in 1..100) {
+            while (x > 0)
+                {
+                println(i)
+            }
+        }
+    }
+}

diktat-rules/src/test/resources/test/paragraph3/braces/LoopsBracesInsideScopeFunctionsTest.kt

@@ -0,0 +1,34 @@
+package test.paragraph3.braces
+
+fun foo1() {
+    str.apply {
+        for (i in 1..100) println(i)
+    }
+}
+
+fun foo2() {
+    str.let { while (x > 0) println(i)
+    }
+}
+
+fun foo3() {
+    str.run {
+        do println(i)
+        while (x > 0)
+    }
+}
+
+fun foo4() {
+    str.with { do println(i)
+    while (x > 0)
+    }
+}
+
+fun foo5() {
+    str.also {
+        for (i in 1..100) {
+            while (x > 0)
+                println(i)
+        }
+    }
+}