1 files changed, 72 insertions, 0 deletions
diff --git a/src/main/java/org/zwobble/couscous/util/asm/StackManipulationSwitch.java b/src/main/java/org/zwobble/couscous/util/asm/StackManipulationSwitch.java
new file mode 100644
index 0000000..2dbcd70
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/util/asm/StackManipulationSwitch.java
@@ -0,0 +1,72 @@
+package org.zwobble.couscous.util.asm;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Ordering;
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.jar.asm.Label;
+import net.bytebuddy.jar.asm.MethodVisitor;
+
+import java.util.List;
+
+import static java.util.Arrays.asList;
+import static org.zwobble.couscous.util.ExtraIterables.lazyMap;
+import static org.zwobble.couscous.util.ExtraLists.list;
+
+public class StackManipulationSwitch implements StackManipulation {
+    public static Case switchCase(int key, StackManipulation stackManipulation) {
+        return new Case(key, stackManipulation);
+    }
+
+    public static class Case {
+        private final int key;
+        private final StackManipulation stackManipulation;
+
+        private Case(int key, StackManipulation stackManipulation) {
+            this.key = key;
+            this.stackManipulation = stackManipulation;
+        }
+    }
+
+    private final StackManipulation defaultCase;
+    private final List<Case> cases;
+
+    public StackManipulationSwitch(StackManipulation defaultCase, List<Case> cases) {
+        this.defaultCase = defaultCase;
+        this.cases = Ordering.<Case>from((first, second) -> first.key - second.key).immutableSortedCopy(cases);
+    }
+
+    @Override
+    public boolean isValid() {
+        return true;
+    }
+
+    @Override
+    public Size apply(MethodVisitor methodVisitor, Implementation.Context implementationContext) {
+        Label defaultLabel = new Label();
+        Label[] caseLabels = Iterables.toArray(lazyMap(cases, switchCase -> new Label()), Label.class);
+
+        return new StackManipulation.Compound(ImmutableList.copyOf(Iterables.concat(
+            list(
+                new StackManipulationLookupTable(
+                    defaultLabel,
+                    cases.stream().mapToInt(switchCase -> switchCase.key).toArray(),
+                    caseLabels
+                )
+            ),
+
+            lazyMap(cases, asList(caseLabels), (switchCase, label) -> new StackManipulation.Compound(
+                new StackManipulationLabel(label),
+                StackManipulationFrame.SAME,
+                switchCase.stackManipulation
+            )),
+
+            list(
+                new StackManipulationLabel(defaultLabel),
+                StackManipulationFrame.SAME,
+                defaultCase
+            )
+        ))).apply(methodVisitor, implementationContext);
+    }
+}