diff --git a/compiler/testData/codegen/box/smartCasts/whenSmartCast.kt b/compiler/testData/codegen/box/smartCasts/whenSmartCast.kt
new file mode 100644
index 00000000000..8069048ce00
--- /dev/null
+++ b/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"
diff --git a/compiler/testData/diagnostics/tests/smartCasts/smartCastOnIf.kt b/compiler/testData/diagnostics/tests/smartCasts/smartCastOnIf.kt
new file mode 100644
index 00000000000..57af6a08b3b
--- /dev/null
+++ b/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
+}
diff --git a/compiler/testData/diagnostics/tests/smartCasts/smartCastOnIf.txt b/compiler/testData/diagnostics/tests/smartCasts/smartCastOnIf.txt
new file mode 100644
index 00000000000..ceb966e9a84
--- /dev/null
+++ b/compiler/testData/diagnostics/tests/smartCasts/smartCastOnIf.txt
@@ -0,0 +1,3 @@
+package
+
+public fun baz(/*0*/ s: kotlin.String?): kotlin.Int
diff --git a/compiler/testData/diagnostics/tests/smartCasts/smartCastOnWhen.kt b/compiler/testData/diagnostics/tests/smartCasts/smartCastOnWhen.kt
new file mode 100644
index 00000000000..ebbe6c09f32
--- /dev/null
+++ b/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
+}
+
diff --git a/compiler/testData/diagnostics/tests/smartCasts/smartCastOnWhen.txt b/compiler/testData/diagnostics/tests/smartCasts/smartCastOnWhen.txt
new file mode 100644
index 00000000000..f777a1470b8
--- /dev/null
+++ b/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
diff --git a/compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java
index 027f49b19d0..e3fb246dab8 100644
--- a/compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java
+++ b/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");
diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/generated/BlackBoxCodegenTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/generated/BlackBoxCodegenTestGenerated.java
index 306736e8f50..7599ceaff3d 100644
--- a/compiler/tests/org/jetbrains/kotlin/codegen/generated/BlackBoxCodegenTestGenerated.java
+++ b/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")
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/VariablesHighlightingVisitor.java b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/VariablesHighlightingVisitor.java
index 0319d07b318..05fd607b049 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/VariablesHighlightingVisitor.java
+++ b/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);
         }
 
diff --git a/idea/testData/checker/infos/SmartCastOnIf.kt b/idea/testData/checker/infos/SmartCastOnIf.kt
new file mode 100644
index 00000000000..a9be41eb19b
--- /dev/null
+++ b/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
+}
diff --git a/idea/testData/checker/infos/SmartCastOnWhen.kt b/idea/testData/checker/infos/SmartCastOnWhen.kt
new file mode 100644
index 00000000000..5034fb4592c
--- /dev/null
+++ b/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
+}
diff --git a/idea/tests/org/jetbrains/kotlin/checkers/PsiCheckerTestGenerated.java b/idea/tests/org/jetbrains/kotlin/checkers/PsiCheckerTestGenerated.java
index e18a51323bc..2ef38e9564d 100644
--- a/idea/tests/org/jetbrains/kotlin/checkers/PsiCheckerTestGenerated.java
+++ b/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");