Simplify dynamic visitor creation

2 files changed, 30 insertions, 36 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 243dde1..0637c44 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeMapper.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeMapper.java
@@ -1,29 +1,26 @@
 package org.zwobble.couscous.ast.visitors;
 
-import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.implementation.bytecode.member.MethodReturn;
 import org.zwobble.couscous.ast.Node;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.function.Function;
 
-public class DynamicNodeMapper<T, R> {
-    public static <T, R> DynamicNodeMapper<T, R> build(Class<T> clazz, Class<R> returnType, String methodName) {
-        return build(clazz, new TypeDescription.ForLoadedType(returnType), methodName);
-    }
-
-    public static <T, R> DynamicNodeMapper<T, R> build(Class<T> clazz, TypeDescription returnType, String methodName) {
-        return new DynamicNodeMapper<>(
-            DynamicNodeVisitor.buildClassSupplier(clazz, Function.class, methodName, MethodReturn.REFERENCE)
-        );
-    }
+import static org.zwobble.couscous.util.ExtraMaps.entry;
 
-    private final Function<T, Function> visitor;
-
-    public DynamicNodeMapper(Function<T, Function> visitor) {
-        this.visitor = visitor;
-    }
+public class DynamicNodeMapper {
+    private static final Map<Map.Entry<Class<?>, String>, Function<?, Function>> VISITOR_BUILDERS
+        = new HashMap<>();
 
-    public Function<Node, R> instantiate(T instance) {
-        return visitor.apply(instance);
+    // TODO: check return type -- at the very least, check the methods are consistent.
+    public static <T, R> Function<Node, R> instantiate(T visitor, String methodName) {
+        Map.Entry<Class<?>, String> key = entry(visitor.getClass(), methodName);
+        Function builder = VISITOR_BUILDERS.get(key);
+        if (builder == null) {
+            builder = DynamicNodeVisitor.buildClassSupplier(visitor.getClass(), Function.class, methodName, MethodReturn.REFERENCE);
+            VISITOR_BUILDERS.put(key, builder);
+        }
+        return (Function<Node, R>) builder.apply(visitor);
     }
 }
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 aed1253..438bcea 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java
@@ -25,9 +25,7 @@ import org.zwobble.couscous.util.asm.StackManipulationSwitch;
 import org.zwobble.couscous.util.asm.TypeDescriptions;
 
 import java.lang.reflect.Method;
-import java.util.List;
-import java.util.Optional;
-import java.util.Set;
+import java.util.*;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -36,6 +34,7 @@ import static java.util.Arrays.asList;
 import static org.zwobble.couscous.util.ExtraIterables.lazyMap;
 import static org.zwobble.couscous.util.ExtraLists.eagerFilter;
 import static org.zwobble.couscous.util.ExtraLists.eagerMap;
+import static org.zwobble.couscous.util.ExtraMaps.entry;
 
 /**
  * Dynamically generate a method to dispatch methods based on node type.
@@ -82,10 +81,18 @@ import static org.zwobble.couscous.util.ExtraLists.eagerMap;
  *   (which affects which node types later stages are required to handle).
  * </p>
  */
-public class DynamicNodeVisitor<T> {
-    public static <T> DynamicNodeVisitor<T> build(Class<T> clazz, String methodName) {
-        Function<T, Consumer> builder = buildClassSupplier(clazz, Consumer.class, methodName, MethodReturn.VOID);
-        return new DynamicNodeVisitor<>(builder);
+public class DynamicNodeVisitor {
+    private static final Map<Map.Entry<Class<?>, String>, Function<?, Consumer>> VISITOR_BUILDERS
+        = new HashMap<>();
+
+    public static <T> Consumer<Node> instantiate(T visitor, String methodName) {
+        Map.Entry<Class<?>, String> key = entry(visitor.getClass(), methodName);
+        Function builder = VISITOR_BUILDERS.get(key);
+        if (builder == null) {
+            builder = buildClassSupplier(visitor.getClass(), Consumer.class, methodName, MethodReturn.VOID);
+            VISITOR_BUILDERS.put(key, builder);
+        }
+        return (Consumer<Node>) builder.apply(visitor);
     }
 
     public static <T> BiConsumer<Node, T> visitor(Class<T> clazz) {
@@ -98,7 +105,7 @@ public class DynamicNodeVisitor<T> {
         return (node, visitor) -> classSupplier.apply(visitor).accept(node);
     }
 
-    public static <T, F> Function<T, F> buildClassSupplier(Class<T> clazz, Class<F> function, String methodName, MethodReturn methodReturn) {
+    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 {
@@ -208,14 +215,4 @@ public class DynamicNodeVisitor<T> {
             Node.class.isAssignableFrom(method.getParameterTypes()[0]) &&
             !method.getParameterTypes()[0].equals(Node.class);
     }
-
-    private final Function<T, Consumer> visitor;
-
-    private DynamicNodeVisitor(Function<T, Consumer> visitor) {
-        this.visitor = visitor;
-    }
-
-    public Consumer<Node> instantiate(T instance) {
-        return visitor.apply(instance);
-    }
 }