QuickMath/src/main/java/io/gitlab/jfronny/quickmath/BytecodeTransformer.java

package io.gitlab.jfronny.quickmath;

import io.gitlab.jfronny.libjf.unsafe.asm.AsmConfig;
import io.gitlab.jfronny.libjf.unsafe.asm.AsmTransformer;
import io.gitlab.jfronny.libjf.unsafe.asm.patch.Patch;
import io.gitlab.jfronny.libjf.unsafe.asm.patch.PatchUtil;
import org.objectweb.asm.*;
import org.objectweb.asm.tree.*;

import java.util.Map;
import java.util.Set;

public class BytecodeTransformer implements AsmConfig {
    private static final String math = "java/lang/Math";
    private static final String random = "java/util/Random";
    private static final String mathUtil = "io/gitlab/jfronny/quickmath/MathUtil";
    private static final String mathHelperIntermediary = "net.minecraft.class_3532";
    private static final String mathHelper = PatchUtil.mapClassNameInternal(mathHelperIntermediary);
    private static final String mojangRandomIntermediary = "net.minecraft.class_5819";
    private static final String mojangRandom = PatchUtil.mapClassNameInternal(mojangRandomIntermediary);
    private static final String mathHelperRandomUuid = mth("method_15378", "(Lnet/minecraft/class_5819;)Ljava/util/UUID;");
    private static final Map<String, String> mth = DMap.of( // Maps methods in mathHelper to QuickMäth MathUtil methods
            mth("method_15374", "(F)F"), "sinM",
            mth("method_15362", "(F)F"), "cosM",
            mth("method_15355", "(F)F"), "sqrtM",
            mth("method_15375", "(D)I"), "floor"
    );
    private static final Map<String, String> rnd = DMap.of( // Maps methods in Minecraft Random to QuickMäth MathUtil methods
            rnd("method_43054", "()I"), "nextInt",
            rnd("method_43048", "(I)I"), "nextInt",
            rnd("method_43055", "()J"), "nextLong",
            rnd("method_43056", "()Z"), "nextBoolean",
            rnd("method_43057", "()F"), "nextFloat",
            rnd("method_43058", "()D"), "random",
            rnd("method_43059", "()D"), "random"
    );
    private static final Map<String, Boolean> stat = Map.of( // Maps QuickMäth MathUtil methods to booleans representing whether to overwrite them
            "sin", Cfg.corruptTrigonometry.contains(Cfg.CorruptionLevel2.FULL),
            "cos", Cfg.corruptTrigonometry.contains(Cfg.CorruptionLevel2.FULL),
            "sinM", Cfg.corruptTrigonometry.contains(Cfg.CorruptionLevel2.MAJOR),
            "cosM", Cfg.corruptTrigonometry.contains(Cfg.CorruptionLevel2.MAJOR),
            //"sqrt", Cfg.corruptGenericMath.contains(Cfg.CorruptionLevel2.MAJOR),
            "sqrtM", Cfg.corruptGenericMath.contains(Cfg.CorruptionLevel2.MAJOR),
            //"floor", Cfg.corruptGenericMath.contains(Cfg.CorruptionLevel2.FULL),
//            "nextInt", Cfg.corruptGenericMath.contains(Cfg.CorruptionLevel2.FULL),
            "nextLong", Cfg.corruptGenericMath.contains(Cfg.CorruptionLevel2.FULL),
            "nextBoolean", Cfg.corruptGenericMath.contains(Cfg.CorruptionLevel2.FULL),
//            "nextFloat", Cfg.corruptGenericMath.contains(Cfg.CorruptionLevel2.FULL),
            "random", Cfg.corruptGenericMath.contains(Cfg.CorruptionLevel2.FULL)
    );

    private static String mth(String method, String descriptor) {
        return PatchUtil.mapMethodName(mathHelperIntermediary, method, descriptor);
    }

    private static String rnd(String method, String descriptor) {
        return PatchUtil.mapMethodName(mojangRandomIntermediary, method, descriptor);
    }

    @Override
    public Set<String> skipClasses() {
        return null;
    }

    @Override
    public Set<Patch> getPatches() {
        return Set.of(this::patchInvokes);
    }

    private void patchInvokes(ClassNode klazz) {
        for (MethodNode method : klazz.methods) {
            if (klazz.name.equals(mathHelper) && method.name.equals(mathHelperRandomUuid)) { // UUIDs still need to work
                if (Cfg.debugAsm) {
                    ModMain.LOGGER.info("Skipped replacing method calls in MathHelper.randomUuid");
                }
                continue;
            }
            for (AbstractInsnNode insn : method.instructions.toArray()) {
                if (insn.getOpcode() == Opcodes.INVOKESTATIC || insn.getOpcode() == Opcodes.INVOKEVIRTUAL || insn.getOpcode() == Opcodes.INVOKEINTERFACE) {
                    String insNew = null;
                    MethodInsnNode mIns = (MethodInsnNode) insn;
                    // Resolve a possible replacement method in QuickMäth MathUtil
                    if (mIns.owner.equals(math)) {
                        if (stat.containsKey(mIns.name))
                            insNew = mIns.name;
                    } else if (mIns.owner.equals(mathHelper)) {
                        if (mth.containsKey(mIns.name))
                            insNew = mth.get(mIns.name);
                    } else if (mIns.owner.equals(mojangRandom)) {
                        if (rnd.containsKey(mIns.name))
                            insNew = rnd.get(mIns.name);
                    } else if (mIns.owner.equals(random)) {
                        insNew = switch (mIns.name) {
                            case "nextInt" -> "nextInt";
                            case "nextLong" -> "nextLong";
                            case "nextBoolean" -> "nextBoolean";
                            case "nextFloat" -> "nextFloat";
                            case "nextDouble", "nextGaussian" -> "random";
                            default -> null;
                        };
                    }
                    // Check whether the method should be replaced
                    if (!klazz.name.equals(mathUtil) && insNew != null && stat.containsKey(insNew) && stat.get(insNew)) {
                        String originalOwner = mIns.owner;
                        String originalName = mIns.name;
                        // Pop the instance when calling an instance method
                        if (mIns.getOpcode() != Opcodes.INVOKESTATIC) {
                            Type[] params = Type.getArgumentTypes(mIns.desc);
                            // This implementation only works with 0 or 1 parameters of category 1 computational types
                            // This means that doubles and longs are unsupported
                            if (params.length > 1)
                                throw new IllegalArgumentException("The quickmeth bytecode transformer does not support more than one argument");
                            for (Type param : params) {
                                if (param.getSize() != 1)
                                    throw new IllegalStateException("The quickmeth bytecode transformer only supports category 1 computational types");
                            }
                            // If a parameter is present, swap the object to the top, then pop
                            if (params.length == 1)
                                method.instructions.insertBefore(mIns, new InsnNode(Opcodes.SWAP));
                            // Pop the object instance, leaving the parameter if it exists
                            method.instructions.insertBefore(mIns, new InsnNode(Opcodes.POP));
                        }
                        // Invoke the static method
                        mIns.setOpcode(Opcodes.INVOKESTATIC);
                        mIns.owner = mathUtil;
                        mIns.name = insNew;
                        mIns.itf = false;

                        if (Cfg.debugAsm) {
                            ModMain.LOGGER.info("Replaced call to L" + originalOwner + ";" + originalName + mIns.desc
                                    + " in " + klazz.name
                                    + " with L" + mIns.owner + ";" + mIns.name + mIns.desc);
                        }
                    }
                }
            }
        }
    }
}