Accurate representation of smart casts on complex expressions

compiler/testData/codegen/box/smartCasts/whenSmartCast.kt

@@ -0,0 +1,9 @@
+fun baz(s: String?): Int {
+    if (s == null) return 0
+    return when(s) {
+        "abc" -> s
+        else -> "xyz"
+    }.length
+}
+
+fun box() = if (baz("abc") == 3 && baz("") == 3 && baz(null) == 0) "OK" else "FAIL"

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

@@ -0,0 +1,10 @@
+fun baz(s: String?): Int {
+    return <!DEBUG_INFO_SMARTCAST!>if (s == null) {
+        ""
+    }
+    else {
+        val u: String? = null
+        if (u == null) return 0
+        u
+    }<!>.length
+}

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

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

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

@@ -0,0 +1,20 @@
+// Smart casts on complex expressions
+fun baz(s: String?): Int {
+    if (s == null) return 0
+    return <!DEBUG_INFO_SMARTCAST!>when(<!DEBUG_INFO_SMARTCAST!>s<!>) {
+        "abc" -> s
+        else -> "xyz"
+    }<!>.length
+}
+
+var ss: String? = null
+
+fun bar(): Int {
+    if (ss == null) return 0
+    // ss cannot be smart casted, so an error here
+    return when(ss) {
+        "abc" -> ss
+        else -> "xyz"
+    }<!UNSAFE_CALL!>.<!>length
+}
+

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

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

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

@@ -15225,6 +15225,18 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("smartCastOnIf.kt")
+            public void testSmartCastOnIf() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/smartCastOnIf.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("smartCastOnWhen.kt")
+            public void testSmartCastOnWhen() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/smartCastOnWhen.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("thisWithLabel.kt")
             public void testThisWithLabel() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/thisWithLabel.kt");

compiler/tests/org/jetbrains/kotlin/codegen/generated/BlackBoxCodegenTestGenerated.java

@@ -7275,6 +7275,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/smartCasts/nullSmartCast.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("whenSmartCast.kt")
+        public void testWhenSmartCast() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/smartCasts/whenSmartCast.kt");
+            doTest(fileName);
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/specialBuiltins")

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

@@ -84,8 +84,14 @@ class VariablesHighlightingVisitor extends AfterAnalysisHighlightingVisitor {
 
         KotlinType smartCast = bindingContext.get(SMARTCAST, expression);
         if (smartCast != null) {
-            holder.createInfoAnnotation(expression, "Smart cast to " +
-                                                    DescriptorRenderer.FQ_NAMES_IN_TYPES.renderType(smartCast))
+            PsiElement target = expression;
+            if (expression instanceof KtIfExpression) {
+                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/SmartCastOnIf.kt

@@ -0,0 +1,10 @@
+fun baz(s: String?): Int {
+    return <info descr="Smart cast to kotlin.String">if</info> (s == null) {
+        ""
+    }
+    else {
+        val u: String? = null
+        if (u == null) return 0
+        u
+    }.length
+}

idea/testData/checker/infos/SmartCastOnWhen.kt

@@ -0,0 +1,7 @@
+fun baz(s: String?): Int {
+    if (s == null) return 0
+    return <info descr="Smart cast to kotlin.String">when</info>(<info descr="Smart cast to kotlin.String">s</info>) {
+        "abc" -> s
+        else -> "xyz"
+    }.length
+}

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

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