Add indigo support to custom block textures

6 files changed, 383 insertions, 0 deletions

diff --git a/javaplugin/src/main/java/moe/nea/firmament/javaplugin/InitReplacer.java b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/InitReplacer.java
new file mode 100644
index 0000000..a0d28ab
--- /dev/null
+++ b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/InitReplacer.java
@@ -0,0 +1,73 @@
+package moe.nea.firmament.javaplugin;
+
+import com.sun.source.tree.ClassTree;
+import com.sun.source.tree.CompilationUnitTree;
+import com.sun.source.tree.VariableTree;
+import com.sun.source.util.TreeScanner;
+import com.sun.tools.javac.code.Symbol;
+import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.tree.TreeMaker;
+import com.sun.tools.javac.util.List;
+import com.sun.tools.javac.util.Names;
+
+public class InitReplacer extends TreeScanner<Void, Void> {
+    private final MappingTree mappingTree;
+    private final TreeMaker treeMaker;
+    private final Names names;
+    private final IntermediaryNameResolutionTask plugin;
+    private Symbol.ClassSymbol classTree;
+    private CompilationUnitTree compilationUnitTree;
+
+    public InitReplacer(MappingTree mappingTree, IntermediaryNameResolutionTask plugin) {
+        this.mappingTree = mappingTree;
+        this.treeMaker = plugin.treeMaker;
+        this.names = plugin.names;
+        this.plugin = plugin;
+    }
+
+    @Override
+    public Void visitClass(ClassTree node, Void unused) {
+        this.classTree = plugin.utils.getSymbol(node);
+        return super.visitClass(node, unused);
+    }
+
+    @Override
+    public Void visitCompilationUnit(CompilationUnitTree node, Void unused) {
+        this.compilationUnitTree = node;
+        return super.visitCompilationUnit(node, unused);
+    }
+
+    @Override
+    public Void visitVariable(VariableTree node, Void unused) {
+        var annotation = node
+            .getModifiers().getAnnotations()
+            .stream()
+            .filter(it -> it.getAnnotationType().toString().equals("IntermediaryName")) // Crazy type-safety!
+            .findAny();
+        if (annotation.isEmpty())
+            return super.visitVariable(node, unused);
+        var jcAnnotation = (JCTree.JCAnnotation) annotation.get();
+        var jcNode = (JCTree.JCVariableDecl) node;
+        if (node.getInitializer() != null) {
+            plugin.utils.reportError(
+                compilationUnitTree.getSourceFile(),
+                jcNode.getInitializer(),
+                "Providing an initializer for a variable is illegal for @IntermediaryName annotated fields"
+            );
+            return super.visitVariable(node, unused);
+        }
+        var target = plugin.utils.getAnnotationValue(jcAnnotation, "value");
+        var targetClass = plugin.utils.resolveClassLiteralExpression(target).tsym.flatName().toString();
+        var intermediaryClass = mappingTree.resolveClassToIntermediary(targetClass);
+        var remapper = treeMaker.Select(treeMaker.This(classTree.type), names.fromString("remapper"));
+        var remappingCall = treeMaker.Apply(
+            List.nil(),
+            treeMaker.Select(remapper, names.fromString("mapClassName")),
+            List.of(treeMaker.Literal("intermediary"),
+                    treeMaker.Literal(intermediaryClass)));
+        jcNode.init = remappingCall;
+        jcNode.mods.annotations = List.filter(jcNode.mods.annotations, jcAnnotation);
+        return super.visitVariable(node, unused);
+    }
+
+}
diff --git a/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryMethodReplacer.java b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryMethodReplacer.java
new file mode 100644
index 0000000..d3040b7
--- /dev/null
+++ b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryMethodReplacer.java
@@ -0,0 +1,72 @@
+package moe.nea.firmament.javaplugin;
+
+import com.sun.source.tree.CompilationUnitTree;
+import com.sun.source.tree.ExpressionTree;
+import com.sun.source.tree.MethodInvocationTree;
+import com.sun.source.util.TreeScanner;
+import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.util.List;
+
+import javax.tools.JavaFileObject;
+
+public class IntermediaryMethodReplacer extends TreeScanner<Void, Void> {
+    private final MappingTree mappings;
+    private final IntermediaryNameResolutionTask plugin;
+    private JavaFileObject sourceFile;
+    private CompilationUnitTree compilationUnit;
+
+    public IntermediaryMethodReplacer(MappingTree mappings, IntermediaryNameResolutionTask plugin) {
+        this.mappings = mappings;
+        this.plugin = plugin;
+    }
+
+
+    @Override
+    public Void visitCompilationUnit(CompilationUnitTree node, Void unused) {
+        sourceFile = node.getSourceFile();
+        compilationUnit = node;
+        return super.visitCompilationUnit(node, unused);
+    }
+
+    public void replaceMethodName(JCTree.JCMethodInvocation node) {
+        var select = node.getMethodSelect();
+        if (!(select instanceof JCTree.JCFieldAccess fieldAccess)) return;
+        if (!fieldAccess.name.contentEquals("methodName")) return;
+        if (!(node.args.head instanceof JCTree.JCMemberReference methodReference)) {
+            plugin.utils.reportError(sourceFile, node, "Please provide a Class::method reference directly (and nothing else)");
+            return;
+        }
+        var clearName = methodReference.name.toString();
+        var classRef = methodReference.expr;
+        var type = plugin.utils.resolveClassName(classRef, compilationUnit);
+        var intermediaryName = mappings.resolveMethodToIntermediary(
+            type.tsym.flatName().toString(),
+            clearName
+        );
+        fieldAccess.name = plugin.names.fromString("id");
+        node.args = List.of(plugin.treeMaker.Literal(intermediaryName));
+    }
+
+    public void replaceClassName(JCTree.JCMethodInvocation node) {
+        var select = node.getMethodSelect();
+        if (!(select instanceof JCTree.JCFieldAccess fieldAccess)) return;
+        if (!fieldAccess.name.contentEquals("className")) return;
+        if (node.getTypeArguments().size() != 1) {
+            plugin.utils.reportError(sourceFile, node, "You need to explicitly provide the class you want the intermediary name for");
+            return;
+        }
+        var head = node.typeargs.head;
+        var resolved = plugin.utils.resolveClassName(head, compilationUnit);
+        var mappedName = mappings.resolveClassToIntermediary(resolved.tsym.flatName().toString());
+        fieldAccess.name = plugin.names.fromString("id");
+        node.typeargs = List.nil();
+        node.args = List.of(plugin.treeMaker.Literal(mappedName));
+    }
+
+    @Override
+    public Void visitMethodInvocation(MethodInvocationTree node, Void unused) {
+        replaceClassName((JCTree.JCMethodInvocation) node);
+        replaceMethodName((JCTree.JCMethodInvocation) node);
+        return super.visitMethodInvocation(node, unused);
+    }
+}
diff --git a/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryNameResolutionPlugin.java b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryNameResolutionPlugin.java
new file mode 100644
index 0000000..ba6a0c5
--- /dev/null
+++ b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryNameResolutionPlugin.java
@@ -0,0 +1,25 @@
+package moe.nea.firmament.javaplugin;
+
+import com.sun.source.util.JavacTask;
+import com.sun.source.util.Plugin;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class IntermediaryNameResolutionPlugin implements Plugin {
+
+    @Override
+    public String getName() {
+        return "IntermediaryNameReplacement";
+    }
+
+    @Override
+    public void init(JavacTask task, String... args) {
+        Map<String, String> argMap = new HashMap<>();
+        for (String arg : args) {
+            String[] parts = arg.split("=", 2);
+            argMap.put(parts[0], parts.length == 2 ? parts[1] : "true");
+        }
+        task.addTaskListener(new IntermediaryNameResolutionTask(this, task, argMap));
+    }
+}
diff --git a/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryNameResolutionTask.java b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryNameResolutionTask.java
new file mode 100644
index 0000000..86a5598
--- /dev/null
+++ b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/IntermediaryNameResolutionTask.java
@@ -0,0 +1,44 @@
+package moe.nea.firmament.javaplugin;
+
+import com.sun.source.util.JavacTask;
+import com.sun.source.util.TaskEvent;
+import com.sun.source.util.TaskListener;
+import com.sun.tools.javac.api.BasicJavacTask;
+import com.sun.tools.javac.tree.TreeMaker;
+import com.sun.tools.javac.util.Names;
+import net.fabricmc.stitch.commands.tinyv2.TinyV2Reader;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Map;
+
+public class IntermediaryNameResolutionTask implements TaskListener {
+    TreeMaker treeMaker;
+    Names names;
+    MappingTree mappings;
+    Utils utils;
+
+    public IntermediaryNameResolutionTask(IntermediaryNameResolutionPlugin intermediaryNameResolutionPlugin, JavacTask task, Map<String, String> argMap) {
+        var context = ((BasicJavacTask) task).getContext();
+        var mappingFile = new File(argMap.get("mappingFile"));
+        System.err.println("Loading mappings from " + mappingFile);
+        try {
+            var tinyV2File = TinyV2Reader.read(mappingFile.toPath());
+            mappings = new MappingTree(tinyV2File, argMap.get("sourceNs"), argMap.getOrDefault("targetNs", "intermediary"));
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        treeMaker = TreeMaker.instance(context);
+        names = Names.instance(context);
+        utils = Utils.instance(context);
+    }
+
+    @Override
+    public void finished(TaskEvent e) {
+        if (e.getKind() != TaskEvent.Kind.ENTER) return;
+        if (e.getCompilationUnit() == null || e.getSourceFile() == null) return;
+        e.getCompilationUnit().accept(new InitReplacer(mappings, this), null);
+        e.getCompilationUnit().accept(new IntermediaryMethodReplacer(mappings, this), null);
+    }
+
+}
diff --git a/javaplugin/src/main/java/moe/nea/firmament/javaplugin/MappingTree.java b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/MappingTree.java
new file mode 100644
index 0000000..7a270b7
--- /dev/null
+++ b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/MappingTree.java
@@ -0,0 +1,48 @@
+package moe.nea.firmament.javaplugin;
+
+import net.fabricmc.stitch.commands.tinyv2.TinyClass;
+import net.fabricmc.stitch.commands.tinyv2.TinyFile;
+import net.fabricmc.stitch.commands.tinyv2.TinyMethod;
+
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class MappingTree {
+
+    private final Map<String, TinyClass> classLookup;
+    private final int targetIndex;
+    private final int sourceIndex;
+
+    public MappingTree(TinyFile tinyV2File, String sourceNamespace, String targetNamespace) {
+        sourceIndex = tinyV2File.getHeader().getNamespaces().indexOf(sourceNamespace);
+        if (sourceIndex < 0)
+            throw new RuntimeException("Could not find source namespace " + sourceNamespace + " in mappings file.");
+        this.classLookup = tinyV2File
+            .getClassEntries()
+            .stream()
+            .collect(Collectors.toMap(it -> it.getClassNames().get(sourceIndex), it -> it));
+        targetIndex = tinyV2File.getHeader().getNamespaces().indexOf(targetNamespace);
+        if (targetIndex < 0)
+            throw new RuntimeException("Could not find target namespace " + targetNamespace + " in mappings file.");
+    }
+
+    public String resolveMethodToIntermediary(String className, String methodName) {
+        var classData = classLookup.get(className.replace(".", "/"));
+        TinyMethod candidate = null;
+        for (TinyMethod method : classData.getMethods()) {
+            if (method.getMethodNames().get(sourceIndex).equals(methodName)) {
+                if (candidate != null) {
+                    throw new RuntimeException("Found two candidates for method " + className + "." + methodName);
+                }
+                candidate = method;
+            }
+        }
+        return candidate.getMethodNames().get(targetIndex);
+    }
+
+    public String resolveClassToIntermediary(String className) {
+        return classLookup.get(className.replace(".", "/"))
+                          .getClassNames().get(targetIndex)
+                          .replace("/", ".");
+    }
+}
diff --git a/javaplugin/src/main/java/moe/nea/firmament/javaplugin/Utils.java b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/Utils.java
new file mode 100644
index 0000000..d9008bd
--- /dev/null
+++ b/javaplugin/src/main/java/moe/nea/firmament/javaplugin/Utils.java
@@ -0,0 +1,121 @@
+package moe.nea.firmament.javaplugin;
+
+import com.sun.source.tree.AnnotationTree;
+import com.sun.source.tree.AssignmentTree;
+import com.sun.source.tree.ClassTree;
+import com.sun.source.tree.CompilationUnitTree;
+import com.sun.source.tree.ExpressionTree;
+import com.sun.source.tree.IdentifierTree;
+import com.sun.source.tree.MemberSelectTree;
+import com.sun.source.tree.Tree;
+import com.sun.tools.javac.comp.Attr;
+import com.sun.tools.javac.code.Symbol;
+import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.comp.AttrContext;
+import com.sun.tools.javac.comp.Enter;
+import com.sun.tools.javac.comp.Env;
+import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.util.Context;
+import com.sun.tools.javac.util.JCDiagnostic;
+import com.sun.tools.javac.util.JavacMessages;
+import com.sun.tools.javac.util.Log;
+
+import javax.tools.JavaFileObject;
+import java.util.ListResourceBundle;
+
+public class Utils {
+    private static final Context.Key<Utils> KEY = new Context.Key<>();
+    private final Log log;
+    private final JCDiagnostic.Factory diagnostics;
+    private final Types types;
+    private final Attr attr;
+    private final Enter enter;
+
+    private Utils(Context context) {
+        context.put(KEY, this);
+        JavacMessages.instance(context).add(l -> new ListResourceBundle() {
+
+            @Override
+            protected Object[][] getContents() {
+                return new Object[][]{
+                    new Object[]{"compiler.err.firmament.generic", "{0}"}
+                };
+            }
+        });
+        log = Log.instance(context);
+        diagnostics = JCDiagnostic.Factory.instance(context);
+        types = Types.instance(context);
+        attr = Attr.instance(context);
+        enter = Enter.instance(context);
+    }
+
+    public static Utils instance(Context context) {
+        var utils = context.get(KEY);
+        if (utils == null) {
+            utils = new Utils(context);
+        }
+        return utils;
+    }
+
+    public Type resolveClassName(ExpressionTree expression) {
+        var tree = (JCTree) expression;
+        return tree.type;
+    }
+
+    public Type resolveClassName(ExpressionTree tree, CompilationUnitTree unit) {
+        return resolveClassName(tree, enter.getTopLevelEnv((JCTree.JCCompilationUnit) unit));
+    }
+
+    public Type resolveClassName(ExpressionTree tree, Env<AttrContext> env) {
+        var t = resolveClassName(tree);
+        if (t != null) return t;
+        return attr.attribType((JCTree) tree, env);
+    }
+
+    public Symbol getSymbol(IdentifierTree tree) {
+        return ((JCTree.JCIdent) tree).sym;
+    }
+
+    public Symbol.ClassSymbol getSymbol(ClassTree tree) {
+        return ((JCTree.JCClassDecl) tree).sym;
+    }
+
+    public ExpressionTree getAnnotationValue(
+        AnnotationTree tree,
+        String name) {
+        // TODO: strip parenthesis
+        for (var argument : tree.getArguments()) {
+            var assignment = (AssignmentTree) argument;
+            if (((IdentifierTree) assignment.getVariable()).getName().toString().equals(name))
+                return assignment.getExpression();
+        }
+        return null;
+    }
+
+    public Type.ClassType resolveClassLiteralExpression(ExpressionTree tree) {
+        if (!(tree instanceof MemberSelectTree select))
+            throw new RuntimeException("Cannot resolve non field access class literal: " + tree);
+        if (!select.getIdentifier().toString().equals("class"))
+            throw new RuntimeException("Class literal " + select + "accessed non .class attribute");
+
+        return (Type.ClassType) resolveClassName(select.getExpression());
+    }
+
+    public void reportError(
+        JavaFileObject file,
+        Tree node,
+        String message
+    ) {
+        var originalSource = log.useSource(file);
+        var error = diagnostics.error(
+            JCDiagnostic.DiagnosticFlag.API,
+            log.currentSource(),
+            ((JCTree) node).pos(),
+            "firmament.generic",
+            message
+        );
+        log.report(error);
+        log.useSource(originalSource);
+    }
+}