From f47258c6b101e9fa7775810b83efe24e6ba50db2 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 16 Apr 2019 17:20:57 +0300
Subject: [PATCH] FIR: make substitution wrapper for class-based type scope

---
 .../kotlin/fir/resolve/ScopeUtils.kt          | 22 +++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ScopeUtils.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ScopeUtils.kt
index 851dfc5d0b8..d6798d3953a 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ScopeUtils.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ScopeUtils.kt
@@ -11,8 +11,11 @@ import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
 import org.jetbrains.kotlin.fir.resolve.transformers.firSafeNullable
 import org.jetbrains.kotlin.fir.resolve.transformers.firUnsafe
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirClassSubstitutionScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirCompositeScope
+import org.jetbrains.kotlin.fir.symbols.ConeTypeParameterSymbol
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
@@ -26,7 +29,7 @@ fun ConeKotlinType.scope(useSiteSession: FirSession, scopeSession: ScopeSession)
             // For ConeClassLikeType they might be a type alias instead of a regular class
             // TODO: support that case and switch back to `firUnsafe` instead of `firSafeNullable`
             val fir = this.lookupTag.toSymbol(useSiteSession)?.firSafeNullable<FirRegularClass>() ?: return null
-            fir.buildUseSiteScope(useSiteSession, scopeSession)
+            wrapSubstitutionScopeIfNeed(useSiteSession, fir.buildUseSiteScope(useSiteSession, scopeSession), scopeSession)
         }
         is ConeTypeParameterType -> {
             // TODO: support LibraryTypeParameterSymbol or get rid of it
@@ -39,8 +42,23 @@ fun ConeKotlinType.scope(useSiteSession: FirSession, scopeSession: ScopeSession)
     }
 }
 
+private fun ConeClassLikeType.wrapSubstitutionScopeIfNeed(
+    session: FirSession,
+    useSiteScope: FirScope,
+    builder: ScopeSession
+): FirScope {
+    if (this.typeArguments.isEmpty()) return useSiteScope
+    val symbol = this.lookupTag.toSymbol(session) as? FirClassSymbol ?: return useSiteScope
+    val regularClass = symbol.fir
+    return builder.getOrBuild(symbol, SubstitutionScopeKey(this)) {
+        @Suppress("UNCHECKED_CAST")
+        val substitution = regularClass.typeParameters.zip(this.typeArguments) { typeParameter, typeArgument ->
+            typeParameter.symbol to (typeArgument as? ConeTypedProjection)?.type
+        }.filter { (_, type) -> type != null }.toMap() as Map<ConeTypeParameterSymbol, ConeKotlinType>
 
-
+        FirClassSubstitutionScope(session, useSiteScope, substitution)
+    }
+}
 
 fun FirRegularClass.defaultType(): ConeClassTypeImpl {
     return ConeClassTypeImpl(