FIR: introduce BodyResolveContext.forBlock, withWhenExpression

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
@@ -659,4 +660,16 @@ class BodyResolveContext(
     fun dropCallableReferenceContext(callableReferenceAccess: FirCallableReferenceAccess) {
         towerDataContextForCallableReferences.remove(callableReferenceAccess)
     }
+
+    inline fun <T> withWhenExpression(whenExpression: FirWhenExpression, crossinline f: () -> T): T {
+        if (whenExpression.subjectVariable == null) return f()
+        return forBlock(f)
+    }
+
+    inline fun <T> forBlock(crossinline f: () -> T): T {
+        return withTowerDataCleanup {
+            addLocalScope(FirLocalScope())
+            f()
+        }
+    }
 }

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

@@ -43,21 +43,6 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
 
     final override val session: FirSession get() = components.session
 
-    protected inline fun <T> withLocalScopeCleanup(crossinline l: () -> T): T {
-        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 open fun needReplacePhase(firDeclaration: FirDeclaration) = true
 
     fun replaceDeclarationResolvePhaseIfNeeded(firDeclaration: FirDeclaration, newResolvePhase: FirResolvePhase) {

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

@@ -41,9 +41,9 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
     }
 
     override fun transformDoWhileLoop(doWhileLoop: FirDoWhileLoop, data: ResolutionMode): CompositeTransformResult<FirStatement> {
-        val context = ResolutionMode.ContextIndependent
         // Do-while has a specific scope structure (its block and condition effectively share the scope)
-        return withNewLocalScope {
+        return context.forBlock {
+            val context = ResolutionMode.ContextIndependent
             doWhileLoop.also(dataFlowAnalyzer::enterDoWhileLoop)
                 .also {
                     transformer.expressionsTransformer.transformBlockInCurrentScope(it.block, context)
@@ -62,10 +62,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         }
         whenExpression.annotations.forEach { it.accept(this, data) }
         dataFlowAnalyzer.enterWhenExpression(whenExpression)
-        return withLocalScopeCleanup with@{
+        return context.withWhenExpression(whenExpression) with@{
-            if (whenExpression.subjectVariable != null) {
-                addNewLocalScope()
-            }
             @Suppress("NAME_SHADOWING")
             var whenExpression = whenExpression.transformSubject(transformer, ResolutionMode.ContextIndependent)
 
@@ -177,7 +174,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
     override fun transformCatch(catch: FirCatch, data: ResolutionMode): CompositeTransformResult<FirCatch> {
         dataFlowAnalyzer.enterCatchClause(catch)
         catch.parameter.transformReturnTypeRef(transformer, ResolutionMode.ContextIndependent)
-        return withNewLocalScope {
+        return context.forBlock {
             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/FirExpressionsResolveTransformer.kt

@@ -302,7 +302,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
     }
 
     override fun transformBlock(block: FirBlock, data: ResolutionMode): CompositeTransformResult<FirStatement> {
-        withNewLocalScope {
+        context.forBlock {
             transformBlockInCurrentScope(block, data)
         }
         return block.compose()
@@ -897,9 +897,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             augmentedArraySetCall.assignCall.explicitReceiver as FirFunctionCall
         )
 
-        val firstResult = withLocalScopeCleanup {
+        val firstResult = augmentedArraySetCall.setGetBlock.transformSingle(transformer, ResolutionMode.ContextIndependent)
-            augmentedArraySetCall.setGetBlock.transformSingle(transformer, ResolutionMode.ContextIndependent)
-        }
         val secondResult = augmentedArraySetCall.assignCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
 
         fun isSuccessful(functionCall: FirFunctionCall): Boolean =