AA FIR: handle smartcasted arg when building call's argument mapping

^KTIJ-25112 Fixed

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

@@ -784,6 +784,18 @@ public class Fe10IdeNormalAnalysisSourceModuleResolveCallTestGenerated extends A
         runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastImplicitExtensionReceiver.kt");
     }
 
+    @Test
+    @TestMetadata("smartCastedArg.kt")
+    public void testSmartCastedArg() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastedArg.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastedNonNullArg.kt")
+    public void testSmartCastedNonNullArg() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastedNonNullArg.kt");
+    }
+
     @Test
     @TestMetadata("substituteOverrideResolution.kt")
     public void testSubstituteOverrideResolution() throws Exception {

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

@@ -1272,9 +1272,11 @@ internal class KtFirCallResolver(
     }
 
     private fun FirExpression.findSourceKtExpressionForCallArgument(): KtExpression? {
+        // For smart-casted expression, refer to the source of the original expression
         // For spread, named, and lambda arguments, the source is the KtValueArgument.
         // For other arguments (including array indices), the source is the KtExpression.
         return when (this) {
+            is FirSmartCastExpression -> originalExpression.realPsi as? KtExpression
             is FirNamedArgumentExpression, is FirSpreadArgumentExpression, is FirLambdaArgumentExpression ->
                 realPsi.safeAs<KtValueArgument>()?.getArgumentExpression()
             else -> realPsi as? KtExpression

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

@@ -784,6 +784,18 @@ public class FirIdeNormalAnalysisSourceModuleResolveCallTestGenerated extends Ab
         runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastImplicitExtensionReceiver.kt");
     }
 
+    @Test
+    @TestMetadata("smartCastedArg.kt")
+    public void testSmartCastedArg() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastedArg.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastedNonNullArg.kt")
+    public void testSmartCastedNonNullArg() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastedNonNullArg.kt");
+    }
+
     @Test
     @TestMetadata("substituteOverrideResolution.kt")
     public void testSubstituteOverrideResolution() throws Exception {

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

@@ -784,6 +784,18 @@ public class FirStandaloneNormalAnalysisSourceModuleResolveCallTestGenerated ext
         runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastImplicitExtensionReceiver.kt");
     }
 
+    @Test
+    @TestMetadata("smartCastedArg.kt")
+    public void testSmartCastedArg() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastedArg.kt");
+    }
+
+    @Test
+    @TestMetadata("smartCastedNonNullArg.kt")
+    public void testSmartCastedNonNullArg() throws Exception {
+        runTest("analysis/analysis-api/testData/components/callResolver/resolveCall/smartCastedNonNullArg.kt");
+    }
+
     @Test
     @TestMetadata("substituteOverrideResolution.kt")
     public void testSubstituteOverrideResolution() throws Exception {

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

@@ -0,0 +1,10 @@
+open class A
+class B : A()
+
+private fun processB(b: B): Int = 2
+
+fun test(a: A) {
+    if (a is B) {
+        <expr>processB(a)</expr>
+    }
+}

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

@@ -0,0 +1,28 @@
+KtSuccessCallInfo:
+  call = KtSimpleFunctionCall:
+    isImplicitInvoke = false
+    partiallyAppliedSymbol = KtPartiallyAppliedSymbol:
+      dispatchReceiver = null
+      extensionReceiver = null
+      signature = KtFunctionLikeSignature:
+        receiverType = null
+        returnType = kotlin.Int
+        symbol = /processB(b: B): kotlin.Int
+        valueParameters = [
+          KtVariableLikeSignature:
+            name = b
+            receiverType = null
+            returnType = B
+            symbol = b: B
+            callableIdIfNonLocal = null
+        ]
+        callableIdIfNonLocal = /processB
+    typeArgumentsMapping = {}
+    argumentMapping = {
+      a -> (KtVariableLikeSignature:
+        name = b
+        receiverType = null
+        returnType = B
+        symbol = b: B
+        callableIdIfNonLocal = null)
+    }

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

@@ -0,0 +1,8 @@
+private fun test(s: String?) {
+    if (s != null) {
+        <expr>foo(s)</expr>
+    }
+}
+
+private fun foo(s: String) {
+}

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

@@ -0,0 +1,28 @@
+KtSuccessCallInfo:
+  call = KtSimpleFunctionCall:
+    isImplicitInvoke = false
+    partiallyAppliedSymbol = KtPartiallyAppliedSymbol:
+      dispatchReceiver = null
+      extensionReceiver = null
+      signature = KtFunctionLikeSignature:
+        receiverType = null
+        returnType = kotlin.Unit
+        symbol = /foo(s: kotlin.String): kotlin.Unit
+        valueParameters = [
+          KtVariableLikeSignature:
+            name = s
+            receiverType = null
+            returnType = kotlin.String
+            symbol = s: kotlin.String
+            callableIdIfNonLocal = null
+        ]
+        callableIdIfNonLocal = /foo
+    typeArgumentsMapping = {}
+    argumentMapping = {
+      s -> (KtVariableLikeSignature:
+        name = s
+        receiverType = null
+        returnType = kotlin.String
+        symbol = s: kotlin.String
+        callableIdIfNonLocal = null)
+    }