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 isCurrentBlockDead get() = blockStack.lastOrNull()?.let { block ->
-    val isLastTerminating get() = insns.lastOrNull()?.isTerminating ?: false
+        // 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 addInsns(vararg insns: AbstractInsnNode) =
-
+        if (isCurrentBlockDead) this else copy(insns = this.insns + insns)
-    fun apply(fn: (Func) -> Func) = fn(this)
 
     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 popExpecting(type: TypeRef): Func {
-
+        assertTopOfStack(type)
-    fun popExpectingAny(vararg types: TypeRef): Func {
-        peekExpectingAny(types = *types)
         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()) =
+    fun assertTopOfStack(type: TypeRef, currBlock: Block? = blockStack.lastOrNull()): Unit {
-        peekExpectingAny(currBlock, type)
+        // If it's dead, we just go with it
-
+        if (!isCurrentBlockDead) {
-    fun peekExpectingAny(currBlock: Block? = blockStack.lastOrNull(), vararg types: TypeRef): TypeRef {
+            if (isStackEmptyForBlock(currBlock)) throw CompileErr.StackMismatch(arrayOf(type), null)
-        if (isStackEmptyForBlock(currBlock)) throw CompileErr.StackMismatch(types, null)
+            if (stack.lastOrNull() != type) throw CompileErr.StackMismatch(arrayOf(type), stack.lastOrNull())
-        val hasExpected = stack.lastOrNull()?.let(types::contains) ?: false
+        }
-        if (!hasExpected) throw CompileErr.StackMismatch(types, stack.lastOrNull())
-        return stack.last()
     }
 
     fun toMethodNode(): MethodNode {
@@ -71,7 +75,8 @@ data class Func(
 
     fun withoutAffectingStack(fn: (Func) -> Func) = fn(this).copy(stack = stack)
 
-    fun stackSwap(currBlock: Block? = blockStack.lastOrNull()) = pop(currBlock).let { (fn, refLast) ->
+    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)
@@ -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,26 +450,38 @@ open class FuncBuilder {
                 java.lang.Double::class.invokeStatic("longBitsToDouble", Double::class, Long::class))
     }
 
-    fun setAsUnconditionalBranchUpToDepth(ctx: FuncContext, fn: Func, depth: Int) =
+    fun setAsUnconditionalBranch(ctx: FuncContext, fn: Func) =
-        (0..depth).fold(fn) { fn, depth ->
+        fn.blockAtDepth(0).let { (fn, block) ->
-            fn.blockAtDepth(depth).let { (fn, block) -> block.unconditionalBranch = true; fn }
+            // 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 ->
+                // 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
+                if (type == null) block.origStack.size
                 else block.origStack.size + 1
             ctx.debug {
                 "Unconditional branch on ${block.insn}, curr stack ${fn.stack}, " +
@@ -479,11 +493,10 @@ open class FuncBuilder {
                     // 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)
+                        if (block.insn !is Node.Instr.Loop) fn.assertTopOfStack(typ, block)
                         block.blockExitVals += i to typ
                     }
-                        fn
+                    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) }
+                defaultBlock.insnType?.typeRef?.let { fn.assertTopOfStack(it) }
-                fn
+                // 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) ->
+        setAsUnconditionalBranch(ctx, fn)
-            block.unconditionalBranch = true
-            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++
                     }
                 }