Iterate on symbols, not on scopes

- This way it is easier to get the correct results

idea/idea-completion/testData/basic/common/primitiveCompletion/classFieldsAndFunctionsFromInnerClassExplicitReceiver.kt

@@ -0,0 +1,18 @@
+class A {
+    fun aa() {}
+    val aaa = 10
+
+    inner class AA {
+        fun bb() {}
+        val bbb = 20
+
+        fun test() {
+            this.<caret>
+        }
+    }
+}
+
+// ABSENT: aa
+// ABSENT: aaa
+// EXIST: bb
+// EXIST: bbb

idea/idea-completion/testData/basic/common/primitiveCompletion/extensionPropertyAndFunctionImplicitReceiver.kt

@@ -0,0 +1,25 @@
+fun Any.anyFun() {}
+val Any.anyVal: Int get() = 10
+
+class A
+fun A.aFun() {}
+val A.aVal: Int get() = 10
+
+class B
+fun B.bFun() {}
+val B.bVal: Int get() = 10
+
+fun test(a: A) {
+    a.run {
+        <caret>
+
+        Unit
+    }
+}
+
+// EXIST: anyFun
+// EXIST: anyVal
+// EXIST: aVal
+// EXIST: aFun
+// ABSENT: bVal
+// ABSENT: bFun

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JSBasicCompletionTestGenerated.java

@@ -2247,6 +2247,11 @@ public class JSBasicCompletionTestGenerated extends AbstractJSBasicCompletionTes
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/classFieldsAndFunctionsFromInnerClass.kt");
             }
 
+            @TestMetadata("classFieldsAndFunctionsFromInnerClassExplicitReceiver.kt")
+            public void testClassFieldsAndFunctionsFromInnerClassExplicitReceiver() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/classFieldsAndFunctionsFromInnerClassExplicitReceiver.kt");
+            }
+
             @TestMetadata("explicitReceiverCompletion.kt")
             public void testExplicitReceiverCompletion() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/explicitReceiverCompletion.kt");
@@ -2257,6 +2262,11 @@ public class JSBasicCompletionTestGenerated extends AbstractJSBasicCompletionTes
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/extensionPropertyAndFunctionExplicitReceiver.kt");
             }
 
+            @TestMetadata("extensionPropertyAndFunctionImplicitReceiver.kt")
+            public void testExtensionPropertyAndFunctionImplicitReceiver() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/extensionPropertyAndFunctionImplicitReceiver.kt");
+            }
+
             @TestMetadata("genericExtensionPropertyAndFunctionExplicitReceiver.kt")
             public void testGenericExtensionPropertyAndFunctionExplicitReceiver() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/genericExtensionPropertyAndFunctionExplicitReceiver.kt");

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmBasicCompletionTestGenerated.java

@@ -2247,6 +2247,11 @@ public class JvmBasicCompletionTestGenerated extends AbstractJvmBasicCompletionT
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/classFieldsAndFunctionsFromInnerClass.kt");
             }
 
+            @TestMetadata("classFieldsAndFunctionsFromInnerClassExplicitReceiver.kt")
+            public void testClassFieldsAndFunctionsFromInnerClassExplicitReceiver() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/classFieldsAndFunctionsFromInnerClassExplicitReceiver.kt");
+            }
+
             @TestMetadata("explicitReceiverCompletion.kt")
             public void testExplicitReceiverCompletion() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/explicitReceiverCompletion.kt");
@@ -2257,6 +2262,11 @@ public class JvmBasicCompletionTestGenerated extends AbstractJvmBasicCompletionT
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/extensionPropertyAndFunctionExplicitReceiver.kt");
             }
 
