FIR IDE: Add simple shortening for qualified calls and properties

idea/idea-fir/tests/org/jetbrains/kotlin/shortenRefs/FirShortenRefsTestGenerated.java

@@ -29,6 +29,64 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/shortenRefsFir"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
     }
 
+    @TestMetadata("idea/testData/shortenRefsFir/calls")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Calls extends AbstractFirShortenRefsTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTestWithMuting, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInCalls() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/shortenRefsFir/calls"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("classInSameFile.kt")
+        public void testClassInSameFile() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/classInSameFile.kt");
+        }
+
+        @TestMetadata("functionInSameFile.kt")
+        public void testFunctionInSameFile() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/functionInSameFile.kt");
+        }
+
+        @TestMetadata("functionInSameFile2.kt")
+        public void testFunctionInSameFile2() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/functionInSameFile2.kt");
+        }
+
+        @TestMetadata("propertyChainCall.kt")
+        public void testPropertyChainCall() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/propertyChainCall.kt");
+        }
+
+        @TestMetadata("propertyInSameFile.kt")
+        public void testPropertyInSameFile() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/propertyInSameFile.kt");
+        }
+
+        @TestMetadata("propertyInSameFile2.kt")
+        public void testPropertyInSameFile2() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/propertyInSameFile2.kt");
+        }
+
+        @TestMetadata("rootPackage.kt")
+        public void testRootPackage() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/rootPackage.kt");
+        }
+
+        @TestMetadata("rootPackageShortenFakeRootPackage.kt")
+        public void testRootPackageShortenFakeRootPackage() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/rootPackageShortenFakeRootPackage.kt");
+        }
+
+        @TestMetadata("selfReferencingFunction.kt")
+        public void testSelfReferencingFunction() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/selfReferencingFunction.kt");
+        }
+    }
+
     @TestMetadata("idea/testData/shortenRefsFir/types")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

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

@@ -12,7 +12,9 @@ import com.intellij.util.containers.addIfNotNull
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractStarImportingScope
@@ -20,7 +22,10 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirExplicitSimpleImportingScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirPackageMemberScope
 import org.jetbrains.kotlin.fir.scopes.processClassifiersByName
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.lowerBoundIfFlexible
@@ -51,15 +56,19 @@ internal class KtFirReferenceShortener(
         resolveFileToBodyResolve(file)
         val firFile = file.getOrBuildFirOfType<FirFile>(firResolveState)
 
-        val typesToImport = mutableListOf<FqName>()
+        val namesToImport = mutableListOf<FqName>()
-        val typesToShorten = mutableListOf<KtUserType>()
 
-        firFile.acceptChildren(TypesCollectingVisitor(typesToImport, typesToShorten))
+        val typesToShorten = mutableListOf<KtUserType>()
+        val callsToShorten = mutableListOf<KtDotQualifiedExpression>()
+
+        firFile.acceptChildren(TypesCollectingVisitor(namesToImport, typesToShorten))
+        firFile.acceptChildren(CallsCollectingVisitor(namesToImport, callsToShorten))
 
         return ShortenCommandImpl(
             file,
-            typesToImport.distinct(),
+            namesToImport.distinct(),
-            typesToShorten.distinct().map { it.createSmartPointer() }
+            typesToShorten.distinct().map { it.createSmartPointer() },
+            callsToShorten.distinct().map { it.createSmartPointer() }
         )
     }
 
@@ -79,6 +88,14 @@ internal class KtFirReferenceShortener(
         return null
     }
 
