package autostack;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.commons.TryCatchBlockSorter;

public class AutostackAgent implements Opcodes, ClassFileTransformer {
    private static boolean atLeastJava7() {
        String propValue = System.getProperty("java.class.version");
        return propValue.startsWith("52.") || propValue.startsWith("51.");
    }

    public static void premain(String agentArguments, Instrumentation instrumentation) {
        instrumentation.addTransformer(new AutostackAgent());
    }

    public byte[] transform(ClassLoader loader, final String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer)
            throws IllegalClassFormatException {
        if (className == null || className.startsWith("java/")
                || className.startsWith("org/lwjgl/vulkan/")
                || className.startsWith("org/lwjgl/system/"))
            return null;
        ClassReader cr = new ClassReader(classfileBuffer);
        final Set<String> stackMethods = new HashSet<String>();
        // Scan all methods that need auto-stack
        cr.accept(new ClassVisitor(ASM5) {
            public MethodVisitor visitMethod(int access, final String methodName, final String methodDesc, String signature, String[] exceptions) {
                MethodVisitor mv = new MethodVisitor(ASM5) {
                    public void visitMethodInsn(int opcode, String owner, String name, String desc, boolean itf) {
                        if (opcode == INVOKESTATIC && owner.startsWith("org/lwjgl/") && (
                                name.equals("mallocStack") ||
                                name.equals("stackGet") ||
                                name.equals("callocStack")) && !itf) {
                            stackMethods.add(methodName + methodDesc);
                        }
                    }
                };
                return mv;
            }
        }, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
        if (stackMethods.isEmpty())
            return null;

        // Now, transform all such methods
        ClassWriter cw = new ClassWriter(cr, 0 | (atLeastJava7() ? ClassWriter.COMPUTE_FRAMES : 0));
        cr.accept(new ClassVisitor(ASM5, cw) {
            public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
                final MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
                if (!stackMethods.contains(name + desc))
                    return mv;
                MethodVisitor tcbs = new TryCatchBlockSorter(mv, access, name, desc, signature, exceptions);
                MethodVisitor own = new MethodVisitor(ASM5, tcbs) {
                    Label tryLabel = new Label();
                    Label finallyLabel = new Label();

                    public void visitCode() {
                        mv.visitCode();
                        mv.visitMethodInsn(INVOKESTATIC, "org/lwjgl/system/MemoryStack", "stackPush", "()Lorg/lwjgl/system/MemoryStack;", false);
                        mv.visitLabel(tryLabel);
                    }

                    public void visitEnd() {
                        mv.visitLabel(finallyLabel);
                        mv.visitMethodInsn(INVOKESTATIC, "org/lwjgl/system/MemoryStack", "stackPop", "()Lorg/lwjgl/system/MemoryStack;", false);
                        mv.visitInsn(RETURN);
                        mv.visitTryCatchBlock(tryLabel, finallyLabel, finallyLabel, "java/lang/Exception");
                        mv.visitEnd();
                    }
                };
                return own;
            }
        }, 0);
        byte[] arr = cw.toByteArray();
        return arr;
    }
}

@@ -489,16 +489,48 @@
 codes;%0D%0A
+import org.objectweb.asm.Type;%0D%0A
 import o
@@ -2815,32 +2815,38 @@
 tor visitMethod(
+final 
 int access, Stri
@@ -2845,32 +2845,38 @@
 ss, String name,
+ final
  String desc, St
@@ -3239,165 +3239,1208 @@
     
-MethodVisitor own = new MethodVisitor(ASM5, tcbs) %7B%0D%0A                    Label tryLabel = new Label();%0D%0A                    Label finallyLabel = new Label();
+int var = ((access & ACC_STATIC) == ACC_STATIC ? 0 : 1);%0D%0A                Type%5B%5D paramTypes = Type.getArgumentTypes(desc);%0D%0A                for (int i = 0; i %3C paramTypes.length; i++) %7B%0D%0A                    Type t = paramTypes%5Bi%5D;%0D%0A                    var += t.getSize();%0D%0A                %7D%0D%0A                final int firstFreeVar = var;%0D%0A                MethodVisitor own = new MethodVisitor(ASM5, tcbs) %7B%0D%0A                    Label tryLabel = new Label();%0D%0A                    Label finallyLabel = new Label();%0D%0A                    int freeVar = firstFreeVar;%0D%0A%0D%0A                    public void visitVarInsn(int opcode, int var) %7B%0D%0A                        freeVar = Math.max(freeVar, var + 1);%0D%0A                        mv.visitVarInsn(opcode, var);%0D%0A                    %7D%0D%0A%0D%0A                    public void visitInsn(int opcode) %7B%0D%0A                        if (opcode %3E= IRETURN && opcode %3C= RETURN) %7B%0D%0A                            mv.visitMethodInsn(INVOKESTATIC, %22org/lwjgl/system/MemoryStack%22, %22stackPop%22, %22()Lorg/lwjgl/system/MemoryStack;%22, false);%0D%0A                            mv.visitInsn(POP);%0D%0A                        %7D%0D%0A                        super.visitInsn(opcode);%0D%0A                    %7D
 %0D%0A%0D%0A
@@ -4666,32 +4666,76 @@
 ack;%22, false);%0D%0A
+                        mv.visitInsn(POP);%0D%0A
                 
@@ -4885,32 +4885,91 @@
 finallyLabel);%0D%0A
+                        mv.visitVarInsn(ASTORE, freeVar);%0D%0A
                 
@@ -5135,22 +5135,124 @@
 sitInsn(
-RETURN
+POP);%0D%0A                        mv.visitVarInsn(ALOAD, freeVar);%0D%0A                        mv.visitInsn(ATHROW
 );%0D%0A    

Handle methods returning non-void