introduce adt methods

Author: Andriy Onyshchuk

dataenum-processor/src/main/java/com/spotify/dataenum/processor/DataEnumProcessor.java

@@ -33,19 +33,24 @@
 import com.squareup.javapoet.ArrayTypeName;
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.TypeSpec;
+import com.sun.source.util.Trees;
 import java.io.IOException;
 import java.util.Collections;
 import java.util.Set;
-import javax.annotation.processing.AbstractProcessor;
-import javax.annotation.processing.Filer;
-import javax.annotation.processing.Messager;
-import javax.annotation.processing.RoundEnvironment;
+import javax.annotation.processing.*;
 import javax.lang.model.SourceVersion;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.TypeElement;
 import javax.tools.Diagnostic;
 
 public class DataEnumProcessor extends AbstractProcessor {
+  private Trees trees;
+
+  @Override
+  public synchronized void init(ProcessingEnvironment processingEnv) {
+    super.init(processingEnv);
+    this.trees = Trees.instance(processingEnv);
+  }
 
   @Override
   public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnvironment) {
@@ -58,7 +63,7 @@ public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnv
     for (Element element : roundEnvironment.getElementsAnnotatedWith(DataEnum.class)) {
       try {
 
-        Spec spec = SpecParser.parse(element, processingEnv);
+        Spec spec = SpecParser.parse(element, new ProcessingContext(processingEnv));
         if (spec == null) {
           continue;
         }

dataenum-processor/src/main/java/com/spotify/dataenum/processor/ProcessingContext.java

@@ -0,0 +1,24 @@
+package com.spotify.dataenum.processor;
+
+import com.spotify.dataenum.dataenum_case;
+import com.sun.source.util.Trees;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.type.TypeMirror;
+import javax.lang.model.util.Elements;
+import javax.lang.model.util.Types;
+
+public final class ProcessingContext {
+  public final ProcessingEnvironment env;
+  public final Trees trees;
+  public final TypeMirror dataenum_class_element;
+
+  public ProcessingContext(ProcessingEnvironment env) {
+    this.env = env;
+    this.trees = Trees.instance(env);
+
+    Types types = env.getTypeUtils();
+    Elements elements = env.getElementUtils();
+    this.dataenum_class_element =
+        elements.getTypeElement(dataenum_case.class.getCanonicalName()).asType();
+  }
+}

dataenum-processor/src/main/java/com/spotify/dataenum/processor/data/OutputSpec.java

@@ -31,7 +31,7 @@ public class OutputSpec extends Spec {
   private final Iterable<OutputValue> outputValues;
 
   public OutputSpec(Spec input, ClassName outputClass, Iterable<OutputValue> outputValues) {
-    super(input.specClass(), input.typeVariables(), input.values());
+    super(input.specClass(), input.typeVariables(), input.values(), input.methods());
     this.outputClass = outputClass;
     this.outputValues = outputValues;
   }

dataenum-processor/src/main/java/com/spotify/dataenum/processor/data/Spec.java

@@ -21,6 +21,7 @@
 
 import com.spotify.dataenum.processor.util.Iterables;
 import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.TypeVariableName;
 
 /**
@@ -33,12 +34,17 @@ public class Spec {
   private final ClassName specClass;
   private final Iterable<TypeVariableName> typeVariables;
   private final Iterable<Value> values;
+  private final Iterable<MethodSpec> methods;
 
   public Spec(
-      ClassName specClass, Iterable<TypeVariableName> typeVariables, Iterable<Value> values) {
+      ClassName specClass,
+      Iterable<TypeVariableName> typeVariables,
+      Iterable<Value> values,
+      Iterable<MethodSpec> methods) {
     this.specClass = specClass;
     this.typeVariables = typeVariables;
     this.values = values;
+    this.methods = methods;
   }
 
   public ClassName specClass() {
@@ -56,4 +62,8 @@ public boolean hasTypeVariables() {
   public Iterable<Value> values() {
     return values;
   }
+
+  public Iterable<MethodSpec> methods() {
+    return methods;
+  }
 }

dataenum-processor/src/main/java/com/spotify/dataenum/processor/generator/method/MethodMethods.java

@@ -0,0 +1,77 @@
+package com.spotify.dataenum.processor.generator.method;
+
+import com.squareup.javapoet.*;
+import com.sun.source.tree.MethodTree;
+import com.sun.source.util.TreePathScanner;
+import com.sun.source.util.Trees;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.ExecutableElement;
+import javax.lang.model.element.Modifier;
+
+public final class MethodMethods {
+  private MethodMethods() {}
+
+  public static MethodSpec.Builder builderFrom(
+      ExecutableElement el,
+      Function<Stream<Modifier>, Stream<Modifier>> adjustModifiers,
+      Function<Stream<? extends AnnotationMirror>, Stream<? extends AnnotationMirror>>
+          adjustAnnotations) {
+    return MethodSpec.methodBuilder(el.getSimpleName().toString())
+        .addModifiers(adjustModifiers.apply(el.getModifiers().stream()).collect(Collectors.toSet()))
+        .addAnnotations(
+            adjustAnnotations
+                .apply(el.getAnnotationMirrors().stream())
+                .map(AnnotationSpec::get)
+                .collect(Collectors.toList()))
+        .addTypeVariables(
+            el.getTypeParameters().stream().map(TypeVariableName::get).collect(Collectors.toList()))
+        .returns(TypeName.get(el.getReturnType()))
+        .addParameters(
+            el.getParameters().stream().map(ParameterSpec::get).collect(Collectors.toList()))
+        .varargs(el.isVarArgs())
+        .addExceptions(
+            el.getThrownTypes().stream().map(TypeName::get).collect(Collectors.toList()));
+  }
+
+  public static CodeBlock codeBlockFrom(ExecutableElement el, Trees trees) {
+    final MethodTree methodTree = MethodLookup.lookupTree(el, trees);
+    return methodTree
+        .getBody()
+        .getStatements()
+        .stream()
+        .map(x -> CodeBlock.of(x.toString()))
+        .reduce(CodeBlock.of(""), (a, b) -> a.toBuilder().add(b).build());
+  }
+
+  private static class MethodLookup extends TreePathScanner<Void, Void> {
+
+    private final Consumer<MethodTree> onMethod;
+
+    private MethodLookup(Consumer<MethodTree> onMethod) {
+      this.onMethod = onMethod;
+    }
+
+    public static MethodTree lookupTree(ExecutableElement methodElement, Trees trees) {
+      assert methodElement.getKind() == ElementKind.METHOD;
+
+      AtomicReference<MethodTree> methodRef = new AtomicReference<>();
+      new MethodLookup(methodRef::set).scan(trees.getPath(methodElement), null);
+      MethodTree method = methodRef.get();
+      assert method != null;
+
+      return method;
+    }
+
+    @Override
+    public Void visitMethod(MethodTree methodTree, Void v) {
+      this.onMethod.accept(methodTree);
+      return super.visitMethod(methodTree, v);
+    }
+  }
+}

dataenum-processor/src/main/java/com/spotify/dataenum/processor/generator/spec/SpecTypeFactory.java

@@ -95,6 +95,8 @@ public static TypeSpec create(
       enumBuilder.addMethod(mapMethods.createAbstractFoldMethod());
     }
 
+    enumBuilder.addMethods(spec.methods());
+
     return enumBuilder.build();
   }
 }