Incorrect fix of null-safety rule

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

@@ -7,6 +7,8 @@ import org.cqfn.diktat.ruleset.utils.*
 
 import com.pinterest.ktlint.core.ast.ElementType
 import com.pinterest.ktlint.core.ast.ElementType.BINARY_EXPRESSION
+import com.pinterest.ktlint.core.ast.ElementType.BLOCK
+import com.pinterest.ktlint.core.ast.ElementType.BREAK
 import com.pinterest.ktlint.core.ast.ElementType.CONDITION
 import com.pinterest.ktlint.core.ast.ElementType.ELSE
 import com.pinterest.ktlint.core.ast.ElementType.IF
@@ -88,6 +90,17 @@ class NullChecksRule(configRules: List<RulesConfig>) : DiktatRule(
                                      binaryExpression: KtBinaryExpression,
                                      isEqualToNull: Boolean) {
         val variableName = binaryExpression.left!!.text
+        listOf(THEN, ELSE).forEach { elementType ->
+            // `break` word in block should be converted to `return`
+            condition.treeParent.findChildByType(elementType)?.let { blockNode ->
+                val correctChild = blockNode.findChildByType(BLOCK) ?: blockNode
+                correctChild.findChildByType(BREAK)?.let {
+                    val newNode = KotlinParser().createNode("return")
+                    correctChild.addChild(newNode, it)
+                    correctChild.removeChild(it)
+                }
+            }
+        }
         val thenCodeLines = condition.extractLinesFromBlock(THEN)
         val elseCodeLines = condition.extractLinesFromBlock(ELSE)
         val text = if (isEqualToNull) {

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

@@ -47,3 +47,26 @@ value
 }
 }
 
+fun foo() {
+    var result: Int? = 10
+    while (result != 0 ) {
+        result?.let {
+goo()
+}
+?: run {
+for(i in 1..10)
+break
+}
+    }
+    while (result != 0) {
+        result = goo()
+        result?.let {
+goo()
+}
+?: run {
+println(123)
+return
+}
+    }
+}
+

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

@@ -48,3 +48,24 @@ fun test() {
     }
 }
 
+fun foo() {
+    var result: Int? = 10
+    while (result != 0 ) {
+        if (result != null) {
+            goo()
+        } else {
+            for(i in 1..10)
+                break
+        }
+    }
+    while (result != 0) {
+        result = goo()
+        if (result != null) {
+            goo()
+        } else {
+            println(123)
+            break
+        }
+    }
+}
+