Implicit receiver smart casts implementation and highlighting

compiler/frontend/src/org/jetbrains/kotlin/checkers/CheckerTestUtil.java

@@ -155,6 +155,12 @@ public class CheckerTestUtil {
                 debugAnnotations.add(new DebugInfoDiagnostic(expression, DebugInfoDiagnosticFactory.SMARTCAST));
             }
         }
+        //noinspection TestOnlyProblems
+        for (KtExpression expression : bindingContext.getSliceContents(BindingContext.IMPLICIT_RECEIVER_SMARTCAST).keySet()) {
+            if (PsiTreeUtil.isAncestor(root, expression, false)) {
+                debugAnnotations.add(new DebugInfoDiagnostic(expression, DebugInfoDiagnosticFactory.IMPLICIT_RECEIVER_SMARTCAST));
+            }
+        }
         return debugAnnotations;
     }
 
@@ -510,6 +516,7 @@ public class CheckerTestUtil {
 
     public static class DebugInfoDiagnosticFactory extends DiagnosticFactory<DebugInfoDiagnostic> {
         public static final DebugInfoDiagnosticFactory SMARTCAST = new DebugInfoDiagnosticFactory("SMARTCAST");
+        public static final DebugInfoDiagnosticFactory IMPLICIT_RECEIVER_SMARTCAST = new DebugInfoDiagnosticFactory("IMPLICIT_RECEIVER_SMARTCAST");
         public static final DebugInfoDiagnosticFactory ELEMENT_WITH_ERROR_TYPE = new DebugInfoDiagnosticFactory("ELEMENT_WITH_ERROR_TYPE");
         public static final DebugInfoDiagnosticFactory UNRESOLVED_WITH_TARGET = new DebugInfoDiagnosticFactory("UNRESOLVED_WITH_TARGET");
         public static final DebugInfoDiagnosticFactory MISSING_UNRESOLVED = new DebugInfoDiagnosticFactory("MISSING_UNRESOLVED");

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -131,6 +131,7 @@ public interface BindingContext {
     WritableSlice<KtExpression, ResolvedCall<FunctionDescriptor>> INDEXED_LVALUE_SET = Slices.createSimpleSlice();
 
     WritableSlice<KtExpression, KotlinType> SMARTCAST = Slices.createSimpleSlice();
+    WritableSlice<KtExpression, KotlinType> IMPLICIT_RECEIVER_SMARTCAST = Slices.createSimpleSlice();
 
     WritableSlice<KtWhenExpression, Boolean> EXHAUSTIVE_WHEN = Slices.createSimpleSlice();
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CandidateResolver.kt

@@ -451,7 +451,7 @@ public class CandidateResolver(
             val dataFlowValue = DataFlowValueFactory.createDataFlowValue(receiverArgument, this)
 
             val smartCastResult = SmartCastManager.checkAndRecordPossibleCast(
-                    dataFlowValue, expectedReceiverParameterType, expression, this, /*recordType =*/ true
+                    dataFlowValue, expectedReceiverParameterType, expression, this, candidateCall.call.calleeExpression, /*recordType =*/true
             )
 
             if (smartCastResult == null) {

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

@@ -40,6 +40,7 @@ import java.util.List;
 import java.util.Set;
 
 import static org.jetbrains.kotlin.diagnostics.Errors.SMARTCAST_IMPOSSIBLE;
+import static org.jetbrains.kotlin.resolve.BindingContext.IMPLICIT_RECEIVER_SMARTCAST;
 import static org.jetbrains.kotlin.resolve.BindingContext.SMARTCAST;
 
 public class SmartCastManager {
@@ -174,6 +175,7 @@ public class SmartCastManager {
             @NotNull KotlinType expectedType,
             @Nullable KtExpression expression,
             @NotNull ResolutionContext c,
+            @Nullable KtExpression calleeExpression,
             boolean recordExpressionType
     ) {
         for (KotlinType possibleType : c.dataFlowInfo.getPossibleTypes(dataFlowValue)) {
@@ -181,6 +183,9 @@ public class SmartCastManager {
                 if (expression != null) {
                     recordCastOrError(expression, possibleType, c.trace, dataFlowValue, recordExpressionType);
                 }
+                else if (calleeExpression != null && dataFlowValue.isPredictable()) {
+                    c.trace.record(IMPLICIT_RECEIVER_SMARTCAST, calleeExpression, possibleType);
+                }
                 return new SmartCastResult(possibleType, dataFlowValue.isPredictable());
             }
         }
@@ -211,7 +216,7 @@ public class SmartCastManager {
 
                 return new SmartCastResult(dataFlowValue.getType(), immanentlyNotNull || dataFlowValue.isPredictable());
             }
-            return checkAndRecordPossibleCast(dataFlowValue, nullableExpectedType, expression, c, recordExpressionType);
+            return checkAndRecordPossibleCast(dataFlowValue, nullableExpectedType, expression, c, calleeExpression, recordExpressionType);
         }
 
         return null;

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

@@ -252,7 +252,7 @@ public class DataFlowAnalyzer {
     ) {
         DataFlowValue dataFlowValue = DataFlowValueFactory.createDataFlowValue(expression, expressionType, c);
 
-        SmartCastResult result = SmartCastManager.checkAndRecordPossibleCast(dataFlowValue, c.expectedType, expression, c, false);
+        SmartCastResult result = SmartCastManager.checkAndRecordPossibleCast(dataFlowValue, c.expectedType, expression, c, null, false);
         return result != null ? result.getResultType() : null;
     }
 

compiler/testData/diagnostics/tests/generics/nullability/smartCastsOnThis.kt

@@ -25,7 +25,7 @@ fun <T : String?> T.foo() {
     <!UNSAFE_CALL!>length<!>
 
     if (this is String) {
-        length
+        <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>length<!>
         this<!UNNECESSARY_SAFE_CALL!>?.<!>length
 
         bar1()

compiler/testData/diagnostics/tests/nullabilityAndSmartCasts/kt2109.kt

@@ -9,7 +9,7 @@ fun A?.bar() {
     if (this == null) {
         return
     }
-    foo()
+    <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>foo<!>()
 }
 
 fun A.baz() {

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

@@ -0,0 +1,15 @@
+open class SuperFoo {
+    public fun bar(): String {
+        if (this is Foo) {
+            <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>superFoo<!>()
+            return baz()
+        }
+        return baz()
+    }
+
+    public fun baz() = "OK"
+}
+
+class Foo : SuperFoo() {
+    public fun superFoo() {}
+}

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

@@ -0,0 +1,20 @@
+package
+
+public final class Foo : SuperFoo {
+    public constructor Foo()
+    public final override /*1*/ /*fake_override*/ fun bar(): kotlin.String
+    public final override /*1*/ /*fake_override*/ fun baz(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public final fun superFoo(): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public open class SuperFoo {
+    public constructor SuperFoo()
+    public final fun bar(): kotlin.String
+    public final fun baz(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

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

@@ -0,0 +1,24 @@
+open class A {
+    class B : A() {
+        val a = "FAIL"
+    }
+
+    fun foo(): String {
+        if (this is B) return <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>a<!>
+        return "OK"
+    }
+}
+
+fun A?.bar() {
+    if (this != null) <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>foo<!>()
+}
+
+fun A.gav() = if (this is A.B) <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>a<!> else ""
+
+class C {
+    fun A?.complex(): String {
+        if (this is A.B) return <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>a<!>
+        else if (this != null) return <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>foo<!>()
+        else return ""
+    }
+}

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

@@ -0,0 +1,29 @@
+package
+
+public fun A?.bar(): kotlin.Unit
+public fun A.gav(): kotlin.String
+
+public open class A {
+    public constructor A()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public final fun foo(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+
+    public final class B : A {
+        public constructor B()
+        public final val a: kotlin.String = "FAIL"
+        public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+        public final override /*1*/ /*fake_override*/ fun foo(): kotlin.String
+        public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+        public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    }
+}
+
+public final class C {
+    public constructor C()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    public final fun A?.complex(): kotlin.String
+}

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

@@ -0,0 +1,11 @@
+open class A {
+    open fun foo() = "FAIL"
+
+    fun bar() = if (this is C) <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>foo<!>() else "FAIL"
+}
+
+open class B : A()
+
+open class C : B() {
+    override fun foo() = "OK"
+}

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

@@ -0,0 +1,28 @@
+package
+
+public open class A {
+    public constructor A()
+    public final fun bar(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open fun foo(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public open class B : A {
+    public constructor B()
+    public final override /*1*/ /*fake_override*/ fun bar(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun foo(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public open class C : B {
+    public constructor C()
+    public final override /*1*/ /*fake_override*/ fun bar(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ fun foo(): kotlin.String
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/testData/diagnostics/tests/syntheticExtensions/javaProperties/SmartCastImplicitReceiver.kt

@@ -1,8 +1,8 @@
 // FILE: KotlinFile.kt
 fun Any.foo(): Int {
     if (this is JavaClass) {
-        something++
-        return x
+        <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>something<!>++
+        return <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>x<!>
     }
     return 0
 }

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

@@ -14586,6 +14586,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("falseReceiverSmartCast.kt")
+            public void testFalseReceiverSmartCast() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/falseReceiverSmartCast.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("fieldExclExcl.kt")
             public void testFieldExclExcl() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/fieldExclExcl.kt");
@@ -14604,6 +14610,18 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("implicitReceiver.kt")
+            public void testImplicitReceiver() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/implicitReceiver.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("implicitToGrandSon.kt")
+            public void testImplicitToGrandSon() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/implicitToGrandSon.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("incDecToNull.kt")
             public void testIncDecToNull() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/smartCasts/incDecToNull.kt");

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

@@ -74,6 +74,11 @@ class VariablesHighlightingVisitor extends AfterAnalysisHighlightingVisitor {
             holder.createInfoAnnotation(expression, "Smart cast to " + DescriptorRenderer.FQ_NAMES_IN_TYPES.renderType(smartCast)).setTextAttributes(
                     KotlinHighlightingColors.SMART_CAST_VALUE);
         }
+        smartCast = bindingContext.get(IMPLICIT_RECEIVER_SMARTCAST, expression);
+        if (smartCast != null) {
+            holder.createInfoAnnotation(expression, "Implicit receiver smart cast to " + DescriptorRenderer.FQ_NAMES_IN_TYPES.renderType(smartCast)).setTextAttributes(
+                    KotlinHighlightingColors.SMART_CAST_VALUE);
+        }
         super.visitExpression(expression);
     }
 

idea/testData/checker/infos/SmartCasts.kt

@@ -1,5 +1,9 @@
 <info descr="null">open</info> class A() {
-  fun foo() {}
+  fun foo() {
+      if (this is B) {
+          <info descr="Implicit receiver smart cast to B">bar</info>()
+      }
+  }
 }
 
 class B() : A() {