+            @TestMetadata("extensionPropertyAndFunctionImplicitReceiver.kt")
+            public void testExtensionPropertyAndFunctionImplicitReceiver() throws Exception {
+                runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/extensionPropertyAndFunctionImplicitReceiver.kt");
+            }
+
             @TestMetadata("genericExtensionPropertyAndFunctionExplicitReceiver.kt")
             public void testGenericExtensionPropertyAndFunctionExplicitReceiver() throws Exception {
                 runTest("idea/idea-completion/testData/basic/common/primitiveCompletion/genericExtensionPropertyAndFunctionExplicitReceiver.kt");

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt

@@ -17,7 +17,6 @@ import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
 import org.jetbrains.kotlin.idea.fir.getFirOfClosestParent
 import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
@@ -27,6 +26,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.types.AbstractTypeChecker
+import org.jetbrains.kotlin.types.model.TypeCheckerProviderContext
 
 class KotlinFirCompletionContributor : CompletionContributor() {
     init {
@@ -53,65 +53,71 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
         val element = nameExpression.getFirOfClosestParent() as? FirQualifiedAccessExpression ?: return
         val towerDataContext = completionContext.getTowerDataContext(nameExpression)
 
-        val receiver = element.explicitReceiver
+        val symbols: Sequence<FirCallableSymbol<*>> = sequence {
-        val explicitReceiverType = receiver?.typeRef?.coneTypeUnsafe<ConeKotlinType>()
+            val explicitReceiver = element.explicitReceiver
-
+            val explicitReceiverType = explicitReceiver?.typeRef?.coneTypeUnsafe<ConeKotlinType>()
-        val scopes: Sequence<FirScope> = sequence {
             val explicitReceiverScope = explicitReceiverType?.scope(completionContext.session, ScopeSession())
-            if (explicitReceiverScope != null) {
-                yield(explicitReceiverScope)
-            }
 
-            yieldAll(towerDataContext.localScopes)
+            val implicitReceivers = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.implicitReceiver }
-
+            val implicitReceiversTypes = implicitReceivers.map { it.type }
-            val implicitReceiversScopes = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.implicitReceiver?.implicitScope }
-            yieldAll(implicitReceiversScopes)
 
+            val localScopes = towerDataContext.localScopes
             val nonLocalScopes = towerDataContext.nonLocalTowerDataElements.mapNotNull { it.scope }
-            yieldAll(nonLocalScopes)
+            val implicitReceiversScopes = implicitReceivers.mapNotNull { it.implicitScope }
-        }
 
-        for (scope in scopes) {
+            val allNonExplicitScopes = localScopes.asSequence() + nonLocalScopes + implicitReceiversScopes
-            for (symbol in scope.collectCallableSymbols()) {
-                val symbolName = symbol.callableId.callableName
-                if (symbolName.isConstructor) continue
 
-                val expectedReceiverType = symbol.fir.receiverTypeRef?.coneTypeUnsafe<ConeKotlinType>()
+            val typeContext = completionContext.session.typeContext
 
-                if (expectedReceiverType != null) {
+            if (explicitReceiverScope != null) {
-                    val receiverTypes = if (explicitReceiverType != null) {
+                yieldAll(explicitReceiverScope.collectCallableSymbols())
-                        listOf(explicitReceiverType)
-                    } else {
-                        towerDataContext.nonLocalTowerDataElements.mapNotNull { it.implicitReceiver?.type }
-                    }
 
-                    val expectedReceiverTypeIsPresent = receiverTypes.any {
+                val allApplicableExtensions = allNonExplicitScopes
-                        AbstractTypeChecker.isSubtypeOf(completionContext.session.typeContext, it, expectedReceiverType)
+                    .flatMap { it.collectCallableSymbols() }
-                    }
+                    .filter { it.isExtension && it.isExtensionThatCanBeCalledOnTypes(listOf(explicitReceiverType), typeContext) }
 
-                    if (expectedReceiverTypeIsPresent) {
+                yieldAll(allApplicableExtensions)
-                        result.addElement(LookupElementBuilder.create(symbolName.toString()))
+            } else {
-                    }
+                val allAvailableSymbols = allNonExplicitScopes
-                } else if (explicitReceiverType == null || symbol.callableId.classId == explicitReceiverType.classId) {
+                    .flatMap { it.collectCallableSymbols() }
-                    result.addElement(LookupElementBuilder.create(symbolName.toString()))
+                    .filter { !it.isExtension || it.isExtensionThatCanBeCalledOnTypes(implicitReceiversTypes, typeContext) }
-                }
+
+                yieldAll(allAvailableSymbols)
             }
         }
+
+        for (symbol in symbols) {
+            val symbolName = symbol.callableId.callableName
+            if (symbolName.isConstructor) continue
+
+            result.addElement(LookupElementBuilder.create(symbolName.toString()))
+        }
     }
 
-    private fun FirScope.collectCallableSymbols(): MutableList<FirCallableSymbol<*>> {
+}
-        val allCallableSymbols = mutableListOf<FirCallableSymbol<*>>()
+
-        fun symbolsCollector(symbol: FirCallableSymbol<*>) {
+private val FirCallableSymbol<*>.isExtension get() = fir.receiverTypeRef != null
-            allCallableSymbols.add(symbol)
+
-        }
+private val Name.isConstructor get() = this == Name.special("<init>")
-
+
-        for (name in getCallableNames()) {
+private fun FirScope.collectCallableSymbols(): MutableList<FirCallableSymbol<*>> {
-            processFunctionsByName(name, ::symbolsCollector)
+    val allCallableSymbols = mutableListOf<FirCallableSymbol<*>>()
-            processPropertiesByName(name, ::symbolsCollector)
+    fun symbolsCollector(symbol: FirCallableSymbol<*>) {
-        }
+        allCallableSymbols.add(symbol)
-
+    }
-        return allCallableSymbols
+
-    }
+    for (name in getCallableNames()) {
-
+        processFunctionsByName(name, ::symbolsCollector)
-    private val Name.isConstructor get() = this == Name.special("<init>")
+        processPropertiesByName(name, ::symbolsCollector)
+    }
+
+    return allCallableSymbols
+}
+
+private fun FirCallableSymbol<*>.isExtensionThatCanBeCalledOnTypes(
+    receivers: List<ConeKotlinType>,
+    session: TypeCheckerProviderContext
+): Boolean {
+    val expectedReceiverType = fir.receiverTypeRef?.coneTypeUnsafe<ConeKotlinType>() ?: return false
+    return receivers.any { AbstractTypeChecker.isSubtypeOf(session, it, expectedReceiverType) }
 }