LibJF/libjf-unsafe-v0/src/main/java/io/gitlab/jfronny/libjf/unsafe/asm/AsmTransformer.java

package io.gitlab.jfronny.libjf.unsafe.asm;

import io.gitlab.jfronny.libjf.Flags;
import io.gitlab.jfronny.libjf.LibJf;
import io.gitlab.jfronny.libjf.unsafe.asm.patch.Patch;
import net.fabricmc.loader.api.FabricLoader;
import net.fabricmc.loader.api.MappingResolver;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.tree.ClassNode;
import org.spongepowered.asm.mixin.MixinEnvironment;
import org.spongepowered.asm.mixin.transformer.IMixinTransformer;
import org.spongepowered.asm.mixin.transformer.ext.IExtensionRegistry;
import org.spongepowered.asm.transformers.MixinClassWriter;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;

public class AsmTransformer implements IMixinTransformer {
    public static AsmTransformer INSTANCE;
    public static final MappingResolver MAPPING_RESOLVER = FabricLoader.getInstance().getMappingResolver();
    public static final String INTERMEDIARY = "intermediary";
    public IMixinTransformer delegate;
    public Set<BakedAsmConfig> asmConfigs;
    private AsmConfig currentConfig = null;
    private boolean export;
    private boolean debugLog;

    public void init() {
        Set<Flags.BooleanFlag> flags = Flags.getBoolFlags("asm.export");
        flags.removeIf(flag -> !flag.value());
        export = !flags.isEmpty();
        if (export) {
            Set<String> flagNames = new LinkedHashSet<>();
            for (Flags.BooleanFlag flag : flags) {
                flagNames.add(flag.source());
            }
            LibJf.LOGGER.info("Exporting ASM due to: " + String.join(", ", flagNames));
        }
        flags = Flags.getBoolFlags("asm.log");
        flags.removeIf(flag -> !flag.value());
        debugLog = !flags.isEmpty();
        if (debugLog) {
            Set<String> flagNames = new LinkedHashSet<>();
            for (Flags.BooleanFlag flag : flags) {
                flagNames.add(flag.source());
            }
            LibJf.LOGGER.info("Logging ASM logs due to: " + String.join(", ", flagNames));
        }
    }

    public boolean debugLogsEnabled() {
        return debugLog;
    }

    @Override
    public void audit(MixinEnvironment environment) {
        delegate.audit(environment);
    }

    @Override
    public List<String> reload(String mixinClass, ClassNode classNode) {
        return delegate.reload(mixinClass, classNode);
    }

    @Override
    public boolean computeFramesForClass(MixinEnvironment environment, String name, ClassNode classNode) {
        return delegate.computeFramesForClass(environment, name, classNode);
    }

    @Override
    public byte[] transformClassBytes(String name, String transformedName, byte[] classBytes) {
        classBytes = delegate.transformClassBytes(name, transformedName, classBytes);
        if (classBytes == null || name == null) return classBytes;
        if (isClassUnmoddable(name, null)) {
            if (debugLogsEnabled()) LibJf.LOGGER.info("Skipping " + name);
            return classBytes;
        }

        ClassNode klazz = new ClassNode();
        ClassReader reader = new ClassReader(classBytes);
        reader.accept(klazz, ClassReader.EXPAND_FRAMES);

        for (AsmConfig config : asmConfigs) {
            currentConfig = config;
            if (!isClassUnmoddable(name, config)) {
                for (Patch patch : config.getPatches()) {
                    try {
                        patch.apply(klazz);
                    }
                    catch (Throwable t) {
                        LibJf.LOGGER.error("Could not apply patch: " + patch.getClass() + " on " + name, t);
                    }
                }
            }
        }
        currentConfig = null;

        ClassWriter writer = new MixinClassWriter(reader, ClassWriter.COMPUTE_FRAMES);
        try {
            klazz.accept(writer);
        }
        catch (NullPointerException t) {
            LibJf.LOGGER.error("Could not transform " + transformedName, t);
            return null;
        }
        classBytes = writer.toByteArray();
        if (export) {
            try {
                Path path = FabricLoader.getInstance().getGameDir().resolve("libjf").resolve("asm").resolve(name.replace('.', '/') + ".class");
                if (!Files.exists(path)) Files.createDirectories(path.getParent());
                Files.write(path, classBytes);
            } catch (IOException e) {
                LibJf.LOGGER.error("Could not export modified bytecode", e);
            }
        }
        return classBytes;
    }

    @Override
    public byte[] transformClass(MixinEnvironment environment, String name, byte[] classBytes) {
        return delegate.transformClass(environment, name, classBytes);
    }

    @Override
    public boolean transformClass(MixinEnvironment environment, String name, ClassNode classNode) {
        return delegate.transformClass(environment, name, classNode);
    }

    @Override
    public byte[] generateClass(MixinEnvironment environment, String name) {
        return delegate.generateClass(environment, name);
    }

    @Override
    public boolean generateClass(MixinEnvironment environment, String name, ClassNode classNode) {
        return delegate.generateClass(environment, name, classNode);
    }

    @Override
    public IExtensionRegistry getExtensions() {
        return delegate.getExtensions();
    }

    public static boolean isClassUnmoddable(String className, AsmConfig config) {
        className = className.replace('/', '.');
        if (className.startsWith("org.objectweb.asm")
                || className.startsWith("org.spongepowered.asm")
                //|| className.startsWith("net.fabricmc.loader")
                //|| className.startsWith("io.gitlab.jfronny.libjf.unsafe.asm")
        )
            return true;
        if (config == null) return false;
        Set<String> classes = config.skipClasses();
        if (classes == null) return false;
        return classes.contains(MAPPING_RESOLVER.unmapClassName(INTERMEDIARY, className));
    }

    public AsmConfig getCurrentConfig() {
        return currentConfig;
    }
}