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

2023-06-02 14:23:40 +00:00
parent 7d6275e228
commit bfd57fd2df
39 changed files with 551 additions and 90 deletions
@@ -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
 )
@@ -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)
@@ -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
@@ -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
+        }
+    }
 }
@@ -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()}")
            }
@@ -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) {