FIR. Refactor smart-cast representation in FIR tree

Make smart-casts non-transparent expression without delegation
to underlying FirQualifiedAccessExpression, as children delegation in
fir tree has unclear semantics
Remove two different kinds of tree nodes for smart-casts

analysis/analysis-api-fe10/tests-gen/org/jetbrains/kotlin/analysis/api/fe10/test/cases/generated/cases/components/callResolver/Fe10IdeNormalAnalysisSourceModuleResolveCallTestGenerated.java

@@ -832,6 +832,12 @@ public class Fe10IdeNormalAnalysisSourceModuleResolveCallTestGenerated extends A
         runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/variableWithMemberInvoke.kt");
     }
 
+    @Test
+    @TestMetadata("whenSelectorSmartCast.kt")
+    public void testWhenSelectorSmartCast() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/whenSelectorSmartCast.kt");
+    }
+
     @Nested
     @TestMetadata("analysis/analysis-api/testData/components/callResolver/resolveCall/assignments")
     @TestDataPath("$PROJECT_ROOT")

analysis/analysis-api-fe10/tests-gen/org/jetbrains/kotlin/analysis/api/fe10/test/cases/generated/cases/components/smartCastProvider/Fe10IdeNormalAnalysisSourceModuleHLSmartCastInfoTestGenerated.java

@@ -46,6 +46,12 @@ public class Fe10IdeNormalAnalysisSourceModuleHLSmartCastInfoTestGenerated exten
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @Test
+    @TestMetadata("bothImplicitReceiversSmartCast.kt")
+    public void testBothImplicitReceiversSmartCast() throws Exception {
+        runTest("analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo/bothImplicitReceiversSmartCast.kt");
+    }
+
     @Test
     @TestMetadata("multiSmartcastAsReceiver_stable.kt")
     public void testMultiSmartcastAsReceiver_stable() throws Exception {

analysis/analysis-api-fe10/tests-gen/org/jetbrains/kotlin/analysis/api/fe10/test/cases/generated/cases/references/Fe10IdeNormalAnalysisSourceModuleReferenceResolveTestGenerated.java

@@ -460,6 +460,18 @@ public class Fe10IdeNormalAnalysisSourceModuleReferenceResolveTestGenerated exte
         runTest("analysis/analysis-api/testData/referenceResolve/SeveralOverrides.kt");
     }
 
+    @Test
+    @TestMetadata("smartCastExpression.kt")
+    public void testSmartCastExpression() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastExpression.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastInvokeReceiver.kt")
+    public void testSmartCastInvokeReceiver() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastInvokeReceiver.kt");
+    }
+
     @Test
     @TestMetadata("SuperTypePrimaryConstructor.kt")
     public void testSuperTypePrimaryConstructor() throws Exception {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt

@@ -201,6 +201,9 @@ internal class KtFirCallResolver(
                 resolveCalleeExpressionOfFunctionCall,
                 resolveFragmentOfCall
             )
+            is FirSmartCastExpression -> originalExpression.toKtCallInfo(
+                psi, resolveCalleeExpressionOfFunctionCall, resolveFragmentOfCall
+            )
             else -> null
         }
     }
@@ -465,7 +468,7 @@ internal class KtFirCallResolver(
                     isImplicitInvoke
                 )
             }
-            is FirExpressionWithSmartcast -> createKtCall(psi, fir.originalExpression, candidate, resolveFragmentOfCall)
+            is FirSmartCastExpression -> (fir.originalExpression as? FirResolvable)?.let { createKtCall(psi, it, candidate, resolveFragmentOfCall) }
             else -> null
         }
     }
