Data flow values for Elvis / parenthesized expressions, smart casts on them

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/smartcasts/DataFlowValueFactory.java

@@ -66,6 +66,21 @@ public class DataFlowValueFactory {
                                    resolutionContext.scope.getOwnerDescriptor());
     }
 
+    private static boolean isComplexExpression(@NotNull KtExpression expression) {
+        if (expression instanceof KtBlockExpression ||
+            expression instanceof KtIfExpression ||
+            expression instanceof KtWhenExpression ||
+            (expression instanceof KtBinaryExpression && ((KtBinaryExpression) expression).getOperationToken() == KtTokens.ELVIS)) {
+
+            return true;
+        }
+        if (expression instanceof KtParenthesizedExpression) {
+            KtExpression deparenthesized = KtPsiUtil.deparenthesize(expression);
+            return deparenthesized != null && isComplexExpression(deparenthesized);
+        }
+        return false;
+    }
+
     @NotNull
     public static DataFlowValue createDataFlowValue(
             @NotNull KtExpression expression,
@@ -96,7 +111,7 @@ public class DataFlowValueFactory {
                                      Nullability.NOT_NULL);
         }
 
-        if (expression instanceof KtBlockExpression || expression instanceof KtIfExpression || expression instanceof KtWhenExpression) {
+        if (isComplexExpression(expression)) {
             return createDataFlowValueForComplexExpression(expression, type);
         }
 

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/BasicExpressionTypingVisitor.java

@@ -173,7 +173,14 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
         if (innerExpression == null) {
             return TypeInfoFactoryKt.noTypeInfo(context);
         }
-        return facade.getTypeInfo(innerExpression, context.replaceScope(context.scope));
+        KotlinTypeInfo result = facade.getTypeInfo(innerExpression, context.replaceScope(context.scope));
+        KotlinType resultType = result.getType();
+        if (resultType != null) {
+            DataFlowValue innerValue = DataFlowValueFactory.createDataFlowValue(innerExpression, resultType, context);
+            DataFlowValue resultValue = DataFlowValueFactory.createDataFlowValue(expression, resultType, context);
+            result = result.replaceDataFlowInfo(result.getDataFlowInfo().assign(resultValue, innerValue));
+        }
+        return result;
     }
 
     @Override
@@ -1278,19 +1285,30 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
         boolean loopBreakContinuePossible = leftTypeInfo.getJumpOutPossible() || rightTypeInfo.getJumpOutPossible();
         KotlinType rightType = rightTypeInfo.getType();
 
-        // Only left argument DFA is taken into account here: we cannot be sure that right argument is executed
+        // Only left argument DFA is taken into account here: we cannot be sure that right argument is joined
+        // (we merge it with right DFA if right argument contains no jump outside)
         DataFlowInfo dataFlowInfo = resolvedCall.getDataFlowInfoForArguments().getInfo(call.getValueArguments().get(1));
