Bytebuffer abstraction (#9)

build.gradle

@@ -21,7 +21,7 @@ buildscript {
 allprojects {
     apply plugin: 'java'
     group 'com.github.cretz.asmble'
-    version '0.4.4-fl'
+    version '0.4.5-fl'
 
     // skips building and running for the specified examples
     ext.skipExamples = ['c-simple', 'go-simple', 'rust-regex']

compiler/src/main/java/asmble/compile/jvm/MemoryBuffer.java

@@ -0,0 +1,40 @@
+package asmble.compile.jvm;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
+/**
+ * The abstraction that describes work with the memory of the virtual machine.
+ */
+public abstract class MemoryBuffer {
+
+    /**
+     * The default implementation of MemoryBuffer that based on java.nio.DirectByteBuffer
+     */
+    public static MemoryBuffer init(int capacity) {
+        return new MemoryByteBuffer(ByteBuffer.allocateDirect(capacity));
+    }
+
+    public abstract int capacity();
+    public abstract int limit();
+    public abstract MemoryBuffer clear();
+    public abstract MemoryBuffer limit(int newLimit);
+    public abstract MemoryBuffer position(int newPosition);
+    public abstract MemoryBuffer order(ByteOrder order);
+    public abstract MemoryBuffer duplicate();
+    public abstract MemoryBuffer put(byte[] arr, int offset, int length);
+    public abstract MemoryBuffer put(byte[] arr);
+    public abstract MemoryBuffer put(int index, byte b);
+    public abstract MemoryBuffer putInt(int index, int n);
+    public abstract MemoryBuffer putLong(int index, long n);
+    public abstract MemoryBuffer putDouble(int index, double n);
+    public abstract MemoryBuffer putShort(int index, short n);
+    public abstract MemoryBuffer putFloat(int index, float n);
+    public abstract byte get(int index);
+    public abstract int getInt(int index);
+    public abstract long getLong(int index);
+    public abstract short getShort(int index);
+    public abstract float getFloat(int index);
+    public abstract double getDouble(int index);
+    public abstract MemoryBuffer get(byte[] arr);
+}

compiler/src/main/java/asmble/compile/jvm/MemoryBufferBuilder.java

@@ -0,0 +1,8 @@
+package asmble.compile.jvm;
+
+/**
+ * Interface to initialize MemoryBuffer
+ */
+public interface MemoryBufferBuilder {
+    MemoryBuffer build(int capacity);
+}

compiler/src/main/kotlin/asmble/cli/ScriptCommand.kt

@@ -1,7 +1,7 @@
 package asmble.cli
 
 import asmble.ast.Script
-import asmble.compile.jvm.javaIdent
+import asmble.compile.jvm.*
 import asmble.run.jvm.LoggerModule
 import asmble.run.jvm.Module
 import asmble.run.jvm.ScriptContext
@@ -56,7 +56,8 @@ abstract class ScriptCommand<T> : Command<T>() {
     fun prepareContext(args: ScriptArgs): ScriptContext {
         var context = ScriptContext(
             packageName = "asmble.temp" + UUID.randomUUID().toString().replace("-", ""),
-            defaultMaxMemPages = args.defaultMaxMemPages
+            defaultMaxMemPages = args.defaultMaxMemPages,
+            memoryBuilder = args.memoryBuilder
         )
         // Compile everything
         context = args.inFiles.foldIndexed(context) { index, ctx, inFile ->
@@ -86,10 +87,23 @@ abstract class ScriptCommand<T> : Command<T>() {
                 throw Exception("Failed loading $inFile - ${e.message}", e)
             }
         }
+
+        val memBuilder = args.memoryBuilder
+
+        // throws ArithmeticException if the result overflows an int
+        val capacity = Math.multiplyExact(args.defaultMaxMemPages, Mem.PAGE_SIZE)
+
         // Do registrations
         context = args.registrations.fold(context) { ctx, (moduleName, className) ->
-            ctx.withModuleRegistered(moduleName,
+            if (memBuilder != null) {
-                Module.Native(Class.forName(className, true, ctx.classLoader).newInstance()))
+                ctx.withModuleRegistered(moduleName,
+                        Module.Native(Class.forName(className, true, ctx.classLoader)
+                                .getConstructor(MemoryBuffer::class.java)
+                                .newInstance(memBuilder.build(capacity))))
+            } else {
+                ctx.withModuleRegistered(moduleName,
+                        Module.Native(Class.forName(className, true, ctx.classLoader).newInstance()))
+            }
         }
         if (args.specTestRegister) context = context.withHarnessRegistered()
 
@@ -113,6 +127,7 @@ abstract class ScriptCommand<T> : Command<T>() {
      * @param specTestRegister If true, registers the spec test harness as 'spectest'
      * @param defaultMaxMemPages The maximum number of memory pages when a module doesn't say
      * @param loggerMemPages The maximum number of memory pages of the logger module.
+     * @param memoryBuilder The builder to initialize new memory class.
      */
     data class ScriptArgs(
         val inFiles: List<String>,
@@ -120,6 +135,7 @@ abstract class ScriptCommand<T> : Command<T>() {
         val disableAutoRegister: Boolean,
         val specTestRegister: Boolean,
         val defaultMaxMemPages: Int,
-        val loggerMemPages: Int
+        val loggerMemPages: Int,
+        val memoryBuilder: MemoryBufferBuilder? = null
     )
 }

compiler/src/main/kotlin/asmble/compile/jvm/ByteBufferMem.kt

@@ -4,32 +4,30 @@ import asmble.ast.Node
 import org.objectweb.asm.Opcodes
 import org.objectweb.asm.Type
 import org.objectweb.asm.tree.*
-import java.nio.Buffer
-import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import kotlin.reflect.KClass
 import kotlin.reflect.KFunction
 
-open class ByteBufferMem(val direct: Boolean = true) : Mem {
+open class ByteBufferMem : Mem {
-    override val memType = ByteBuffer::class.ref
+    override val memType: TypeRef = MemoryBuffer::class.ref
 
-    override fun limitAndCapacity(instance: Any) =
+    override fun limitAndCapacity(instance: Any): Pair<Int, Int> =
-        if (instance !is ByteBuffer) error("Unrecognized memory instance: $instance")
+        if (instance !is MemoryBuffer) error("Unrecognized memory instance: $instance")
         else instance.limit() to instance.capacity()
 
     override fun create(func: Func) = func.popExpecting(Int::class.ref).addInsns(
-        (if (direct) ByteBuffer::allocateDirect else ByteBuffer::allocate).invokeStatic()
+        (MemoryBuffer::init).invokeStatic()
     ).push(memType)
 
     override fun init(func: Func, initial: Int) = func.popExpecting(memType).addInsns(
         // Set the limit to initial
         (initial * Mem.PAGE_SIZE).const,
-        forceFnType<ByteBuffer.(Int) -> Buffer>(ByteBuffer::limit).invokeVirtual(),
+        forceFnType<MemoryBuffer.(Int) -> MemoryBuffer>(MemoryBuffer::limit).invokeVirtual(),
-        TypeInsnNode(Opcodes.CHECKCAST, ByteBuffer::class.ref.asmName),
+        TypeInsnNode(Opcodes.CHECKCAST, memType.asmName),
         // Set it to use little endian
         ByteOrder::LITTLE_ENDIAN.getStatic(),
-        forceFnType<ByteBuffer.(ByteOrder) -> ByteBuffer>(ByteBuffer::order).invokeVirtual()
+        forceFnType<MemoryBuffer.(ByteOrder) -> MemoryBuffer>(MemoryBuffer::order).invokeVirtual()
-    ).push(ByteBuffer::class.ref)
+    ).push(memType)
 
     override fun data(func: Func, bytes: ByteArray, buildOffset: (Func) -> Func) =
         // Sadly there is no absolute bulk put, so we need to fake one. Ref:
@@ -42,10 +40,10 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
         // where we could call put directly, but it too is negligible for now.
         // Note, with this approach, the mem not be left on the stack for future data() calls which is fine.
         func.popExpecting(memType).
-            addInsns(ByteBuffer::duplicate.invokeVirtual()).
+            addInsns(MemoryBuffer::duplicate.invokeVirtual()).
             let(buildOffset).popExpecting(Int::class.ref).
             addInsns(
-                forceFnType<ByteBuffer.(Int) -> Buffer>(ByteBuffer::position).invokeVirtual(),
+                forceFnType<MemoryBuffer.(Int) -> MemoryBuffer>(MemoryBuffer::position).invokeVirtual(),
                 TypeInsnNode(Opcodes.CHECKCAST, memType.asmName)
             ).addInsns(
                 // We're going to do this as an LDC string in ISO-8859 and read it back at runtime. However,
@@ -61,7 +59,7 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
                             "getBytes", "(Ljava/lang/String;)[B", false),
                         0.const,
                         bytes.size.const,
-                        forceFnType<ByteBuffer.(ByteArray, Int, Int) -> ByteBuffer>(ByteBuffer::put).invokeVirtual()
+                        forceFnType<MemoryBuffer.(ByteArray, Int, Int) -> MemoryBuffer>(MemoryBuffer::put).invokeVirtual()
                     )
                 }.toList()
             ).addInsns(
@@ -69,7 +67,7 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
             )
 
     override fun currentMemory(ctx: FuncContext, func: Func) = func.popExpecting(memType).addInsns(
-        forceFnType<ByteBuffer.() -> Int>(ByteBuffer::limit).invokeVirtual(),
+        forceFnType<MemoryBuffer.() -> Int>(MemoryBuffer::limit).invokeVirtual(),
         Mem.PAGE_SIZE.const,
         InsnNode(Opcodes.IDIV)
     ).push(Int::class.ref)
@@ -86,10 +84,10 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
             val okLim = LabelNode()
             val node = MethodNode(
                 Opcodes.ACC_PRIVATE + Opcodes.ACC_STATIC + Opcodes.ACC_SYNTHETIC,
-                "\$\$growMemory", "(Ljava/nio/ByteBuffer;I)I", null, null
+                "\$\$growMemory", "(Lasmble/compile/jvm/MemoryBuffer;I)I", null, null
             ).addInsns(
                 VarInsnNode(Opcodes.ALOAD, 0), // [mem]
-                forceFnType<ByteBuffer.() -> Int>(ByteBuffer::limit).invokeVirtual(), // [lim]
+                forceFnType<MemoryBuffer.() -> Int>(MemoryBuffer::limit).invokeVirtual(), // [lim]
                 InsnNode(Opcodes.DUP), // [lim, lim]
                 VarInsnNode(Opcodes.ALOAD, 0), // [lim, lim, mem]
                 InsnNode(Opcodes.SWAP), // [lim, mem, lim]
@@ -102,7 +100,7 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
                 InsnNode(Opcodes.LADD), // [lim, mem, newlimL]
                 InsnNode(Opcodes.DUP2), // [lim, mem, newlimL, newlimL]
                 VarInsnNode(Opcodes.ALOAD, 0), // [lim, mem, newlimL, newlimL, mem]
-                ByteBuffer::capacity.invokeVirtual(), // [lim, mem, newlimL, newlimL, cap]
+                MemoryBuffer::capacity.invokeVirtual(), // [lim, mem, newlimL, newlimL, cap]
                 InsnNode(Opcodes.I2L), // [lim, mem, newlimL, newlimL, capL]
                 InsnNode(Opcodes.LCMP), // [lim, mem, newlimL, cmpres]
                 JumpInsnNode(Opcodes.IFLE, okLim), // [lim, mem, newlimL]
@@ -111,7 +109,7 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
                 InsnNode(Opcodes.IRETURN),
                 okLim, // [lim, mem, newlimL]
                 InsnNode(Opcodes.L2I), // [lim, mem, newlim]
-                forceFnType<ByteBuffer.(Int) -> Buffer>(ByteBuffer::limit).invokeVirtual(), // [lim, mem]
+                forceFnType<MemoryBuffer.(Int) -> MemoryBuffer>(MemoryBuffer::limit).invokeVirtual(), // [lim, mem]
                 InsnNode(Opcodes.POP), // [lim]
                 Mem.PAGE_SIZE.const, // [lim, pagesize]
                 InsnNode(Opcodes.IDIV), // [limpages]
@@ -125,7 +123,7 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
         // Ug, some tests expect this to be a runtime failure so we feature flagged it
         if (ctx.cls.eagerFailLargeMemOffset)
             require(insn.offset <= Int.MAX_VALUE, { "Offsets > ${Int.MAX_VALUE} unsupported" }).let { this }
-        fun Func.load(fn: ByteBuffer.(Int) -> Any, retClass: KClass<*>) =
+        fun Func.load(fn: MemoryBuffer.(Int) -> Any, retClass: KClass<*>) =
             this.popExpecting(Int::class.ref).let { func ->
                 // No offset means we'll access it directly
                 (if (insn.offset == 0L) func else {
@@ -141,9 +139,9 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
                     }
                 }).popExpecting(memType).addInsns((fn as KFunction<*>).invokeVirtual())
             }.push(retClass.ref)
-        fun Func.loadI32(fn: ByteBuffer.(Int) -> Any) =
+        fun Func.loadI32(fn: MemoryBuffer.(Int) -> Any) =
             this.load(fn, Int::class)
-        fun Func.loadI64(fn: ByteBuffer.(Int) -> Any) =
+        fun Func.loadI64(fn: MemoryBuffer.(Int) -> Any) =
             this.load(fn, Long::class)
         /* Ug: https://youtrack.jetbrains.com/issue/KT-17064
         fun Func.toUnsigned(fn: KFunction<*>) =
@@ -163,33 +161,33 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
         // Had to move this in here instead of as first expr because of https://youtrack.jetbrains.com/issue/KT-8689
         return when (insn) {
             is Node.Instr.I32Load ->
-                func.loadI32(ByteBuffer::getInt)
+                func.loadI32(MemoryBuffer::getInt)
             is Node.Instr.I64Load ->
-                func.loadI64(ByteBuffer::getLong)
+                func.loadI64(MemoryBuffer::getLong)
             is Node.Instr.F32Load ->
-                func.load(ByteBuffer::getFloat, Float::class)
+                func.load(MemoryBuffer::getFloat, Float::class)
             is Node.Instr.F64Load ->
-                func.load(ByteBuffer::getDouble, Double::class)
+                func.load(MemoryBuffer::getDouble, Double::class)
             is Node.Instr.I32Load8S ->
-                func.loadI32(ByteBuffer::get)
+                func.loadI32(MemoryBuffer::get)
             is Node.Instr.I32Load8U ->
-                func.loadI32(ByteBuffer::get).toUnsigned32(java.lang.Byte::class, "toUnsignedInt", Byte::class)
+                func.loadI32(MemoryBuffer::get).toUnsigned32(java.lang.Byte::class, "toUnsignedInt", Byte::class)
             is Node.Instr.I32Load16S ->
-                func.loadI32(ByteBuffer::getShort)
+                func.loadI32(MemoryBuffer::getShort)
             is Node.Instr.I32Load16U ->
-                func.loadI32(ByteBuffer::getShort).toUnsigned32(java.lang.Short::class, "toUnsignedInt", Short::class)
+                func.loadI32(MemoryBuffer::getShort).toUnsigned32(java.lang.Short::class, "toUnsignedInt", Short::class)
             is Node.Instr.I64Load8S ->
-                func.loadI32(ByteBuffer::get).i32ToI64()
+                func.loadI32(MemoryBuffer::get).i32ToI64()
             is Node.Instr.I64Load8U ->
-                func.loadI32(ByteBuffer::get).toUnsigned64(java.lang.Byte::class, "toUnsignedLong", Byte::class)
+                func.loadI32(MemoryBuffer::get).toUnsigned64(java.lang.Byte::class, "toUnsignedLong", Byte::class)
             is Node.Instr.I64Load16S ->
-                func.loadI32(ByteBuffer::getShort).i32ToI64()
+                func.loadI32(MemoryBuffer::getShort).i32ToI64()
             is Node.Instr.I64Load16U ->
-                func.loadI32(ByteBuffer::getShort).toUnsigned64(java.lang.Short::class, "toUnsignedLong", Short::class)
+                func.loadI32(MemoryBuffer::getShort).toUnsigned64(java.lang.Short::class, "toUnsignedLong", Short::class)
             is Node.Instr.I64Load32S ->
-                func.loadI32(ByteBuffer::getInt).i32ToI64()
+                func.loadI32(MemoryBuffer::getInt).i32ToI64()
             is Node.Instr.I64Load32U ->
-                func.loadI32(ByteBuffer::getInt).toUnsigned64(java.lang.Integer::class, "toUnsignedLong", Int::class)
+                func.loadI32(MemoryBuffer::getInt).toUnsigned64(java.lang.Integer::class, "toUnsignedLong", Int::class)
             else -> throw IllegalArgumentException("Unknown load op $insn")
         }
     }
@@ -224,12 +222,12 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
                         popExpecting(Int::class.ref).
                         popExpecting(memType).
                         addInsns(fn).
-                        push(ByteBuffer::class.ref)
+                        push(memType)
                 }
             // Ug, I hate these as strings but can't introspect Kotlin overloads
             fun bufStoreFunc(name: String, valType: KClass<*>) =
-                MethodInsnNode(Opcodes.INVOKEVIRTUAL, ByteBuffer::class.ref.asmName, name,
+                MethodInsnNode(Opcodes.INVOKEVIRTUAL, memType.asmName, name,
-                    ByteBuffer::class.ref.asMethodRetDesc(Int::class.ref, valType.ref), false)
+                    memType.asMethodRetDesc(Int::class.ref, valType.ref), false)
             fun Func.changeI64ToI32() =
                 this.popExpecting(Long::class.ref).push(Int::class.ref)
             when (insn) {

compiler/src/main/kotlin/asmble/compile/jvm/MemoryByteBuffer.kt

@@ -0,0 +1,122 @@
+package asmble.compile.jvm
+
+import java.nio.ByteBuffer
+import java.nio.ByteOrder
+
+/**
+ * The default implementation of MemoryBuffer that based on java.nio.ByteBuffer
+ */
+open class MemoryByteBuffer(val bb: ByteBuffer) : MemoryBuffer() {
+    override fun put(arr: ByteArray): MemoryBuffer {
+        bb.put(arr)
+        return this
+    }
+
+    override fun clear(): MemoryBuffer {
+        bb.clear()
+        return this
+    }
+
+    override fun get(arr: ByteArray): MemoryBuffer {
+        bb.get(arr)
+        return this
+    }
+
+    override fun putLong(index: Int, n: Long): MemoryBuffer {
+        bb.putLong(index, n)
+        return this
+    }
+
+    override fun putDouble(index: Int, n: Double): MemoryBuffer {
+        bb.putDouble(index, n)
+        return this
+    }
+
+    override fun putShort(index: Int, n: Short): MemoryBuffer {
+        bb.putShort(index, n)
+        return this
+    }
+
+    override fun putFloat(index: Int, n: Float): MemoryBuffer {
+        bb.putFloat(index, n)
+        return this
+    }
+
+    override fun put(index: Int, b: Byte): MemoryBuffer {
+        bb.put(index, b)
+        return this
+    }
+
+    override fun putInt(index: Int, n: Int): MemoryBuffer {
+        bb.putInt(index, n)
+        return this
+    }
+
+    override fun capacity(): Int {
+        return bb.capacity()
+    }
+
+    override fun limit(): Int {
+        return bb.limit()
+    }
+
+    override fun limit(newLimit: Int): MemoryBuffer {
+        bb.limit(newLimit)
+        return this
+    }
+
+    override fun position(newPosition: Int): MemoryBuffer {
+        bb.position(newPosition)
+        return this
+    }
+
+    override fun order(order: ByteOrder): MemoryBuffer {
+        bb.order(order)
+        return this
+    }
+
+    override fun duplicate(): MemoryBuffer {
+        return MemoryByteBuffer(bb.duplicate())
+    }
+
+    override fun put(arr: ByteArray, offset: Int, length: Int): MemoryBuffer {
+        bb.put(arr, offset, length)
+        return this
+    }
+
+    override fun getInt(index: Int): Int {
+        return bb.getInt(index)
+    }
+
+    override fun get(index: Int): Byte {
+        return bb.get(index)
+    }
+
+    override fun getLong(index: Int): Long {
+        return bb.getLong(index)
+    }
+
+    override fun getShort(index: Int): Short {
+        return bb.getShort(index)
+    }
+
+    override fun getFloat(index: Int): Float {
+        return bb.getFloat(index)
+    }
+
+    override fun getDouble(index: Int): Double {
+        return bb.getDouble(index)
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other)
+            return true
+        if (other !is MemoryByteBuffer)
+            return false
+        return bb == other.bb
+    }
+
+    override fun hashCode(): Int {
+        return bb.hashCode()
+    }
+}

compiler/src/main/kotlin/asmble/run/jvm/ScriptContext.kt

@@ -43,7 +43,8 @@ data class ScriptContext(
         ScriptContext.SimpleClassLoader(ScriptContext::class.java.classLoader, logger),
     val exceptionTranslator: ExceptionTranslator = ExceptionTranslator,
     val defaultMaxMemPages: Int = 1,
-    val includeBinaryInCompiledClass: Boolean = false
+    val includeBinaryInCompiledClass: Boolean = false,
+    val memoryBuilder: MemoryBufferBuilder? = null
 ) : Logger by logger {
 
     fun withHarnessRegistered(out: PrintWriter = PrintWriter(System.out, true)) =

compiler/src/main/kotlin/asmble/run/jvm/TestHarness.kt

@@ -3,6 +3,8 @@ package asmble.run.jvm
 import asmble.annotation.WasmExport
 import asmble.annotation.WasmExternalKind
 import asmble.compile.jvm.Mem
+import asmble.compile.jvm.MemoryBuffer
+import asmble.compile.jvm.MemoryByteBuffer
 import java.io.PrintWriter
 import java.lang.invoke.MethodHandle
 import java.nio.ByteBuffer
@@ -17,10 +19,10 @@ open class TestHarness(val out: PrintWriter) {
     val global_f32 = 666.6f
     val global_f64 = 666.6
     val table = arrayOfNulls<MethodHandle>(10)
-    val memory = ByteBuffer.
+    val memory = MemoryByteBuffer(ByteBuffer.
         allocateDirect(2 * Mem.PAGE_SIZE).
         order(ByteOrder.LITTLE_ENDIAN).
-        limit(Mem.PAGE_SIZE) as ByteBuffer
+        limit(Mem.PAGE_SIZE) as ByteBuffer) as MemoryBuffer
 
     // Note, we have all of these overloads because my import method
     // resolver is simple right now and only finds exact methods via

compiler/src/test/kotlin/asmble/compile/jvm/LargeDataTest.kt

@@ -5,7 +5,6 @@ import asmble.ast.Node
 import asmble.run.jvm.ScriptContext
 import asmble.util.get
 import org.junit.Test
-import java.nio.ByteBuffer
 import java.util.*
 import kotlin.test.assertEquals
 
@@ -35,9 +34,9 @@ class LargeDataTest : TestBase() {
         val cls = ScriptContext.SimpleClassLoader(javaClass.classLoader, logger).fromBuiltContext(ctx)
         // Instantiate it, get the memory out, and check it
         val field = cls.getDeclaredField("memory").apply { isAccessible = true }
-        val buf = field[cls.newInstance()] as ByteBuffer
+        val buf = field[cls.newInstance()] as MemoryByteBuffer
         // Grab all + 1 and check values
-        val bytesActual = ByteArray(70001).also { buf.get(0, it) }
+        val bytesActual = ByteArray(70001).also { buf.bb.get(0, it) }
         bytesActual.forEachIndexed { index, byte ->
             assertEquals(if (index == 70000) 0.toByte() else bytesExpected[index], byte)
         }

compiler/src/test/kotlin/asmble/run/jvm/LargeFuncTest.kt

@@ -4,10 +4,10 @@ import asmble.TestBase
 import asmble.ast.Node
 import asmble.compile.jvm.AstToAsm
 import asmble.compile.jvm.ClsContext
+import asmble.compile.jvm.MemoryByteBuffer
 import org.junit.Assert
 import org.junit.Test
 import org.objectweb.asm.MethodTooLargeException
-import java.nio.ByteBuffer
 import java.util.*
 import kotlin.test.assertEquals
 
@@ -62,7 +62,7 @@ class LargeFuncTest : TestBase() {
         // Run someFunc
         cls.getMethod("someFunc").invoke(inst)
         // Get the memory out
-        val mem = cls.getMethod("getMemory").invoke(inst) as ByteBuffer
+        val mem = cls.getMethod("getMemory").invoke(inst) as MemoryByteBuffer
         // Read out the mem values
         (0 until numInsnChunks).forEach { assertEquals(it * (it - 1), mem.getInt(it * 4)) }
     }

compiler/src/test/kotlin/asmble/run/jvm/TestRunner.kt

@@ -3,6 +3,8 @@ package asmble.run.jvm
 import asmble.BaseTestUnit
 import asmble.TestBase
 import asmble.annotation.WasmModule
+import asmble.compile.jvm.MemoryBufferBuilder
+import asmble.compile.jvm.MemoryByteBuffer
 import asmble.io.AstToBinary
 import asmble.io.AstToSExpr
 import asmble.io.ByteWriter
@@ -12,6 +14,7 @@ import org.junit.Test
 import java.io.ByteArrayOutputStream
 import java.io.OutputStreamWriter
 import java.io.PrintWriter
+import java.nio.ByteBuffer
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
 
@@ -40,7 +43,10 @@ abstract class TestRunner<out T : BaseTestUnit>(val unit: T) : TestBase() {
             adjustContext = { it.copy(eagerFailLargeMemOffset = false) },
             defaultMaxMemPages = unit.defaultMaxMemPages,
             // Include the binary data so we can check it later
-            includeBinaryInCompiledClass = true
+            includeBinaryInCompiledClass = true,
+            memoryBuilder = MemoryBufferBuilder { it ->
+                MemoryByteBuffer(ByteBuffer.allocateDirect(it))
+            }
         ).withHarnessRegistered(PrintWriter(OutputStreamWriter(out, Charsets.UTF_8), true))
 
         // This will fail assertions as necessary