[WASM] Fix interop adapter for long strings

2022-06-27 21:21:15 +02:00
parent 6e75e4a56e
commit 6f88e9b16f
12 changed files with 136 additions and 71 deletions
@@ -185,6 +185,7 @@ class WasmSymbols(
    val nullableDoubleIeee754Equals = getInternalFunction("nullableDoubleIeee754Equals")

    val unsafeGetScratchRawMemory = getInternalFunction("unsafeGetScratchRawMemory")
+    val unsafeGetScratchRawMemorySize = getInternalFunction("unsafeGetScratchRawMemorySize")
    val startCoroutineUninterceptedOrReturnIntrinsics =
        (0..2).map { getInternalFunction("startCoroutineUninterceptedOrReturnIntrinsic$it") }

@@ -465,6 +465,10 @@ class BodyGenerator(
                body.buildConstI32Symbol(context.scratchMemAddr)
            }

+            wasmSymbols.unsafeGetScratchRawMemorySize -> {
+                body.buildConstI32Symbol(WasmSymbol(context.scratchMemSizeInBytes))
+            }
+
            else -> {
                return false
            }
@@ -16,6 +16,7 @@ interface WasmBaseCodegenContext {
    val backendContext: WasmBackendContext

    val scratchMemAddr: WasmSymbol<Int>
+    val scratchMemSizeInBytes: Int

    fun referenceFunction(irFunction: IrFunctionSymbol): WasmSymbol<WasmFunction>
    fun referenceGlobalField(irField: IrFieldSymbol): WasmSymbol<WasmGlobal>
@@ -60,7 +60,7 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {
    val initFunctions = mutableListOf<FunWithPriority>()

    val scratchMemAddr = WasmSymbol<Int>()
-    private val scratchMemSizeInBytes = 65_536
+    val scratchMemSizeInBytes = 65_536

    open class ReferencableElements<Ir, Wasm : Any> {
        val unbound = mutableMapOf<Ir, WasmSymbol<Wasm>>()
@@ -25,6 +25,9 @@ class WasmModuleCodegenContextImpl(
    override val scratchMemAddr: WasmSymbol<Int>
        get() = wasmFragment.scratchMemAddr

+    override val scratchMemSizeInBytes: Int
+        get() = wasmFragment.scratchMemSizeInBytes
+
    override fun transformType(irType: IrType): WasmType {
        return with(typeTransformer) { irType.toWasmValueType() }
    }
@@ -0,0 +1,19 @@
+// FILE: externals.js
+
+function id(str) {
+    return str
+}
+
+// FILE: externals.kt
+
+external fun id(str: String): String
+
+fun box(): String {
+    var x = "1234567890"
+    for (i in 1 until 20) {
+        x += x
+        val stringFromJs = id(x)
+        if (!stringFromJs.equals(x)) return "FAIL"
+    }
+    return "OK"
+}
@@ -55,6 +55,12 @@ public class IrCodegenWasmJsInteropJsTestGenerated extends AbstractIrCodegenWasm
        runTest("compiler/testData/codegen/boxWasmJsInterop/jsExport.kt");
    }

+    @Test
+    @TestMetadata("longStrings.kt")
+    public void testLongStrings() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmJsInterop/longStrings.kt");
+    }
+
    @Test
    @TestMetadata("types.kt")
    public void testTypes() throws Exception {
@@ -55,6 +55,11 @@ public class IrCodegenWasmJsInteropWasmTestGenerated extends AbstractIrCodegenWa
        runTest("compiler/testData/codegen/boxWasmJsInterop/jsExport.kt");
    }

+    @TestMetadata("longStrings.kt")
+    public void testLongStrings() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmJsInterop/longStrings.kt");
+    }
+
    @TestMetadata("nullableExternRefs.kt")
    public void testNullableExternRefs() throws Exception {
        runTest("compiler/testData/codegen/boxWasmJsInterop/nullableExternRefs.kt");
@@ -91,4 +91,8 @@ public class String private constructor(internal val chars: WasmCharArray) : Com
    }
 }

-internal fun stringLiteral(startAddr: Int, length: Int) = String.unsafeFromCharArray(unsafeRawMemoryToWasmCharArray(startAddr, length))
+internal fun stringLiteral(startAddr: Int, length: Int): String {
+    val array = WasmCharArray(length)
+    unsafeRawMemoryToWasmCharArray(startAddr, 0, length, array)
+    return String.unsafeFromCharArray(array)
+}
@@ -138,49 +138,87 @@ internal fun anyToExternRef(x: Any): ExternalInterfaceType {
        x.asWasmExternRef()
 }

-@JsFun("""(addr) => {
-    const mem16 = new Uint16Array(wasmExports.memory.buffer);
-    const mem32 = new Int32Array(wasmExports.memory.buffer);
-    const len = mem32[addr / 4];
-    const str_start_addr = (addr + 4) / 2;
-    const slice = mem16.slice(str_start_addr, str_start_addr + len);
-    return String.fromCharCode.apply(null, slice);
-}
-""")
-internal external fun importStringFromWasm(addr: Int): ExternalInterfaceType
-
@JsFun("x => x.length")
 internal external fun stringLength(x: ExternalInterfaceType): Int

-internal fun convertJsStringToKotlinString(x: ExternalInterfaceType): String {
-    val length = stringLength(x)
-    val addr = unsafeGetScratchRawMemory(INT_SIZE_BYTES + length * CHAR_SIZE_BYTES)
-    jsWriteStringIntoMemory(x, addr)
-    return String.unsafeFromCharArray(unsafeRawMemoryToWasmCharArray(addr, length))
+// kotlin string to js string export
+// TODO Uint16Array may work with byte endian different with Wasm (i.e. little endian)
+@JsFun("""(address, length, prefix) => {
+    const mem16 = new Uint16Array(wasmExports.memory.buffer, address, length);
+    const str = String.fromCharCode.apply(null, mem16);
+    return (prefix == null) ? str : prefix + str;
 }
-
-
-//language=js
-@JsFun(
-""" (str, addr) => { 
-    const memory = new DataView(wasmExports.memory.buffer);
-    for (var i = 0; i < str.length; i++) {
-        memory.setInt16(addr + i * 2, str.charCodeAt(i), true);
-    }
-}
-"""
-)
-internal external fun jsWriteStringIntoMemory(str: ExternalInterfaceType, addr: Int)
-
+""")
+internal external fun importStringFromWasm(address: Int, length: Int, prefix: ExternalInterfaceType?): ExternalInterfaceType

 internal fun kotlinToJsStringAdapter(x: String?): ExternalInterfaceType? {
    // Using nullable String to represent default value
    // for parameters with default values
-    if (x == null)
-        return null
-    return importStringFromWasm(exportString(x))
+    if (x == null) return null
+    if (x.isEmpty()) return jsEmptyString()
+
+    val srcArray = x.chars
+    val stringLength = srcArray.len()
+    val maxStringLength = unsafeGetScratchRawMemorySize() / CHAR_SIZE_BYTES
+
+    val memBuffer = unsafeGetScratchRawMemory(stringLength.coerceAtMost(maxStringLength) * CHAR_SIZE_BYTES)
+
+    var result: ExternalInterfaceType? = null
+    var srcStartIndex = 0
+    while (srcStartIndex < stringLength - maxStringLength) {
+        unsafeWasmCharArrayToRawMemory(srcArray, srcStartIndex, maxStringLength, memBuffer)
+        result = importStringFromWasm(memBuffer, maxStringLength, result)
+        srcStartIndex += maxStringLength
+    }
+
+    unsafeWasmCharArrayToRawMemory(srcArray, srcStartIndex, stringLength - srcStartIndex, memBuffer)
+    return importStringFromWasm(memBuffer, stringLength - srcStartIndex, result)
 }

+// js string to kotlin string import
+// TODO Uint16Array may work with byte endian different with Wasm (i.e. little endian)
+//language=js
+@JsFun(
+    """ (src, srcOffset, srcLength, dstAddr) => {
+        const mem16 = new Uint16Array(wasmExports.memory.buffer, dstAddr, srcLength);
+        let arrayIndex = 0;
+        let srcIndex = srcOffset;
+        while (arrayIndex < srcLength) {
+            mem16.set([src.charCodeAt(srcIndex)], arrayIndex);
+            srcIndex++;
+            arrayIndex++;
+        }
+    }
+"""
+)
+internal external fun jsExportStringToWasm(src: ExternalInterfaceType, srcOffset: Int, srcLength: Int, dstAddr: Int)
+
+internal fun jsToKotlinStringAdapter(x: ExternalInterfaceType): String {
+    val stringLength = stringLength(x)
+    val dstArray = WasmCharArray(stringLength)
+    if (stringLength == 0) {
+        return String.unsafeFromCharArray(dstArray)
+    }
+    val maxStringLength = unsafeGetScratchRawMemorySize() / CHAR_SIZE_BYTES
+
+    val memBuffer = unsafeGetScratchRawMemory(stringLength.coerceAtMost(maxStringLength) * CHAR_SIZE_BYTES)
+
+    var srcStartIndex = 0
+    while (srcStartIndex < stringLength - maxStringLength) {
+        jsExportStringToWasm(x, srcStartIndex, maxStringLength, memBuffer)
+        unsafeRawMemoryToWasmCharArray(memBuffer, srcStartIndex, maxStringLength, dstArray)
+        srcStartIndex += maxStringLength
+    }
+
+    jsExportStringToWasm(x, srcStartIndex, stringLength - srcStartIndex, memBuffer)
+    unsafeRawMemoryToWasmCharArray(memBuffer, srcStartIndex, stringLength - srcStartIndex, dstArray)
+    return String.unsafeFromCharArray(dstArray)
+}
+
+
+@JsFun("() => ''")
+internal external fun jsEmptyString(): ExternalInterfaceType
+
@JsFun("() => true")
 internal external fun jsTrue(): ExternalInterfaceType

@@ -196,9 +234,6 @@ internal fun kotlinToJsAnyAdapter(x: Any): ExternalInterfaceType =
 internal fun jsToKotlinAnyAdapter(x: ExternalInterfaceType): Any? =
    externRefToAny(x)

-internal fun jsToKotlinStringAdapter(x: ExternalInterfaceType): String =
-    convertJsStringToKotlinString(x)
-
 internal fun jsToKotlinByteAdapter(x: Int): Byte = x.toByte()
 internal fun jsToKotlinShortAdapter(x: Int): Short = x.toShort()
 internal fun jsToKotlinCharAdapter(x: Int): Char = x.toChar()
@@ -1,21 +0,0 @@
-/*
- * Copyright 2010-2021 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
-
-// This is called when exported function returns a string. It writes [i32 length, [i16 chars ...]] into a temporary raw memory area and
-// returns pointer to the start of it.
-// Note: currently there is a single temporary raw memory area so it's not possible to export more than one string at a time.
-internal fun exportString(src: String?): Int {
-    if (src == null)
-        throw IllegalArgumentException("Exporting null string")
-
-    val retAddr = unsafeGetScratchRawMemory(INT_SIZE_BYTES + src.length * CHAR_SIZE_BYTES)
-    wasm_i32_store(retAddr, src.length)
-    unsafeWasmCharArrayToRawMemory(src.chars, retAddr + INT_SIZE_BYTES)
-    return retAddr
-}
-
-// See importStringToJs for the JS-side import for strings
@@ -6,14 +6,17 @@
 package kotlin.wasm.internal

 internal const val CHAR_SIZE_BYTES = 2
-internal const val INT_SIZE_BYTES = 4

-internal fun unsafeRawMemoryToChar(addr: Int) = wasm_i32_load16_u(addr).toChar()
-
-internal fun unsafeRawMemoryToWasmCharArray(startAddr: Int, length: Int): WasmCharArray {
-    val result = WasmCharArray(length)
-    result.fill(length) { unsafeRawMemoryToChar(startAddr + it * CHAR_SIZE_BYTES) }
-    return result
+internal fun unsafeRawMemoryToWasmCharArray(srcAddr: Int, dstOffset: Int, dstLength: Int, dst: WasmCharArray) {
+    var curAddr = srcAddr
+    val srcAddrEndOffset = srcAddr + dstLength * CHAR_SIZE_BYTES
+    var dstIndex = dstOffset
+    while (curAddr < srcAddrEndOffset) {
+        val char = wasm_i32_load16_u(curAddr).toChar()
+        dst.set(dstIndex, char)
+        curAddr += CHAR_SIZE_BYTES
+        dstIndex++
+    }
 }

 // Returns a pointer into a temporary scratch segment in the raw wasm memory. Aligned by 4.
@@ -22,14 +25,19 @@ internal fun unsafeRawMemoryToWasmCharArray(startAddr: Int, length: Int): WasmCh
 internal fun unsafeGetScratchRawMemory(sizeBytes: Int): Int =
    implementedAsIntrinsic

+@ExcludedFromCodegen
+internal fun unsafeGetScratchRawMemorySize(): Int =
+    implementedAsIntrinsic
+
 // Assumes there is enough space at the destination, fails with wasm trap otherwise.
-internal fun unsafeWasmCharArrayToRawMemory(src: WasmCharArray, dstAddr: Int) {
+internal fun unsafeWasmCharArrayToRawMemory(src: WasmCharArray, srcOffset: Int, srcLength: Int, dstAddr: Int) {
    var curAddr = dstAddr
-    var i = 0
-    while (i < src.len()) {
-        wasm_i32_store16(curAddr, src.get(i))
+    val srcEndOffset = srcOffset + srcLength
+    var srcIndex = srcOffset
+    while (srcIndex < srcEndOffset) {
+        wasm_i32_store16(curAddr, src.get(srcIndex))
        curAddr += CHAR_SIZE_BYTES
-        i++
+        srcIndex++
    }
 }