Large table jump support

src/main/kotlin/asmble/compile/jvm/ClsContext.kt

@@ -8,6 +8,7 @@ import org.objectweb.asm.Type
 import org.objectweb.asm.tree.ClassNode
 import org.objectweb.asm.tree.MethodInsnNode
 import org.objectweb.asm.tree.MethodNode
+import java.util.*
 
 data class ClsContext(
     val packageName: String,
@@ -24,7 +25,8 @@ data class ClsContext(
     val eagerFailLargeMemOffset: Boolean = true,
     val preventMemIndexOverflow: Boolean = false,
     val accurateNanBits: Boolean = true,
-    val checkSignedDivIntegerOverflow: Boolean = true
+    val checkSignedDivIntegerOverflow: Boolean = true,
+    val jumpTableChunkSize: Int = 5000
 ) : Logger by logger {
     val importFuncs: List<Node.Import> by lazy { mod.imports.filter { it.kind is Node.Import.Kind.Func } }
     val importGlobals: List<Node.Import> by lazy { mod.imports.filter { it.kind is Node.Import.Kind.Global } }
@@ -92,4 +94,14 @@ data class ClsContext(
     val divAssertL get() = syntheticFunc("assertLDiv", SyntheticFuncBuilder::buildLDivAssertion)
 
     val indirectBootstrap get() = syntheticFunc("indirectBootstrap", SyntheticFuncBuilder::buildIndirectBootstrap)
+
+    // Builds a method that takes an int and returns a depth int
+    fun largeTableJumpCall(table: Node.Instr.BrTable): MethodInsnNode {
+        val namePrefix = "largeTable" + UUID.randomUUID().toString().replace("-", "")
+        val methods = syntheticFuncBuilder.buildLargeTableJumps(this, namePrefix, table)
+        cls.methods.addAll(methods)
+        return methods.first().let { method ->
+            MethodInsnNode(Opcodes.INVOKESTATIC, thisRef.asmName, method.name, method.desc, false)
+        }
+    }
 }

src/main/kotlin/asmble/compile/jvm/FuncBuilder.kt

@@ -515,14 +515,53 @@ open class FuncBuilder {
                     throw CompileErr.TableTargetMismatch(defaultBlock.labelTypes, targetBlock.labelTypes)
                 fn to (blocks + targetBlock)
             }.let { (fn, targetBlocks) ->
+                // If it's large, we need to handle it differently
+                if (insn.targetTable.size > ctx.cls.jumpTableChunkSize) {
+                    applyLargeBrTable(ctx, fn, insn, defaultBlock, targetBlocks)
+                } else {
                     // In some cases, the target labels is empty. We need to make 0 goto
                     // the default as well.
                     val targetLabelsArr =
                         if (targetBlocks.isNotEmpty()) targetBlocks.map(Func.Block::requiredLabel).toTypedArray()
                         else arrayOf(defaultBlock.requiredLabel)
-                fn.popExpecting(Int::class.ref).addInsns(TableSwitchInsnNode(0, targetLabelsArr.size - 1,
+                    fn.popExpecting(Int::class.ref).
-                    defaultBlock.requiredLabel, *targetLabelsArr))
+                        addInsns(TableSwitchInsnNode(0, targetLabelsArr.size - 1,
-            }.popExpectingMulti(defaultBlock.labelTypes).markUnreachable()
+                            defaultBlock.requiredLabel, *targetLabelsArr)).
+                        popExpectingMulti(defaultBlock.labelTypes).
+                        markUnreachable()
+                }
+            }
+        }
+
+    fun applyLargeBrTable(
+        ctx: FuncContext,
+        fn: Func,
+        insn: Node.Instr.BrTable,
+        defaultBlock: Func.Block,
+        targetBlocks: List<Func.Block>
+    ): Func {
+        // We build a method call to get our set of depths, then we do a table switch
+        // on the depths. There may be holes in the depths, which we'll fill in w/ the
+        // default label. And we'll make the default label unreachable.
+        val depthToBlock = mutableListOf<LabelNode?>()
+        fun addBlock(depth: Int, block: Func.Block) {
+            if (depthToBlock.getOrNull(depth) == null) {
+                for (i in depthToBlock.size..depth) depthToBlock.add(null)
+                depthToBlock[depth] = block.requiredLabel
+            }
+        }
+        insn.targetTable.forEachIndexed { index, targetDepth -> addBlock(targetDepth, targetBlocks[index]) }
+        addBlock(insn.default, defaultBlock)
+
+        val unreachableLabel = LabelNode()
+        return fn.popExpecting(Int::class.ref).
+            addInsns(
+                ctx.cls.largeTableJumpCall(insn),
+                TableSwitchInsnNode(0, depthToBlock.size - 1, unreachableLabel,
+                    *depthToBlock.map { it ?: unreachableLabel }.toTypedArray()),
+                unreachableLabel
+            ).addInsns(UnsupportedOperationException::class.athrow("Unreachable")).
+            markUnreachable()
     }
 
     fun needsToPopBeforeJumping(ctx: FuncContext, fn: Func, block: Func.Block): Boolean {

src/main/kotlin/asmble/compile/jvm/InsnReworker.kt

@@ -181,7 +181,8 @@ open class InsnReworker {
         is Node.Instr.Unreachable, is Node.Instr.Nop, is Node.Instr.Block,
         is Node.Instr.Loop, is Node.Instr.If, is Node.Instr.Else,
         is Node.Instr.End, is Node.Instr.Br, is Node.Instr.BrIf,
-        is Node.Instr.BrTable, is Node.Instr.Return -> NOP
+        is Node.Instr.Return -> NOP
+        is Node.Instr.BrTable -> POP_PARAM
         is Node.Instr.Call -> ctx.funcTypeAtIndex(insn.index).let {
             // All calls pop params and any return is a push
             (POP_PARAM * it.params.size) + (if (it.ret == null) NOP else PUSH_RESULT)

src/main/kotlin/asmble/compile/jvm/SyntheticFuncBuilder.kt

@@ -1,5 +1,6 @@
 package asmble.compile.jvm
 
+import asmble.ast.Node
 import org.objectweb.asm.ClassReader
 import org.objectweb.asm.Opcodes
 import org.objectweb.asm.tree.*
@@ -23,6 +24,63 @@ open class SyntheticFuncBuilder {
         ).addInsns(*helperMeth.instructions.toArray())
     }
 
+    // Guaranteed that the first method result can be called to get the proper index.
+    // Caller needs to make sure namePrefix is unique.
+    fun buildLargeTableJumps(ctx: ClsContext, namePrefix: String, table: Node.Instr.BrTable): List<MethodNode> {
+        // Sadly really large table jumps need to be broken up into multiple methods
+        // because Java has method size limits. What we are going to do here is make
+        // a method that takes an int, then a table switch node that returns the depth.
+        // The default will chain to another method if there are more to handle.
+
+        // Build a bunch of chunk views...first of each is start, second is sub list
+        val chunks = (0 until Math.ceil(table.targetTable.size / ctx.jumpTableChunkSize.toDouble()).toInt()).
+            fold(emptyList<Pair<Int, List<Int>>>()) { chunks, chunkNum ->
+                val start = chunkNum * ctx.jumpTableChunkSize
+                chunks.plusElement(start to table.targetTable.subList(start,
+                    Math.min(table.targetTable.size, (chunkNum + 1) * ctx.jumpTableChunkSize)))
+            }
+        // Go over the chunks, backwards, building the jump methods, then flip em back
+        return chunks.asReversed().fold(emptyList<MethodNode>()) { methods, (start, chunk) ->
+            val defaultLabel = LabelNode()
+            val method = largeTableJumpMethod(ctx, namePrefix, start, chunk, defaultLabel)
+            // If we are the last chunk, default is what table default is
+            methods + if (methods.isEmpty()) method.addInsns(
+                defaultLabel, table.default.const, InsnNode(Opcodes.IRETURN)
+            ) else method.addInsns(
+                // Otherwise, the default label just calls the prev
+                defaultLabel,
+                VarInsnNode(Opcodes.ILOAD, 0),
+                methods.last().let { other ->
+                    MethodInsnNode(Opcodes.INVOKESTATIC, ctx.thisRef.asmName, other.name, other.desc, false)
+                },
+                InsnNode(Opcodes.IRETURN)
+            )
+        }.reversed()
+    }
+
+    private fun largeTableJumpMethod(
+        ctx: ClsContext,
+        namePrefix: String,
+        startIndex: Int,
+        targets: List<Int>,
+        defaultLabel: LabelNode
+    ): MethodNode {
+        val labelsByTargets = mutableMapOf<Int, LabelNode>()
+        return MethodNode(
+            Opcodes.ACC_PRIVATE + Opcodes.ACC_STATIC + Opcodes.ACC_SYNTHETIC,
+            "${namePrefix}_${startIndex}_until_${startIndex + targets.size}", "(I)I", null, null
+        ).addInsns(
+            VarInsnNode(Opcodes.ILOAD, 0),
+            TableSwitchInsnNode(startIndex, (startIndex + targets.size) - 1, defaultLabel,
+                *targets.map { labelsByTargets.getOrPut(it) { LabelNode() } }.toTypedArray()
+            )
+        ).also { method ->
+            labelsByTargets.forEach { (target, label) ->
+                method.addInsns(label, target.const, InsnNode(Opcodes.IRETURN))
+            }
+        }
+    }
+
     fun buildIDivAssertion(ctx: ClsContext, name: String) =
         LabelNode().let { safeLabel ->
             LabelNode().let { overflowLabel ->

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

@@ -264,7 +264,7 @@ data class ScriptContext(
         } catch (e: NoSuchMethodException) {
             // Try any method w/ the proper annotation
             module.cls.methods.forEach { method ->
-                if (method.getAnnotation(WasmName::class.java)?.name == import.field) {
+                if (method.getAnnotation(WasmName::class.java)?.value == import.field) {
                     val handle = MethodHandles.lookup().unreflect(method).bindTo(module.instance)
                     if (handle.type() == methodType) return handle
                 }

src/main/kotlin/asmble/run/jvm/annotation/WasmName.kt

@@ -1,4 +1,4 @@
 package asmble.run.jvm.annotation
 
 @Target(AnnotationTarget.PROPERTY_GETTER, AnnotationTarget.FUNCTION)
-annotation class WasmName(val name: String)
+annotation class WasmName(val value: String)