Smart completion after "by" and "in": non-imported operators supported

idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt

@@ -140,6 +140,7 @@ object IfConditionAdditionalData : ExpectedInfo.AdditionalData
 class ExpectedInfos(
         private val bindingContext: BindingContext,
         private val resolutionFacade: ResolutionFacade,
+        private val indicesHelper: KotlinIndicesHelper?,
         private val useHeuristicSignatures: Boolean = true,
         private val useOuterCallsExpectedTypeCount: Int = 0
 ) {
@@ -207,7 +208,7 @@ class ExpectedInfos(
 
         if (useOuterCallsExpectedTypeCount > 0 && results.any(::makesSenseToUseOuterCallExpectedType)) {
             val callExpression = (call.callElement as? KtExpression)?.getQualifiedExpressionForSelectorOrThis() ?: return results
-            val expectedFuzzyTypes = ExpectedInfos(bindingContext, resolutionFacade, useHeuristicSignatures, useOuterCallsExpectedTypeCount - 1)
+            val expectedFuzzyTypes = ExpectedInfos(bindingContext, resolutionFacade, indicesHelper, useHeuristicSignatures, useOuterCallsExpectedTypeCount - 1)
                     .calculate(callExpression)
                     .mapNotNull { it.fuzzyType }
             if (expectedFuzzyTypes.isEmpty() || expectedFuzzyTypes.any { it.freeParameters.isNotEmpty() }) return results
@@ -587,7 +588,7 @@ class ExpectedInfos(
 
         val leftOperandType = binaryExpression.left?.let { bindingContext.getType(it) } ?: return null
         val scope = expressionWithType.getResolutionScope(bindingContext, resolutionFacade)
-        val detector = TypesWithContainsDetector(scope, leftOperandType)
+        val detector = TypesWithContainsDetector(scope, indicesHelper, leftOperandType)
 
         val byTypeFilter = object : ByTypeFilter {
             override fun matchingSubstitutor(descriptorType: FuzzyType): TypeSubstitutor? {
@@ -608,8 +609,8 @@ class ExpectedInfos(
                             ?: return null
 
         val explicitPropertyType = property.returnType?.check { propertyDeclaration.typeReference != null }
-        val typesWithGetDetector = TypesWithGetValueDetector(scope, propertyOwnerType, explicitPropertyType)
-        val typesWithSetDetector = if (property.isVar) TypesWithSetValueDetector(scope, propertyOwnerType) else null
+        val typesWithGetDetector = TypesWithGetValueDetector(scope, indicesHelper, propertyOwnerType, explicitPropertyType)
+        val typesWithSetDetector = if (property.isVar) TypesWithSetValueDetector(scope, indicesHelper, propertyOwnerType) else null
 
         val byTypeFilter = object : ByTypeFilter {
             override fun matchingSubstitutor(descriptorType: FuzzyType): TypeSubstitutor? {

idea/idea-core/src/org/jetbrains/kotlin/idea/core/KotlinIndicesHelper.kt

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.idea.core
 
-import com.intellij.codeInsight.CodeInsightSettings
 import com.intellij.codeInsight.JavaProjectCodeInsightSettings
 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.psi.PsiFile
@@ -26,10 +25,7 @@ import com.intellij.psi.search.PsiShortNamesCache
 import com.intellij.psi.stubs.StringStubIndexExtension
 import com.intellij.util.Processor
 import com.intellij.util.indexing.IdFilter
-import org.jetbrains.kotlin.descriptors.CallableDescriptor
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.descriptors.ClassKind
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.caches.resolve.getJavaFieldDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.getJavaMethodDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.resolveImportReference
@@ -78,12 +74,12 @@ class KotlinIndicesHelper(
     }
 
     fun getTopLevelCallablesByName(name: String): Collection<CallableDescriptor> {
-        val declarations = HashSet<KtCallableDeclaration>()
+        val declarations = LinkedHashSet<KtCallableDeclaration>()
         declarations.addTopLevelNonExtensionCallablesByName(KotlinFunctionShortNameIndex.getInstance(), name)
         declarations.addTopLevelNonExtensionCallablesByName(KotlinPropertyShortNameIndex.getInstance(), name)
         return declarations
                 .flatMap { it.resolveToDescriptorsWithHack() }
-                .filter { it.extensionReceiverParameter == null && descriptorFilter(it) }
+                .filter { descriptorFilter(it) }
     }
 
     private fun MutableSet<KtCallableDeclaration>.addTopLevelNonExtensionCallablesByName(
@@ -93,6 +89,15 @@ class KotlinIndicesHelper(
         index.get(name, project, scope).filterTo(this) { it.parent is KtFile && it.receiverTypeReference == null }
     }
 
+    fun getTopLevelExtensionOperatorsByName(name: String): Collection<FunctionDescriptor> {
+        return KotlinFunctionShortNameIndex.getInstance().get(name, project, scope)
+                .filter { it.parent is KtFile && it.receiverTypeReference != null && it.hasModifier(KtTokens.OPERATOR_KEYWORD) }
+                .flatMap { it.resolveToDescriptorsWithHack() }
+                .filterIsInstance<FunctionDescriptor>()
+                .filter { descriptorFilter(it) }
+                .distinct()
+    }
+
     fun processTopLevelCallables(nameFilter: (String) -> Boolean, processor: (CallableDescriptor) -> Unit) {
         fun processIndex(index: StringStubIndexExtension<out KtCallableDeclaration>) {
             for (key in index.getAllKeys(project)) {

idea/idea-core/src/org/jetbrains/kotlin/idea/core/TypesWithOperatorDetector.kt

@@ -30,22 +30,30 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.util.isValidOperator
 import java.util.*
 
-//TODO: support not imported extensions
 abstract class TypesWithOperatorDetector(
         private val name: Name,
-        private val scope: LexicalScope
+        private val scope: LexicalScope,
+        private val indicesHelper: KotlinIndicesHelper?
 ) {
     protected abstract fun isSuitableByType(function: FunctionDescriptor, freeTypeParams: Collection<TypeParameterDescriptor>): Boolean
 
     private val cache = HashMap<FuzzyType, FunctionDescriptor?>()
 
-    private val typesWithExtension: Map<KotlinType, FunctionDescriptor> by lazy {
+    private val typesWithExtensionFromScope: Map<KotlinType, FunctionDescriptor> by lazy {
         scope.collectFunctions(name, NoLookupLocation.FROM_IDE)
                 .filter { it.extensionReceiverParameter != null && it.isValidOperator() && isSuitableByType(it, it.typeParameters) }
                 .map { it.extensionReceiverParameter!!.type to it }
                 .toMap()
     }
 
+    private val typesWithExtensionFromIndices: Map<KotlinType, FunctionDescriptor> by lazy {
+        indicesHelper?.getTopLevelExtensionOperatorsByName(name.asString())
+                ?.filter { it.extensionReceiverParameter != null && it.isValidOperator() && isSuitableByType(it, it.typeParameters) }
+                ?.map { it.extensionReceiverParameter!!.type to it }
+                ?.filter { it.first !in typesWithExtensionFromScope.keys }
+                ?.toMap() ?: emptyMap()
+    }
+
     fun findOperator(type: FuzzyType): FunctionDescriptor? {
         if (cache.containsKey(type)) {
             return cache[type]
@@ -64,19 +72,21 @@ abstract class TypesWithOperatorDetector(
             if (memberFunction != null) return memberFunction
         }
 
-        for ((typeWithExtension, operator) in typesWithExtension) {
+        for ((typeWithExtension, operator) in typesWithExtensionFromScope + typesWithExtensionFromIndices) {
             if (type.checkIsSubtypeOf(typeWithExtension) != null) {
                 return operator //TODO: substitution
             }
         }
+
         return null
     }
 }
 
 class TypesWithContainsDetector(
         scope: LexicalScope,
+        indicesHelper: KotlinIndicesHelper?,
         private val argumentType: KotlinType
-) : TypesWithOperatorDetector(OperatorNameConventions.CONTAINS, scope) {
+) : TypesWithOperatorDetector(OperatorNameConventions.CONTAINS, scope, indicesHelper) {
 
     override fun isSuitableByType(function: FunctionDescriptor, freeTypeParams: Collection<TypeParameterDescriptor>): Boolean {
         val parameter = function.valueParameters.single()
@@ -87,9 +97,10 @@ class TypesWithContainsDetector(
 
 class TypesWithGetValueDetector(
         scope: LexicalScope,
+        indicesHelper: KotlinIndicesHelper?,
         private val propertyOwnerType: KotlinType,
         private val propertyType: KotlinType?
-) : TypesWithOperatorDetector(OperatorNameConventions.GET_VALUE, scope) {
+) : TypesWithOperatorDetector(OperatorNameConventions.GET_VALUE, scope, indicesHelper) {
 
     override fun isSuitableByType(function: FunctionDescriptor, freeTypeParams: Collection<TypeParameterDescriptor>): Boolean {
         val paramType = FuzzyType(function.valueParameters.first().type, freeTypeParams)
@@ -106,8 +117,9 @@ class TypesWithGetValueDetector(
 
 class TypesWithSetValueDetector(
         scope: LexicalScope,
+        indicesHelper: KotlinIndicesHelper?,
         private val propertyOwnerType: KotlinType
-) : TypesWithOperatorDetector(OperatorNameConventions.SET_VALUE, scope) {
+) : TypesWithOperatorDetector(OperatorNameConventions.SET_VALUE, scope, indicesHelper) {
 
     override fun isSuitableByType(function: FunctionDescriptor, freeTypeParams: Collection<TypeParameterDescriptor>): Boolean {
         val paramType = FuzzyType(function.valueParameters.first().type, freeTypeParams)