Accurate representation of smart casts on complex expressions

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");