path: root/src/main/java/ftbsc/geb/processor/GEBProcessor.java

package ftbsc.geb.processor;

import com.squareup.javapoet.*;
import ftbsc.geb.api.IEvent;
import ftbsc.geb.api.IEventCancelable;
import ftbsc.geb.api.IEventDispatcher;
import ftbsc.geb.api.IListener;
import ftbsc.geb.api.annotations.Listen;
import ftbsc.geb.exceptions.BadListenerArgumentsException;
import ftbsc.geb.exceptions.MissingInterfaceException;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.tools.Diagnostic;
import javax.tools.FileObject;
import javax.tools.JavaFileObject;
import javax.tools.StandardLocation;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.*;
import java.util.stream.Collectors;

/**
 * GEB's {@link javax.annotation.processing.Processor annotation processor},
 * which takes care of generating the {@link IEventDispatcher dispatchers}.
 */
@SupportedAnnotationTypes("ftbsc.geb.api.annotations.Listen")
public class GEBProcessor extends AbstractProcessor {

   /**
    * A {@link Map} tying each event class to a {@link Set} of listeners.
    */
   private final Map<TypeMirror, Set<ListenerContainer>> listenerMap = new HashMap<>();

   /**
    * A {@link Set} containing the fully-qualified names of the generated classes.
    */
   private final Set<String> generatedClasses = new HashSet<>();

   /**
    * A {@link TypeMirror} representing the {@link IListener} interface.
    */
   private TypeMirror listenerInterface;

   /**
    * A {@link TypeMirror} representing the {@link IEvent} interface.
    */
   private TypeMirror eventInterface;

   /**
    * A {@link TypeMirror} representing the {@link IEventCancelable} interface.
    */
   private TypeMirror cancelableEventInterface;

   /**
    * A {@link TypeMirror} representing the {@link IEventDispatcher} interface.
    */
   private TypeMirror dispatcherInterface;

   /**
    * Initializes the processor with the given environment.
    * Also takes carae of initializing the TypeMirror "constants" for later use.
    * @param env the environment
    */
   @Override
   public synchronized void init(ProcessingEnvironment env) {
      super.init(env);
      this.listenerInterface = env.getElementUtils()
         .getTypeElement("ftbsc.geb.api.IListener").asType();
      this.eventInterface = env.getElementUtils()
         .getTypeElement("ftbsc.geb.api.IEvent").asType();
      this.dispatcherInterface = env.getElementUtils()
         .getTypeElement("ftbsc.geb.api.IEventDispatcher").asType();
      this.cancelableEventInterface = env.getElementUtils()
         .getTypeElement("ftbsc.geb.api.IEventCancelable").asType();
   }

