[FIR] Implement primitive tower for type resolving

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirTypeResolver.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.scopes.FirIterableScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.service
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
@@ -15,8 +16,8 @@ import org.jetbrains.kotlin.fir.types.FirUserTypeRef
 
 interface FirTypeResolver {
 
-    fun resolveType(typeRef: FirTypeRef, scope: FirScope): ConeKotlinType
+    fun resolveType(typeRef: FirTypeRef, scope: FirIterableScope): ConeKotlinType
-    fun resolveToSymbol(typeRef: FirTypeRef, scope: FirScope): FirClassifierSymbol<*>?
+    fun resolveToSymbol(typeRef: FirTypeRef, scope: FirIterableScope): FirClassifierSymbol<*>?
 
     companion object {
         fun getInstance(session: FirSession): FirTypeResolver = session.service()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/impl/FirTypeResolverImpl.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.resolve.FirQualifierResolver
 import org.jetbrains.kotlin.fir.resolve.FirTypeResolver
 import org.jetbrains.kotlin.fir.resolve.constructType
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
+import org.jetbrains.kotlin.fir.scopes.FirIterableScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
@@ -41,7 +42,7 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver {
 
     override fun resolveToSymbol(
         typeRef: FirTypeRef,
-        scope: FirScope
+        scope: FirIterableScope
     ): FirClassifierSymbol<*>? {
         return when (typeRef) {
             is FirResolvedTypeRef -> typeRef.coneTypeSafe<ConeLookupTagBasedType>()?.lookupTag?.let(symbolProvider::getSymbolByLookupTag)
@@ -50,22 +51,25 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver {
                 val qualifierResolver = FirQualifierResolver.getInstance(session)
 
                 var resolvedSymbol: FirClassifierSymbol<*>? = null
-                scope.processClassifiersByName(typeRef.qualifier.first().name) { symbol ->
+                for (typeScope in scope.scopes) {
-                    if (resolvedSymbol != null) return@processClassifiersByName
+                    typeScope.processClassifiersByName(typeRef.qualifier.first().name) { symbol ->
-                    resolvedSymbol = when (symbol) {
+                        if (resolvedSymbol != null) return@processClassifiersByName
-                        is FirClassLikeSymbol<*> -> {
+                        resolvedSymbol = when (symbol) {
-                            if (typeRef.qualifier.size == 1) {
+                            is FirClassLikeSymbol<*> -> {
-                                symbol
+                                if (typeRef.qualifier.size == 1) {
-                            } else {
+                                    symbol
-                                qualifierResolver.resolveSymbolWithPrefix(typeRef.qualifier, symbol.classId)
+                                } else {
+                                    qualifierResolver.resolveSymbolWithPrefix(typeRef.qualifier, symbol.classId)
+                                }
                             }
+                            is FirTypeParameterSymbol -> {
+                                assert(typeRef.qualifier.size == 1)
+                                symbol
+                            }
+                            else -> error("!")
                         }
-                        is FirTypeParameterSymbol -> {
-                            assert(typeRef.qualifier.size == 1)
-                            symbol
-                        }
-                        else -> error("!")
                     }
+                    if (resolvedSymbol != null) break
                 }
 
                 // TODO: Imports
@@ -89,7 +93,7 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver {
     private fun createFunctionalType(typeRef: FirFunctionTypeRef): ConeClassLikeType {
         val parameters =
             listOfNotNull((typeRef.receiverTypeRef as FirResolvedTypeRef?)?.type) +
-                    typeRef.valueParameters.map { it.returnTypeRef.coneTypeUnsafe<ConeKotlinType>() } +
+                    typeRef.valueParameters.map { it.returnTypeRef.coneTypeUnsafe() } +
                     listOf(typeRef.returnTypeRef.coneTypeUnsafe())
         return ConeClassLikeTypeImpl(
             resolveBuiltInQualified(KotlinBuiltIns.getFunctionClassId(typeRef.parametersCount), session).toLookupTag(),
@@ -100,7 +104,7 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver {
 
     override fun resolveType(
         typeRef: FirTypeRef,
-        scope: FirScope
+        scope: FirIterableScope
     ): ConeKotlinType {
         return when (typeRef) {
             is FirResolvedTypeRef -> typeRef.type

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirAbstractTreeTransformerWithSuperTypes.kt

@@ -25,9 +25,7 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
         val result = l()
         val size = towerScope.scopes.size
         assert(size >= sizeBefore)
-        repeat(size - sizeBefore) {
+        towerScope.dropLastScopes(size - sizeBefore)
-            towerScope.scopes.let { it.removeAt(it.size - 1) }
-        }
         return result
     }
 
@@ -37,25 +35,26 @@ abstract class FirAbstractTreeTransformerWithSuperTypes(
     ): CompositeTransformResult<FirStatement> {
         return withScopeCleanup {
             // ? Is it Ok to use original file session here ?
-            lookupSuperTypes(regularClass, lookupInterfaces = false, deep = true, useSiteSession = session)
+            val superTypes = lookupSuperTypes(regularClass, lookupInterfaces = false, deep = true, useSiteSession = session).asReversed()
-                .asReversed().mapNotNullTo(towerScope.scopes) {
+            for (superType in superTypes) {
-                    session.firSymbolProvider.getNestedClassifierScope(it.lookupTag.classId)
+                session.firSymbolProvider.getNestedClassifierScope(superType.lookupTag.classId)?.let {
+                    towerScope.addScope(it)
                 }
+            }
             regularClass.addTypeParametersScope()
             val companionObject = regularClass.companionObject
             if (companionObject != null) {
-                towerScope.scopes += nestedClassifierScope(companionObject)
+                towerScope.addScope(nestedClassifierScope(companionObject))
             }
-            towerScope.scopes += nestedClassifierScope(regularClass)
+            towerScope.addScope(nestedClassifierScope(regularClass))
 
             transformElement(regularClass, data)
         }
     }
 
     protected fun FirMemberDeclaration.addTypeParametersScope() {
-        val scopes = towerScope.scopes
         if (typeParameters.isNotEmpty()) {
-            scopes += FirMemberTypeParameterScope(this)
+            towerScope.addScope(FirMemberTypeParameterScope(this))
         }
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSpecificTypeResolverTransformer.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.FirTypeResolver
-import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.FirIterableScope
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.FirResolvedFunctionTypeRefImpl
 import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
@@ -16,7 +16,7 @@ import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
 
 class FirSpecificTypeResolverTransformer(
-    private val towerScope: FirScope,
+    private val towerScope: FirIterableScope,
     override val session: FirSession
 ) : FirAbstractTreeTransformer<Nothing?>(phase = FirResolvePhase.SUPER_TYPES) {
     private val typeResolver = FirTypeResolver.getInstance(session)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -11,18 +11,15 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.diagnostics.FirSimpleDiagnostic
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.scopes.FirIterableScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.createImportingScopes
 import org.jetbrains.kotlin.fir.scopes.impl.FirMemberTypeParameterScope
 import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.FirErrorTypeRefImpl
 import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -376,41 +373,8 @@ sealed class SupertypeComputationStatus {
 private typealias ImmutableList<E> = javaslang.collection.List<E>
 
 private class FirImmutableCompositeScope(
-    private val scopes: ImmutableList<FirScope>
+    override val scopes: ImmutableList<FirScope>
-) : FirScope() {
+) : FirIterableScope() {
-
     fun childScope(newScope: FirScope?) = newScope?.let { FirImmutableCompositeScope(scopes.push(newScope)) } ?: this
     fun childScope(newScopes: Collection<FirScope>) = FirImmutableCompositeScope(scopes.pushAll(newScopes))
-
-    override fun processClassifiersByName(
-        name: Name,
-        processor: (FirClassifierSymbol<*>) -> Unit
-    ) {
-        for (scope in scopes) {
-            scope.processClassifiersByName(name, processor)
-        }
-    }
-
-    private inline fun <T> processComposite(
-        process: FirScope.(Name, (T) -> Unit) -> Unit,
-        name: Name,
-        noinline processor: (T) -> Unit
-    ) {
-        val unique = mutableSetOf<T>()
-        for (scope in scopes) {
-            scope.process(name) {
-                if (unique.add(it)) {
-                    processor(it)
-                }
-            }
-        }
-    }
-
-    override fun processFunctionsByName(name: Name, processor: (FirFunctionSymbol<*>) -> Unit) {
-        return processComposite(FirScope::processFunctionsByName, name, processor)
-    }
-
-    override fun processPropertiesByName(name: Name, processor: (FirCallableSymbol<*>) -> Unit) {
-        return processComposite(FirScope::processPropertiesByName, name, processor)
-    }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirTypeResolveTransformer.kt

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.scopes.addImportingScopes
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/Scopes.kt

@@ -36,9 +36,5 @@ fun createImportingScopes(
     )
 }
 
-fun FirCompositeScope.addImportingScopes(file: FirFile, session: FirSession, scopeSession: ScopeSession) {
-    scopes.addImportingScopes(file, session, scopeSession)
-}
-
 private val PACKAGE_MEMBER = scopeSessionKey<FqName, FirPackageMemberScope>()
 

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

@@ -5,50 +5,31 @@
 
 package org.jetbrains.kotlin.fir.scopes.impl
 
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.scopes.FirIterableScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.scopes.createImportingScopes
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
-import org.jetbrains.kotlin.name.Name
 
 class FirCompositeScope(
-    val scopes: MutableList<FirScope>,
+    private val scopeList: MutableList<FirScope>,
     private val reversedPriority: Boolean = false
-) : FirScope() {
+) : FirIterableScope() {
-    constructor(vararg scopes: FirScope) : this(scopes.toMutableList())
+    override val scopes get() = if (reversedPriority) scopeList.asReversed() else scopeList
 
-    override fun processClassifiersByName(
+    fun addImportingScopes(file: FirFile, session: FirSession, scopeSession: ScopeSession) {
-        name: Name,
+        scopeList += createImportingScopes(file, session, scopeSession)
-        processor: (FirClassifierSymbol<*>) -> Unit
+    }
-    ) {
+
-        val scopes = if (reversedPriority) scopes.asReversed() else scopes
+    fun addScope(scope: FirScope) {
-        for (scope in scopes) {
+        scopeList += scope
-            scope.processClassifiersByName(name, processor)
+    }
+
+    fun dropLastScopes(number: Int) {
+        repeat(number) {
+            scopeList.removeAt(scopeList.size - 1)
         }
     }
 
-    private inline fun <T> processComposite(
-        process: FirScope.(Name, (T) -> Unit) -> Unit,
-        name: Name,
-        noinline processor: (T) -> Unit
-    ) {
-        val unique = mutableSetOf<T>()
-        val scopes = if (reversedPriority) scopes.asReversed() else scopes
-        for (scope in scopes) {
-            scope.process(name) {
-                if (unique.add(it)) {
-                    processor(it)
-                }
-            }
-        }
-    }
-
-    override fun processFunctionsByName(name: Name, processor: (FirFunctionSymbol<*>) -> Unit) {
-        return processComposite(FirScope::processFunctionsByName, name, processor)
-    }
-
-    override fun processPropertiesByName(name: Name, processor: (FirCallableSymbol<*>) -> Unit) {
-        return processComposite(FirScope::processPropertiesByName, name, processor)
-    }
-
 }

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

@@ -6,29 +6,17 @@
 package org.jetbrains.kotlin.fir.scopes.impl
 
 import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
+import org.jetbrains.kotlin.fir.scopes.FirIterableScope
 import org.jetbrains.kotlin.fir.scopes.FirScope
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
-import org.jetbrains.kotlin.name.Name
 
 class FirTypeResolveScopeForBodyResolve(
     private val topLevelScopes: List<FirScope>,
     private val implicitReceiverStack: ImplicitReceiverStack,
     private val localScopes: List<FirScope>
-) : FirScope() {
+) : FirIterableScope() {
-    override fun processClassifiersByName(
+    override val scopes: Iterable<FirScope>
-        name: Name,
+        get() = localScopes.asReversed() + implicitReceiverStack.receiversAsReversed().mapNotNull {
-        processor: (FirClassifierSymbol<*>) -> Unit
+            (it as? ImplicitDispatchReceiverValue)?.implicitScope
-    ) {
+        } + topLevelScopes.asReversed()
-        for (scope in localScopes.asReversed()) {
-            scope.processClassifiersByName(name, processor)
-        }
-        for (receiverValue in implicitReceiverStack.receiversAsReversed()) {
-            if (receiverValue is ImplicitExtensionReceiverValue) continue
-            receiverValue.implicitScope?.processClassifiersByName(name, processor)
-        }
-        for (scope in topLevelScopes.asReversed()) {
-            scope.processClassifiersByName(name, processor)
-        }
-    }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/scopes/FirIterableScope.kt

@@ -0,0 +1,47 @@
+/*
+ * 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.fir.scopes
+
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.name.Name
+
+abstract class FirIterableScope : FirScope() {
+    abstract val scopes: Iterable<FirScope>
+
+    override fun processClassifiersByName(
+        name: Name,
+        processor: (FirClassifierSymbol<*>) -> Unit
+    ) {
+        for (scope in scopes) {
+            scope.processClassifiersByName(name, processor)
+        }
+    }
+
+    private inline fun <T> processComposite(
+        process: FirScope.(Name, (T) -> Unit) -> Unit,
+        name: Name,
+        noinline processor: (T) -> Unit
+    ) {
+        val unique = mutableSetOf<T>()
+        for (scope in scopes) {
+            scope.process(name) {
+                if (unique.add(it)) {
+                    processor(it)
+                }
+            }
+        }
+    }
+
+    override fun processFunctionsByName(name: Name, processor: (FirFunctionSymbol<*>) -> Unit) {
+        return processComposite(FirScope::processFunctionsByName, name, processor)
+    }
+
+    override fun processPropertiesByName(name: Name, processor: (FirCallableSymbol<*>) -> Unit) {
+        return processComposite(FirScope::processPropertiesByName, name, processor)
+    }
+}