From 5bcd854bee2ef97cbc6c262a17ede2f63cb7b7ab Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 2 Sep 2021 13:43:09 +0300
Subject: [PATCH] [FIR] Store intersection overrides in global session
 component

---
 .../fir/scopes/FirFakeOverrideStorage.kt      | 21 +++++++-
 .../scopes/impl/FirTypeIntersectionScope.kt   | 52 +++++++++++--------
 2 files changed, 49 insertions(+), 24 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirFakeOverrideStorage.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirFakeOverrideStorage.kt
index 10b35614b2a..5dc2aa896e3 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirFakeOverrideStorage.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirFakeOverrideStorage.kt
@@ -13,7 +13,10 @@ import org.jetbrains.kotlin.fir.caches.createCache
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
 import org.jetbrains.kotlin.fir.resolve.ScopeSessionKey
 import org.jetbrains.kotlin.fir.scopes.impl.FirClassSubstitutionScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirTypeIntersectionScope
+import org.jetbrains.kotlin.fir.scopes.impl.MemberWithBaseScope
 import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
 
 class FirFakeOverrideStorage(val session: FirSession) : FirSessionComponent {
     private val cachesFactory = session.firCachesFactory
@@ -36,7 +39,23 @@ class FirFakeOverrideStorage(val session: FirSession) : FirSessionComponent {
                 }
             }
     }
+
+    class IntersectionOverrideCache(cachesFactory: FirCachesFactory) {
+        val intersectionOverrides: FirCache<FirCallableSymbol<*>, MemberWithBaseScope<FirCallableSymbol<*>>, ContextForIntersectionOverrideConstruction<*>> =
+            cachesFactory.createCache { mostSpecific, context ->
+                val (intersectionScope, extractedOverrides, scopeForMostSpecific) = context
+                intersectionScope.createIntersectionOverride(extractedOverrides, mostSpecific, scopeForMostSpecific)
+            }
+    }
+
+    data class ContextForIntersectionOverrideConstruction<D : FirCallableSymbol<*>>(
+        val intersectionScope: FirTypeIntersectionScope,
+        val extractedOverrides: List<MemberWithBaseScope<D>>,
+        val scopeForMostSpecific: FirTypeScope
+    )
+
+    val intersectionOverrideCacheByScope: FirCache<ConeKotlinType, IntersectionOverrideCache, Nothing?> =
+        cachesFactory.createCache { _ -> IntersectionOverrideCache(cachesFactory) }
 }
 
-
 val FirSession.fakeOverrideStorage: FirFakeOverrideStorage by FirSession.sessionComponentAccessor()
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirTypeIntersectionScope.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirTypeIntersectionScope.kt
index ee96ea0bd2d..bc044a02cbc 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirTypeIntersectionScope.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirTypeIntersectionScope.kt
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isExpect
 import org.jetbrains.kotlin.fir.declarations.utils.modality