@@ -707,7 +710,7 @@ internal class KtFirCallResolver(
     private fun FirExpression.toKtReceiverValue(): KtReceiverValue? {
         val psi = psi
         return when (this) {
-            is FirExpressionWithSmartcast -> {
+            is FirSmartCastExpression -> {
                 val result = originalExpression.toKtReceiverValue()
                 if (result != null && isStable) {
                     KtSmartCastedReceiverValue(result, smartcastType.coneType.asKtType())

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirExpressionTypeProvider.kt

@@ -231,11 +231,12 @@ internal class KtFirExpressionTypeProvider(
         }
 
         when (val fir = expression.getOrBuildFir(analysisSession.firResolveSession)) {
-            is FirExpressionWithSmartcastToNothing -> if (fir.isStable) {
-                return DefiniteNullability.DEFINITELY_NULL
-            }
-            is FirExpressionWithSmartcast -> if (fir.isStable && fir.isNotNullable()) {
-                return DefiniteNullability.DEFINITELY_NOT_NULL
+            is FirSmartCastExpression -> if (fir.isStable) {
+                if (fir.smartcastTypeWithoutNullableNothing != null) {
+                    return DefiniteNullability.DEFINITELY_NULL
+                } else if (fir.isNotNullable()) {
+                    return DefiniteNullability.DEFINITELY_NOT_NULL
+                }
             }
             is FirExpression -> if (fir.isNotNullable()) {
                 return DefiniteNullability.DEFINITELY_NOT_NULL

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSmartcastProvider.kt

@@ -12,10 +12,7 @@ import org.jetbrains.kotlin.analysis.api.components.KtSmartCastProvider
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
-import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
-import org.jetbrains.kotlin.fir.expressions.FirImplicitInvokeCall
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.fir.expressions.FirSafeCallExpression
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.fir.types.isStableSmartcast
@@ -45,15 +42,15 @@ internal class KtFirSmartcastProvider(
                     this is KtParenthesizedExpression
         }
 
-    private fun getMatchingFirExpressionWithSmartCast(expression: KtExpression): FirExpressionWithSmartcast? {
+    private fun getMatchingFirExpressionWithSmartCast(expression: KtExpression): FirSmartCastExpression? {
         if (!expression.isExplicitSmartCastInfoTarget) return null
 
         val possibleFunctionCall = expression.getPossiblyQualifiedCallExpressionForCallee() ?: expression
 
         return when (val firExpression = possibleFunctionCall.getOrBuildFir(analysisSession.firResolveSession)) {
-            is FirExpressionWithSmartcast -> firExpression
-            is FirSafeCallExpression -> firExpression.selector as? FirExpressionWithSmartcast
-            is FirImplicitInvokeCall -> firExpression.explicitReceiver as? FirExpressionWithSmartcast
+            is FirSmartCastExpression -> firExpression
+            is FirSafeCallExpression -> firExpression.selector as? FirSmartCastExpression
+            is FirImplicitInvokeCall -> firExpression.explicitReceiver as? FirSmartCastExpression
             else -> null
         }
     }
@@ -63,7 +60,7 @@ internal class KtFirSmartcastProvider(
         return getSmartCastedInfo(firSmartCastExpression)
     }
 
-    private fun getSmartCastedInfo(expression: FirExpressionWithSmartcast): KtSmartCastInfo? {
+    private fun getSmartCastedInfo(expression: FirSmartCastExpression): KtSmartCastInfo? {
         val type = expression.smartcastType.coneTypeSafe<ConeKotlinType>()?.asKtType() ?: return null
         return KtSmartCastInfo(type, expression.isStable, token)
     }
@@ -82,6 +79,7 @@ internal class KtFirSmartcastProvider(
         return when (val firExpression = wholeExpression.getOrBuildFir(analysisSession.firResolveSession)) {
             is FirQualifiedAccessExpression -> firExpression
             is FirSafeCallExpression -> firExpression.selector as? FirQualifiedAccessExpression
+            is FirSmartCastExpression -> firExpression.originalExpression as? FirQualifiedAccessExpression
             else -> null
         }
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirReferenceResolveHelper.kt

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.resolve.dfa.unwrapSmartcastExpression
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -188,7 +189,10 @@ internal object FirReferenceResolveHelper {
         if (expression.isSyntheticOperatorReference()) return emptyList()
         val symbolBuilder = analysisSession.firSymbolBuilder
         val adjustedResolutionExpression = adjustResolutionExpression(expression)
-        val fir = adjustedResolutionExpression.getOrBuildFir(analysisSession.firResolveSession)
+        val fir = when (val baseFir = adjustedResolutionExpression.getOrBuildFir(analysisSession.firResolveSession)) {
+            is FirSmartCastExpression -> baseFir.originalExpression
+            else -> baseFir
+        }
         val session = analysisSession.firResolveSession.useSiteFirSession
         return when (fir) {
             is FirResolvedTypeRef -> getSymbolsForResolvedTypeRef(fir, expression, session, symbolBuilder)
@@ -361,7 +365,7 @@ internal object FirReferenceResolveHelper {
             return listOfNotNull((fir.dispatchReceiver.typeRef as? FirResolvedTypeRef)?.toTargetSymbol(session, symbolBuilder))
         }
         val implicitInvokeReceiver = if (fir is FirImplicitInvokeCall) {
-            fir.explicitReceiver as? FirQualifiedAccessExpression
+            fir.explicitReceiver?.unwrapSmartcastExpression() as? FirQualifiedAccessExpression
         } else {
             null
         }

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/callResolver/FirIdeNormalAnalysisSourceModuleResolveCallTestGenerated.java

@@ -832,6 +832,12 @@ public class FirIdeNormalAnalysisSourceModuleResolveCallTestGenerated extends Ab
         runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/variableWithMemberInvoke.kt");
     }
 
+    @Test
+    @TestMetadata("whenSelectorSmartCast.kt")
+    public void testWhenSelectorSmartCast() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/whenSelectorSmartCast.kt");
+    }
+
     @Nested
     @TestMetadata("analysis/analysis-api/testData/components/callResolver/resolveCall/assignments")
     @TestDataPath("$PROJECT_ROOT")

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/smartCastProvider/FirIdeDependentAnalysisSourceModuleHLSmartCastInfoTestGenerated.java

@@ -46,6 +46,12 @@ public class FirIdeDependentAnalysisSourceModuleHLSmartCastInfoTestGenerated ext
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @Test
+    @TestMetadata("bothImplicitReceiversSmartCast.kt")
+    public void testBothImplicitReceiversSmartCast() throws Exception {
+        runTest("analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo/bothImplicitReceiversSmartCast.kt");
+    }
+
     @Test
     @TestMetadata("multiSmartcastAsReceiver_stable.kt")
     public void testMultiSmartcastAsReceiver_stable() throws Exception {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/smartCastProvider/FirIdeNormalAnalysisSourceModuleHLSmartCastInfoTestGenerated.java

@@ -46,6 +46,12 @@ public class FirIdeNormalAnalysisSourceModuleHLSmartCastInfoTestGenerated extend
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @Test
+    @TestMetadata("bothImplicitReceiversSmartCast.kt")
+    public void testBothImplicitReceiversSmartCast() throws Exception {
+        runTest("analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo/bothImplicitReceiversSmartCast.kt");
+    }
+
     @Test
     @TestMetadata("multiSmartcastAsReceiver_stable.kt")
     public void testMultiSmartcastAsReceiver_stable() throws Exception {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeDependentAnalysisSourceModuleReferenceResolveTestGenerated.java

@@ -460,6 +460,18 @@ public class FirIdeDependentAnalysisSourceModuleReferenceResolveTestGenerated ex
         runTest("analysis/analysis-api/testData/referenceResolve/SeveralOverrides.kt");
     }
 
+    @Test
+    @TestMetadata("smartCastExpression.kt")
+    public void testSmartCastExpression() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastExpression.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastInvokeReceiver.kt")
+    public void testSmartCastInvokeReceiver() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastInvokeReceiver.kt");
+    }
+
     @Test
     @TestMetadata("SuperTypePrimaryConstructor.kt")
     public void testSuperTypePrimaryConstructor() throws Exception {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeNormalAnalysisLibrarySourceModuleReferenceResolveTestGenerated.java

@@ -460,6 +460,18 @@ public class FirIdeNormalAnalysisLibrarySourceModuleReferenceResolveTestGenerate
         runTest("analysis/analysis-api/testData/referenceResolve/SeveralOverrides.kt");
     }
 
+    @Test
+    @TestMetadata("smartCastExpression.kt")
+    public void testSmartCastExpression() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastExpression.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastInvokeReceiver.kt")
+    public void testSmartCastInvokeReceiver() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastInvokeReceiver.kt");
+    }
+
     @Test
     @TestMetadata("SuperTypePrimaryConstructor.kt")
     public void testSuperTypePrimaryConstructor() throws Exception {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeNormalAnalysisSourceModuleReferenceResolveTestGenerated.java

@@ -460,6 +460,18 @@ public class FirIdeNormalAnalysisSourceModuleReferenceResolveTestGenerated exten
         runTest("analysis/analysis-api/testData/referenceResolve/SeveralOverrides.kt");
     }
 
+    @Test
+    @TestMetadata("smartCastExpression.kt")
+    public void testSmartCastExpression() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastExpression.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastInvokeReceiver.kt")
+    public void testSmartCastInvokeReceiver() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastInvokeReceiver.kt");
+    }
+
     @Test
     @TestMetadata("SuperTypePrimaryConstructor.kt")
     public void testSuperTypePrimaryConstructor() throws Exception {

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/smartCastProvider/AbstractHLSmartCastInfoTest.kt

@@ -24,6 +24,12 @@ abstract class AbstractHLSmartCastInfoTest : AbstractAnalysisApiSingleFileTest()
                     appendLine("expression: ${expression.text}")
                     appendLine("isStable: ${smartCastInfo?.isStable}")
                     appendLine("smartCastType: ${smartCastInfo?.smartCastType?.render()}")
+
+                    val receiverSmartCasts = expression.getImplicitReceiverSmartCast()
+                    for (receiverSmartCast in receiverSmartCasts) {
+                        appendLine("receiver: ${receiverSmartCast.kind}")
+                        appendLine("    smartCastType: ${receiverSmartCast.type.render()}")
+                    }
                 }
             }
         }

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/components/callResolver/FirStandaloneNormalAnalysisSourceModuleResolveCallTestGenerated.java

@@ -832,6 +832,12 @@ public class FirStandaloneNormalAnalysisSourceModuleResolveCallTestGenerated ext
         runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/variableWithMemberInvoke.kt");
     }
 
+    @Test
+    @TestMetadata("whenSelectorSmartCast.kt")
+    public void testWhenSelectorSmartCast() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/whenSelectorSmartCast.kt");
+    }
+
     @Nested
     @TestMetadata("analysis/analysis-api/testData/components/callResolver/resolveCall/assignments")
     @TestDataPath("$PROJECT_ROOT")

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/components/smartCastProvider/FirStandaloneNormalAnalysisSourceModuleHLSmartCastInfoTestGenerated.java

@@ -46,6 +46,12 @@ public class FirStandaloneNormalAnalysisSourceModuleHLSmartCastInfoTestGenerated
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @Test
+    @TestMetadata("bothImplicitReceiversSmartCast.kt")
+    public void testBothImplicitReceiversSmartCast() throws Exception {
+        runTest("analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo/bothImplicitReceiversSmartCast.kt");
+    }
+
     @Test
     @TestMetadata("multiSmartcastAsReceiver_stable.kt")
     public void testMultiSmartcastAsReceiver_stable() throws Exception {

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/references/FirStandaloneNormalAnalysisSourceModuleReferenceResolveTestGenerated.java

@@ -460,6 +460,18 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceResolveTestGenerate
         runTest("analysis/analysis-api/testData/referenceResolve/SeveralOverrides.kt");
     }
 
+    @Test
+    @TestMetadata("smartCastExpression.kt")
+    public void testSmartCastExpression() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastExpression.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastInvokeReceiver.kt")
+    public void testSmartCastInvokeReceiver() throws Exception {
+        runTest("analysis/analysis-api/testData/referenceResolve/smartCastInvokeReceiver.kt");
+    }
+
     @Test
     @TestMetadata("SuperTypePrimaryConstructor.kt")
     public void testSuperTypePrimaryConstructor() throws Exception {

analysis/analysis-api/testData/components/callResolver/resolveCall/whenSelectorSmartCast.kt

@@ -0,0 +1,7 @@
+fun test(obj: Any): String {
+    return when {
+        obj !is Iterable<*> -> "not iterable"
+        <expr>obj</expr> !is Collection<*> -> "not collection"
+        else -> "unknown"
+    }
+}

analysis/analysis-api/testData/components/callResolver/resolveCall/whenSelectorSmartCast.txt

@@ -0,0 +1,14 @@
+KtSuccessCallInfo:
+  call = KtSimpleVariableAccessCall:
+    partiallyAppliedSymbol = KtPartiallyAppliedSymbol:
+      dispatchReceiver = null
+      extensionReceiver = null
+      signature = KtVariableLikeSignature:
+        name = obj
+        receiverType = null
+        returnType = kotlin.Any
+        symbol = obj: kotlin.Any
+        callableIdIfNonLocal = null
+    simpleAccess = Read:
+
+    typeArgumentsMapping = {}

analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo/bothImplicitReceiversSmartCast.kt

@@ -0,0 +1,18 @@
+interface Foo {
+    fun Bar.foo()
+}
+
+interface Bar {}
+
+inline fun <T, R> myWith(argument: T, lambda: T.() -> R): R {
+    argument.lambda()
+}
+
+fun Any.action(other: Any) {
+    if (this is Foo) {
+        with(other) {
+            this as Bar
+            <expr>foo</expr>()
+        }
+    }
+}

analysis/analysis-api/testData/components/smartCastProvider/smartCastInfo/bothImplicitReceiversSmartCast.txt

@@ -0,0 +1,7 @@
+expression: foo
+isStable: null
+smartCastType: null
+receiver: DISPATCH
+    smartCastType: Foo
+receiver: EXTENSION
+    smartCastType: Bar

analysis/analysis-api/testData/referenceResolve/smartCastExpression.kt

@@ -0,0 +1,9 @@
+interface Foo {
+    fun function()
+}
+
+fun foo(parameter: Any) {
+    if (parameter is Foo) {
+        <caret>parameter.function()
+    }
+}

analysis/analysis-api/testData/referenceResolve/smartCastExpression.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: parameter: kotlin.Any

analysis/analysis-api/testData/referenceResolve/smartCastInvokeReceiver.kt

@@ -0,0 +1,9 @@
+interface WithInvoke {
+    operator fun invoke() {}
+}
+
+fun foo(parameter: Any) {
+    if (parameter is WithInvoke) {
+        <caret>parameter()
+    }
+}

analysis/analysis-api/testData/referenceResolve/smartCastInvokeReceiver.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: parameter: kotlin.Any

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FirElementsRecorder.kt

@@ -9,10 +9,7 @@ import org.jetbrains.kotlin.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.DuplicatedFirSourceElementsException
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isErrorElement
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.expressions.FirConstExpression
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
-import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
-import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
+import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildConstExpression
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
@@ -115,6 +112,7 @@ internal open class FirElementsRecorder : FirVisitor<Unit, MutableMap<KtElement,
                         it.kind == KtFakeSourceElementKind.ImplicitConstructor ||
                         it.kind == KtFakeSourceElementKind.DesugaredPrefixNameReference ||
                         it.kind == KtFakeSourceElementKind.DesugaredPostfixNameReference ||
+                        it.kind == KtFakeSourceElementKind.SmartCastExpression ||
                         it.isSourceForCompoundAccess(element)
             }.psi as? KtElement
             ?: return

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerFirTestdataTestGenerated.java

@@ -4111,6 +4111,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
                 public void testSafeCalls() throws Exception {
                     runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/safeCalls.kt");
                 }
+
+                @Test
+                @TestMetadata("unstableSmartCastOnSafeCallArgument.kt")
+                public void testUnstableSmartCastOnSafeCallArgument() throws Exception {
+                    runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/unstableSmartCastOnSafeCallArgument.kt");
+                }
             }
 
             @Nested