[Wasm] Make specialisations for closured primitive values

Fixed KT-66065
2024-02-23 18:03:08 +01:00
parent 4eba0075bb
commit 60d425e2c7
7 changed files with 114 additions and 189 deletions
@@ -122,8 +122,7 @@ class WasmBackendContext(
    val fieldInitFunction = createInitFunction("fieldInit")
    val mainCallsWrapperFunction = createInitFunction("mainCallsWrapper")

-    override val sharedVariablesManager =
-        WasmSharedVariablesManager(this, irBuiltIns, internalPackageFragment)
+    override val sharedVariablesManager = WasmSharedVariablesManager(this)

    val wasmSymbols: WasmSymbols = WasmSymbols(this@WasmBackendContext, symbolTable)
    override val reflectionSymbols: ReflectionSymbols get() = wasmSymbols.reflectionSymbols
@@ -154,6 +154,22 @@ class WasmSymbols(
    fun findVoidConsumer(type: IrType): IrSimpleFunctionSymbol =
        consumePrimitiveIntoVoid[type] ?: consumeAnyIntoVoid

+    private val closureBoxAnyClass = getInternalClass("ClosureBoxAny")
+
+    private val closureBoxClasses = mapOf(
+        context.irBuiltIns.booleanType to getInternalClass("ClosureBoxBoolean"),
+        context.irBuiltIns.byteType to getInternalClass("ClosureBoxByte"),
+        context.irBuiltIns.shortType to getInternalClass("ClosureBoxShort"),
+        context.irBuiltIns.charType to getInternalClass("ClosureBoxChar"),
+        context.irBuiltIns.intType to getInternalClass("ClosureBoxInt"),
+        context.irBuiltIns.longType to getInternalClass("ClosureBoxLong"),
+        context.irBuiltIns.floatType to getInternalClass("ClosureBoxFloat"),
+        context.irBuiltIns.doubleType to getInternalClass("ClosureBoxDouble")
+    )
+
+    fun findClosureBoxClass(type: IrType): IrClassSymbol =
+        closureBoxClasses[type] ?: closureBoxAnyClass
+
    val equalityFunctions by lazy {
        mapOf(
            context.irBuiltIns.booleanType to getInternalFunction("wasm_i32_eq"),
@@ -6,15 +6,7 @@
 package org.jetbrains.kotlin.backend.wasm.lower

 import org.jetbrains.kotlin.backend.common.ir.SharedVariablesManager
-import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
-import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.ir.IrBuiltIns
-import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
-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.builders.declarations.buildClass
-import org.jetbrains.kotlin.ir.builders.declarations.buildValueParameter
+import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrVariableImpl
 import org.jetbrains.kotlin.ir.expressions.IrExpression
@@ -23,53 +15,47 @@ import org.jetbrains.kotlin.ir.expressions.IrSetValue
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
-import org.jetbrains.kotlin.ir.symbols.impl.IrConstructorSymbolImpl
-import org.jetbrains.kotlin.ir.symbols.impl.IrFieldSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
-import org.jetbrains.kotlin.ir.types.IrType
-import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
+import org.jetbrains.kotlin.ir.types.classOrFail
+import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.defaultType
-import org.jetbrains.kotlin.ir.util.parentAsClass
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.name.SpecialNames
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance

-/**
- * This is a copy of an old version of JS lowering, because JS did platform-specific optimization incompatible with Wasm.
- * TODO: Revisit
- */
-class WasmSharedVariablesManager(val context: JsCommonBackendContext, val builtIns: IrBuiltIns, val implicitDeclarationsFile: IrPackageFragment) : SharedVariablesManager {
+class WasmSharedVariablesManager(val context: WasmBackendContext) : SharedVariablesManager {
    override fun declareSharedVariable(originalDeclaration: IrVariable): IrVariable {
-        val initializer = originalDeclaration.initializer ?: IrConstImpl.constNull(
+        val initializer = originalDeclaration.initializer ?: IrConstImpl.defaultValueForType(
            originalDeclaration.startOffset,
            originalDeclaration.endOffset,
-            builtIns.nothingNType
+            originalDeclaration.type
        )

-        val constructorSymbol = closureBoxConstructorDeclaration.symbol
+        val boxClass = context.wasmSymbols.findClosureBoxClass(originalDeclaration.type)
+        val constructorSymbol = boxClass.constructors.first()

        val irCall =
            IrConstructorCallImpl(
-                initializer.startOffset,
-                initializer.endOffset,
-                closureBoxType,
-                constructorSymbol,
-                closureBoxConstructorDeclaration.parentAsClass.typeParameters.size,
-                closureBoxConstructorDeclaration.typeParameters.size,
-                closureBoxConstructorDeclaration.valueParameters.size
+                startOffset = initializer.startOffset,
+                endOffset = initializer.endOffset,
+                type = boxClass.defaultType,
+                symbol = constructorSymbol,
+                typeArgumentsCount = boxClass.owner.typeParameters.size,
+                constructorTypeArgumentsCount = constructorSymbol.owner.typeParameters.size,
+                valueArgumentsCount = constructorSymbol.owner.valueParameters.size
            ).apply {
                putValueArgument(0, initializer)
            }

        return IrVariableImpl(
-            originalDeclaration.startOffset,
-            originalDeclaration.endOffset,
-            originalDeclaration.origin,
-            IrVariableSymbolImpl(),
-            originalDeclaration.name,
-            irCall.type,
-            false,
-            false,
-            false
+            startOffset = originalDeclaration.startOffset,
+            endOffset = originalDeclaration.endOffset,
+            origin = originalDeclaration.origin,
+            symbol = IrVariableSymbolImpl(),
+            name = originalDeclaration.name,
+            type = irCall.type,
+            isVar = false,
+            isConst = false,
+            isLateinit = false
        ).also {
            it.parent = originalDeclaration.parent
            it.initializer = irCall
@@ -79,150 +65,66 @@ class WasmSharedVariablesManager(val context: JsCommonBackendContext, val builtI
    override fun defineSharedValue(originalDeclaration: IrVariable, sharedVariableDeclaration: IrVariable) = sharedVariableDeclaration

    override fun getSharedValue(sharedVariableSymbol: IrValueSymbol, originalGet: IrGetValue): IrExpression {
-        val getField = IrGetFieldImpl(
-            originalGet.startOffset, originalGet.endOffset,
-            closureBoxFieldDeclaration.symbol,
-            closureBoxFieldDeclaration.type,
-            IrGetValueImpl(
-                originalGet.startOffset,
-                originalGet.endOffset,
-                closureBoxType,
-                sharedVariableSymbol,
-                originalGet.origin
-            ),
-            originalGet.origin
-        )
+        val boxClass = sharedVariableSymbol.owner.type.classOrFail.owner
+        val valueProperty = boxClass.declarations.firstIsInstance<IrProperty>()
+
+        check(valueProperty.name.asString() == "value")
+        val propertyGetter = valueProperty.getter!!
+
+        val propertyGet = IrCallImpl(
+            startOffset = originalGet.startOffset,
+            endOffset = originalGet.endOffset,
+            type = propertyGetter.returnType,
+            symbol = propertyGetter.symbol,
+            typeArgumentsCount = 0,
+            valueArgumentsCount = 0,
+            origin = originalGet.origin
+        ).also {
+            it.dispatchReceiver = IrGetValueImpl(
+                startOffset = originalGet.startOffset,
+                endOffset = originalGet.endOffset,
+                type = boxClass.defaultType,
+                symbol = sharedVariableSymbol,
+                origin = originalGet.origin
+            )
+        }

        return IrTypeOperatorCallImpl(
-            originalGet.startOffset,
-            originalGet.endOffset,
-            originalGet.type,
-            IrTypeOperator.IMPLICIT_CAST,
-            originalGet.type,
-            getField
+            startOffset = originalGet.startOffset,
+            endOffset = originalGet.endOffset,
+            type = originalGet.type,
+            operator = IrTypeOperator.IMPLICIT_CAST,
+            typeOperand = originalGet.type,
+            argument = propertyGet
        )
    }

-    override fun setSharedValue(sharedVariableSymbol: IrValueSymbol, originalSet: IrSetValue): IrExpression =
-        IrSetFieldImpl(
-            originalSet.startOffset,
-            originalSet.endOffset,
-            closureBoxFieldDeclaration.symbol,
-            IrGetValueImpl(
-                originalSet.startOffset,
-                originalSet.endOffset,
-                closureBoxType,
-                sharedVariableSymbol,
-                originalSet.origin
-            ),
-            originalSet.value,
-            originalSet.type,
-            originalSet.origin
-        )
+    override fun setSharedValue(sharedVariableSymbol: IrValueSymbol, originalSet: IrSetValue): IrExpression {
+        val boxClass = sharedVariableSymbol.owner.type.classOrFail.owner
+        val valueProperty = boxClass.declarations.firstIsInstance<IrProperty>()

-    private val boxTypeName = "\$closureBox\$"
+        check(valueProperty.name.asString() == "value")
+        val propertySetter = valueProperty.setter!!

-    private val closureBoxClassDeclaration by lazy {
-        createClosureBoxClassDeclaration()
-    }
-
-    private val closureBoxConstructorDeclaration by lazy {
-        createClosureBoxConstructorDeclaration()
-    }
-
-    private val closureBoxFieldDeclaration by lazy {
-        closureBoxPropertyDeclaration
-    }
-
-    private val closureBoxPropertyDeclaration by lazy {
-        createClosureBoxPropertyDeclaration()
-    }
-
-    private lateinit var closureBoxType: IrType
-
-    private fun createClosureBoxClassDeclaration(): IrClass {
-        val declaration = context.irFactory.buildClass {
-            origin = JsLoweredDeclarationOrigin.JS_CLOSURE_BOX_CLASS_DECLARATION
-            name = Name.identifier(boxTypeName)
-            visibility = DescriptorVisibilities.PUBLIC
-            modality = Modality.FINAL
-            isCompanion = false
-            isInner = false
-            isData = false
-            isExternal = false
-            isValue = false
-            isExpect = false
-            isFun = false
-        }
-
-        declaration.parent = implicitDeclarationsFile
-        // TODO: substitute
-        closureBoxType = IrSimpleTypeImpl(declaration.symbol, false, emptyList(), emptyList())
-        declaration.thisReceiver = buildValueParameter(declaration) {
-            name = Name.identifier("_this_")
-            index = -1
-            type = closureBoxType
-        }
-        implicitDeclarationsFile.declarations += declaration
-
-        return declaration
-    }
-
-    private fun createClosureBoxPropertyDeclaration(): IrField {
-        val symbol = IrFieldSymbolImpl()
-        val fieldName = Name.identifier("v")
-        return context.irFactory.createField(
-            startOffset = UNDEFINED_OFFSET,
-            endOffset = UNDEFINED_OFFSET,
-            origin = IrDeclarationOrigin.FIELD_FOR_OUTER_THIS,
-            name = fieldName,
-            visibility = DescriptorVisibilities.PUBLIC,
-            symbol = symbol,
-            type = builtIns.anyNType,
-            isFinal = false,
-            isStatic = false,
-            isExternal = false,
+        val propertySet = IrCallImpl(
+            startOffset = originalSet.startOffset,
+            endOffset = originalSet.endOffset,
+            type = propertySetter.returnType,
+            symbol = propertySetter.symbol,
+            typeArgumentsCount = 0,
+            valueArgumentsCount = 1,
+            origin = originalSet.origin
        ).also {
-            it.parent = closureBoxClassDeclaration
-            closureBoxClassDeclaration.declarations += it
+            it.dispatchReceiver = IrGetValueImpl(
+                startOffset = originalSet.startOffset,
+                endOffset = originalSet.endOffset,
+                type = boxClass.defaultType,
+                symbol = sharedVariableSymbol,
+                origin = originalSet.origin
+            )
+            it.putValueArgument(0, originalSet.value)
        }
-    }

-    private fun createClosureBoxConstructorDeclaration(): IrConstructor {
-        val symbol = IrConstructorSymbolImpl()
-
-        val declaration = context.irFactory.createConstructor(
-            startOffset = UNDEFINED_OFFSET,
-            endOffset = UNDEFINED_OFFSET,
-            origin = JsLoweredDeclarationOrigin.JS_CLOSURE_BOX_CLASS_DECLARATION,
-            name = SpecialNames.INIT,
-            visibility = DescriptorVisibilities.PUBLIC,
-            isInline = false,
-            isExpect = false,
-            returnType = closureBoxClassDeclaration.defaultType,
-            symbol = symbol,
-            isPrimary = true,
-            isExternal = false,
-        )
-
-        declaration.parent = closureBoxClassDeclaration
-
-        val parameterDeclaration = createClosureBoxConstructorParameterDeclaration(declaration)
-
-        declaration.valueParameters += parameterDeclaration
-
-        val receiver = JsIrBuilder.buildGetValue(closureBoxClassDeclaration.thisReceiver!!.symbol)
-        val value = JsIrBuilder.buildGetValue(parameterDeclaration.symbol)
-
-        val setField = JsIrBuilder.buildSetField(closureBoxFieldDeclaration.symbol, receiver, value, builtIns.unitType)
-
-        declaration.body = context.irFactory.createBlockBody(UNDEFINED_OFFSET, UNDEFINED_OFFSET, listOf(setField))
-
-        closureBoxClassDeclaration.declarations += declaration
-        return declaration
-    }
-
-    private fun createClosureBoxConstructorParameterDeclaration(irConstructor: IrConstructor): IrValueParameter {
-        return JsIrBuilder.buildValueParameter(irConstructor,"p", 0, closureBoxPropertyDeclaration.type)
+        return propertySet
    }
 }
@@ -38,19 +38,10 @@ fun f(block: () -> Unit) {
 // test.kt:8 box

 // EXPECTATIONS WASM
-// test.kt:4 $box (12, 12, 4)
-// Runtime.kt:70 $kotlin.wasm.internal.getBoxedBoolean (8, 8)
-// Runtime.kt:73 $kotlin.wasm.internal.getBoxedBoolean (8, 35)
-// Standard.kt:71 $kotlin.wasm.internal.getBoxedBoolean (0, 0, 0, 0)
-// Standard.kt:95 $kotlin.wasm.internal.getBoxedBoolean (4, 4)
-// Standard.kt:98 $kotlin.wasm.internal.getBoxedBoolean (4, 10, 4, 10)
-// Standard.kt:74 $kotlin.wasm.internal.getBoxedBoolean (15, 7)
-// Standard.kt:99 $kotlin.wasm.internal.getBoxedBoolean (11, 4, 11, 4)
-// Runtime.kt:74 $kotlin.wasm.internal.getBoxedBoolean (5, 5)
+// test.kt:4 $box (12, 12, 12, 4)
+// ClosureBoxes.kt:9 $kotlin.wasm.internal.ClosureBoxBoolean.<init> (33, 52)
 // test.kt:5 $box (6, 6, 4)
 // test.kt:11 $f
 // test.kt:6 $box$lambda.invoke (8, 12, 8, 16)
-// Runtime.kt:71 $kotlin.wasm.internal.getBoxedBoolean (8, 33)
-// Standard.kt:68 $kotlin.wasm.internal.getBoxedBoolean (25, 25, 25, 25, 45, 38)
 // test.kt:12 $f
 // test.kt:8 $box
@@ -45,7 +45,7 @@ fun box() {
 // EXPECTATIONS WASM
 // test.kt:18 $box
 // test.kt:10 $flaf
-// test.kt:7 $flaf (46, 46, 57)
+// test.kt:7 $flaf (46, 57)
 // test.kt:6 $foo (39, 46)
 // test.kt:7 $flaf$lambda.invoke (52, 55)
 // kotlin-package.kt:6 $flaf$lambda.invoke
@@ -77,6 +77,7 @@ fun box() {
 // Array.kt:75 $kotlin.Array.set (8, 20, 27, 16, 8, 20, 27, 16)
 // Array.kt:76 $kotlin.Array.set (5, 5)
 // String.kt:149 $kotlin.stringLiteral (11, 4, 11, 4)
+// ClosureBoxes.kt:8 $kotlin.wasm.internal.ClosureBoxAny.<init> (29, 44)
 // test.kt:22 $box
 // test.kt:9 $box (4, 4)
 // test.kt:15 $<no name provided>.<init>
@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2024 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 kotlin.wasm.internal
+
+internal class ClosureBoxAny(var value: Any)
+internal class ClosureBoxBoolean(var value: Boolean)
+internal class ClosureBoxByte(var value: Byte)
+internal class ClosureBoxShort(var value: Short)
+internal class ClosureBoxChar(var value: Char)
+internal class ClosureBoxInt(var value: Int)
+internal class ClosureBoxLong(var value: Long)
+internal class ClosureBoxFloat(var value: Float)
+internal class ClosureBoxDouble(var value: Double)