MixinEngine.java

package xyz.phanta.forgeveryfast.coremod.engine;

import net.minecraft.launchwrapper.IClassTransformer;
import net.minecraft.launchwrapper.LaunchClassLoader;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Handle;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.*;
import xyz.phanta.forgeveryfast.coremod.engine.injector.MixinInjector;

import javax.annotation.Nullable;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Iterator;
import java.util.ListIterator;
import java.util.Map;
import java.util.regex.Pattern;

public class MixinEngine implements IClassTransformer {

    private static final String T_MIXINJECT = "Lxyz/phanta/forgeveryfast/coremod/engine/Mixinject;";

    @SuppressWarnings("NotNullFieldNotInitialized")
    private static LaunchClassLoader classLoader;
    private static final Map<String, MixinTypeEntry> mixinCodeCache = new HashMap<>();

    public static byte[] inject(byte[] targetCode, String mixinClassName) {
        MixinTypeEntry mixin = loadMixinClass(mixinClassName);
        if (mixin == null) {
            System.out.println("FVF -- No definition found for mixin class: " + mixinClassName);
            return targetCode;
        }
        try {
            ClassNode target = new ClassNode();
            ClassReader targetReader = new ClassReader(targetCode);
            targetReader.accept(target, 0);
            mixin.modify(target);
            ClassWriter targetWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
            target.accept(targetWriter);
            System.out.println("FVF -- Injected mixin class: " + mixin.name);
            byte[] x = targetWriter.toByteArray();
            Files.write(Paths.get("mixin_debug/" + target.name.substring(target.name.lastIndexOf('/') + 1) + ".class"), x);
            return x;
        } catch (Exception e) {
            System.out.println("FVF -- Exception encountered while injecting mixin class: " + mixin.name);
            e.printStackTrace(System.out);
            return targetCode;
        }
    }

    public static void injectClassLoader(LaunchClassLoader classLoader) {
        MixinEngine.classLoader = classLoader;
    }

    @Nullable
    private static MixinTypeEntry loadMixinClass(String name) {
        MixinTypeEntry mixin = mixinCodeCache.get(name);
        if (mixin == null) {
            try {
                byte[] code = classLoader.getClassBytes(name);
                ClassNode mixinClass = new ClassNode();
                ClassReader mixinReader = new ClassReader(code);
                mixinReader.accept(mixinClass, 0);
                mixin = new MixinTypeEntry(name, mixinClass);
                mixinCodeCache.put(name, mixin);
            } catch (Exception e) {
                System.out.println("FVF -- Exception encountered while loading mixin class: " + name);
                e.printStackTrace(System.out);
            }
        }
        return mixin;
    }

    @Override
    public byte[] transform(String name, String transformedName, byte[] code) {
        MixinTypeEntry mixin = mixinCodeCache.get(transformedName);
        return mixin != null ? mixin.classCode : code;
    }

    private static class MixinTypeEntry {

        final String name;
        final ClassNode mixinClass;
        final byte[] classCode;
        final Map<String, MixinSite> injectionSites = new HashMap<>();

        MixinTypeEntry(String name, ClassNode mixinClass) {
            this.name = name;
            this.mixinClass = mixinClass;
            ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
            mixinClass.accept(writer);
            this.classCode = writer.toByteArray();
            for (MethodNode method : mixinClass.methods) {
                if (method.invisibleAnnotations != null) {
                    InjectSite siteType = null;
                    String args = "";
                    String targetName = "";
                    int ignoreLocals = -1;
                    for (AnnotationNode annot : method.invisibleAnnotations) {
                        if (annot.desc.equals(T_MIXINJECT)) {
                            for (int i = 0; i < annot.values.size(); i += 2) {
                                switch ((String)annot.values.get(i)) {
                                    case "value":
                                        siteType = InjectSite.valueOf(((String[])annot.values.get(i + 1))[1]);
                                        break;
                                    case "args":
                                        args = (String)annot.values.get(i + 1);
                                        break;
                                    case "target":
                                        targetName = (String)annot.values.get(i + 1);
                                        break;
                                    case "ignoreLocals":
                                        ignoreLocals = (Integer)annot.values.get(i + 1);
                                        break;
                                }
                            }
                            break;
                        }
                    }
                    if (siteType != null) {
                        injectionSites.put(
                                targetName.isEmpty() ? getMethodKey(method) : getMethodKey(targetName, method.desc),
                                new MixinSite(siteType.createInjector(args), mixinClass, method, ignoreLocals));
                    }
                }
            }
        }

