FIR: Cleanup local scope tracking

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

@@ -48,12 +48,22 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
         return context.withTowerDataCleanup(l)
     }
 
+    protected inline fun <T> withNewLocalScope(crossinline l: () -> T): T {
+        return context.withTowerDataCleanup {
+            addNewLocalScope()
+            l()
+        }
+    }
+
+    protected fun addNewLocalScope() {
+        context.addLocalScope(FirLocalScope())
+    }
+
     protected fun addLocalScope(localScope: FirLocalScope?) {
         if (localScope == null) return
         context.addLocalScope(localScope)
     }
 
-
     @OptIn(PrivateForInline::class)
     internal inline fun <T> withFullBodyResolve(crossinline l: () -> T): T {
         if (!implicitTypeOnly) return l()

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

@@ -6,8 +6,8 @@
 package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import org.jetbrains.kotlin.fir.FirTargetElement
-import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
+import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyExpressionBlock
@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSyntheticCallGenerator
 import org.jetbrains.kotlin.fir.resolve.transformers.FirWhenExhaustivenessTransformer
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
-import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
@@ -57,7 +56,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         dataFlowAnalyzer.enterWhenExpression(whenExpression)
         return withLocalScopeCleanup with@{
             if (whenExpression.subjectVariable != null) {
-                addLocalScope(FirLocalScope())
+                addNewLocalScope()
             }
             @Suppress("NAME_SHADOWING")
             var whenExpression = whenExpression.transformSubject(transformer, ResolutionMode.ContextIndependent)
@@ -169,8 +168,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
     override fun transformCatch(catch: FirCatch, data: ResolutionMode): CompositeTransformResult<FirCatch> {
         dataFlowAnalyzer.enterCatchClause(catch)
         catch.parameter.transformReturnTypeRef(transformer, ResolutionMode.ContextIndependent)
-        return withLocalScopeCleanup {
+        return withNewLocalScope {
-            addLocalScope(FirLocalScope())
             catch.transformParameter(transformer, ResolutionMode.ContextIndependent)
             catch.transformBlock(transformer, ResolutionMode.ContextDependent)
         }.also { dataFlowAnalyzer.exitCatchClause(it) }.compose()

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

@@ -488,8 +488,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         function: FirFunction<F>,
         data: ResolutionMode
     ): CompositeTransformResult<FirStatement> {
-        return withLocalScopeCleanup {
+        return withNewLocalScope {
-            addLocalScope(FirLocalScope())
             val functionIsNotAnalyzed = transformerPhase != function.resolvePhase
             if (functionIsNotAnalyzed) {
                 dataFlowAnalyzer.enterFunction(function)
@@ -519,7 +518,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         return withLocalScopeCleanup {
             dataFlowAnalyzer.enterInitBlock(anonymousInitializer)
             addLocalScope(context.getPrimaryConstructorParametersScope())
-            addLocalScope(FirLocalScope())
+            addNewLocalScope()
             val result =
                 transformDeclarationContent(anonymousInitializer, ResolutionMode.ContextIndependent).single as FirAnonymousInitializer
             val graph = dataFlowAnalyzer.exitInitBlock(result)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformer.kt

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeContractDescriptionError
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
-import org.jetbrains.kotlin.fir.scopes.impl.FirLocalScope
 import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
@@ -117,8 +116,7 @@ class FirContractResolveTransformer(
             val contractDescription = owner.contractDescription
             require(contractDescription is FirRawContractDescription)
             val valueParameters = owner.valueParameters
-            val contractCall = withLocalScopeCleanup {
+            val contractCall = withNewLocalScope {
-                addLocalScope(FirLocalScope())
                 for (valueParameter in valueParameters) {
                     context.storeVariable(valueParameter)
                 }

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

@@ -18,7 +18,7 @@ import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.Name
 
-class FirLocalScope(
+class FirLocalScope private constructor(
     properties: PersistentMap<Name, FirVariableSymbol<*>>,
     functions: PersistentMultimap<Name, FirFunctionSymbol<*>>,
     classes: PersistentMap<Name, FirRegularClassSymbol>