WithValue injector

ScratchPrelaunch.java

package io.github.daomephsta;

import org.spongepowered.asm.mixin.injection.struct.InjectionInfo;

import io.github.daomephsta.injector.WithValueInjectionInfo;
import net.fabricmc.loader.api.entrypoint.PreLaunchEntrypoint;

public class ScratchPrelaunch implements PreLaunchEntrypoint
{
    @Override
    public void onPreLaunch()
    {
        InjectionInfo.register(WithValueInjectionInfo.class);
    }
}

WithValue.java

package io.github.daomephsta.injector;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

import org.spongepowered.asm.mixin.injection.At;

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface WithValue
{
    public String[] method();
    
    public At[] at();
}

WithValueInjectionInfo.java

package io.github.daomephsta.injector;

import org.objectweb.asm.tree.AnnotationNode;
import org.objectweb.asm.tree.MethodNode;
import org.spongepowered.asm.mixin.injection.code.Injector;
import org.spongepowered.asm.mixin.injection.struct.InjectionInfo;
import org.spongepowered.asm.mixin.injection.struct.InjectionInfo.HandlerPrefix;
import org.spongepowered.asm.mixin.transformer.MixinTargetContext;

@InjectionInfo.AnnotationType(WithValue.class)
@HandlerPrefix("withValue")
public class WithValueInjectionInfo extends InjectionInfo
{
    public WithValueInjectionInfo(MixinTargetContext mixin, MethodNode method, AnnotationNode annotation)
    {
        super(mixin, method, annotation);
    }

    @Override
    protected Injector parseInjector(AnnotationNode injectAnnotation)
    {
        return new WithValueInjector(this);
    }
}

WithValueInjector.java

package io.github.daomephsta.injector;

import static org.objectweb.asm.Opcodes.ARETURN;
import static org.objectweb.asm.Opcodes.DUP;
import static org.objectweb.asm.Opcodes.GETFIELD;
import static org.objectweb.asm.Opcodes.GETSTATIC;
import static org.objectweb.asm.Opcodes.INVOKEINTERFACE;
import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
import static org.objectweb.asm.Opcodes.INVOKESTATIC;
import static org.objectweb.asm.Opcodes.INVOKEVIRTUAL;
import static org.objectweb.asm.Opcodes.NEW;

import java.lang.reflect.Modifier;
import java.util.Optional;
import java.util.stream.Stream;

import org.objectweb.asm.Type;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.FieldInsnNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.TypeInsnNode;
import org.spongepowered.asm.mixin.injection.code.Injector;
import org.spongepowered.asm.mixin.injection.struct.InjectionInfo;
import org.spongepowered.asm.mixin.injection.struct.InjectionNodes.InjectionNode;
import org.spongepowered.asm.mixin.injection.struct.Target;
import org.spongepowered.asm.mixin.injection.throwables.InjectionError;

public class WithValueInjector extends Injector
{
    public WithValueInjector(InjectionInfo info)
    {
        super(info, "@WithValue");
    }

    @Override
    protected void inject(Target target, InjectionNode node)
    {
        MethodNode handler = info.getMethod();
        Type handlerType = Type.getMethodType(handler.desc);
        if (target.isStatic && !Modifier.isStatic(handler.access))
            throw new InjectionError("Handler must be static like its target");
        if (handlerType.getReturnType() != Type.VOID_TYPE)
            throw new InjectionError("Return type of handler must be void");
        if (handlerType.getArgumentTypes().length != 1)
            throw new InjectionError("Handler must have exactly one parameter");
        
        InsnList insns = new InsnList();
        insns.add(new InsnNode(DUP));
        insns.add(new MethodInsnNode(getInvokeOpcode(handler), info.getClassNode().name, handler.name, handler.desc));
        target.insns.insert(findValueNode(target, handlerType, node.getCurrentTarget()).orElseThrow(
            () -> new InjectionError("Injection target does not push a value onto the stack, " + 
            "or the value is compile-time constant")), insns);
        info.addCallbackInvocation(handler);
    }
    
    private Optional<AbstractInsnNode> findValueNode(Target target, Type handlerType, AbstractInsnNode node)
    {
        switch (node.getOpcode())
        {
            case INVOKEVIRTUAL:
            case INVOKESPECIAL:
            case INVOKESTATIC:
            case INVOKEINTERFACE:
                checkArgumentType(handlerType, Type.getReturnType(((MethodInsnNode) node).desc));
                return Optional.of(node);
            case GETFIELD:
            case GETSTATIC:
                checkArgumentType(handlerType, Type.getType(((FieldInsnNode) node).desc));
                return Optional.of(node);
            case NEW:
            {
                TypeInsnNode typeInsn = (TypeInsnNode) node;
                checkArgumentType(handlerType, Type.getObjectType(typeInsn.desc));
                return Optional.of(target.findInitNodeFor(typeInsn));
            }
            case ARETURN:
                return Stream.iterate(node, n -> n.getPrevious() != null, AbstractInsnNode::getPrevious)
                    .map(n -> findValueNode(target, handlerType, n))
                    .filter(Optional::isPresent)
                    .findFirst()
                    .orElseThrow(() -> new InjectionError("Could not find value source for ARETURN"));
            default:
                return Optional.empty();
        }
    }
    
    private void checkArgumentType(Type handlerType, Type actualType)
    {
        Type expectedType = handlerType.getArgumentTypes()[0];
        if (!expectedType.equals(actualType))
            throw new InjectionError("Injection target pushes value of type " + actualType + ". Expected type " + expectedType);
    }

    private int getInvokeOpcode(MethodNode handler)
    {
        return Modifier.isStatic(handler.access) ? INVOKESTATIC : INVOKEVIRTUAL;
    }
}

WithValueTest.java

package io.github.daomephsta.mixin;

import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.injection.At;

import com.google.common.collect.ImmutableMap;

import io.github.daomephsta.injector.WithValue;
import net.minecraft.block.Block;
import net.minecraft.block.Blocks;
import net.minecraft.item.AxeItem;

@Mixin(AxeItem.class)
public class WithValueTest
{
    @WithValue(method = "<clinit>", at = @At(value = "NEW", target = "com/google/common/collect/ImmutableMap$Builder"))
    private static void addStrippable(ImmutableMap.Builder<Block, Block> strippable)
    {
        strippable.put(Blocks.STICKY_PISTON, Blocks.PISTON);
    }
}