LibJF/libjf-devutil/src/main/java/io/gitlab/jfronny/libjf/devutil/mixin/ArgumentTypesMixin.java

package io.gitlab.jfronny.libjf.devutil.mixin;

import com.mojang.brigadier.arguments.ArgumentType;
import net.minecraft.command.argument.*;
import net.minecraft.command.argument.serialize.ArgumentSerializer;
import net.minecraft.registry.Registry;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Redirect;

@Mixin(ArgumentTypes.class)
public abstract class ArgumentTypesMixin {
    @Shadow
    private static <A extends ArgumentType<?>, T extends ArgumentSerializer.ArgumentTypeProperties<A>> ArgumentSerializer<A, T> register(Registry<ArgumentSerializer<?, ?>> registry, String id, Class<? extends A> clazz, ArgumentSerializer<A, T> serializer) {
        throw new RuntimeException("Mixin not applied properly");
    }

    @Redirect(method = "register(Lnet/minecraft/registry/Registry;)Lnet/minecraft/command/argument/serialize/ArgumentSerializer;", at = @At(value = "INVOKE", target = "Lnet/minecraft/command/argument/ArgumentTypes;register(Lnet/minecraft/registry/Registry;Ljava/lang/String;Ljava/lang/Class;Lnet/minecraft/command/argument/serialize/ArgumentSerializer;)Lnet/minecraft/command/argument/serialize/ArgumentSerializer;"))
    private static <A extends ArgumentType<?>, T extends ArgumentSerializer.ArgumentTypeProperties<A>> ArgumentSerializer<A, T> libjf$redirectRegister(Registry<ArgumentSerializer<?, ?>> registry, String id, Class<? extends A> clazz, ArgumentSerializer<A, T> serializer) {
        if (clazz != TestFunctionArgumentType.class && clazz != TestClassArgumentType.class)
            return register(registry, id, clazz, serializer);
        return null;
    }
}