Simple stage layers + enforce model

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

@@ -112,7 +112,7 @@ fun DeclarationContainerLoweringPass.runOnFilePostfix(irFile: IrFile) {
     this.lower(irFile as IrDeclarationContainer)
 }
 
-fun BodyLoweringPass.runOnFilePostfix(irFile: IrFile, withLocalDeclarations: Boolean = false) {
+fun BodyLoweringPass.runOnFilePostfix(irFile: IrFile, withLocalDeclarations: Boolean = false, allowDeclarationModification: Boolean = false) {
     ArrayList(irFile.declarations).forEach {
         it.accept(object : IrElementVisitor<Unit, IrDeclaration?> {
             override fun visitElement(element: IrElement, data: IrDeclaration?) {
@@ -131,8 +131,13 @@ fun BodyLoweringPass.runOnFilePostfix(irFile: IrFile, withLocalDeclarations: Boo
 
             override fun visitBody(body: IrBody, data: IrDeclaration?) {
                 if (withLocalDeclarations) body.acceptChildren(this, null)
-
-                lower(body, data!!)
+                if (allowDeclarationModification) {
+                    lower(body, data!!)
+                } else {
+                    stageController.bodyLowering {
+                        lower(body, data!!)
+                    }
+                }
             }
         }, null)
     }
@@ -160,7 +165,9 @@ interface DeclarationTransformer: FileLoweringPass {
 }
 
 fun DeclarationTransformer.transformFlatRestricted(declaration: IrDeclaration): List<IrDeclaration>? {
-    return transformFlat(declaration)
+    return stageController.restrictTo(declaration) {
+        transformFlat(declaration)
+    }
 }
 
 fun DeclarationTransformer.toFileLoweringPass(): FileLoweringPass {

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

@@ -366,7 +366,9 @@ fun IrType.remapTypeParameters(
 
 /* Copied from K/N */
 fun IrDeclarationContainer.addChild(declaration: IrDeclaration) {
-    this.declarations += declaration
+    stageController.unrestrictDeclarationListsAccess {
+        this.declarations += declaration
+    }
     declaration.accept(SetDeclarationsParentVisitor, this)
 }
 
@@ -443,6 +445,7 @@ fun IrClass.simpleFunctions() = declarations.flatMap {
     }
 }
 
+
 fun IrClass.createParameterDeclarations() {
     assert(thisReceiver == null)
 

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

@@ -64,7 +64,8 @@ open class InitializersLowering(context: CommonBackendContext) : InitializersLow
 
 abstract class InitializersLoweringBase(open val context: CommonBackendContext) {
     protected fun extractInitializers(irClass: IrClass, filter: (IrDeclaration) -> Boolean) =
-        irClass.declarations.filter(filter).mapNotNull {
+        // TODO What about fields that were added by lowerings? e.g. captured outer class or locals?
+        ArrayList(irClass.declarations).mapNotNull { if (it is IrProperty) it.backingField else it }.filter(filter).mapNotNull {
             when (it) {
                 is IrField -> handleField(irClass, it)
                 is IrAnonymousInitializer -> handleAnonymousInitializer(it)

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

@@ -34,7 +34,9 @@ class InnerClassesLowering(val context: BackendContext) : DeclarationTransformer
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
 
         if (declaration is IrClass && declaration.isInner) {
-            declaration.declarations += context.declarationFactory.getOuterThisField(declaration)
+            stageController.unrestrictDeclarationListsAccess {
+                declaration.declarations += context.declarationFactory.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

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 //This lower takes part of old LocalDeclarationLowering job to pop up local classes from functions
 open class LocalClassPopupLowering(val context: BackendContext) : BodyLoweringPass {
     override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, withLocalDeclarations = true)
+        runOnFilePostfix(irFile, withLocalDeclarations = true, allowDeclarationModification = true)
     }
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {

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

@@ -95,6 +95,10 @@ class LocalDeclarationsLowering(
 ) :
     BodyLoweringPass {
 
+    override fun lower(irFile: IrFile) {
+        runOnFilePostfix(irFile, allowDeclarationModification = true)
+    }
+
     object DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE :
         IrDeclarationOriginImpl("FIELD_FOR_CAPTURED_VALUE", isSynthetic = true)
 

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

@@ -194,12 +194,12 @@ class JsIntrinsics(private val irBuiltIns: IrBuiltIns, val context: JsIrBackendC
     val anyConstructorSymbol = anyClassSymbol.constructors.single()
 
     val jsObjectClassSymbol = getInternalClassWithoutPackage("kotlin.js.JsObject")
-    val jsObjectConstructorSymbol by lazy { jsObjectClassSymbol.constructors.single() }
+    val jsObjectConstructorSymbol by lazy2 { jsObjectClassSymbol.constructors.single() }
 
-    val uByteClassSymbol by lazy { getInternalClassWithoutPackage("kotlin.UByte") }
-    val uShortClassSymbol by lazy { getInternalClassWithoutPackage("kotlin.UShort") }
-    val uIntClassSymbol by lazy { getInternalClassWithoutPackage("kotlin.UInt") }
-    val uLongClassSymbol by lazy { getInternalClassWithoutPackage("kotlin.ULong") }
+    val uByteClassSymbol by lazy2 { getInternalClassWithoutPackage("kotlin.UByte") }
+    val uShortClassSymbol by lazy2 { getInternalClassWithoutPackage("kotlin.UShort") }
+    val uIntClassSymbol by lazy2 { getInternalClassWithoutPackage("kotlin.UInt") }
+    val uLongClassSymbol by lazy2 { getInternalClassWithoutPackage("kotlin.ULong") }
 
     val unreachable = defineUnreachableIntrinsic()
 
@@ -274,16 +274,16 @@ class JsIntrinsics(private val irBuiltIns: IrBuiltIns, val context: JsIrBackendC
 
     // TODO move CharSequence-related stiff to IntrinsifyCallsLowering
     val charSequenceClassSymbol = context.symbolTable.referenceClass(context.getClass(FqName("kotlin.CharSequence")))
-    val charSequenceLengthPropertyGetterSymbol by lazy {
+    val charSequenceLengthPropertyGetterSymbol by lazy2 {
         with(charSequenceClassSymbol.owner.declarations) {
             filterIsInstance<IrProperty>().firstOrNull { it.name.asString() == "length" }?.getter ?:
             filterIsInstance<IrFunction>().first { it.name.asString() == "<get-length>" }
         }.symbol
     }
-    val charSequenceGetFunctionSymbol by lazy {
+    val charSequenceGetFunctionSymbol by lazy2 {
         charSequenceClassSymbol.owner.declarations.filterIsInstance<IrFunction>().single { it.name.asString() == "get" }.symbol
     }
-    val charSequenceSubSequenceFunctionSymbol by lazy {
+    val charSequenceSubSequenceFunctionSymbol by lazy2 {
         charSequenceClassSymbol.owner.declarations.filterIsInstance<IrFunction>().single { it.name.asString() == "subSequence" }.symbol
     }
 

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

@@ -87,7 +87,7 @@ class JsIrBackendContext(
     val declarationLevelJsModules = mutableListOf<IrDeclarationWithName>()
 
     val internalPackageFragmentDescriptor = EmptyPackageFragmentDescriptor(builtIns.builtInsModule, FqName("kotlin.js.internal"))
-    val implicitDeclarationFile by lazy {
+    val implicitDeclarationFile by lazy2 {
         IrFileImpl(object : SourceManager.FileEntry {
             override val name = "<implicitDeclarations>"
             override val maxOffset = UNDEFINED_OFFSET
@@ -281,27 +281,27 @@ class JsIrBackendContext(
     val suiteFun = getFunctions(FqName("kotlin.test.suite")).singleOrNull()?.let { symbolTable.referenceSimpleFunction(it) }
     val testFun = getFunctions(FqName("kotlin.test.test")).singleOrNull()?.let { symbolTable.referenceSimpleFunction(it) }
 
-    val coroutineImplLabelPropertyGetter by lazy { ir.symbols.coroutineImpl.getPropertyGetter("state")!!.owner }
-    val coroutineImplLabelPropertySetter by lazy { ir.symbols.coroutineImpl.getPropertySetter("state")!!.owner }
-    val coroutineImplResultSymbolGetter by lazy { ir.symbols.coroutineImpl.getPropertyGetter("result")!!.owner }
-    val coroutineImplResultSymbolSetter by lazy { ir.symbols.coroutineImpl.getPropertySetter("result")!!.owner }
-    val coroutineImplExceptionPropertyGetter by lazy { ir.symbols.coroutineImpl.getPropertyGetter("exception")!!.owner }
-    val coroutineImplExceptionPropertySetter by lazy { ir.symbols.coroutineImpl.getPropertySetter("exception")!!.owner }
-    val coroutineImplExceptionStatePropertyGetter by lazy { ir.symbols.coroutineImpl.getPropertyGetter("exceptionState")!!.owner }
-    val coroutineImplExceptionStatePropertySetter by lazy { ir.symbols.coroutineImpl.getPropertySetter("exceptionState")!!.owner }
+    val coroutineImplLabelPropertyGetter by lazy2 { ir.symbols.coroutineImpl.getPropertyGetter("state")!!.owner }
+    val coroutineImplLabelPropertySetter by lazy2 { ir.symbols.coroutineImpl.getPropertySetter("state")!!.owner }
+    val coroutineImplResultSymbolGetter by lazy2 { ir.symbols.coroutineImpl.getPropertyGetter("result")!!.owner }
+    val coroutineImplResultSymbolSetter by lazy2 { ir.symbols.coroutineImpl.getPropertySetter("result")!!.owner }
+    val coroutineImplExceptionPropertyGetter by lazy2 { ir.symbols.coroutineImpl.getPropertyGetter("exception")!!.owner }
+    val coroutineImplExceptionPropertySetter by lazy2 { ir.symbols.coroutineImpl.getPropertySetter("exception")!!.owner }
+    val coroutineImplExceptionStatePropertyGetter by lazy2 { ir.symbols.coroutineImpl.getPropertyGetter("exceptionState")!!.owner }
+    val coroutineImplExceptionStatePropertySetter by lazy2 { ir.symbols.coroutineImpl.getPropertySetter("exceptionState")!!.owner }
 
-    val primitiveClassProperties by lazy {
+    val primitiveClassProperties by lazy2 {
         primitiveClassesObject.owner.declarations.filterIsInstance<IrProperty>()
     }
 
-    val primitiveClassFunctionClass by lazy {
+    val primitiveClassFunctionClass by lazy2 {
         primitiveClassesObject.owner.declarations
             .filterIsInstance<IrSimpleFunction>()
             .find { it.name == Name.identifier("functionClass") }!!
     }
 
-    val throwableConstructors by lazy { throwableClass.owner.declarations.filterIsInstance<IrConstructor>().map { it.symbol } }
-    val defaultThrowableCtor by lazy { throwableConstructors.single { !it.owner.isPrimary && it.owner.valueParameters.size == 0 } }
+    val throwableConstructors by lazy2 { throwableClass.owner.declarations.filterIsInstance<IrConstructor>().map { it.symbol } }
+    val defaultThrowableCtor by lazy2 { throwableConstructors.single { !it.owner.isPrimary && it.owner.valueParameters.size == 0 } }
 
     private fun referenceOperators(): Map<Name, MutableMap<IrClassifierSymbol, IrSimpleFunctionSymbol>> {
         val primitiveIrSymbols = irBuiltIns.primitiveIrTypes.map { it.classifierOrFail as IrClassSymbol }
@@ -352,3 +352,5 @@ class JsIrBackendContext(
         print(message)
     }
 }
+
+fun <T> lazy2(fn: () -> T) = lazy { stageController.withInitialIr(fn) }

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

@@ -498,77 +498,82 @@ private val objectUsageLoweringPhase = makeJsModulePhase(
     description = "Transform IrGetObjectValue into instance generator call"
 )
 
+val phaseList: List<CompilerPhase<JsIrBackendContext, IrModuleFragment, IrModuleFragment>> = listOf(
+    scriptRemoveReceiverLowering,
+    validateIrBeforeLowering,
+    testGenerationPhase,
+    expectDeclarationsRemovingPhase,
+    stripTypeAliasDeclarationsPhase,
+    arrayConstructorPhase,
+    functionInliningPhase,
+    copyInlineFunctionBodyLoweringPhase,
+    createScriptFunctionsPhase,
+    provisionalFunctionExpressionPhase,
+    lateinitNullableFieldsPhase,
+    lateinitDeclarationLoweringPhase,
+    lateinitUsageLoweringPhase,
+    tailrecLoweringPhase,
+    enumClassConstructorLoweringPhase,
+    enumClassConstructorBodyLoweringPhase,
+    sharedVariablesLoweringPhase,
+    localDelegatedPropertiesLoweringPhase,
+    localDeclarationsLoweringPhase,
+    localClassExtractionPhase,
+    innerClassesLoweringPhase,
+    innerClassesMemberBodyLoweringPhase,
+    innerClassConstructorCallsLoweringPhase,
+    propertiesLoweringPhase,
+    primaryConstructorLoweringPhase,
+    delegateToPrimaryConstructorLoweringPhase,
+    annotationConstructorLowering,
+    initializersLoweringPhase,
+    initializersCleanupLoweringPhase,
+    // Common prefix ends
+    enumEntryInstancesLoweringPhase,
+    enumEntryInstancesBodyLoweringPhase,
+    enumClassCreateInitializerLoweringPhase,
+    enumEntryCreateGetInstancesFunsLoweringPhase,
+    enumSyntheticFunsLoweringPhase,
+    enumUsageLoweringPhase,
+    enumEntryRemovalLoweringPhase,
+    suspendFunctionsLoweringPhase,
+    suspendLambdasRemovalLoweringPhase,
+    returnableBlockLoweringPhase,
+    forLoopsLoweringPhase,
+    primitiveCompanionLoweringPhase,
+    propertyAccessorInlinerLoweringPhase,
+    foldConstantLoweringPhase,
+    privateMembersLoweringPhase,
+    privateMemberUsagesLoweringPhase,
+    callableReferenceLoweringPhase,
+    defaultArgumentStubGeneratorPhase,
+    defaultArgumentPatchOverridesPhase,
+    defaultParameterInjectorPhase,
+    defaultParameterCleanerPhase,
+    jsDefaultCallbackGeneratorPhase,
+    removeInlineFunctionsWithReifiedTypeParametersLoweringPhase,
+    throwableSuccessorsLoweringPhase,
+    varargLoweringPhase,
+    multipleCatchesLoweringPhase,
+    bridgesConstructionPhase,
+    typeOperatorLoweringPhase,
+    secondaryConstructorLoweringPhase,
+    secondaryFactoryInjectorLoweringPhase,
+    classReferenceLoweringPhase,
+    inlineClassDeclarationLoweringPhase,
+    inlineClassUsageLoweringPhase,
+    autoboxingTransformerPhase,
+    blockDecomposerLoweringPhase,
+    constLoweringPhase,
+    objectDeclarationLoweringPhase,
+    objectUsageLoweringPhase,
+    callsLoweringPhase,
+    validateIrAfterLowering
+)
+
 val jsPhases = namedIrModulePhase(
     name = "IrModuleLowering",
     description = "IR module lowering",
-    lower = scriptRemoveReceiverLowering then
-            validateIrBeforeLowering then
-            testGenerationPhase then
-            expectDeclarationsRemovingPhase then
-            stripTypeAliasDeclarationsPhase then
-            arrayConstructorPhase then
-            functionInliningPhase then
-            copyInlineFunctionBodyLoweringPhase then
-            createScriptFunctionsPhase then
-            provisionalFunctionExpressionPhase then
-            lateinitNullableFieldsPhase then
-            lateinitDeclarationLoweringPhase then
-            lateinitUsageLoweringPhase then
-            tailrecLoweringPhase then
-            enumClassConstructorLoweringPhase then
-            enumClassConstructorBodyLoweringPhase then
-            sharedVariablesLoweringPhase then
-            localDelegatedPropertiesLoweringPhase then
-            localDeclarationsLoweringPhase then
-            localClassExtractionPhase then
-            innerClassesLoweringPhase then
-            innerClassesMemberBodyLoweringPhase then
-            innerClassConstructorCallsLoweringPhase then
-            propertiesLoweringPhase then
-            primaryConstructorLoweringPhase then
-            delegateToPrimaryConstructorLoweringPhase then
-            annotationConstructorLowering then
-            initializersLoweringPhase then
-            initializersCleanupLoweringPhase then
-            // Common prefix ends
-            enumEntryInstancesLoweringPhase then
-            enumEntryInstancesBodyLoweringPhase then
-            enumClassCreateInitializerLoweringPhase then
-            enumEntryCreateGetInstancesFunsLoweringPhase then
-            enumSyntheticFunsLoweringPhase then
-            enumUsageLoweringPhase then
-            enumEntryRemovalLoweringPhase then
-            suspendFunctionsLoweringPhase then
-            suspendLambdasRemovalLoweringPhase then
-            returnableBlockLoweringPhase then
-            forLoopsLoweringPhase then
-            primitiveCompanionLoweringPhase then
-            propertyAccessorInlinerLoweringPhase then
-            foldConstantLoweringPhase then
-            privateMembersLoweringPhase then
-            privateMemberUsagesLoweringPhase then
-            callableReferenceLoweringPhase then
-            defaultArgumentStubGeneratorPhase then
-            defaultArgumentPatchOverridesPhase then
-            defaultParameterInjectorPhase then
-            defaultParameterCleanerPhase then
-            jsDefaultCallbackGeneratorPhase then
-            removeInlineFunctionsWithReifiedTypeParametersLoweringPhase then
-            throwableSuccessorsLoweringPhase then
-            varargLoweringPhase then
-            multipleCatchesLoweringPhase then
-            bridgesConstructionPhase then
-            typeOperatorLoweringPhase then
-            secondaryConstructorLoweringPhase then
-            secondaryFactoryInjectorLoweringPhase then
-            classReferenceLoweringPhase then
-            inlineClassDeclarationLoweringPhase then
-            inlineClassUsageLoweringPhase then
-            autoboxingTransformerPhase then
-            blockDecomposerLoweringPhase then
-            constLoweringPhase then
-            objectDeclarationLoweringPhase then
-            objectUsageLoweringPhase then
-            callsLoweringPhase then
-            validateIrAfterLowering
-)
+    lower = phaseList.drop(1).fold(phaseList[0]) { acc: CompilerPhase<JsIrBackendContext, IrModuleFragment, IrModuleFragment>, phase ->
+        acc.then(phase)
+    })

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

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2010-2020 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.ir.backend.js
+
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.declarations.impl.IrPersistingElementBase
+
+open class MutableController(override var currentStage: Int = 0) : StageController {
+
+    override var bodiesEnabled: Boolean = true
+
+    override fun <T> withStage(stage: Int, fn: () -> T): T {
+        val prevStage = currentStage
+        currentStage = stage
+        try {
+            return fn()
+        } finally {
+            currentStage = prevStage
+        }
+    }
+
+    override fun <T> withInitialIr(block: () -> T): T = restrictionImpl(null) { withStage(0, block) }
+
+    override fun <T> withInitialStateOf(declaration: IrDeclaration, block: () -> T): T = withStage((declaration as? IrPersistingElementBase<*>)?.createdOn ?: 0, block)
+
+    private var restricted: Boolean = false
+
+    private var restrictedToDeclaration: IrDeclaration? = null
+
+    override fun <T> restrictTo(declaration: IrDeclaration, fn: () -> T): T = restrictionImpl(declaration, fn)
+
+    private fun <T> restrictionImpl(declaration: IrDeclaration?, fn: () -> T): T {
+        val prev = restrictedToDeclaration
+        restrictedToDeclaration = declaration
+        val wereBodiesEnabled = bodiesEnabled
+        bodiesEnabled = false
+        val wasRestricted = restricted
+        restricted = true
+        val wereDeclarationListsRestricted = declarationListsRestricted
+        declarationListsRestricted = true
+        try {
+            return fn()
+        } finally {
+            restrictedToDeclaration = prev
+            bodiesEnabled = wereBodiesEnabled
+            restricted = wasRestricted
+            declarationListsRestricted = wereDeclarationListsRestricted
+        }
+    }
+
+    override fun <T> bodyLowering(fn: () -> T): T {
+        val wereBodiesEnabled = bodiesEnabled
+        bodiesEnabled = true
+        val wasRestricted = restricted
+        restricted = true
+        val wereDeclarationListsRestricted = declarationListsRestricted
+        declarationListsRestricted = true
+        try {
+            return fn()
+        } finally {
+            bodiesEnabled = wereBodiesEnabled
+            restricted = wasRestricted
+            declarationListsRestricted = wereDeclarationListsRestricted
+        }
+    }
+
+    override fun canModify(element: IrElement): Boolean {
+        return !restricted || restrictedToDeclaration === element || element is IrPersistingElementBase<*> && element.createdOn == currentStage
+    }
+
+    private var declarationListsRestricted = false
+
+    override fun <T> unrestrictDeclarationListsAccess(fn: () -> T): T {
+        val wereDeclarationListsRestricted = declarationListsRestricted
+        declarationListsRestricted = false
+        try {
+            return fn()
+        } finally {
+            declarationListsRestricted = wereDeclarationListsRestricted
+        }
+    }
+
+    override fun canAccessDeclarationsOf(irClass: IrClass): Boolean {
+        return !declarationListsRestricted || irClass.visibility == Visibilities.LOCAL /*|| currentStage == (irClass as? IrPersistingElementBase<*>)?.createdOn ?: 0*/
+    }
+}

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.ir.backend.js
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analyzer.AbstractAnalyzerWithCompilerReport
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
+import org.jetbrains.kotlin.backend.common.phaser.PhaserState
 import org.jetbrains.kotlin.backend.common.phaser.invokeToplevel
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.ir.backend.js.export.isExported
@@ -16,10 +17,7 @@ import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.IrModuleToJsTransf
 import org.jetbrains.kotlin.ir.backend.js.utils.JsMainFunctionDetector
 import org.jetbrains.kotlin.ir.backend.js.utils.NameTables
 import org.jetbrains.kotlin.ir.backend.js.utils.isJsExport
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
-import org.jetbrains.kotlin.ir.declarations.IrField
-import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.util.ExternalDependenciesGenerator
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
 import org.jetbrains.kotlin.ir.util.generateTypicalIrProviderList
@@ -59,6 +57,8 @@ fun compile(
     generateFullJs: Boolean = true,
     generateDceJs: Boolean = false
 ): CompilerResult {
+    stageController = object : StageController {}
+
     val (moduleFragment: IrModuleFragment, dependencyModules, irBuiltIns, symbolTable, deserializer) =
         loadIr(project, mainModule, analyzer, configuration, allDependencies, friendDependencies)
 
@@ -93,7 +93,14 @@ fun compile(
 
     moveBodilessDeclarationsToSeparatePlace(context, moduleFragment)
 
-    jsPhases.invokeToplevel(phaseConfig, context, moduleFragment)
+    val controller = MutableController()
+    stageController = controller
+
+    val phaserState = PhaserState<IrModuleFragment>()
+    phaseList.forEachIndexed { index, phase ->
+        controller.currentStage = index + 1
+        phase.invoke(phaseConfig, phaserState, context, moduleFragment)
+    }
 
     val transformer = IrModuleToJsTransformer(context, mainFunction, mainArguments)
     return transformer.generateModule(moduleFragment, generateFullJs, generateDceJs)

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

@@ -68,7 +68,9 @@ 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)
-                        (container.parent as IrDeclarationContainer).declarations += initFunction
+                        stageController.unrestrictDeclarationListsAccess {
+                            (container.parent as IrDeclarationContainer).declarations += initFunction
+                        }
                         initFunction
                     } else {
                         container

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

@@ -74,7 +74,11 @@ class EnumUsageLowering(val context: JsIrBackendContext) : BodyLoweringPass {
             ).also {
                 descriptor.bind(it)
                 it.parent = irClass
-                irClass.declarations += it
+
+                // TODO need a way to emerge local declarations from BodyLoweringPass
+                stageController.unrestrictDeclarationListsAccess {
+                    irClass.declarations += it
+                }
             }
         }
     }
@@ -268,7 +272,7 @@ class EnumClassConstructorBodyTransformer(val context: JsCommonBackendContext) :
     ) :
         IrElementTransformerVoid() {
 
-        private val enumEntries = irClass.declarations.filterIsInstance<IrEnumEntry>()
+        private val enumEntries = irClass.enumEntries
 
         private val builder = context.createIrBuilder(irClass.symbol)
 
@@ -384,8 +388,11 @@ class EnumClassCreateInitializerLowering(val context: JsIrBackendContext) : Decl
             declaration.initEntryInstancesFun = initEntryInstancesFun
 
             // TODO Why not move to upper level?
-            declaration.declarations += entryInstancesInitializedVar
-            declaration.declarations += initEntryInstancesFun
+            // TODO Also doesn't fit the transformFlat-ish API
+            stageController.unrestrictDeclarationListsAccess {
+                declaration.declarations += entryInstancesInitializedVar
+                declaration.declarations += initEntryInstancesFun
+            }
 
             return null
         }
@@ -411,7 +418,7 @@ class EnumClassCreateInitializerLowering(val context: JsIrBackendContext) : Decl
                     +irIfThen(irGetField(null, entryInstancesInitializedField), irReturnUnit())
                     +irSetField(null, entryInstancesInitializedField, irBoolean(true))
 
-                    irClass.declarations.filterIsInstance<IrEnumEntry>().forEach { entry ->
+                    irClass.enumEntries.forEach { entry ->
                         entry.correspondingField?.let { instanceField ->
                             +irSetField(null, instanceField, entry.initializerExpression!!.expression)
                         }
@@ -444,9 +451,11 @@ class EnumEntryCreateGetInstancesFunsLowering(val context: JsIrBackendContext):
 
                 // TODO prettify
                 entryGetInstanceFun.parent = irClass.parent
-                (irClass.parent as IrDeclarationContainer).declarations += entryGetInstanceFun
+                stageController.unrestrictDeclarationListsAccess {
+                    (irClass.parent as IrDeclarationContainer).declarations += entryGetInstanceFun
+                }
 
-                return listOf(declaration)
+                return listOf(declaration) // TODO not null?
             }
         }
 
@@ -492,10 +501,6 @@ class EnumSyntheticFunctionsLowering(val context: JsIrBackendContext): Declarati
 
     private val throwISESymbol = context.throwISEsymbol
 
-    // TODO cache
-    private val IrClass.enumEntries: List<IrEnumEntry>
-        get() = declarations.filterIsInstance<IrEnumEntry>()
-
     private fun createEnumValueOfBody(valueOfFun: IrFunction, irClass: IrClass): IrBlockBody {
         val nameParameter = valueOfFun.valueParameters[0]
 
@@ -536,6 +541,9 @@ class EnumSyntheticFunctionsLowering(val context: JsIrBackendContext): Declarati
     }
 }
 
+private val IrClass.enumEntries: List<IrEnumEntry>
+    get() = declarations.filterIsInstance<IrEnumEntry>()
+
 // Should be applied recursively
 class EnumClassRemoveEntriesLowering(val context: JsIrBackendContext) : DeclarationTransformer {
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {

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

@@ -37,6 +37,7 @@ class ObjectDeclarationLowering(
 ) : DeclarationTransformer {
 
     private var IrClass.instanceField by context.mapping.objectToInstanceField
+    private var IrClass.syntheticPrimaryConstructor by context.mapping.classToSyntheticPrimaryConstructor
 
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
         if (declaration !is IrClass || declaration.kind != ClassKind.OBJECT || declaration.isEffectivelyExternal())
@@ -55,7 +56,7 @@ class ObjectDeclarationLowering(
 
         declaration.instanceField = instanceField
 
-        val primaryConstructor = declaration.primaryConstructor!! // TODO find a way to trigger this correctly
+        val primaryConstructor = declaration.primaryConstructor ?: declaration.syntheticPrimaryConstructor!!
 
         getInstanceFun.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
             statements += context.createIrBuilder(getInstanceFun.symbol).irBlockBody(getInstanceFun) {

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.ir.expressions.IrInstanceInitializerCall
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrDelegatingConstructorCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
+import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
@@ -27,7 +28,7 @@ class PrimaryConstructorLowering(context: JsCommonBackendContext) : DeclarationT
 
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
         if (declaration is IrClass) {
-            val constructors = declaration.declarations.filterIsInstance<IrConstructor>()
+            val constructors = declaration.constructors
 
             if (constructors.any { it.isPrimary }) return null
 
@@ -42,10 +43,13 @@ class PrimaryConstructorLowering(context: JsCommonBackendContext) : DeclarationT
     private val unitType = context.irBuiltIns.unitType
 
     private fun createPrimaryConstructor(irClass: IrClass): IrConstructor {
-        val declaration = irClass.addConstructor {
-            origin = SYNTHETIC_PRIMARY_CONSTRUCTOR
-            isPrimary = true
-            visibility = Visibilities.PRIVATE
+        // TODO better API for declaration creation. This case doesn't fit the usual transformFlat-like API.
+        val declaration = stageController.unrestrictDeclarationListsAccess {
+            irClass.addConstructor {
+                origin = SYNTHETIC_PRIMARY_CONSTRUCTOR
+                isPrimary = true
+                visibility = Visibilities.PRIVATE
+            }
         }
 
         declaration.body = irClass.run {

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

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.ir.types.isPrimitiveType
 import org.jetbrains.kotlin.ir.types.isString
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.irCall
+import org.jetbrains.kotlin.ir.util.properties
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 
@@ -51,9 +52,12 @@ class PrimitiveCompanionLowering(val context: JsIrBackendContext) : BodyLowering
         val actualCompanion = getActualPrimitiveCompanion(companion)
             ?: return null
 
-        return actualCompanion.declarations
-            .filterIsInstance<IrSimpleFunction>()
-            .single { it.name == function.name }
+        for (p in actualCompanion.properties) {
+            p.getter?.let { if (it.name == function.name) return it }
+            p.setter?.let { if (it.name == function.name) return it }
+        }
+
+        error("Accessor not found")
     }
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {

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

@@ -23,9 +23,21 @@ interface StageController {
 
     val bodiesEnabled: Boolean get() = true
 
+    fun <T> withInitialIr(block: () -> T): T = block()
+
+    fun <T> withInitialStateOf(declaration: IrDeclaration, block: () -> T): T = block()
+
+    fun <T> restrictTo(declaration: IrDeclaration, fn: () -> T): T = fn()
+
     fun <T> bodyLowering(fn: () -> T): T = fn()
 
     fun canModify(element: IrElement): Boolean = true
 
-    fun canAccessDeclarationsOf(declaration: IrClass): Boolean = true
+    fun <T> unrestrictDeclarationListsAccess(fn: () -> T): T = fn()
+
+    fun canAccessDeclarationsOf(irClass: IrClass): Boolean = true
+}
+
+inline fun <T> withInitialIr(noinline fn: () -> T): T {
+    return stageController.withInitialIr(fn)
 }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/lazy/IrLazyTypeAlias.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrTypeAlias
 import org.jetbrains.kotlin.ir.declarations.IrTypeParameter
+import org.jetbrains.kotlin.ir.declarations.withInitialIr
 import org.jetbrains.kotlin.ir.symbols.IrTypeAliasSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.DeclarationStubGenerator
@@ -47,8 +48,10 @@ class IrLazyTypeAlias(
     }
 
     override val expandedType: IrType by lazy {
-        typeTranslator.buildWithScope(this) {
-            descriptor.expandedType.toIrType()
+        withInitialIr {
+            typeTranslator.buildWithScope(this) {
+                descriptor.expandedType.toIrType()
+            }
         }
     }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/lazy/IrLazyTypeParameter.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrTypeParameter
 import org.jetbrains.kotlin.ir.declarations.IrTypeParametersContainer
+import org.jetbrains.kotlin.ir.declarations.withInitialIr
 import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.DeclarationStubGenerator
@@ -57,9 +58,11 @@ class IrLazyTypeParameter(
     }
 
     override val superTypes: MutableList<IrType> by lazy {
-        typeTranslator.buildWithScope(this.parent as IrTypeParametersContainer) {
-            val descriptor = symbol.descriptor
-            descriptor.upperBounds.mapTo(arrayListOf()) { it.toIrType() }
+        withInitialIr {
+            typeTranslator.buildWithScope(this.parent as IrTypeParametersContainer) {
+                val descriptor = symbol.descriptor
+                descriptor.upperBounds.mapTo(arrayListOf()) { it.toIrType() }
+            }
         }
     }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/lazy/lazyUtil.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.ir.declarations.lazy
 
+import org.jetbrains.kotlin.ir.declarations.withInitialIr
 import kotlin.reflect.KProperty
 
 internal fun <T> lazyVar(initializer: () -> T): UnsafeLazyVar<T> = UnsafeLazyVar(initializer)
@@ -17,7 +18,7 @@ internal class UnsafeLazyVar<T>(initializer: () -> T) {
     private val value: T
         get() {
             if (!isInitialized) {
-                _value = initializer!!()
+                withInitialIr { _value = initializer!!() }
                 isInitialized = true
                 initializer = null
             }