FIR DFA: unwrap transparent expressions in more places

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolver/resolverUtils.kt

@@ -58,9 +58,7 @@ internal open class StubBodyResolveTransformerComponents(
             override fun receiverUpdated(symbol: FirBasedSymbol<*>, types: Set<ConeKotlinType>?) =
                 error("Should not be called")
 
-            override fun getTypeUsingSmartcastInfo(
+            override fun getTypeUsingSmartcastInfo(expression: FirExpression): Pair<PropertyStability, MutableList<ConeKotlinType>>? =
-                symbol: FirBasedSymbol<*>,
+                null
-                expression: FirExpression
-            ): Pair<PropertyStability, MutableList<ConeKotlinType>>? = null
         }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt

@@ -27,12 +27,10 @@ import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.JumpNode
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertyAccessorSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.types.ConstantValueKind
-import org.jetbrains.kotlin.utils.addIfNotNull
 
 object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
 
@@ -110,7 +108,7 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         }
 
         val conditionStatements = effectDeclaration.condition.buildTypeStatements(
-            function, logicSystem, variableStorage, flow, context
+            function, logicSystem, variableStorage, context
         ) ?: return false
 
         for ((realVar, requiredTypeStatement) in conditionStatements) {
@@ -142,20 +140,16 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
     private fun ConeBooleanExpression.buildTypeStatements(
         function: FirFunction,
         logicSystem: LogicSystem<*>,
-        variableStorage: VariableStorageImpl,
+        variableStorage: VariableStorage,
-        flow: Flow,
         context: CheckerContext
     ): TypeStatements? {
-        fun getOrCreateRealVariable(arg: ConeValueParameterReference): RealVariable? {
+        fun ConeValueParameterReference.toVariable(): RealVariable {
-            val parameterSymbol = function.getParameterSymbol(arg.parameterIndex, context)
+            val parameterSymbol = function.getParameterSymbol(parameterIndex, context)
-
+            return variableStorage.getLocalVariable(parameterSymbol)
-            @OptIn(SymbolInternals::class)
+                ?: RealVariable(
-            val parameter = parameterSymbol.fir
+                    Identifier(parameterSymbol, null, null),
-            return variableStorage.getOrCreateRealVariable(flow, parameterSymbol, parameter)?.takeIf {
+                    parameterIndex < 0, null, parameterIndex + 1, PropertyStability.STABLE_VALUE
-                it.stability == PropertyStability.STABLE_VALUE ||
+                )
-                        // TODO: consider removing the part below
-                        it.stability == PropertyStability.LOCAL_VAR
-            }
         }
 
         fun ConeBooleanExpression.toTypeStatements(inverted: Boolean): TypeStatements? = when (this) {
@@ -170,10 +164,9 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
                 }
             }
             is ConeIsInstancePredicate ->
-                if (isNegated == inverted) getOrCreateRealVariable(arg)?.let { it typeEq type }?.singleton() else mapOf()
+                if (isNegated == inverted) (arg.toVariable() typeEq type).singleton() else mapOf()
             is ConeIsNullPredicate ->
-                getOrCreateRealVariable(arg)?.nullabilityStatement(context.session.builtinTypes, isNull = isNegated == inverted)
+                arg.toVariable().nullabilityStatement(context.session.builtinTypes, isNull = isNegated == inverted).singleton()
-                    ?.singleton()
             is ConeLogicalNot -> arg.toTypeStatements(!inverted)
             else -> null
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -153,41 +153,14 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     // ----------------------------------- Requests -----------------------------------
 
     fun isAccessToUnstableLocalVariable(expression: FirExpression): Boolean {
-        val qualifiedAccessExpression = when (expression) {
+        val analyzer = context.firLocalVariableAssignmentAnalyzer ?: return false
-            is FirSmartCastExpression -> expression.originalExpression as FirQualifiedAccessExpression
+        val realFir = expression.unwrapElement() as? FirQualifiedAccessExpression ?: return false
-            is FirQualifiedAccessExpression -> expression
+        return analyzer.isAccessToUnstableLocalVariable(realFir)
-            is FirWhenSubjectExpression -> {
-                val whenExpression = expression.whenRef.value
-                when {
-                    whenExpression.subjectVariable != null -> return true
-                    else -> whenExpression.subject as? FirQualifiedAccessExpression
-                }
-            }
-            else -> null
-        } ?: return false
-        return context.firLocalVariableAssignmentAnalyzer?.isAccessToUnstableLocalVariable(qualifiedAccessExpression) == true
     }
 
-    fun getTypeUsingSmartcastInfo(whenSubjectExpression: FirWhenSubjectExpression): Pair<PropertyStability, MutableList<ConeKotlinType>>? {
+    open fun getTypeUsingSmartcastInfo(expression: FirExpression): Pair<PropertyStability, MutableList<ConeKotlinType>>? {
-        val symbol = whenSubjectExpression.symbol ?: return null
-        return getTypeUsingSmartcastInfo(symbol, whenSubjectExpression)
-    }
-
-    fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): Pair<PropertyStability, MutableList<ConeKotlinType>>? {
-        /*
-         * DataFlowAnalyzer holds variables only for declarations that have some smartcast (or can have)
-         * If there is no useful information there is no data flow variable also
-         */
-        val symbol: FirBasedSymbol<*> = qualifiedAccessExpression.symbol ?: return null
-        return getTypeUsingSmartcastInfo(symbol, qualifiedAccessExpression)
-    }
-
-    protected open fun getTypeUsingSmartcastInfo(
-        symbol: FirBasedSymbol<*>,
-        expression: FirExpression
-    ): Pair<PropertyStability, MutableList<ConeKotlinType>>? {
         val flow = graphBuilder.lastNode.flow
-        val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(flow, symbol, expression) ?: return null
+        val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(flow, expression) ?: return null
         return flow.getType(variable)?.takeIf { it.isNotEmpty() }?.let { variable.stability to it.toMutableList() }
     }
 
@@ -968,7 +941,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             exitVariableInitialization(flow, assignment.rValue, property, assignment, hasExplicitType = false)
         } else {
             // TODO: add unstable smartcast for non-local var
-            val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(flow, property.symbol, assignment)
+            val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(flow, assignment)
             if (variable != null) {
                 logicSystem.recordNewAssignment(flow, variable, context.newAssignmentIndex())
             }
@@ -994,10 +967,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             logicSystem.recordNewAssignment(flow, propertyVariable, context.newAssignmentIndex())
         }
 
-        variableStorage.getOrCreateRealVariable(flow, initializer.symbol, initializer.unwrapSmartcastExpression())
+        variableStorage.getOrCreateRealVariable(flow, initializer)
             ?.let { initializerVariable ->
                 val isInitializerStable =
-                    initializerVariable.isStable || (initializerVariable.hasLocalStability && initializer.isAccessToStableVariable())
+                    initializerVariable.isStable || (initializerVariable.hasLocalStability && !isAccessToUnstableLocalVariable(initializer))
 
                 if (!hasExplicitType && isInitializerStable && (propertyVariable.hasLocalStability || propertyVariable.isStable)) {
                     logicSystem.addLocalVariableAlias(flow, propertyVariable, initializerVariable)
@@ -1026,9 +999,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    private fun FirExpression.isAccessToStableVariable(): Boolean =
-        !isAccessToUnstableLocalVariable(this)
-
     private val RealVariable.isStable get() = stability == PropertyStability.STABLE_VALUE
     private val RealVariable.hasLocalStability get() = stability == PropertyStability.LOCAL_VAR
 
@@ -1280,7 +1250,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     private fun resetReceivers(flow: FLOW) {
         receiverStack.forEach {
-            variableStorage.getRealVariable(flow, it.boundSymbol, it.receiverExpression)?.let { variable ->
+            variableStorage.getLocalVariable(it.boundSymbol)?.let { variable ->
                 receiverUpdated(it.boundSymbol, flow.getType(variable))
             }
         }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/VariableStorage.kt

@@ -9,8 +9,9 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 
 abstract class VariableStorage {
-    abstract fun getRealVariableWithoutUnwrappingAlias(flow: Flow, symbol: FirBasedSymbol<*>?, fir: FirElement): RealVariable?
+    abstract fun getRealVariableWithoutUnwrappingAlias(flow: Flow, fir: FirElement): RealVariable?
-    abstract fun getRealVariable(flow: Flow, symbol: FirBasedSymbol<*>?, fir: FirElement): RealVariable?
+    abstract fun getRealVariable(flow: Flow, fir: FirElement): RealVariable?
+    abstract fun getLocalVariable(symbol: FirBasedSymbol<*>): RealVariable?
     abstract fun getSyntheticVariable(fir: FirElement): SyntheticVariable?
     abstract fun getVariable(flow: Flow, fir: FirElement): DataFlowVariable?
 }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/VariableStorageImpl.kt

@@ -59,15 +59,6 @@ class VariableStorageImpl(private val session: FirSession) : VariableStorage() {
         return flow.unwrapVariable(getOrCreateRealVariableWithoutUnwrappingAlias(flow, symbol, fir, stability))
     }
 
-    private fun FirElement.unwrapElement(): FirElement = when (this) {
-        is FirWhenSubjectExpression -> whenRef.value.let { it.subjectVariable ?: it.subject }?.unwrapElement() ?: this
-        is FirSmartCastExpression -> originalExpression.unwrapElement()
-        is FirSafeCallExpression -> selector.unwrapElement()
-        is FirCheckedSafeCallSubject -> originalReceiverRef.value.unwrapElement()
-        is FirCheckNotNullCall -> argument.unwrapElement()
-        else -> this
-    }
-
     private fun getIdentifierBySymbol(
         flow: Flow,
         symbol: FirBasedSymbol<*>,
@@ -128,9 +119,12 @@ class VariableStorageImpl(private val session: FirSession) : VariableStorage() {
         }
     }
 
-    @JvmName("getOrCreateRealVariableOrNull")
+    fun getOrCreateRealVariable(flow: Flow, fir: FirElement): RealVariable? {
-    fun getOrCreateRealVariable(flow: Flow, symbol: FirBasedSymbol<*>?, fir: FirElement): RealVariable? =
+        val realFir = fir.unwrapElement()
-        symbol.getStability(fir)?.let { getOrCreateRealVariable(flow, symbol!!, fir, it) }
+        val symbol = realFir.symbol ?: return null
+        val stability = symbol.getStability(realFir) ?: return null
+        return getOrCreateRealVariable(flow, symbol, realFir, stability)
+    }
 
     fun createSyntheticVariable(fir: FirElement): SyntheticVariable =
         SyntheticVariable(fir, counter++).also { syntheticVariables[fir] = it }
@@ -146,30 +140,26 @@ class VariableStorageImpl(private val session: FirSession) : VariableStorage() {
         }
     }
 
-    override fun getRealVariableWithoutUnwrappingAlias(flow: Flow, symbol: FirBasedSymbol<*>?, fir: FirElement): RealVariable? {
+    override fun getRealVariableWithoutUnwrappingAlias(flow: Flow, fir: FirElement): RealVariable? {
         val realFir = fir.unwrapElement()
-        return symbol.takeIf { it.getStability(realFir) != null }?.let {
+        val symbol = realFir.symbol ?: return null
-            _realVariables[getIdentifierBySymbol(flow, it, realFir.unwrapElement())]
+        if (symbol.getStability(realFir) == null) return null
-        }
+        return _realVariables[getIdentifierBySymbol(flow, symbol, realFir)]
     }
 
-    override fun getRealVariable(flow: Flow, symbol: FirBasedSymbol<*>?, fir: FirElement): RealVariable? {
+    override fun getRealVariable(flow: Flow, fir: FirElement): RealVariable? {
-        return getRealVariableWithoutUnwrappingAlias(flow, symbol, fir)?.let { flow.unwrapVariable(it) }
+        return getRealVariableWithoutUnwrappingAlias(flow, fir)?.let { flow.unwrapVariable(it) }
     }
 
+    override fun getLocalVariable(symbol: FirBasedSymbol<*>): RealVariable? =
+        _realVariables[Identifier(symbol, null, null)]
+
     override fun getSyntheticVariable(fir: FirElement): SyntheticVariable? {
         return syntheticVariables[fir.unwrapElement()]
     }
 
     override fun getVariable(flow: Flow, fir: FirElement): DataFlowVariable? {
-        val realFir = fir.unwrapElement()
+        return getRealVariable(flow, fir) ?: getSyntheticVariable(fir)
-        val symbol = realFir.symbol
-        val stability = symbol.getStability(fir)
-        return if (stability != null) {
-            getRealVariable(flow, symbol, realFir)
-        } else {
-            getSyntheticVariable(fir)
-        }
     }
 
     fun removeRealVariable(symbol: FirBasedSymbol<*>) {

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt

@@ -114,6 +114,16 @@ internal val FirResolvable.symbol: FirBasedSymbol<*>?
         else -> null
     }
 
+@DfaInternals
+fun FirElement.unwrapElement(): FirElement = when (this) {
+    is FirWhenSubjectExpression -> whenRef.value.let { it.subjectVariable ?: it.subject }?.unwrapElement() ?: this
+    is FirSmartCastExpression -> originalExpression.unwrapElement()
+    is FirSafeCallExpression -> selector.unwrapElement()
+    is FirCheckedSafeCallSubject -> originalReceiverRef.value.unwrapElement()
+    is FirCheckNotNullCall -> argument.unwrapElement()
+    else -> this
+}
+
 fun FirExpression.unwrapSmartcastExpression(): FirExpression =
     when (this) {
         is FirSmartCastExpression -> originalExpression

compiler/testData/diagnostics/testsWithStdLib/inference/kt38737.fir.kt

@@ -1,13 +0,0 @@
-private const val dateRangeStart: String = "2020-01-01"
-private const val dateRangeEnd: String = "2020-05-01"
-
-private fun String?.toIconList(): List<String> = when (this) {
-    null -> listOf("DATE_IS_NULL")
-    <!ARGUMENT_TYPE_MISMATCH!>in dateRangeStart..dateRangeEnd<!> -> emptyList()
-    else -> listOf("DATE_IS_OUT_OF_RANGE")
-}
-
-fun main() {
-    println("2019-12-31".toIconList())
-    println(null.toIconList())
-}

compiler/testData/diagnostics/testsWithStdLib/inference/kt38737.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 private const val dateRangeStart: String = "2020-01-01"
 private const val dateRangeEnd: String = "2020-05-01"
 

compiler/testData/ir/irText/expressions/whenWithSubjectVariable.fir.ir.txt

@@ -0,0 +1,47 @@
+FILE fqName:<root> fileName:/whenWithSubjectVariable.kt
+  FUN name:foo visibility:public modality:FINAL <> () returnType:kotlin.Any
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun foo (): kotlin.Any declared in <root>'
+        CONST Int type=kotlin.Int value=1
+  FUN name:test visibility:public modality:FINAL <> () returnType:kotlin.Int
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun test (): kotlin.Int declared in <root>'
+        BLOCK type=kotlin.Int origin=WHEN
+          VAR name:y type:kotlin.Any [val]
+            CALL 'public final fun foo (): kotlin.Any declared in <root>' type=kotlin.Any origin=null
+          WHEN type=kotlin.Int origin=WHEN
+            BRANCH
+              if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
+                arg0: GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+                arg1: CONST Int type=kotlin.Int value=42
+              then: CONST Int type=kotlin.Int value=1
+            BRANCH
+              if: TYPE_OP type=kotlin.Boolean origin=INSTANCEOF typeOperand=kotlin.String
+                GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: CALL 'public open fun <get-length> (): kotlin.Int declared in kotlin.String' type=kotlin.Int origin=GET_PROPERTY
+                $this: TYPE_OP type=kotlin.String origin=IMPLICIT_CAST typeOperand=kotlin.String
+                  GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+            BRANCH
+              if: TYPE_OP type=kotlin.Boolean origin=NOT_INSTANCEOF typeOperand=kotlin.Int
+                GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: CONST Int type=kotlin.Int value=2
+            BRANCH
+              if: CALL 'public open fun contains (value: kotlin.Int): kotlin.Boolean [operator] declared in kotlin.ranges.IntRange' type=kotlin.Boolean origin=IN
+                $this: CALL 'public final fun rangeTo (other: kotlin.Int): kotlin.ranges.IntRange [operator] declared in kotlin.Int' type=kotlin.ranges.IntRange origin=RANGE
+                  $this: CONST Int type=kotlin.Int value=0
+                  other: CONST Int type=kotlin.Int value=10
+                value: TYPE_OP type=kotlin.Int origin=IMPLICIT_CAST typeOperand=kotlin.Int
+                  GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: CONST Int type=kotlin.Int value=3
+            BRANCH
+              if: CALL 'public final fun not (): kotlin.Boolean [operator] declared in kotlin.Boolean' type=kotlin.Boolean origin=null
+                $this: CALL 'public open fun contains (value: kotlin.Int): kotlin.Boolean [operator] declared in kotlin.ranges.IntRange' type=kotlin.Boolean origin=IN
+                  $this: CALL 'public final fun rangeTo (other: kotlin.Int): kotlin.ranges.IntRange [operator] declared in kotlin.Int' type=kotlin.ranges.IntRange origin=RANGE
+                    $this: CONST Int type=kotlin.Int value=10
+                    other: CONST Int type=kotlin.Int value=20
+                  value: TYPE_OP type=kotlin.Int origin=IMPLICIT_CAST typeOperand=kotlin.Int
+                    GET_VAR 'val y: kotlin.Any [val] declared in <root>.test' type=kotlin.Any origin=null
+              then: CONST Int type=kotlin.Int value=4
+            BRANCH
+              if: CONST Boolean type=kotlin.Boolean value=true
+              then: CONST Int type=kotlin.Int value=-1

compiler/testData/ir/irText/expressions/whenWithSubjectVariable.fir.kt.txt

@@ -0,0 +1,17 @@
+fun foo(): Any {
+  return 1
+}
+
+fun test(): Int {
+  return { // BLOCK
+    val y: Any = foo()
+    when {
+      EQEQ(arg0 = y, arg1 = 42) -> 1
+      y is String -> y /*as String */.<get-length>()
+      y !is Int -> 2
+      0.rangeTo(other = 10).contains(value = y /*as Int */) -> 3
+      10.rangeTo(other = 20).contains(value = y /*as Int */).not() -> 4
+      else -> -1
+    }
+  }
+}

compiler/testData/ir/irText/expressions/whenWithSubjectVariable.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_K2: JVM_IR
 // !LANGUAGE: +VariableDeclarationInWhenSubject
 
 fun foo(): Any = 1
@@ -11,4 +10,4 @@ fun test() =
         in 0..10 -> 3
         !in 10..20 -> 4
         else -> -1
-    }
+    }