KTIJ-27434 [AA] Unwrap safe call expression when resolving implicit invoke calls

`FirSafeCallExpression` does not implement `FirResolvable`, so we need
to unwrap it to correctly resolve calls like `foo?.bar()` when `bar` is
a lambda with receiver parameter or a functional type property

^KTIJ-27434 Fixed

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

@@ -1123,6 +1123,42 @@ public class Fe10IdeNormalAnalysisSourceModuleReferenceResolveTestGenerated exte
             runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaNoParRCurly.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_regularCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_safeCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_regularCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_safeCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_regularCall.kt")
+        public void testLambdaWithExtensionParameter_property_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_safeCall.kt")
+        public void testLambdaWithExtensionParameter_property_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_safeCall.kt");
+        }
+
         @Test
         @TestMetadata("noParams.kt")
         public void testNoParams() throws Exception {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirUtils.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.getAnnotationsByClassId
 import org.jetbrains.kotlin.fir.declarations.getStringArgument
@@ -25,6 +26,7 @@ import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
+import org.jetbrains.kotlin.fir.expressions.FirSafeCallExpression
 import org.jetbrains.kotlin.fir.expressions.arguments
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
@@ -130,3 +132,6 @@ fun FirAnnotationContainer.getJvmNameFromAnnotation(session: FirSession, target:
             ?.takeIf { target == null || call.useSiteTarget == target }
     }
 }
+
+internal fun FirElement.unwrapSafeCall(): FirElement =
+    (this as? FirSafeCallExpression)?.selector ?: this

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.analysis.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.analysis.api.fir.buildSymbol
 import org.jetbrains.kotlin.analysis.api.fir.getCandidateSymbols
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirPackageSymbol
+import org.jetbrains.kotlin.analysis.api.fir.unwrapSafeCall
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirSafe
@@ -322,9 +323,10 @@ internal object FirReferenceResolveHelper {
     ): Collection<KtSymbol> {
         val parentAsCall = expression.parent as? KtCallExpression
         if (parentAsCall != null) {
-            val firResolvable = parentAsCall.getOrBuildFirSafe<FirResolvable>(analysisSession.firResolveSession)
+            val firCall = parentAsCall.getOrBuildFir(analysisSession.firResolveSession)?.unwrapSafeCall()
-            if (firResolvable != null) {
+
-                return getSymbolsByResolvable(firResolvable, expression, session, symbolBuilder)
+            if (firCall is FirResolvable) {
+                return getSymbolsByResolvable(firCall, expression, session, symbolBuilder)
             }
         }
         return fir.toTargetSymbol(session, symbolBuilder)

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

@@ -1123,6 +1123,42 @@ public class FirIdeDependentAnalysisSourceModuleReferenceResolveTestGenerated ex
             runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaNoParRCurly.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_regularCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_safeCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_regularCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_safeCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_regularCall.kt")
+        public void testLambdaWithExtensionParameter_property_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_safeCall.kt")
+        public void testLambdaWithExtensionParameter_property_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_safeCall.kt");
+        }
+
         @Test
         @TestMetadata("noParams.kt")
         public void testNoParams() throws Exception {

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

@@ -1047,6 +1047,42 @@ public class FirIdeNormalAnalysisLibrarySourceModuleReferenceResolveTestGenerate
             runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaNoParRCurly.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_regularCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_safeCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_regularCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_safeCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_regularCall.kt")
+        public void testLambdaWithExtensionParameter_property_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_safeCall.kt")
+        public void testLambdaWithExtensionParameter_property_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_safeCall.kt");
+        }
+
         @Test
         @TestMetadata("noParams.kt")
         public void testNoParams() throws Exception {

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

@@ -1123,6 +1123,42 @@ public class FirIdeNormalAnalysisSourceModuleReferenceResolveTestGenerated exten
             runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaNoParRCurly.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_regularCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_safeCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_regularCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_safeCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_regularCall.kt")
+        public void testLambdaWithExtensionParameter_property_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_safeCall.kt")
+        public void testLambdaWithExtensionParameter_property_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_safeCall.kt");
+        }
+
         @Test
         @TestMetadata("noParams.kt")
         public void testNoParams() throws Exception {

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

@@ -1123,6 +1123,42 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceResolveTestGenerate
             runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaNoParRCurly.kt");
         }
 
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_regularCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_extensionProperty_safeCall.kt")
+        public void testLambdaWithExtensionParameter_extensionProperty_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_regularCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_parameter_safeCall.kt")
+        public void testLambdaWithExtensionParameter_parameter_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_safeCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_regularCall.kt")
+        public void testLambdaWithExtensionParameter_property_regularCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_regularCall.kt");
+        }
+
+        @Test
+        @TestMetadata("lambdaWithExtensionParameter_property_safeCall.kt")
+        public void testLambdaWithExtensionParameter_property_safeCall() throws Exception {
+            runTest("analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_safeCall.kt");
+        }
+
         @Test
         @TestMetadata("noParams.kt")
         public void testNoParams() throws Exception {

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_regularCall.kt

@@ -0,0 +1,9 @@
+package test
+
+class Foo
+
+val Foo.action: () -> Unit get() = {}
+
+fun usage(foo: Foo) {
+    foo.acti<caret>on()
+}

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_regularCall.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: (in test) val test.Foo.action: () -> kotlin.Unit

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_safeCall.kt

@@ -0,0 +1,9 @@
+package test
+
+class Foo
+
+val Foo.action: () -> Unit get() = {}
+
+fun usage(foo: Foo?) {
+    foo?.acti<caret>on()
+}

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_extensionProperty_safeCall.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: (in test) val test.Foo.action: () -> kotlin.Unit

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_regularCall.kt

@@ -0,0 +1,7 @@
+package test
+
+class Foo
+
+fun usage(foo: Foo, action: Foo.() -> Unit) {
+    foo.acti<caret>on()
+}

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_regularCall.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: action: test.Foo.() -> kotlin.Unit

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_safeCall.kt

@@ -0,0 +1,7 @@
+package test
+
+class Foo
+
+fun usage(foo: Foo?, action: Foo.() -> Unit) {
+    foo?.acti<caret>on()
+}

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_parameter_safeCall.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: action: test.Foo.() -> kotlin.Unit

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_regularCall.kt

@@ -0,0 +1,9 @@
+package test
+
+class Foo {
+    val action: () -> Unit = {}
+}
+
+fun usage(foo: Foo) {
+    foo.acti<caret>on()
+}

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_regularCall.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: (in test.Foo) val action: () -> kotlin.Unit

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_safeCall.kt

@@ -0,0 +1,9 @@
+package test
+
+class Foo {
+    val action: () -> Unit = {}
+}
+
+fun usage(foo: Foo?) {
+    foo?.acti<caret>on()
+}

analysis/analysis-api/testData/referenceResolve/invoke/lambdaWithExtensionParameter_property_safeCall.txt

@@ -0,0 +1,2 @@
+Resolved to:
+0: (in test.Foo) val action: () -> kotlin.Unit