FIR scopes: use use-site member scopes for imports

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirAbstractImportingScope.kt

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.scopes.impl
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
+import org.jetbrains.kotlin.fir.resolve.calls.TowerScopeLevel
+import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.ProcessorAction
+import org.jetbrains.kotlin.fir.symbols.CallableId
+import org.jetbrains.kotlin.fir.symbols.ConeCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.ConeFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.ConeVariableSymbol
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addToStdlib.cast
+
+abstract class FirAbstractImportingScope(session: FirSession, lookupInFir: Boolean) : FirAbstractProviderBasedScope(session, lookupInFir) {
+
+
+    protected val scopeCache = mutableMapOf<ClassId, FirScope>()
+
+    fun <T : ConeCallableSymbol> processCallables(
+        import: FirResolvedImport,
+        name: Name,
+        token: TowerScopeLevel.Token<T>,
+        processor: (ConeCallableSymbol) -> ProcessorAction
+    ): ProcessorAction {
+        val callableId = CallableId(import.packageFqName, import.relativeClassName, name)
+
+        val classId = import.resolvedClassId
+        if (classId != null) {
+
+            val scope = scopeCache.getOrPut(classId) {
+                provider.getClassUseSiteMemberScope(classId, session) ?: error("No scope for $classId")
+            }
+            val action = when (token) {
+                TowerScopeLevel.Token.Functions -> scope.processFunctionsByName(
+                    callableId.callableName,
+                    processor.cast()
+                )
+                TowerScopeLevel.Token.Properties -> scope.processPropertiesByName(
+                    callableId.callableName,
+                    processor.cast()
+                )
+                else -> ProcessorAction.NEXT
+            }
+            if (action.stop()) {
+                return ProcessorAction.STOP
+            }
+        } else {
+            val symbols = provider.getTopLevelCallableSymbols(callableId.packageName, callableId.callableName)
+
+            for (symbol in symbols) {
+                if (processor(symbol).stop()) {
+                    return ProcessorAction.STOP
+                }
+            }
+        }
+
+        return ProcessorAction.NEXT
+    }
+
+    protected abstract fun <T : ConeCallableSymbol> processCallables(
+        name: Name,
+        token: TowerScopeLevel.Token<T>,
+        processor: (ConeCallableSymbol) -> ProcessorAction
+    ): ProcessorAction
+
+    final override fun processFunctionsByName(name: Name, processor: (ConeFunctionSymbol) -> ProcessorAction): ProcessorAction {
+        return processCallables(
+            name,
+            TowerScopeLevel.Token.Functions
+        ) { if (it is ConeFunctionSymbol) processor(it) else ProcessorAction.NEXT }
+    }
+
+    final override fun processPropertiesByName(name: Name, processor: (ConeVariableSymbol) -> ProcessorAction): ProcessorAction {
+        return processCallables(
+            name,
+            TowerScopeLevel.Token.Properties
+        ) { if (it is ConeVariableSymbol) processor(it) else ProcessorAction.NEXT }
+    }
+
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirAbstractSimpleImportingScope.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.fir.scopes.impl
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedImportImpl
 import org.jetbrains.kotlin.fir.resolve.FirSymbolProvider
-import org.jetbrains.kotlin.fir.resolve.getCallableSymbols
+import org.jetbrains.kotlin.fir.resolve.calls.TowerScopeLevel
 import org.jetbrains.kotlin.fir.scopes.FirPosition
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
@@ -16,7 +16,7 @@ import org.jetbrains.kotlin.fir.symbols.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
-abstract class FirAbstractSimpleImportingScope(val session: FirSession) : FirScope {
+abstract class FirAbstractSimpleImportingScope(session: FirSession) : FirAbstractImportingScope(session, lookupInFir = true) {
 
     protected abstract val simpleImports: Map<Name, List<FirResolvedImportImpl>>
 
@@ -41,34 +41,19 @@ abstract class FirAbstractSimpleImportingScope(val session: FirSession) : FirSco
         return true
     }
 
-
+    override fun <T : ConeCallableSymbol> processCallables(
-    override fun processFunctionsByName(name: Name, processor: (ConeFunctionSymbol) -> ProcessorAction): ProcessorAction {
-        return processCallables(name, processor)
-    }
-
-    override fun processPropertiesByName(name: Name, processor: (ConeVariableSymbol) -> ProcessorAction): ProcessorAction {
-        return processCallables(name, processor)
-    }
-
-    private inline fun <reified T : ConeCallableSymbol> processCallables(
         name: Name,
-        processor: (T) -> ProcessorAction
+        token: TowerScopeLevel.Token<T>,
+        processor: (ConeCallableSymbol) -> ProcessorAction
     ): ProcessorAction {
+
+
         val imports = simpleImports[name] ?: return ProcessorAction.NEXT
         if (imports.isEmpty()) return ProcessorAction.NEXT
-        val provider = FirSymbolProvider.getInstance(session)
-        for (import in imports) {
-            val importedName = import.importedName ?: continue
-            val callableId = CallableId(
-                import.packageFqName,
-                import.relativeClassName,
-                importedName
-            )
 
-            for (symbol in provider.getCallableSymbols(callableId).filterIsInstance<T>()) {
+        for (import in imports) {
-                if (!processor(symbol)) {
+            if (processCallables(import, name, token, processor).stop()) {
-                    return ProcessorAction.NEXT
+                return ProcessorAction.STOP
-                }
             }
         }
         return ProcessorAction.NEXT

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirAbstractStarImportingScope.kt

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.fir.scopes.impl
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
-import org.jetbrains.kotlin.fir.resolve.getCallableSymbols
+import org.jetbrains.kotlin.fir.resolve.calls.TowerScopeLevel
 import org.jetbrains.kotlin.fir.scopes.FirPosition
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.fir.symbols.*
@@ -16,7 +16,7 @@ import org.jetbrains.kotlin.name.Name
 
 abstract class FirAbstractStarImportingScope(
     session: FirSession, lookupInFir: Boolean = true
-) : FirAbstractProviderBasedScope(session, lookupInFir) {
+) : FirAbstractImportingScope(session, lookupInFir) {
 
     protected abstract val starImports: List<FirResolvedImport>
 
@@ -40,28 +40,16 @@ abstract class FirAbstractStarImportingScope(
         return true
     }
 
-    override fun processFunctionsByName(name: Name, processor: (ConeFunctionSymbol) -> ProcessorAction): ProcessorAction {
-        return processCallables(name, processor)
-    }
 
-    override fun processPropertiesByName(name: Name, processor: (ConeVariableSymbol) -> ProcessorAction): ProcessorAction {
+    override fun <T : ConeCallableSymbol> processCallables(
-        return processCallables(name, processor)
-    }
-
-    private inline fun <reified T : ConeCallableSymbol> processCallables(
         name: Name,
-        processor: (T) -> ProcessorAction
+        token: TowerScopeLevel.Token<T>,
+        processor: (ConeCallableSymbol) -> ProcessorAction
     ): ProcessorAction {
         for (import in starImports) {
-            val callableId = CallableId(import.packageFqName, import.relativeClassName, name)
+            if (processCallables(import, name, token, processor).stop()) {
-            val symbols = provider.getCallableSymbols(callableId).filterIsInstance<T>()
+                return ProcessorAction.STOP
-
-            for (symbol in symbols) {
-                if (processor(symbol) == ProcessorAction.STOP) {
-                    return ProcessorAction.STOP
-                }
             }
-
         }
         return ProcessorAction.NEXT
     }