   /**
    * The starting point of the processor.
    * It calls {@link #processListener(Element)} on all elements annotated with
    * the {@link Listen} annotation.
    * @param annotations the annotation types requested to be processed
    * @param env environment for information about the current and prior round
    * @return whether the set of annotation types are claimed by this processor
    */
   @Override
   public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment env) {
      boolean claimed = false;
      for(TypeElement ann : annotations) {
         if(ann.getQualifiedName().contentEquals(Listen.class.getName())) {
            claimed = true;
            for(Element e : env.getElementsAnnotatedWith(ann))
               this.processListener(e);
            if(!this.listenerMap.isEmpty()) {
               this.generateClasses();
               this.generateServiceProvider();
            }
         }
      }
      return claimed;
   }

   /**
    * Sets the supported source version to the latest one.
    * It's either that or constant warnings, and the processor is simple enough.
    * @return the latest source version
    */
   @Override
   public SourceVersion getSupportedSourceVersion() {
      return SourceVersion.latest();
   }

   /**
    * Verifies that the annotated method is valid and, if it is, adds it to
    * the list. See the annotation's javadoc for details on what's considered
    * a valid listener.
    * @see Listen
    * @param target the {@link Element} that was annotated with {@link Listen}
    */
   private void processListener(Element target) {
      ExecutableElement listener = (ExecutableElement) target; //this cast will never fail

      //ensure the parent is instance of IListener
      TypeMirror parentType = listener.getEnclosingElement().asType();
      if(!this.processingEnv.getTypeUtils().isAssignable(parentType, this.listenerInterface))
         throw new MissingInterfaceException(
            listener.getEnclosingElement().getSimpleName().toString(),
            listener.getSimpleName().toString());

      //ensure the listener method has only one parameter
      List<? extends VariableElement> params = listener.getParameters();
      if(listener.getParameters().size() != 1)
         throw new BadListenerArgumentsException.Count(
            listener.getEnclosingElement().getSimpleName().toString(),
            listener.getSimpleName().toString(),
            params.size());

      //ensure said parameter implements IEvent
      TypeMirror event = params.get(0).asType();
      if(!this.processingEnv.getTypeUtils().isAssignable(event, this.eventInterface))
         throw new BadListenerArgumentsException.Type(
            listener.getEnclosingElement().getSimpleName().toString(),
            listener.getSimpleName().toString(),
            params.get(0).getSimpleName().toString());

      //warn about return type
      if(!listener.getReturnType().getKind().equals(TypeKind.VOID))
         this.processingEnv.getMessager().printMessage(Diagnostic.Kind.WARNING, String.format(
            "The method %s::%s has a return type: please note that it will be ignored.",
            listener.getEnclosingElement().getSimpleName().toString(),
            listener.getSimpleName().toString()));

      this.listenerMap.computeIfAbsent(event, k -> new HashSet<>())
         .add(new ListenerContainer(listener));
   }

   /**
    * Uses JavaPoet to generate the classes dispatcher classes.
    */
   private void generateClasses() {
      this.listenerMap.forEach((event, listeners) -> {
         TypeElement eventClass = (TypeElement) this.processingEnv.getTypeUtils().asElement(event);
         boolean cancelable = this.processingEnv.getTypeUtils().isAssignable(event, this.cancelableEventInterface);

         ParameterSpec eventParam = ParameterSpec.builder(TypeName.get(this.eventInterface), "event").build();
         ParameterSpec listenersParam = ParameterSpec.builder(
            ParameterizedTypeName.get(
               ClassName.get("java.util", "Map"),
               ParameterizedTypeName.get(
                  ClassName.get("java.lang", "Class"),
                  WildcardTypeName.subtypeOf(TypeName.get(this.listenerInterface))
               ),
               ParameterizedTypeName.get(
                  ClassName.get("java.util", "Set"),
                  ClassName.get(this.listenerInterface)
               )
            ),
            "listeners"
         ).build();

         MethodSpec.Builder callListenersBuilder = MethodSpec.methodBuilder("callListeners")
            .addModifiers(Modifier.PUBLIC)
            .addAnnotation(Override.class)
            .addAnnotation(AnnotationSpec.builder(SuppressWarnings.class) // because why not
               .addMember("value" , "{$S}", "unchecked").build())
            .addParameter(eventParam)
            .addParameter(listenersParam)
            .returns(boolean.class);

         // reorder the injectors to follow priority
         Map<TypeMirror, Integer> done = new HashMap<>();
         List<ListenerContainer> ordered = listeners.stream().sorted(Comparator.comparingInt(
            container -> container.annotation.priority()
         )).collect(Collectors.toList());

         // get all the relevant injectors
         for(int i = 0; i < ordered.size(); i++) {
            ListenerContainer listener = ordered.get(i);
            if(!done.containsKey(listener.parent)) {
               done.put(listener.parent, i);
               String varName = String.format("listener%d", i);
               callListenersBuilder.addStatement(
                  "java.util.Set<$T> $L = $N.get($T.class)",
                  this.listenerInterface,
                  varName,
                  listenersParam,
                  listener.parent
               );
            }
         }

         for(ListenerContainer listener : ordered) {
            String varName = String.format("listener%d", done.get(listener.parent));
            callListenersBuilder
               .addStatement("if($L != null) { for($T l : $L) {", varName, this.listenerInterface, varName)
               .addStatement(
                  "if(l != null) (($T) l).$L(($T) $N); } }",
                  listener.parent,
                  listener.method.getSimpleName().toString(),
                  event,
                  eventParam
            );
            if(cancelable) callListenersBuilder
               .addStatement("if((($T) $N).isCanceled()) return false", this.cancelableEventInterface, eventParam);
         }

         callListenersBuilder.addStatement("return false");

         MethodSpec eventType = MethodSpec.methodBuilder("eventType")
            .addModifiers(Modifier.PUBLIC)
            .addAnnotation(Override.class)
            .returns(Class.class)
            .addStatement("return $T.class", event)
            .build();

         String clazzName = String.format("%sDispatcher", eventClass.getSimpleName());
         TypeSpec clazz = TypeSpec.classBuilder(clazzName)
            .addModifiers(Modifier.PUBLIC)
            .addSuperinterface(this.dispatcherInterface)
            .addMethod(callListenersBuilder.build())
            .addMethod(eventType)
            .build();

         String packageName = "ftbsc.geb.generated";
         JavaFile javaFile = JavaFile.builder(packageName, clazz).build();
         String resultingClassName = String.format("%s.%s", packageName, clazzName);

         try {
            JavaFileObject injectorFile = this.processingEnv.getFiler().createSourceFile(resultingClassName);
            PrintWriter out = new PrintWriter(injectorFile.openWriter());
            javaFile.writeTo(out);
            out.close();
         } catch(IOException e) {
            throw new RuntimeException(e);
         }

         this.generatedClasses.add(resultingClassName);
      });
   }

   /**
    * Generates the Service Provider file for the dispatchers.
    */
   private void generateServiceProvider() {
      try {
         FileObject serviceProvider = processingEnv.getFiler().createResource(StandardLocation.CLASS_OUTPUT, "",
            "META-INF/services/ftbsc.geb.api.IEventDispatcher");
         PrintWriter out = new PrintWriter(serviceProvider.openWriter());
         this.generatedClasses.forEach(out::println);
         out.close();
      } catch(IOException e) {
         throw new RuntimeException(e);
      }
   }

   /**
    * A container class to carry information about a listener class.
    */
   private static class ListenerContainer {
      /**
       * The actual listener, the annotated method.
       */
      public final ExecutableElement method;

      /**
       * The parent which implements {@link IListener}.
       */
      public final TypeMirror parent;

      /**
       * The {@link Listen} annotation on the method.
       */
      public final Listen annotation;

      /**
       * The public constructor.
       * @param method the annotated method, assumed to be valid
       *               and already checked
       */
      public ListenerContainer(ExecutableElement method) {
         this.method = method;
         this.parent = method.getEnclosingElement().asType();
         this.annotation = method.getAnnotation(Listen.class);
      }
   }
}