package com.intellij.compiler.notNullVerification;

import org.objectweb.asm.*;

import java.util.ArrayList;

/**
 * @author ven
 * @noinspection HardCodedStringLiteral
 */
public class NotNullVerifyingInstrumenter extends ClassAdapter {

  private boolean myIsModification = false;
  private String myClassName;

  public NotNullVerifyingInstrumenter(final ClassVisitor classVisitor) {
    super(classVisitor);
  }

  public boolean isModification() {
    return myIsModification;
  }

  public void visit(final int version,
                    final int access,
                    final String name,
                    final String signature,
                    final String superName,
                    final String[] interfaces) {
    super.visit(version, access, name, signature, superName, interfaces);
    myClassName = name;
  }

  public MethodVisitor visitMethod(
    final int access,
    final String name,
    final String desc,
    final String signature,
    final String[] exceptions) {
    final Type[] args = Type.getArgumentTypes(desc);
    final Type returnType = Type.getReturnType(desc);
    MethodVisitor v = cv.visitMethod(access,
                                     name,
                                     desc,
                                     signature,
                                     exceptions);
    return new MethodAdapter(v) {

      private ArrayList myNotNullParams = new ArrayList();
      private boolean myIsNotNull = false;

      public AnnotationVisitor visitParameterAnnotation(
        final int parameter,
        final String anno,
        final boolean visible) {
        AnnotationVisitor av;
        av = mv.visitParameterAnnotation(parameter,
                                         anno,
                                         visible);
        if (args[parameter].getSort() == Type.OBJECT &&
            anno.equals("Lorg/jetbrains/annotations/NotNull;")) {
          myNotNullParams.add(new Integer(parameter));
        }
        return av;
      }

      public AnnotationVisitor visitAnnotation(String anno,
                                               boolean isRuntime) {
        final AnnotationVisitor av = mv.visitAnnotation(anno, isRuntime);
        if (returnType.getSort() == Type.OBJECT &&
            anno.equals("Lorg/jetbrains/annotations/NotNull;")) {
          myIsNotNull = true;
        }

        return av;
      }

      public void visitCode() {
        for (int p = 0; p < myNotNullParams.size(); ++p) {
          int var = ((access & Opcodes.ACC_STATIC) == 0) ? 1 : 0;
          int param = ((Integer)myNotNullParams.get(p)).intValue();
          for (int i = 0; i < param; ++i) {
            var += args[i].getSize();
          }
          mv.visitVarInsn(Opcodes.ALOAD, var);

          generateConditionalThrow("Argument " + param + " for @NotNull parameter of " + myClassName + "." + name + " must not be null",
                                   "java/lang/IllegalArgumentException");
        }
      }

      public void visitInsn(int opcode) {
        if (opcode == Opcodes.ARETURN && myIsNotNull) {
          mv.visitInsn(Opcodes.DUP);
          generateConditionalThrow("@NotNull method " + myClassName + "." + name + " must not return null",
                                   "java/lang/IllegalStateException");
        }

        mv.visitInsn(opcode);
      }

      private void generateConditionalThrow(final String descr, final String exceptionClass) {
        String exceptionParamClass = "(Ljava/lang/String;)V";
        Label end = new Label();
        mv.visitJumpInsn(Opcodes.IFNONNULL, end);

        mv.visitTypeInsn(Opcodes.NEW, exceptionClass);
        mv.visitInsn(Opcodes.DUP);
        mv.visitLdcInsn(descr);
        mv.visitMethodInsn(Opcodes.INVOKESPECIAL,
                           exceptionClass,
                           "<init>",
                           exceptionParamClass);
        mv.visitInsn(Opcodes.ATHROW);
        mv.visitLabel(end);

        myIsModification = true;
      }
    };
  }
}

@@ -1473,16 +1473,49 @@
 = false;
+%0A      public Label myThrowLabel;
 %0A%0A      
@@ -2818,33 +2818,172 @@
     
-generateConditionalThrow(
+Label end = new Label();%0A          mv.visitJumpInsn(Opcodes.IFNONNULL, end);%0A%0A          generateThrow(%22java/lang/IllegalArgumentException%22,%0A                        
 %22Arg
@@ -3075,24 +3075,31 @@
 ot be null%22,
+ end);%0A
 %0A           
@@ -3091,32 +3091,35 @@
 );%0A%0A        
+%7D%0A%0A
         
             
@@ -3106,35 +3106,281 @@
         
-                   
+if (myIsNotNull) %7B%0A          Label codeStart = new Label();%0A          mv.visitJumpInsn(Opcodes.GOTO, codeStart);%0A%0A          myThrowLabel = new Label();%0A          mv.visitLabel(myThrowLabel);%0A          //generate throw for method null return%0A          generateThrow(
 %22java/la
@@ -3385,32 +3385,29 @@
 lang/Illegal
-Argument
+State
 Exception%22);
@@ -3396,32 +3396,140 @@
 lStateException%22
+, %22@NotNull method %22 + myClassName + %22.%22 + name + %22 must not return null%22,%0A                        codeStart
 );%0A        %7D%0A   
@@ -3671,32 +3671,34 @@
 DUP);%0A          
+/*
 generateConditio
@@ -3849,16 +3849,76 @@
 ption%22);
+*/%0A          mv.visitJumpInsn(Opcodes.IFNULL, myThrowLabel);
 %0A       
@@ -3983,35 +3983,24 @@
 oid generate
-Conditional
 Throw(final 
@@ -4006,21 +4006,30 @@
  String 
-descr
+exceptionClass
 , final 
@@ -4035,30 +4035,38 @@
  String 
-exceptionClass
+descr, final Label end
 ) %7B%0A    
@@ -4127,92 +4127,8 @@
 V%22;%0A
-        Label end = new Label();%0A        mv.visitJumpInsn(Opcodes.IFNONNULL, end);%0A%0A
     

rework instrumenter