OptiNotFine / src / main / java / dev / redstudio / optinotfine / asm / OptiNotFineTransformer.java

package dev.redstudio.optinotfine.asm;

import dev.redstudio.optinotfine.config.OptiNotFineConfig;
import net.minecraft.client.resources.AbstractResourcePack;
import net.minecraft.launchwrapper.IClassTransformer;
import net.minecraftforge.fml.relauncher.FMLLaunchHandler;
import org.objectweb.asm.*;
import org.objectweb.asm.tree.*;

import java.lang.reflect.Field;
import java.util.function.BiPredicate;
import java.util.function.Function;

/// Transformer for OptiNotFine.
///
/// @author Luna Mira Lage (Desoroxxx)
/// @since 1.0
public final class OptiNotFineTransformer implements IClassTransformer {

	private static final String RESOURCE_PACK_FILE_FIELD_NAME = FMLLaunchHandler.isDeobfuscatedEnvironment() ? "resourcePackFile" : "field_110597_b";

	public static final String[][] BUFFER_ALLOCATED_PATHS = {
			{"jdk.internal.access.SharedSecrets", "getJavaNioAccess", "getDirectBufferPool", "getMemoryUsed"},
			{"jdk.internal.misc.SharedSecrets", "getJavaNioAccess", "getDirectBufferPool", "getMemoryUsed"},
			{"sun.misc.SharedSecrets", "getJavaNioAccess", "getDirectBufferPool", "getMemoryUsed"}
	};

	@Override
	public byte[] transform(final String name, final String transformedName, final byte[] basicClass) {
		if (basicClass == null)
			return null;

		switch (transformedName) {
			case "net.minecraftforge.fml.client.FMLClientHandler":
				return transformFMLClientHandler(basicClass);
			case "net.minecraft.client.resources.AbstractResourcePack":
				return transformAbstractResourcePack(basicClass);
			case "net.optifine.util.NativeMemory":
				return transformNativeMemory(basicClass);
		}

		if (!OptiNotFineConfig.stopLogSpam)
			return basicClass;

		switch (transformedName) {
			case "net.optifine.shaders.SMCLog":
				return transformSMCLog(basicClass);
			case "net.optifine.config.ConnectedParser":
				return transformConnectedParser(basicClass);
			case "net.optifine.shaders.ItemAliases":
				return transformAliases(basicClass, "loadItemAliases", "item");
			case "net.optifine.shaders.BlockAliases":
				return transformAliases(basicClass, "loadBlockAliases", "block");
			case "net.optifine.shaders.EntityAliases":
				return transformAliases(basicClass, "loadEntityAliases", "entity");
			case "net.optifine.shaders.config.MacroExpressionResolver":
				return stripConfigWarn(basicClass, "getExpression", "(Ljava/lang/String;)Lnet/optifine/expr/IExpression;");
			case "net.optifine.shaders.config.ShaderPackParser":
				return stripConfigWarn(basicClass, "collectShaderOptions", "(Lnet/optifine/shaders/IShaderPack;Ljava/lang/String;Ljava/util/Map;)V");
		}

		return basicClass;
	}

	/// Add our branding to the main menu screen
	private static byte[] transformFMLClientHandler(final byte[] basicClass) {
		return transform(basicClass, classWriter -> targetMethod(classWriter, "getAdditionalBrandingInformation", methodVisitor ->
				new MethodVisitor(Opcodes.ASM5, methodVisitor) {
					@Override
					public void visitLdcInsn(final Object value) {
						super.visitLdcInsn(value.equals("Optifine %s") ? "OptiNotFine 1.0-Dev-1 on %s" : value);
					}
				}));
	}

	/// Add a missing getter method to [AbstractResourcePack].
	///
	/// That getter method was added by Cleanroom but is accidentally removed due to OptiFine's binary patching.
	///
	/// **Warning:** We currently do not check whether this method already exists as we have a guarantee that it does not.
	/// The reason for that is since our transformer runs only when OptiFine is here, the method is definitely removed.
	private static byte[] transformAbstractResourcePack(final byte[] basicClass) {
		return transform(basicClass, classWriter -> new ClassVisitor(Opcodes.ASM5, classWriter) {
			@Override
			public void visitEnd() {
				final MethodVisitor methodVisitor = visitMethod(Opcodes.ACC_PUBLIC, "getResourcePackFile", "()Ljava/io/File;", null, null);

				methodVisitor.visitCode();
				methodVisitor.visitVarInsn(Opcodes.ALOAD, 0);
				methodVisitor.visitFieldInsn(Opcodes.GETFIELD, "net/minecraft/client/resources/AbstractResourcePack", RESOURCE_PACK_FILE_FIELD_NAME, "Ljava/io/File;");
				methodVisitor.visitInsn(Opcodes.ARETURN);
				methodVisitor.visitMaxs(1, 1);
				methodVisitor.visitEnd();

				super.visitEnd();
			}
		});
	}

