[Wasm][Stdlib] Add public APIs for linear memory access

Needed for interop with APIs that use linear memory.
2022-12-29 19:25:23 +00:00
parent fbf06b5495
commit 9bc6b420a9
20 changed files with 702 additions and 56 deletions
@@ -4,6 +4,7 @@
      <w>anyref</w>
      <w>dataref</w>
      <w>ushr</w>
+      <w>wasi</w>
      <w>wasm</w>
    </words>
  </dictionary>
@@ -211,7 +211,6 @@ class WasmSymbols(
    val nullableDoubleIeee754Equals = getInternalFunction("nullableDoubleIeee754Equals")

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

@@ -568,15 +568,9 @@ class BodyGenerator(
            }

            wasmSymbols.unsafeGetScratchRawMemory -> {
-                // TODO: This drops size of the allocated segment. Instead we can check that it's in bounds for better error messages.
-                body.buildDrop()
                body.buildConstI32Symbol(context.scratchMemAddr)
            }

-            wasmSymbols.unsafeGetScratchRawMemorySize -> {
-                body.buildConstI32Symbol(WasmSymbol(context.scratchMemSizeInBytes))
-            }
-
            wasmSymbols.wasmArrayCopy -> {
                val immediate = WasmImmediate.GcType(
                    context.referenceGcType(call.getTypeArgument(0)!!.getRuntimeClass(irBuiltIns).symbol)
@@ -891,6 +885,8 @@ class BodyGenerator(
                                WasmImmediate.HeapType(WasmHeapType.Type(getReferenceGcType()))
                            WasmImmediateKind.TYPE_IDX ->
                                WasmImmediate.TypeIdx(getReferenceGcType())
+                            WasmImmediateKind.MEMORY_IDX ->
+                                WasmImmediate.MemoryIdx(0)

                            else ->
                                error("Immediate $imm is unsupported")
@@ -65,7 +65,6 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {
    val initFunctions = mutableListOf<FunWithPriority>()

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

    val stringPoolSize = WasmSymbol<Int>()

@@ -127,11 +126,8 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {
            currentDataSectionAddress += typeInfoElement.sizeInBytes
        }

-        // Reserve some memory to pass complex exported types (like strings). It's going to be accessible through 'unsafeGetScratchRawMemory'
-        // runtime call from stdlib.
        currentDataSectionAddress = alignUp(currentDataSectionAddress, INT_SIZE_BYTES)
        scratchMemAddr.bind(currentDataSectionAddress)
-        currentDataSectionAddress += scratchMemSizeInBytes

        bind(classIds.unbound, klassIds)
        interfaceId.unbound.onEachIndexed { index, entry -> entry.value.bind(index) }
@@ -176,9 +172,10 @@ class WasmCompiledModuleFragment(val irBuiltIns: IrBuiltIns) {

        val typeInfoSize = currentDataSectionAddress
        val memorySizeInPages = (typeInfoSize / 65_536) + 1
-        val memory = WasmMemory(WasmLimits(memorySizeInPages.toUInt(), memorySizeInPages.toUInt()))
+        val memory = WasmMemory(WasmLimits(memorySizeInPages.toUInt(), null /* "unlimited" */))

        // Need to export the memory in order to pass complex objects to the host language.
+        // Export name "memory" is a WASI ABI convention.
        exports += WasmExport.Memory("memory", memory)

        val importedFunctions = functions.elements.filterIsInstance<WasmFunction.Imported>()
@@ -28,9 +28,6 @@ class WasmModuleCodegenContext(
    val stringPoolSize: WasmSymbol<Int>
        get() = wasmFragment.stringPoolSize

-    val scratchMemSizeInBytes: Int
-        get() = wasmFragment.scratchMemSizeInBytes
-
    fun transformType(irType: IrType): WasmType {
        return with(typeTransformer) { irType.toWasmValueType() }
    }
@@ -64,4 +64,4 @@ versions.protobuf-relocated=2.6.1-1
 versions.r8=2.2.64
 versions.robolectric=4.0
 versions.nodejs=18.12.1
-versions.v8=10.9.194
+versions.v8=11.1.31
@@ -1,4 +1,4 @@
-// EXPECTED_REACHABLE_NODES: 1273
+// EXPECTED_REACHABLE_NODES: 1490
 // DONT_TARGET_EXACT_BACKEND: WASM
 // WASM_MUTE_REASON: UNSUPPORTED_JS_INTEROP
 // KJS_WITH_FULL_RUNTIME
@@ -9,7 +9,10 @@ import kotlin.js.Date

 fun box(): String {
    assertEquals("1970-01-01T00:00:00.000Z", Date(0L).toISOString())
-    assertEquals("10/4/1995, 12:00:00 AM", Date(1995, 9, 4, 0, 0, 0, 0L).toLocaleString("en-US"))
+    assertTrue(Date(1995, 9, 4, 0, 0, 0, 0L).toLocaleString("en-US") in listOf(
+        "10/4/1995, 12:00:00 AM",
+        "10/4/1995, 12:00:00 AM"  // New v8 uses NNBSP
+    ))
    assertEquals(812764800000.0, Date.UTC(1995, 9, 4, 0, 0, 0, 0L))

    return "OK"
@@ -6,6 +6,8 @@
 package kotlin.wasm.internal

 import kotlin.wasm.internal.reftypes.anyref
+import kotlin.wasm.unsafe.withScopedMemoryAllocator
+import kotlin.wasm.unsafe.UnsafeWasmMemoryApi

 internal external interface ExternalInterfaceType

@@ -188,20 +190,23 @@ internal fun kotlinToJsStringAdapter(x: String?): ExternalInterfaceType? {

    val srcArray = x.chars
    val stringLength = srcArray.len()
-    val maxStringLength = unsafeGetScratchRawMemorySize() / CHAR_SIZE_BYTES
+    val maxStringLength = STRING_INTEROP_MEM_BUFFER_SIZE / CHAR_SIZE_BYTES

-    val memBuffer = unsafeGetScratchRawMemory(stringLength.coerceAtMost(maxStringLength) * CHAR_SIZE_BYTES)
+    @OptIn(UnsafeWasmMemoryApi::class)
+    withScopedMemoryAllocator { allocator ->
+        val memBuffer = allocator.allocate(stringLength.coerceAtMost(maxStringLength) * CHAR_SIZE_BYTES).address.toInt()

-    var result: ExternalInterfaceType? = null
-    var srcStartIndex = 0
-    while (srcStartIndex < stringLength - maxStringLength) {
-        unsafeWasmCharArrayToRawMemory(srcArray, srcStartIndex, maxStringLength, memBuffer)
-        result = importStringFromWasm(memBuffer, maxStringLength, result)
-        srcStartIndex += maxStringLength
+        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)
    }
-
-    unsafeWasmCharArrayToRawMemory(srcArray, srcStartIndex, stringLength - srcStartIndex, memBuffer)
-    return importStringFromWasm(memBuffer, stringLength - srcStartIndex, result)
 }

 internal fun jsCheckIsNullOrUndefinedAdapter(x: ExternalInterfaceType?): ExternalInterfaceType? =
@@ -225,25 +230,31 @@ internal fun jsCheckIsNullOrUndefinedAdapter(x: ExternalInterfaceType?): Externa
 )
 internal external fun jsExportStringToWasm(src: ExternalInterfaceType, srcOffset: Int, srcLength: Int, dstAddr: Int)

+private const val STRING_INTEROP_MEM_BUFFER_SIZE = 65_536 // 1 page 4KiB
+
 internal fun jsToKotlinStringAdapter(x: ExternalInterfaceType): String {
    val stringLength = stringLength(x)
    val dstArray = WasmCharArray(stringLength)
    if (stringLength == 0) {
        return dstArray.createString()
    }
-    val maxStringLength = unsafeGetScratchRawMemorySize() / CHAR_SIZE_BYTES

-    val memBuffer = unsafeGetScratchRawMemory(stringLength.coerceAtMost(maxStringLength) * CHAR_SIZE_BYTES)
+    @OptIn(UnsafeWasmMemoryApi::class)
+    withScopedMemoryAllocator { allocator ->
+        val maxStringLength = STRING_INTEROP_MEM_BUFFER_SIZE / CHAR_SIZE_BYTES
+        val memBuffer = allocator.allocate(stringLength.coerceAtMost(maxStringLength) * CHAR_SIZE_BYTES).address.toInt()

-    var srcStartIndex = 0
-    while (srcStartIndex < stringLength - maxStringLength) {
-        jsExportStringToWasm(x, srcStartIndex, maxStringLength, memBuffer)
-        unsafeRawMemoryToWasmCharArray(memBuffer, srcStartIndex, maxStringLength, dstArray)
-        srcStartIndex += maxStringLength
+        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)
    }

-    jsExportStringToWasm(x, srcStartIndex, stringLength - srcStartIndex, memBuffer)
-    unsafeRawMemoryToWasmCharArray(memBuffer, srcStartIndex, stringLength - srcStartIndex, dstArray)
    return dstArray.createString()
 }

@@ -19,14 +19,10 @@ internal fun unsafeRawMemoryToWasmCharArray(srcAddr: Int, dstOffset: Int, dstLen
    }
 }

-// Returns a pointer into a temporary scratch segment in the raw wasm memory. Aligned by 4.
-// Note: currently there is single such segment for a whole wasm module, so use with care.
+// Returns starting address of unused linear memory.
@ExcludedFromCodegen
-internal fun unsafeGetScratchRawMemory(sizeBytes: Int): Int =
-    implementedAsIntrinsic
-
-@ExcludedFromCodegen
-internal fun unsafeGetScratchRawMemorySize(): Int =
+@PublishedApi
+internal fun unsafeGetScratchRawMemory(): Int =
    implementedAsIntrinsic

 // Assumes there is enough space at the destination, fails with wasm trap otherwise.
@@ -0,0 +1,18 @@
+/*
+ * Copyright 2010-2022 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.unsafe
+
+import kotlin.annotation.AnnotationTarget.*
+
+/**
+ * This annotation marks APIs for working with unmanaged WebAssembly linear memory.
+ *
+ * Any usage of a declaration annotated with `@UnsafeWasmMemoryApi` must be accepted either by
+ * annotating that usage with the [OptIn] annotation, e.g. `@OptIn(UnsafeWasmMemoryApi::class)`,
+ * or by using the compiler argument `-opt-in=kotlin.wasm.unsafe.UnsafeWasmMemoryApi`.
+ */
+@RequiresOptIn("Unsafe APIs to access to WebAssembly linear memory")
+public annotation class UnsafeWasmMemoryApi
@@ -0,0 +1,73 @@
+/*
+ * Copyright 2010-2022 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.unsafe
+
+import kotlin.wasm.internal.WasmOp
+import kotlin.wasm.internal.implementedAsIntrinsic
+
+/**
+ * Linear memory pointer type.
+ * Corresponds to `i32` type on 32-bit Wasm architecture.
+ */
+@UnsafeWasmMemoryApi
+public value class Pointer public constructor(public val address: UInt) {
+
+    /** Adds an [Int] to the address of this [Pointer] */
+    public operator fun plus(other: Int): Pointer =
+        Pointer(address + other.toUInt())
+
+    /** Subtracts an [Int] from the address of this [Pointer] */
+    public operator fun minus(other: Int): Pointer =
+        Pointer(address - other.toUInt())
+
+    /** Adds an [UInt] to the address of this [Pointer] */
+    public operator fun plus(other: UInt): Pointer =
+        Pointer(address + other)
+
+    /** Subtracts an [UInt] from the address of this [Pointer] */
+    public operator fun minus(other: UInt): Pointer =
+        Pointer(address - other)
+
+    /** Load a Byte (8 bit) value */
+    @WasmOp(WasmOp.I32_LOAD8_S)
+    public fun loadByte(): Byte =
+        implementedAsIntrinsic
+
+    /** Load a Short (16 bit) value */
+    @WasmOp(WasmOp.I32_LOAD16_S)
+    public fun loadShort(): Short =
+        implementedAsIntrinsic
+
+    /** Load an Int (32 bit) value */
+    @WasmOp(WasmOp.I32_LOAD)
+    public fun loadInt(): Int =
+        implementedAsIntrinsic
+
+    /** Load a Long (64 bit) value */
+    @WasmOp(WasmOp.I64_LOAD)
+    public fun loadLong(): Long =
+        implementedAsIntrinsic
+
+    /** Store a Byte (8 bit) [value] */
+    @WasmOp(WasmOp.I32_STORE8)
+    public fun storeByte(value: Byte): Unit =
+        implementedAsIntrinsic
+
+    /** Store a Short (16 bit) [value] */
+    @WasmOp(WasmOp.I32_STORE16)
+    public fun storeShort(value: Short): Unit =
+        implementedAsIntrinsic
+
+    /** Store an Int (32 bit) [value] */
+    @WasmOp(WasmOp.I32_STORE)
+    public fun storeInt(value: Int): Unit =
+        implementedAsIntrinsic
+
+    /** Store a Long (64 bit) [value] */
+    @WasmOp(WasmOp.I64_STORE)
+    public fun storeLong(value: Long): Unit =
+        implementedAsIntrinsic
+}
@@ -0,0 +1,156 @@
+/*
+ * Copyright 2010-2022 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.unsafe
+
+import kotlin.wasm.internal.WasmOp
+import kotlin.wasm.internal.implementedAsIntrinsic
+import kotlin.wasm.internal.unsafeGetScratchRawMemory
+import kotlin.contracts.*
+
+/**
+ * WebAssembly linear memory allocator.
+ */
+@UnsafeWasmMemoryApi
+public abstract class MemoryAllocator {
+    /**
+     * Allocates a block of uninitialized linear memory of the given [size] in bytes.
+     *
+     * @return an address of allocated memory. It is guaranteed to be a multiple of 8.
+     */
+    public abstract fun allocate(size: Int): Pointer
+}
+
+/**
+ * Runs the [block] of code, providing it a temporary [MemoryAllocator] as an argument, and returns the result of this block.
+ *
+ * Frees all memory allocated with the provided allocator after running the [block].
+ *
+ * This function is intened to facilitate the exchange of values with outside world through linear memory.
+ * For example:
+ *
+ * ```
+ * val buffer_size = ...
+ * withScopedMemoryAllocator { allocator ->
+ *     val buffer_address = allocator.allocate(buffer_size)
+ *     importedWasmFunctionThatWritesToBuffer(buffer_address, buffer_size)
+ *     return readDataFromBufferIntoManagedKotlinMemory(buffer_address, buffer_size)
+ * }
+ * ```
+ *
+ * WARNING! Addresses allocated inside the [block] function become invalid after exiting the function.
+ *
+ * WARNING! A nested call to [withScopedMemoryAllocator] will temporarily disable the allocator from the outer scope
+ *   for the duration of the call. Calling [MemoryAllocator.allocate] on a disabled allocator
+ *   will throw [IllegalStateException].
+ *
+ * WARNING! Accessing the allocator outside of the [block] scope will throw [IllegalStateException].
+ */
+@UnsafeWasmMemoryApi
+public inline fun <T> withScopedMemoryAllocator(
+    block: (allocator: MemoryAllocator) -> T
+): T {
+    contract { callsInPlace(block, InvocationKind.EXACTLY_ONCE) }
+    val allocator = createAllocatorInTheNewScope()
+    val result = try {
+        block(allocator)
+    } finally {
+        allocator.destroy()
+        currentAllocator = allocator.parent
+    }
+    return result
+}
+
+@PublishedApi
+@UnsafeWasmMemoryApi
+internal fun createAllocatorInTheNewScope(): ScopedMemoryAllocator {
+    val allocator = currentAllocator?.createChild() ?:
+        ScopedMemoryAllocator(unsafeGetScratchRawMemory(), parent = null)
+    currentAllocator = allocator
+    return allocator
+}
+
+
+@PublishedApi
+@UnsafeWasmMemoryApi
+internal var currentAllocator: ScopedMemoryAllocator? = null
+
+@PublishedApi
+@UnsafeWasmMemoryApi
+internal class ScopedMemoryAllocator(
+    startAddress: Int,
+    // Allocator from parent scope or null for top-level scope.
+    @PublishedApi
+    internal var parent: ScopedMemoryAllocator?,
+) : MemoryAllocator() {
+    // true if allocator is out of scope
+    private var destroyed = false
+    // true if child allocator is active
+    private var suspended = false
+    // all memory is available starting from this address
+    private var availableAddress: ULong = startAddress.toULong()
+
+    override fun allocate(size: Int): Pointer {
+        check(!destroyed) { "ScopedMemoryAllocator is destroyed when out of scope" }
+        check(!suspended) { "ScopedMemoryAllocator is suspended when nested allocators are used" }
+
+        // Pad available address to align it to 8
+        // 8 is a max alignment number currently needed for Wasm component model canonical ABI
+        val align = 8uL
+        val result = (availableAddress + align - 1uL) and (align - 1uL).inv()
+        check(result % 8uL == 0uL)
+
+        availableAddress = result + size.toULong()
+
+        if (availableAddress > UInt.MAX_VALUE.toULong()) {
+            error("Out of linear memory. All available address space (4gb) is used.")
+        }
+
+        val currentMaxSize = wasmMemorySize().toULong() * WASM_PAGE_SIZE_IN_BYTES.toULong()
+        if (availableAddress >= currentMaxSize) {
+
+            val numPagesToGrow =
+                (availableAddress - currentMaxSize) / WASM_PAGE_SIZE_IN_BYTES.toULong() + 2uL
+
+            if (wasmMemoryGrow(numPagesToGrow.toInt()) == -1) {
+                error("Out of linear memory. memory.grow returned -1")
+            }
+        }
+
+        check(availableAddress < wasmMemorySize().toULong() * WASM_PAGE_SIZE_IN_BYTES.toULong())
+
+        return Pointer(result.toUInt())
+    }
+
+    @PublishedApi
+    internal fun createChild(): ScopedMemoryAllocator {
+        val child = ScopedMemoryAllocator(availableAddress.toInt(), parent = this)
+        suspended = true
+        return child
+    }
+
+    @PublishedApi
+    internal fun destroy() {
+        destroyed = true
+        parent?.suspended = false
+    }
+}
+
+private const val WASM_PAGE_SIZE_IN_BYTES = 65_536  // 64 KiB
+
+/**
+ * Current linear memory size in pages
+ */
+@WasmOp(WasmOp.MEMORY_SIZE)
+internal fun wasmMemorySize(): Int =
+    implementedAsIntrinsic
+
+/**
+ * Grow memory by a given delta (in pages).
+ * Return the previous size, or -1 if enough memory cannot be allocated.
+ */
+@WasmOp(WasmOp.MEMORY_GROW)
+internal fun wasmMemoryGrow(delta: Int): Int =
+    implementedAsIntrinsic
@@ -0,0 +1,184 @@
+package test.wasm.unsafe
+
+import kotlin.wasm.unsafe.*
+import kotlin.test.*
+
+@OptIn(UnsafeWasmMemoryApi::class)
+class MemoryAccessTestTest {
+    @Test
+    fun testPointer() {
+        val p: Pointer = Pointer(19u)
+        assertEquals(p.address, 19u)
+        assertEquals((p + 10u).address, 29u)
+        assertEquals((p - 10u).address, 9u)
+        assertEquals((p + 10).address, 29u)
+        assertEquals((p - 10).address, 9u)
+    }
+
+    fun <T> testLoadStore(values: List<T>, typeSize: Int, store: (Pointer, T) -> Unit, load: (Pointer) -> T) {
+        withScopedMemoryAllocator { a ->
+            // Memory layout: [Long1][T][Long2]
+            val ptrToLong1 = a.allocate(24)
+            val ptrToT = ptrToLong1 + 8
+            val ptrToLong2 = ptrToT + typeSize
+
+            for (x in values) {
+                val prevLong = ptrToLong1.loadLong()
+                val nextLong = ptrToLong2.loadLong()
+
+                store(ptrToT, x)
+                assertEquals(load(ptrToT), x)
+
+                assertEquals(ptrToLong1.loadLong(), prevLong)
+                assertEquals(ptrToLong2.loadLong(), nextLong)
+            }
+        }
+    }
+
+    @Test
+    fun testByte() {
+        val bytes = listOf<Byte>(0, Byte.MIN_VALUE, Byte.MAX_VALUE, -91, -21, -47, -72, -42, 118, 120, 125, 21, -43)
+        testLoadStore(bytes, 1, { p, v -> p.storeByte(v) }, { it.loadByte() })
+    }
+
+    @Test
+    fun testShort() {
+        val shorts = listOf<Short>(
+            0, Short.MIN_VALUE, Short.MAX_VALUE, 26350, 17667, 5437, 1381,
+            21183, 26042, -25961, -22913, 9128, -10684
+        )
+        testLoadStore(shorts, 2, { p, v -> p.storeShort(v) }, { it.loadShort() })
+    }
+
+    @Test
+    fun testInt() {
+        val ints = listOf<Int>(
+            0, Int.MIN_VALUE, Int.MAX_VALUE,
+            1348618689, -299556943, -394977414, -621300994, 1034622853, -1010496662,
+            -2102993550, 199131417, 407819728, -1093382545
+        )
+        testLoadStore(ints, 4, { p, v -> p.storeInt(v) }, { it.loadInt() })
+    }
+
+    @Test
+    fun testLong() {
+        val longs = listOf<Long>(
+            0, Long.MIN_VALUE, Long.MAX_VALUE,
+            6964777768087685094, -3965399897925814666, 6876207943944046195, 7675081221595661767,
+            -3388229176969119769, 4265730675328983821, -4893379785828386453, -7516879919690485136,
+            8512965883914804069, 6155050932825287650
+        )
+        testLoadStore(longs, 8, { p, v -> p.storeLong(v) }, { it.loadLong() })
+    }
+
+    @Test
+    fun testAccessingWithDifferentTypes() {
+        withScopedMemoryAllocator { a ->
+            val size = 16
+            val sizeU = size.toUInt()
+            val pointer = a.allocate(size)
+            val addr = pointer.address
+
+            fun fillWith(value: Byte) {
+                for (ptr in addr..<addr + sizeU) {
+                    Pointer(ptr).storeByte(value)
+                }
+            }
+
+            fun fillWith(value: Short) {
+                for (ptr in addr..<addr + sizeU step 2) {
+                    Pointer(ptr).storeShort(value)
+                }
+            }
+
+            fun fillWith(value: Int) {
+                for (ptr in addr..<addr + sizeU step 4) {
+                    Pointer(ptr).storeInt(value)
+                }
+            }
+
+            fun fillWith(value: Long) {
+                for (ptr in addr..<addr + sizeU step 8) {
+                    Pointer(ptr).storeLong(value)
+                }
+            }
+
+            fun checkMem(
+                bytes: List<Byte>,
+                shorts: List<Short>,
+                ints: List<Int>,
+                longs: List<Long>
+            ) {
+                assertEquals(bytes.size, size)
+                assertEquals(shorts.size, size / 2)
+                assertEquals(ints.size, size / 4)
+                assertEquals(longs.size, size / 8)
+                for (i in 0..<size) {
+                    assertEquals((pointer + i).loadByte(), bytes[i])
+                }
+                for (i in 0..<size / 2) {
+                    assertEquals((pointer + i * 2).loadShort(), shorts[i])
+                }
+                for (i in 0..<size / 4) {
+                    assertEquals((pointer + i * 4).loadInt(), ints[i])
+                }
+                for (i in 0..<size / 8) {
+                    assertEquals((pointer + i * 8).loadLong(), longs[i])
+                }
+            }
+
+            fun checkZero() {
+                checkMem(
+                    bytes = List(size) { 0 },
+                    shorts = List(size / 2) { 0 },
+                    ints = List(size / 4) { 0 },
+                    longs = List(size / 8) { 0L }
+                )
+            }
+
+            fillWith(0.toByte())
+            checkZero()
+            fillWith(0x0F.toByte())
+            checkMem(
+                bytes = List(size) { 0x0F },
+                shorts = List(size / 2) { 0x0F0F },
+                ints = List(size / 4) { 0x0F0F0F0F },
+                longs = List(size / 8) { 0x0F0F0F0F0F0F0F0FL }
+            )
+
+            fillWith(0.toShort())
+            checkZero()
+            fillWith(0xABCDu.toShort())
+            checkMem(
+                bytes = mutableListOf<Byte>().also { list ->
+                    repeat(size / 2) {
+                        // little-endian
+                        list += 0xCD.toByte()
+                        list += 0xAB.toByte()
+                    }
+                },
+                shorts = List(size / 2) { 0xABCDu.toShort() },
+                ints = List(size / 4) { 0xABCDABCDu.toInt() },
+                longs = List(size / 8) { 0xABCDABCDABCDABCDuL.toLong() }
+            )
+            fillWith(0.toInt())
+            checkZero()
+            fillWith(0xFFFFFFFFu.toInt())
+            checkMem(
+                bytes = List(size) { 0xFFu.toByte() },
+                shorts = List(size / 2) { 0xFFFFu.toShort() },
+                ints = List(size / 4) { 0xFFFFFFFFu.toInt() },
+                longs = List(size / 8) { 0xFFFFFFFFFFFFFFFFuL.toLong() }
+            )
+            fillWith(0L)
+            checkZero()
+            fillWith(0x1212121212121212L)
+            checkMem(
+                bytes = List(size) { 0x012 },
+                shorts = List(size / 2) { 0x1212 },
+                ints = List(size / 4) { 0x12121212 },
+                longs = List(size / 8) { 0x1212121212121212L }
+            )
+        }
+    }
+}
@@ -0,0 +1,217 @@
+package test.wasm.unsafe
+
+import kotlin.wasm.unsafe.*
+import kotlin.test.*
+
+@JsFun("(a, b) => a + b")
+private external fun jsConcatStrings(a: String, b: String): String
+
+@OptIn(UnsafeWasmMemoryApi::class)
+class MemoryAllocationTest {
+    val pageSize = 65_536
+
+    @Test
+    fun testWasmMemorySizeGrow() {
+        val s1 = wasmMemorySize()
+        val grow_res = wasmMemoryGrow(10)
+        val s2 = wasmMemorySize()
+        assertNotEquals(grow_res, -1)
+        assertEquals(grow_res, s1)
+        assertEquals(s2 - s1, 10)
+    }
+
+    @Test
+    fun testScopedAllocator() {
+        val sizes = listOf<Int>(1, 1, 2, 3, 8, 10, 305, 12_747, 31_999)
+        val allocations = mutableListOf<Pointer>()
+        withScopedMemoryAllocator { a ->
+            for (size in sizes) {
+                allocations += a.allocate(size)
+            }
+        }
+
+        val allocations2 = mutableListOf<Pointer>()
+        withScopedMemoryAllocator { a ->
+            for (size in sizes) {
+                allocations2 += a.allocate(size)
+            }
+        }
+
+        // Test that we run withScopedMemoryAllocator body
+        assertEquals(allocations.size, sizes.size)
+
+        // Memory should be reusing in different scopes
+        // NOTE: This is current impl detail and can be changed
+        assertEquals(allocations, allocations2)
+
+        // Allocations are aligned
+        assertTrue(allocations.all { it.address % 8u == 0u })
+
+        // Allocations are different
+        assertTrue(allocations.distinct().size == allocations.size)
+
+        // Allocations do not intersect
+        for (i1 in 0..<sizes.size) {
+            for (i2 in (i1 + 1)..<sizes.size) {
+                val a1 = allocations[i1].address
+                val a2 = allocations[i2].address
+                val size1 = sizes[i1].toUInt()
+                val size2 = sizes[i2].toUInt()
+                assertTrue(a1 !in a2..<a2 + size2)
+                assertTrue(a1 + size1 - 1u !in a2..<a2 + size2)
+            }
+        }
+    }
+
+    @Test
+    fun testScopedAllocatorGrowsMemory() {
+        // Allocations past current memory size should grow memory
+        val memSizes = mutableListOf<Int>(wasmMemorySize())
+        withScopedMemoryAllocator { a ->
+            var allocatedAddress = a.allocate(pageSize)
+            var allocationSize = pageSize
+
+            repeat(10) {
+                var currPagesUsed = (allocatedAddress.address.toInt() + allocationSize + 1) / pageSize
+                var currPagesAvailable = wasmMemorySize()
+                assertTrue(currPagesAvailable > currPagesUsed)
+                // Allocate 10 pages past max page
+                allocationSize = (currPagesAvailable - currPagesUsed + 10) * pageSize
+                allocatedAddress = a.allocate(allocationSize)
+                memSizes += wasmMemorySize()
+            }
+        }
+        assertTrue(memSizes.size == 11)
+        assertEquals(memSizes.distinct(), memSizes)
+        assertEquals(memSizes.sorted(), memSizes)
+    }
+
+    @Test
+    fun nestedAllocators() {
+        val sizes = listOf<Int>(1, 1, 2, 3, 8, 10, 305, 12_747, 31_999)
+        var allocations1: List<Pointer>
+        var allocations1_1: List<Pointer>
+        var allocations1_2: List<Pointer>
+        var allocations2: List<Pointer>
+
+        withScopedMemoryAllocator { allocator0 ->
+            allocations1 = sizes.map { size -> allocator0.allocate(size) }
+            withScopedMemoryAllocator { allocator0_0 ->
+                allocations1_1 = sizes.map { size -> allocator0_0.allocate(size) }
+            }
+            withScopedMemoryAllocator { allocator0_1 ->
+                allocations1_2 = sizes.map { size -> allocator0_1.allocate(size) }
+            }
+        }
+        withScopedMemoryAllocator { allocator0 ->
+            allocations2 = sizes.map { size -> allocator0.allocate(size) }
+        }
+
+        // Check that all allocations happened and distinct
+        listOf(
+            allocations1,
+            allocations1_1,
+            allocations1_2,
+            allocations2
+        ).forEach { allocation ->
+            assertEquals(allocation.size, sizes.size)
+            assertEquals(allocation.distinct().size, allocation.size)
+        }
+
+        // Impl detail: sibling scopes use the same memory
+        assertEquals(allocations1, allocations2)
+        assertEquals(allocations1_1, allocations1_2)
+
+        // Impl detaiol: allocator in child scope allocates new memory
+        val max1: UInt = allocations1.maxOf { it.address }
+        val min1_1: UInt = allocations1_1.minOf { it.address }
+        assertTrue(max1 < min1_1)
+    }
+
+    @Test
+    fun testJsIntropInsideAllocations() {
+        withScopedMemoryAllocator { allocator ->
+            assertEquals(jsConcatStrings("str1", "str2"), "str1str2")
+            allocator.allocate(10)
+        }
+    }
+
+    @Test
+    fun testNestedAllocatorThrows() {
+        var leakedAllocator1: MemoryAllocator? = null
+        var leakedAllocator2: MemoryAllocator? = null
+        var leakedAllocator3: MemoryAllocator? = null
+
+        withScopedMemoryAllocator { allocator1 ->
+            leakedAllocator1 = allocator1
+            allocator1.allocate(100)
+            // 2-level nesting
+            withScopedMemoryAllocator { allocator2 ->
+                leakedAllocator2 = allocator2
+                allocator2.allocate(100)
+                assertFailsWith<IllegalStateException> {
+                    allocator1.allocate(100)
+                }
+                // 3-level nesting
+                withScopedMemoryAllocator { allocator3 ->
+                    leakedAllocator3 = allocator3
+                    allocator3.allocate(100)
+                    assertFailsWith<IllegalStateException> {
+                        allocator1.allocate(100)
+                    }
+                    assertFailsWith<IllegalStateException> {
+                        allocator2.allocate(100)
+                    }
+                    allocator3.allocate(100)
+                }
+                assertFailsWith<IllegalStateException> {
+                    leakedAllocator3?.allocate(100)
+                }
+                // now it is legal to use allocator1 since we're in its immediate scope
+                allocator2.allocate(100)
+            }
+            assertFailsWith<IllegalStateException> {
+                leakedAllocator2?.allocate(100)
+            }
+            allocator1.allocate(100)
+        }
+
+        for (leakedAllocator in listOf(leakedAllocator1, leakedAllocator2, leakedAllocator3)) {
+            assertFailsWith<IllegalStateException> {
+                leakedAllocator?.allocate(100)
+            }
+        }
+    }
+
+    @Test
+    fun testScopedAllocatorThrows() {
+        assertFailsWith<IllegalStateException> {
+            var leakedAllocator: MemoryAllocator? = null
+            withScopedMemoryAllocator { allocator ->
+                leakedAllocator = allocator
+            }
+            leakedAllocator?.allocate(10)
+        }
+
+        assertFailsWith<IllegalStateException> {
+            var leakedAllocator: MemoryAllocator? = null
+            try {
+                withScopedMemoryAllocator { allocator ->
+                    leakedAllocator = allocator
+                    throw Error()
+                }
+            } catch (e: Throwable) {
+            }
+            leakedAllocator?.allocate(10)
+        }
+
+        assertFailsWith<IllegalStateException> {
+            fun foo(): MemoryAllocator {
+                withScopedMemoryAllocator { allocator ->
+                    return allocator  // non-local return
+                }
+            }
+            foo().allocate(10)
+        }
+    }
+}
@@ -14,7 +14,9 @@ with(org.jetbrains.kotlin.gradle.targets.js.nodejs.NodeJsRootPlugin.apply(rootPr
 }

 with(org.jetbrains.kotlin.gradle.targets.js.d8.D8RootPlugin.apply(rootProject)) {
-    version = "10.9.194"
+    // Test that we can set the version and it is a String.
+    // But use the default version since update this place every time anyway.
+    version = (version as String)
 }

 with(org.jetbrains.kotlin.gradle.targets.js.yarn.YarnPlugin.apply(rootProject)) {
@@ -41,7 +41,7 @@ open class D8RootExtension(@Transient val rootProject: Project) : ConfigurationP
        fi;
    done;
    */
-    var version by Property("10.9.194")
+    var version by Property("11.1.31")
    var edition by Property("rel") // rel or dbg

    val setupTaskProvider: TaskProvider<out Copy>
@@ -69,9 +69,7 @@ sealed class WasmImmediate {

    class ValTypeVector(val value: List<WasmType>) : WasmImmediate()

-    class MemoryIdx(val value: WasmSymbol<WasmMemory>) : WasmImmediate() {
-        constructor(value: WasmMemory) : this(WasmSymbol(value))
-    }
+    class MemoryIdx(val value: Int) : WasmImmediate()

    class DataIdx(val value: WasmSymbol<Int>) : WasmImmediate() {
        constructor(value: Int) : this(WasmSymbol(value))
@@ -425,7 +425,7 @@ class WasmBinaryToIR(val b: MyByteReader) {
                WasmImmediateKind.LOCAL_IDX -> WasmImmediate.LocalIdx(locals[b.readVarUInt32AsInt()])
                WasmImmediateKind.GLOBAL_IDX -> WasmImmediate.GlobalIdx(globalByIdx(b.readVarUInt32AsInt()))
                WasmImmediateKind.TYPE_IDX -> WasmImmediate.TypeIdx(functionTypes[b.readVarUInt32AsInt()])
-                WasmImmediateKind.MEMORY_IDX -> WasmImmediate.MemoryIdx(memoryByIdx(b.readVarUInt32AsInt()))
+                WasmImmediateKind.MEMORY_IDX -> WasmImmediate.MemoryIdx(b.readVarUInt32AsInt())
                WasmImmediateKind.DATA_IDX -> WasmImmediate.DataIdx(b.readVarUInt32AsInt())
                WasmImmediateKind.TABLE_IDX -> WasmImmediate.TableIdx(b.readVarUInt32AsInt())
                WasmImmediateKind.LABEL_IDX -> WasmImmediate.LabelIdx(b.readVarUInt32AsInt())
@@ -243,7 +243,7 @@ class WasmIrToBinary(
            is WasmImmediate.LocalIdx -> appendLocalReference(x.value.owner)
            is WasmImmediate.GlobalIdx -> appendModuleFieldReference(x.value.owner)
            is WasmImmediate.TypeIdx -> appendModuleFieldReference(x.value.owner)
-            is WasmImmediate.MemoryIdx -> appendModuleFieldReference(x.value.owner)
+            is WasmImmediate.MemoryIdx -> b.writeVarUInt32(x.value)
            is WasmImmediate.DataIdx -> b.writeVarUInt32(x.value.owner)
            is WasmImmediate.TableIdx -> b.writeVarUInt32(x.value.owner)
            is WasmImmediate.LabelIdx -> b.writeVarUInt32(x.value)
@@ -124,7 +124,7 @@ class WasmIrToText : SExpressionBuilder() {
            is WasmImmediate.LocalIdx -> appendLocalReference(x.value.owner)
            is WasmImmediate.GlobalIdx -> appendModuleFieldReference(x.value.owner)
            is WasmImmediate.TypeIdx -> sameLineList("type") { appendModuleFieldReference(x.value.owner) }
-            is WasmImmediate.MemoryIdx -> appendModuleFieldIdIfNotNull(x.value.owner)
+            is WasmImmediate.MemoryIdx -> appendIdxIfNotZero(x.value)
            is WasmImmediate.DataIdx -> appendElement(x.value.toString())
            is WasmImmediate.TableIdx -> appendElement(x.value.toString())
            is WasmImmediate.LabelIdx -> appendElement(x.value.toString())
@@ -523,9 +523,7 @@ class WasmIrToText : SExpressionBuilder() {
        appendElement("$${local.id}_${sanitizeWatIdentifier(local.name)}")
    }

-    fun appendModuleFieldIdIfNotNull(field: WasmNamedModuleField) {
-        val id = field.id
-            ?: error("${field::class} ${field.name} ID is unlinked")
+    fun appendIdxIfNotZero(id: Int) {
        if (id != 0) appendElement(id.toString())
    }