diff --git a/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/RefinedTypeForDataFlowTypeAttribute.kt b/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/RefinedTypeForDataFlowTypeAttribute.kt
new file mode 100644
index 0000000..a7d30fb
--- /dev/null
+++ b/compiler/fir/cones/src/org/jetbrains/kotlin/fir/types/RefinedTypeForDataFlowTypeAttribute.kt
@@ -0,0 +1,38 @@
+/*
+ * Copyright 2010-2025 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.types
+
+import kotlin.reflect.KClass
+
+/**
+ * [RefinedTypeForDataFlowTypeAttribute] stores the abbreviated type ([coneType]) of its owning expanded type. In the compiler, it is used exclusively
+ * for rendering the abbreviated typed in place of the expanded type. The Analysis API uses this attribute for additional purposes such as
+ * navigation to the type alias.
+ *
+ * The abbreviated type may not always be resolvable from a use-site session. For example, the owning expanded type may come from a library
+ * `L2` with a dependency on another library `L1` containing the type alias declaration. If the use-site session depends on `L2` but not on
+ * `L1`, the abbreviated type won't be resolvable even if the expanded type is.
+ */
+class RefinedTypeForDataFlowTypeAttribute(
+    override val coneType: ConeKotlinType,
+) : ConeAttributeWithConeType<RefinedTypeForDataFlowTypeAttribute>() {
+    override fun union(other: RefinedTypeForDataFlowTypeAttribute?): RefinedTypeForDataFlowTypeAttribute? = null
+    override fun intersect(other: RefinedTypeForDataFlowTypeAttribute?): RefinedTypeForDataFlowTypeAttribute? = null
+    override fun add(other: RefinedTypeForDataFlowTypeAttribute?): RefinedTypeForDataFlowTypeAttribute = other ?: this
+    override fun isSubtypeOf(other: RefinedTypeForDataFlowTypeAttribute?): Boolean = true
+    override fun toString(): String = "{${coneType.renderForDebugging()}=}"
+    override fun copyWith(newType: ConeKotlinType): RefinedTypeForDataFlowTypeAttribute = RefinedTypeForDataFlowTypeAttribute(newType)
+
+    override val key: KClass<out RefinedTypeForDataFlowTypeAttribute>
+        get() = RefinedTypeForDataFlowTypeAttribute::class
+    override val keepInInferredDeclarationType: Boolean
+        get() = true
+}
+
+val ConeAttributes.refinedTypeForDataFlow: RefinedTypeForDataFlowTypeAttribute? by ConeAttributes.attributeAccessor<RefinedTypeForDataFlowTypeAttribute>()
+
+val ConeKotlinType.refinedTypeForDataFlowOrSelf: ConeKotlinType
+    get() = attributes.refinedTypeForDataFlow?.coneType ?: this
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt
index 02b697f..8685e68 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt
@@ -246,10 +246,6 @@
      * @See org.jetbrains.kotlin.types.expressions.ControlStructureTypingUtils.createKnownTypeParameterSubstitutorForSpecialCall
      */
     private fun Candidate.isSyntheticFunctionCallThatShouldUseEqualityConstraint(expectedType: ConeKotlinType): Boolean {
-        // If we're inside an assignment's RHS, we mustn't add an equality constraint because it might prevent smartcasts.
-        // Example: val x: String? = null; x = if (foo) "" else throw Exception()
-        if (components.context.isInsideAssignmentRhs) return false
-
         val symbol = symbol as? FirCallableSymbol ?: return false
         if (symbol.origin != FirDeclarationOrigin.Synthetic.FakeFunction ||
             expectedType.isUnitOrAnyWithArbitraryNullability() ||
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt
index e189c44..ce573d9 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt
@@ -57,6 +57,7 @@
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.TransformData
 import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.resolve.calls.NewCommonSuperTypeCalculator
 import org.jetbrains.kotlin.resolve.calls.inference.model.InferredEmptyIntersection
 import org.jetbrains.kotlin.resolve.calls.tower.ApplicabilityDetail
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
@@ -1137,22 +1138,52 @@
         data: ExpectedArgumentType?,
     ): FirStatement {
         return transformSyntheticCall(whenExpression, data).apply {
+            if (isProperlyExhaustive) {
+                extracted(this, branches.map { it.result.resultType })
+            }
             replaceReturnTypeIfNotExhaustive(session)
         }
     }
 
+    private fun extracted(expression: FirExpression, branchesTypes: List<ConeKotlinType>) {
+        val refinedTypeForDataFlow = with(NewCommonSuperTypeCalculator) {
+            with(session.typeContext) {
+                commonSuperType(branchesTypes)
+            }
+        } as ConeKotlinType
+
+        val currentType = expression.resultType
+        if (!refinedTypeForDataFlow.isUnitOrFlexibleUnit && !currentType.isUnitOrFlexibleUnit &&
+            !currentType.isSubtypeOf(refinedTypeForDataFlow, session)
+        ) {
+            expression.resultType = currentType.withAttributes(
+                currentType.attributes.add(RefinedTypeForDataFlowTypeAttribute(refinedTypeForDataFlow))
+            )
+        }
+    }
+
     override fun transformTryExpression(
         tryExpression: FirTryExpression,
         data: ExpectedArgumentType?,
     ): FirStatement {
-        return transformSyntheticCall(tryExpression, data)
+        return transformSyntheticCall(tryExpression, data).also { transformed ->
+            extracted(
+                transformed,
+                buildList {
+                    add(transformed.tryBlock.resultType)
+                    transformed.catches.mapTo(this) { it.block.resultType }
+                }
+            )
+        }
     }
 
     override fun transformCheckNotNullCall(
         checkNotNullCall: FirCheckNotNullCall,
         data: ExpectedArgumentType?,
     ): FirStatement {
-        return transformSyntheticCall(checkNotNullCall, data)
+        return transformSyntheticCall(checkNotNullCall, data).also {
+            extracted(it, listOf(it.argumentList.arguments[0].resultType.makeConeTypeDefinitelyNotNullOrNotNull(session.typeContext)))
+        }
     }
 
     override fun transformElvisExpression(
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
index 78287c3..1890119 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
@@ -84,20 +84,6 @@
     @set:PrivateForInline
     var inferenceSession: FirInferenceSession = FirInferenceSession.DEFAULT
 
-    @set:PrivateForInline
-    var isInsideAssignmentRhs: Boolean = false
-
-    @OptIn(PrivateForInline::class)
-    inline fun <R> withAssignmentRhs(block: () -> R): R {
-        val oldMode = this.isInsideAssignmentRhs
-        this.isInsideAssignmentRhs = true
-        return try {
-            block()
-        } finally {
-            this.isInsideAssignmentRhs = oldMode
-        }
-    }
-
     @OptIn(PrivateForInline::class)
     inline fun withClassHeader(clazz: FirRegularClass, action: () -> Unit) {
         withSwitchedTowerDataModeForStaticNestedClass(clazz) {
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
index 33da4ce..9005c2f 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt
@@ -1300,14 +1300,12 @@
             }
         }
 
-        val result = context.withAssignmentRhs {
-            variableAssignment.transformRValue(
-                transformer,
-                withExpectedType(
-                    variableAssignment.lValue.resolvedType.toFirResolvedTypeRef(),
-                ),
-            )
-        }
+        val result = variableAssignment.transformRValue(
+            transformer,
+            withExpectedType(
+                variableAssignment.lValue.resolvedType.toFirResolvedTypeRef(),
+            ),
+        )
 
         // for cases like
         // buildSomething { tVar = "" // Should infer TV from String assignment }
diff --git a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/model.kt b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/model.kt
index 5e64eb3..ca0da559 100644
--- a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/model.kt
+++ b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/model.kt
@@ -10,6 +10,7 @@
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.types.ConeErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.refinedTypeForDataFlowOrSelf
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract
 
@@ -53,7 +54,7 @@
     MutableTypeStatement(this, lowerTypes = linkedSetOf(DfaType.BooleanLiteral(boolean)))
 
 infix fun RealVariable.typeEq(type: ConeKotlinType): MutableTypeStatement =
-    MutableTypeStatement(this, if (type is ConeErrorType) linkedSetOf() else linkedSetOf(type))
+    MutableTypeStatement(this, if (type is ConeErrorType) linkedSetOf() else linkedSetOf(type.refinedTypeForDataFlowOrSelf))
 
 infix fun RealVariable.typeNotEq(type: ConeKotlinType): MutableTypeStatement =
     MutableTypeStatement(this, lowerTypes = if (type is ConeErrorType) linkedSetOf() else linkedSetOf(DfaType.Cone(type)))
diff --git a/compiler/testData/codegen/box/inference/ifWithAssignmentAndNothingBranch.kt b/compiler/testData/codegen/box/inference/ifWithAssignmentAndNothingBranch.kt
index 16d8552..59f03b9 100644
--- a/compiler/testData/codegen/box/inference/ifWithAssignmentAndNothingBranch.kt
+++ b/compiler/testData/codegen/box/inference/ifWithAssignmentAndNothingBranch.kt
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_K2: ANY
 // ISSUE: KT-78127
 
 fun <T : Any> materialize(): T {
@@ -7,8 +6,6 @@
 
 var b = true
 
-var s: Nothing? = null
-
 fun box(): String {
     var x: String? = null
 
diff --git a/compiler/testData/codegen/box/inference/tryCatchAtAssignment.kt b/compiler/testData/codegen/box/inference/tryCatchAtAssignment.kt
index 5420a45..b4caa9d 100644
--- a/compiler/testData/codegen/box/inference/tryCatchAtAssignment.kt
+++ b/compiler/testData/codegen/box/inference/tryCatchAtAssignment.kt
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_K2: ANY
 // ISSUE: KT-77008
 
 interface I {
diff --git a/compiler/testData/codegen/box/inference/whenAtAssignment.kt b/compiler/testData/codegen/box/inference/whenAtAssignment.kt
index 3a77239..e1df20ab 100644
--- a/compiler/testData/codegen/box/inference/whenAtAssignment.kt
+++ b/compiler/testData/codegen/box/inference/whenAtAssignment.kt
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_K2: ANY
 // ISSUE: KT-77008
 
 interface I {