Extension callables index is built correctly when using alias imports

idea/idea-analysis/src/org/jetbrains/kotlin/idea/stubindex/IndexUtils.kt

@@ -16,6 +16,9 @@
 
 package org.jetbrains.kotlin.idea.stubindex
 
+import com.google.common.collect.HashMultimap
+import com.google.common.collect.Multimap
+import com.intellij.openapi.util.Key
 import com.intellij.psi.stubs.IndexSink
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.stubs.KotlinCallableStubBase
@@ -36,22 +39,25 @@ private fun <TDeclaration : JetCallableDeclaration> JetTypeElement.index(declara
 
     when (this) {
         is JetUserType -> {
-            //TODO: aliases
-            var typeName = getReferencedName() ?: return
+            var referenceName = getReferencedName() ?: return
 
             if (declaration is JetNamedFunction) {
-                val typeParameter = declaration.getTypeParameters().firstOrNull { it.getName() == typeName }
+                val typeParameter = declaration.getTypeParameters().firstOrNull { it.getName() == referenceName }
                 if (typeParameter != null) {
                     val bound = typeParameter.getExtendsBound()
                     if (bound != null) {
                         bound.getTypeElement()?.index(declaration, sink)
                         return
                     }
-                    typeName = "Any"
+                    occurrence("Any")
+                    return
                 }
             }
 
-            occurrence(typeName)
+            occurrence(referenceName)
+
+            val aliasNames = declaration.getContainingJetFile().aliasImportMap()[referenceName]
+            aliasNames.forEach { occurrence(it) }
         }
 
         is JetNullableType -> getInnerType()?.index(declaration, sink)
@@ -64,3 +70,30 @@ private fun <TDeclaration : JetCallableDeclaration> JetTypeElement.index(declara
         else -> occurrence("Any")
     }
 }
+
+private class CachedAliasImportData(val map: Multimap<String, String>, val fileModificationStamp: Long)
+
+private val ALIAS_IMPORT_DATA_KEY = Key<CachedAliasImportData>("ALIAS_IMPORT_MAP_KEY")
+
+private fun JetFile.aliasImportMap(): Multimap<String, String> {
+    val cached = getUserData(ALIAS_IMPORT_DATA_KEY)
+    val modificationStamp = getModificationStamp()
+    if (cached != null && modificationStamp == cached.fileModificationStamp) {
+        return cached.map
+    }
+
+    val data = CachedAliasImportData(buildAliasImportMap(), modificationStamp)
+    putUserData(ALIAS_IMPORT_DATA_KEY, cached)
+    return data.map
+}
+
+private fun JetFile.buildAliasImportMap(): Multimap<String, String> {
+    val map = HashMultimap.create<String, String>()
+    val importList = getImportList() ?: return map
+    for (import in importList.getImports()) {
+        val aliasName = import.getAliasName() ?: continue
+        val name = import.getImportPath()?.fqnPart()?.shortName()?.asString() ?: continue
+        map.put(aliasName, name)
+    }
+    return map
+}

idea/testData/completion/basic/multifile/NotImportedExtensionFunctionAndAlias/NotImportedExtensionFunctionAndAlias.dependency1.kt

@@ -0,0 +1,9 @@
+package second
+
+import third.Dependency as MyDependency
+
+fun MyDependency?.helloFun() {
+}
+
+fun <T: MyDependency> T.helloFunGeneric() {
+}

idea/testData/completion/basic/multifile/NotImportedExtensionFunctionAndAlias/NotImportedExtensionFunctionAndAlias.dependency2.kt

@@ -0,0 +1,3 @@
+package third
+
+class Dependency

idea/testData/completion/basic/multifile/NotImportedExtensionFunctionAndAlias/NotImportedExtensionFunctionAndAlias.kt

@@ -0,0 +1,9 @@
+package first
+
+fun firstFun(x: third.Dependency) {
+    x.hello<caret>
+}
+
+// EXIST: helloFun
+// EXIST: helloFunGeneric
+// NUMBER: 2

idea/tests/org/jetbrains/kotlin/completion/MultiFileJvmBasicCompletionTestGenerated.java

@@ -149,6 +149,12 @@ public class MultiFileJvmBasicCompletionTestGenerated extends AbstractMultiFileJ
         doTest(fileName);
     }
 
+    @TestMetadata("NotImportedExtensionFunctionAndAlias")
+    public void testNotImportedExtensionFunctionAndAlias() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/basic/multifile/NotImportedExtensionFunctionAndAlias/");
+        doTest(fileName);
+    }
+
     @TestMetadata("NotImportedExtensionProperty")
     public void testNotImportedExtensionProperty() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/completion/basic/multifile/NotImportedExtensionProperty/");