[K/JS, K/Wasm] Optimize simple objects declaration and usage ^Fixed KT-58797

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

@@ -43,7 +43,7 @@ open class PropertyAccessorInlineLowering(
         private fun canBeInlined(callee: IrSimpleFunction): Boolean {
             val property = callee.correspondingPropertySymbol?.owner ?: return false
 
-            // Some devirtualization required here
+            // Some de-virtualization required here
             if (!property.isSafeToInline(container)) return false
 
             val parent = property.parent

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

@@ -1000,12 +1000,27 @@ private val es6PrimaryConstructorUsageOptimizationLowering = makeBodyLoweringPha
     prerequisite = setOf(es6BoxParameterOptimization, es6PrimaryConstructorOptimizationLowering)
 )
 
+private val purifyObjectInstanceGetters = makeDeclarationTransformerPhase(
+    ::PurifyObjectInstanceGettersLowering,
+    name = "PurifyObjectInstanceGettersLowering",
+    description = "[Optimization] Make object instance getter functions pure whenever it's possible",
+)
+
+private val inlineObjectsWithPureInitialization = makeBodyLoweringPhase(
+    ::InlineObjectsWithPureInitializationLowering,
+    name = "InlineObjectsWithPureInitializationLowering",
+    description = "[Optimization] Inline object instance fields getters whenever it's possible",
+    prerequisite = setOf(purifyObjectInstanceGetters)
+)
+
 val optimizationLoweringList = listOf<Lowering>(
     es6CollectConstructorsWhichNeedBoxParameterLowering,
     es6CollectPrimaryConstructorsWhichCouldBeOptimizedLowering,
     es6BoxParameterOptimization,
     es6PrimaryConstructorOptimizationLowering,
     es6PrimaryConstructorUsageOptimizationLowering,
+    purifyObjectInstanceGetters,
+    inlineObjectsWithPureInitialization
 )
 
 val jsOptimizationPhases = SameTypeNamedCompilerPhase(

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

@@ -35,6 +35,8 @@ class JsMapping : DefaultMapping() {
 
     val suspendArityStore = DefaultDelegateFactory.newDeclarationToDeclarationCollectionMapping<IrClass, Collection<IrSimpleFunction>>()
 
+    val objectsWithPureInitialization = DefaultDelegateFactory.newDeclarationToValueMapping<IrClass, Boolean>()
+
     val inlineFunctionsBeforeInlining = DefaultDelegateFactory.newDeclarationToDeclarationMapping<IrFunction, IrFunction>()
 
     // Wasm mappings

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

@@ -199,7 +199,7 @@ internal class JsUsefulDeclarationProcessor(
         !isExternal && !isExpect && !isBuiltInClass(this)
 
     override fun processConstructedClassDeclaration(declaration: IrDeclaration) {
-        if (declaration in result) return
+        if (declaration.isReachable()) return
 
         super.processConstructedClassDeclaration(declaration)
 

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

@@ -108,7 +108,7 @@ abstract class UsefulDeclarationProcessor(
             contagiousReachableDeclarations.add(this as IrOverridableDeclaration<*>)
         }
 
-        if (this !in result) {
+        if (!isReachable()) {
             result.add(this)
             queue.addLast(this)
 
@@ -181,7 +181,7 @@ abstract class UsefulDeclarationProcessor(
     }
 
     protected open fun processConstructedClassDeclaration(declaration: IrDeclaration) {
-        if (declaration in result) return
+        if (declaration.isReachable()) return
 
         fun IrOverridableDeclaration<*>.findOverriddenContagiousDeclaration(): IrOverridableDeclaration<*>? {
             for (overriddenSymbol in this.overriddenSymbols) {
@@ -287,6 +287,8 @@ abstract class UsefulDeclarationProcessor(
 
         return result
     }
+
+    protected fun IrDeclaration.isReachable(): Boolean = this in result
 }
 
 private data class ReachabilityInfo(

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

@@ -143,14 +143,22 @@ object JsIrBuilder {
     fun buildSetVariable(symbol: IrVariableSymbol, value: IrExpression, type: IrType) =
         IrSetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, type, symbol, value, JsStatementOrigins.SYNTHESIZED_STATEMENT)
 
-    fun buildGetField(symbol: IrFieldSymbol, receiver: IrExpression?, superQualifierSymbol: IrClassSymbol? = null, type: IrType? = null) =
+    fun buildGetField(
+        symbol: IrFieldSymbol,
+        receiver: IrExpression? = null,
+        superQualifierSymbol: IrClassSymbol? = null,
+        type: IrType? = null,
+        startOffset: Int = UNDEFINED_OFFSET,
+        endOffset: Int = UNDEFINED_OFFSET,
+        origin: IrStatementOrigin? = JsStatementOrigins.SYNTHESIZED_STATEMENT
+    ) =
         IrGetFieldImpl(
-            UNDEFINED_OFFSET,
-            UNDEFINED_OFFSET,
+            startOffset,
+            endOffset,
             symbol,
             type ?: symbol.owner.type,
             receiver,
-            JsStatementOrigins.SYNTHESIZED_STATEMENT,
+            origin,
             superQualifierSymbol
         )
 

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

@@ -245,6 +245,9 @@ private tailrec fun IrExpression.isGetUnit(irBuiltIns: IrBuiltIns): Boolean =
                 else -> false
             }
 
+        is IrConstructorCall ->
+            this.type == irBuiltIns.unitType
+
         is IrGetObjectValue ->
             this.symbol == irBuiltIns.unitClass
 

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2023 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.lower
+
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
+import org.jetbrains.kotlin.ir.backend.js.JsCommonBackendContext
+import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
+import org.jetbrains.kotlin.ir.backend.js.utils.isObjectInstanceGetter
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.types.classOrNull
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+
+class InlineObjectsWithPureInitializationLowering(val context: JsCommonBackendContext) : BodyLoweringPass {
+    private val IrClass.instanceField by context.mapping.objectToInstanceField
+    private val IrClass.hasPureInitialization by context.mapping.objectsWithPureInitialization
+
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        irBody.transformChildrenVoid(object : IrElementTransformerVoid() {
+            override fun visitCall(expression: IrCall): IrExpression {
+                if (!expression.symbol.owner.isObjectInstanceGetter()) return super.visitCall(expression)
+                val objectToCreate = expression.symbol.owner.returnType.classOrNull?.owner ?: error("Expect return type of an object getter is an object type")
+                if (objectToCreate.hasPureInitialization != true) return super.visitCall(expression)
+                val instanceFieldForObject = objectToCreate.instanceField ?: error("An instance field for an object should exist")
+                return JsIrBuilder.buildGetField(
+                    instanceFieldForObject.symbol,
+                    startOffset = expression.startOffset,
+                    endOffset = expression.endOffset,
+                    origin = expression.origin
+                )
+            }
+        })
+    }
+}

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.ir.backend.js.lower
 
 import org.jetbrains.kotlin.backend.common.BodyLoweringPass
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
+import org.jetbrains.kotlin.ir.backend.js.JsStatementOrigins
 import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.expressions.IrBody
@@ -28,7 +29,15 @@ class InvokeStaticInitializersLowering(val context: JsIrBackendContext) : BodyLo
 
         val instance = context.mapping.objectToGetInstanceFunction[companionObject] ?: return
 
-        val getInstanceCall = IrCallImpl(irClass.startOffset, irClass.endOffset, context.irBuiltIns.unitType, instance.symbol, 0, 0)
+        val getInstanceCall = IrCallImpl(
+            irClass.startOffset,
+            irClass.endOffset,
+            context.irBuiltIns.unitType,
+            instance.symbol,
+            0,
+            0,
+            JsStatementOrigins.SYNTHESIZED_STATEMENT
+        )
 
         (irBody as IrStatementContainer).statements.add(0, getInstanceCall)
     }

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

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2010-2023 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.lower
+
+import org.jetbrains.kotlin.backend.common.DeclarationTransformer
+import org.jetbrains.kotlin.backend.common.getOrPut
+import org.jetbrains.kotlin.backend.common.ir.isPure
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.backend.js.JsCommonBackendContext
+import org.jetbrains.kotlin.ir.backend.js.JsLoweredDeclarationOrigin
+import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
+import org.jetbrains.kotlin.ir.backend.js.utils.isObjectInstanceField
+import org.jetbrains.kotlin.ir.backend.js.utils.isObjectInstanceGetter
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrExpressionBodyImpl
+import org.jetbrains.kotlin.ir.types.classOrNull
+import org.jetbrains.kotlin.ir.util.*
+
+class PurifyObjectInstanceGettersLowering(val context: JsCommonBackendContext) : DeclarationTransformer {
+    private var IrClass.instanceField by context.mapping.objectToInstanceField
+
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        if (
+            !declaration.isObjectInstanceGetter() &&
+            !declaration.isObjectInstanceField() &&
+            !declaration.isObjectConstructor()
+        ) return null
+
+        return when (declaration) {
+            is IrSimpleFunction -> declaration.purifyObjectGetterIfPossible()
+            is IrField -> declaration.purifyObjectInstanceFieldIfPossible()
+            is IrConstructor -> declaration.removeInstanceFieldInitializationIfPossible()
+            else -> error("Unexpected IR type ${declaration::class.qualifiedName}")
+        }
+    }
+
+    private fun IrConstructor.removeInstanceFieldInitializationIfPossible(): List<IrDeclaration>? {
+        if (parentAsClass.isPureObject()) {
+            (body as? IrBlockBody)?.statements?.removeIf {
+                it is IrSetField && it.symbol.owner.isObjectInstanceField()
+            }
+        }
+
+        return null
+    }
+
+    private fun IrSimpleFunction.purifyObjectGetterIfPossible(): List<IrDeclaration>? {
+        val objectToCreate = returnType.classOrNull?.owner ?: return null
+
+        if (objectToCreate.isPureObject()) {
+            val body = (body as? IrBlockBody) ?: return null
+            val instanceField = objectToCreate.instanceField ?: error("Expect the object instance field to be created")
+
+            body.statements.clear()
+            body.statements += JsIrBuilder.buildReturn(
+                symbol,
+                JsIrBuilder.buildGetField(instanceField.symbol),
+                objectToCreate.defaultType
+            )
+        }
+
+        return null
+    }
+
+    private fun IrField.purifyObjectInstanceFieldIfPossible(): List<IrDeclaration>? {
+        val objectToCreate = type.classOrNull?.owner ?: return null
+
+        if (objectToCreate.isPureObject()) {
+            val objectConstructor = objectToCreate.primaryConstructor ?: error("Object should contain a primary constructor")
+            initializer = IrExpressionBodyImpl(JsIrBuilder.buildConstructorCall(objectConstructor.symbol))
+        }
+
+        return null
+    }
+
+    private fun IrDeclaration.isObjectConstructor(): Boolean {
+        return this is IrConstructor && parentAsClass.isObject
+    }
+
+    private fun IrClass.isPureObject(): Boolean {
+        return context.mapping.objectsWithPureInitialization.getOrPut(this) {
+            superClass == null && primaryConstructor?.body?.statements?.all { it.isPureStatementForObjectInitialization(this@isPureObject) } != false
+        }
+    }
+
+    private fun IrStatement.isPureStatementForObjectInitialization(owner: IrClass): Boolean {
+        return (
+                // Only objects which don't have a class parent
+                (this is IrDelegatingConstructorCall && symbol.owner.parent == context.irBuiltIns.anyClass.owner) ||
+                        (this is IrExpression && isPure(anyVariable = true, checkFields = false, context = context)) ||
+                        (this is IrContainerExpression && statements.all { it.isPureStatementForObjectInitialization(owner) }) ||
+                        (this is IrVariable && initializer?.isPureStatementForObjectInitialization(owner) != false) ||
+                        // Only fields of the objects are safe to not save an intermediate state of another class/object/global
+                        (this is IrGetField && receiver?.isPureStatementForObjectInitialization(owner) == true) ||
+                        (this is IrSetField && receiver?.isPureStatementForObjectInitialization(owner) == true && value.isPureStatementForObjectInitialization(owner)) ||
+                        // Only current object could be initialized inside the object constructor, so we need to ignore it as an effect
+                        (this is IrSetField && symbol.owner.isObjectInstanceField()) ||
+                        (this is IrSetValue && symbol.owner.isLocal && value.isPureStatementForObjectInitialization(owner))
+                )
+
+    }
+}

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

@@ -10,8 +10,10 @@ import org.jetbrains.kotlin.backend.common.phaser.invokeToplevel
 import org.jetbrains.kotlin.ir.backend.js.dce.eliminateDeadDeclarations
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.JsIrProgramFragment
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.js.backend.ast.JsClass
 import org.jetbrains.kotlin.js.backend.ast.JsFunction
 import org.jetbrains.kotlin.js.backend.ast.RecursiveJsVisitor
+import org.jetbrains.kotlin.js.inline.clean.ClassPostProcessor
 import org.jetbrains.kotlin.js.inline.clean.FunctionPostProcessor
 
 fun optimizeProgramByIr(
@@ -24,11 +26,20 @@ fun optimizeProgramByIr(
 }
 
 fun optimizeFragmentByJsAst(fragment: JsIrProgramFragment) {
-    fragment.declarations.statements.forEach {
-        it.accept(object : RecursiveJsVisitor() {
-            override fun visitFunction(x: JsFunction) {
-                FunctionPostProcessor(x).apply()
-            }
-        })
+    val optimizer = object : RecursiveJsVisitor() {
+        override fun visitFunction(x: JsFunction) {
+            FunctionPostProcessor(x).apply()
+        }
+
+        override fun visitClass(x: JsClass) {
+            super.visitClass(x)
+            ClassPostProcessor(x).apply()
+        }
+    }
+
+    fragment.declarations.statements.forEach { it.accept(optimizer) }
+    fragment.classes.values.forEach { klass ->
+        klass.postDeclarationBlock.statements.forEach { it.accept(optimizer)}
+        klass.preDeclarationBlock.statements.forEach { it.accept(optimizer)}
     }
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsExpressionTransformer.kt

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.ir.backend.js.JsStatementOrigins
 import org.jetbrains.kotlin.ir.backend.js.utils.JsGenerationContext
 import org.jetbrains.kotlin.ir.backend.js.utils.Namer
 import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
-import org.jetbrains.kotlin.ir.backend.js.utils.isUnitInstanceFunction
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.types.isString
@@ -130,6 +129,12 @@ class IrElementToJsExpressionTransformer : BaseIrElementToJsNodeTransformer<JsEx
         }
         val fieldName = context.getNameForField(field)
         return jsElementAccess(fieldName, expression.receiver?.accept(this, context)).withSource(expression, context)
+            .apply {
+                if (field.origin == IrDeclarationOrigin.FIELD_FOR_OBJECT_INSTANCE && expression.origin == JsStatementOrigins.SYNTHESIZED_STATEMENT) {
+                    synthetic = true
+                    sideEffects = SideEffectKind.PURE
+                }
+            }
     }
 
     override fun visitGetValue(expression: IrGetValue, context: JsGenerationContext): JsExpression {

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

@@ -112,3 +112,18 @@ fun irEmpty(context: JsIrBackendContext): IrExpression {
     return JsIrBuilder.buildComposite(context.dynamicType, emptyList())
 }
 
+fun IrDeclaration.isObjectInstanceGetter(): Boolean {
+    return this is IrSimpleFunction && isObjectInstanceGetter()
+}
+
+fun IrSimpleFunction.isObjectInstanceGetter(): Boolean {
+    return origin == JsLoweredDeclarationOrigin.OBJECT_GET_INSTANCE_FUNCTION
+}
+
+fun IrDeclaration.isObjectInstanceField(): Boolean {
+    return this is IrField && isObjectInstanceField()
+}
+
+fun IrField.isObjectInstanceField(): Boolean {
+    return origin == IrDeclarationOrigin.FIELD_FOR_OBJECT_INSTANCE
+}

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

@@ -107,6 +107,9 @@ fun IrBody.prependFunctionCall(
 fun JsCommonBackendContext.findUnitGetInstanceFunction(): IrSimpleFunction =
     mapping.objectToGetInstanceFunction[irBuiltIns.unitClass.owner]!!
 
+fun JsCommonBackendContext.findUnitInstanceField(): IrField =
+    mapping.objectToInstanceField[irBuiltIns.unitClass.owner]!!
+
 fun IrDeclaration.isImportedFromModuleOnly(): Boolean {
     return isTopLevel && isEffectivelyExternal() && (getJsModule() != null && !isJsNonModule() || (parent as? IrAnnotationContainer)?.getJsModule() != null)
 }

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

@@ -610,6 +610,20 @@ private val unitToVoidLowering = makeWasmModulePhase(
     description = "Replace some Unit's with Void's"
 )
 
+private val purifyObjectInstanceGettersLoweringPhase = makeWasmModulePhase(
+    ::PurifyObjectInstanceGettersLowering,
+    name = "PurifyObjectInstanceGettersLowering",
+    description = "[Optimization] Make object instance getter functions pure whenever it's possible",
+    prerequisite = setOf(objectDeclarationLoweringPhase, objectUsageLoweringPhase)
+)
+
+private val inlineObjectsWithPureInitializationLoweringPhase = makeWasmModulePhase(
+    ::InlineObjectsWithPureInitializationLowering,
+    name = "InlineObjectsWithPureInitializationLowering",
+    description = "[Optimization] Inline object instance fields getters whenever it's possible",
+    prerequisite = setOf(purifyObjectInstanceGettersLoweringPhase)
+)
+
 val wasmPhases = SameTypeNamedCompilerPhase(
     name = "IrModuleLowering",
     description = "IR module lowering",
@@ -707,7 +721,6 @@ val wasmPhases = SameTypeNamedCompilerPhase(
             associatedObjectsLowering then
 
             objectDeclarationLoweringPhase then
-            fieldInitializersLoweringPhase then
             genericReturnTypeLowering then
             unitToVoidLowering then
 
@@ -715,8 +728,12 @@ val wasmPhases = SameTypeNamedCompilerPhase(
             builtInsLoweringPhase0 then
 
             autoboxingTransformerPhase then
-            explicitlyCastExternalTypesPhase then
+
             objectUsageLoweringPhase then
+            purifyObjectInstanceGettersLoweringPhase then
+            fieldInitializersLoweringPhase then
+
+            explicitlyCastExternalTypesPhase then
             typeOperatorLoweringPhase then
 
             // Clean up built-ins after type operator lowering
@@ -724,5 +741,6 @@ val wasmPhases = SameTypeNamedCompilerPhase(
 
             virtualDispatchReceiverExtractionPhase then
             staticMembersLoweringPhase then
+            inlineObjectsWithPureInitializationLoweringPhase then
             validateIrAfterLowering
 )

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/dce/Dce.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.backend.js.utils.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.util.primaryConstructor
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
@@ -30,7 +31,7 @@ fun eliminateDeadDeclarations(modules: List<IrModuleFragment>, context: WasmBack
         dumpReachabilityInfoToFile
     ).collectDeclarations(rootDeclarations = buildRoots(modules, context))
 
-    val remover = WasmUselessDeclarationsRemover(usefulDeclarations)
+    val remover = WasmUselessDeclarationsRemover(context, usefulDeclarations)
     modules.onAllFiles {
         acceptVoid(remover)
     }
@@ -58,8 +59,8 @@ private fun buildRoots(modules: List<IrModuleFragment>, context: WasmBackendCont
     add(context.irBuiltIns.throwableClass.owner)
     add(context.mainCallsWrapperFunction)
     add(context.fieldInitFunction)
-    // TODO move Unit related optimization on IR level and make unit usages explicit 
-    add(context.findUnitGetInstanceFunction())
+    add(context.findUnitInstanceField())
+    add(context.irBuiltIns.unitClass.owner.primaryConstructor!!)
 
     // Remove all functions used to call a kotlin closure from JS side, reachable ones will be added back later.
     removeAll(context.closureCallExports.values)

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/dce/WasmUsefulDeclarationProcessor.kt

@@ -22,8 +22,6 @@ internal class WasmUsefulDeclarationProcessor(
     dumpReachabilityInfoToFile: String?
 ) : UsefulDeclarationProcessor(printReachabilityInfo, removeUnusedAssociatedObjects = false, dumpReachabilityInfoToFile) {
 
-    private val unitGetInstance: IrSimpleFunction = context.findUnitGetInstanceFunction()
-
     // The mapping from function for wrapping a kotlin closure/lambda with JS closure to function used to call a kotlin closure from JS side.
     private val kotlinClosureToJsClosureConvertFunToKotlinClosureCallFun =
         context.kotlinClosureToJsConverters.entries.associate { (k, v) -> v to context.closureCallExports[k] }
@@ -49,6 +47,22 @@ internal class WasmUsefulDeclarationProcessor(
             super.visitVararg(expression, data)
         }
 
+        override fun visitSetField(expression: IrSetField, data: IrDeclaration) {
+            if (!expression.symbol.owner.isObjectInstanceField()) {
+                super.visitSetField(expression, data)
+            }
+        }
+
+        override fun visitGetField(expression: IrGetField, data: IrDeclaration) {
+            val field = expression.symbol.owner
+
+            if (field.isObjectInstanceField()) {
+                field.type.classOrFail.owner.primaryConstructor?.enqueue(field, "object lazy initialization")
+            }
+
+            super.visitGetField(expression, data)
+        }
+
         private fun tryToProcessIntrinsicCall(from: IrDeclaration, call: IrCall): Boolean = when (call.symbol) {
             context.wasmSymbols.unboxIntrinsic -> {
                 val fromType = call.getTypeArgument(0)
@@ -77,9 +91,6 @@ internal class WasmUsefulDeclarationProcessor(
             super.visitCall(expression, data)
 
             val function: IrFunction = expression.symbol.owner.realOverrideTarget
-            if (function.returnType == context.irBuiltIns.unitType) {
-                unitGetInstance.enqueue(data, "function Unit return type")
-            }
 
             if (tryToProcessIntrinsicCall(data, expression)) return
             if (function.hasWasmNoOpCastAnnotation()) return

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/dce/WasmUselessDeclarationsRemover.kt

@@ -5,17 +5,21 @@
 
 package org.jetbrains.kotlin.backend.wasm.dce
 
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrDeclaration
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationContainer
-import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.backend.js.utils.isObjectInstanceField
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrSetField
+import org.jetbrains.kotlin.ir.types.classOrFail
+import org.jetbrains.kotlin.ir.util.primaryConstructor
 import org.jetbrains.kotlin.ir.util.transformFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
 
 class WasmUselessDeclarationsRemover(
+    private val context: WasmBackendContext,
     private val usefulDeclarations: Set<IrDeclaration>
 ) : IrElementVisitorVoid {
     override fun visitElement(element: IrElement) {
@@ -30,6 +34,14 @@ class WasmUselessDeclarationsRemover(
         process(declaration)
     }
 
+    override fun visitSimpleFunction(declaration: IrSimpleFunction) {
+        if (declaration == context.fieldInitFunction) {
+            declaration.removeUnusedObjectsInitializers()
+        }
+
+        super.visitSimpleFunction(declaration)
+    }
+
     // TODO bring back the primary constructor fix
     private fun process(container: IrDeclarationContainer) {
         container.declarations.transformFlat { member ->
@@ -41,4 +53,10 @@ class WasmUselessDeclarationsRemover(
             }
         }
     }
+
+    private fun IrSimpleFunction.removeUnusedObjectsInitializers() {
+        (body as? IrBlockBody)?.statements?.removeIf {
+            it is IrSetField && it.symbol.owner.isObjectInstanceField() && it.symbol.owner !in usefulDeclarations
+        }
+    }
 }

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/BodyGenerator.kt

@@ -14,10 +14,7 @@ import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.backend.js.lower.PrimaryConstructorLowering
-import org.jetbrains.kotlin.ir.backend.js.utils.erasedUpperBound
-import org.jetbrains.kotlin.ir.backend.js.utils.findUnitGetInstanceFunction
-import org.jetbrains.kotlin.ir.backend.js.utils.isDispatchReceiver
-import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
+import org.jetbrains.kotlin.ir.backend.js.utils.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
@@ -34,7 +31,6 @@ class BodyGenerator(
     val context: WasmModuleCodegenContext,
     val functionContext: WasmFunctionCodegenContext,
     private val hierarchyDisjointUnions: DisjointUnions<IrClassSymbol>,
-    private val isGetUnitFunction: Boolean,
 ) : IrElementVisitorVoid {
     val body: WasmExpressionBuilder = functionContext.bodyGen
 
@@ -44,16 +40,13 @@ class BodyGenerator(
     private val irBuiltIns: IrBuiltIns = backendContext.irBuiltIns
 
     private val unitGetInstance by lazy { backendContext.findUnitGetInstanceFunction() }
-    fun WasmExpressionBuilder.buildGetUnit() {
-        buildInstr(
-            WasmOp.GET_UNIT,
-            SourceLocation.NoLocation("GET_UNIT"),
-            WasmImmediate.FuncIdx(context.referenceFunction(unitGetInstance.symbol))
-        )
-    }
+    private val unitInstanceField by lazy { backendContext.findUnitInstanceField() }
 
-    private val anyConstructor by lazy {
-        wasmSymbols.any.constructors.first { it.owner.valueParameters.isEmpty() }
+    fun WasmExpressionBuilder.buildGetUnit() {
+        buildGetGlobal(
+            context.referenceGlobalField(unitInstanceField.symbol),
+            SourceLocation.NoLocation("GET_UNIT")
+        )
     }
 
     // Generates code for the given IR element. Leaves something on the stack unless expression was of the type Void.
@@ -184,6 +177,7 @@ class BodyGenerator(
         val field: IrField = expression.symbol.owner
         val receiver: IrExpression? = expression.receiver
         val location = expression.getSourceLocation()
+
         if (receiver != null) {
             generateExpression(receiver)
             if (backendContext.inlineClassesUtils.isClassInlineLike(field.parentAsClass)) {
@@ -378,12 +372,6 @@ class BodyGenerator(
             return
         }
 
-        // Get unit is a special case because it is the only function which returns the real unit instance.
-        if (call.symbol == unitGetInstance.symbol) {
-            body.buildGetUnit()
-            return
-        }
-
         // Some intrinsics are a special case because we want to remove them completely, including their arguments.
         if (!backendContext.configuration.getNotNull(JSConfigurationKeys.WASM_ENABLE_ARRAY_RANGE_CHECKS)) {
             if (call.symbol == wasmSymbols.rangeCheck) {
@@ -473,8 +461,9 @@ class BodyGenerator(
         }
 
         // Unit types don't cross function boundaries
-        if (function.returnType.isUnit())
+        if (function.returnType.isUnit() && function !is IrConstructor) {
             body.buildGetUnit()
+        }
     }
 
     private fun generateRefNullCast(fromType: IrType, toType: IrType, location: SourceLocation) {
@@ -701,14 +690,12 @@ class BodyGenerator(
 
     private fun visitFunctionReturn(expression: IrReturn) {
         val returnType = expression.returnTargetSymbol.owner.returnType(backendContext)
-        if (returnType == irBuiltIns.unitType && expression.returnTargetSymbol.owner != unitGetInstance) {
-            generateAsStatement(expression.value)
-        } else {
-            if (isGetUnitFunction) {
-                generateExpression(expression.value)
-            } else {
-                generateWithExpectedType(expression.value, returnType)
-            }
+        val isGetUnitFunction = expression.returnTargetSymbol.owner == unitGetInstance
+
+        when {
+            isGetUnitFunction -> generateExpression(expression.value)
+            returnType == irBuiltIns.unitType -> generateAsStatement(expression.value)
+            else -> generateWithExpectedType(expression.value, returnType)
         }
 
         if (functionContext.irFunction is IrConstructor) {
@@ -839,11 +826,7 @@ class BodyGenerator(
         if (!seenElse && resultType != null) {
             assert(expression.type != irBuiltIns.nothingType)
             if (expression.type.isUnit()) {
-                if (isGetUnitFunction) {
-                    generateDefaultInitializerForType(resultType, body)
-                } else {
-                    body.buildGetUnit()
-                }
+                body.buildGetUnit()
             } else {
                 error("'When' without else branch and non Unit type: ${expression.type.dumpKotlinLike()}")
             }

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/ir2wasm/DeclarationGenerator.kt

@@ -38,6 +38,7 @@ class DeclarationGenerator(
     private val irBuiltIns: IrBuiltIns = backendContext.irBuiltIns
 
     private val unitGetInstanceFunction: IrSimpleFunction by lazy { backendContext.findUnitGetInstanceFunction() }
+    private val unitPrimaryConstructor: IrConstructor? by lazy { backendContext.irBuiltIns.unitClass.owner.primaryConstructor }
 
     override fun visitElement(element: IrElement) {
         error("Unexpected element of type ${element::class}")
@@ -98,12 +99,11 @@ class DeclarationGenerator(
         // Generate function type
         val watName = declaration.fqNameWhenAvailable.toString()
         val irParameters = declaration.getEffectiveValueParameters()
-        val resultType =
-            when {
-                // Unit_getInstance returns true Unit reference instead of "void"
-                declaration == unitGetInstanceFunction -> context.transformType(declaration.returnType)
-                else -> context.transformResultType(declaration.returnType)
-            }
+        val resultType = when (declaration) {
+            // Unit_getInstance returns true Unit reference instead of "void"
+            unitGetInstanceFunction, unitPrimaryConstructor -> context.transformType(declaration.returnType)
+            else -> context.transformResultType(declaration.returnType)
+        }
 
         val wasmFunctionType =
             WasmFunctionType(
@@ -149,7 +149,6 @@ class DeclarationGenerator(
             context = context,
             functionContext = functionCodegenContext,
             hierarchyDisjointUnions = hierarchyDisjointUnions,
-            isGetUnitFunction = declaration == unitGetInstanceFunction
         )
 
         if (declaration is IrConstructor) {

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/types/irTypes.kt

@@ -103,6 +103,9 @@ val IrType.classOrNull: IrClassSymbol?
             else -> null
         }
 
+val IrType.classOrFail: IrClassSymbol
+    get() = classOrNull ?: error("Expect type to be a class type")
+
 val IrType.classFqName: FqName?
     get() = classOrNull?.owner?.fqNameWhenAvailable