package com.github.sulir.runtimesearch.agent.transformer;

import com.github.sulir.runtimesearch.agent.Check;
import com.github.sulir.runtimesearch.agent.analyzer.StackAnalyzer;
import com.github.sulir.runtimesearch.agent.analyzer.StringValue;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.*;
import org.objectweb.asm.tree.analysis.AnalyzerException;

public class MethodTransformer implements Opcodes {
    public static final String CHECK_CLASS = Type.getType(Check.class).getInternalName();

    private final String className;
    private final MethodNode method;
    private final InsnList instructions;

    public MethodTransformer(String className, MethodNode method) {
        this.className = className;
        this.method = method;
        this.instructions = method.instructions;
    }

    public void transform() throws AnalyzerException {
        StackAnalyzer analyzer = new StackAnalyzer(className, method);
        analyzer.analyze();

        int i = 0;
        for (AbstractInsnNode instruction : instructions) {
            switch (instruction.getOpcode()) {
                case LDC:
                case ALOAD:
                case AALOAD:
                case GETSTATIC:
                case GETFIELD:
                    if (analyzer.getStackTopAfter(i) == StringValue.MAYBE_STRING)
                        instructions.insert(instruction, generateInstrumentation());
                    break;
                case INVOKEVIRTUAL:
                case INVOKESTATIC:
                case INVOKEINTERFACE:
                    Type returnType = Type.getType(((MethodInsnNode) instruction).desc).getReturnType();
                    if (new ObjectType(returnType).canBeString())
                        instructions.insert(instruction, generateInstrumentation());
                    break;
                case INVOKEDYNAMIC:
                    returnType = Type.getType(((InvokeDynamicInsnNode) instruction).desc).getReturnType();
                    if (new ObjectType(returnType).canBeString())
                        instructions.insert(instruction, generateInstrumentation());
                    break;
                case INVOKESPECIAL:
                    MethodInsnNode invokespecial = (MethodInsnNode) instruction;

                    if (new ObjectType(Type.getType(invokespecial.desc).getReturnType()).canBeString())
                        instructions.insert(instruction, generateInstrumentation());
                    else if (invokespecial.owner.equals(ObjectType.STRING) && invokespecial.name.equals("<init>"))
                        instructions.insert(instruction, generateInvokestatic());
                    break;
                case NEW:
                    if (analyzer.getStackTopAfter(i) == StringValue.MAYBE_STRING)
                        instructions.insert(instruction, generateDup());
            }
            i++;
        }
    }

    private AbstractInsnNode generateDup() {
        return new InsnNode(Opcodes.DUP);
    }

    private AbstractInsnNode generateInvokestatic() {
        return new MethodInsnNode(Opcodes.INVOKESTATIC, CHECK_CLASS, "perform", "(Ljava/lang/Object;)V");
    }

    private InsnList generateInstrumentation() {
        InsnList instructions = new InsnList();

        instructions.add(generateDup());
        instructions.add(generateInvokestatic());

        return instructions;
    }
}

@@ -2632,57 +2632,37 @@
 stru
-ctions.insert(instruction, generateInvokestatic()
+mentConstructor(invokespecial
 );%0A 
@@ -2703,37 +2703,35 @@
     
+%7D%0A
     
-case NEW:%0A
+ 
        
+i++;%0A
         
     
@@ -2730,188 +2730,664 @@
     
+%7D%0A
     
- if (analyzer.getStackTopAfter(i) == StringValue.MAYBE_STRING)%0A                        instructions.insert(instruction, generateDup());%0A            %7D%0A            i++;%0A        %7D
+%7D%0A%0A    private void instrumentConstructor(MethodInsnNode invokespecial) %7B%0A        Type%5B%5D arguments = Type.getArgumentTypes(invokespecial.desc);%0A        InsnList copy = new InsnList();%0A%0A        for (int i = arguments.length - 1; i %3E= 0; i--)%0A            copy.add(new VarInsnNode(arguments%5Bi%5D.getOpcode(Opcodes.ISTORE), method.maxLocals + i));%0A        copy.add(generateDup());%0A        for (int i = 0; i %3C arguments.length; i++)%0A            copy.add(new VarInsnNode(arguments%5Bi%5D.getOpcode(Opcodes.ILOAD), method.maxLocals + i));%0A%0A        instructions.insertBefore(invokespecial, copy);%0A        instructions.insert(invokespecial, generateInvokestatic());
 %0A   

Fix VerifyError when dup_x1 is between new and invokespecial