Enums

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

@@ -158,9 +158,6 @@ class JsIrBackendContext(
     private val coroutinePackage = module.getPackage(COROUTINE_PACKAGE_FQNAME)
     private val coroutineIntrinsicsPackage = module.getPackage(COROUTINE_INTRINSICS_PACKAGE_FQNAME)
 
-    val enumEntryToGetInstanceFunction = mutableMapOf<IrEnumEntrySymbol, IrSimpleFunction>()
-    val objectToGetInstanceFunction = mutableMapOf<IrClassSymbol, IrSimpleFunction>()
-    val enumEntryExternalToInstanceField = mutableMapOf<IrEnumEntrySymbol, IrField>()
     val callableReferencesCache = mutableMapOf<CallableReferenceKey, IrSimpleFunction>()
 
     val intrinsics = JsIntrinsics(irBuiltIns, this)

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

@@ -150,10 +150,45 @@ private val enumClassConstructorLoweringPhase = makeJsModulePhase(
     description = "Transform Enum Class into regular Class"
 )
 
-private val enumClassLoweringPhase = makeJsModulePhase(
+private val enumClassConstructorBodyLoweringPhase = makeJsModulePhase(
-    ::EnumClassLowering,
+    ::EnumClassConstructorBodyTransformer,
-    name = "EnumClassLowering",
+    name = "EnumClassConstructorBodyLowering",
-    description = "Transform Enum Class into regular Class",
+    description = "Transform Enum Class into regular Class"
+)
+
+
+private val enumEntryInstancesLoweringPhase = makeJsModulePhase(
+    ::EnumEntryInstancesLowering,
+    name = "EnumEntryInstancesLowering",
+    description = "Create instance variable for each enum entry initialized with `null`",
+    prerequisite = setOf(enumClassConstructorLoweringPhase)
+)
+
+private val enumEntryInstancesBodyLoweringPhase = makeJsModulePhase(
+    ::EnumEntryInstancesBodyLowering,
+    name = "EnumEntryInstancesBodyLowering",
+    description = "Insert enum entry field initialization into correxposnding class constructors",
+    prerequisite = setOf(enumEntryInstancesLoweringPhase)
+)
+
+private val enumClassCreateInitializerLoweringPhase = makeJsModulePhase(
+    ::EnumClassCreateInitializerLowering,
+    name = "EnumClassCreateInitializerLowering",
+    description = "Create initializer for enum entries",
+    prerequisite = setOf(enumClassConstructorLoweringPhase)
+)
+
+private val enumEntryCreateGetInstancesFunsLoweringPhase = makeJsModulePhase(
+    ::EnumEntryCreateGetInstancesFunsLowering,
+    name = "EnumEntryCreateGetInstancesFunsLowering",
+    description = "Create enumEntry_getInstance functions",
+    prerequisite = setOf(enumClassConstructorLoweringPhase)
+)
+
+private val enumSyntheticFunsLoweringPhase = makeJsModulePhase(
+    ::EnumSyntheticFunctionsLowering,
+    name = "EnumSyntheticFunctionsLowering",
+    description = "Implement `valueOf` and `values`",
     prerequisite = setOf(enumClassConstructorLoweringPhase)
 )
 
@@ -161,7 +196,14 @@ private val enumUsageLoweringPhase = makeJsModulePhase(
     ::EnumUsageLowering,
     name = "EnumUsageLowering",
     description = "Replace enum access with invocation of corresponding function",
-    prerequisite = setOf(enumClassLoweringPhase)
+    prerequisite = setOf(enumEntryCreateGetInstancesFunsLoweringPhase)
+)
+
+private val enumEntryRemovalLoweringPhase = makeJsModulePhase(
+    ::EnumClassRemoveEntriesLowering,
+    name = "EnumEntryRemovalLowering",
+    description = "Replace enum entry with corresponding class",
+    prerequisite = setOf(enumUsageLoweringPhase)
 )
 
 private val sharedVariablesLoweringPhase = makeJsModulePhase(
@@ -469,6 +511,7 @@ val jsPhases = namedIrModulePhase(
             lateinitUsageLoweringPhase then
             tailrecLoweringPhase then
             enumClassConstructorLoweringPhase then
+            enumClassConstructorBodyLoweringPhase then
             sharedVariablesLoweringPhase then
             localDelegatedPropertiesLoweringPhase then
             localDeclarationsLoweringPhase then
@@ -483,8 +526,13 @@ val jsPhases = namedIrModulePhase(
             initializersLoweringPhase then
             initializersCleanupLoweringPhase then
             // Common prefix ends
-            enumClassLoweringPhase then
+            enumEntryInstancesLoweringPhase then
+            enumEntryInstancesBodyLoweringPhase then
+            enumClassCreateInitializerLoweringPhase then
+            enumEntryCreateGetInstancesFunsLoweringPhase then
+            enumSyntheticFunsLoweringPhase then
             enumUsageLoweringPhase then
+            enumEntryRemovalLoweringPhase then
             suspendFunctionsLoweringPhase then
             returnableBlockLoweringPhase then
             forLoopsLoweringPhase then

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

@@ -19,4 +19,12 @@ class JsMapping : DefaultMapping() {
     val objectToInstanceField = newMapping<IrClass, IrField>()
     val classToSyntheticPrimaryConstructor = newMapping<IrClass, IrConstructor>()
     val privateMemberToCorrespondingStatic = newMapping<IrFunction, IrSimpleFunction>()
+
+    val enumEntryToGetInstanceFun = newMapping<IrEnumEntry, IrSimpleFunction>()
+    val enumEntryToInstanceField = newMapping<IrEnumEntry, IrField>()
+    val enumConstructorToNewConstructor = newMapping<IrConstructor, IrConstructor>()
+    val enumClassToCorrespondingEnumEntry = newMapping<IrClass, IrEnumEntry>()
+    val enumConstructorOldToNewValueParameters = newMapping<IrValueDeclaration, IrValueParameter>()
+    val enumEntryToCorrespondingField = newMapping<IrEnumEntry, IrField>()
+    val enumClassToInitEntryInstancesFun = newMapping<IrClass, IrSimpleFunction>()
 }

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

@@ -5,16 +5,19 @@
 
 package org.jetbrains.kotlin.ir.backend.js.lower
 
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
-import org.jetbrains.kotlin.backend.common.DeclarationContainerLoweringPass
+import org.jetbrains.kotlin.backend.common.DeclarationTransformer
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.getOrPut
 import org.jetbrains.kotlin.backend.common.ir.copyParameterDeclarationsFrom
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irBlockBody
 import org.jetbrains.kotlin.backend.common.lower.irIfThen
+import org.jetbrains.kotlin.backend.common.lower.parents
+import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Visibilities
-import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.backend.js.JsCommonBackendContext
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.backend.js.utils.toJsArrayLiteral
@@ -26,11 +29,7 @@ import org.jetbrains.kotlin.ir.declarations.impl.IrFieldImpl
 import org.jetbrains.kotlin.ir.descriptors.WrappedClassConstructorDescriptor
 import org.jetbrains.kotlin.ir.descriptors.WrappedFieldDescriptor
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrVarargImpl
-import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
-import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrConstructorSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrFieldSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrType
@@ -40,11 +39,12 @@ import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
-import java.util.*
 
-class EnumUsageLowering(val context: JsIrBackendContext) : FileLoweringPass {
+class EnumUsageLowering(val context: JsIrBackendContext) : BodyLoweringPass {
-    override fun lower(irFile: IrFile) {
+    private var IrEnumEntry.getInstanceFun by context.mapping.enumEntryToGetInstanceFun
-        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        irBody.transformChildrenVoid(object : IrElementTransformerVoid() {
             override fun visitGetEnumValue(expression: IrGetEnumValue): IrExpression {
                 val enumEntry = expression.symbol.owner
                 val klass = enumEntry.parent as IrClass
@@ -54,7 +54,7 @@ class EnumUsageLowering(val context: JsIrBackendContext) : FileLoweringPass {
     }
 
     private fun lowerExternalEnumEntry(enumEntry: IrEnumEntry, klass: IrClass) =
-        context.enumEntryExternalToInstanceField.getOrPut(enumEntry.symbol) { createFieldForEntry(enumEntry, klass) }.let {
+        context.mapping.enumEntryToInstanceField.getOrPut(enumEntry) { createFieldForEntry(enumEntry, klass) }.let {
             JsIrBuilder.buildGetField(it.symbol, classAsReceiver(klass), null, klass.defaultType)
         }
 
@@ -80,25 +80,7 @@ class EnumUsageLowering(val context: JsIrBackendContext) : FileLoweringPass {
     }
 
     private fun lowerEnumEntry(enumEntry: IrEnumEntry, klass: IrClass) =
-        context.enumEntryToGetInstanceFunction.getOrPut(enumEntry.symbol) {
+        enumEntry.getInstanceFun!!.run { JsIrBuilder.buildCall(symbol) }
-            JsIrBuilder.buildFunction(
-                createEntryAccessorName(klass.name.identifier, enumEntry),
-                returnType = enumEntry.getType(klass),
-                parent = klass
-            )
-        }.run { JsIrBuilder.buildCall(symbol) }
-}
-
-class EnumClassLowering(val context: JsIrBackendContext) : DeclarationContainerLoweringPass {
-    override fun lower(irDeclarationContainer: IrDeclarationContainer) {
-        irDeclarationContainer.transformDeclarationsFlat { declaration ->
-            if (declaration is IrClass && declaration.isEnumClass &&
-                !declaration.descriptor.isExpect && !declaration.isEffectivelyExternal()
-            ) {
-                EnumClassTransformer(context, declaration).transform()
-            } else null
-        }
-    }
 }
 
 
@@ -107,52 +89,30 @@ private fun createEntryAccessorName(enumName: String, enumEntry: IrEnumEntry) =
 
 private fun IrEnumEntry.getType(irClass: IrClass) = (correspondingClass ?: irClass).defaultType
 
-class EnumClassConstructorLowering(val context: CommonBackendContext) : DeclarationContainerLoweringPass {
+// Should be applied recursively
-    override fun lower(irDeclarationContainer: IrDeclarationContainer) {
+class EnumClassConstructorLowering(val context: JsCommonBackendContext) : DeclarationTransformer {
-        irDeclarationContainer.transformDeclarationsFlat { declaration ->
-            if (declaration is IrClass && declaration.isEnumClass &&
-                !declaration.descriptor.isExpect && !declaration.isEffectivelyExternal()
-            ) {
-                EnumClassConstructorTransformer(context, declaration).transform()
-            } else null
-        }
-    }
-}
 
-class EnumClassConstructorTransformer(val context: CommonBackendContext, private val irClass: IrClass) {
+    private var IrConstructor.newConstructor by context.mapping.enumConstructorToNewConstructor
-    private val builder = context.createIrBuilder(irClass.symbol)
+    private var IrClass.correspondingEntry by context.mapping.enumClassToCorrespondingEnumEntry
-    private val enumEntries = irClass.declarations.filterIsInstance<IrEnumEntry>()
+    private var IrValueDeclaration.valueParameter by context.mapping.enumConstructorOldToNewValueParameters
-    private val loweredEnumConstructors = HashMap<IrConstructorSymbol, IrConstructor>()
 
-    fun transform(): List<IrDeclaration> {
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
-        // Add `name` and `ordinal` parameters to enum class constructors
+        (declaration.parent as? IrClass)?.let { irClass ->
-        lowerEnumConstructorsSignature()
+            if (!irClass.isEnumClass || irClass.descriptor.isExpect || irClass.isEffectivelyExternal()) return null
 
-        // Pass these parameters to delegating constructor calls
-        lowerEnumConstructorsBody()
-
-        // The first step creates a new `IrConstructor` with new `IrValueParameter`s so references to old `IrValueParameter`s must be replaced with new ones.
-        fixReferencesToConstructorParameters()
-
-        // Lower `IrEnumConstructorCall`s inside of enum entry class constructors to corresponding `IrDelegatingConstructorCall`s.
-        // Add `name` and `ordinal` parameters.
-        lowerEnumEntryClassConstructors()
-
-        // Lower `IrEnumConstructorCall`s to corresponding `IrCall`s.
-        // Add `name` and `ordinal` constant parameters only for calls to the "enum class" constructors ("enum entry class" constructors
-        // already delegate these parameters)
-        lowerEnumEntryInitializerExpression()
-
-        return listOf(irClass)
-    }
-
-    private fun lowerEnumConstructorsSignature() {
-        irClass.declarations.transform { declaration ->
             if (declaration is IrConstructor) {
-                transformEnumConstructor(declaration, irClass)
+                // Add `name` and `ordinal` parameters to enum class constructors
-            } else
+                return listOf(transformEnumConstructor(declaration, irClass))
-                declaration
+            }
+
+            if (declaration is IrEnumEntry) {
+                declaration.correspondingClass?.let { klass ->
+                    klass.correspondingEntry = declaration
+                }
+            }
         }
+
+        return null
     }
 
     private fun transformEnumConstructor(enumConstructor: IrConstructor, enumClass: IrClass): IrConstructor {
@@ -177,17 +137,94 @@ class EnumClassConstructorTransformer(val context: CommonBackendContext, private
             valueParameters += JsIrBuilder.buildValueParameter("name", 0, context.irBuiltIns.stringType).also { it.parent = this }
             valueParameters += JsIrBuilder.buildValueParameter("ordinal", 1, context.irBuiltIns.intType).also { it.parent = this }
             copyParameterDeclarationsFrom(enumConstructor)
-            body = enumConstructor.body
-            loweredEnumConstructors[enumConstructor.symbol] = this
 
-            this.acceptVoid(PatchDeclarationParentsVisitor(enumClass))
+            val newConstructor = this
+            enumConstructor.newConstructor = this
+
+            enumConstructor.body?.let { oldBody ->
+                body = IrBlockBodyImpl(oldBody.startOffset, oldBody.endOffset) {
+                    statements += (oldBody as IrBlockBody).statements
+
+                    context.fixReferencesToConstructorParameters(enumClass, this)
+
+                    acceptVoid(PatchDeclarationParentsVisitor(enumClass))
+
+                    body = this
+                }
+            }
+
+            // TODO except for `fixReferencesToConstructorParameters` this code seems to be obsolete
+            val oldParameters = enumConstructor.valueParameters
+            val newParameters = valueParameters
+            oldParameters.forEach { old ->
+                // TODO Match by index?
+                val new = newParameters.single { it.name == old.name }
+                old.valueParameter = new
+
+                old.defaultValue?.let { default ->
+                    new.defaultValue = IrExpressionBodyImpl(default.startOffset, default.endOffset) {
+                        expression = default.expression
+                        expression.patchDeclarationParents(newConstructor)
+                        context.fixReferencesToConstructorParameters(enumClass, this)
+                    }
+                }
+            }
+        }
+    }
+}
+
+// The first step creates a new `IrConstructor` with new `IrValueParameter`s so references to old `IrValueParameter`s must be replaced with new ones.
+private fun JsCommonBackendContext.fixReferencesToConstructorParameters(irClass: IrClass, body: IrBody) {
+    body.transformChildrenVoid(object : IrElementTransformerVoid() {
+        private val builder = createIrBuilder(irClass.symbol)
+
+        override fun visitGetValue(expression: IrGetValue): IrExpression {
+            mapping.enumConstructorOldToNewValueParameters[expression.symbol.owner]?.let {
+                return builder.irGet(it)
+            }
+
+            return super.visitGetValue(expression)
+        }
+    })
+}
+
+class EnumClassConstructorBodyTransformer(val context: JsCommonBackendContext) : BodyLoweringPass {
+
+    private var IrConstructor.newConstructor by context.mapping.enumConstructorToNewConstructor
+    private var IrClass.correspondingEntry by context.mapping.enumClassToCorrespondingEnumEntry
+
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+
+        (container.parent as? IrClass)?.let { irClass ->
+
+            // TODO Don't apply to everything
+            context.fixReferencesToConstructorParameters(irClass, irBody)
+
+            if (container is IrConstructor) {
+
+                if (irClass.goodEnum) {
+                    // Pass new parameters to delegating constructor calls
+                    lowerEnumConstructorsBody(container)
+                }
+
+                irClass.correspondingEntry?.let { enumEntry ->
+                    // Lower `IrEnumConstructorCall`s inside of enum entry class constructors to corresponding `IrDelegatingConstructorCall`s.
+                    // Add `name` and `ordinal` parameters.
+                    lowerEnumEntryClassConstructors(irClass, enumEntry, container)
+                }
+            }
+
+            if (container is IrEnumEntry) {
+                // Lower `IrEnumConstructorCall`s to corresponding `IrCall`s.
+                // Add `name` and `ordinal` constant parameters only for calls to the "enum class" constructors ("enum entry class" constructors
+                // already delegate these parameters)
+                lowerEnumEntryInitializerExpression(irClass, container)
+            }
         }
     }
 
-    private fun lowerEnumConstructorsBody() {
+    private fun lowerEnumConstructorsBody(constructor: IrConstructor) {
-        irClass.declarations.filterIsInstance<IrConstructor>().forEach {
+        IrEnumClassConstructorTransformer(constructor).transformBody()
-            IrEnumClassConstructorTransformer(it).transformBody()
-        }
     }
 
     private inner class IrEnumClassConstructorTransformer(val constructor: IrConstructor) : IrElementTransformerVoid() {
@@ -205,11 +242,7 @@ class EnumClassConstructorTransformer(val context: CommonBackendContext, private
             }
 
         override fun visitDelegatingConstructorCall(expression: IrDelegatingConstructorCall): IrExpression {
-            var delegatingConstructor = expression.symbol.owner
+            val delegatingConstructor = expression.symbol.owner. let { it.newConstructor ?: it }
-            val constructorWasTransformed = delegatingConstructor.symbol in loweredEnumConstructors
-
-            if (constructorWasTransformed)
-                delegatingConstructor = loweredEnumConstructors[delegatingConstructor.symbol]!!
 
             return builder.irDelegatingConstructorCall(delegatingConstructor).apply {
                 var valueArgIdx = 0
@@ -223,41 +256,25 @@ class EnumClassConstructorTransformer(val context: CommonBackendContext, private
         }
     }
 
-    private fun fixReferencesToConstructorParameters() {
-        val fromOldToNewParameter = mutableMapOf<IrValueParameterSymbol, IrValueParameter>()
 
-        loweredEnumConstructors.forEach { (oldCtorSymbol, newCtor) ->
+    private fun lowerEnumEntryClassConstructors(irClass: IrClass, entry: IrEnumEntry, constructor: IrConstructor) {
-            val oldParameters = oldCtorSymbol.owner.valueParameters
+        constructor.transformChildrenVoid(IrEnumEntryClassConstructorTransformer(irClass, entry, true))
-            val newParameters = newCtor.valueParameters
-
-            oldParameters.forEach { old ->
-                fromOldToNewParameter[old.symbol] = newParameters.single { it.name == old.name }
-            }
-        }
-
-        irClass.transformChildrenVoid(object : IrElementTransformerVoid() {
-            override fun visitGetValue(expression: IrGetValue): IrExpression {
-                fromOldToNewParameter[expression.symbol]?.let {
-                    return builder.irGet(it)
-                }
-
-                return super.visitGetValue(expression)
-            }
-        })
     }
 
-
+    private inner class IrEnumEntryClassConstructorTransformer(
-    private fun lowerEnumEntryClassConstructors() {
+        val irClass: IrClass,
-        for (entry in enumEntries) {
+        val entry: IrEnumEntry,
-            entry.correspondingClass?.constructors?.forEach {
+        val isInsideConstructor: Boolean
-                it.transformChildrenVoid(IrEnumEntryClassConstructorTransformer(entry, true))
+    ) :
-            }
-        }
-    }
-
-    private inner class IrEnumEntryClassConstructorTransformer(val entry: IrEnumEntry, val isInsideConstructor: Boolean) :
         IrElementTransformerVoid() {
 
+        private val enumEntries = irClass.declarations.filterIsInstance<IrEnumEntry>()
+
+        private val builder = context.createIrBuilder(irClass.symbol)
+
+        private fun IrEnumEntry.getNameExpression() = builder.irString(this.name.identifier)
+        private fun IrEnumEntry.getOrdinalExpression() = builder.irInt(enumEntries.indexOf(this))
+
         private fun buildConstructorCall(constructor: IrConstructor) =
             if (isInsideConstructor)
                 builder.irDelegatingConstructorCall(constructor)
@@ -266,11 +283,11 @@ class EnumClassConstructorTransformer(val context: CommonBackendContext, private
 
         override fun visitEnumConstructorCall(expression: IrEnumConstructorCall): IrExpression {
             var constructor = expression.symbol.owner
-            val constructorWasTransformed = constructor.symbol in loweredEnumConstructors
+            val constructorWasTransformed = constructor.newConstructor != null
 
             // Enum entry class constructors are not transformed
             if (constructorWasTransformed)
-                constructor = loweredEnumConstructors[constructor.symbol]!!
+                constructor = constructor.newConstructor!!
 
             return buildConstructorCall(constructor).apply {
                 var valueArgIdx = 0
@@ -287,64 +304,209 @@ class EnumClassConstructorTransformer(val context: CommonBackendContext, private
         }
     }
 
-    private fun lowerEnumEntryInitializerExpression() {
+    private fun lowerEnumEntryInitializerExpression(irClass: IrClass, entry: IrEnumEntry) {
-        for (entry in enumEntries) {
+        entry.initializerExpression =
-            entry.initializerExpression =
+            entry.initializerExpression?.transform(IrEnumEntryClassConstructorTransformer(irClass, entry, false), null)
-                entry.initializerExpression?.transform(IrEnumEntryClassConstructorTransformer(entry, false), null)
-        }
     }
-
-    private fun IrEnumEntry.getNameExpression() = builder.irString(this.name.identifier)
-    private fun IrEnumEntry.getOrdinalExpression() = builder.irInt(enumEntries.indexOf(this))
 }
 
 //-------------------------------------------------------
 
+private val IrClass.goodEnum: Boolean
+    get() = isEnumClass && !descriptor.isExpect && !isEffectivelyExternal()
 
-class EnumClassTransformer(val context: JsIrBackendContext, private val irClass: IrClass) {
+class EnumEntryInstancesLowering(val context: JsIrBackendContext) : DeclarationTransformer {
-    private val builder = context.createIrBuilder(irClass.symbol)
-    private val enumEntries = irClass.declarations.filterIsInstance<IrEnumEntry>()
-    private val enumName = irClass.name.identifier
-    private val throwISESymbol = context.throwISEsymbol
 
-    fun transform(): List<IrDeclaration> {
+    private var IrEnumEntry.correspondingField by context.mapping.enumEntryToCorrespondingField
 
-        // Create instance variable for each enum entry initialized with `null`
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
-        val entryInstances = createEnumEntryInstanceVariables()
+        if (declaration is IrEnumEntry) {
-
+            val irClass = declaration.parentAsClass
-        // Create boolean flag that indicates if entry instances were initialized.
+            if (irClass.goodEnum) {
-        val entryInstancesInitializedVar = createEntryInstancesInitializedVar()
+                // Create instance variable for each enum entry initialized with `null`
-
+                return listOf(declaration, createEnumEntryInstanceVariable(irClass, declaration))
-        // Create function that initializes all enum entry instances using `IrEnumEntry.initializationExpression`.
+            }
-        // It should be called on the first `IrGetEnumValue`, consecutive calls to this function will do nothing.
+        }
-        val initEntryInstancesFun = createInitEntryInstancesFun(entryInstancesInitializedVar, entryInstances)
+        return null
-
-        // Create entry instance getters. These are used to lower `IrGetEnumValue`.
-        val entryGetInstanceFuns = createGetEntryInstanceFuns(initEntryInstancesFun, entryInstances)
-
-        // Create body for `values` and `valueOf` functions
-        lowerSyntheticFunctions()
-
-        // Remove IrEnumEntry nodes from class declarations. Replace them with corresponding class declarations (if they have them).
-        replaceIrEntriesWithCorrespondingClasses()
-
-        return listOf(irClass) + entryInstances + listOf(entryInstancesInitializedVar, initEntryInstancesFun) + entryGetInstanceFuns
     }
 
+    private fun createEnumEntryInstanceVariable(irClass: IrClass, enumEntry: IrEnumEntry): IrField {
+        val enumName = irClass.name.identifier
 
-    private fun createEnumValueOfBody(): IrBody {
+        val result = buildField {
-        val valueOfFun = findFunctionDescriptorForMemberWithSyntheticBodyKind(IrSyntheticBodyKind.ENUM_VALUEOF)
+            name = Name.identifier("${enumName}_${enumEntry.name.identifier}_instance")
+            type = enumEntry.getType(irClass).makeNullable()
+            isStatic = true
+        }.apply {
+            parent = irClass
+            val builder = context.createIrBuilder(irClass.symbol)
+            initializer = builder.run { irExprBody(irImplicitCast(irNull(), type)) }
+        }
+
+        enumEntry.correspondingField = result
+
+        return result
+    }
+}
+
+class EnumEntryInstancesBodyLowering(val context: JsIrBackendContext) : BodyLoweringPass {
+
+    private var IrEnumEntry.correspondingField by context.mapping.enumEntryToCorrespondingField
+
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        if (container is IrConstructor && container.constructedClass.kind == ClassKind.ENUM_ENTRY) {
+            val entryClass = container.constructedClass
+            val enum = entryClass.parentAsClass
+            if (enum.goodEnum) {
+                val entry = enum.declarations.filterIsInstance<IrEnumEntry>().find { it.correspondingClass === entryClass }!!
+                (irBody as IrBlockBody).statements.add(0, context.createIrBuilder(container.symbol).run {
+                    irSetField(null, entry.correspondingField!!, irGet(entryClass.thisReceiver!!))
+                })
+            }
+        }
+    }
+}
+
+class EnumClassCreateInitializerLowering(val context: JsIrBackendContext) : DeclarationTransformer {
+
+    private var IrEnumEntry.correspondingField by context.mapping.enumEntryToCorrespondingField
+    private var IrClass.initEntryInstancesFun: IrSimpleFunction? by context.mapping.enumClassToInitEntryInstancesFun
+
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        if (declaration is IrClass && declaration.goodEnum) {
+            // Create boolean flag that indicates if entry instances were initialized.
+            val entryInstancesInitializedVar = createEntryInstancesInitializedVar(declaration)
+
+            // Create function that initializes all enum entry instances using `IrEnumEntry.initializationExpression`.
+            // It should be called on the first `IrGetEnumValue`, consecutive calls to this function will do nothing.
+            val initEntryInstancesFun = createInitEntryInstancesFun(declaration, entryInstancesInitializedVar)
+
+            declaration.initEntryInstancesFun = initEntryInstancesFun
+
+            // TODO Why not move to upper level?
+            declaration.declarations += entryInstancesInitializedVar
+            declaration.declarations += initEntryInstancesFun
+
+            return null
+        }
+
+        return null
+    }
+
+    private fun createEntryInstancesInitializedVar(irClass: IrClass): IrField = buildField {
+        val enumName = irClass.name.identifier
+        name = Name.identifier("${enumName}_entriesInitialized")
+        type = context.irBuiltIns.booleanType
+        isStatic = true
+    }.apply {
+        parent = irClass
+        val builder = context.createIrBuilder(irClass.symbol)
+        initializer = builder.run { irExprBody(irBoolean(false)) }
+    }
+
+    private fun createInitEntryInstancesFun(irClass: IrClass, entryInstancesInitializedField: IrField): IrSimpleFunction =
+        buildFunction(irClass, "${irClass.name.identifier}_initEntries", context.irBuiltIns.unitType).also {
+            it.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+                statements += context.createIrBuilder(it.symbol).irBlockBody(it) {
+                    +irIfThen(irGetField(null, entryInstancesInitializedField), irReturnUnit())
+                    +irSetField(null, entryInstancesInitializedField, irBoolean(true))
+
+                    irClass.declarations.filterIsInstance<IrEnumEntry>().forEach { entry ->
+                        entry.correspondingField?.let { instanceField ->
+                            +irSetField(null, instanceField, entry.initializerExpression!!.expression)
+                        }
+                    }
+                }.also {
+                    // entry.initializerExpression can have local declarations
+                    it.acceptVoid(PatchDeclarationParentsVisitor(irClass))
+                }.statements
+            }
+        }
+}
+
+private fun buildFunction(
+    irClass: IrClass,
+    name: String,
+    returnType: IrType
+) = JsIrBuilder.buildFunction(name, returnType, irClass)
+
+class EnumEntryCreateGetInstancesFunsLowering(val context: JsIrBackendContext): DeclarationTransformer {
+
+    private var IrEnumEntry.correspondingField by context.mapping.enumEntryToCorrespondingField
+    private var IrClass.initEntryInstancesFun: IrSimpleFunction? by context.mapping.enumClassToInitEntryInstancesFun
+
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        if (declaration is IrEnumEntry) {
+            val irClass = declaration.parentAsClass
+            if (irClass.goodEnum) {
+                // Create entry instance getters. These are used to lower `IrGetEnumValue`.
+                val entryGetInstanceFun = createGetEntryInstanceFun(irClass, declaration, irClass.initEntryInstancesFun!!)
+
+                // TODO prettify
+                entryGetInstanceFun.parent = irClass.parent
+                (irClass.parent as IrDeclarationContainer).declarations += entryGetInstanceFun
+
+                return listOf(declaration)
+            }
+        }
+
+        return null
+    }
+
+    private fun createGetEntryInstanceFun(irClass: IrClass, enumEntry: IrEnumEntry, initEntryInstancesFun: IrSimpleFunction): IrSimpleFunction {
+
+        return context.mapping.enumEntryToGetInstanceFun.getOrPut(enumEntry) { buildFunction(irClass, createEntryAccessorName(irClass.name.identifier, enumEntry), enumEntry.getType(irClass)) }
+            .also {
+                it.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+                    statements += context.createIrBuilder(it.symbol).irBlockBody(it) {
+                        +irCall(initEntryInstancesFun)
+                        +irReturn(irGetField(null, enumEntry.correspondingField!!))
+                    }.statements
+                }
+            }
+    }
+}
+
+class EnumSyntheticFunctionsLowering(val context: JsIrBackendContext): DeclarationTransformer {
+
+    private var IrEnumEntry.getInstanceFun by context.mapping.enumEntryToGetInstanceFun
+
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        if (declaration is IrSimpleFunction) {
+            (declaration.body as? IrSyntheticBody)?.let { body ->
+                val kind = body.kind
+
+                declaration.parents.filterIsInstance<IrClass>().firstOrNull { it.goodEnum }?.let { irClass ->
+                    declaration.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+                        statements += when (kind) {
+                            IrSyntheticBodyKind.ENUM_VALUES -> createEnumValuesBody(declaration, irClass)
+                            IrSyntheticBodyKind.ENUM_VALUEOF -> createEnumValueOfBody(declaration, irClass)
+                        }.statements
+                    }
+                }
+            }
+        }
+
+        return null
+    }
+
+    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]
-        val entryInstanceToFunction = context.enumEntryToGetInstanceFunction
 
         return context.createIrBuilder(valueOfFun.symbol).run {
             irBlockBody {
                 +irReturn(
                     irWhen(
                         irClass.defaultType,
-                        enumEntries.map {
+                        irClass.enumEntries.map {
                             irBranch(
-                                irEquals(irString(it.name.identifier), irGet(nameParameter)), irCall(entryInstanceToFunction[it.symbol]!!)
+                                irEquals(irString(it.name.identifier), irGet(nameParameter)), irCall(it.getInstanceFun!!)
                             )
                         } + irElseBranch(irCall(throwISESymbol))
                     )
@@ -353,131 +515,35 @@ class EnumClassTransformer(val context: JsIrBackendContext, private val irClass:
         }
     }
 
-    private fun createEnumValuesBody(): IrBody {
+    private fun List<IrExpression>.toArrayLiteral(arrayType: IrType, elementType: IrType): IrExpression {
-        val valuesFun = findFunctionDescriptorForMemberWithSyntheticBodyKind(IrSyntheticBodyKind.ENUM_VALUES)
+        val irVararg = IrVarargImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, arrayType, elementType, this)
-        val entryInstanceToFunction = context.enumEntryToGetInstanceFunction
 
+        return IrCallImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, arrayType, context.intrinsics.arrayLiteral).apply {
+            putValueArgument(0, irVararg)
+        }
+    }
+
+    private fun createEnumValuesBody(valuesFun: IrFunction, irClass: IrClass): IrBlockBody {
         val backendContext = context
         return context.createIrBuilder(valuesFun.symbol).run {
             irBlockBody {
                 +irReturn(
-                    enumEntries.map { irCall(entryInstanceToFunction[it.symbol]!!) }
+                    irClass.enumEntries.map { irCall(it.getInstanceFun!!) }
                         .toJsArrayLiteral(backendContext, valuesFun.returnType, irClass.defaultType)
                 )
             }
         }
     }
+}
 
-
+// Should be applied recursively
-    private fun lowerEnumEntryClassConstructors(entryInstances: List<IrVariable>) {
+class EnumClassRemoveEntriesLowering(val context: JsIrBackendContext) : DeclarationTransformer {
-        for ((entry, instance) in enumEntries.zip(entryInstances)) {
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
-            entry.correspondingClass?.constructors?.forEach {
+        // Remove IrEnumEntry nodes from class declarations. Replace them with corresponding class declarations (if they have them).
-
+        if (declaration is IrEnumEntry && !declaration.descriptor.isExpect && !declaration.isEffectivelyExternal()) {
-                // Initialize entry instance at the beginning of constructor so it can be used inside constructor body
+            return listOfNotNull(declaration.correspondingClass)
-                (it.body as? IrBlockBody)?.apply {
-                    statements.add(0, context.createIrBuilder(it.symbol).run {
-                        irSetVar(instance.symbol, irGet(entry.correspondingClass!!.thisReceiver!!))
-                    })
-                }
-            }
-        }
-    }
-
-    private fun createEnumEntryInstanceVariables() = enumEntries.map { enumEntry ->
-        val result = buildField {
-            name = Name.identifier("${enumName}_${enumEntry.name.identifier}_instance")
-            type = enumEntry.getType(irClass).makeNullable()
-            isStatic = true
-        }.apply {
-            parent = irClass
-            initializer = builder.run { irExprBody(irImplicitCast(irNull(), type)) }
         }
 
-        enumEntry.correspondingClass?.constructors?.forEach {
+        return null
-            // Initialize entry instance at the beginning of constructor so it can be used inside constructor body
-            (it.body as? IrBlockBody)?.apply {
-                statements.add(0, context.createIrBuilder(it.symbol).run {
-                    irSetField(null, result, irGet(enumEntry.correspondingClass!!.thisReceiver!!))
-                })
-            }
-        }
-
-        result
-    }
-
-    private fun replaceIrEntriesWithCorrespondingClasses() {
-        irClass.transformDeclarationsFlat {
-            listOfNotNull(if (it is IrEnumEntry) it.correspondingClass else it)
-        }
-    }
-
-    private fun lowerSyntheticFunctions() {
-        irClass.transformChildrenVoid(object : IrElementTransformerVoid() {
-            override fun visitSyntheticBody(body: IrSyntheticBody): IrBody {
-                return when (body.kind) {
-                    IrSyntheticBodyKind.ENUM_VALUES -> createEnumValuesBody()
-                    IrSyntheticBodyKind.ENUM_VALUEOF -> createEnumValueOfBody()
-                }
-            }
-        })
-    }
-
-    private fun createGetEntryInstanceFuns(
-        initEntryInstancesFun: IrSimpleFunction,
-        entryInstances: List<IrField>
-    ) = enumEntries.mapIndexed { index, enumEntry ->
-        context.enumEntryToGetInstanceFunction.getOrPut(enumEntry.symbol) {
-            buildFunction(createEntryAccessorName(enumName, enumEntry), enumEntry.getType(irClass))
-        }.apply {
-            body = context.createIrBuilder(symbol).irBlockBody(this) {
-                +irCall(initEntryInstancesFun)
-                +irReturn(irGetField(null, entryInstances[index]))
-            }
-        }
-    }
-
-    private fun createInitEntryInstancesFun(
-        entryInstancesInitializedField: IrField,
-        entryInstances: List<IrField>
-    ) = buildFunction("${enumName}_initEntries") {
-        +irIfThen(irGetField(null, entryInstancesInitializedField), irReturnUnit())
-        +irSetField(null, entryInstancesInitializedField, irBoolean(true))
-        for ((entry, instanceField) in enumEntries.zip(entryInstances)) {
-            +irSetField(null, instanceField, entry.initializerExpression!!.expression)
-        }
-    }.also {
-        // entry.initializerExpression can have local declarations
-        it.acceptVoid(PatchDeclarationParentsVisitor(irClass))
-    }
-
-    private fun createEntryInstancesInitializedVar(): IrField = buildField {
-        name = Name.identifier("${enumName}_entriesInitialized")
-        type = context.irBuiltIns.booleanType
-        isStatic = true
-    }.apply {
-        parent = irClass
-        initializer = builder.run { irExprBody(irBoolean(false)) }
-    }
-
-
-    private fun findFunctionDescriptorForMemberWithSyntheticBodyKind(kind: IrSyntheticBodyKind): IrFunction =
-        irClass.declarations.asSequence().filterIsInstance<IrFunction>()
-            .first {
-                it.body.let { body ->
-                    body is IrSyntheticBody && body.kind == kind
-                }
-            }
-
-    private fun buildFunction(
-        name: String,
-        returnType: IrType = context.irBuiltIns.unitType
-    ) = JsIrBuilder.buildFunction(name, returnType, irClass)
-
-    private fun buildFunction(
-        name: String,
-        returnType: IrType = context.irBuiltIns.unitType,
-        bodyBuilder: IrBlockBodyBuilder.() -> Unit
-    ) = JsIrBuilder.buildFunction(name, returnType, irClass).also {
-        it.body = context.createIrBuilder(it.symbol).irBlockBody(it, bodyBuilder)
     }
 }