Feature. Auto-fixes for null-checks

diktat-rules/src/main/kotlin/org/cqfn/diktat/ruleset/rules/NullChecksRule.kt

@@ -9,14 +9,17 @@ import com.pinterest.ktlint.core.Rule
 import com.pinterest.ktlint.core.ast.ElementType
 import com.pinterest.ktlint.core.ast.ElementType.BINARY_EXPRESSION
 import com.pinterest.ktlint.core.ast.ElementType.CONDITION
+import com.pinterest.ktlint.core.ast.ElementType.ELSE
 import com.pinterest.ktlint.core.ast.ElementType.IF
 import com.pinterest.ktlint.core.ast.ElementType.IF_KEYWORD
 import com.pinterest.ktlint.core.ast.ElementType.NULL
 import com.pinterest.ktlint.core.ast.ElementType.REFERENCE_EXPRESSION
+import com.pinterest.ktlint.core.ast.ElementType.THEN
 import com.pinterest.ktlint.core.ast.ElementType.VALUE_ARGUMENT
 import com.pinterest.ktlint.core.ast.ElementType.VALUE_ARGUMENT_LIST
 import com.pinterest.ktlint.core.ast.parent
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode
+import org.jetbrains.kotlin.com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtIfExpression
 
@@ -72,21 +75,60 @@ class NullChecksRule(private val configRules: List<RulesConfig>) : Rule("null-ch
                     ElementType.EQEQ, ElementType.EQEQEQ ->
                         warnAndFixOnNullCheck(condition, true,
                             "use '.let/.also/?:/e.t.c' instead of ${condition.text}") {
-                            // todo implement fixer
+                            fixNullInIfCondition(node, condition, true)
                         }
                     // `!==` and `!==` comparison can be fixed with `.let/also` operators
                     ElementType.EXCLEQ, ElementType.EXCLEQEQEQ ->
                         warnAndFixOnNullCheck(condition, true,
                             "use '.let/.also/?:/e.t.c' instead of ${condition.text}") {
-                            // todo implement fixer
+                            fixNullInIfCondition(node, condition, false)
                         }
                     else -> return
                 }
             }
         }
     }
 
-    @Suppress("COMMENT_WHITE_SPACE")
+    @Suppress("UnsafeCallOnNullableType")
+    private fun fixNullInIfCondition(condition: ASTNode,
+                                     binaryExpression: KtBinaryExpression,
+                                     isEqualToNull: Boolean) {
+        val variableName = binaryExpression.left!!.text
+        val thenCodeLines = condition.extractLinesFromBlock(THEN)
+        val elseCodeLines = condition.extractLinesFromBlock(ELSE)
+        val text = if (isEqualToNull) {
+            if (elseCodeLines.isNullOrEmpty()) {
+                "$variableName ?: run {\n${thenCodeLines?.joinToString(separator = "\n")}\n}"
+            } else {
+                """
+                    |$variableName?.let {
+                    |${elseCodeLines.joinToString(separator = "\n")}
+                    |}
+                    |?: run {
+                    |${thenCodeLines?.joinToString(separator = "\n")}
+                    |}
+                """.trimMargin()
+            }
+        } else {
+            if (elseCodeLines.isNullOrEmpty()) {
+                "$variableName?.let {\n${thenCodeLines?.joinToString(separator = "\n")}\n}"
+            } else {
+                """
+                    |$variableName?.let {
+                    |${thenCodeLines?.joinToString(separator = "\n")}
+                    |}
+                    |?: run {
+                    |${elseCodeLines.joinToString(separator = "\n")}
+                    |}
+                """.trimMargin()
+            }
+        }
+        val tree = KotlinParser().createNode(text)
+        condition.treeParent.treeParent.addChild(tree, condition.treeParent)
+        condition.treeParent.treeParent.removeChild(condition.treeParent)
+    }
+
+    @Suppress("COMMENT_WHITE_SPACE", "UnsafeCallOnNullableType")
     private fun nullCheckInOtherStatements(binaryExprNode: ASTNode) {
         val condition = (binaryExprNode.psi as KtBinaryExpression)
         if (isNullCheckBinaryExpression(condition)) {
@@ -104,17 +146,31 @@ class NullChecksRule(private val configRules: List<RulesConfig>) : Rule("null-ch
                     if (listOf(ElementType.EXCLEQ, ElementType.EXCLEQEQEQ).contains(condition.operationToken)) {
                         warnAndFixOnNullCheck(
                             condition,
-                            false,
+                            true,
                             "use 'requireNotNull' instead of require(${condition.text})"
                         ) {
-                            // todo implement fixer
+                            val variableName = (binaryExprNode.psi as KtBinaryExpression).left!!.text
+                            val newMethod = KotlinParser().createNode("requireNotNull($variableName)")
+                            grandParent.treeParent.treeParent.addChild(newMethod, grandParent.treeParent)
+                            grandParent.treeParent.treeParent.removeChild(grandParent.treeParent)
                         }
                     }
                 }
             }
         }
     }
 
+    @Suppress("WRONG_INDENTATION")
+    private fun ASTNode.extractLinesFromBlock(type: IElementType): List<String>? =
+                treeParent
+                .getFirstChildWithType(type)
+                ?.text
+                ?.trim('{', '}')
+                ?.split("\n")
+                ?.filter { it.isNotBlank() }
+                ?.map { it.trim() }
+                ?.toList()
+
     @Suppress("UnsafeCallOnNullableType")
     private fun isNullCheckBinaryExpression(condition: KtBinaryExpression): Boolean =
             // check that binary expression has `null` as right or left operand

diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter4/NullChecksRuleFixTest.kt

@@ -0,0 +1,22 @@
+package org.cqfn.diktat.ruleset.chapter4
+
+import org.cqfn.diktat.ruleset.rules.NullChecksRule
+import org.cqfn.diktat.util.FixTestBase
+
+import generated.WarningNames
+import org.junit.jupiter.api.Tag
+import org.junit.jupiter.api.Test
+
+class NullChecksRuleFixTest : FixTestBase("test/paragraph4/null_checks", ::NullChecksRule) {
+    @Test
+    @Tag(WarningNames.AVOID_NULL_CHECKS)
+    fun `should fix if conditions`() {
+        fixAndCompare("IfConditionNullCheckExpected.kt", "IfConditionNullCheckTest.kt")
+    }
+
+    @Test
+    @Tag(WarningNames.AVOID_NULL_CHECKS)
+    fun `should fix require function`() {
+        fixAndCompare("RequireFunctionExpected.kt", "RequireFunctionTest.kt")
+    }
+}

diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/chapter4/NullChecksRuleWarnTest.kt

@@ -41,6 +41,9 @@ class NullChecksRuleWarnTest : LintTestBase(::NullChecksRule) {
                 |         println("null")
                 |         return
                 |     }
+                |     myVar ?: kotlin.run { 
+                |       println("null")
+                |     }
                 | }
                 """.trimMargin(),
             LintError(3, 11, ruleId, Warnings.AVOID_NULL_CHECKS.warnText() +
@@ -183,7 +186,7 @@ class NullChecksRuleWarnTest : LintTestBase(::NullChecksRule) {
                 | }
                 """.trimMargin(),
             LintError(2, 14, ruleId, Warnings.AVOID_NULL_CHECKS.warnText() +
-                    " use 'requireNotNull' instead of require(myVar != null)", false),
+                    " use 'requireNotNull' instead of require(myVar != null)", true),
         )
     }
 }

diktat-rules/src/test/kotlin/org/cqfn/diktat/ruleset/smoke/DiktatSmokeTest.kt

@@ -91,6 +91,12 @@ class DiktatSmokeTest : FixTestBase("test/smoke/src/main/kotlin",
         fixAndCompare("DefaultPackageExpected.kt", "DefaultPackageTest.kt")
     }
 
+    @Test
+    @Tag("DiktatRuleSetProvider")
+    fun `smoke test #7`() {
+        fixAndCompare("Example7Expected.kt", "Example7Test.kt")
+    }
+
     @Test
     @Tag("DiktatRuleSetProvider")
     fun `smoke test #6`() {

diktat-rules/src/test/resources/test/paragraph4/null_checks/IfConditionNullCheckExpected.kt

@@ -0,0 +1,33 @@
+package test.paragraph4.null_checks
+
+fun test() {
+    val some: Int? = null
+    some ?: run {
+println("some")
+bar()
+}
+
+    some?.let {
+println("some")
+bar()
+}
+
+    if (some == null && true) {
+        print("asd")
+    }
+
+    some?.let {
+print("qwe")
+}
+?: run {
+print("asd")
+}
+
+    some?.let {
+print("qqq")
+}
+?: run {
+print("www")
+}
+}
+

diktat-rules/src/test/resources/test/paragraph4/null_checks/IfConditionNullCheckTest.kt

@@ -0,0 +1,31 @@
+package test.paragraph4.null_checks
+
+fun test() {
+    val some: Int? = null
+    if (some == null) {
+        println("some")
+        bar()
+    }
+
+    if (some != null) {
+        println("some")
+        bar()
+    }
+
+    if (some == null && true) {
+        print("asd")
+    }
+
+    if (some == null) {
+        print("asd")
+    } else {
+        print("qwe")
+    }
+
+    if (some != null) {
+        print("qqq")
+    } else {
+        print("www")
+    }
+}
+

diktat-rules/src/test/resources/test/paragraph4/null_checks/RequireFunctionExpected.kt

@@ -0,0 +1,8 @@
+package test.paragraph4.null_checks
+
+fun test() {
+    val some: Int? = null
+
+    requireNotNull(some)
+}
+

diktat-rules/src/test/resources/test/paragraph4/null_checks/RequireFunctionTest.kt

@@ -0,0 +1,8 @@
+package test.paragraph4.null_checks
+
+fun test() {
+    val some: Int? = null
+
+    require(some != null)
+}
+

diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example7Expected.kt

@@ -0,0 +1,23 @@
+package org.cqfn.diktat
+
+fun foo() {
+    val prop: Int = 0
+
+    prop ?: run {
+        println("prop is null")
+        bar()
+    }
+
+    prop?.let {
+        baz()
+        gaz()
+    }
+
+    prop?.let {
+        doAnotherSmth()
+    }
+        ?: run {
+            doSmth()
+        }
+}
+

diktat-rules/src/test/resources/test/smoke/src/main/kotlin/Example7Test.kt

@@ -0,0 +1,22 @@
+package org.cqfn.diktat
+
+fun foo() {
+    val prop: Int? = null
+
+    if (prop == null) {
+        println("prop is null")
+        bar()
+    }
+
+    if (prop != null) {
+        baz()
+        gaz()
+    }
+
+    if (prop == null) {
+        doSmth()
+    } else {
+        doAnotherSmth()
+    }
+}
+