Make smart casts work with variable declared in when subject

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/util/CallMaker.java

@@ -324,6 +324,12 @@ public class CallMaker {
         return new ExpressionValueArgument(expression, expression, true);
     }
 
+    @NotNull
+    public static ValueArgument makeExternalValueArgument(@NotNull KtExpression expression, @NotNull KtElement reportErrorsOn) {
+        return new ExpressionValueArgument(expression, reportErrorsOn, true);
+    }
+
+
     @NotNull
     public static Call makePropertyCall(@Nullable Receiver explicitReceiver, @Nullable ASTNode callOperationNode, @NotNull KtSimpleNameExpression nameExpression) {
         return makeCallWithExpressions(nameExpression, explicitReceiver, callOperationNode, nameExpression, Collections.emptyList());

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/PatternMatchingTypingVisitor.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.diagnostics.Errors.*
 import org.jetbrains.kotlin.psi.*
@@ -30,9 +31,7 @@ import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.calls.checkers.RttiExpressionInformation
 import org.jetbrains.kotlin.resolve.calls.checkers.RttiOperation
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency.INDEPENDENT
-import org.jetbrains.kotlin.resolve.calls.smartcasts.ConditionalDataFlowInfo
+import org.jetbrains.kotlin.resolve.calls.smartcasts.*
-import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
-import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue
 import org.jetbrains.kotlin.resolve.calls.util.CallMaker
 import org.jetbrains.kotlin.resolve.checkers.PrimitiveNumericComparisonCallChecker
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
@@ -86,6 +85,84 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
     override fun visitWhenExpression(expression: KtWhenExpression, context: ExpressionTypingContext) =
         visitWhenExpression(expression, context, false)
 
+    private sealed class Subject(
+            val element: KtElement?,
+            val typeInfo: KotlinTypeInfo?,
+            val scopeWithSubject: LexicalScope?,
+            val type: KotlinType = typeInfo?.type ?: ErrorUtils.createErrorType("Unknown type")
+    ) {
+
+        protected abstract fun createDataFlowValue(contextAfterSubject: ExpressionTypingContext, builtIns: KotlinBuiltIns): DataFlowValue
+        abstract fun makeValueArgument(): ValueArgument
+        abstract val valueExpression: KtExpression
+
+        private var _dataFlowValue: DataFlowValue? = null
+        val dataFlowValue get() = _dataFlowValue!!
+
+        fun initDataFlowValue(contextAfterSubject: ExpressionTypingContext, builtIns: KotlinBuiltIns) {
+            _dataFlowValue = createDataFlowValue(contextAfterSubject, builtIns)
+        }
+
+        val dataFlowInfo get() = typeInfo?.dataFlowInfo
+
+        val jumpOutPossible get() = typeInfo?.jumpOutPossible ?: false
+
+        class Expression(
+                val expression: KtExpression,
+                typeInfo: KotlinTypeInfo
+        ) : Subject(expression, typeInfo, null) {
+            override fun createDataFlowValue(contextAfterSubject: ExpressionTypingContext, builtIns: KotlinBuiltIns) =
+                    DataFlowValueFactory.createDataFlowValue(expression, type, contextAfterSubject)
+
+            override fun makeValueArgument(): ValueArgument =
+                    CallMaker.makeExternalValueArgument(expression)
+
+            override val valueExpression: KtExpression
+                get() = expression
+        }
+
+
+        class Variable(
+                val variable: KtProperty,
+                val descriptor: VariableDescriptor,
+                typeInfo: KotlinTypeInfo,
+                scopeWithSubject: LexicalScope
+        ) : Subject(variable, typeInfo, scopeWithSubject) {
+            private val initializer = variable.initializer!!
+
+            override fun createDataFlowValue(contextAfterSubject: ExpressionTypingContext, builtIns: KotlinBuiltIns) =
+                    DataFlowValue(
+                            IdentifierInfo.Variable(
+                                    descriptor,
+                                    DataFlowValue.Kind.STABLE_VALUE,
+                                    contextAfterSubject.trace.bindingContext[BindingContext.BOUND_INITIALIZER_VALUE, descriptor]
+                            ),
+                            descriptor.type
+                    )
+
+            override fun makeValueArgument(): ValueArgument =
+                    CallMaker.makeExternalValueArgument(
+                            KtPsiFactory(variable.project, true).createExpression(variable.name!!),
+                            initializer
+                    )
+
+            override val valueExpression: KtExpression
+                get() = initializer
+        }
+
+
+        class None : Subject(null, null, null) {
+            override fun createDataFlowValue(contextAfterSubject: ExpressionTypingContext, builtIns: KotlinBuiltIns) =
+                    DataFlowValue.nullValue(builtIns)
+
+            override fun makeValueArgument(): ValueArgument =
+                    error("Should not be called for Subject.None")
+
+            override val valueExpression: KtExpression
+                get() = error("Should not be called for Subject.None")
+        }
+    }
+
     fun visitWhenExpression(
         expression: KtWhenExpression,
         contextWithExpectedType: ExpressionTypingContext,
@@ -104,63 +181,67 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
         val subjectExpression = expression.subjectExpression
         val subjectVariable = expression.subjectVariable
 
-        val subjectTypeInfo: KotlinTypeInfo?
+        val subject = when {
-        val scopeWithSubjectVariable: LexicalScope?
+            subjectVariable != null -> {
+                if (!components.languageVersionSettings.supportsFeature(LanguageFeature.VariableDeclarationInWhenSubject)) {
+                    trace.report(Errors.UNSUPPORTED_FEATURE.on(
+                            subjectVariable,
+                            Pair(LanguageFeature.VariableDeclarationInWhenSubject, components.languageVersionSettings)
+                    ))
+                }
 
-        if (subjectVariable != null) {
+                val scopeWithSubjectVariable = ExpressionTypingUtils.newWritableScopeImpl(contextBeforeSubject, LexicalScopeKind.WHEN, components.overloadChecker)
-            if (!components.languageVersionSettings.supportsFeature(LanguageFeature.VariableDeclarationInWhenSubject)) {
+
-                trace.report(Errors.UNSUPPORTED_FEATURE.on(
+                // Destructuring causes SOE in UAST :(
-                        subjectVariable,
+                val resolveResult = components.localVariableResolver.process(subjectVariable, contextBeforeSubject, contextBeforeSubject.scope, facade)
-                        Pair(LanguageFeature.VariableDeclarationInWhenSubject, components.languageVersionSettings)
+                val typeInfo = resolveResult.first
-                ))
+                val descriptor = resolveResult.second
+
+                scopeWithSubjectVariable.addVariableDescriptor(descriptor)
+
+                // NB typeInfo returned by 'localVariableResolver.process(...)' treats local variable declaration as a statement,
+                // so 'typeInfo' above it has type 'kotlin.Unit'.
+                // Propagate declared variable type as a "subject expression" type.
+                Subject.Variable(subjectVariable, descriptor, typeInfo.replaceType(descriptor.type), scopeWithSubjectVariable)
             }
 
-            scopeWithSubjectVariable = ExpressionTypingUtils.newWritableScopeImpl(contextBeforeSubject, LexicalScopeKind.WHEN, components.overloadChecker)
+            subjectExpression != null ->
+                Subject.Expression(subjectExpression, facade.getTypeInfo(subjectExpression, contextBeforeSubject))
 
-            val (typeInfo, descriptor) = components.localVariableResolver.process(subjectVariable, contextBeforeSubject, contextBeforeSubject.scope, facade)
+            else -> Subject.None()
-            // NB typeInfo returned by 'localVariableResolver.process(...)' treats local variable declaration as a statement,
-            // so 'typeInfo' above it has type 'kotlin.Unit'.
-            // Propagate declared variable type as a "subject expression" type.
-            subjectTypeInfo = typeInfo.replaceType(descriptor.type)
-            scopeWithSubjectVariable.addVariableDescriptor(descriptor)
-        }
-        else {
-            scopeWithSubjectVariable = null
-
-            subjectTypeInfo = subjectExpression?.let { facade.getTypeInfo(it, contextBeforeSubject) }
         }
 
         val contextAfterSubject = run {
-            var context = contextBeforeSubject
+            var result = contextBeforeSubject
-            if (subjectTypeInfo != null) {
+            subject.scopeWithSubject?.let { result = result.replaceScope(it) }
-                context = context.replaceDataFlowInfo(subjectTypeInfo.dataFlowInfo)
+            subject.dataFlowInfo?.let { result = result.replaceDataFlowInfo(it) }
-            }
+            result
-            if (scopeWithSubjectVariable != null) {
-                context = context.replaceScope(scopeWithSubjectVariable)
-            }
-            context
         }
 
-        val contextWithExpectedTypeAndSubjectVariable = scopeWithSubjectVariable?.let { contextWithExpectedType.replaceScope(it) } ?:
+        val contextWithExpectedTypeAndSubjectVariable =
-                                                        contextWithExpectedType
+                subject.scopeWithSubject?.let { contextWithExpectedType.replaceScope(it) } ?:
+                contextWithExpectedType
 
-        val subjectType = subjectTypeInfo?.type ?: ErrorUtils.createErrorType("Unknown type")
+//        val subjectType = subjectTypeInfo?.type ?: ErrorUtils.createErrorType("Unknown type")
-        val jumpOutPossibleInSubject: Boolean = subjectTypeInfo?.jumpOutPossible ?: false
+//        val jumpOutPossibleInSubject: Boolean = subjectTypeInfo?.jumpOutPossible ?: false
-        val subjectDataFlowValue = subjectExpression?.let {
+//        val subjectDataFlowValue = subjectExpression?.let {
-            facade.components.dataFlowValueFactory.createDataFlowValue(it, subjectType, contextAfterSubject)
+//            facade.components.dataFlowValueFactory.createDataFlowValue(it, subjectType, contextAfterSubject)
-        } ?: DataFlowValue.nullValue(components.builtIns)
+//        } ?: DataFlowValue.nullValue(components.builtIns)
+        subject.initDataFlowValue(contextAfterSubject, components.builtIns)
 
         val possibleTypesForSubject = subjectTypeInfo?.dataFlowInfo?.getStableTypes(
             subjectDataFlowValue, components.languageVersionSettings
         ) ?: emptySet()
         checkSmartCastsInSubjectIfRequired(expression, contextBeforeSubject, subjectType, possibleTypesForSubject)
+        val possibleTypesForSubject = subject.typeInfo?.dataFlowInfo?.getStableTypes(subject.dataFlowValue) ?: emptySet()
+        checkSmartCastsInSubjectIfRequired(expression, contextBeforeSubject, subject.type, possibleTypesForSubject)
 
-        val dataFlowInfoForEntries = analyzeConditionsInWhenEntries(expression, contextAfterSubject, subjectDataFlowValue, subjectType)
+        val dataFlowInfoForEntries = analyzeConditionsInWhenEntries(expression, contextAfterSubject, subject)
         val whenReturnType = inferTypeForWhenExpression(expression, contextWithExpectedTypeAndSubjectVariable, contextAfterSubject, dataFlowInfoForEntries)
         val whenResultValue = whenReturnType?.let { facade.components.dataFlowValueFactory.createDataFlowValue(expression, it, contextAfterSubject) }
 
         val branchesTypeInfo =
-            joinWhenExpressionBranches(expression, contextAfterSubject, whenReturnType, jumpOutPossibleInSubject, whenResultValue)
+                joinWhenExpressionBranches(expression, contextAfterSubject, whenReturnType, subject.jumpOutPossible, whenResultValue)
 
         val isExhaustive = WhenChecker.isWhenExhaustive(expression, trace)
 
@@ -219,20 +300,17 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
     }
 
     private fun analyzeConditionsInWhenEntries(
-        expression: KtWhenExpression,
+            expression: KtWhenExpression,
-        contextAfterSubject: ExpressionTypingContext,
+            contextAfterSubject: ExpressionTypingContext,
-        subjectDataFlowValue: DataFlowValue,
+            subject: Subject
-        subjectType: KotlinType
     ): ArrayList<DataFlowInfo> {
-        val subjectExpression = expression.subjectExpression
-
         val argumentDataFlowInfos = ArrayList<DataFlowInfo>()
         var inputDataFlowInfo = contextAfterSubject.dataFlowInfo
         for (whenEntry in expression.entries) {
             val conditionsInfo = analyzeWhenEntryConditions(
-                whenEntry,
+                    whenEntry,
-                contextAfterSubject.replaceDataFlowInfo(inputDataFlowInfo),
+                    contextAfterSubject.replaceDataFlowInfo(inputDataFlowInfo),
-                subjectExpression, subjectType, subjectDataFlowValue
+                    subject
             )
             inputDataFlowInfo = inputDataFlowInfo.and(conditionsInfo.elseInfo)
 
@@ -273,12 +351,11 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
                 }
 
             currentDataFlowInfo =
-                    if (entryType != null && KotlinBuiltIns.isNothing(entryType))
+                    when {
-                        currentDataFlowInfo
+                        entryType != null && KotlinBuiltIns.isNothing(entryType) -> currentDataFlowInfo
-                    else if (currentDataFlowInfo != null)
+                        currentDataFlowInfo != null -> currentDataFlowInfo.or(entryDataFlowInfo)
-                        currentDataFlowInfo.or(entryDataFlowInfo)
+                        else -> entryDataFlowInfo
-                    else
+                    }
-                        entryDataFlowInfo
 
             jumpOutPossible = jumpOutPossible or entryTypeInfo.jumpOutPossible
         }
@@ -342,11 +419,9 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
     }
 
     private fun analyzeWhenEntryConditions(
-        whenEntry: KtWhenEntry,
+            whenEntry: KtWhenEntry,
-        context: ExpressionTypingContext,
+            context: ExpressionTypingContext,
-        subjectExpression: KtExpression?,
+            subject: Subject
-        subjectType: KotlinType,
-        subjectDataFlowValue: DataFlowValue
     ): ConditionalDataFlowInfo {
         if (whenEntry.isElse) {
             return ConditionalDataFlowInfo(context.dataFlowInfo)
@@ -355,10 +430,7 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
         var entryInfo: ConditionalDataFlowInfo? = null
         var contextForCondition = context
         for (condition in whenEntry.conditions) {
-            val conditionInfo = checkWhenCondition(
+            val conditionInfo = checkWhenCondition(subject, condition, contextForCondition)
-                subjectExpression, subjectType, condition,
-                contextForCondition, subjectDataFlowValue
-            )
             entryInfo = entryInfo?.let {
                 ConditionalDataFlowInfo(it.thenInfo.or(conditionInfo.thenInfo), it.elseInfo.and(conditionInfo.elseInfo))
             } ?: conditionInfo
@@ -370,23 +442,24 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
     }
 
     private fun checkWhenCondition(
-        subjectExpression: KtExpression?,
+            subject: Subject,
-        subjectType: KotlinType,
+            condition: KtWhenCondition,
-        condition: KtWhenCondition,
+            context: ExpressionTypingContext
-        context: ExpressionTypingContext,
-        subjectDataFlowValue: DataFlowValue
     ): ConditionalDataFlowInfo {
         var newDataFlowInfo = noChange(context)
+
         condition.accept(object : KtVisitorVoid() {
             override fun visitWhenConditionInRange(condition: KtWhenConditionInRange) {
                 val rangeExpression = condition.rangeExpression ?: return
-                if (subjectExpression == null) {
+
+                if (subject is Subject.None) {
                     context.trace.report(EXPECTED_CONDITION.on(condition))
                     val dataFlowInfo = facade.getTypeInfo(rangeExpression, context).dataFlowInfo
                     newDataFlowInfo = ConditionalDataFlowInfo(dataFlowInfo)
                     return
                 }
-                val argumentForSubject = CallMaker.makeExternalValueArgument(subjectExpression)
+
+                val argumentForSubject = subject.makeValueArgument()
                 val typeInfo = facade.checkInExpression(
                     condition, condition.operationReference,
                     argumentForSubject, rangeExpression, context
@@ -400,43 +473,40 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
             }
 
             override fun visitWhenConditionIsPattern(condition: KtWhenConditionIsPattern) {
-                if (subjectExpression == null) {
+                if (subject is Subject.None) {
                     context.trace.report(EXPECTED_CONDITION.on(condition))
                 }
-                val typeReference = condition.typeReference
+
-                if (typeReference != null) {
+                val typeReference = condition.typeReference ?: return
-                    val result = checkTypeForIs(context, condition, condition.isNegated, subjectType, typeReference, subjectDataFlowValue)
+
-                    newDataFlowInfo = if (condition.isNegated) {
+                val result = checkTypeForIs(context, condition, condition.isNegated, subject.type, typeReference, subject.dataFlowValue)
-                        ConditionalDataFlowInfo(result.elseInfo, result.thenInfo)
+                newDataFlowInfo = if (condition.isNegated) ConditionalDataFlowInfo(result.elseInfo, result.thenInfo) else result
-                    } else {
+                val rhsType = context.trace[BindingContext.TYPE, typeReference]
-                        result
+
-                    }
+                if (subject !is Subject.None) {
-                    val rhsType = context.trace[BindingContext.TYPE, typeReference]
+                    val rttiInformation = RttiExpressionInformation(
-                    if (subjectExpression != null) {
+                            subject = subject.element!!,
-                        val rttiInformation = RttiExpressionInformation(
+                            sourceType = subject.type,
-                            subject = subjectExpression,
-                            sourceType = subjectType,
                             targetType = rhsType,
                             operation = if (condition.isNegated) RttiOperation.NOT_IS else RttiOperation.IS
-                        )
+                    )
-                        components.rttiExpressionCheckers.forEach {
+                    components.rttiExpressionCheckers.forEach {
-                            it.check(rttiInformation, condition, context.trace)
+                        it.check(rttiInformation, condition, context.trace)
-                        }
                     }
                 }
+
             }
 
             override fun visitWhenConditionWithExpression(condition: KtWhenConditionWithExpression) {
-                val expression = condition.expression
+                val expression = condition.expression ?: return
-                if (expression != null) {
+                val basicDataFlowInfo = checkTypeForExpressionCondition(
-                    val basicDataFlowInfo = checkTypeForExpressionCondition(
+                    context, expression, subject.type, subject is Subject.None, subject.dataFlowValue
-                        context, expression, subjectType, subjectExpression, subjectDataFlowValue
+                )
-                    )
+                val moduleDescriptor = DescriptorUtils.getContainingModule(context.scope.ownerDescriptor)
-                    val moduleDescriptor = DescriptorUtils.getContainingModule(context.scope.ownerDescriptor)
+                val dataFlowInfoFromES = components.effectSystem.getDataFlowInfoWhenEquals(
-                    val dataFlowInfoFromES =
+                    subject.valueExpression, expression, context.trace, moduleDescriptor
-                        components.effectSystem.getDataFlowInfoWhenEquals(subjectExpression, expression, context.trace, moduleDescriptor)
+                )
-                    newDataFlowInfo = basicDataFlowInfo.and(dataFlowInfoFromES)
+                newDataFlowInfo = basicDataFlowInfo.and(dataFlowInfoFromES)
-                }
             }
 
             override fun visitKtElement(element: KtElement) {

compiler/testData/diagnostics/tests/when/withSubjectVariable/smartCastsOnSubjectVariable.kt

@@ -0,0 +1,8 @@
+// !LANGUAGE: +VariableDeclarationInWhenSubject
+
+fun test(x: Any?) =
+        when (val y = x) {
+            is String -> "String, length = ${<!DEBUG_INFO_SMARTCAST!>y<!>.length}"
+            null -> "Null"
+            else -> "Any, hashCode = ${<!DEBUG_INFO_SMARTCAST!>y<!>.hashCode()}"
+        }

compiler/testData/diagnostics/tests/when/withSubjectVariable/smartCastsOnSubjectVariable.txt

@@ -0,0 +1,3 @@
+package
+
+public fun test(/*0*/ x: kotlin.Any?): kotlin.String

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

@@ -22478,6 +22478,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("smartCastsOnSubjectVariable.kt")
+                public void testSmartCastsOnSubjectVariable() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/when/withSubjectVariable/smartCastsOnSubjectVariable.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("unsupportedFeature.kt")
                 public void testUnsupportedFeature() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/when/withSubjectVariable/unsupportedFeature.kt");

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsUsingJavacTestGenerated.java

@@ -22478,6 +22478,12 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                     doTest(fileName);
                 }
 
+                @TestMetadata("smartCastsOnSubjectVariable.kt")
+                public void testSmartCastsOnSubjectVariable() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/when/withSubjectVariable/smartCastsOnSubjectVariable.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("unsupportedFeature.kt")
                 public void testUnsupportedFeature() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/when/withSubjectVariable/unsupportedFeature.kt");