Generate dynamic visitor

6 files changed, 155 insertions, 2 deletions

diff --git a/src/main/java/org/zwobble/couscous/util/ExtraIterables.java b/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
index a0c4593..d8635c9 100644
--- a/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
+++ b/src/main/java/org/zwobble/couscous/util/ExtraIterables.java
@@ -26,7 +26,7 @@ public class ExtraIterables {
         }
     }
 
-    public static <T1, T2, R> Iterable<R> map(Iterable<T1> iterable1, Iterable<T2> iterable2, BiFunction<T1, T2, R> function) {
+    public static <T1, T2, R> Iterable<R> lazyMap(Iterable<T1> iterable1, Iterable<T2> iterable2, BiFunction<T1, T2, R> function) {
         return new Iterable<R>() {
             @Override
             public Iterator<R> iterator() {
diff --git a/src/main/java/org/zwobble/couscous/util/ExtraMaps.java b/src/main/java/org/zwobble/couscous/util/ExtraMaps.java
index f25f06f..a659a01 100644
--- a/src/main/java/org/zwobble/couscous/util/ExtraMaps.java
+++ b/src/main/java/org/zwobble/couscous/util/ExtraMaps.java
@@ -29,7 +29,7 @@ public class ExtraMaps {
     }
 
     public static <T, K, V> Map<K, V> toMap(Iterable<K> keys, Iterable<V> values) {
-        return toMap(ExtraIterables.map(keys, values, ExtraMaps::entry));
+        return toMap(ExtraIterables.lazyMap(keys, values, ExtraMaps::entry));
     }
 
     public static <T, K, V> Map<K, V> toMap(Iterable<Map.Entry<K, V>> entries) {
diff --git a/src/main/java/org/zwobble/couscous/util/asm/StackManipulationFrame.java b/src/main/java/org/zwobble/couscous/util/asm/StackManipulationFrame.java
new file mode 100644
index 0000000..186c752
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/util/asm/StackManipulationFrame.java
@@ -0,0 +1,25 @@
+package org.zwobble.couscous.util.asm;
+
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.implementation.bytecode.StackSize;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+
+public class StackManipulationFrame implements StackManipulation {
+    public static final StackManipulation SAME = new StackManipulationFrame();
+
+    private StackManipulationFrame() {
+    }
+
+    @Override
+    public boolean isValid() {
+        return true;
+    }
+
+    @Override
+    public Size apply(MethodVisitor methodVisitor, Implementation.Context implementationContext) {
+        methodVisitor.visitFrame(Opcodes.F_SAME, 0, new Object[0], 0, new Object[0]);
+        return StackSize.ZERO.toIncreasingSize();
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/util/asm/StackManipulationLabel.java b/src/main/java/org/zwobble/couscous/util/asm/StackManipulationLabel.java
new file mode 100644
index 0000000..13d27bb
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/util/asm/StackManipulationLabel.java
@@ -0,0 +1,26 @@
+package org.zwobble.couscous.util.asm;
+
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.implementation.bytecode.StackSize;
+import net.bytebuddy.jar.asm.Label;
+import net.bytebuddy.jar.asm.MethodVisitor;
+
+public class StackManipulationLabel implements StackManipulation {
+    private final Label label;
+
+    public StackManipulationLabel(Label label) {
+        this.label = label;
+    }
+
+    @Override
+    public boolean isValid() {
+        return true;
+    }
+
+    @Override
+    public Size apply(MethodVisitor methodVisitor, Implementation.Context implementationContext) {
+        methodVisitor.visitLabel(label);
+        return StackSize.ZERO.toIncreasingSize();
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/util/asm/StackManipulationLookupTable.java b/src/main/java/org/zwobble/couscous/util/asm/StackManipulationLookupTable.java
new file mode 100644
index 0000000..f3fc2c2
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/util/asm/StackManipulationLookupTable.java
@@ -0,0 +1,30 @@
+package org.zwobble.couscous.util.asm;
+
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.implementation.bytecode.StackSize;
+import net.bytebuddy.jar.asm.Label;
+import net.bytebuddy.jar.asm.MethodVisitor;
+
+public class StackManipulationLookupTable implements StackManipulation {
+    private final Label dflt;
+    private final int[] keys;
+    private final Label[] labels;
+
+    public StackManipulationLookupTable(Label dflt, int[] keys, Label[] labels) {
+        this.dflt = dflt;
+        this.keys = keys;
+        this.labels = labels;
+    }
+
+    @Override
+    public boolean isValid() {
+        return true;
+    }
+
+    @Override
+    public Size apply(MethodVisitor methodVisitor, Implementation.Context implementationContext) {
+        methodVisitor.visitLookupSwitchInsn(dflt, keys, labels);
+        return StackSize.SINGLE.toDecreasingSize();
+    }
+}
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);
+    }
+}