Complete main func splitting impl with test

compiler/src/main/kotlin/asmble/ast/opt/SplitLargeFunc.kt

@@ -7,8 +7,8 @@ import asmble.ast.Stack
 // the most instructions off into its own function.
 open class SplitLargeFunc(
     val minSetLength: Int = 5,
-    val maxSetLength: Int = 20,
-    val maxParamCount: Int = 10
+    val maxSetLength: Int = 40,
+    val maxParamCount: Int = 30
 ) {
 
     // Null if no replacement. Second value is number of instructions saved. fnIndex must map to actual func,
@@ -58,8 +58,9 @@ open class SplitLargeFunc(
         // only have a certain set of insns that can be broken off. It can also only change the stack by 0 or 1
         // value while never dipping below the starting stack. We also store the index they started at.
         var insnSets = emptyList<InsnSet>()
-        fn.instructions.foldIndexed(null as List<Node.Instr>?) { index, lastInsns, insn ->
-            if (!insn.canBeMoved) null else (lastInsns ?: emptyList()).plus(insn).also { fullNewInsnSet ->
+        // Pair in fold keyed by insn index
+        fn.instructions.foldIndexed(null as List<Pair<Int, Node.Instr>>?) { index, lastInsns, insn ->
+            if (!insn.canBeMoved) null else (lastInsns ?: emptyList()).plus(index to insn).also { fullNewInsnSet ->
                 // Get all final instructions between min and max size and with allowed param count (i.e. const count)
                 val trailingInsnSet = fullNewInsnSet.takeLast(maxSetLength)
 
@@ -67,10 +68,14 @@ open class SplitLargeFunc(
                 insnSets += (minSetLength..maxSetLength).
                     asSequence().
                     flatMap { trailingInsnSet.asSequence().windowed(it) }.
-                    filter { it.count { it is Node.Instr.Args.Const<*> } <= maxParamCount }.
-                    mapNotNull { newInsnSet ->
+                    filter { it.count { it.second is Node.Instr.Args.Const<*> } <= maxParamCount }.
+                    mapNotNull { newIndexedInsnSet ->
                         // Before adding, make sure it qualifies with the stack
-                        InsnSet(index + 1 - newInsnSet.size, newInsnSet, null).withStackValueIfValid(stack)
+                        InsnSet(
+                            startIndex = newIndexedInsnSet.first().first,
+                            insns = newIndexedInsnSet.map { it.second },
+                            valueAddedToStack = null
+                        ).withStackValueIfValid(stack)
                     }
             }
         }
@@ -110,7 +115,7 @@ open class SplitLargeFunc(
 
         // First, make sure the stack after the last insn is the same as the first or the same + 1 val
         val startingStack = stack.insnApplies[startIndex].stackAtBeginning!!
-        val endingStack = stack.insnApplies.getOrNull(startIndex + insns.size + 1)?.stackAtBeginning ?: stack.current!!
+        val endingStack = stack.insnApplies.getOrNull(startIndex + insns.size)?.stackAtBeginning ?: stack.current!!
         if (endingStack.size != startingStack.size && endingStack.size != startingStack.size + 1) return null
         if (endingStack.take(startingStack.size) != startingStack) return null
 
@@ -159,7 +164,8 @@ open class SplitLargeFunc(
         val range: IntRange,
         val preCallConsts: List<Node.Instr>
     ) {
-        val insnsSaved get() = range.last - range.first + 1 - preCallConsts.size
+        // Subtract one because there is a call after this
+        val insnsSaved get() = (range.last + 1) - range.first - 1 - preCallConsts.size
         fun overlaps(o: Replacement) = range.contains(o.range.first) || range.contains(o.range.last) ||
             o.range.contains(range.first) || o.range.contains(range.last)
     }
@@ -170,7 +176,7 @@ open class SplitLargeFunc(
         val replacements: List<Replacement>
     ) {
         // Replacement pieces saved (with one added for the invocation) less new func instructions
-        val insnsSaved get() = replacements.sumBy { 1 + it.insnsSaved } - newFunc.instructions.size
+        val insnsSaved get() = replacements.sumBy { it.insnsSaved } - newFunc.instructions.size
     }
 
     companion object : SplitLargeFunc()

compiler/src/test/kotlin/asmble/ast/opt/SplitLargeFuncTest.kt

@@ -23,7 +23,7 @@ class SplitLargeFuncTest : TestBase() {
                 funcs = listOf(Node.Func(
                     type = Node.Type.Func(params = emptyList(), ret = null),
                     locals = emptyList(),
-                    instructions = (0 until 500).flatMap {
+                    instructions = (0 until 501).flatMap {
                         listOf<Node.Instr>(
                             Node.Instr.I32Const(it * 4),
                             // Let's to i * (i = 1)
@@ -47,17 +47,43 @@ class SplitLargeFuncTest : TestBase() {
         )
         // Compile it
         AstToAsm.fromModule(ctx)
-        val unsplitCls = ScriptContext.SimpleClassLoader(javaClass.classLoader, logger).fromBuiltContext(ctx)
-        val unsplitInst = unsplitCls.newInstance()
+        val cls = ScriptContext.SimpleClassLoader(javaClass.classLoader, logger).fromBuiltContext(ctx)
+        val inst = cls.newInstance()
         // Run someFunc
-        unsplitCls.getMethod("someFunc").invoke(unsplitInst)
+        cls.getMethod("someFunc").invoke(inst)
         // Get the memory out
-        val unsplitMem = unsplitCls.getMethod("getMemory").invoke(unsplitInst) as ByteBuffer
+        val mem = cls.getMethod("getMemory").invoke(inst) as ByteBuffer
         // Read out the mem values
-        (0 until 500).forEach { assertEquals(it * (it - 1), unsplitMem.getInt(it * 4)) }
+        (0 until 501).forEach { assertEquals(it * (it - 1), mem.getInt(it * 4)) }
+
         // Now split it
         val (splitMod, insnsSaved) = SplitLargeFunc.apply(ctx.mod, 0) ?: error("Nothing could be split")
-        println("SAVED! $insnsSaved NEW FUNC - " + splitMod.funcs.last())
-        // TODO: the rest
+        // Count insns and confirm it is as expected
+        val origInsnCount = ctx.mod.funcs.sumBy { it.instructions.size }
+        val newInsnCount = splitMod.funcs.sumBy { it.instructions.size }
+        assertEquals(origInsnCount - newInsnCount, insnsSaved)
+        // Compile it
+        val splitCtx = ClsContext(
+            packageName = "test",
+            className = "Temp" + UUID.randomUUID().toString().replace("-", ""),
+            logger = logger,
+            mod = splitMod
+        )
+        AstToAsm.fromModule(splitCtx)
+        val splitCls = ScriptContext.SimpleClassLoader(javaClass.classLoader, logger).fromBuiltContext(splitCtx)
+        val splitInst = splitCls.newInstance()
+        // Run someFunc
+        splitCls.getMethod("someFunc").invoke(splitInst)
+        // Get the memory out and compare it
+        val splitMem = splitCls.getMethod("getMemory").invoke(splitInst) as ByteBuffer
+        assertEquals(mem, splitMem)
+        // Dump some info
+        logger.debug {
+            val orig = ctx.mod.funcs.first()
+            val (new, split) = splitMod.funcs.let { it.first() to it.last() }
+            "Split complete, from single func with ${orig.instructions.size} insns to func " +
+                "with ${new.instructions.size} insns + delegated func " +
+                "with ${split.instructions.size} insns and ${split.type.params.size} params"
+        }
     }
 }