Make FirRegularTowerDataContexts immutable, SpecialTower.. just a storage

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

@@ -18,30 +18,69 @@ enum class FirTowerDataMode {
     SPECIAL,
 }
 
-class FirRegularTowerDataContexts(
+class FirRegularTowerDataContexts private constructor(
-    forMemberDeclarations: FirTowerDataContext,
+    private val modeMap: EnumMap<FirTowerDataMode, FirTowerDataContext>,
-    forNestedClasses: FirTowerDataContext? = null,
+    val primaryConstructorPureParametersScope: FirLocalScope?,
-    forCompanionObject: FirTowerDataContext? = null,
+    val primaryConstructorAllParametersScope: FirLocalScope?,
-    forConstructorHeaders: FirTowerDataContext? = null,
+    val mode: FirTowerDataMode,
-    forEnumEntries: FirTowerDataContext? = null,
-    val primaryConstructorPureParametersScope: FirLocalScope? = null,
-    val primaryConstructorAllParametersScope: FirLocalScope? = null,
 ) {
-    private val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
+    constructor(
+        forMemberDeclarations: FirTowerDataContext,
+        forNestedClasses: FirTowerDataContext? = null,
+        forCompanionObject: FirTowerDataContext? = null,
+        forConstructorHeaders: FirTowerDataContext? = null,
+        forEnumEntries: FirTowerDataContext? = null,
+        forSpecial: FirTowerDataContext? = null,
+        primaryConstructorPureParametersScope: FirLocalScope? = null,
+        primaryConstructorAllParametersScope: FirLocalScope? = null,
+    ) : this(
+        enumMap(forMemberDeclarations, forNestedClasses, forCompanionObject, forConstructorHeaders, forEnumEntries, forSpecial),
+        primaryConstructorPureParametersScope,
+        primaryConstructorAllParametersScope,
+        FirTowerDataMode.MEMBER_DECLARATION
+    )
 
-    init {
+    val currentContext: FirTowerDataContext?
-        modeMap[FirTowerDataMode.MEMBER_DECLARATION] = forMemberDeclarations
+        get() = modeMap[mode]
-        modeMap[FirTowerDataMode.NESTED_CLASS] = forNestedClasses
+
-        modeMap[FirTowerDataMode.COMPANION_OBJECT] = forCompanionObject
+    fun copy(newContext: FirTowerDataContext): FirRegularTowerDataContexts {
-        modeMap[FirTowerDataMode.CONSTRUCTOR_HEADER] = forConstructorHeaders
+        val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
-        modeMap[FirTowerDataMode.ENUM_ENTRY] = forEnumEntries
+        modeMap.putAll(this.modeMap)
+        modeMap[mode] = newContext
+        return FirRegularTowerDataContexts(modeMap, primaryConstructorPureParametersScope, primaryConstructorAllParametersScope, mode)
     }
 
-    var mode: FirTowerDataMode = FirTowerDataMode.MEMBER_DECLARATION
+    fun copy(newMode: FirTowerDataMode): FirRegularTowerDataContexts {
+        if (newMode == mode) return this
+        return FirRegularTowerDataContexts(modeMap, primaryConstructorPureParametersScope, primaryConstructorAllParametersScope, newMode)
+    }
 
-    var currentContext: FirTowerDataContext?
+    fun copyWithSpecial(newContext: FirTowerDataContext): FirRegularTowerDataContexts {
-        get() = modeMap[mode]
+        val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
-        set(value) {
+        modeMap.putAll(this.modeMap)
-            modeMap[mode] = value
+        modeMap[FirTowerDataMode.SPECIAL] = newContext
+        return FirRegularTowerDataContexts(
+            modeMap, primaryConstructorPureParametersScope, primaryConstructorAllParametersScope, FirTowerDataMode.SPECIAL
+        )
+    }
+
+    companion object {
+        private fun enumMap(
+            forMemberDeclarations: FirTowerDataContext,
+            forNestedClasses: FirTowerDataContext?,
+            forCompanionObject: FirTowerDataContext?,
+            forConstructorHeaders: FirTowerDataContext?,
+            forEnumEntries: FirTowerDataContext?,
+            forSpecial: FirTowerDataContext?,
+        ): EnumMap<FirTowerDataMode, FirTowerDataContext> {
+            val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
+            modeMap[FirTowerDataMode.MEMBER_DECLARATION] = forMemberDeclarations
+            modeMap[FirTowerDataMode.NESTED_CLASS] = forNestedClasses
+            modeMap[FirTowerDataMode.COMPANION_OBJECT] = forCompanionObject
+            modeMap[FirTowerDataMode.CONSTRUCTOR_HEADER] = forConstructorHeaders
+            modeMap[FirTowerDataMode.ENUM_ENTRY] = forEnumEntries
+            modeMap[FirTowerDataMode.SPECIAL] = forSpecial
+            return modeMap
         }
+    }
 }

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

@@ -12,23 +12,38 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 class FirSpecialTowerDataContexts {
     val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext> = mutableMapOf()
     val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext> = mutableMapOf()
-    var currentContext: FirTowerDataContext? = null
 
-    fun setAnonymousFunctionContextIfAny(symbol: FirAnonymousFunctionSymbol): Boolean {
+    fun getAnonymousFunctionContext(symbol: FirAnonymousFunctionSymbol): FirTowerDataContext? {
-        val context = towerDataContextForAnonymousFunctions[symbol]
+        return towerDataContextForAnonymousFunctions[symbol]
-        if (context != null) {
-            currentContext = context
-            return true
-        }
-        return false
     }
 
-    fun setCallableReferenceContextIfAny(access: FirCallableReferenceAccess): Boolean {
+    fun getCallableReferenceContext(access: FirCallableReferenceAccess): FirTowerDataContext? {
-        val context = towerDataContextForCallableReferences[access]
+        return towerDataContextForCallableReferences[access]
-        if (context != null) {
+    }
-            currentContext = context
+
-            return true
+    fun storeAnonymousFunctionContext(symbol: FirAnonymousFunctionSymbol, context: FirTowerDataContext) {
-        }
+        towerDataContextForAnonymousFunctions[symbol] = context
-        return false
+    }
+
+    fun dropAnonymousFunctionContext(symbol: FirAnonymousFunctionSymbol) {
+        towerDataContextForAnonymousFunctions.remove(symbol)
+    }
+
+    fun storeCallableReferenceContext(access: FirCallableReferenceAccess, context: FirTowerDataContext) {
+        towerDataContextForCallableReferences[access] = context
+    }
+
+    fun dropCallableReferenceContext(access: FirCallableReferenceAccess) {
+        towerDataContextForCallableReferences.remove(access)
+    }
+
+    fun putAll(contexts: FirSpecialTowerDataContexts) {
+        towerDataContextForCallableReferences.putAll(contexts.towerDataContextForCallableReferences)
+        towerDataContextForAnonymousFunctions.putAll(contexts.towerDataContextForAnonymousFunctions)
+    }
+
+    fun clear() {
+        towerDataContextForAnonymousFunctions.clear()
+        towerDataContextForCallableReferences.clear()
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -52,32 +52,27 @@ class BodyResolveContext(
     @set:PrivateForInline
     lateinit var file: FirFile
 
+    @PrivateForInline
     var regularTowerDataContexts = FirRegularTowerDataContexts(forMemberDeclarations = FirTowerDataContext())
-        private set
 
+    @PrivateForInline
     val specialTowerDataContexts = FirSpecialTowerDataContexts()
 
+    @OptIn(PrivateForInline::class)
     val towerDataContext: FirTowerDataContext
         get() = regularTowerDataContexts.currentContext
-            ?: specialTowerDataContexts.currentContext
+            ?: throw AssertionError("No regular data context found, towerDataMode = $towerDataMode")
-            ?: throw AssertionError("Not current data context found, towerDataMode = $towerDataMode")
 
     @OptIn(PrivateForInline::class)
     var towerDataMode: FirTowerDataMode
         get() = regularTowerDataContexts.mode
         set(value) {
-            regularTowerDataContexts.mode = value
+            regularTowerDataContexts = regularTowerDataContexts.copy(newMode = value)
         }
 
     val implicitReceiverStack: ImplicitReceiverStack
         get() = towerDataContext.implicitReceiverStack
 
-    private val towerDataContextForAnonymousFunctions: MutableMap<FirAnonymousFunctionSymbol, FirTowerDataContext>
-        get() = specialTowerDataContexts.towerDataContextForAnonymousFunctions
-
-    private val towerDataContextForCallableReferences: MutableMap<FirCallableReferenceAccess, FirTowerDataContext>
-        get() = specialTowerDataContexts.towerDataContextForCallableReferences
-
     @set:PrivateForInline
     var containers: ArrayDeque<FirDeclaration> = ArrayDeque()
 
@@ -97,6 +92,7 @@ class BodyResolveContext(
     val topClassDeclaration: FirRegularClass?
         get() = containingClassDeclarations.lastOrNull()
 
+    @OptIn(PrivateForInline::class)
     private inline fun <T> withNewTowerDataForClass(newContexts: FirRegularTowerDataContexts, f: () -> T): T {
         val old = regularTowerDataContexts
         regularTowerDataContexts = newContexts
@@ -178,17 +174,12 @@ class BodyResolveContext(
 
     @PrivateForInline
     fun replaceTowerDataContext(newContext: FirTowerDataContext) {
-        if (towerDataMode == FirTowerDataMode.SPECIAL) {
+        regularTowerDataContexts = regularTowerDataContexts.copy(newContext)
-            specialTowerDataContexts.currentContext = newContext
-        } else {
-            regularTowerDataContexts.currentContext = newContext
-        }
     }
 
     @PrivateForInline
     fun clear() {
-        towerDataContextForAnonymousFunctions.clear()
+        specialTowerDataContexts.clear()
-        towerDataContextForCallableReferences.clear()
         dataFlowAnalyzerContext.reset()
     }
 
@@ -286,9 +277,11 @@ class BodyResolveContext(
 
     // ANALYSIS PUBLIC API
 
+    @OptIn(PrivateForInline::class)
     fun getPrimaryConstructorPureParametersScope(): FirLocalScope? =
         regularTowerDataContexts.primaryConstructorPureParametersScope
 
+    @OptIn(PrivateForInline::class)
     fun getPrimaryConstructorAllParametersScope(): FirLocalScope? =
         regularTowerDataContexts.primaryConstructorAllParametersScope
 
@@ -317,8 +310,9 @@ class BodyResolveContext(
     @OptIn(PrivateForInline::class)
     inline fun <T> withAnonymousFunctionTowerDataContext(symbol: FirAnonymousFunctionSymbol, f: () -> T): T {
         return withTowerModeCleanup {
-            if (specialTowerDataContexts.setAnonymousFunctionContextIfAny(symbol)) {
+            val newContext = specialTowerDataContexts.getAnonymousFunctionContext(symbol)
-                regularTowerDataContexts.mode = FirTowerDataMode.SPECIAL
+            if (newContext != null) {
+                regularTowerDataContexts = regularTowerDataContexts.copyWithSpecial(newContext)
             }
             f()
         }
@@ -327,15 +321,17 @@ class BodyResolveContext(
     @OptIn(PrivateForInline::class)
     inline fun <T> withCallableReferenceTowerDataContext(access: FirCallableReferenceAccess, f: () -> T): T {
         return withTowerModeCleanup {
-            if (specialTowerDataContexts.setCallableReferenceContextIfAny(access)) {
+            val newContext = specialTowerDataContexts.getCallableReferenceContext(access)
-                regularTowerDataContexts.mode = FirTowerDataMode.SPECIAL
+            if (newContext != null) {
+                regularTowerDataContexts = regularTowerDataContexts.copyWithSpecial(newContext)
             }
             f()
         }
     }
 
+    @OptIn(PrivateForInline::class)
     fun dropContextForAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
-        towerDataContextForAnonymousFunctions.remove(anonymousFunction.symbol)
+        specialTowerDataContexts.dropAnonymousFunctionContext(anonymousFunction.symbol)
     }
 
     @OptIn(PrivateForInline::class)
@@ -346,8 +342,7 @@ class BodyResolveContext(
         BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses, outerLocalClassForNested).apply {
             file = this@BodyResolveContext.file
             fileImportsScope += this@BodyResolveContext.fileImportsScope
-            towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
+            specialTowerDataContexts.putAll(this@BodyResolveContext.specialTowerDataContexts)
-            towerDataContextForCallableReferences.putAll(this@BodyResolveContext.towerDataContextForCallableReferences)
             containers = this@BodyResolveContext.containers
             containingClassDeclarations = ArrayDeque(this@BodyResolveContext.containingClassDeclarations)
             containingClass = this@BodyResolveContext.containingClass
@@ -477,6 +472,7 @@ class BodyResolveContext(
             statics,
             scopeForConstructorHeader,
             scopeForEnumEntries,
+            forSpecial = null,
             primaryConstructorPureParametersScope,
             primaryConstructorAllParametersScope
         )
@@ -573,7 +569,7 @@ class BodyResolveContext(
         f: () -> T
     ): T {
         if (mode !is ResolutionMode.LambdaResolution) {
-            towerDataContextForAnonymousFunctions[anonymousFunction.symbol] = towerDataContext
+            specialTowerDataContexts.storeAnonymousFunctionContext(anonymousFunction.symbol, towerDataContext)
         }
         if (mode is ResolutionMode.ContextDependent || mode is ResolutionMode.ContextDependentDelegate) {
             return f()
@@ -756,12 +752,14 @@ class BodyResolveContext(
         }
     }
 
+    @OptIn(PrivateForInline::class)
     fun storeCallableReferenceContext(callableReferenceAccess: FirCallableReferenceAccess) {
-        towerDataContextForCallableReferences[callableReferenceAccess] = towerDataContext
+        specialTowerDataContexts.storeCallableReferenceContext(callableReferenceAccess, towerDataContext)
     }
 
+    @OptIn(PrivateForInline::class)
     fun dropCallableReferenceContext(callableReferenceAccess: FirCallableReferenceAccess) {
-        towerDataContextForCallableReferences.remove(callableReferenceAccess)
+        specialTowerDataContexts.dropCallableReferenceContext(callableReferenceAccess)
     }
 
     fun <T> withWhenExpression(whenExpression: FirWhenExpression, session: FirSession, f: () -> T): T {