@@ -34,12 +35,15 @@ class FirTypeIntersectionScope private constructor(
     private val absentProperties: MutableSet<Name> = mutableSetOf()
     private val absentClassifiers: MutableSet<Name> = mutableSetOf()
 
-    private val typeCheckerState = session.typeContext.newTypeCheckerState(false, false)
+    private val typeCheckerState = session.typeContext.newTypeCheckerState(
+        errorTypesEqualToAnything = false,
+        stubTypesEqualToAnything = false
+    )
 
-    private val overriddenSymbols: MutableMap<FirCallableSymbol<*>, Collection<MemberWithBaseScope<out FirCallableSymbol<*>>>> =
-        mutableMapOf()
+    private val overriddenSymbols: MutableMap<FirCallableSymbol<*>, Collection<MemberWithBaseScope<FirCallableSymbol<*>>>> = mutableMapOf()
 
-    private val intersectionOverrides: MutableMap<FirCallableSymbol<*>, MemberWithBaseScope<out FirCallableSymbol<*>>> = mutableMapOf()
+    private val intersectionOverrides =
+        session.fakeOverrideStorage.intersectionOverrideCacheByScope.getValue(dispatchReceiverType).intersectionOverrides
 
     private val callableNamesCached by lazy(LazyThreadSafetyMode.PUBLICATION) {
         scopes.flatMapTo(mutableSetOf()) { it.getCallableNames() }
@@ -107,23 +111,10 @@ class FirTypeIntersectionScope private constructor(
             val baseMembersForIntersection = extractedOverrides.calcBaseMembersForIntersectionOverride()
             if (baseMembersForIntersection.size > 1) {
                 val (mostSpecific, scopeForMostSpecific) = selectMostSpecificMember(baseMembersForIntersection)
-                val intersectionOverride = intersectionOverrides.getOrPut(mostSpecific) {
-                    val newModality = chooseIntersectionOverrideModality(extractedOverrides)
-                    val newVisibility = chooseIntersectionVisibility(extractedOverrides)
-                    val extractedOverridesSymbols = extractedOverrides.map { it.member }
-                    @Suppress("UNCHECKED_CAST")
-                    when (mostSpecific) {
-                        is FirNamedFunctionSymbol -> {
-                            createIntersectionOverride(mostSpecific, extractedOverridesSymbols, newModality, newVisibility)
-                        }
-                        is FirPropertySymbol -> {
-                            createIntersectionOverride(mostSpecific, extractedOverridesSymbols, newModality, newVisibility)
-                        }
-                        else -> {
-                            throw IllegalStateException("Should not be here")
-                        }
-                    }.withScope(scopeForMostSpecific)
-                }
+                val intersectionOverride = intersectionOverrides.getValue(
+                    mostSpecific,
+                    FirFakeOverrideStorage.ContextForIntersectionOverrideConstruction(this, extractedOverrides, scopeForMostSpecific)
+                )
                 overriddenSymbols[intersectionOverride.member] = extractedOverrides
                 @Suppress("UNCHECKED_CAST")
                 processor(intersectionOverride.member as D)
@@ -152,6 +143,21 @@ class FirTypeIntersectionScope private constructor(
         } while (true)
     }
 
+    fun <D : FirCallableSymbol<*>> createIntersectionOverride(
+        extractedOverrides: List<MemberWithBaseScope<D>>,
+        mostSpecific: D,
+        scopeForMostSpecific: FirTypeScope
+    ): MemberWithBaseScope<FirCallableSymbol<*>> {
+        val newModality = chooseIntersectionOverrideModality(extractedOverrides)
+        val newVisibility = chooseIntersectionVisibility(extractedOverrides)
+        val extractedOverridesSymbols = extractedOverrides.map { it.member }
+        return when (mostSpecific) {
+            is FirNamedFunctionSymbol -> createIntersectionOverride(mostSpecific, extractedOverridesSymbols, newModality, newVisibility)
+            is FirPropertySymbol -> createIntersectionOverride(mostSpecific, extractedOverridesSymbols, newModality, newVisibility)
+            else -> throw IllegalStateException("Should not be here")
+        }.withScope(scopeForMostSpecific)
+    }
+
     private fun <S : FirCallableSymbol<*>>
             MutableList<MemberWithBaseScope<S>>.calcBaseMembersForIntersectionOverride(): List<MemberWithBaseScope<S>> {
         if (size == 1) return this
@@ -502,7 +508,7 @@ class FirTypeIntersectionScope private constructor(
 
     @Suppress("UNCHECKED_CAST")
     private fun <S : FirCallableSymbol<*>> getDirectOverriddenSymbols(symbol: S): Collection<MemberWithBaseScope<S>> {
-        val intersectionOverride = intersectionOverrides[symbol]
+        val intersectionOverride = intersectionOverrides.getValueIfComputed(symbol)
         val allDirectOverridden = overriddenSymbols[symbol].orEmpty() + intersectionOverride?.let {
             overriddenSymbols[it.member]
         }.orEmpty()
@@ -565,7 +571,7 @@ class FirTypeIntersectionScope private constructor(
     }
 }
 
-private class MemberWithBaseScope<D : FirCallableSymbol<*>>(val member: D, val baseScope: FirTypeScope) {
+class MemberWithBaseScope<out D : FirCallableSymbol<*>>(val member: D, val baseScope: FirTypeScope) {
     operator fun component1() = member
     operator fun component2() = baseScope