[JS IR] Cleanup StageController & Lowerings magic a bit

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/Lower.kt

@@ -138,10 +138,9 @@ private class DeclarationContainerLoweringVisitor(
 
 fun BodyLoweringPass.runOnFilePostfix(
     irFile: IrFile,
-    withLocalDeclarations: Boolean = false,
+    withLocalDeclarations: Boolean = false
-    allowDeclarationModification: Boolean = false
 ) {
-    val visitor = BodyLoweringVisitor(this, withLocalDeclarations, allowDeclarationModification)
+    val visitor = BodyLoweringVisitor(this, withLocalDeclarations)
     for (declaration in ArrayList(irFile.declarations)) {
         try {
             declaration.accept(visitor, null)
@@ -158,11 +157,8 @@ fun BodyLoweringPass.runOnFilePostfix(
     }
 }
 
-fun BodyAndScriptBodyLoweringPass.runOnFilePostfix(
+fun BodyAndScriptBodyLoweringPass.runOnFilePostfix(irFile: IrFile) {
-    irFile: IrFile,
+    val visitor = ScriptBodyLoweringVisitor(this)
-    allowDeclarationModification: Boolean = false
-) {
-    val visitor = ScriptBodyLoweringVisitor(this, allowDeclarationModification)
     for (declaration in ArrayList(irFile.declarations)) {
         declaration.accept(visitor, null)
     }
@@ -171,7 +167,6 @@ fun BodyAndScriptBodyLoweringPass.runOnFilePostfix(
 private open class BodyLoweringVisitor(
     private val loweringPass: BodyLoweringPass,
     private val withLocalDeclarations: Boolean,
-    private val allowDeclarationModification: Boolean,
 ) : IrElementVisitor<Unit, IrDeclaration?> {
     override fun visitElement(element: IrElement, data: IrDeclaration?) {
         element.acceptChildren(this, data)
@@ -191,13 +186,7 @@ private open class BodyLoweringVisitor(
         if (withLocalDeclarations) body.acceptChildren(this, null)
         val stageController = data!!.factory.stageController
         stageController.restrictTo(data) {
-            if (allowDeclarationModification) {
+            loweringPass.lower(body, data)
-                loweringPass.lower(body, data)
-            } else {
-                stageController.bodyLowering {
-                    loweringPass.lower(body, data)
-                }
-            }
         }
     }
 
@@ -208,25 +197,15 @@ private open class BodyLoweringVisitor(
 }
 
 private class ScriptBodyLoweringVisitor(
-    private val loweringPass: BodyAndScriptBodyLoweringPass,
+    private val loweringPass: BodyAndScriptBodyLoweringPass
-    private val allowDeclarationModification: Boolean
+) : BodyLoweringVisitor(loweringPass, false) {
-) : BodyLoweringVisitor(loweringPass, false, allowDeclarationModification) {
 
     override fun visitClass(declaration: IrClass, data: IrDeclaration?) {
         declaration.thisReceiver?.accept(this, declaration)
         declaration.typeParameters.forEach { it.accept(this, declaration) }
         ArrayList(declaration.declarations).forEach { it.accept(this, declaration) }
         if (declaration.origin == IrDeclarationOrigin.SCRIPT_CLASS) {
-            val stageController = declaration.factory.stageController
+            loweringPass.lowerScriptBody(declaration, declaration)
-            stageController.restrictTo(declaration) {
-                if (allowDeclarationModification) {
-                    loweringPass.lowerScriptBody(declaration, declaration)
-                } else {
-                    stageController.bodyLowering {
-                        loweringPass.lowerScriptBody(declaration, declaration)
-                    }
-                }
-            }
         }
     }
 }
@@ -267,13 +246,6 @@ interface DeclarationTransformer : FileLoweringPass {
             element.acceptChildrenVoid(this)
         }
 
-        override fun visitBody(body: IrBody) {
-            if (transformer.withLocalDeclarations) {
-                super.visitBody(body)
-            }
-            // else stop
-        }
-
         override fun visitFunction(declaration: IrFunction) {
             declaration.acceptChildrenVoid(this)
 
@@ -322,7 +294,9 @@ interface DeclarationTransformer : FileLoweringPass {
             declaration.typeParameters.forEach { it.accept(this, null) }
             ArrayList(declaration.declarations).forEach { it.accept(this, null) }
 
-            declaration.declarations.transformFlat(transformer::transformFlatRestricted)
+            declaration.declarations.transformFlat {
+                transformer.transformFlatRestricted(it)
+            }
         }
 
         override fun visitScript(declaration: IrScript) {

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt

@@ -356,9 +356,7 @@ fun IrType.remapTypeParameters(
 
 /* Copied from K/N */
 fun IrDeclarationContainer.addChild(declaration: IrDeclaration) {
-    declaration.factory.stageController.unrestrictDeclarationListsAccess {
+    this.declarations += declaration
-        this.declarations += declaration
-    }
     declaration.setDeclarationsParent(this)
 }
 

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/InnerClassesLowering.kt

@@ -33,9 +33,7 @@ class InnerClassesLowering(val context: BackendContext, private val innerClasses
 
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
         if (declaration is IrClass && declaration.isInner) {
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
+            declaration.declarations += innerClassesSupport.getOuterThisField(declaration)
-                declaration.declarations += innerClassesSupport.getOuterThisField(declaration)
-            }
         } else if (declaration is IrConstructor) {
             val irClass = declaration.parentAsClass
             if (!irClass.isInner) return null

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalClassPopupLowering.kt

@@ -20,7 +20,7 @@ open class LocalClassPopupLowering(
     val recordExtractedLocalClasses: BackendContext.(IrClass) -> Unit = {},
 ) : BodyLoweringPass {
     override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, withLocalDeclarations = true, allowDeclarationModification = true)
+        runOnFilePostfix(irFile, withLocalDeclarations = true)
     }
 
     private data class ExtractedLocalClass(

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt

@@ -85,7 +85,7 @@ class LocalDeclarationsLowering(
     BodyLoweringPass {
 
     override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, allowDeclarationModification = true)
+        runOnFilePostfix(irFile)
     }
 
     object DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE :

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/inline/LocalClasses.kt

@@ -34,7 +34,7 @@ import org.jetbrains.kotlin.ir.visitors.*
 
 class LocalClassesInInlineLambdasLowering(val context: CommonBackendContext) : BodyLoweringPass {
     override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, allowDeclarationModification = true)
+        runOnFilePostfix(irFile)
     }
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {
@@ -84,7 +84,7 @@ class LocalClassesInInlineLambdasLowering(val context: CommonBackendContext) : B
 
 class LocalClassesInInlineFunctionsLowering(val context: CommonBackendContext) : BodyLoweringPass {
     override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, allowDeclarationModification = true)
+        runOnFilePostfix(irFile)
     }
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/Dce.kt

@@ -33,18 +33,11 @@ fun eliminateDeadDeclarations(
     removeUnusedAssociatedObjects: Boolean = true,
 ) {
 
-    val allRoots = context.irFactory.stageController.withInitialIr { buildRoots(modules, context) }
+    val allRoots = buildRoots(modules, context)
 
     val usefulDeclarations = usefulDeclarations(allRoots, context, removeUnusedAssociatedObjects)
 
-    context.irFactory.stageController.unrestrictDeclarationListsAccess {
+    processUselessDeclarations(modules, usefulDeclarations, context, removeUnusedAssociatedObjects)
-        processUselessDeclarations(
-            modules,
-            usefulDeclarations,
-            context,
-            removeUnusedAssociatedObjects,
-        )
-    }
 }
 
 private fun IrField.isConstant(): Boolean {
@@ -276,32 +269,30 @@ fun usefulDeclarations(
     }
 
     // use withInitialIr to avoid ConcurrentModificationException in dce-driven lowering when adding roots' nested declarations (members)
-    context.irFactory.stageController.withInitialIr {
+    // Add roots
-        // Add roots
+    roots.forEach {
-        roots.forEach {
+        it.enqueue(null, null, altFromFqn = "<ROOT>")
-            it.enqueue(null, null, altFromFqn = "<ROOT>")
+    }
-        }
 
-        // Add roots' nested declarations
+    // Add roots' nested declarations
-        roots.forEach {
+    roots.forEach {
-            it.acceptVoid(
+        it.acceptVoid(
-                object : IrElementVisitorVoid {
+            object : IrElementVisitorVoid {
-                    override fun visitElement(element: IrElement) {
+                override fun visitElement(element: IrElement) {
-                        element.acceptChildrenVoid(this)
+                    element.acceptChildrenVoid(this)
-                    }
-
-                    override fun visitBody(body: IrBody) {
-                        // Skip
-                    }
-
-                    override fun visitDeclaration(declaration: IrDeclarationBase) {
-                        if (declaration !== it) declaration.enqueue(it, "roots' nested declaration")
-
-                        super.visitDeclaration(declaration)
-                    }
                 }
-            )
+
-        }
+                override fun visitBody(body: IrBody) {
+                    // Skip
+                }
+
+                override fun visitDeclaration(declaration: IrDeclarationBase) {
+                    if (declaration !== it) declaration.enqueue(it, "roots' nested declaration")
+
+                    super.visitDeclaration(declaration)
+                }
+            }
+        )
     }
 
     val toStringMethod =

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/WholeWorldStageController.kt

@@ -5,8 +5,6 @@
 
 package org.jetbrains.kotlin.ir.backend.js
 
-import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
-import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.StageController
 import org.jetbrains.kotlin.ir.util.IdSignature
@@ -20,8 +18,6 @@ class WholeWorldStageController : StageController() {
     override var currentDeclaration: IrDeclaration? = null
     private var index: Int = 0
 
-    var declarationListsRestricted: Boolean = false
-
     override fun <T> restrictTo(declaration: IrDeclaration, fn: () -> T): T {
         val previousCurrentDeclaration = currentDeclaration
         val previousIndex = index
@@ -37,56 +33,20 @@ class WholeWorldStageController : StageController() {
         }
     }
 
-    override fun createSignature(parentSignature: IdSignature): IdSignature {
+    override fun <T> withInitialIr(block: () -> T): T {
-        return IdSignature.LoweredDeclarationSignature(parentSignature, currentStage, index++)
-    }
-
-    override fun <T> withStage(stage: Int, fn: () -> T): T {
         val oldStage = currentStage
-        currentStage = stage
+        currentStage = 0
         val oldCurrentDeclaration = currentDeclaration
         currentDeclaration = null
         return try {
-            fn()
+            block()
         } finally {
             currentStage = oldStage
             currentDeclaration = oldCurrentDeclaration
         }
     }
 
-    override fun <T> withInitialIr(block: () -> T): T {
+    override fun createSignature(parentSignature: IdSignature): IdSignature {
-        return withStage(0) {
+        return IdSignature.LoweredDeclarationSignature(parentSignature, currentStage, index++)
-            declarationRestriction(true, block)
-        }
-    }
-
-    override fun <T> bodyLowering(fn: () -> T): T {
-        val wasRestricted = declarationListsRestricted
-        declarationListsRestricted = true
-
-        return try {
-            fn()
-        } finally {
-            declarationListsRestricted = wasRestricted
-        }
-    }
-
-    override fun <T> unrestrictDeclarationListsAccess(fn: () -> T): T {
-        return declarationRestriction(false, fn)
-    }
-
-    private fun <T> declarationRestriction(value: Boolean, fn: () -> T): T{
-        val wasRestricted = declarationListsRestricted
-        declarationListsRestricted = value
-
-        return try {
-            fn()
-        } finally {
-            declarationListsRestricted = wasRestricted
-        }
-    }
-
-    override fun canAccessDeclarationsOf(irClass: IrClass): Boolean {
-        return !declarationListsRestricted || irClass.visibility == DescriptorVisibilities.LOCAL/* && irClass !in context.extractedLocalClasses*/
     }
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/BlockDecomposerLowering.kt

@@ -72,9 +72,7 @@ abstract class AbstractBlockDecomposerLowering(
                     val lastStatement = newBody.statements.last()
                     val actualParent = if (newBody.statements.size > 1 || lastStatement !is IrReturn || lastStatement.value != expression) {
                         expression = JsIrBuilder.buildCall(initFunction.symbol, expression.type)
-                        context.irFactory.stageController.unrestrictDeclarationListsAccess {
+                        (container.parent as IrDeclarationContainer).declarations += initFunction
-                            (container.parent as IrDeclarationContainer).declarations += initFunction
-                        }
                         initFunction
                     } else {
                         container

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/EnumClassLowering.kt

@@ -364,10 +364,8 @@ class EnumClassCreateInitializerLowering(val context: JsCommonBackendContext) :
 
             // TODO Why not move to upper level?
             // TODO Also doesn't fit the transformFlat-ish API
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
+            declaration.declarations += entryInstancesInitializedVar
-                declaration.declarations += entryInstancesInitializedVar
+            declaration.declarations += initEntryInstancesFun
-                declaration.declarations += initEntryInstancesFun
-            }
 
             return null
         }
@@ -428,9 +426,7 @@ class EnumEntryCreateGetInstancesFunsLowering(val context: JsCommonBackendContex
 
                 // TODO prettify
                 entryGetInstanceFun.parent = irClass.parent
-                context.irFactory.stageController.unrestrictDeclarationListsAccess {
+                (irClass.parent as IrDeclarationContainer).declarations += entryGetInstanceFun
-                    (irClass.parent as IrDeclarationContainer).declarations += entryGetInstanceFun
-                }
 
                 return listOf(declaration) // TODO not null?
             }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/EscapedIdentifiersLowering.kt

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
-import org.jetbrains.kotlin.ir.backend.js.utils.hasStableJsName
 import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/ExternalEnumUsagesLowering.kt

@@ -10,7 +10,10 @@ import org.jetbrains.kotlin.backend.common.getOrPut
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.builders.declarations.buildField
-import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrEnumEntry
+import org.jetbrains.kotlin.ir.declarations.IrField
 import org.jetbrains.kotlin.ir.expressions.IrBody
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrGetEnumValue
@@ -51,10 +54,6 @@ class ExternalEnumUsagesLowering(val context: JsIrBackendContext) : BodyLowering
             isStatic = true
         }.also {
             it.parent = irClass
-
+            irClass.declarations += it
-            // TODO need a way to emerge local declarations from BodyLoweringPass
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
-                irClass.declarations += it
-            }
         }
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/JsAnnotationImplementationLowering.kt

@@ -38,9 +38,7 @@ class JsAnnotationImplementationTransformer(val jsContext: JsIrBackendContext) :
 
     override fun visitClassNew(declaration: IrClass): IrStatement {
         if (declaration.isAnnotationClass) {
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
+            implementGeneratedFunctions(declaration, declaration)
-                implementGeneratedFunctions(declaration, declaration)
-            }
         }
         return super.visitClassNew(declaration)
     }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/JsSingleAbstractMethodLowering.kt

@@ -45,9 +45,7 @@ class JsSingleAbstractMethodLowering(context: CommonBackendContext) : SingleAbst
 
         for (wrapper in cachedImplementations.values + inlineCachedImplementations.values) {
             val parentClass = wrapper.parent as IrDeclarationContainer
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
+            parentClass.declarations += wrapper
-                parentClass.declarations += wrapper
-            }
         }
     }
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/PrimaryConstructorLowering.kt

@@ -43,13 +43,10 @@ class PrimaryConstructorLowering(val context: JsCommonBackendContext) : Declarat
     private val unitType = context.irBuiltIns.unitType
 
     private fun createPrimaryConstructor(irClass: IrClass): IrConstructor {
-        // TODO better API for declaration creation. This case doesn't fit the usual transformFlat-like API.
+        val declaration = irClass.addConstructor {
-        val declaration = context.irFactory.stageController.unrestrictDeclarationListsAccess {
+            origin = SYNTHETIC_PRIMARY_CONSTRUCTOR
-            irClass.addConstructor {
+            isPrimary = true
-                origin = SYNTHETIC_PRIMARY_CONSTRUCTOR
+            visibility = DescriptorVisibilities.PRIVATE
-                isPrimary = true
-                visibility = DescriptorVisibilities.PRIVATE
-            }
         }
 
         declaration.body = irClass.run {

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/PersistentApi.kt

@@ -8,19 +8,9 @@ package org.jetbrains.kotlin.ir.declarations
 import org.jetbrains.kotlin.ir.util.IdSignature
 
 open class StageController(open val currentStage: Int = 0) {
-    open fun <T> withStage(stage: Int, fn: () -> T): T = fn()
-
-    open val bodiesEnabled: Boolean get() = true
-
-    open fun <T> withInitialIr(block: () -> T): T = block()
-
     open fun <T> restrictTo(declaration: IrDeclaration, fn: () -> T): T = fn()
 
-    open fun <T> bodyLowering(fn: () -> T): T = fn()
+    open fun <T> withInitialIr(block: () -> T): T = block()
-
-    open fun <T> unrestrictDeclarationListsAccess(fn: () -> T): T = fn()
-
-    open fun canAccessDeclarationsOf(irClass: IrClass): Boolean = true
 
     // Used in JS IC. Declarations created during lowerings need meaningful signatures.
     open fun createSignature(parentSignature: IdSignature): IdSignature? = null