New Mixin Fixer system

src/main/java/me/modmuss50/optifabric/compat/IMixinFixer.java

@@ -0,0 +1,83 @@
+package me.modmuss50.optifabric.compat;
+
+import com.google.common.collect.Lists;
+import me.modmuss50.optifabric.util.ASMUtils;
+import me.modmuss50.optifabric.util.MixinInternals;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+import org.objectweb.asm.tree.*;
+import org.spongepowered.asm.mixin.extensibility.IMixinInfo;
+import org.spongepowered.asm.mixin.transformer.ClassInfo;
+
+import java.lang.reflect.Modifier;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public interface IMixinFixer {
+	void fix(IMixinInfo mixinInfo, ClassNode mixinNode);
+
+	default int getIndex(MethodNode method, boolean afterCallback, boolean afterSequence, String... sequence) {
+		String desc = method.desc;
+		int offset = 0;
+		if (afterCallback) {
+			List<Type> params = Lists.newArrayList(Type.getArgumentTypes(desc));
+			for (Type type : params) {
+				offset++;
+				if (type.toString().startsWith("Lorg/spongepowered/asm/mixin/injection/callback/CallbackInfo")) {
+					break;
+				}
+			}
+			desc = params.subList(offset, params.size()).stream().map(Type::toString).collect(Collectors.joining(""));
+		}
+		if (afterSequence) offset++;
+		desc = desc.split(String.join("", sequence))[afterSequence ? 1 : 0];
+		if (!desc.contains("(")) desc = "(" + desc;
+		if (!desc.contains(")")) desc = desc + ")V";
+		return Type.getArgumentTypes(desc).length + offset;
+	}
+
+	default void insertParams(MethodNode method, IMixinInfo mixinInfo, int index, String... params) {
+		List<Type> newDesc = Arrays.stream(Type.getArgumentTypes(method.desc)).collect(Collectors.toList());
+		newDesc.addAll(index, Arrays.stream(params).map(Type::getType).collect(Collectors.toList()));
+		int shiftBy = 0;
+		for (String param : params) {
+			shiftBy++;
+			if (ASMUtils.isWideType(param)) shiftBy++;
+		}
+		method.maxLocals += shiftBy;
+
+		for (int i = 0; i < params.length; i++) {
+			method.parameters.add(index + i, new ParameterNode("syn_" + i, Opcodes.ACC_SYNTHETIC));
+		}
+
+		for (int i = index; i > 0; i--) {
+			if (ASMUtils.isWideType(newDesc.get(i))) {
+				index++;
+			}
+		}
+		if (!Modifier.isStatic(method.access)) index++;
+
+		//shift locals (not mandatory)
+		for (LocalVariableNode local : method.localVariables) {
+			if (local.index >= index) {
+				local.index += shiftBy;
+			}
+		}
+		//shift instructions
+		for (AbstractInsnNode insn : method.instructions) {
+			if (insn instanceof VarInsnNode && ((VarInsnNode) insn).var >= index) {
+				((VarInsnNode) insn).var += shiftBy;
+			} else if (insn instanceof IincInsnNode && ((IincInsnNode) insn).var >= index) {
+				((IincInsnNode) insn).var += shiftBy;
+			}
+		}
+
+		ClassInfo info = MixinInternals.getClassInfoFor(mixinInfo);
+		Set<ClassInfo.Method> methods = MixinInternals.getClassInfoMethods(info);
+		methods.removeIf(meth -> method.name.equals(meth.getOriginalName()) && method.desc.equals(meth.getOriginalDesc()));
+		method.desc = Type.getMethodDescriptor(Type.getReturnType(method.desc), newDesc.toArray(new Type[0]));
+		methods.add(info.new Method(method, true));
+	}
+}

src/main/java/me/modmuss50/optifabric/compat/MixinFixerExtension.java

@@ -0,0 +1,48 @@
+package me.modmuss50.optifabric.compat;
+
+import me.modmuss50.optifabric.util.MixinInternals;
+import org.apache.commons.lang3.tuple.Pair;
+import org.objectweb.asm.tree.ClassNode;
+import org.spongepowered.asm.mixin.MixinEnvironment;
+import org.spongepowered.asm.mixin.extensibility.IMixinInfo;
+import org.spongepowered.asm.mixin.transformer.ext.IExtension;
+import org.spongepowered.asm.mixin.transformer.ext.ITargetClassContext;
+
+import java.util.Collections;
+import java.util.Set;
+import java.util.WeakHashMap;
+
+public class MixinFixerExtension implements IExtension {
+	private static final Set<ClassNode> PREPARED_MIXINS = Collections.newSetFromMap(new WeakHashMap<>());
+
+	@Override
+	public boolean checkActive(MixinEnvironment environment) {
+		return true;
+	}
+
+	@Override
+	public void preApply(ITargetClassContext context) {
+		for (Pair<IMixinInfo, ClassNode> pair : MixinInternals.getMixinsFor(context)) {
+			prepareMixin(pair.getLeft(), pair.getRight());
+		}
+	}
+
+	@Override
+	public void postApply(ITargetClassContext context) {
+
+	}
+
+	@Override
+	public void export(MixinEnvironment env, String name, boolean force, ClassNode classNode) {
+
+	}
+
+	private static void prepareMixin(IMixinInfo mixinInfo, ClassNode mixinNode) {
+		if (PREPARED_MIXINS.contains(mixinNode)) {
+			// Don't scan the whole class again.
+			return;
+		}
+		ModMixinFixer.INSTANCE.getFixers(mixinInfo.getClassName()).forEach(transformer -> transformer.fix(mixinInfo, mixinNode));
+		PREPARED_MIXINS.add(mixinNode);
+	}
+}

