Better handling of dead code

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

@@ -42,7 +42,7 @@ open class ByteBufferMem(val direct: Boolean = true) : Mem {
         addInsns(bytes.withIndex().flatMap { (index, byte) ->
             listOf(InsnNode(Opcodes.DUP), index.const, byte.toInt().const, InsnNode(Opcodes.BASTORE))
         }).
-        apply(buildOffset).popExpecting(Int::class.ref).
+        let(buildOffset).popExpecting(Int::class.ref).
         // BOO! https://discuss.kotlinlang.org/t/overload-resolution-ambiguity-function-reference-requiring-local-var/2425
         addInsns(
             bytes.size.const,

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

@@ -3,7 +3,6 @@ package asmble.compile.jvm
 import asmble.ast.Node
 import org.objectweb.asm.Opcodes
 import org.objectweb.asm.tree.*
-import java.util.*
 
 data class Func(
     val name: String,
@@ -18,14 +17,19 @@ data class Func(
 ) {
 
     val desc: String get() = ret.asMethodRetDesc(*params.toTypedArray())
-    val isLastUnconditionalJump get() = insns.lastOrNull()?.isUnconditionalJump ?: false
-    val isLastTerminating get() = insns.lastOrNull()?.isTerminating ?: false
+    val isCurrentBlockDead get() = blockStack.lastOrNull()?.let { block ->
+        // It's dead if it's marked unconditional or it's an unconditional
+        // if/else and we are in that if/else area
+        block.unconditionalBranch ||
+            (block.unconditionalBranchInIf && !block.hasElse) ||
+            (block.unconditionalBranchInElse && block.hasElse)
+    } ?: false
 
-    fun addInsns(insns: List<AbstractInsnNode>) = copy(insns = this.insns + insns)
+    fun addInsns(insns: List<AbstractInsnNode>) =
+        if (isCurrentBlockDead) this else copy(insns = this.insns + insns)
 
-    fun addInsns(vararg insns: AbstractInsnNode) = copy(insns = this.insns + insns)
-
-    fun apply(fn: (Func) -> Func) = fn(this)
+    fun addInsns(vararg insns: AbstractInsnNode) =
+        if (isCurrentBlockDead) this else copy(insns = this.insns + insns)
 
     fun push(vararg types: TypeRef) = copy(stack = stack + types)
 
@@ -33,10 +37,8 @@ data class Func(
 
     fun popExpectingMulti(vararg types: TypeRef) = types.reversed().fold(this, Func::popExpecting)
 
-    fun popExpecting(type: TypeRef) = popExpectingAny(type)
-
-    fun popExpectingAny(vararg types: TypeRef): Func {
-        peekExpectingAny(types = *types)
+    fun popExpecting(type: TypeRef): Func {
+        assertTopOfStack(type)
         return pop().first
     }
 
@@ -47,18 +49,20 @@ data class Func(
     }
 
     fun pop(currBlock: Block? = blockStack.lastOrNull()): Pair<Func, TypeRef> {
-        if (isStackEmptyForBlock(currBlock)) throw CompileErr.StackMismatch(emptyArray(), null)
+        if (isStackEmptyForBlock(currBlock)) {
+            // Just fake it if dead
+            if (isCurrentBlockDead) return this to Int::class.ref
+            throw CompileErr.StackMismatch(emptyArray(), null)
+        }
         return copy(stack = stack.dropLast(1)) to stack.last()
     }
 
-    fun peekExpecting(type: TypeRef, currBlock: Block? = blockStack.lastOrNull()) =
-        peekExpectingAny(currBlock, type)
-
-    fun peekExpectingAny(currBlock: Block? = blockStack.lastOrNull(), vararg types: TypeRef): TypeRef {
-        if (isStackEmptyForBlock(currBlock)) throw CompileErr.StackMismatch(types, null)
-        val hasExpected = stack.lastOrNull()?.let(types::contains) ?: false
-        if (!hasExpected) throw CompileErr.StackMismatch(types, stack.lastOrNull())
-        return stack.last()
+    fun assertTopOfStack(type: TypeRef, currBlock: Block? = blockStack.lastOrNull()): Unit {
+        // If it's dead, we just go with it
+        if (!isCurrentBlockDead) {
+            if (isStackEmptyForBlock(currBlock)) throw CompileErr.StackMismatch(arrayOf(type), null)
+            if (stack.lastOrNull() != type) throw CompileErr.StackMismatch(arrayOf(type), stack.lastOrNull())
+        }
     }
 
     fun toMethodNode(): MethodNode {
@@ -71,25 +75,26 @@ data class Func(
 
     fun withoutAffectingStack(fn: (Func) -> Func) = fn(this).copy(stack = stack)
 
-    fun stackSwap(currBlock: Block? = blockStack.lastOrNull()) = pop(currBlock).let { (fn, refLast) ->
-        fn.pop(currBlock).let { (fn, refFirst) ->
-            (if (refFirst.stackSize == 2) {
-                if (refLast.stackSize == 2)
-                    // If they are both 2, dup2_x2 + pop2
-                    fn.addInsns(InsnNode(Opcodes.DUP2_X2), InsnNode(Opcodes.POP2))
-                else
-                    // If only the first one is, dup_x2 + pop
-                    fn.addInsns(InsnNode(Opcodes.DUP_X2), InsnNode(Opcodes.POP))
-            } else {
-                if (refLast.stackSize == 2)
-                   // If the first is not 2 but the last is, dup_2x1, pop2
-                    fn.addInsns(InsnNode(Opcodes.DUP2_X1), InsnNode(Opcodes.POP2))
-                else
-                    // If neither are 2, just swap
-                    fn.addInsns(InsnNode(Opcodes.SWAP))
-            }).push(refLast).push(refFirst)
+    fun stackSwap(currBlock: Block? = blockStack.lastOrNull()) =
+        if (isCurrentBlockDead) this else pop(currBlock).let { (fn, refLast) ->
+            fn.pop(currBlock).let { (fn, refFirst) ->
+                (if (refFirst.stackSize == 2) {
+                    if (refLast.stackSize == 2)
+                        // If they are both 2, dup2_x2 + pop2
+                        fn.addInsns(InsnNode(Opcodes.DUP2_X2), InsnNode(Opcodes.POP2))
+                    else
+                        // If only the first one is, dup_x2 + pop
+                        fn.addInsns(InsnNode(Opcodes.DUP_X2), InsnNode(Opcodes.POP))
+                } else {
+                    if (refLast.stackSize == 2)
+                       // If the first is not 2 but the last is, dup_2x1, pop2
+                        fn.addInsns(InsnNode(Opcodes.DUP2_X1), InsnNode(Opcodes.POP2))
+                    else
+                        // If neither are 2, just swap
+                        fn.addInsns(InsnNode(Opcodes.SWAP))
+                }).push(refLast).push(refFirst)
+            }
         }
-    }
 
     fun pushBlock(insn: Node.Instr) = copy(blockStack = blockStack + Block(insn, insns.size, stack))
 
@@ -123,6 +128,8 @@ data class Func(
         open val requiredEndStack: List<TypeRef>? get() = null
         open val hasElse: Boolean get() = false
         open val unconditionalBranch: Boolean get() = false
+        open val unconditionalBranchInIf: Boolean get() = false
+        open val unconditionalBranchInElse: Boolean get() = false
         // First val is the insn, second is the type
         open val blockExitVals: List<Pair<Node.Instr, TypeRef?>> = emptyList()
         fun withLabel(label: LabelNode) = WithLabel(insn, startIndex, origStack, label)
@@ -138,6 +145,8 @@ data class Func(
             override var requiredEndStack: List<TypeRef>? = null
             override var hasElse = false
             override var unconditionalBranch = false
+            override var unconditionalBranchInIf = false
+            override var unconditionalBranchInElse = false
         }
     }
 }

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

@@ -87,7 +87,9 @@ open class FuncBuilder {
 
     fun applyNodeInsn(ctx: FuncContext, fn: Func, i: Node.Instr, index: Int) = when (i) {
         is Node.Instr.Unreachable ->
-            fn.addInsns(UnsupportedOperationException::class.athrow("Unreachable"))
+            fn.addInsns(UnsupportedOperationException::class.athrow("Unreachable")).let { fn ->
+                setAsUnconditionalBranch(ctx, fn)
+            }
         is Node.Instr.Nop ->
             fn.addInsns(InsnNode(Opcodes.NOP))
         is Node.Instr.Block ->
@@ -448,42 +450,53 @@ open class FuncBuilder {
                 java.lang.Double::class.invokeStatic("longBitsToDouble", Double::class, Long::class))
     }
 
-    fun setAsUnconditionalBranchUpToDepth(ctx: FuncContext, fn: Func, depth: Int) =
-        (0..depth).fold(fn) { fn, depth ->
-            fn.blockAtDepth(depth).let { (fn, block) -> block.unconditionalBranch = true; fn }
+    fun setAsUnconditionalBranch(ctx: FuncContext, fn: Func) =
+        fn.blockAtDepth(0).let { (fn, block) ->
+            // We stop at the first conditional branch that:
+            // * Has an else but doesn't have an initial unconditional branch
+            // * Doesn't have an else
+            if (block.insn is Node.Instr.If) {
+                if (block.hasElse) block.unconditionalBranchInElse = true
+                else block.unconditionalBranchInIf = true
+
+                if (block.hasElse && !block.unconditionalBranchInIf) return fn
+                if (!block.hasElse) return fn
+            }
+            block.unconditionalBranch = true
+            fn
         }
 
     fun applyBr(ctx: FuncContext, fn: Func, i: Node.Instr.Br) =
-        setAsUnconditionalBranchUpToDepth(ctx, fn, i.relativeDepth).let { fn ->
-            fn.blockAtDepth(i.relativeDepth).let { (fn, block) ->
-                // We have to pop all unnecessary values per the spec
-                val type = block.insnType?.typeRef
-                fun pop(fn: Func): Func {
-                    // Have to swap first if there is a type expected
-                    return (if (type != null) fn.stackSwap(block) else fn).let { fn ->
-                        fn.pop().let { (fn, poppedType) ->
-                            fn.addInsns(InsnNode(if (poppedType.stackSize == 2) Opcodes.POP2 else Opcodes.POP))
-                        }
+        fn.blockAtDepth(i.relativeDepth).let { (fn, block) ->
+            // We have to pop all unnecessary values per the spec
+            val type = block.insnType?.typeRef
+            fun pop(fn: Func): Func {
+                // Have to swap first if there is a type expected
+                // Note, we check stack size because dead code is allowed to do some crazy
+                // things.
+                return (if (type != null && fn.stack.size > 1) fn.stackSwap(block) else fn).let { fn ->
+                    fn.pop().let { (fn, poppedType) ->
+                        fn.addInsns(InsnNode(if (poppedType.stackSize == 2) Opcodes.POP2 else Opcodes.POP))
                     }
                 }
-                val expectedStackSize =
-                    if (block.insn is Node.Instr.Loop || type == null) block.origStack.size
-                    else block.origStack.size + 1
-                ctx.debug {
-                    "Unconditional branch on ${block.insn}, curr stack ${fn.stack}, " +
-                        " orig stack ${block.origStack}, expected stack size $expectedStackSize" }
-                val popCount = Math.max(0, fn.stack.size - expectedStackSize)
-                (0 until popCount).fold(fn) { fn, _ -> pop(fn) }.let { fn ->
-                    fn.addInsns(JumpInsnNode(Opcodes.GOTO, block.label)).let { fn ->
-                        // We only peek here, because we keep items on the stack for
-                        // dead code that uses it
-                        block.insnType?.typeRef?.let { typ ->
-                            // Loop breaks don't have to type check
-                            if (block.insn !is Node.Instr.Loop) fn.peekExpecting(typ, block)
-                            block.blockExitVals += i to typ
-                        }
-                        fn
+            }
+            val expectedStackSize =
+                if (type == null) block.origStack.size
+                else block.origStack.size + 1
+            ctx.debug {
+                "Unconditional branch on ${block.insn}, curr stack ${fn.stack}, " +
+                    " orig stack ${block.origStack}, expected stack size $expectedStackSize" }
+            val popCount = Math.max(0, fn.stack.size - expectedStackSize)
+            (0 until popCount).fold(fn) { fn, _ -> pop(fn) }.let { fn ->
+                fn.addInsns(JumpInsnNode(Opcodes.GOTO, block.label)).let { fn ->
+                    // We only peek here, because we keep items on the stack for
+                    // dead code that uses it
+                    block.insnType?.typeRef?.let { typ ->
+                        // Loop breaks don't have to type check
+                        if (block.insn !is Node.Instr.Loop) fn.assertTopOfStack(typ, block)
+                        block.blockExitVals += i to typ
                     }
+                    setAsUnconditionalBranch(ctx, fn)
                 }
             }
         }
@@ -496,7 +509,7 @@ open class FuncBuilder {
                 val toLabel = if (needsPopBeforeJump) LabelNode() else block.label
                 fn.addInsns(JumpInsnNode(Opcodes.IFNE, toLabel)).let { fn ->
                     block.insnType?.typeRef?.let {
-                        fn.peekExpecting(it)
+                        fn.assertTopOfStack(it)
                         block.blockExitVals += i to it
                     }
                     if (needsPopBeforeJump) buildPopBeforeJump(ctx, fn, block, toLabel)
@@ -508,11 +521,9 @@ open class FuncBuilder {
     // Can compile quite cleanly as a table switch on the JVM
     fun applyBrTable(ctx: FuncContext, fn: Func, insn: Node.Instr.BrTable) =
         fn.blockAtDepth(insn.default).let { (fn, defaultBlock) ->
-            defaultBlock.unconditionalBranch = true
             defaultBlock.insnType?.typeRef?.let { defaultBlock.blockExitVals += insn to it }
             insn.targetTable.fold(fn to emptyList<LabelNode>()) { (fn, labels), targetDepth ->
                 fn.blockAtDepth(targetDepth).let { (fn, targetBlock) ->
-                    targetBlock.unconditionalBranch = true
                     targetBlock.insnType?.typeRef?.let { targetBlock.blockExitVals += insn to it }
                     fn to (labels + targetBlock.label)
                 }
@@ -527,8 +538,9 @@ open class FuncBuilder {
             }.let { fn ->
                 // We only peek here, because we keep items on the stack for
                 // dead code that uses it
-                defaultBlock.insnType?.typeRef?.let { fn.peekExpecting(it) }
-                fn
+                defaultBlock.insnType?.typeRef?.let { fn.assertTopOfStack(it) }
+                // I think we're only going to mark ourselves dead for now
+                setAsUnconditionalBranch(ctx, fn)
             }
         }
 
@@ -560,9 +572,11 @@ open class FuncBuilder {
         return fn.addInsns(JumpInsnNode(Opcodes.GOTO, resumeLabel), tempLabel).withoutAffectingStack { fn ->
             (requiredStackCount until fn.stack.size).fold(fn) { fn, index ->
                 if (fn.stack.size == 1) {
-                    fn.addInsns(InsnNode(if (fn.stack.last().stackSize == 2) Opcodes.POP2 else Opcodes.POP)).pop(block).first
+                    fn.addInsns(InsnNode(if (fn.stack.last().stackSize == 2) Opcodes.POP2 else Opcodes.POP)).
+                        pop(block).first
                 } else fn.stackSwap(block).let { fn ->
-                    fn.addInsns(InsnNode(if (fn.stack.last().stackSize == 2) Opcodes.POP2 else Opcodes.POP)).pop(block).first
+                    fn.addInsns(InsnNode(if (fn.stack.last().stackSize == 2) Opcodes.POP2 else Opcodes.POP)).
+                        pop(block).first
                 }
             }
         }.addInsns(
@@ -577,7 +591,7 @@ open class FuncBuilder {
         val label = LabelNode()
         fn.peekIf().label = label
         ctx.debug { "Else block for ${block.insn}, orig stack ${block.origStack}" }
-        assertValidBlockEnd(ctx, fn, block)
+        if (!block.unconditionalBranchInIf) assertValidBlockEnd(ctx, fn, block)
         block.hasElse = true
         fn.addInsns(JumpInsnNode(Opcodes.GOTO, block.label), label).copy(stack = block.origStack)
     }
@@ -622,6 +636,8 @@ open class FuncBuilder {
     }
 
     fun assertValidBlockEnd(ctx: FuncContext, fn: Func, block: Func.Block) {
+        // If it's dead, who cares
+        if (block.unconditionalBranch || (block.unconditionalBranchInIf && !block.hasElse)) return
         // Go over each exit and make sure it did the right thing
         block.blockExitVals.forEach {
             require(it.second == block.insnType?.typeRef) { "Block exit val was $it, expected ${block.insnType}" }
@@ -1193,10 +1209,7 @@ open class FuncBuilder {
             fn.popExpecting(Double::class.ref).addInsns(InsnNode(Opcodes.DRETURN))
     }.let { fn ->
         if (fn.stack.isNotEmpty()) throw CompileErr.UnusedStackOnReturn(fn.stack)
-        fn.blockAtDepth(0).let { (fn, block) ->
-            block.unconditionalBranch = true
-            fn
-        }
+        setAsUnconditionalBranch(ctx, fn)
     }
 
     companion object : FuncBuilder()

src/main/kotlin/asmble/io/SExprToAst.kt

@@ -466,7 +466,7 @@ open class SExprToAst {
                     addMaybeExport(impExp, Node.ExternalKind.FUNCTION, funcCount++)
                     mod = mod.copy(funcs = mod.funcs + fn)
                 }
-                "export" -> mod = mod.copy(exports = mod.exports + toExport(exp, nameMap))
+                "export" -> mod = mod.copy(exports = mod.exports + toExport(it, nameMap))
                 "global" -> toGlobal(it, nameMap).also { (_, glb, impExp) ->
                     addMaybeExport(impExp, Node.ExternalKind.GLOBAL, globalCount++)
                     mod = mod.copy(globals = mod.globals + glb)
@@ -570,7 +570,6 @@ open class SExprToAst {
                 var instrAlign = 0
                 if (exp.vals.size > offset + count) exp.vals[offset + count].symbolStr().also {
                     if (it != null && it.startsWith("offset=")) {
-                        // TODO: unsigned ints everywhere!
                         instrOffset = UnsignedInteger.valueOf(it.substring(7)).toLong()
                         count++
                     }
@@ -578,6 +577,9 @@ open class SExprToAst {
                 if (exp.vals.size > offset + count) exp.vals[offset + count].symbolStr().also {
                     if (it != null && it.startsWith("align=")) {
                         instrAlign = it.substring(6).toInt()
+                        require(instrAlign > 0 && instrAlign and (instrAlign - 1) == 0) {
+                            "Alignment expected to be positive power of 2, but got $instrAlign"
+                        }
                         count++
                     }
                 }