ObjectLowering

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

@@ -426,14 +426,14 @@ private val staticMembersLoweringPhase = makeJsModulePhase(
     description = "Move static member declarations to top-level"
 )
 
-private val objectDeclarationLoweringPhase = makeCustomJsModulePhase(
-    { context, module -> ObjectDeclarationLowering(context, context.objectToGetInstanceFunction).lower(module) },
+private val objectDeclarationLoweringPhase = makeJsModulePhase(
+    ::ObjectDeclarationLowering,
     name = "ObjectDeclarationLowering",
     description = "Create lazy object instance generator functions"
 )
 
-private val objectUsageLoweringPhase = makeCustomJsModulePhase(
-    { context, module -> ObjectUsageLowering(context, context.objectToGetInstanceFunction).lower(module) },
+private val objectUsageLoweringPhase = makeJsModulePhase(
+    ::ObjectUsageLowering,
     name = "ObjectUsageLowering",
     description = "Transform IrGetObjectValue into instance generator call"
 )

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

@@ -15,4 +15,6 @@ class JsMapping : DefaultMapping() {
     val originalInnerClassPrimaryConstructorByClass = newMapping<IrClass, IrConstructor>()
     val secondaryConstructorToDelegate = newMapping<IrConstructor, IrSimpleFunction>()
     val secondaryConstructorToFactory = newMapping<IrConstructor, IrSimpleFunction>()
+    val objectToGetInstanceFunction = newMapping<IrClass, IrSimpleFunction>()
+    val objectToInstanceField = newMapping<IrClass, IrField>()
 }

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

@@ -5,94 +5,103 @@
 
 package org.jetbrains.kotlin.ir.backend.js.lower
 
-import org.jetbrains.kotlin.backend.common.CommonBackendContext
-import org.jetbrains.kotlin.backend.common.DeclarationContainerLoweringPass
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
+import org.jetbrains.kotlin.backend.common.DeclarationTransformer
+import org.jetbrains.kotlin.backend.common.getOrPut
 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.descriptors.ClassKind
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
-import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
+import org.jetbrains.kotlin.ir.backend.js.JsCommonBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.buildField
-import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationContainer
-import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrBody
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrGetObjectValue
-import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl
 import org.jetbrains.kotlin.ir.types.makeNullable
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.isEffectivelyExternal
+import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.util.primaryConstructor
-import org.jetbrains.kotlin.ir.util.transformDeclarationsFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 
 class ObjectDeclarationLowering(
-    val context: CommonBackendContext,
-    val objectToGetInstanceFunction: MutableMap<IrClassSymbol, IrSimpleFunction>
-) : DeclarationContainerLoweringPass {
-    override fun lower(irDeclarationContainer: IrDeclarationContainer) {
-        irDeclarationContainer.transformDeclarationsFlat { declaration ->
-            if (declaration !is IrClass || declaration.kind != ClassKind.OBJECT || declaration.isEffectivelyExternal())
-                return@transformDeclarationsFlat null
+    val context: JsCommonBackendContext
+) : DeclarationTransformer {
 
-            val getInstanceFun = getOrCreateGetInstanceFunction(objectToGetInstanceFunction, declaration)
+    private var IrClass.instanceField by context.mapping.objectToInstanceField
 
-            val instanceField = buildField {
-                name = Name.identifier(declaration.name.asString() + "_instance")
-                type = declaration.defaultType.makeNullable()
-                isStatic = true
-            }.apply {
-                parent = declaration.parent
-                initializer = null  // Initialized with 'undefined'
-            }
-            val primaryConstructor = declaration.primaryConstructor!!
-            val body = primaryConstructor.body as IrBlockBody
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        if (declaration !is IrClass || declaration.kind != ClassKind.OBJECT || declaration.isEffectivelyExternal())
+            return null
 
-            // Initialize instance field in the beginning of the constructor because it can be used inside the constructor later
-            val initInstanceField = context.createIrBuilder(primaryConstructor.symbol).buildStatement(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
-                irSetField(null, instanceField, irGet(declaration.thisReceiver!!))
-            }
-            body.statements.add(0, initInstanceField)
+        val getInstanceFun = context.getOrCreateGetInstanceFunction(declaration)
 
-            getInstanceFun.body = context.createIrBuilder(getInstanceFun.symbol).irBlockBody(getInstanceFun) {
+        val instanceField = buildField {
+            name = Name.identifier(declaration.name.asString() + "_instance")
+            type = declaration.defaultType.makeNullable()
+            isStatic = true
+        }.apply {
+            parent = declaration.parent
+            initializer = null  // Initialized with 'undefined'
+        }
+
+        declaration.instanceField = instanceField
+
+        val primaryConstructor = declaration.primaryConstructor!! // TODO find a way to trigger this correctly
+
+        getInstanceFun.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+            statements += context.createIrBuilder(getInstanceFun.symbol).irBlockBody(getInstanceFun) {
                 +irIfThen(
                     irEqualsNull(irGetField(null, instanceField)),
                     // Instance field initialized inside constructor
                     irCallConstructor(primaryConstructor.symbol, emptyList())
                 )
                 +irReturn(irGetField(null, instanceField))
-            }
-
-            listOf(declaration, instanceField, getInstanceFun)
+            }.statements
         }
+
+        return listOf(declaration, instanceField, getInstanceFun)
     }
 }
 
 class ObjectUsageLowering(
-    val context: CommonBackendContext,
-    val objectToGetInstanceFunction: MutableMap<IrClassSymbol, IrSimpleFunction>
-) : FileLoweringPass {
-    override fun lower(irFile: IrFile) {
-        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+    val context: JsCommonBackendContext
+) : BodyLoweringPass {
+
+    private var IrClass.instanceField by context.mapping.objectToInstanceField
+
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        if (container is IrConstructor && container.isPrimary) {
+            val irClass = container.parentAsClass
+            irClass.instanceField?.let { instanceField ->
+                // Initialize instance field in the beginning of the constructor because it can be used inside the constructor later
+                val initInstanceField = context.createIrBuilder(container.symbol).buildStatement(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+                    irSetField(null, instanceField, irGet(irClass.thisReceiver!!))
+                }
+                (irBody as IrBlockBody).statements.add(0, initInstanceField)
+            }
+        }
+
+        irBody.transformChildrenVoid(object : IrElementTransformerVoid() {
             override fun visitGetObjectValue(expression: IrGetObjectValue): IrExpression {
                 val obj: IrClass = expression.symbol.owner
                 if (obj.isEffectivelyExternal()) return expression
-                return JsIrBuilder.buildCall(getOrCreateGetInstanceFunction(objectToGetInstanceFunction, obj).symbol)
+                return JsIrBuilder.buildCall(context.getOrCreateGetInstanceFunction(obj).symbol)
             }
         })
     }
 }
 
-fun getOrCreateGetInstanceFunction(objectToGetInstanceFunction: MutableMap<IrClassSymbol, IrSimpleFunction>, obj: IrClass) =
-    objectToGetInstanceFunction.getOrPut(obj.symbol) {
+fun JsCommonBackendContext.getOrCreateGetInstanceFunction(obj: IrClass) =
+    mapping.objectToGetInstanceFunction.getOrPut(obj) {
         JsIrBuilder.buildFunction(
             obj.name.asString() + "_getInstance",
             returnType = obj.defaultType,

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt

@@ -332,14 +332,14 @@ private val builtInsLoweringPhase = makeWasmModulePhase(
     description = "Lower IR buildins"
 )
 
-private val objectDeclarationLoweringPhase = makeCustomWasmModulePhase(
-    { context, module -> ObjectUsageLowering(context, context.objectToGetInstanceFunction).lower(module) },
+private val objectDeclarationLoweringPhase = makeWasmModulePhase(
+    ::ObjectUsageLowering,
     name = "ObjectDeclarationLowering",
     description = "Create lazy object instance generator functions"
 )
 
-private val objectUsageLoweringPhase = makeCustomWasmModulePhase(
-    { context, module -> ObjectUsageLowering(context, context.objectToGetInstanceFunction).lower(module) },
+private val objectUsageLoweringPhase = makeWasmModulePhase(
+    ::ObjectUsageLowering,
     name = "ObjectUsageLowering",
     description = "Transform IrGetObjectValue into instance generator call"
 )