Generate dynamic visitor

2 files changed, 291 insertions, 0 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/NodeTypes.java b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
new file mode 100644
index 0000000..a82cd1c
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/NodeTypes.java
@@ -0,0 +1,86 @@
+package org.zwobble.couscous.ast;
+
+import com.google.common.collect.ImmutableMap;
+
+import java.util.Map;
+
+public class NodeTypes {
+    public static final int ANNOTATION = 10000;
+    public static final int ARRAY = 10001;
+    public static final int ASSIGNMENT = 10002;
+    public static final int CAST = 10003;
+    public static final int CLASS = 10004;
+    public static final int CONSTRUCTOR_CALL = 10005;
+    public static final int CONSTRUCTOR = 10006;
+    public static final int ENUM = 10007;
+    public static final int EXCEPTION_HANDLER = 10008;
+    public static final int EXPRESSION_STATEMENT = 10009;
+    public static final int FIELD_ACCESS = 10010;
+    public static final int FIELD_DECLARATION = 10011;
+    public static final int FORMAL_ARGUMENT = 10012;
+    public static final int FORMAL_TYPE_PARAMETER = 10013;
+    public static final int IF_STATEMENT = 10014;
+    public static final int INSTANCE_OF = 10015;
+    public static final int INSTANCE_RECEIVER = 10016;
+    public static final int INTERFACE = 10017;
+    public static final int LITERAL = 10018;
+    public static final int LOCAL_VARIABLE_DECLARATION = 10019;
+    public static final int METHOD_CALL = 10020;
+    public static final int METHOD = 10031;
+    public static final int OPERATION = 10021;
+    public static final int RETURN = 10022;
+    public static final int STATIC_RECEIVER = 10023;
+    public static final int TERNARY_CONDITIONAL = 10024;
+    public static final int THIS_REFERENCE = 10025;
+    public static final int THROW = 10026;
+    public static final int TRY = 10027;
+    public static final int TYPE_COERCION = 10028;
+    public static final int VARIABLE_REFERENCE = 10029;
+    public static final int WHILE = 10030;
+
+    private NodeTypes() {
+    }
+
+    private static final Map<Class<? extends Node>, Integer> NODE_TYPES = ImmutableMap.<Class<? extends Node>, Integer>builder()
+        .put(AnnotationNode.class, ANNOTATION)
+        .put(ArrayNode.class, ARRAY)
+        .put(AssignmentNode.class, ASSIGNMENT)
+        .put(CastNode.class, CAST)
+        .put(ClassNode.class, CLASS)
+        .put(ConstructorCallNode.class, CONSTRUCTOR_CALL)
+        .put(ConstructorNode.class, CONSTRUCTOR)
+        .put(EnumNode.class, ENUM)
+        //.put(ExceptionHandlerNode.class, EXCEPTION_HANDLER)
+        .put(ExpressionStatementNode.class, EXPRESSION_STATEMENT)
+        .put(FieldAccessNode.class, FIELD_ACCESS)
+        .put(FieldDeclarationNode.class, FIELD_DECLARATION)
+        .put(FormalArgumentNode.class, FORMAL_ARGUMENT)
+        .put(FormalTypeParameterNode.class, FORMAL_TYPE_PARAMETER)
+        .put(IfStatementNode.class, IF_STATEMENT)
+        .put(InstanceOfNode.class, INSTANCE_OF)
+        .put(InstanceReceiver.class, INSTANCE_RECEIVER)
+        .put(InterfaceNode.class, INTERFACE)
+        .put(LiteralNode.class, LITERAL)
+        .put(LocalVariableDeclarationNode.class, LOCAL_VARIABLE_DECLARATION)
+        .put(MethodCallNode.class, METHOD_CALL)
+        .put(MethodNode.class, METHOD)
+        .put(OperationNode.class, OPERATION)
+        .put(ReturnNode.class, RETURN)
+        .put(StaticReceiver.class, STATIC_RECEIVER)
+        .put(TernaryConditionalNode.class, TERNARY_CONDITIONAL)
+        .put(ThisReferenceNode.class, THIS_REFERENCE)
+        .put(ThrowNode.class, THROW)
+        .put(TryNode.class, TRY)
+        .put(TypeCoercionNode.class, TYPE_COERCION)
+        .put(VariableReferenceNode.class, VARIABLE_REFERENCE)
+        .put(WhileNode.class, WHILE)
+        .build();
+
+    public static int forClass(Class<?> nodeClass) {
+        if (NODE_TYPES.containsKey(nodeClass)) {
+            return NODE_TYPES.get(nodeClass);
+        } else {
+            throw new UnsupportedOperationException("Unknown node class: " + nodeClass);
+        }
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java b/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java
new file mode 100644
index 0000000..f9e406d
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/DynamicNodeVisitor.java
@@ -0,0 +1,205 @@
+package org.zwobble.couscous.ast.visitors;
+
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.modifier.FieldManifestation;
+import net.bytebuddy.description.modifier.Visibility;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.scaffold.InstrumentedType;
+import net.bytebuddy.implementation.Implementation;
+import net.bytebuddy.implementation.bytecode.ByteCodeAppender;
+import net.bytebuddy.implementation.bytecode.StackManipulation;
+import net.bytebuddy.implementation.bytecode.assign.TypeCasting;
+import net.bytebuddy.implementation.bytecode.member.FieldAccess;
+import net.bytebuddy.implementation.bytecode.member.MethodInvocation;
+import net.bytebuddy.implementation.bytecode.member.MethodReturn;
+import net.bytebuddy.implementation.bytecode.member.MethodVariableAccess;
+import net.bytebuddy.jar.asm.Label;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.matcher.ElementMatchers;
+import org.zwobble.couscous.ast.Node;
+import org.zwobble.couscous.ast.NodeTypes;
+import org.zwobble.couscous.util.asm.StackManipulationSwitch;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.List;
+import java.util.function.Consumer;
+
+import static java.util.Arrays.asList;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+import static org.zwobble.couscous.util.ExtraLists.eagerFilter;
+import static org.zwobble.couscous.util.ExtraLists.eagerMap;
+
+/**
+ * Dynamically generate a method to dispatch methods based on node type.
+ *
+ * <p>For instance, let's say you've defined two methods in a class:</p>
+ *
+ * <blockquote><pre>
+ *     void visit(LiteralNode);
+ *     void visit(VariableReferenceNode);
+ * </pre></blockquote>
+ *
+ * Calling:
+ *
+ * <blockquote><pre>
+ *     DynamicNodeVisitor.build(NodeWriter.class, "visit").instantiate(nodeWriter)
+ * </pre></blockquote>
+ *
+ * will generate a function accepting a node equivalent to:
+ *
+ * <blockquote><pre>
+ *     switch (node.nodeType()) {
+ *         case NodeTypes.LITERAL:
+ *             nodeWriter.visit((LiteralNode) node);
+ *             return;
+ *         case NodeTypes.VARIABLE_REFERENCE:
+ *             nodeWriter.visit((VariableReferenceNode) node);
+ *             return;
+ *         default;
+ *             return;
+ *     }
+ * </pre></blockquote>
+ *
+ * <p>
+ *   Note that {@code build} is comparatively expensive,
+ *   so should normally be called once separately from {@code instantiate}.
+ * </p>
+ *
+ * <p>
+ *   In answer to "why not define a visitor interface instead?":
+ *   this allows visitors to be defined over a subset of nodes,
+ *   rather than requiring a visit method to be defined for every node.
+ *   Further, it allows extension of the Node interface,
+ *   making it easier to mix and match transformations that may add or remove particular node types
+ *   (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) {
+        List<Method> visitMethods = eagerFilter(
+            asList(clazz.getMethods()),
+            method -> isVisitMethod(methodName, method)
+        );
+        Label literalLabel = new Label();
+
+        Class<? extends Consumer> visitor = new ByteBuddy()
+            .subclass(Consumer.class)
+
+            .defineField("visitor", clazz, Visibility.PRIVATE, FieldManifestation.FINAL)
+
+            .defineConstructor(Visibility.PUBLIC)
+            .withParameters(clazz)
+            .intercept(new Implementation() {
+                @Override
+                public InstrumentedType prepare(InstrumentedType instrumentedType) {
+                    return instrumentedType;
+                }
+
+                @Override
+                public ByteCodeAppender appender(Target implementationTarget) {
+                    return new ByteCodeAppender() {
+                        @Override
+                        public Size apply(MethodVisitor methodVisitor, Context implementationContext, MethodDescription instrumentedMethod) {
+                            StackManipulation.Size size = new StackManipulation.Compound(
+                                // this()
+                                MethodVariableAccess.REFERENCE.loadOffset(0),
+                                MethodInvocation.invoke(new TypeDescription.ForLoadedType(Object.class)
+                                    .getDeclaredMethods()
+                                    .filter(isConstructor().and(takesArguments(0))).getOnly()),
+
+                                // this.visitor = visitor
+                                MethodVariableAccess.REFERENCE.loadOffset(0),
+                                MethodVariableAccess.REFERENCE.loadOffset(1),
+                                FieldAccess.forField(implementationTarget.getInstrumentedType()
+                                    .getDeclaredFields()
+                                    .filter(ElementMatchers.named("visitor"))
+                                    .getOnly()).putter(),
+
+                                // return
+                                MethodReturn.VOID
+                            ).apply(methodVisitor, implementationContext);
+                            return new Size(size.getMaximalSize(), instrumentedMethod.getStackSize());
+                        }
+                    };
+                }
+            })
+
+            .method(ElementMatchers.named("accept"))
+            .intercept(new Implementation() {
+                @Override
+                public InstrumentedType prepare(InstrumentedType instrumentedType) {
+                    return instrumentedType;
+                }
+
+                @Override
+                public ByteCodeAppender appender(Target target) {
+                    return new ByteCodeAppender() {
+                        @Override
+                        public Size apply(MethodVisitor methodVisitor, Context context, MethodDescription instrumentedMethod) {
+                            MethodDescription typeMethod = new TypeDescription.ForLoadedType(Node.class)
+                                .getDeclaredMethods()
+                                .filter(ElementMatchers.named("type").and(ElementMatchers.takesArguments(0)))
+                                .getOnly();
+                            StackManipulation.Size size = new StackManipulation.Compound(
+                                MethodVariableAccess.REFERENCE.loadOffset(1),
+                                MethodInvocation.invoke(typeMethod),
+                                new StackManipulationSwitch(
+                                    // TODO: raise exception in default case
+                                    MethodReturn.VOID,
+                                    eagerMap(visitMethods, method -> {
+                                        Class<?> nodeClass = method.getParameterTypes()[0];
+                                        return StackManipulationSwitch.switchCase(
+                                            NodeTypes.forClass(nodeClass),
+                                            new StackManipulation.Compound(
+                                                MethodVariableAccess.REFERENCE.loadOffset(0),
+                                                FieldAccess.forField(target.getInstrumentedType()
+                                                    .getDeclaredFields()
+                                                    .filter(ElementMatchers.named("visitor"))
+                                                    .getOnly()).getter(),
+                                                MethodVariableAccess.REFERENCE.loadOffset(1),
+                                                TypeCasting.to(new TypeDescription.ForLoadedType(nodeClass)),
+                                                MethodInvocation.invoke(new MethodDescription.ForLoadedMethod(method)),
+                                                MethodReturn.VOID
+                                            )
+                                        );
+                                    })
+                                )
+                            ).apply(methodVisitor, context);
+                            return new Size(size.getMaximalSize(), instrumentedMethod.getStackSize());
+                        }
+                    };
+                }
+            })
+            .make()
+            .load(DynamicNodeVisitor.class.getClassLoader())
+            .getLoaded();
+
+        return new DynamicNodeVisitor<T>(clazz, visitor);
+    }
+
+    private static boolean isVisitMethod(String methodName, Method method) {
+        return
+            method.getName().equals(methodName) &&
+            method.getParameterCount() == 1 &&
+            Node.class.isAssignableFrom(method.getParameterTypes()[0]);
+    }
+
+    private final Class<T> clazz;
+    private final Class<? extends Consumer> visitor;
+
+    public DynamicNodeVisitor(Class<T> clazz, Class<? extends Consumer> visitor) {
+        this.clazz = clazz;
+        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);
+        }
+    }
+}