[FIR] Store substitution overrides in global session component

2021-09-02 13:16:48 +03:00
parent e67a78abaf
commit 38a2af0c92
6 changed files with 86 additions and 43 deletions
@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.fir.resolve.providers.impl.FirTypeResolverImpl
 import org.jetbrains.kotlin.fir.resolve.transformers.FirPhaseCheckingPhaseManager
 import org.jetbrains.kotlin.fir.symbols.FirPhaseManager
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClassIndex
+import org.jetbrains.kotlin.fir.scopes.FirFakeOverrideStorage
 import org.jetbrains.kotlin.fir.scopes.impl.FirDeclaredMemberScopeProvider
 import org.jetbrains.kotlin.fir.types.FirCorrespondingSupertypesCache
 import org.jetbrains.kotlin.incremental.components.LookupTracker
@@ -52,6 +53,8 @@ fun FirSession.registerCommonComponents(languageVersionSettings: LanguageVersion
    register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.create(this))
    register(FirPredicateBasedProvider::class, FirPredicateBasedProvider.create(this))
    register(GeneratedClassIndex::class, GeneratedClassIndex.create())
+
+    register(FirFakeOverrideStorage::class, FirFakeOverrideStorage(this))
 }

@OptIn(SessionConfiguration::class)
@@ -77,7 +77,7 @@ private fun wrapSubstitutionScopeIfNeed(
        val platformSubstitution = createSubstitution(platformTypeParameters, declaration.defaultType(), session)
        val substitutor = substitutorByMap(platformSubstitution, session)
        FirClassSubstitutionScope(
-            session, useSiteMemberScope, substitutor,
+            session, useSiteMemberScope, PLATFORM_TYPE_PARAMETERS_SUBSTITUTION_SCOPE_KEY, substitutor,
            dispatchReceiverTypeForSubstitutedMembers = derivedClass.defaultType(),
            skipPrivateMembers = true,
        )
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2021 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.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.FirCachesFactory
+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.symbols.impl.*
+
+class FirFakeOverrideStorage(val session: FirSession) : FirSessionComponent {
+    private val cachesFactory = session.firCachesFactory
+
+    val substitutionOverrideCacheByScope: FirCache<ScopeSessionKey<*, *>, SubstitutionOverrideCache, Nothing?> =
+        cachesFactory.createCache { _ -> SubstitutionOverrideCache(session.firCachesFactory) }
+
+    class SubstitutionOverrideCache(cachesFactory: FirCachesFactory) {
+        val overridesForFunctions: FirCache<FirNamedFunctionSymbol, FirNamedFunctionSymbol, FirClassSubstitutionScope> =
+            cachesFactory.createCache { original, scope -> scope.createSubstitutionOverrideFunction(original) }
+        val overridesForConstructors: FirCache<FirConstructorSymbol, FirConstructorSymbol, FirClassSubstitutionScope> =
+            cachesFactory.createCache { original, scope -> scope.createSubstitutionOverrideConstructor(original) }
+        val overridesForVariables: FirCache<FirVariableSymbol<*>, FirVariableSymbol<*>, FirClassSubstitutionScope> =
+            cachesFactory.createCache { original, scope ->
+                when (original) {
+                    is FirPropertySymbol -> scope.createSubstitutionOverrideProperty(original)
+                    is FirFieldSymbol -> scope.createSubstitutionOverrideField(original)
+                    is FirAccessorSymbol -> scope.createSubstitutionOverrideAccessor(original)
+                    else -> error("symbol $original is not overridable")
+                }
+            }
+    }
+}
+
+
+val FirSession.fakeOverrideStorage: FirFakeOverrideStorage by FirSession.sessionComponentAccessor()
@@ -153,11 +153,12 @@ private fun FirClass.scopeForClassImpl(
    val basicScope = unsubstitutedScope(useSiteSession, scopeSession, withForcedTypeCalculator = false)
    if (substitutor == ConeSubstitutor.Empty) return basicScope

+    val key = ConeSubstitutionScopeKey(classFirDispatchReceiver.symbol.toLookupTag(), isFromExpectClass, substitutor)
    return scopeSession.getOrBuild(
-        this, ConeSubstitutionScopeKey(classFirDispatchReceiver.symbol.toLookupTag(), isFromExpectClass, substitutor)
+        this, key
    ) {
        FirClassSubstitutionScope(
-            useSiteSession, basicScope, substitutor, classFirDispatchReceiver.defaultType(),
+            useSiteSession, basicScope, key, substitutor, classFirDispatchReceiver.defaultType(),
            skipPrivateMembers, makeExpect = isFromExpectClass
        )
    }
@@ -7,20 +7,16 @@ package org.jetbrains.kotlin.fir.scopes.impl

 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
-import org.jetbrains.kotlin.fir.declarations.FirTypeParameterRef
-import org.jetbrains.kotlin.fir.declarations.FirTypeParameterRefsOwner
+import org.jetbrains.kotlin.fir.caches.contains
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.dispatchReceiverClassOrNull
 import org.jetbrains.kotlin.fir.originalForSubstitutionOverride
+import org.jetbrains.kotlin.fir.resolve.ScopeSessionKey
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.chain
-import org.jetbrains.kotlin.fir.scopes.FakeOverrideSubstitution
-import org.jetbrains.kotlin.fir.scopes.FirTypeScope
-import org.jetbrains.kotlin.fir.scopes.ProcessorAction
+import org.jetbrains.kotlin.fir.scopes.*
 import org.jetbrains.kotlin.fir.symbols.ensureResolved
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
@@ -33,21 +29,28 @@ import org.jetbrains.kotlin.utils.addToStdlib.runIf
 class FirClassSubstitutionScope(
    private val session: FirSession,
    private val useSiteMemberScope: FirTypeScope,
+    key: ScopeSessionKey<*, *>,
    private val substitutor: ConeSubstitutor,
    private val dispatchReceiverTypeForSubstitutedMembers: ConeClassLikeType,
    private val skipPrivateMembers: Boolean,
    private val makeExpect: Boolean = false
 ) : FirTypeScope() {
+    companion object {
+        private val FirVariableSymbol<*>.isOverridable: Boolean
+            get() = when (this) {
+                is FirPropertySymbol,
+                is FirFieldSymbol,
+                is FirAccessorSymbol -> true
+                else -> false
+            }
+    }

-    private val substitutionOverrideFunctions = mutableMapOf<FirNamedFunctionSymbol, FirNamedFunctionSymbol>()
-    private val substitutionOverrideConstructors = mutableMapOf<FirConstructorSymbol, FirConstructorSymbol>()
-    private val substitutionOverrideVariables = mutableMapOf<FirVariableSymbol<*>, FirVariableSymbol<*>>()
-
+    private val substitutionOverrideCache = session.fakeOverrideStorage.substitutionOverrideCacheByScope.getValue(key, null)
    private val newOwnerClassId = dispatchReceiverTypeForSubstitutedMembers.lookupTag.classId

    override fun processFunctionsByName(name: Name, processor: (FirNamedFunctionSymbol) -> Unit) {
        useSiteMemberScope.processFunctionsByName(name) process@{ original ->
-            val function = substitutionOverrideFunctions.getOrPut(original) { createSubstitutionOverrideFunction(original) }
+            val function = substitutionOverrideCache.overridesForFunctions.getValue(original, this)
            processor(function)
        }

@@ -59,16 +62,18 @@ class FirClassSubstitutionScope(
        processor: (FirNamedFunctionSymbol, FirTypeScope) -> ProcessorAction
    ): ProcessorAction =
        processDirectOverriddenWithBaseScope(
-            functionSymbol, processor, FirTypeScope::processDirectOverriddenFunctionsWithBaseScope, substitutionOverrideFunctions
-        )
+            functionSymbol,
+            processor,
+            FirTypeScope::processDirectOverriddenFunctionsWithBaseScope,
+        ) { it in substitutionOverrideCache.overridesForFunctions }

    private inline fun <reified D : FirCallableSymbol<*>> processDirectOverriddenWithBaseScope(
        callableSymbol: D,
        noinline processor: (D, FirTypeScope) -> ProcessorAction,
        processDirectOverriddenCallablesWithBaseScope: FirTypeScope.(D, ((D, FirTypeScope) -> ProcessorAction)) -> ProcessorAction,
-        fakeOverridesMap: Map<out FirCallableSymbol<*>, FirCallableSymbol<*>>
+        originalInCache: (D) -> Boolean
    ): ProcessorAction {
-        val original = callableSymbol.originalForSubstitutionOverride?.takeIf { it in fakeOverridesMap }
+        val original = callableSymbol.originalForSubstitutionOverride?.takeIf { originalInCache(it) }
            ?: return useSiteMemberScope.processDirectOverriddenCallablesWithBaseScope(callableSymbol, processor)

        if (!processor(original, useSiteMemberScope)) return ProcessorAction.STOP
@@ -78,23 +83,12 @@ class FirClassSubstitutionScope(

    override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
        return useSiteMemberScope.processPropertiesByName(name) process@{ original ->
-            when (original) {
-                is FirPropertySymbol -> {
-                    val property = substitutionOverrideVariables.getOrPut(original) { createSubstitutionOverrideProperty(original) }
-                    processor(property)
-                }
-                is FirFieldSymbol -> {
-                    val field = substitutionOverrideVariables.getOrPut(original) { createSubstitutionOverrideField(original) }
-                    processor(field)
-                }
-                is FirAccessorSymbol -> {
-                    val accessor = substitutionOverrideVariables.getOrPut(original) { createSubstitutionOverrideAccessor(original) }
-                    processor(accessor)
-                }
-                else -> {
-                    processor(original)
-                }
+            val symbol = if (original.isOverridable) {
+                substitutionOverrideCache.overridesForVariables.getValue(original, this)
+            } else {
+                original
            }
+            processor(symbol)
        }
    }

@@ -104,8 +98,7 @@ class FirClassSubstitutionScope(
    ): ProcessorAction =
        processDirectOverriddenWithBaseScope(
            propertySymbol, processor, FirTypeScope::processDirectOverriddenPropertiesWithBaseScope,
-            substitutionOverrideVariables
-        )
+        ) { it in substitutionOverrideCache.overridesForVariables }

    override fun processClassifiersByNameWithSubstitution(name: Name, processor: (FirClassifierSymbol<*>, ConeSubstitutor) -> Unit) {
        useSiteMemberScope.processClassifiersByNameWithSubstitution(name) { symbol, substitutor ->
@@ -121,7 +114,7 @@ class FirClassSubstitutionScope(
        return substitutor.substituteOrNull(this)
    }

-    private fun createSubstitutionOverrideFunction(original: FirNamedFunctionSymbol): FirNamedFunctionSymbol {
+    fun createSubstitutionOverrideFunction(original: FirNamedFunctionSymbol): FirNamedFunctionSymbol {
        if (substitutor == ConeSubstitutor.Empty) return original
        val member = original.fir
        if (skipPrivateMembers && member.visibility == Visibilities.Private) return original
@@ -157,7 +150,7 @@ class FirClassSubstitutionScope(
        )
    }

-    private fun createSubstitutionOverrideConstructor(original: FirConstructorSymbol): FirConstructorSymbol {
+    fun createSubstitutionOverrideConstructor(original: FirConstructorSymbol): FirConstructorSymbol {
        if (substitutor == ConeSubstitutor.Empty) return original
        val constructor = original.fir

@@ -177,7 +170,7 @@ class FirClassSubstitutionScope(
        ).symbol
    }

-    private fun createSubstitutionOverrideProperty(original: FirPropertySymbol): FirPropertySymbol {
+    fun createSubstitutionOverrideProperty(original: FirPropertySymbol): FirPropertySymbol {
        if (substitutor == ConeSubstitutor.Empty) return original
        val member = original.fir
        if (skipPrivateMembers && member.visibility == Visibilities.Private) return original
@@ -231,7 +224,7 @@ class FirClassSubstitutionScope(
        return SubstitutedData(newTypeParameters, newReceiverType, newReturnType, substitutor, fakeOverrideSubstitution)
    }

-    private fun createSubstitutionOverrideField(original: FirFieldSymbol): FirFieldSymbol {
+    fun createSubstitutionOverrideField(original: FirFieldSymbol): FirFieldSymbol {
        if (substitutor == ConeSubstitutor.Empty) return original
        val member = original.fir
        if (skipPrivateMembers && member.visibility == Visibilities.Private) return original
@@ -244,7 +237,7 @@ class FirClassSubstitutionScope(
        return FirFakeOverrideGenerator.createSubstitutionOverrideField(session, member, original, newReturnType, newOwnerClassId)
    }

-    private fun createSubstitutionOverrideAccessor(original: FirAccessorSymbol): FirAccessorSymbol {
+    fun createSubstitutionOverrideAccessor(original: FirAccessorSymbol): FirAccessorSymbol {
        if (substitutor == ConeSubstitutor.Empty) return original
        val member = original.fir as FirSyntheticProperty
        if (skipPrivateMembers && member.visibility == Visibilities.Private) return original
@@ -282,7 +275,7 @@ class FirClassSubstitutionScope(

    override fun processDeclaredConstructors(processor: (FirConstructorSymbol) -> Unit) {
        useSiteMemberScope.processDeclaredConstructors process@{ original ->
-            val constructor = substitutionOverrideConstructors.getOrPut(original) { createSubstitutionOverrideConstructor(original) }
+            val constructor = substitutionOverrideCache.overridesForConstructors.getValue(original, this)
            processor(constructor)
        }
    }
@@ -13,3 +13,7 @@ abstract class FirCache<in K : Any, out V, in CONTEXT> {
@Suppress("NOTHING_TO_INLINE")
 inline fun <K : Any, V> FirCache<K, V, Nothing?>.getValue(key: K): V =
    getValue(key, null)
+
+operator fun <K : Any, V> FirCache<K, V, Nothing>.contains(key: K): Boolean {
+    return getValueIfComputed(key) != null
+}