Generate computer methods at compile time with an annotation processor (#7829)

Author: thiakil

.gitignore

@@ -22,9 +22,5 @@ MDK.zip
 *.blend2
 /bin
 /logs
-/src/datagen/*/resources
-/src/datagen/*/annotation_generated
-!/src/datagen/main/resources/PersistingMekanismDatagenCache.js
-!/src/datagen/main/resources/META-INF/coremods.json
 /annotation-processor/bin
 libs/

annotation-processor/build.gradle

@@ -3,7 +3,7 @@ plugins {
 }
 
 group 'mekanism.annotation-processor'
-version '1.1.1'
+version '2.0.0'
 java.toolchain.languageVersion = JavaLanguageVersion.of("${java_version}")
 
 repositories {
@@ -24,5 +24,6 @@ repositories {
 dependencies {
     implementation "com.blamejared.crafttweaker:CraftTweaker-forge-${minecraft_version}:${crafttweaker_version}"
     //Version of GSON used by vanilla (and thus packed and already downloaded)
-    implementation group: 'com.google.code.gson', name: 'gson', version: '2.8.9'
-}
+    implementation group: 'com.google.code.gson', name: 'gson', version: '2.10'
+    implementation "com.squareup:javapoet:1.13.0"
+}

annotation-processor/src/main/java/mekanism/AnnotationParamScanner.java

@@ -1,12 +1,33 @@
 package mekanism;
 
+import javax.lang.model.element.Element;
 import java.lang.annotation.Annotation;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Set;
 
-public record AnnotationParamScanner(String targetFile, Set<Class<? extends Annotation>> supportedAnnotations) {
+public record AnnotationParamScanner(String targetFile, Set<Class<? extends Annotation>> supportedAnnotations, List<Element> originatingElements) {
 
     @SafeVarargs
     public AnnotationParamScanner(String targetFile, Class<? extends Annotation>... supportedAnnotations) {
-        this(targetFile, Set.of(supportedAnnotations));
+        this(targetFile, Set.of(supportedAnnotations), new ArrayList<>());
     }
-}
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+
+        return (o instanceof AnnotationParamScanner that) &&
+               targetFile.equals(that.targetFile) &&
+               supportedAnnotations.equals(that.supportedAnnotations);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = targetFile.hashCode();
+        result = 31 * result + supportedAnnotations.hashCode();
+        return result;
+    }
+}

annotation-processor/src/main/java/mekanism/ComputerMethodProcessor.java

@@ -0,0 +1,63 @@
+package mekanism;
+
+import com.squareup.javapoet.JavaFile;
+import mekanism.builder.ComputerHandlerBuilder;
+
+import javax.annotation.processing.AbstractProcessor;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.annotation.processing.RoundEnvironment;
+import javax.annotation.processing.SupportedAnnotationTypes;
+import javax.annotation.processing.SupportedSourceVersion;
+import javax.lang.model.SourceVersion;
+import javax.lang.model.element.Element;
+import javax.lang.model.element.TypeElement;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Annotation processor to generate ComputerMethodFactory subclasses for computer methods.
+ * Must only produce files directly related to the annotated elements, and not based on other elements.
+ * This is due to being marked as an Isolating processor in Gradle.
+ * <p>
+ * Technically we violate this with the Wrapping methods, but if new methods are added to a wrapper,
+ * a manually triggered full rebuild should catch it.
+ */
+@SupportedSourceVersion(SourceVersion.RELEASE_17)
+@SupportedAnnotationTypes({
+        MekAnnotationProcessors.ANNOTATION_COMPUTER_METHOD,
+        MekAnnotationProcessors.ANNOTATION_SYNTHETIC_COMPUTER_METHOD,
+        MekAnnotationProcessors.ANNOTATION_WRAPPING_COMPUTER_METHOD
+})
+public class ComputerMethodProcessor extends AbstractProcessor {
+    @Override
+    public synchronized void init(ProcessingEnvironment processingEnv) {
+        super.init(processingEnv);
+        ComputerHandlerBuilder.init(processingEnv.getElementUtils(), processingEnv.getTypeUtils());
+    }
+
+    @Override
+    public boolean process(Set<? extends TypeElement> annotatedTypes, RoundEnvironment roundEnvironment) {
+        //map annotated elements to multimap by enclosing class
+        Map<TypeElement, List<Element>> annotatedElementsByParent = new HashMap<>();
+        for (Element element : roundEnvironment.getElementsAnnotatedWithAny(annotatedTypes.toArray(new TypeElement[0]))) {
+            annotatedElementsByParent.computeIfAbsent((TypeElement) element.getEnclosingElement(), i -> new ArrayList<>()).add(element);
+        }
+
+        annotatedElementsByParent.forEach(this::processTypeWithAnnotations);
+
+        return true;
+    }
+
+    private void processTypeWithAnnotations(TypeElement containingType, List<Element> annotatedElements) {
+        JavaFile factoryFile = new ComputerHandlerBuilder(containingType, processingEnv).build(annotatedElements);
+        try {
+            factoryFile.writeTo(processingEnv.getFiler());
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

annotation-processor/src/main/java/mekanism/MekAnnotationProcessors.java

@@ -0,0 +1,20 @@
+package mekanism;
+
+import com.squareup.javapoet.ClassName;
+
+public class MekAnnotationProcessors {
+    public static final String MODULE_OPTION = "mekanismModule";
+
+    //Packages
+    public static final String COMPUTER_INTEGRATION_PACKAGE = "mekanism.common.integration.computer";
+    public static final String ANNOTATION_PACKAGE = COMPUTER_INTEGRATION_PACKAGE + ".annotation";
+
+    //annotations
+    public static final String COMPUTER_METHOD_FACTORY_ANNOTATION_CLASSNAME = ANNOTATION_PACKAGE + ".MethodFactory";
+    public static final String ANNOTATION_COMPUTER_METHOD = ANNOTATION_PACKAGE + ".ComputerMethod";
+    public static final String ANNOTATION_SYNTHETIC_COMPUTER_METHOD = ANNOTATION_PACKAGE + ".SyntheticComputerMethod";
+    public static final String ANNOTATION_WRAPPING_COMPUTER_METHOD = ANNOTATION_PACKAGE + ".WrappingComputerMethod";
+
+    //class names
+    public static final ClassName COMPUTER_METHOD_FACTORY_ANNOTATION = ClassName.bestGuess(COMPUTER_METHOD_FACTORY_ANNOTATION_CLASSNAME);
+}

annotation-processor/src/main/java/mekanism/MethodFactoryProcessor.java

@@ -0,0 +1,136 @@
+package mekanism;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.CodeBlock;
+import com.squareup.javapoet.JavaFile;
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.TypeSpec;
+import java.io.Writer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import javax.annotation.processing.AbstractProcessor;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.annotation.processing.RoundEnvironment;
+import javax.annotation.processing.SupportedAnnotationTypes;
+import javax.annotation.processing.SupportedOptions;
+import javax.annotation.processing.SupportedSourceVersion;
+import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.Element;
+import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.Modifier;
+import javax.lang.model.element.TypeElement;
+import javax.tools.StandardLocation;
+import mekanism.visitors.AnnotationHelper;
+
+import javax.lang.model.SourceVersion;
+import javax.lang.model.type.TypeKind;
+import javax.lang.model.type.TypeMirror;
+import javax.lang.model.util.Types;
+import javax.tools.Diagnostic;
+import java.io.IOException;
+
+/**
+ * Gathering (Gradle) annotation processor which generates a ComputerMethodRegistry for the Factories generated
+ * by the {@link ComputerMethodProcessor} processor.
+ */
+@SupportedSourceVersion(SourceVersion.RELEASE_17)
+@SupportedAnnotationTypes(MekAnnotationProcessors.COMPUTER_METHOD_FACTORY_ANNOTATION_CLASSNAME)
+@SupportedOptions(MekAnnotationProcessors.MODULE_OPTION)
+public class MethodFactoryProcessor extends AbstractProcessor {
+    private String mekModule;
+    private final ClassName factoryRegistry = ClassName.get(MekAnnotationProcessors.COMPUTER_INTEGRATION_PACKAGE, "FactoryRegistry");
+    private final ClassName methodRegistryInterface = ClassName.get(MekAnnotationProcessors.COMPUTER_INTEGRATION_PACKAGE, "IComputerMethodRegistry");
+    private final MethodSpec.Builder registryInit = MethodSpec.methodBuilder("register")
+          .addModifiers(Modifier.PUBLIC)
+          .addAnnotation(Override.class);
+
+    @Override
+    public synchronized void init(ProcessingEnvironment processingEnv) {
+        super.init(processingEnv);
+        mekModule = processingEnv.getOptions().getOrDefault(MekAnnotationProcessors.MODULE_OPTION, "value_not_supplied");
+    }
+
+    @Override
+    public boolean process(Set<? extends TypeElement> annotatedTypes, RoundEnvironment roundEnvironment) {
+        TypeMirror methodFactoryType = processingEnv.getElementUtils().getTypeElement(MekAnnotationProcessors.COMPUTER_METHOD_FACTORY_ANNOTATION_CLASSNAME).asType();
+        TypeSpec.Builder registryType = TypeSpec.classBuilder("ComputerMethodRegistry_" + mekModule)
+              .addModifiers(Modifier.PUBLIC)
+              .addSuperinterface(methodRegistryInterface);
+
+        //this should only ever be 1 annotation
+        for (Element element : roundEnvironment.getElementsAnnotatedWithAny(annotatedTypes.toArray(new TypeElement[0]))) {
+            if (element instanceof TypeElement factoryTypeEl) {
+                //get the annotation mirror for @MethodFactory
+                AnnotationMirror annotationMirror = null;
+                for (AnnotationMirror am : factoryTypeEl.getAnnotationMirrors()) {
+                    if (typeUtils().isSameType(am.getAnnotationType(), methodFactoryType)) {
+                        annotationMirror = am;
+                        break;
+                    }
+                }
+                if (annotationMirror == null) {
+                    processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, "Couldn't find annotation mirror", factoryTypeEl);
+                    continue;
+                }
+                registryType.addOriginatingElement(factoryTypeEl);
+                AnnotationHelper helper = new AnnotationHelper(processingEnv.getElementUtils(), annotationMirror);
+                addHandlerToRegistry((TypeElement) typeUtils().asElement(helper.getClassValue("target")), ClassName.get(factoryTypeEl));
+            }
+        }
+
+        if (!registryType.originatingElements.isEmpty()) {
+            registryType.addMethod(registryInit.build());
+            TypeSpec registrySpec = registryType.build();
+            String packageName = "mekanism.generated." + mekModule;
+            try {
+                JavaFile.builder(packageName, registrySpec).build().writeTo(processingEnv.getFiler());
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            try(Writer serviceWriter = processingEnv.getFiler().createResource(StandardLocation.CLASS_OUTPUT, "","META-INF/services/"+methodRegistryInterface.canonicalName()).openWriter()) {
+                serviceWriter.write(packageName+"."+registrySpec.name);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+        return false;
+    }
+
+    /**
+     * Gather superclasses for handledType and add a register call
+     * @param handledType the subject type of the handler
+     * @param factoryClassName the factory's class name
+     */
+    private void addHandlerToRegistry(TypeElement handledType, ClassName factoryClassName) {
+        //gather all superclasses (in mekanism package)
+        List<ClassName> superClasses = new ArrayList<>();
+        TypeMirror superClass = handledType.getSuperclass();
+        TypeElement superTypeElement;
+        do {
+            superTypeElement = (TypeElement) typeUtils().asElement(superClass);
+            if (superTypeElement == null) {
+                break;
+            }
+            ClassName clazz = ClassName.get(superTypeElement);
+            if (clazz.canonicalName().startsWith("mekanism")) {
+                superClasses.add(0, clazz);
+            }
+        } while ((superClass = superTypeElement.getSuperclass()).getKind() != TypeKind.NONE);
+
+        //add register call to the factory
+        String registerName = handledType.getKind() == ElementKind.INTERFACE ? "registerInterface" : "register";
+        CodeBlock.Builder registerStatement = CodeBlock.builder()
+                .add("$T.$L($T.class, $T::new", factoryRegistry, registerName, typeUtils().erasure(handledType.asType()), factoryClassName);
+        //add all super classes, so we don't have to calculate at runtime
+        for (ClassName cls : superClasses) {
+            registerStatement.add(", $T.class", cls);
+        }
+        registerStatement.add(")");
+        registryInit.addStatement(registerStatement.build());
+    }
+
+    private Types typeUtils() {
+        return processingEnv.getTypeUtils();
+    }
+}