src/main/java/me/modmuss50/optifabric/compat/ModMixinFixer.java

@@ -0,0 +1,20 @@
+package me.modmuss50.optifabric.compat;
+
+import java.util.*;
+
+public class ModMixinFixer {
+	public static final ModMixinFixer INSTANCE = new ModMixinFixer();
+
+	private final Map<String, List<IMixinFixer>> classFixes = new HashMap<>();
+
+	private ModMixinFixer() {
+	}
+
+	public void addFixer(String mixinClass, IMixinFixer fixer) {
+		classFixes.computeIfAbsent(mixinClass, s -> new ArrayList<>()).add(fixer);
+	}
+
+	public List<IMixinFixer> getFixers(String className) {
+		return classFixes.getOrDefault(className.replace('.', '/'), Collections.emptyList());
+	}
+}

src/main/java/me/modmuss50/optifabric/compat/OptifabricMixinPlugin.java

@@ -0,0 +1,20 @@
+package me.modmuss50.optifabric.compat;
+
+import me.modmuss50.optifabric.util.MixinInternals;
+import org.objectweb.asm.tree.ClassNode;
+import org.spongepowered.asm.mixin.extensibility.IMixinInfo;
+
+public class OptifabricMixinPlugin extends EmptyMixinPlugin {
+	@Override
+	public void onLoad(String mixinPackage) {
+		MixinInternals.registerExtension(new MixinFixerExtension());
+	}
+
+	@Override
+	public void preApply(String targetClassName, ClassNode targetClass, String mixinClassName, IMixinInfo mixinInfo) {
+	}
+
+	@Override
+	public void postApply(String targetClassName, ClassNode targetClass, String mixinClassName, IMixinInfo mixinInfo) {
+	}
+}

src/main/java/me/modmuss50/optifabric/util/ASMUtils.java

