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>"))
                        instrumentConstructor(invokespecial);
                    break;
            }
            i++;
        }
    }

    private void instrumentConstructor(MethodInsnNode invokespecial) {
        Type[] arguments = Type.getArgumentTypes(invokespecial.desc);
        InsnList copy = new InsnList();

        for (int i = arguments.length - 1; i >= 0; i--)
            copy.add(new VarInsnNode(arguments[i].getOpcode(Opcodes.ISTORE), method.maxLocals + i));
        copy.add(generateDup());
        for (int i = 0; i < arguments.length; i++)
            copy.add(new VarInsnNode(arguments[i].getOpcode(Opcodes.ILOAD), method.maxLocals + i));

        instructions.insertBefore(invokespecial, copy);
        instructions.insert(invokespecial, generateInvokestatic());
    }

    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;
    }
}

@@ -1377,36 +1377,27 @@
       instru
-ctions.inser
+men
 t(instructio
@@ -1401,35 +1401,8 @@
 tion
-, generateInstrumentation()
 );%0A 
@@ -1733,36 +1733,27 @@
       instru
-ctions.inser
+men
 t(instructio
@@ -1757,35 +1757,8 @@
 tion
-, generateInstrumentation()
 );%0A 
@@ -2018,36 +2018,27 @@
       instru
-ctions.inser
+men
 t(instructio
@@ -2042,35 +2042,8 @@
 tion
-, generateInstrumentation()
 );%0A 
@@ -2324,60 +2324,26 @@
 stru
-ctions.insert(instruction, generateInstrumentation()
+ment(invokespecial
 );%0A 
@@ -2585,32 +2585,262 @@
        %7D%0A    %7D%0A%0A
+    private void instrument(AbstractInsnNode instruction) %7B%0A        AbstractInsnNode dup = new InsnNode(Opcodes.DUP);%0A        instructions.insert(instruction, dup);%0A        instructions.insert(dup, generateInvokestatic());%0A    %7D%0A%0A
     private void
@@ -3175,36 +3175,48 @@
    copy.add(
-generateDup(
+new InsnNode(Opcodes.DUP
 ));%0A        
@@ -3494,102 +3494,8 @@
  %7D%0A%0A
-    private AbstractInsnNode generateDup() %7B%0A        return new InsnNode(Opcodes.DUP);%0A    %7D%0A%0A
     
@@ -3659,233 +3659,7 @@
    %7D
-%0A%0A    private InsnList generateInstrumentation() %7B%0A        InsnList instructions = new InsnList();%0A%0A        instructions.add(generateDup());%0A        instructions.add(generateInvokestatic());%0A%0A        return instructions;%0A    %7D
 %0A%7D%0A

Refactor instruction instrumentation