FIR IDE: introduce composite scopes

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/scopes/KtScope.kt

@@ -22,6 +22,10 @@ interface KtScope : ValidityOwner {
     fun containsName(name: Name): Boolean
 }
 
+interface KtCompositeScope : KtScope {
+    val subScopes: List<KtScope>
+}
+
 interface KtMemberScope : KtScope {
     val owner: KtClassOrObjectSymbol
 }

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/scopes/KtScopeProvider.kt

@@ -13,4 +13,5 @@ abstract class KtScopeProvider : ValidityOwner {
     abstract fun getMemberScope(classSymbol: KtClassOrObjectSymbol): KtMemberScope
     abstract fun getDeclaredMemberScope(classSymbol: KtClassOrObjectSymbol): KtDeclaredMemberScope
     abstract fun getPackageScope(packageSymbol: KtPackageSymbol): KtPackageScope
+    abstract fun getCompositeScope(subScopes: List<KtScope>): KtCompositeScope
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirCompositeScope.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2010-2020 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.
+ */
+
+package org.jetbrains.kotlin.idea.frontend.api.fir.scopes
+
+import org.jetbrains.kotlin.idea.frontend.api.ValidityOwner
+import org.jetbrains.kotlin.idea.frontend.api.ValidityOwnerByValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtCompositeScope
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScope
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassLikeSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+import org.jetbrains.kotlin.name.Name
+
+
+// todo do we need caches here?
+@OptIn(ExperimentalStdlibApi::class)
+class KtFirCompositeScope(
+    override val subScopes: List<KtScope>,
+    override val token: ValidityOwner
+) : KtCompositeScope, ValidityOwnerByValidityToken {
+    override fun getAllNames(): Set<Name> = withValidityAssertion {
+        buildSet {
+            subScopes.flatMapTo(this) { it.getAllNames() }
+        }
+    }
+
+    override fun getCallableNames(): Set<Name> = withValidityAssertion {
+        buildSet {
+            subScopes.flatMapTo(this) { it.getCallableNames() }
+        }
+    }
+
+    override fun getClassLikeSymbolNames(): Set<Name> = withValidityAssertion {
+        buildSet {
+            subScopes.flatMapTo(this) { it.getClassLikeSymbolNames() }
+        }
+    }
+
+    override fun getAllSymbols(): Sequence<KtSymbol> = withValidityAssertion {
+        sequence {
+            subScopes.forEach { yieldAll(it.getAllSymbols()) }
+        }
+    }
+
+    override fun getCallableSymbols(): Sequence<KtSymbol> = withValidityAssertion {
+        sequence {
+            subScopes.forEach { yieldAll(it.getCallableSymbols()) }
+        }
+    }
+
+    override fun getClassClassLikeSymbols(): Sequence<KtClassLikeSymbol> = withValidityAssertion {
+        sequence {
+            subScopes.forEach { yieldAll(it.getClassClassLikeSymbols()) }
+        }
+    }
+
+    override fun containsName(name: Name): Boolean = withValidityAssertion {
+        subScopes.any { it.containsName(name) }
+    }
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirScopeProvider.kt

@@ -12,10 +12,7 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirPackageSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
-import org.jetbrains.kotlin.idea.frontend.api.scopes.KtDeclaredMemberScope
-import org.jetbrains.kotlin.idea.frontend.api.scopes.KtMemberScope
-import org.jetbrains.kotlin.idea.frontend.api.scopes.KtPackageScope
-import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeProvider
+import org.jetbrains.kotlin.idea.frontend.api.scopes.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
@@ -52,4 +49,8 @@ internal class KtFirScopeProvider(
             KtFirPackageScope(packageSymbol, token, builder, session)
         }
     }
+
+    override fun getCompositeScope(subScopes: List<KtScope>): KtCompositeScope = withValidityAssertion {
+        KtFirCompositeScope(subScopes, token)
+    }
 }