KT-60957 [Analysis API] Properly handle source PSI for implicit invoke calls in KtFirReferenceShortener

For implicit invoke operator calls, there are two instances of
`FirResolvedNamedReference`. One of them references `invoke` function
and cannot be used to analyze the property access, and the second one
has `source == null`. Luckily, the parent of the second reference is
`FirPropertyAccessExpression`, which has the correct source

^KT-60957 Fixed

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

@@ -985,7 +985,7 @@ private class ElementsToShortenCollector(
     }
 
     private fun processPropertyAccess(firPropertyAccess: FirPropertyAccessExpression) {
-        val propertyReferenceExpression = (firPropertyAccess.psi as? KtDotQualifiedExpression)?.selectorExpression as? KtNameReferenceExpression ?: return
+        val propertyReferenceExpression = firPropertyAccess.correspondingNameReference ?: return
         if (!propertyReferenceExpression.textRange.intersects(selection)) return
 
         val qualifiedProperty = propertyReferenceExpression.getQualifiedElement() as? KtDotQualifiedExpression ?: return
@@ -1014,6 +1014,19 @@ private class ElementsToShortenCollector(
         )
     }
 
+    private val FirPropertyAccessExpression.correspondingNameReference: KtNameReferenceExpression?
+        get() {
+            val nameReference = when (val sourcePsi = psi) {
+                // usual `foo.bar.baz` case
+                is KtDotQualifiedExpression -> sourcePsi.selectorExpression
+
+                // short `foo` case, or implicit invoke call like `foo.bar.baz()`
+                else -> sourcePsi
+            }
+
+            return nameReference as? KtNameReferenceExpression
+        }
+
     private val FirPropertyAccessExpression.referencedSymbol: FirVariableSymbol<*>?
         get() = (calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirVariableSymbol<*>
 

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

@@ -460,6 +460,12 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusOperator.kt");
     }
 
+    @Test
+    @TestMetadata("variable_invokeOperator.kt")
+    public void testVariable_invokeOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variable_invokeOperator.kt");
+    }
+
     @Nested
     @TestMetadata("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses")
     @TestDataPath("$PROJECT_ROOT")

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

@@ -460,6 +460,12 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenera
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variableAssignment_plusOperator.kt");
     }
 
+    @Test
+    @TestMetadata("variable_invokeOperator.kt")
+    public void testVariable_invokeOperator() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variable_invokeOperator.kt");
+    }
+
     @Nested
     @TestMetadata("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses")
     @TestDataPath("$PROJECT_ROOT")

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variable_invokeOperator.kt

@@ -0,0 +1,15 @@
+// FILE: main.kt
+package test
+
+fun usage() {
+    <expr>dependency.foo()</expr>
+}
+
+// FILE: dependency.kt
+package dependency
+
+var foo: Foo = Foo()
+
+class Foo {
+    operator fun invoke() {}
+}

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/variable_invokeOperator.txt

@@ -0,0 +1,7 @@
+Before shortening: dependency.foo()
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[qualifier] dependency.foo()
+with SHORTEN_AND_STAR_IMPORT:
+[qualifier] dependency.foo()