2 files changed, 37 insertions, 27 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeMapper.java b/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeMapper.java
index 598c909..243dde1 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeMapper.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeMapper.java
@@ -4,7 +4,6 @@ import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.implementation.bytecode.member.MethodReturn;
 import org.zwobble.couscous.ast.Node;
 
-import java.lang.reflect.InvocationTargetException;
 import java.util.function.Function;
 
 public class DynamicNodeMapper<T, R> {
@@ -14,24 +13,17 @@ public class DynamicNodeMapper<T, R> {
 
     public static <T, R> DynamicNodeMapper<T, R> build(Class<T> clazz, TypeDescription returnType, String methodName) {
         return new DynamicNodeMapper<>(
-            clazz,
-            DynamicNodeVisitor.buildClass(clazz, Function.class, methodName, MethodReturn.REFERENCE)
+            DynamicNodeVisitor.buildClassSupplier(clazz, Function.class, methodName, MethodReturn.REFERENCE)
         );
     }
 
-    private final Class<T> clazz;
-    private final Class<? extends Function> visitor;
+    private final Function<T, Function> visitor;
 
-    public DynamicNodeMapper(Class<T> clazz, Class<? extends Function> visitor) {
-        this.clazz = clazz;
+    public DynamicNodeMapper(Function<T, Function> visitor) {
         this.visitor = visitor;
     }
 
     public Function<Node, R> instantiate(T instance) {
-        try {
-            return visitor.getConstructor(clazz).newInstance(instance);
-        } catch (InstantiationException | IllegalAccessException | NoSuchMethodException | InvocationTargetException exception) {
-            throw new RuntimeException(exception);
-        }
+        return visitor.apply(instance);
     }
 }
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java b/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java
index e473cf9..a64b7bc 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java
@@ -23,11 +23,11 @@ import org.zwobble.couscous.util.asm.Implementations;
 import org.zwobble.couscous.util.asm.StackManipulationSwitch;
 import org.zwobble.couscous.util.asm.TypeDescriptions;
 
-import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import static java.util.Arrays.asList;
 import static org.zwobble.couscous.util.ExtraLists.eagerFilter;
@@ -80,13 +80,37 @@ import static org.zwobble.couscous.util.ExtraLists.eagerMap;
  */
 public class DynamicNodeVisitor<T> {
     public static <T> DynamicNodeVisitor<T> build(Class<T> clazz, String methodName) {
-        return new DynamicNodeVisitor<>(
-            clazz,
-            buildClass(clazz, Consumer.class, methodName, MethodReturn.VOID)
-        );
+        Function<T, Consumer> builder = buildClassSupplier(clazz, Consumer.class, methodName, MethodReturn.VOID);
+        return new DynamicNodeVisitor<>(builder);
+    }
+
+    public static <T, F> Function<T, F> buildClassSupplier(Class<T> clazz, Class<F> function, String methodName, MethodReturn methodReturn) {
+        Class<? extends F> visitorClass = buildClass(clazz, function, methodName, methodReturn);
+
+        try {
+            return new ByteBuddy()
+                .subclass(Function.class)
+
+                .method(ElementMatchers.isAbstract())
+                .intercept(Implementations.stackManipulation(target -> new StackManipulation.Compound(
+                    TypeCreation.of(new TypeDescription.ForLoadedType(visitorClass)),
+                    Duplication.SINGLE,
+                    MethodVariableAccess.REFERENCE.loadOffset(1),
+                    TypeCasting.to(new TypeDescription.ForLoadedType(clazz)),
+                    MethodInvocation.invoke(TypeDescriptions.findConstructor(visitorClass, clazz)),
+                    MethodReturn.REFERENCE
+                )))
+
+                .make()
+                .load(DynamicNodeVisitor.class.getClassLoader())
+                .getLoaded()
+                .newInstance();
+        } catch (InstantiationException | IllegalAccessException exception) {
+            throw new RuntimeException(exception);
+        }
     }
 
-    public static <T, F> Class<? extends F> buildClass(Class<T> clazz, Class<F> function, String methodName, MethodReturn methodReturn) {
+    private static <T, F> Class<? extends F> buildClass(Class<T> clazz, Class<F> function, String methodName, MethodReturn methodReturn) {
         List<Method> visitMethods = eagerFilter(
             asList(clazz.getMethods()),
             method -> isVisitMethod(methodName, method)
@@ -171,19 +195,13 @@ public class DynamicNodeVisitor<T> {
             !method.getParameterTypes()[0].equals(Node.class);
     }
 
-    private final Class<T> clazz;
-    private final Class<? extends Consumer> visitor;
+    private final Function<T, Consumer> visitor;
 
-    public DynamicNodeVisitor(Class<T> clazz, Class<? extends Consumer> visitor) {
-        this.clazz = clazz;
+    private DynamicNodeVisitor(Function<T, Consumer> visitor) {
         this.visitor = visitor;
     }
 
     public Consumer<Node> instantiate(T instance) {
-        try {
-            return visitor.getConstructor(clazz).newInstance(instance);
-        } catch (InstantiationException | IllegalAccessException | NoSuchMethodException | InvocationTargetException exception) {
-            throw new RuntimeException(exception);
-        }
+        return visitor.apply(instance);
     }
 }