	/// Add `jdk.internal.access.SharedSecrets` to the paths OptiFine will search to see the native memory usage.
	///
	/// We can safely access this because [Imagine Breaker](https://github.com/Rongmario/ImagineBreaker) will remove the module system.
	private static byte[] transformNativeMemory(final byte[] basicClass) {
		final ClassNode classNode = new ClassNode();
		new ClassReader(basicClass).accept(classNode, 0);

		final String fieldName = findFieldName(BUFFER_ALLOCATED_PATHS);

		for (final MethodNode methodNode : classNode.methods) {
			if (!methodNode.name.equals("<clinit>"))
				continue;

			for (AbstractInsnNode instruction = methodNode.instructions.getFirst(); instruction != null; instruction = instruction.getNext()) {
				if (!isTargetCall(instruction))
					continue;

				replaceArrayArgument(methodNode.instructions, instruction, fieldName);
				break;
			}

			break;
		}

		final ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_MAXS);
		classNode.accept(classWriter);
		return classWriter.toByteArray();
	}

	private static String findFieldName(final Object fieldValue) {
		try {
			for (final Field field : OptiNotFineTransformer.class.getFields())
				if (field.get(null) == fieldValue)
					return field.getName();
		} catch (final IllegalAccessException illegalAccessException) {
			throw new RuntimeException(illegalAccessException);
		}

		throw new IllegalStateException("Field not found for the given value, something is very wrong");
	}

	private static boolean isTargetCall(final AbstractInsnNode instructionNode) {
		if (!(instructionNode instanceof MethodInsnNode))
			return false;

		final MethodInsnNode methodInstruction = (MethodInsnNode) instructionNode;

		return methodInstruction.getOpcode() == Opcodes.INVOKESTATIC && methodInstruction.owner.equals("net/optifine/util/NativeMemory") && methodInstruction.name.equals("makeLongSupplier");
	}

	private static void replaceArrayArgument(final InsnList instructions, final AbstractInsnNode methodCall, final String fieldName) {
		for (AbstractInsnNode current = methodCall.getPrevious(); current != null; current = current.getPrevious()) {
			if (!(current instanceof TypeInsnNode))
				continue;

			final TypeInsnNode typeInstruction = (TypeInsnNode) current;

			if (typeInstruction.getOpcode() != Opcodes.ANEWARRAY || !typeInstruction.desc.equals("[Ljava/lang/String;"))
				continue;

			removeInstructionsBetween(instructions, current.getPrevious(), methodCall);
			instructions.insertBefore(methodCall, new FieldInsnNode(Opcodes.GETSTATIC, Type.getInternalName(OptiNotFineTransformer.class), fieldName, "[[Ljava/lang/String;"));
			return;
		}

		throw new IllegalStateException("Could not find array argument for method call");
	}

	private static void removeInstructionsBetween(final InsnList instructions, final AbstractInsnNode start, final AbstractInsnNode end) {
		if (start == null)
			return;

		AbstractInsnNode current = start;
		while (current != end) {
			final AbstractInsnNode next = current.getNext();
			instructions.remove(current);
			current = next;
		}
	}

	/// Changes the logger call of `SMCLog.info` from `info` to `debug` as most of it is.
	private static byte[] transformSMCLog(final byte[] basicClass) {
		final String owner = "org/apache/logging/log4j/Logger";
		return transform(basicClass, classWriter ->
				targetMethod(classWriter, "info", methodVisitor ->
						methodCallReplacer(methodVisitor, Opcodes.INVOKEINTERFACE, owner, "info", (methodVisitor1, descriptor, isInterface) ->
								methodVisitor1.visitMethodInsn(Opcodes.INVOKEINTERFACE, owner, "debug", descriptor, isInterface))
				));
	}

	/// Strip a few specific `Config#warn()`.
	private static byte[] transformConnectedParser(final byte[] basicClass) {
		return transform(basicClass, classWriter -> new ClassVisitor(Opcodes.ASM5, classWriter) {
			@Override
			public MethodVisitor visitMethod(final int access, final String name, final String descriptor, final String signature, final String[] exceptions) {
				final MethodVisitor methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions);

				final String triggerString;
				if (name.equals("parseItems")) {
					triggerString = "Item not found: ";
				} else if (name.equals("parseBlockPart")) {
					triggerString = "Block not found for name: ";
				} else {
					return methodVisitor;
				}

				return triggerThenStripWarn(methodVisitor, triggerString, Opcodes.INVOKEVIRTUAL, "net/optifine/config/ConnectedParser", 2);
			}
		});
	}

	/// Strip a specific `Config#warn()`.
	private static byte[] transformAliases(final byte[] basicClass, final String methodName, final String warnPrefix) {
		final String trigger = "[Shaders] Invalid " + warnPrefix + " ID mapping: ";

		return transform(basicClass, classWriter ->
				targetMethod(classWriter, methodName, methodVisitor ->
						triggerThenStripWarn(methodVisitor, trigger, Opcodes.INVOKESTATIC, "Config", 1)
				));
	}

	private static MethodVisitor triggerThenStripWarn(final MethodVisitor parent, final String triggerString, final int targetOpcode, final String warnOwner, final int popCount) {
		return new MethodVisitor(Opcodes.ASM5, parent) {
			private boolean stripNextWarn = false;

			@Override
			public void visitLdcInsn(final Object value) {
				if (triggerString.equals(value))
					stripNextWarn = true;

				super.visitLdcInsn(value);
			}

			@Override
			public void visitMethodInsn(final int opcode, final String owner, final String name, final String descriptor, final boolean isInterface) {
				if (stripNextWarn && opcode == targetOpcode && owner.equals(warnOwner) && name.equals("warn")) {
					stripNextWarn = false;

					for (int i = 0; i < popCount; i++)
						super.visitInsn(Opcodes.POP);

					return;
				}

				super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
			}
		};
	}

	/// Strip any `Config#warn()` calls from a method.
	private static byte[] stripConfigWarn(final byte[] basicClass, final String targetMethod, final String targetDescriptor) {
		return transform(basicClass, classWriter ->
				targetMethod(classWriter, (name, descriptor) -> name.equals(targetMethod) && descriptor.equals(targetDescriptor), methodVisitor ->
						methodCallReplacer(methodVisitor, Opcodes.INVOKESTATIC, "Config", "warn", (methodVisitor1, descriptor, isInterface) ->
								methodVisitor1.visitInsn(Opcodes.POP)
						)));
	}

	private static MethodVisitor methodCallReplacer(final MethodVisitor parent, final int targetOpcode, final String targetOwner, final String targetName, final MethodCallReplacement replacement) {
		return new MethodVisitor(Opcodes.ASM5, parent) {
			@Override
			public void visitMethodInsn(final int opcode, final String owner, final String name, final String descriptor, final boolean isInterface) {
				if (opcode == targetOpcode && owner.equals(targetOwner) && name.equals(targetName)) {
					replacement.apply(this, descriptor, isInterface);
				} else {
					super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
				}
			}
		};
	}

	@FunctionalInterface
	interface MethodCallReplacement {
		void apply(final MethodVisitor methodVisitor, final String descriptor, final boolean isInterface);
	}

	private static byte[] transform(final byte[] basicClass, final Function<ClassWriter, ClassVisitor> visitorFactory) {
		final ClassReader classReader = new ClassReader(basicClass);
		final ClassWriter classWriter = new ClassWriter(classReader, 0);

		classReader.accept(visitorFactory.apply(classWriter), 0);

		return classWriter.toByteArray();
	}

	private static ClassVisitor targetMethod(final ClassWriter classWriter, final String methodName, final MethodVisitorFactory factory) {
		return targetMethod(classWriter, (name, descriptor) -> name.equals(methodName), factory);
	}

	private static ClassVisitor targetMethod(final ClassWriter classWriter, final BiPredicate<String, String> matcher, final MethodVisitorFactory factory) {
		return new ClassVisitor(Opcodes.ASM5, classWriter) {
			@Override
			public MethodVisitor visitMethod(final int access, final String name, final String descriptor, final String signature, final String[] exceptions) {
				final MethodVisitor methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions);
				return matcher.test(name, descriptor) ? factory.create(methodVisitor) : methodVisitor;
			}
		};
	}

	@FunctionalInterface
	interface MethodVisitorFactory {
		MethodVisitor create(final MethodVisitor parent);
	}
}