Add IDEA data-flow analysis to guess nullability

Add "if return..." folding to "return if"

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/IfThenUtils.kt

@@ -63,6 +63,11 @@ fun KtExpression.unwrapBlockOrParenthesis(): KtExpression {
     return innerExpression
 }
 
+fun KtExpression?.isTrivialStatementBody(): Boolean = when (this?.unwrapBlockOrParenthesis()) {
+    is KtIfExpression, is KtBlockExpression -> false
+    else -> true
+}
+
 fun KtExpression?.isNullExpression(): Boolean = this?.unwrapBlockOrParenthesis()?.node?.elementType == KtNodeTypes.NULL
 
 fun KtExpression?.isNullExpressionOrEmptyBlock(): Boolean = this.isNullExpression() || this is KtBlockExpression && this.statements.isEmpty()

idea/src/org/jetbrains/kotlin/idea/j2k/IdeaJavaToKotlinServices.kt

@@ -16,8 +16,13 @@
 
 package org.jetbrains.kotlin.idea.j2k
 
+import com.intellij.codeInspection.dataFlow.DfaUtil
+import com.intellij.codeInspection.dataFlow.Nullness
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiVariable
 import org.jetbrains.kotlin.j2k.*
+import org.jetbrains.kotlin.j2k.ast.Nullability
 
 object IdeaJavaToKotlinServices : JavaToKotlinConverterServices {
     override val referenceSearcher: ReferenceSearcher
@@ -31,8 +36,25 @@ object IdeaJavaToKotlinServices : JavaToKotlinConverterServices {
 
     override val docCommentConverter: DocCommentConverter
         get() = IdeaDocCommentConverter
+
+    override val javaDataFlowAnalyzerFacade: JavaDataFlowAnalyzerFacade
+        get() = IdeaJavaDataFlowAnalyzerFacade
 }
 
 object IdeaSuperMethodSearcher : SuperMethodsSearcher {
     override fun findDeepestSuperMethods(method: PsiMethod) = method.findDeepestSuperMethods().asList()
+}
+
+private object IdeaJavaDataFlowAnalyzerFacade : JavaDataFlowAnalyzerFacade {
+    override fun variableNullability(variable: PsiVariable, context: PsiElement): Nullability =
+            DfaUtil.checkNullness(variable, context).toNullability()
+
+    override fun methodNullability(method: PsiMethod): Nullability =
+            DfaUtil.inferMethodNullity(method).toNullability()
+
+    private fun Nullness.toNullability() = when (this) {
+        Nullness.UNKNOWN -> Nullability.Default
+        Nullness.NOT_NULL -> Nullability.NotNull
+        Nullness.NULLABLE -> Nullability.Nullable
+    }
 }

idea/src/org/jetbrains/kotlin/idea/j2k/J2kPostProcessings.kt

@@ -27,8 +27,11 @@ import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.core.setVisibility
 import org.jetbrains.kotlin.idea.inspections.RedundantSamConstructorInspection
 import org.jetbrains.kotlin.idea.intentions.*
+import org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions.FoldIfToReturnAsymmetricallyIntention
+import org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions.FoldIfToReturnIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions.IfThenToElvisIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions.IfThenToSafeAccessIntention
+import org.jetbrains.kotlin.idea.intentions.branchedTransformations.isTrivialStatementBody
 import org.jetbrains.kotlin.idea.intentions.conventionNameCalls.ReplaceGetOrSetInspection
 import org.jetbrains.kotlin.idea.intentions.conventionNameCalls.ReplaceGetOrSetIntention
 import org.jetbrains.kotlin.idea.quickfix.RemoveModifierFix
@@ -76,9 +79,14 @@ object J2KPostProcessingRegistrar {
         _processings.add(RemoveRedundantCastToNullableProcessing())
 
         registerIntentionBasedProcessing(ConvertToExpressionBodyIntention(convertEmptyToUnit = false)) { it is KtPropertyAccessor }
+
+        registerIntentionBasedProcessing(FoldInitializerAndIfToElvisIntention())
+
+        registerIntentionBasedProcessing(FoldIfToReturnIntention()) { it.then.isTrivialStatementBody() && it.`else`.isTrivialStatementBody() }
+        registerIntentionBasedProcessing(FoldIfToReturnAsymmetricallyIntention()) { it.then.isTrivialStatementBody() && (KtPsiUtil.skipTrailingWhitespacesAndComments(it) as KtReturnExpression).returnedExpression.isTrivialStatementBody() }
+
         registerIntentionBasedProcessing(IfThenToSafeAccessIntention())
         registerIntentionBasedProcessing(IfThenToElvisIntention())
-        registerIntentionBasedProcessing(FoldInitializerAndIfToElvisIntention())
         registerIntentionBasedProcessing(SimplifyNegatedBinaryExpressionIntention())
         registerIntentionBasedProcessing(ReplaceGetOrSetIntention(), additionalChecker = ReplaceGetOrSetInspection.additionalChecker)
         registerIntentionBasedProcessing(AddOperatorModifierIntention())

j2k/src/org/jetbrains/kotlin/j2k/JavaToKotlinConverterServices.kt

@@ -23,6 +23,7 @@ interface JavaToKotlinConverterServices {
     val superMethodsSearcher: SuperMethodsSearcher
     val resolverForConverter: ResolverForConverter
     val docCommentConverter: DocCommentConverter
+    val javaDataFlowAnalyzerFacade: JavaDataFlowAnalyzerFacade
 }
 
 object EmptyJavaToKotlinServices: JavaToKotlinConverterServices {
@@ -37,6 +38,9 @@ object EmptyJavaToKotlinServices: JavaToKotlinConverterServices {
 
     override val docCommentConverter: DocCommentConverter
         get() = EmptyDocCommentConverter
+
+    override val javaDataFlowAnalyzerFacade: JavaDataFlowAnalyzerFacade
+        get() = JavaDataFlowAnalyzerFacade.Default
 }
 
 interface SuperMethodsSearcher {

j2k/src/org/jetbrains/kotlin/j2k/TypeConverter.kt

@@ -27,6 +27,19 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.types.TypeUtils
 import java.util.*
 
+interface JavaDataFlowAnalyzerFacade {
+
+    fun variableNullability(variable: PsiVariable, context: PsiElement): Nullability
+
+    fun methodNullability(method: PsiMethod): Nullability
+
+    object Default : JavaDataFlowAnalyzerFacade {
+        override fun methodNullability(method: PsiMethod): Nullability = Nullability.Default
+
+        override fun variableNullability(variable: PsiVariable, context: PsiElement): Nullability = Nullability.Default
+    }
+}
+
 class TypeConverter(val converter: Converter) {
     private val typesBeingConverted = HashSet<PsiType>()
 
@@ -74,6 +87,9 @@ class TypeConverter(val converter: Converter) {
     fun variableNullability(variable: PsiVariable): Nullability
             = nullabilityFlavor.forVariableType(variable, true)
 
+    fun variableReferenceNullability(variable: PsiVariable, reference: PsiReferenceExpression): Nullability
+            = nullabilityFlavor.forVariableReference(variable, reference)
+
     fun methodNullability(method: PsiMethod): Nullability
             = nullabilityFlavor.forMethodReturnType(method)
 
@@ -181,6 +197,8 @@ class TypeConverter(val converter: Converter) {
             return value
         }
 
+        abstract fun fromDataFlowForMethod(method: PsiMethod): T
+
         private fun forMethodReturnTypeNoCache(method: PsiMethod): T {
             val returnType = method.returnType ?: return default
 
@@ -197,6 +215,9 @@ class TypeConverter(val converter: Converter) {
             value = fromTypeHeuristics(returnType)
             if (value != default) return value
 
+            value = fromDataFlowForMethod(method)
+            if (value != default) return value
+
             if (!converter.inConversionScope(method)) return default // do not analyze body and usages of methods out of our conversion scope
 
             val body = method.body
@@ -227,6 +248,17 @@ class TypeConverter(val converter: Converter) {
     }
 
     private val nullabilityFlavor = object : TypeFlavor<Nullability>(Nullability.Default) {
+        fun forVariableReference(variable: PsiVariable, reference: PsiReferenceExpression): Nullability {
+            assert(reference.resolve() == variable)
+            val dataFlowUtil = converter.services.javaDataFlowAnalyzerFacade
+
+            return dataFlowUtil.variableNullability(variable, reference).takeIf { it != default } ?:
+                   variableNullability(variable)
+        }
+
+        override fun fromDataFlowForMethod(method: PsiMethod): Nullability =
+                converter.services.javaDataFlowAnalyzerFacade.methodNullability(method)
+
         override val forEnumConstant: Nullability
             get() = Nullability.NotNull
 
@@ -361,7 +393,9 @@ class TypeConverter(val converter: Converter) {
 
                 is PsiParenthesizedExpression -> expression?.nullability() ?: Nullability.Default
 
+                is PsiCallExpression -> resolveMethod()?.let { methodNullability(it) } ?: Nullability.Default
 
+                is PsiReferenceExpression -> (resolve() as? PsiVariable)?.let { variableReferenceNullability(it, this) } ?: Nullability.Default
             //TODO: some other cases
 
                 else -> Nullability.Default
@@ -370,6 +404,8 @@ class TypeConverter(val converter: Converter) {
     }
 
     private val mutabilityFlavor = object : TypeFlavor<Mutability>(Mutability.Default) {
+        override fun fromDataFlowForMethod(method: PsiMethod): Mutability = Mutability.Default
+
         override val forEnumConstant: Mutability get() = Mutability.NonMutable
 
         override fun fromType(type: PsiType): Mutability {

j2k/testData/fileOrElement/callChainExpression/libraryMethodCallFromInstance-settings.kt

@@ -3,7 +3,7 @@
 internal class Library {
     fun call() {}
 
-    val string: String?
+    val string: String
         get() = ""
 }
 
@@ -11,9 +11,9 @@ internal class User {
     fun main() {
         val lib: Library = Library()
         lib.call()
-        lib.string!!.isEmpty()
+        lib.string.isEmpty()
 
         Library().call()
-        Library().string!!.isEmpty()
+        Library().string.isEmpty()
     }
 }

j2k/testData/fileOrElement/detectProperties/GetterTypeNotMatch.kt

@@ -1,7 +1,6 @@
-// ERROR: Type mismatch: inferred type is String? but Any was expected
 internal class A {
     private val s: String? = null
 
-    val value: Any
+    val value: Any?
         get() = s
 }

j2k/testData/fileOrElement/ifStatement/singleLine.kt

@@ -1 +1 @@
-if (true) return 1 else return 0
+return if (true) 1 else 0

j2k/testData/fileOrElement/issues/kt-824-isDir.kt

@@ -11,9 +11,9 @@ object Test {
             return false
         }
         val result = true
-        if (parent.isDirectory) {
-            return true
+        return if (parent.isDirectory) {
+            true
         } else
-            return false
+            false
     }
 }

j2k/testData/fileOrElement/issues/kt-852.kt

@@ -8,9 +8,7 @@ internal class Test {
         if (s1.isEmpty() && s2.isEmpty())
             return "OK"
 
-        if (s1.isEmpty() && s2.isEmpty() && s3.isEmpty())
-            return "OOOK"
+        return if (s1.isEmpty() && s2.isEmpty() && s3.isEmpty()) "OOOK" else ""
 
-        return ""
     }
 }

j2k/testData/fileOrElement/nullability/MethodReturnsNull.kt

@@ -1,9 +1,9 @@
 internal class C {
     fun foo(b: Boolean): String? {
-        if (b) {
-            return "abc"
+        return if (b) {
+            "abc"
         } else {
-            return null
+            null
         }
     }
 }

j2k/testData/fileOrElement/nullability/NullableVariableDotAccess.kt

@@ -1,5 +1,4 @@
 fun foo(s: String?, b: Boolean): Int {
     if (s == null) println("null")
-    if (b) return s!!.length
-    return 10
+    return if (b) s!!.length else 10
 }

j2k/testData/fileOrElement/nullability/Overrides.kt

@@ -25,7 +25,7 @@ internal class C : Base(), I {
         return ""
     }
 
-    override fun zoo(o: Any?): String? {
+    override fun zoo(o: Any?): String {
         return ""
     }
 }

j2k/testData/fileOrElement/nullability/autoNotNull.kt

@@ -1,8 +1,5 @@
 internal class A {
     fun foo(s: String?): Int {
-        if (s != null) {
-            return s.length
-        }
-        return -1
+        return s?.length ?: -1
     }
 }

j2k/testData/fileOrElement/nullability/nullableField.kt

@@ -1,4 +1,3 @@
-// ERROR: Type mismatch: inferred type is String? but String was expected
 // ERROR: Type mismatch: inferred type is Int? but Int was expected
 // ERROR: Type inference failed. Please try to specify type arguments explicitly.
 // ERROR: Using 'remove(Int): T' is an error. Use removeAt(index) instead.
@@ -50,7 +49,7 @@ class Test {
 
     }
 
-    fun foo2(s: String) {
+    fun foo2(s: String?) {
 
     }
 

j2k/testData/fileOrElement/nullability/nullableInitializer1.kt

@@ -1,20 +1,18 @@
-// ERROR: Type mismatch: inferred type is String? but String was expected
-// ERROR: Type mismatch: inferred type is String? but String was expected
 class Test {
     fun nullableString(p: Int): String? {
         return if (p > 0) "response" else null
     }
 
-    private var nullableInitializerField: String = nullableString(3)
+    private var nullableInitializerField = nullableString(3)
     private val nullableInitializerFieldFinal = nullableString(3)
-    var nullableInitializerPublicField: String = nullableString(3)
+    var nullableInitializerPublicField = nullableString(3)
 
     fun testProperty() {
         nullableInitializerField = "aaa"
 
-        nullableInitializerField[0]
+        nullableInitializerField!![0]
         nullableInitializerFieldFinal!![0]
-        nullableInitializerPublicField[0]
+        nullableInitializerPublicField!![0]
     }
 
     fun testLocalVariable() {

j2k/testData/fileOrElement/switch/nestedIf.kt

@@ -1,9 +1,9 @@
 fun foo(i: Int, j: Int): String {
     when (i) {
-        0 -> if (j > 0) {
-            return "1"
+        0 -> return if (j > 0) {
+            "1"
         } else {
-            return "2"
+            "2"
         }
         1 -> return "3"
         else -> return "4"

j2k/testData/fileOrElement/switch/nestedIfFallback.kt

@@ -1,10 +1,9 @@
 fun foo(i: Int, j: Int): String {
     when (i) {
         0 -> {
-            if (j > 0) {
-                return "1"
-            }
-            return "2"
+            return if (j > 0) {
+                "1"
+            } else "2"
         }
         1 -> return "2"
         else -> return "3"

j2k/testData/fileOrElement/typeCastExpression/castNullable.kt

@@ -1,7 +1,6 @@
 internal class A {
     private fun foo(o: Any?, b: Boolean): String? {
-        if (b) return o as String?
-        return ""
+        return if (b) o as String? else ""
     }
 
     fun bar() {

j2k/testData/multiFile/NullabilityByDFa/SomeServiceUsage.java

@@ -0,0 +1,53 @@
+
+public class SomeServiceUsage {
+
+    public SomeService getService() {
+        return SomeService.getInstanceNotNull();
+    }
+
+    public SomeService getServiceNullable() {
+        return SomeService.getInstanceNullable();
+    }
+
+    // elvis
+    public SomeService getServiceNotNullByDataFlow() {
+        SomeService s = SomeService.getInstanceNullable();
+        return s == null ? SomeService.getInstanceNotNull() : s;
+    }
+
+    // nullable, bang-bang
+    public String aString1() {
+        return getServiceNullable().nullableString();
+    }
+
+    // nullable
+    public String aString2() {
+        return getService().nullableString();
+    }
+
+    // not nullable
+    public String aString3() {
+        return getService().notNullString();
+    }
+
+    // nullable, no bang-bang
+    public String aString4() {
+        return getServiceNotNullByDataFlow().nullableString();
+    }
+
+    // not nullable, no bang-bang
+    public String aString5() {
+        return getServiceNotNullByDataFlow().notNullString();
+    }
+
+    // nullable, safe-call
+    public String aString6() {
+        SomeService s = getServiceNullable();
+        if (s != null) {
+            return s.nullableString();
+        } else {
+            return null;
+        }
+    }
+
+}

j2k/testData/multiFile/NullabilityByDFa/SomeServiceUsage.kt

@@ -0,0 +1,47 @@
+class SomeServiceUsage {
+
+    val service: SomeService
+        get() = SomeService.getInstanceNotNull()
+
+    val serviceNullable: SomeService?
+        get() = SomeService.getInstanceNullable()
+
+    // elvis
+    val serviceNotNullByDataFlow: SomeService
+        get() {
+            val s = SomeService.getInstanceNullable()
+            return s ?: SomeService.getInstanceNotNull()
+        }
+
+    // nullable, bang-bang
+    fun aString1(): String? {
+        return serviceNullable!!.nullableString()
+    }
+
+    // nullable
+    fun aString2(): String? {
+        return service.nullableString()
+    }
+
+    // not nullable
+    fun aString3(): String {
+        return service.notNullString()
+    }
+
+    // nullable, no bang-bang
+    fun aString4(): String? {
+        return serviceNotNullByDataFlow.nullableString()
+    }
+
+    // not nullable, no bang-bang
+    fun aString5(): String {
+        return serviceNotNullByDataFlow.notNullString()
+    }
+
+    // nullable, safe-call
+    fun aString6(): String? {
+        val s = serviceNullable
+        return s?.nullableString()
+    }
+
+}

j2k/testData/multiFile/NullabilityByDFa/external/JavaFile.java

@@ -0,0 +1,26 @@
+
+public class SomeService {
+
+    public static SomeService getInstanceNotNull() {
+        return new SomeService();
+    }
+
+    public static SomeService getInstanceNullable() {
+        if (Math.random() > 0.5)
+            return null;
+        return new SomeService();
+    }
+
+    public String nullableString() {
+        if (Math.random() < 0.5)
+            return null;
+        return Math.random() + "";
+    }
+
+    public String notNullString() {
+        String s = nullableString();
+        if (s != null)
+            return s;
+        return "null";
+    }
+}

j2k/testData/multiFile/NullabilityByDFa/external/JavaFile.java.expected

@@ -0,0 +1,26 @@
+
+public class SomeService {
+
+    public static SomeService getInstanceNotNull() {
+        return new SomeService();
+    }
+
+    public static SomeService getInstanceNullable() {
+        if (Math.random() > 0.5)
+            return null;
+        return new SomeService();
+    }
+
+    public String nullableString() {
+        if (Math.random() < 0.5)
+            return null;
+        return Math.random() + "";
+    }
+
+    public String notNullString() {
+        String s = nullableString();
+        if (s != null)
+            return s;
+        return "null";
+    }
+}

j2k/tests/org/jetbrains/kotlin/j2k/JavaToKotlinConverterMultiFileTestGenerated.java

@@ -66,6 +66,12 @@ public class JavaToKotlinConverterMultiFileTestGenerated extends AbstractJavaToK
         doTest(fileName);
     }
 
+    @TestMetadata("NullabilityByDFa")
+    public void testNullabilityByDFa() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("j2k/testData/multiFile/NullabilityByDFa/");
+        doTest(fileName);
+    }
+
     @TestMetadata("ProtectedVisibility")
     public void testProtectedVisibility() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("j2k/testData/multiFile/ProtectedVisibility/");