[WASM] Optimise interop adapters

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.backend.wasm
 
 import org.jetbrains.kotlin.backend.common.ir.Symbols
-import org.jetbrains.kotlin.backend.common.serialization.proto.IrCall
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.builtins.isFunctionType
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
@@ -16,11 +15,10 @@ import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.backend.js.ReflectionSymbols
-import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
-import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
-import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
+import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
-import org.jetbrains.kotlin.ir.symbols.*
+import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
-import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.ir.util.constructors
@@ -28,7 +26,6 @@ import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
-import java.lang.IllegalArgumentException
 
 @OptIn(ObsoleteDescriptorBasedAPI::class)
 class WasmSymbols(
@@ -167,6 +164,7 @@ class WasmSymbols(
     val booleanAnd = getInternalFunction("wasm_i32_and")
     val refEq = getInternalFunction("wasm_ref_eq")
     val refIsNull = getInternalFunction("wasm_ref_is_null")
+    val externRefIsNull = getInternalFunction("wasm_externref_is_null")
     val refTest = getInternalFunction("wasm_ref_test")
     val refCast = getInternalFunction("wasm_ref_cast")
     val wasmArrayCopy = getInternalFunction("wasm_array_copy")
@@ -191,7 +189,6 @@ class WasmSymbols(
     val wasmIsInterface = getInternalFunction("wasmIsInterface")
 
     val nullableEquals = getInternalFunction("nullableEquals")
-    val ensureNotNull = getInternalFunction("ensureNotNull")
     val anyNtoString = getInternalFunction("anyNtoString")
 
     val nullableFloatIeee754Equals = getInternalFunction("nullableFloatIeee754Equals")
@@ -275,6 +272,8 @@ class WasmSymbols(
         val kotlinToJsStringAdapter = getInternalFunction("kotlinToJsStringAdapter")
         val kotlinToJsAnyAdapter = getInternalFunction("kotlinToJsAnyAdapter")
 
+        val jsCheckIsNullOrUndefinedAdapter = getInternalFunction("jsCheckIsNullOrUndefinedAdapter")
+
         val jsToKotlinStringAdapter = getInternalFunction("jsToKotlinStringAdapter")
         val jsToKotlinAnyAdapter = getInternalFunction("jsToKotlinAnyAdapter")
 

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

@@ -108,9 +108,15 @@ fun compileWasm(
 fun WasmCompiledModuleFragment.generateJs(): String {
     //language=js
     val runtime = """
-    
     const externrefBoxes = new WeakMap();
-    """.trimIndent()
+    function tryGetOrSetExternrefBox(ref, ifNotCached) {
+        if (ref == null) return null;
+        if (typeof ref !== 'object') return ifNotCached;
+        const cachedBox = externrefBoxes.get(ref);
+        if (cachedBox !== void 0) return cachedBox;
+        externrefBoxes.set(ref, ifNotCached);
+        return ifNotCached;
+    }    """.trimIndent()
 
     val jsCodeBody = jsFuns.joinToString(",\n") { "\"" + it.importName + "\" : " + it.jsCode }
     val jsCodeBodyIndented = jsCodeBody.prependIndent("    ")

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/BuiltInsLowering.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.wasm.WasmBackendContext
 import org.jetbrains.kotlin.config.AnalysisFlags
 import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.ir.backend.js.utils.erasedUpperBound
 import org.jetbrains.kotlin.ir.backend.js.utils.isEqualsInheritedFromAny
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.IrFile
@@ -63,10 +64,12 @@ class BuiltInsLowering(val context: WasmBackendContext) : FileLoweringPass {
                     }
                 }
                 if (lhs.isNullConst()) {
-                    return builder.irCall(symbols.refIsNull).apply { putValueArgument(0, rhs) }
+                    val refIsNull = if (rhsType.erasedUpperBound?.isExternal == true) symbols.externRefIsNull else symbols.refIsNull
+                    return builder.irCall(refIsNull).apply { putValueArgument(0, rhs) }
                 }
                 if (rhs.isNullConst()) {
-                    return builder.irCall(symbols.refIsNull).apply { putValueArgument(0, lhs) }
+                    val refIsNull = if (lhsType.erasedUpperBound?.isExternal == true) symbols.externRefIsNull else symbols.refIsNull
+                    return builder.irCall(refIsNull).apply { putValueArgument(0, lhs) }
                 }
                 if (!lhsType.isNullable()) {
                     return irCall(call, lhsType.findEqualsMethod().symbol, argumentsAsReceivers = true)
@@ -81,17 +84,20 @@ class BuiltInsLowering(val context: WasmBackendContext) : FileLoweringPass {
             }
 
             irBuiltins.checkNotNullSymbol -> {
+                val arg = call.getValueArgument(0)!!
 
-                // Workaround: v8 doesnt support ref.cast-ing unreachable very well.
+                if (arg.isNullConst()) {
-                run {
+                    return builder.irCall(symbols.throwNullPointerException)
-                    val arg = call.getValueArgument(0)!!
-                    if (arg.isNullConst()) {
-                        return builder.irCall(symbols.wasmUnreachable, irBuiltins.nothingType)
-                    }
                 }
 
-                return irCall(call, symbols.ensureNotNull).also {
+                return builder.irComposite {
-                    it.putTypeArgument(0, call.type)
+                    val temporary = irTemporary(arg)
+                    +builder.irIfNull(
+                        type = arg.type.makeNotNull(),
+                        subject = irGet(temporary),
+                        thenPart = builder.irCall(symbols.throwNullPointerException),
+                        elsePart = irGet(temporary)
+                    )
                 }
             }
             in symbols.comparisonBuiltInsToWasmIntrinsics.keys -> {

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/JsInteropFunctionsLowering.kt

@@ -286,38 +286,63 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
         return SendKotlinObjectToJsAdapter(this)
     }
 
+    private fun createNullableAdapter(notNullType: IrType, isPrimitive: Boolean, valueAdapter: InteropTypeAdapter?): InteropTypeAdapter? {
+        return if (isPrimitive) { //nullable primitive should be checked and adapt to target type
+            val externRefToPrimitiveAdapter = when (notNullType) {
+                builtIns.floatType -> adapters.externRefToKotlinFloatAdapter.owner
+                builtIns.doubleType -> adapters.externRefToKotlinDoubleAdapter.owner
+                builtIns.longType -> adapters.externRefToKotlinLongAdapter.owner
+                builtIns.booleanType -> adapters.externRefToKotlinBooleanAdapter.owner
+                else -> adapters.externRefToKotlinIntAdapter.owner
+            }
+            val externalToPrimitiveAdapter = FunctionBasedAdapter(externRefToPrimitiveAdapter)
+            NullOrAdapter(
+                adapter = valueAdapter?.let { CombineAdapter(it, externalToPrimitiveAdapter) } ?: externalToPrimitiveAdapter
+            )
+        } else { //nullable reference should not be checked
+            val nullableValueAdapter = valueAdapter?.let(::NullOrAdapter)
+            if (isExternalType(notNullType)) {
+                val undefinedToNullAdapter = FunctionBasedAdapter(adapters.jsCheckIsNullOrUndefinedAdapter.owner)
+                nullableValueAdapter
+                    ?.let { CombineAdapter(it, undefinedToNullAdapter) }
+                    ?: undefinedToNullAdapter
+            } else {
+                nullableValueAdapter
+            }
+        }
+    }
+
+    private fun createNotNullAdapter(notNullType: IrType, isPrimitive: Boolean, valueAdapter: InteropTypeAdapter?): InteropTypeAdapter? {
+        // !nullable primitive checked by wasm signature
+        if (isPrimitive) return valueAdapter
+
+        // !nullable reference should be null checked
+        // notNullAdapter((undefined -> null)!!)
+        val nullCheckedValueAdapter = valueAdapter?.let(::CheckNotNullAndAdapter)
+            ?: CheckNotNullNoAdapter(notNullType)
+
+        // kotlin types could not take undefined value so just take null-checked value
+        if (!isExternalType(notNullType)) return nullCheckedValueAdapter
+
+        // js value should convert undefined into null and the null-checked
+        return CombineAdapter(
+            outerAdapter = nullCheckedValueAdapter,
+            innerAdapter = FunctionBasedAdapter(adapters.jsCheckIsNullOrUndefinedAdapter.owner)
+        )
+    }
+
     private fun IrType.jsToKotlinAdapterIfNeeded(isReturn: Boolean): InteropTypeAdapter? {
         if (isReturn && this == builtIns.unitType)
             return null
 
         val notNullType = makeNotNull()
-        val notNullAdapter = notNullType.jsToKotlinAdapterIfNeededNotNullable(isReturn)
+        val valueAdapter = notNullType.jsToKotlinAdapterIfNeededNotNullable(isReturn)
-        val isPrimitive = notNullAdapter?.fromType?.isPrimitiveType() ?: notNullType.isPrimitiveType()
+        val isPrimitive = valueAdapter?.fromType?.isPrimitiveType() ?: notNullType.isPrimitiveType()
 
-        return if (isNullable()) {
+        return if (isNullable())
-            if (isPrimitive) { //nullable primitive should be checked and adapt to target type
+            createNullableAdapter(notNullType, isPrimitive, valueAdapter)
-                val externRefToPrimitiveAdapter = when (notNullType) {
+        else
-                    builtIns.floatType -> adapters.externRefToKotlinFloatAdapter.owner
+            createNotNullAdapter(notNullType, isPrimitive, valueAdapter)
-                    builtIns.doubleType -> adapters.externRefToKotlinDoubleAdapter.owner
-                    builtIns.longType -> adapters.externRefToKotlinLongAdapter.owner
-                    builtIns.booleanType -> adapters.externRefToKotlinBooleanAdapter.owner
-                    else -> adapters.externRefToKotlinIntAdapter.owner
-                }
-                val externalToPrimitiveAdapter = FunctionBasedAdapter(externRefToPrimitiveAdapter)
-                NullOrAdapter(
-                    adapter = notNullAdapter?.let { CombineAdapter(it, externalToPrimitiveAdapter) } ?: externalToPrimitiveAdapter
-                )
-            } else { //nullable reference should not be checked
-                notNullAdapter?.let(::NullOrAdapter)
-            }
-        } else {
-            if (isPrimitive) { // !nullable primitive checked by wasm signature
-                notNullAdapter
-            } else { // !nullable reference should be null checked
-                notNullAdapter?.let(::CheckNotNullAndAdapter)
-                    ?: CheckNotNullNoAdapter(this)
-            }
-        }
     }
 
     private fun IrType.jsToKotlinAdapterIfNeededNotNullable(isReturn: Boolean): InteropTypeAdapter? {
@@ -667,13 +692,22 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
     /**
      * Current V8 Wasm GC mandates dataref type instead of structs and arrays
      */
+
+    /**
+     * Effectively `value!!`
+     */
     inner class CheckNotNullNoAdapter(type: IrType) : InteropTypeAdapter {
         override val fromType: IrType = type.makeNullable()
         override val toType: IrType = type.makeNotNull()
         override fun adapt(expression: IrExpression, builder: IrBuilderWithScope): IrExpression {
-            return builder.irCall(context.irBuiltIns.checkNotNullSymbol).also {
+            return builder.irComposite {
-                it.putValueArgument(0, expression)
+                val tmp = irTemporary(expression)
-                it.putTypeArgument(0, fromType)
+                +irIfNull(
+                    type = toType,
+                    subject = irGet(tmp),
+                    thenPart = builder.irCall(symbols.throwNullPointerException),
+                    elsePart = irGet(tmp)
+                )
             }
         }
     }

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/WasmTypeOperatorLowering.kt

@@ -18,10 +18,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.backend.js.utils.erasedUpperBound
 import org.jetbrains.kotlin.ir.builders.*
-import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrTypeParameter
-import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
@@ -208,6 +205,11 @@ class WasmBaseTypeOperatorTransformer(val context: WasmBackendContext) : IrEleme
             return value
         }
 
+        if (value.isNullConst() && fromClass.isExternal != toClass.isExternal) {
+            value.type = toType
+            return value
+        }
+
         if (fromClass.isExternal && !toClass.isExternal) {
             val narrowingToAny = builder.irCall(symbols.jsInteropAdapters.jsToKotlinAnyAdapter).also {
                 it.putValueArgument(0, value)

compiler/testData/codegen/boxInline/nonLocalReturns/tryFinally/kt16417.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: WASM
 // NO_CHECK_LAMBDA_INLINING
 // !LANGUAGE: +ProperFinally
 // FILE: 1.kt

compiler/testData/codegen/boxWasmJsInterop/nullableExternRefs.kt

@@ -63,7 +63,6 @@ fun box(): String {
     assertTrue(isJsNull(jsNull))
 
     assertFalse(isJsUndefined(null))
-    assertTrue(isJsUndefined(jsUndefined))
 
     checkNPE(::getJsNullAsNonNullable)
     checkNPE(::getJsUndefinedAsNonNullable)

libraries/stdlib/wasm/internal/kotlin/wasm/internal/ExternalWrapper.kt

@@ -114,11 +114,7 @@ private external fun doubleToExternref(x: Double): ExternalInterfaceType
 @JsFun("(lhs, rhs) => lhs === rhs")
 private external fun externrefEquals(lhs: ExternalInterfaceType, rhs: ExternalInterfaceType): Boolean
 
-@JsFun("(ref) => (typeof ref !== 'object' || ref === null) ? null : (externrefBoxes.get(ref) ?? null)")
+private external fun tryGetOrSetExternrefBox(ref: ExternalInterfaceType, ifNotCached: JsExternalBox): JsExternalBox?
-private external fun getExternrefBoxOrNull(ref: ExternalInterfaceType): JsExternalBox?
-
-@JsFun("(ref, box) => { if (typeof ref === 'object' && ref !== null) { externrefBoxes.set(ref, box); } }")
-private external fun setExternrefBox(ref: ExternalInterfaceType, box: JsExternalBox)
 
 @WasmNoOpCast
 @Suppress("unused")
@@ -130,11 +126,11 @@ private fun ExternalInterfaceType.externAsWasmAnyref(): anyref =
     implementedAsIntrinsic
 
 @WasmOp(WasmOp.EXTERN_EXTERNALIZE)
-private fun Any?.asWasmExternRef(): ExternalInterfaceType =
+private fun Any.asWasmExternRef(): ExternalInterfaceType =
     implementedAsIntrinsic
 
 @JsFun("(ref) => ref == null")
-internal external fun isNullish(ref: ExternalInterfaceType): Boolean
+internal external fun isNullish(ref: ExternalInterfaceType?): Boolean
 
 internal fun externRefToAny(ref: ExternalInterfaceType): Any? {
     // If ref is an instance of kotlin class -- return it cased to Any
@@ -146,15 +142,10 @@ internal fun externRefToAny(ref: ExternalInterfaceType): Any? {
         }
     }
 
-    if (isNullish(ref))
+    // If we have Null in notNullRef -- return null
-        return null
-
     // If we already have a box -- return it,
-    // otherwise -- create a new box and remember it.
+    // otherwise -- remember new box and return it.
-    return getExternrefBoxOrNull(ref)
+    return tryGetOrSetExternrefBox(ref, JsExternalBox(ref))
-        ?: JsExternalBox(ref).also {
-            setExternrefBox(ref, it)
-        }
 }
 
 
@@ -182,7 +173,7 @@ internal fun kotlinToJsStringAdapter(x: String?): ExternalInterfaceType? {
     // Using nullable String to represent default value
     // for parameters with default values
     if (x == null) return null
-    if (x.isEmpty()) return jsEmptyString()
+    if (x.isEmpty()) return jsEmptyString
 
     val srcArray = x.chars
     val stringLength = srcArray.len()
@@ -202,6 +193,9 @@ internal fun kotlinToJsStringAdapter(x: String?): ExternalInterfaceType? {
     return importStringFromWasm(memBuffer, stringLength - srcStartIndex, result)
 }
 
+internal fun jsCheckIsNullOrUndefinedAdapter(x: ExternalInterfaceType?): ExternalInterfaceType? =
+    x.takeIf { !isNullish(it) }
+
 // js string to kotlin string import
 // TODO Uint16Array may work with byte endian different with Wasm (i.e. little endian)
 //language=js
@@ -244,19 +238,23 @@ internal fun jsToKotlinStringAdapter(x: ExternalInterfaceType): String {
 
 
 @JsFun("() => ''")
-internal external fun jsEmptyString(): ExternalInterfaceType
+private external fun getJsEmptyString(): ExternalInterfaceType
 
 @JsFun("() => true")
-internal external fun jsTrue(): ExternalInterfaceType
+private external fun getJsTrue(): ExternalInterfaceType
 
 @JsFun("() => false")
-internal external fun jsFalse(): ExternalInterfaceType
+private external fun getJsFalse(): ExternalInterfaceType
 
-internal fun kotlinToJsAnyAdapter(x: Any): ExternalInterfaceType =
+private val jsEmptyString by lazy(::getJsEmptyString)
-    anyToExternRef(x)
+private val jsTrue by lazy(::getJsTrue)
+private val jsFalse by lazy(::getJsFalse)
 
-internal fun jsToKotlinAnyAdapter(x: ExternalInterfaceType): Any? =
+internal fun kotlinToJsAnyAdapter(x: Any?): ExternalInterfaceType? =
-    externRefToAny(x)
+    if (x == null) null else anyToExternRef(x)
+
+internal fun jsToKotlinAnyAdapter(x: ExternalInterfaceType?): Any? =
+    if (x == null) null else externRefToAny(x)
 
 internal fun jsToKotlinByteAdapter(x: Int): Byte = x.toByte()
 internal fun jsToKotlinShortAdapter(x: Int): Short = x.toShort()
@@ -282,7 +280,7 @@ internal fun kotlinIntToExternRefAdapter(x: Int): ExternalInterfaceType =
     intToExternref(x)
 
 internal fun kotlinBooleanToExternRefAdapter(x: Boolean): ExternalInterfaceType =
-    if (x) jsTrue() else jsFalse()
+    if (x) jsTrue else jsFalse
 
 internal fun kotlinLongToExternRefAdapter(x: Long): ExternalInterfaceType =
     longToExternref(x)

libraries/stdlib/wasm/internal/kotlin/wasm/internal/Runtime.kt

@@ -65,9 +65,6 @@ internal fun nullableDoubleIeee754Equals(lhs: Double?, rhs: Double?): Boolean {
     return wasm_f64_eq(lhs, rhs)
 }
 
-
-internal fun <T : Any> ensureNotNull(v: T?): T = if (v == null) THROW_NPE() else v
-
 @ExcludedFromCodegen
 internal fun <T, R> boxIntrinsic(x: T): R =
     implementedAsIntrinsic

libraries/stdlib/wasm/internal/kotlin/wasm/internal/WasmInstructions.kt

@@ -13,6 +13,7 @@ import kotlin.wasm.internal.reftypes.anyref
 import kotlin.wasm.internal.reftypes.dataref
 import kotlin.wasm.internal.reftypes.funcref
 import kotlin.wasm.internal.reftypes.i31ref
+import kotlin.wasm.internal.ExternalInterfaceType
 
 @WasmOp(WasmOp.UNREACHABLE)
 internal fun wasm_unreachable(): Nothing =
@@ -384,4 +385,7 @@ internal external fun wasm_ref_is_i31(x: anyref): Boolean
 internal external fun wasm_ref_as_data(x: anyref): dataref
 
 @WasmOp(WasmOp.REF_AS_I31)
-internal external fun wasm_ref_as_i31(x: anyref): i31ref
+internal external fun wasm_ref_as_i31(x: anyref): i31ref
+
+@WasmOp(WasmOp.REF_IS_NULL)
+internal external fun wasm_externref_is_null(a: ExternalInterfaceType?): Boolean