        void modify(ClassNode targetClass) {
            ListIterator<MethodNode> iter = targetClass.methods.listIterator();
            while (iter.hasNext()) {
                MethodNode method = iter.next();
                MixinSite site = injectionSites.get(getMethodKey(method));
                if (site != null) {
                    method = site.inject(method);
                    refactorMemberOwners(targetClass, method.instructions);
                    iter.set(method);
                }
            }
            for (MethodNode method : mixinClass.methods) {
                if ((method.access & Opcodes.ACC_SYNTHETIC) != 0) {
                    refactorMemberOwners(targetClass, method.instructions);
                    targetClass.methods.add(method);
                }
            }
        }

        private void refactorMemberOwners(ClassNode targetClass, InsnList insns) {
            Iterator<AbstractInsnNode> iter = insns.iterator();
            while (iter.hasNext()) {
                AbstractInsnNode insn = iter.next();
                if (insn instanceof FieldInsnNode) {
                    FieldInsnNode fieldInsn = (FieldInsnNode)insn;
                    if (fieldInsn.owner.equals(mixinClass.name)) {
                        fieldInsn.owner = targetClass.name;
                    }
                } else if (insn instanceof MethodInsnNode) {
                    MethodInsnNode methodInsn = (MethodInsnNode)insn;
                    if (methodInsn.owner.equals(mixinClass.name)) {
                        methodInsn.owner = targetClass.name;
                    }
                } else if (insn instanceof InvokeDynamicInsnNode) {
                    InvokeDynamicInsnNode dynvokeInsn = (InvokeDynamicInsnNode)insn;
                    String mixinClassPattern = Pattern.quote(mixinClass.name);
                    dynvokeInsn.desc = dynvokeInsn.desc.replaceAll(mixinClassPattern, targetClass.name);
                    for (int i = 0; i < dynvokeInsn.bsmArgs.length; i++) {
                        Object arg = dynvokeInsn.bsmArgs[i];
                        if (arg instanceof Handle) {
                            Handle handle = (Handle)arg;
                            if (handle.getOwner().equals(mixinClass.name)) {
                                dynvokeInsn.bsmArgs[i] = new Handle(handle.getTag(), targetClass.name,
                                        handle.getName(), handle.getDesc(), handle.isInterface());
                            }
                        }
                    }
                } else if (insn instanceof LabelNode) {
                    LabelNode labelInsn = (LabelNode)insn;
                    labelInsn.resetLabel();
                }
            }
        }

        private static String getMethodKey(MethodNode method) {
            return getMethodKey(method.name, method.desc);
        }

        private static String getMethodKey(String name, String desc) {
            return name + desc.substring(0, desc.lastIndexOf(')'));
        }

        private static class MixinSite {

            final MixinInjector injectionStrategy;
            final ClassNode mixinClass;
            final MethodNode mixinMethod;

            private MixinSite(MixinInjector injectionStrategy, ClassNode mixinClass, MethodNode mixinMethod, int ignoreLocals) {
                this.injectionStrategy = injectionStrategy;
                this.mixinClass = mixinClass;
                this.mixinMethod = mixinMethod;
                if (ignoreLocals > -1) {
                    processLocalRemoval(ignoreLocals);
                }
            }

            private void processLocalRemoval(int ignoreLocals) {
                int maxIgnoredLine = ((LineNumberNode)mixinMethod.instructions.getFirst()).line + ignoreLocals;
                Iterator<AbstractInsnNode> iter = mixinMethod.instructions.iterator();
                while (iter.hasNext()) {
                    AbstractInsnNode insn = iter.next();
                    if (insn instanceof LineNumberNode && ((LineNumberNode)insn).line >= maxIgnoredLine) {
                        break;
                    }
                    iter.remove();
                }
            }

            MethodNode inject(MethodNode targetMethod) {
                return injectionStrategy.inject(targetMethod, mixinMethod);
            }

        }

    }

}