@@ -9,9 +9,18 @@
 import java.util.zip.ZipFile;
 
 import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.ClassNode;
 
 public class ASMUtils {
+	public static boolean isWideType(Type type) {
+		return type.equals(Type.LONG_TYPE) || type.equals(Type.DOUBLE_TYPE);
+	}
+
+	public static boolean isWideType(String type) {
+		return isWideType(Type.getType(type));
+	}
+
 	public static ClassNode readClass(byte[] bytes) {
 		return readClass(new ClassReader(Objects.requireNonNull(bytes, "Cannot read null class bytes")));
 	}

src/main/java/me/modmuss50/optifabric/util/MixinInternals.java

@@ -0,0 +1,104 @@
+package me.modmuss50.optifabric.util;
+
+import org.apache.commons.lang3.tuple.Pair;
+import org.objectweb.asm.tree.ClassNode;
+import org.spongepowered.asm.mixin.MixinEnvironment;
+import org.spongepowered.asm.mixin.extensibility.IMixinInfo;
+import org.spongepowered.asm.mixin.transformer.ClassInfo;
+import org.spongepowered.asm.mixin.transformer.IMixinTransformer;
+import org.spongepowered.asm.mixin.transformer.ext.Extensions;
+import org.spongepowered.asm.mixin.transformer.ext.IExtension;
+import org.spongepowered.asm.mixin.transformer.ext.ITargetClassContext;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.*;
+
+@SuppressWarnings("unchecked")
+public class MixinInternals {
+    private static final Field TARGET_CLASS_CONTEXT_MIXINS_FIELD;
+    private static final Method MIXIN_INFO_GET_STATE_METHOD;
+    private static final Field STATE_CLASS_NODE_FIELD;
+    private static final Field STATE_CLASS_INFO_FIELD;
+    private static final Field CLASS_INFO_METHODS_FIELD;
+    private static final Field EXTENSIONS_FIELD;
+    private static final Field ACTIVE_EXTENSIONS_FIELD;
+
+    static {
+        try {
+            Class<?> TargetClassContext = Class.forName("org.spongepowered.asm.mixin.transformer.TargetClassContext");
+            TARGET_CLASS_CONTEXT_MIXINS_FIELD = TargetClassContext.getDeclaredField("mixins");
+            TARGET_CLASS_CONTEXT_MIXINS_FIELD.setAccessible(true);
+            Class<?> MixinInfo = Class.forName("org.spongepowered.asm.mixin.transformer.MixinInfo");
+            MIXIN_INFO_GET_STATE_METHOD = MixinInfo.getDeclaredMethod("getState");
+            MIXIN_INFO_GET_STATE_METHOD.setAccessible(true);
+            Class<?> State = Class.forName("org.spongepowered.asm.mixin.transformer.MixinInfo$State");
+            STATE_CLASS_NODE_FIELD = State.getDeclaredField("classNode");
+            STATE_CLASS_NODE_FIELD.setAccessible(true);
+            STATE_CLASS_INFO_FIELD = State.getDeclaredField("classInfo");
+            STATE_CLASS_INFO_FIELD.setAccessible(true);
+            CLASS_INFO_METHODS_FIELD = ClassInfo.class.getDeclaredField("methods");
+            CLASS_INFO_METHODS_FIELD.setAccessible(true);
+            EXTENSIONS_FIELD = Extensions.class.getDeclaredField("extensions");
+            EXTENSIONS_FIELD.setAccessible(true);
+            ACTIVE_EXTENSIONS_FIELD = Extensions.class.getDeclaredField("activeExtensions");
+            ACTIVE_EXTENSIONS_FIELD.setAccessible(true);
+        } catch (ClassNotFoundException | NoSuchFieldException | NoSuchMethodException e) {
+            throw new RuntimeException("Failed to access some mixin internals!", e);
+        }
+    }
+
+    public static List<Pair<IMixinInfo, ClassNode>> getMixinsFor(ITargetClassContext context) {
+        try {
+            List<Pair<IMixinInfo, ClassNode>> result = new ArrayList<>();
+            SortedSet<IMixinInfo> mixins = (SortedSet<IMixinInfo>) TARGET_CLASS_CONTEXT_MIXINS_FIELD.get(context);
+            for (IMixinInfo mixin : mixins) {
+                result.add(Pair.of(mixin, getClassNode(mixin)));
+            }
+            return result;
+        } catch (IllegalAccessException e) {
+            throw new RuntimeException("Failed to use mixin internals!", e);
+        }
+    }
+
+    public static ClassInfo getClassInfoFor(IMixinInfo mixinInfo) {
+        try {
+            Object state = MIXIN_INFO_GET_STATE_METHOD.invoke(mixinInfo);
+			return (ClassInfo) STATE_CLASS_INFO_FIELD.get(state);
+        } catch (IllegalAccessException | InvocationTargetException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public static Set<ClassInfo.Method> getClassInfoMethods(ClassInfo classInfo) {
+        try {
+            return (Set<ClassInfo.Method>) CLASS_INFO_METHODS_FIELD.get(classInfo);
+        } catch (IllegalAccessException e) {
+            throw new RuntimeException("Failed to use mixin internals!", e);
+        }
+    }
+
+    public static void registerExtension(IExtension extension) {
+        try {
+            IMixinTransformer transformer = (IMixinTransformer) MixinEnvironment.getDefaultEnvironment().getActiveTransformer();
+            Extensions extensions = (Extensions) transformer.getExtensions();
+            List<IExtension> extensionsList = (List<IExtension>) EXTENSIONS_FIELD.get(extensions);
+            extensionsList.add(extension);
+            List<IExtension> activeExtensions = new ArrayList<>((List<IExtension>) ACTIVE_EXTENSIONS_FIELD.get(extensions));
+            activeExtensions.add(extension);
+            ACTIVE_EXTENSIONS_FIELD.set(extensions, Collections.unmodifiableList(activeExtensions));
+        } catch (IllegalAccessException e) {
+            throw new RuntimeException("Failed to use mixin internals!", e);
+        }
+    }
+
+    private static ClassNode getClassNode(IMixinInfo mixin) {
+        try {
+            Object state = MIXIN_INFO_GET_STATE_METHOD.invoke(mixin);
+            return (ClassNode) STATE_CLASS_NODE_FIELD.get(state);
+        } catch (IllegalAccessException | InvocationTargetException e) {
+            throw new RuntimeException("Failed to use mixin internals!", e);
+        }
+    }
+}

src/main/resources/optifabric.mixins.json

@@ -2,6 +2,7 @@
   "required": true,
   "package": "me.modmuss50.optifabric.mixin",
   "compatibilityLevel": "JAVA_8",
+  "plugin": "me.modmuss50.optifabric.compat.OptifabricMixinPlugin",
   "mixins": [
     "MixinTitleScreen",
     "CrashReportMixin"