+
+        KotlinType type = resolvedCall.getResultingDescriptor().getReturnType();
+        if (type == null || rightType == null) return TypeInfoFactoryKt.noTypeInfo(dataFlowInfo);
         if (leftType != null) {
             DataFlowValue leftValue = createDataFlowValue(left, leftType, context);
             DataFlowInfo rightDataFlowInfo = resolvedCall.getDataFlowInfoForArguments().getResultInfo();
+            boolean jumpInRight = KotlinBuiltIns.isNothing(rightType);
+            DataFlowValue nullValue = DataFlowValue.nullValue(components.builtIns);
             // left argument is considered not-null if it's not-null also in right part or if we have jump in right part
-            if ((rightType != null && KotlinBuiltIns.isNothingOrNullableNothing(rightType) && !rightType.isMarkedNullable())
+            if (jumpInRight || !rightDataFlowInfo.getPredictableNullability(leftValue).canBeNull()) {
-                || !rightDataFlowInfo.getPredictableNullability(leftValue).canBeNull()) {
+                dataFlowInfo = dataFlowInfo.disequate(leftValue, nullValue);
-                dataFlowInfo = dataFlowInfo.disequate(leftValue, DataFlowValue.nullValue(components.builtIns));
+            }
+            DataFlowValue resultValue = DataFlowValueFactory.createDataFlowValue(expression, type, context);
+            dataFlowInfo = dataFlowInfo.assign(resultValue, leftValue).disequate(resultValue, nullValue);
+            if (!jumpInRight) {
+                DataFlowValue rightValue = DataFlowValueFactory.createDataFlowValue(right, rightType, context);
+                rightDataFlowInfo = rightDataFlowInfo.assign(resultValue, rightValue);
+                dataFlowInfo = dataFlowInfo.or(rightDataFlowInfo);
             }
         }
-        KotlinType type = resolvedCall.getResultingDescriptor().getReturnType();
+
-        if (type == null || rightType == null) return TypeInfoFactoryKt.noTypeInfo(dataFlowInfo);
 
         // Sometimes return type for special call for elvis operator might be nullable,
         // but result is not nullable if the right type is not nullable

compiler/testData/diagnostics/tests/smartCasts/elvisExprNotNull.kt

@@ -0,0 +1,37 @@
+fun foo(s: Any?): String {
+    val t = when {
+        // To resolve: String U Nothing? = String?
+        s is String -> s
+        else -> null
+    } ?: ""
+    // Ideally we should have smart cast to String here
+    return <!TYPE_MISMATCH!>t<!>
+}
+
+fun bar(s: Any?): String {
+    // To resolve: String U Nothing? = String?
+    val t = (if (s == null) {
+        null
+    }
+    else {
+        val u: Any? = null
+        if (u !is String) return ""
+        u
+    }) ?: "xyz"
+    // Ideally we should have smart cast to String here
+    return <!TYPE_MISMATCH!>t<!>
+}
+
+fun baz(s: String?, r: String?): String {
+    val t = r ?: when {
+        s != null -> s
+        else -> ""
+    }
+    return <!DEBUG_INFO_SMARTCAST!>t<!>
+}
+
+fun withNull(s: String?): String {
+    val t = s ?: null
+    // Error: nullable
+    return <!TYPE_MISMATCH!>t<!>
+}

compiler/testData/diagnostics/tests/smartCasts/elvisExprNotNull.txt

@@ -0,0 +1,6 @@
+package
+
+public fun bar(/*0*/ s: kotlin.Any?): kotlin.String
+public fun baz(/*0*/ s: kotlin.String?, /*1*/ r: kotlin.String?): kotlin.String
+public fun foo(/*0*/ s: kotlin.Any?): kotlin.String
+public fun withNull(/*0*/ s: kotlin.String?): kotlin.String

compiler/testData/diagnostics/tests/smartCasts/smartCastOnElvis.kt

@@ -0,0 +1,15 @@
+fun foo(s: String) = s.length
+
+fun baz(s: String?, r: String?): Int {
+    return foo(<!DEBUG_INFO_SMARTCAST!>r ?: when {
+        s != null -> s
+        else -> ""
+    }<!>)
+}
+
+fun bar(s: String?, r: String?): Int {
+    return <!DEBUG_INFO_SMARTCAST!>(r ?: when {
+        s != null -> s
+        else -> ""
+    })<!>.length
+}

compiler/testData/diagnostics/tests/smartCasts/smartCastOnElvis.txt

@@ -0,0 +1,5 @@
+package
+
+public fun bar(/*0*/ s: kotlin.String?, /*1*/ r: kotlin.String?): kotlin.Int
+public fun baz(/*0*/ s: kotlin.String?, /*1*/ r: kotlin.String?): kotlin.Int
+public fun foo(/*0*/ s: kotlin.String): kotlin.Int

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -14907,6 +14907,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("elvisExprNotNull.kt")
+            public void testElvisExprNotNull() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/elvisExprNotNull.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("elvisNothingRHS.kt")
             public void testElvisNothingRHS() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/elvisNothingRHS.kt");
@@ -15225,6 +15231,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("smartCastOnElvis.kt")
+            public void testSmartCastOnElvis() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/smartCastOnElvis.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("smartCastOnIf.kt")
             public void testSmartCastOnIf() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/smartCastOnIf.kt");

idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/VariablesHighlightingVisitor.java

@@ -67,6 +67,27 @@ class VariablesHighlightingVisitor extends AfterAnalysisHighlightingVisitor {
         super.visitParameter(parameter);
     }
 
+    @NotNull
+    private static PsiElement getSmartCastTarget(@NotNull KtExpression expression) {
+        PsiElement target = expression;
+        if (target instanceof KtParenthesizedExpression) {
+            target = KtPsiUtil.deparenthesize((KtParenthesizedExpression) target);
+            if (target == null) {
+                target = expression;
+            }
+        }
+        if (target instanceof KtIfExpression) {
+            target = ((KtIfExpression) target).getIfKeyword();
+        }
+        else if (target instanceof KtWhenExpression) {
+            target = ((KtWhenExpression) target).getWhenKeyword();
+        }
+        else if (target instanceof KtBinaryExpression) {
+            target = ((KtBinaryExpression) target).getOperationReference();
+        }
+        return target;
+    }
+
     @Override
     public void visitExpression(@NotNull KtExpression expression) {
         KotlinType implicitSmartCast = bindingContext.get(IMPLICIT_RECEIVER_SMARTCAST, expression);
@@ -84,14 +105,8 @@ class VariablesHighlightingVisitor extends AfterAnalysisHighlightingVisitor {
 
         KotlinType smartCast = bindingContext.get(SMARTCAST, expression);
         if (smartCast != null) {
-            PsiElement target = expression;
+            holder.createInfoAnnotation(getSmartCastTarget(expression),
-            if (expression instanceof KtIfExpression) {
+                                        "Smart cast to " + DescriptorRenderer.FQ_NAMES_IN_TYPES.renderType(smartCast))
-                target = ((KtIfExpression) expression).getIfKeyword();
-            }
-            else if (expression instanceof KtWhenExpression) {
-                target = ((KtWhenExpression) expression).getWhenKeyword();
-            }
-            holder.createInfoAnnotation(target, "Smart cast to " + DescriptorRenderer.FQ_NAMES_IN_TYPES.renderType(smartCast))
                   .setTextAttributes(KotlinHighlightingColors.SMART_CAST_VALUE);
         }
 

idea/testData/checker/infos/smartCastOnElvis.kt

@@ -0,0 +1,15 @@
+fun foo(s: String) = s.length
+
+fun baz(s: String?, r: String?): Int {
+    return foo(r <info descr="Smart cast to kotlin.String">?:</info> when {
+        s != null -> s
+        else -> ""
+    })
+}
+
+fun bar(s: String?, r: String?): Int {
+    return (r <info descr="Smart cast to kotlin.String">?:</info> when {
+        s != null -> s
+        else -> ""
+    }).length
+}

idea/tests/org/jetbrains/kotlin/checkers/PsiCheckerTestGenerated.java

@@ -724,6 +724,12 @@ public class PsiCheckerTestGenerated extends AbstractPsiCheckerTest {
             doTestWithInfos(fileName);
         }
 
+        @TestMetadata("smartCastOnElvis.kt")
+        public void testSmartCastOnElvis() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/checker/infos/smartCastOnElvis.kt");
+            doTestWithInfos(fileName);
+        }
+
         @TestMetadata("SmartCastOnIf.kt")
         public void testSmartCastOnIf() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/checker/infos/SmartCastOnIf.kt");