FIR IDE: Implement simple importing of the functions

This is not a complete algorithm, but it already works in many cases

Disable some tests that not yet work

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -76,6 +76,36 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
             runTest("idea/testData/shortenRefsFir/calls/functionInSameFile2.kt");
         }
 
+        @TestMetadata("functionInSameFileAmbiguous.kt")
+        public void testFunctionInSameFileAmbiguous() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/functionInSameFileAmbiguous.kt");
+        }
+
+        @TestMetadata("notImportedTopLevelFunctionAmbiguous.kt")
+        public void testNotImportedTopLevelFunctionAmbiguous() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionAmbiguous.kt");
+        }
+
+        @TestMetadata("notImportedTopLevelFunctionConflictsWithImported.kt")
+        public void testNotImportedTopLevelFunctionConflictsWithImported() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionConflictsWithImported.kt");
+        }
+
+        @TestMetadata("notImportedTopLevelFunctionMissingArg.kt")
+        public void testNotImportedTopLevelFunctionMissingArg() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionMissingArg.kt");
+        }
+
+        @TestMetadata("notImportedTopLevelFunctionNoArgs.kt")
+        public void testNotImportedTopLevelFunctionNoArgs() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionNoArgs.kt");
+        }
+
+        @TestMetadata("notImportedTopLevelTypeConstructorNoArgs.kt")
+        public void testNotImportedTopLevelTypeConstructorNoArgs() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/notImportedTopLevelTypeConstructorNoArgs.kt");
+        }
+
         @TestMetadata("propertyChainCall.kt")
         public void testPropertyChainCall() throws Exception {
             runTest("idea/testData/shortenRefsFir/calls/propertyChainCall.kt");

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

@@ -17,18 +17,20 @@ import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
+import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeAmbiguityError
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.getFunctions
 import org.jetbrains.kotlin.fir.scopes.impl.FirAbstractStarImportingScope
 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.CallableId
 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.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.lowerBoundIfFlexible
@@ -91,8 +93,8 @@ internal class KtFirReferenceShortener(
         return null
     }
 
-    private fun findSingleFunctionInScopesByName(scopes: List<FirScope>, name: Name): FirNamedFunctionSymbol? {
-        return scopes.asSequence().mapNotNull { it.getSingleFunctionByName(name) }.singleOrNull()
+    private fun findFunctionsInScopes(scopes: List<FirScope>, name: Name): List<FirNamedFunctionSymbol> {
+        return scopes.flatMap { it.getFunctions(name) }
     }
 
     private fun findSinglePropertyInScopesByName(scopes: List<FirScope>, name: Name): FirVariableSymbol<*>? {
@@ -252,13 +254,16 @@ internal class KtFirReferenceShortener(
             val callExpression = functionCall.psi as? KtCallExpression ?: return
             val qualifiedCallExpression = callExpression.getDotQualifiedExpressionForSelector() ?: return
 
-            val resolvedNamedReference = functionCall.calleeReference as? FirResolvedNamedReference ?: return
-            val callableId = (resolvedNamedReference.resolvedSymbol as? FirCallableSymbol<*>)?.callableId ?: return
+            val calleeReference = functionCall.calleeReference
+            val callableId = findUnambiguousReferencedCallableId(calleeReference) ?: return
 
             val scopes = findScopesAtPosition(callExpression, namesToImport) ?: return
-            val singleAvailableCallable = findSingleFunctionInScopesByName(scopes, callableId.callableName)
+            val availableCallables = findFunctionsInScopes(scopes, callableId.callableName)
 
-            if (singleAvailableCallable?.callableId == callableId) {
+            if (availableCallables.isEmpty()) {
+                val additionalImport = callableId.asImportableFqName() ?: return
+                addElementToImportAndShorten(additionalImport, qualifiedCallExpression)
+            } else if (availableCallables.all { it.callableId == callableId }) {
                 addElementToShorten(qualifiedCallExpression)
             }
         }
@@ -274,6 +279,38 @@ internal class KtFirReferenceShortener(
             return receiverType.classKind != ClassKind.OBJECT
         }
 
+        private fun findUnambiguousReferencedCallableId(namedReference: FirNamedReference): CallableId? {
+            val unambiguousSymbol = when (namedReference) {
+                is FirResolvedNamedReference -> namedReference.resolvedSymbol
+                is FirErrorNamedReference -> {
+                    val candidateSymbol = namedReference.candidateSymbol
+                    if (candidateSymbol !is FirErrorFunctionSymbol) {
+                        candidateSymbol
+                    } else {
+                        getSingleUnambiguousCandidate(namedReference)
+                    }
+                }
+                else -> null
+            }
+
+            return (unambiguousSymbol as? FirCallableSymbol<*>)?.callableId
+        }
+
+        /**
+         * If [namedReference] is ambiguous and all candidates point to the callables with same callableId,
+         * returns the first candidate; otherwise returns null.
+         */
+        private fun getSingleUnambiguousCandidate(namedReference: FirErrorNamedReference): FirCallableSymbol<*>? {
+            val coneAmbiguityError = namedReference.diagnostic as? ConeAmbiguityError ?: return null
+
+            val candidates = coneAmbiguityError.candidates.map { it as FirCallableSymbol<*> }
+            require(candidates.isNotEmpty()) { "Cannot have zero candidates" }
+
+            val distinctCandidates = candidates.distinctBy { it.callableId }
+            return distinctCandidates.singleOrNull()
+                ?: error("Expected all candidates to have same callableId, but got: ${distinctCandidates.map { it.callableId }}")
+        }
+
         override fun visitResolvedQualifier(resolvedQualifier: FirResolvedQualifier) {
             super.visitResolvedQualifier(resolvedQualifier)
 
@@ -358,6 +395,8 @@ private class ShortenCommandImpl(
     }
 }
 
+private fun CallableId.asImportableFqName(): FqName? = if (classId == null) packageName.child(callableName) else null
+
 private fun KtElement.getDotQualifiedExpressionForSelector(): KtDotQualifiedExpression? =
     getQualifiedExpressionForSelector() as? KtDotQualifiedExpression
 

idea/testData/shortenRefsFir/calls/functionInSameFileAmbiguous.kt

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+fun foo(i: Int) {}
+
+fun foo(s: String) {}
+
+fun usage() {
+    <selection>test.foo()</selection>
+}

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

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+fun foo(i: Int) {}
+
+fun foo(s: String) {}
+
+fun usage() {
+    foo()
+}

idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionAmbiguous.dependency.kt

@@ -0,0 +1,5 @@
+package dependency
+
+fun foo(s: String) {}
+
+fun foo(i: Int) {}

idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionAmbiguous.kt

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

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

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

idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionConflictsWithImported.dependency.kt

@@ -0,0 +1,3 @@
+package dependency
+
+fun foo() {}

idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionConflictsWithImported.kt

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

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

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

idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionMissingArg.dependency.kt

@@ -0,0 +1,3 @@
+package dependency
+
+fun foo(a: Any) {}

idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionMissingArg.kt

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

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

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

idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionNoArgs.dependency.kt

@@ -0,0 +1,3 @@
+package dependency
+
+fun foo() {}

idea/testData/shortenRefsFir/calls/notImportedTopLevelFunctionNoArgs.kt

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

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

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

idea/testData/shortenRefsFir/calls/notImportedTopLevelTypeConstructorNoArgs.dependency.kt

@@ -0,0 +1,3 @@
+package dependency
+
+class Foo

idea/testData/shortenRefsFir/calls/notImportedTopLevelTypeConstructorNoArgs.kt

@@ -0,0 +1,6 @@
+// FIR_IGNORE
+package test
+
+fun usage() {
+    <selection>dependency.Foo()</selection>
+}

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

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

idea/testData/shortenRefsFir/calls/rootPackageShortenFakeRootPackage.kt

@@ -1,4 +1,4 @@
-// FIR_COMPARISON
+// FIR_IGNORE
 package test
 
 fun test() {}

idea/testData/shortenRefsFir/types/ParameterTypeConflictingTopLevelClassUsed.kt

@@ -1,4 +1,4 @@
-// FIR_COMPARISON
+// FIR_IGNORE
 package test
 
 class Foo