Bugfix for fix mode of AVOID_NULL_CHECK (#1535)

* Bugfix for fix mode of AVOID_NULL_CHECK ### What's done: * Added a check for assignment operator when deciding if to put run block or not * Added tests (#1293)
saveourtool · Oct 27, 2022 · 3beabb0 · 3beabb0
1 parent 5ade536
commit 3beabb0
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 14 deletions.
diff --git a/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter4/NullChecksRule.kt b/diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/chapter4/NullChecksRule.kt
@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtIfExpression
+import org.jetbrains.kotlin.psi.KtPsiUtil
 import org.jetbrains.kotlin.psi.psiUtil.blockExpressionsOrSingle
 
 /**
@@ -119,9 +120,10 @@ class NullChecksRule(configRules: List<RulesConfig>) : DiktatRule(
     }
 
     @Suppress("UnsafeCallOnNullableType", "TOO_LONG_FUNCTION")
-    private fun fixNullInIfCondition(condition: ASTNode,
-                                     binaryExpression: KtBinaryExpression,
-                                     isEqualToNull: Boolean
+    private fun fixNullInIfCondition(
+        condition: ASTNode,
+        binaryExpression: KtBinaryExpression,
+        isEqualToNull: Boolean
     ) {
         val variableName = binaryExpression.left!!.text
         val thenFromExistingCode = condition.extractLinesFromBlock(THEN)
@@ -138,26 +140,41 @@ class NullChecksRule(configRules: List<RulesConfig>) : DiktatRule(
         } else {
             elseFromExistingCode
         }
-        val numberOfStatementsInElseBlock = if (isEqualToNull) {
-            (condition.treeParent.psi as KtIfExpression).then?.blockExpressionsOrSingle()?.count() ?: 0
-        } else {
-            (condition.treeParent.psi as KtIfExpression).`else`?.blockExpressionsOrSingle()?.count() ?: 0
-        }
 
-        val elseEditedCodeLines = getEditedElseCodeLines(elseCodeLines, numberOfStatementsInElseBlock)
+        val (numberOfStatementsInElseBlock, isAssignmentInNewElseBlock) = (condition.treeParent.psi as KtIfExpression)
+            .let {
+                if (isEqualToNull) {
+                    it.then
+                } else {
+                    it.`else`
+                }
+            }
+            ?.blockExpressionsOrSingle()
+            ?.let { elements ->
+                elements.count() to elements.any { element ->
+                    KtPsiUtil.isAssignment(element)
+                }
+            }
+            ?: Pair(0, false)
+
+        val elseEditedCodeLines = getEditedElseCodeLines(elseCodeLines, numberOfStatementsInElseBlock, isAssignmentInNewElseBlock)
         val thenEditedCodeLines = getEditedThenCodeLines(variableName, thenCodeLines, elseEditedCodeLines)
 
         val text = "$thenEditedCodeLines $elseEditedCodeLines"
         val tree = KotlinParser().createNode(text)
-        condition.treeParent.treeParent.addChild(tree, condition.treeParent)
-        condition.treeParent.treeParent.removeChild(condition.treeParent)
+        val ifNode = condition.treeParent
+        ifNode.treeParent.replaceChild(ifNode, tree)
     }
 
-    private fun getEditedElseCodeLines(elseCodeLines: List<String>?, numberOfStatementsInElseBlock: Int): String = when {
+    private fun getEditedElseCodeLines(
+        elseCodeLines: List<String>?,
+        numberOfStatementsInElseBlock: Int,
+        isAssignment: Boolean,
+    ): String = when {
         // else { "null"/empty } -> ""
         elseCodeLines == null || elseCodeLines.singleOrNull() == "null" -> ""
         // else { bar() } -> ?: bar()
-        numberOfStatementsInElseBlock == 1 -> "?: ${elseCodeLines.joinToString(postfix = "\n", separator = "\n")}"
+        numberOfStatementsInElseBlock == 1 && !isAssignment -> "?: ${elseCodeLines.joinToString(postfix = "\n", separator = "\n")}"
         // else { ... } -> ?: run { ... }
         else -> getDefaultCaseElseCodeLines(elseCodeLines)
     }

diff --git a/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/smoke/DiktatSmokeTestBase.kt b/diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/smoke/DiktatSmokeTestBase.kt
@@ -311,7 +311,7 @@ abstract class DiktatSmokeTestBase : FixTestBase("test/smoke/src/main/kotlin",
 
     companion object {
         const val DEFAULT_CONFIG_PATH = "../diktat-analysis.yml"
-        private const val TEST_TIMEOUT_SECONDS = 20L
+        private const val TEST_TIMEOUT_SECONDS = 25L
         val unfixedLintErrors: MutableList<LintError> = mutableListOf()
 
         // by default using same yml config as for diktat code style check, but allow to override

diff --git a/diktat-rules/src/test/resources/test/paragraph4/null_checks/IfConditionNullCheckExpected.kt b/diktat-rules/src/test/resources/test/paragraph4/null_checks/IfConditionNullCheckExpected.kt
@@ -79,3 +79,25 @@ foo()
 } ?: boo()
 }
 
+fun nullCheckWithAssumption() {
+    val a: Int? = 5
+    a?.let {
+foo()
+} ?: run {
+a = 5
+}
+    a?.let {
+foo()
+} ?: run {
+a = 5
+}
+    a?.let {
+a = 5
+} ?: foo()
+    a?.let {
+a = 5
+} ?: foo()
+    a?.let {
+foo()
+} ?: a == 5
+}
diff --git a/diktat-rules/src/test/resources/test/paragraph4/null_checks/IfConditionNullCheckTest.kt b/diktat-rules/src/test/resources/test/paragraph4/null_checks/IfConditionNullCheckTest.kt
@@ -111,3 +111,31 @@ fun reversedCheckSmartCases() {
     }
 }
 
+fun nullCheckWithAssumption() {
+    val a: Int? = 5
+    if (a != null) {
+        foo()
+    } else {
+        a = 5
+    }
+    if (a == null) {
+        a = 5
+    } else {
+        foo()
+    }
+    if (a != null) {
+        a = 5
+    } else {
+        foo()
+    }
+    if (a == null) {
+        foo()
+    } else {
+        a = 5
+    }
+    if (a != null) {
+        foo()
+    } else {
+        a == 5
+    }
+}