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;

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

  public boolean isModification() {
    return myIsModification;
  }


  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 must not be null",
                                   "java/lang/IllegalArgumentException");
        }
      }

      public void visitInsn(int opcode) {
        if (opcode == Opcodes.ARETURN && myIsNotNull) {
          mv.visitInsn(Opcodes.DUP);
          generateConditionalThrow("@NotNull method 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;
      }
    };
  }

}

@@ -271,24 +271,54 @@
 ion = false;
+%0A  private String myClassName;
 %0A%0A  public N
@@ -480,16 +480,371 @@
 n;%0A  %7D%0A%0A
+  public void visit(final int version,%0A                    final int access,%0A                    final String name,%0A                    final String signature,%0A                    final String superName,%0A                    final String%5B%5D interfaces) %7B%0A    super.visit(version, access, name, signature, superName, interfaces);%0A    myClassName = name;%0A  %7D%0A
 %0A  publi
@@ -2852,16 +2852,52 @@
 arameter
+ of %22 + myClassName + %22.%22 + name + %22
  must no
@@ -3187,16 +3187,49 @@
 l method
+ %22 + myClassName + %22.%22 + name + %22
  must no
@@ -4026,11 +4026,10 @@
  %7D;%0A  %7D%0A
-%0A
 %7D%0A

include class and method name in instrumented @NotNull exceptions