+    private fun findSingleFunctionInScopesByName(scopes: List<FirScope>, name: Name): FirNamedFunctionSymbol? {
+        return scopes.asSequence().mapNotNull { it.getSingleFunctionByName(name) }.singleOrNull()
+    }
+
+    private fun findSinglePropertyInScopesByName(scopes: List<FirScope>, name: Name): FirVariableSymbol<*>? {
+        return scopes.asSequence().mapNotNull { it.getSinglePropertyByName(name) }.singleOrNull()
+    }
+
     private fun resolveFileToBodyResolve(file: KtFile) {
         for (declaration in file.declarations) {
             declaration.getOrBuildFir(firResolveState) // temporary hack, resolves declaration to BODY_RESOLVE stage
@@ -98,12 +115,20 @@ internal class KtFirReferenceShortener(
     }
 
     @OptIn(ExperimentalStdlibApi::class)
-    private fun findScopesAtPosition(targetTypeReference: KtElement, newImports: List<FqName>): List<FirScope>? {
+    private fun FirScope.getSingleFunctionByName(name: Name): FirNamedFunctionSymbol? =
-        val towerDataContext = firResolveState.getTowerDataContextForElement(targetTypeReference) ?: return null
+        buildList { processFunctionsByName(name, this::add) }.singleOrNull()
+
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun FirScope.getSinglePropertyByName(name: Name): FirVariableSymbol<*>? =
+        buildList { processPropertiesByName(name, this::add) }.singleOrNull()
+
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun findScopesAtPosition(position: KtElement, newImports: List<FqName>): List<FirScope>? {
+        val towerDataContext = firResolveState.getTowerDataContextForElement(position) ?: return null
 
         val result = buildList<FirScope> {
             addAll(towerDataContext.nonLocalTowerDataElements.mapNotNull { it.scope })
-            addIfNotNull(createFakeImportingScope(targetTypeReference.project, newImports))
+            addIfNotNull(createFakeImportingScope(position.project, newImports))
             addAll(towerDataContext.localScopes)
         }
 
@@ -128,7 +153,7 @@ internal class KtFirReferenceShortener(
     }
 
     private inner class TypesCollectingVisitor(
-        private val typesToImport: MutableList<FqName>,
+        private val namesToImport: MutableList<FqName>,
         private val typesToShorten: MutableList<KtUserType>,
     ) : FirVisitorVoid() {
         override fun visitElement(element: FirElement) {
@@ -157,7 +182,7 @@ internal class KtFirReferenceShortener(
             val allClassIds = generateSequence(wholeClassifierId) { it.outerClassId }
             val allTypeElements = generateSequence(wholeTypeElement) { it.qualifier }
 
-            val positionScopes = findScopesAtPosition(wholeTypeElement, typesToImport) ?: return
+            val positionScopes = findScopesAtPosition(wholeTypeElement, namesToImport) ?: return
 
             for ((classId, typeElement) in allClassIds.zip(allTypeElements)) {
                 val firstFoundClass = findFirstClassifierInScopesByName(positionScopes, classId.shortClassName)?.classId
@@ -184,16 +209,59 @@ internal class KtFirReferenceShortener(
         }
 
         private fun addTypeToImportAndShorten(classFqName: FqName, mostTopLevelTypeElement: KtUserType) {
-            typesToImport.add(classFqName)
+            namesToImport.add(classFqName)
             typesToShorten.add(mostTopLevelTypeElement)
         }
     }
+
+    private inner class CallsCollectingVisitor(
+        private val namesToImport: List<FqName>,
+        private val callsToShorten: MutableList<KtDotQualifiedExpression>
+    ) : FirVisitorVoid() {
+        override fun visitElement(element: FirElement) {
+            element.acceptChildren(this)
+        }
+
+        override fun visitResolvedNamedReference(resolvedNamedReference: FirResolvedNamedReference) {
+            super.visitResolvedNamedReference(resolvedNamedReference)
+
+            val referenceExpression = resolvedNamedReference.psi as? KtNameReferenceExpression
+            val qualifiedProperty = referenceExpression?.parent as? KtDotQualifiedExpression ?: return
+
+            val propertyId = (resolvedNamedReference.resolvedSymbol as? FirCallableSymbol<*>)?.callableId ?: return
+
+            val scopes = findScopesAtPosition(qualifiedProperty, namesToImport) ?: return
+            val singleAvailableProperty = findSinglePropertyInScopesByName(scopes, propertyId.callableName)
+
+            if (singleAvailableProperty?.callableId == propertyId) {
+                callsToShorten.add(qualifiedProperty)
+            }
+        }
+
+        override fun visitFunctionCall(functionCall: FirFunctionCall) {
+            super.visitFunctionCall(functionCall)
+
+            val callExpression = functionCall.psi as? KtCallExpression ?: return
+            val qualifiedCallExpression = callExpression.parent as? KtDotQualifiedExpression ?: return
+
+            val resolvedNamedReference = functionCall.calleeReference as? FirResolvedNamedReference ?: return
+            val callableId = (resolvedNamedReference.resolvedSymbol as? FirCallableSymbol<*>)?.callableId ?: return
+
+            val scopes = findScopesAtPosition(callExpression, namesToImport) ?: return
+            val singleAvailableCallable = findSingleFunctionInScopesByName(scopes, callableId.callableName)
+
+            if (singleAvailableCallable?.callableId == callableId) {
+                callsToShorten.add(qualifiedCallExpression)
+            }
+        }
+    }
 }
 
 private class ShortenCommandImpl(
     val targetFile: KtFile,
     val importsToAdd: List<FqName>,
-    val typesToShorten: List<SmartPsiElementPointer<KtUserType>>
+    val typesToShorten: List<SmartPsiElementPointer<KtUserType>>,
+    val callsToShorten: List<SmartPsiElementPointer<KtDotQualifiedExpression>>,
 ) : ShortenCommand {
 
     override fun invokeShortening() {
@@ -207,8 +275,18 @@ private class ShortenCommandImpl(
             val type = typePointer.element ?: continue
             type.deleteQualifier()
         }
+
+        for (callPointer in callsToShorten) {
+            val call = callPointer.element ?: continue
+            call.deleteQualifier()
+        }
     }
 }
 
 private tailrec fun KtTypeElement?.unwrapNullable(): KtTypeElement? =
     if (this is KtNullableType) this.innerType.unwrapNullable() else this
+
+private fun KtDotQualifiedExpression.deleteQualifier(): KtExpression? {
+    val selectorExpression = selectorExpression ?: return null
+    return this.replace(selectorExpression) as KtExpression
+}

idea/testData/shortenRefsFir/calls/classInSameFile.kt

@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+class A
+
+fun usage(a: <selection>test.A</selection>) {}

idea/testData/shortenRefsFir/calls/classInSameFile.kt.after

@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+class A
+
+fun usage(a: A) {}

idea/testData/shortenRefsFir/calls/functionInSameFile.kt

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+fun foo() {}
+
+fun usage() {
+    <selection>test.foo()</selection>
+}

idea/testData/shortenRefsFir/calls/functionInSameFile.kt.after

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+fun foo() {}
+
+fun usage() {
+    foo()
+}

idea/testData/shortenRefsFir/calls/functionInSameFile2.kt

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test.test2
+
+fun foo() {}
+
+fun usage() {
+    <selection>test.test2.foo()</selection>
+}

idea/testData/shortenRefsFir/calls/functionInSameFile2.kt.after

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test.test2
+
+fun foo() {}
+
+fun usage() {
+    foo()
+}

idea/testData/shortenRefsFir/calls/propertyChainCall.kt

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test.test2
+
+val foo = 1
+
+fun usage() {
+    <selection>test.test2.foo.toString()</selection>
+}

idea/testData/shortenRefsFir/calls/propertyChainCall.kt.after

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test.test2
+
+val foo = 1
+
+fun usage() {
+    foo.toString()
+}

idea/testData/shortenRefsFir/calls/propertyInSameFile.kt

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+val foo = 1
+
+fun usage() {
+    <selection>test.foo</selection>
+}

idea/testData/shortenRefsFir/calls/propertyInSameFile.kt.after

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+val foo = 1
+
+fun usage() {
+    foo
+}

idea/testData/shortenRefsFir/calls/propertyInSameFile2.kt

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test.test2
+
+val foo = 1
+
+fun usage() {
+    <selection>test.test2.foo</selection>
+}

idea/testData/shortenRefsFir/calls/propertyInSameFile2.kt.after

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test.test2
+
+val foo = 1
+
+fun usage() {
+    foo
+}

idea/testData/shortenRefsFir/calls/rootPackage.kt

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+fun test() {}
+
+fun usage() {
+    <selection>_root_ide_package_.test.test()</selection>
+}

idea/testData/shortenRefsFir/calls/rootPackage.kt.after

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+fun test() {}
+
+fun usage() {
+    test()
+}

idea/testData/shortenRefsFir/calls/rootPackageShortenFakeRootPackage.kt

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+fun test() {}
+
+fun usage() {
+    fun test() {}
+
+    <selection>_root_ide_package_.test.test()</selection>
+}

idea/testData/shortenRefsFir/calls/rootPackageShortenFakeRootPackage.kt.after

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+fun test() {}
+
+fun usage() {
+    fun test() {}
+
+    test.test()
+}

idea/testData/shortenRefsFir/calls/selfReferencingFunction.kt

@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test.test1
+
+fun foo() {
+    <selection>test.test1.foo()</selection>
+}

idea/testData/shortenRefsFir/calls/selfReferencingFunction.kt.after

@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test.test1
+
+fun foo() {
+    foo()
+}