Optimize AOT BR_TABLE to not include default in table

dylibso · Jul 5, 2024 · 4676aa1 · 4676aa1
1 parent 7e4782b
commit 4676aa1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/aot/src/main/java/com/dylibso/chicory/aot/AotMachine.java b/aot/src/main/java/com/dylibso/chicory/aot/AotMachine.java
@@ -750,14 +750,22 @@ private void compileBody(
                     exitBlockDepth = ins.depth();
                     ctx.popStackSize();
                     emitInvokeStatic(asm, CHECK_INTERRUPTION);
+                    // skip table switch if it only has a default
+                    if (ins.labelTable().length == 1) {
+                        asm.visitInsn(Opcodes.POP);
+                        emitUnwindStack(asm, type, body, ins, ins.labelTable()[0], ctx);
+                        asm.visitJumpInsn(Opcodes.GOTO, labels.get(ins.labelTable()[0]));
+                        break;
+                    }
                     // collect unique target labels
                     Map<Integer, Label> targets = new HashMap<>();
-                    Label[] table = new Label[ins.labelTable().length];
+                    Label[] table = new Label[ins.labelTable().length - 1];
                     for (int i = 0; i < table.length; i++) {
                         table[i] = targets.computeIfAbsent(ins.labelTable()[i], x -> new Label());
                     }
                     // table switch using the last entry of the label table as the default
-                    Label defaultLabel = table[table.length - 1];
+                    int defaultTarget = ins.labelTable()[ins.labelTable().length - 1];
+                    Label defaultLabel = targets.computeIfAbsent(defaultTarget, x -> new Label());
                     asm.visitTableSwitchInsn(0, table.length - 1, defaultLabel, table);
                     // generate separate unwinds for each target
                     targets.forEach(

diff --git a/aot/src/test/resources/com/dylibso/chicory/approvals/ApprovalTest.verifyBrTable.approved.txt b/aot/src/test/resources/com/dylibso/chicory/approvals/ApprovalTest.verifyBrTable.approved.txt
@@ -21,7 +21,6 @@ public final class com/dylibso/chicory/$gen/CompiledModule {
       0: L0
       1: L1
       2: L2
-      3: L3
       default: L3
    L2
    FRAME SAME