Instrument @NotNull assertions into the compiled JVM bytecode #SCL-22085

README.md

@@ -439,6 +439,26 @@ Generate JVM bytecode to assert that a method is called on the correct IDEA thre
 
 See: [IntelliJ IDEA ThreadingAssertions.java](https://github.com/JetBrains/intellij-community/blob/5758eb99b4a1971ebe75cda755693cc930949465/platform/core-api/src/com/intellij/util/concurrency/ThreadingAssertions.java)
 
+#### `instrumentNotNullAnnotations :: SettingKey[Boolean]`
+
+**Default**: `false`
+
+Generate JVM bytecode to assert that values of `@NotNull`-annotated method parameters and return values are not null.
+A null value passed for an annotated parameter throws an `IllegalArgumentException`; a null value returned from an
+annotated method throws an `IllegalStateException`. The set of annotation classes is configured with the
+`notNullAnnotations` setting.
+
+This is the same instrumentation that IntelliJ IDEA applies to its own codebase and that the
+[IntelliJ Platform Gradle Plugin](https://plugins.jetbrains.com/docs/intellij/tools-intellij-platform-gradle-plugin-tasks.html#instrumentCode)
+performs as part of its `instrumentCode` task. Kotlin classes are skipped, as kotlinc generates its own
+nullability assertions.
+
+#### `notNullAnnotations :: SettingKey[Seq[String]]`
+
+**Default**: `Seq("org.jetbrains.annotations.NotNull")`
+
+Fully qualified names of the annotation classes instrumented by `instrumentNotNullAnnotations`.
+
 #### `packageOutputDir :: SettingKey[File]`
 
 **Default**: `target.value / "plugin" / intellijPluginName.in(ThisBuild).value.removeSpaces`

build.sbt

@@ -133,6 +133,9 @@ lazy val ideaSupport = (project in file("ideaSupport"))
       "org.ow2.asm" % "asm" % "9.10.1",
       "io.get-coursier" %% "coursier" % "2.1.24",
       "commons-io" % "commons-io" % "2.22.0",
+
+      // Used to compile the test fixtures of the @NotNull instrumentation tests
+      "org.jetbrains" % "annotations" % "26.1.0" % Test
     ),
   )
 

ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/AuxiliaryMethodGenerator.java

@@ -0,0 +1,282 @@
+// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package com.intellij.compiler.notNullVerification.
+// The only changes are the package name, the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm)
+// and Opcodes.API_VERSION (a JetBrains ASM addition) -> Opcodes.ASM9.
+package org.jetbrains.sbtidea.instrumentation.notNullVerification;
+
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.objectweb.asm.Opcodes.AASTORE;
+import static org.objectweb.asm.Opcodes.ACC_INTERFACE;
+import static org.objectweb.asm.Opcodes.ACC_PRIVATE;
+import static org.objectweb.asm.Opcodes.ACC_STATIC;
+import static org.objectweb.asm.Opcodes.ACC_SYNTHETIC;
+import static org.objectweb.asm.Opcodes.ANEWARRAY;
+import static org.objectweb.asm.Opcodes.ASM9;
+import static org.objectweb.asm.Opcodes.ATHROW;
+import static org.objectweb.asm.Opcodes.BIPUSH;
+import static org.objectweb.asm.Opcodes.DUP;
+import static org.objectweb.asm.Opcodes.DUP_X1;
+import static org.objectweb.asm.Opcodes.GOTO;
+import static org.objectweb.asm.Opcodes.ICONST_0;
+import static org.objectweb.asm.Opcodes.ICONST_1;
+import static org.objectweb.asm.Opcodes.ICONST_2;
+import static org.objectweb.asm.Opcodes.ICONST_3;
+import static org.objectweb.asm.Opcodes.ICONST_4;
+import static org.objectweb.asm.Opcodes.ICONST_5;
+import static org.objectweb.asm.Opcodes.ILOAD;
+import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
+import static org.objectweb.asm.Opcodes.INVOKESTATIC;
+import static org.objectweb.asm.Opcodes.NEW;
+import static org.objectweb.asm.Opcodes.SIPUSH;
+import static org.objectweb.asm.Opcodes.SWAP;
+
+final class AuxiliaryMethodGenerator {
+  private static final String STRING_CLASS_NAME = "java/lang/String";
+  private static final String OBJECT_CLASS_NAME = "java/lang/Object";
+  private static final String CONSTRUCTOR_NAME = "<init>";
+  private static final String EXCEPTION_INIT_SIGNATURE = "(L" + STRING_CLASS_NAME + ";)V";
+  private static final String REPORTING_METHOD_DESC = "(I)V";
+
+  private final ClassReader myOriginalClass;
+  private final boolean myIsInterface;
+  private final List<ReportingPlace> myReportingPlaces = new ArrayList<>();
+  private String myReportingMethod;
+  private int myMaxArgCount;
+
+  AuxiliaryMethodGenerator(ClassReader originalClass) {
+    myOriginalClass = originalClass;
+    myIsInterface = (myOriginalClass.getAccess() & ACC_INTERFACE) == ACC_INTERFACE;
+  }
+
+  private String getReportingMethodName() {
+    if (myReportingMethod == null) {
+      myReportingMethod = suggestUniqueName();
+    }
+    return myReportingMethod;
+  }
+
+  private String suggestUniqueName() {
+    Set<String> existingMethods = populateExistingMethods();
+    for (int i = 0;; i++) {
+      String name = "$$$reportNull$$$" + i;
+      if (!existingMethods.contains(name)) {
+        return name;
+      }
+    }
+  }
+
+  private Set<String> populateExistingMethods() {
+    final Set<String> existingMethods = new HashSet<>();
+    myOriginalClass.accept(new ClassVisitor(ASM9) {
+      @Override
+      public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
+        existingMethods.add(name);
+        return null;
+      }
+    }, ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
+    return existingMethods;
+  }
+
+  void reportError(MethodVisitor mv, String className, String exceptionClass, String descrPattern, String[] args) {
+    myMaxArgCount = Math.max(myMaxArgCount, args.length);
+
+    int index = myReportingPlaces.size();
+    myReportingPlaces.add(new ReportingPlace(exceptionClass, descrPattern, args));
+    pushIntConstant(mv, index);
+
+    mv.visitMethodInsn(INVOKESTATIC, className, getReportingMethodName(), REPORTING_METHOD_DESC, myIsInterface);
+  }
+
+  private static void pushIntConstant(MethodVisitor mv, int i) {
+    if (i <= 5) {
+      mv.visitInsn(getSmallIntConstantInstruction(i));
+    }
+    else if (i <= Byte.MAX_VALUE) {
+      mv.visitIntInsn(BIPUSH, i);
+    }
+    else {
+      mv.visitIntInsn(SIPUSH, i);
+    }
+  }
+
+  private static int getSmallIntConstantInstruction(int i) {
+    switch (i) {
+      case 0: return ICONST_0;
+      case 1: return ICONST_1;
+      case 2: return ICONST_2;
+      case 3: return ICONST_3;
+      case 4: return ICONST_4;
+      case 5: return ICONST_5;
+      default: throw new AssertionError(i);
+    }
+  }
+
+  void generateReportingMethod(ClassVisitor cw) {
+    if (myReportingPlaces.isEmpty()) return;
+
+    MethodVisitor mv = cw.visitMethod(ACC_PRIVATE | ACC_SYNTHETIC | ACC_STATIC, getReportingMethodName(), REPORTING_METHOD_DESC, null, null);
+    pushExceptionMessage(mv);
+    createExceptionObject(mv);
+    mv.visitInsn(ATHROW);
+    mv.visitMaxs(0, 0);
+  }
+
+  private void createExceptionObject(final MethodVisitor mv) {
+    new SwitchGenerator<String>() {
+      @Override
+      void generateCaseBody(String exceptionClass) {
+        mv.visitTypeInsn(NEW, exceptionClass);
+        mv.visitInsn(DUP_X1);
+        mv.visitInsn(SWAP);
+
+        mv.visitMethodInsn(INVOKESPECIAL, exceptionClass, CONSTRUCTOR_NAME, EXCEPTION_INIT_SIGNATURE, false);
+      }
+
+      @Override
+      String getSwitchedValue(ReportingPlace place) {
+        return place.exceptionClass;
+      }
+    }.generateSwitch(mv);
+  }
+
+  private void pushExceptionMessage(MethodVisitor mv) {
+    pushFormatPattern(mv);
+
+    createFormatArgArray(mv);
+    for (int i = 0; i < myMaxArgCount; i++) {
+      pushFormatArg(mv, i);
+    }
+
+    //noinspection SpellCheckingInspection
+    mv.visitMethodInsn(INVOKESTATIC, STRING_CLASS_NAME, "format", "(Ljava/lang/String;[Ljava/lang/Object;)Ljava/lang/String;", false);
+  }
+
+  private void createFormatArgArray(final MethodVisitor mv) {
+    new SwitchGenerator<Integer>(){
+      @Override
+      void generateCaseBody(Integer argCount) {
+        pushIntConstant(mv, argCount);
+      }
+
+      @Override
+      Integer getSwitchedValue(ReportingPlace place) {
+        return place.args.length;
+      }
+    }.generateSwitch(mv);
+    mv.visitTypeInsn(ANEWARRAY, OBJECT_CLASS_NAME);
+  }
+
+  private void pushFormatArg(final MethodVisitor mv, final int index) {
+    new SwitchGenerator<String>() {
+      @Override
+      protected String getSwitchedValue(ReportingPlace place) {
+        return index < place.args.length ? place.args[index] : null;
+      }
+
+      @Override
+      void generateCaseBody(String value) {
+        if (value != null) {
+          mv.visitInsn(DUP);
+          pushIntConstant(mv, index);
+          mv.visitLdcInsn(value);
+          mv.visitInsn(AASTORE);
+        }
+      }
+    }.generateSwitch(mv);
+  }
+
+  private void pushFormatPattern(final MethodVisitor mv) {
+    new SwitchGenerator<String>() {
+      @Override
+      protected String getSwitchedValue(ReportingPlace place) {
+        return place.descrPattern;
+      }
+
+      @Override
+      void generateCaseBody(String descrPattern) {
+        mv.visitLdcInsn(descrPattern);
+      }
+    }.generateSwitch(mv);
+  }
+
+  private static class ReportingPlace {
+    final String exceptionClass;
+    final String descrPattern;
+    final String[] args;
+
+    ReportingPlace(String exceptionClass, String descrPattern, String[] args) {
+      this.exceptionClass = exceptionClass;
+      this.descrPattern = descrPattern;
+      this.args = args;
+    }
+  }
+
+  private abstract class SwitchGenerator<T> {
+    void generateSwitch(MethodVisitor mv) {
+      Label[] labels = getCaseLabels();
+      if (labels == null) {
+        // all places behave in a same way, don't bother with switch
+        generateCaseBody(getSwitchedValue(myReportingPlaces.get(0)));
+      } else {
+        reallyGenerateSwitch(mv, labels, deduplicateLabels(labels));
+      }
+    }
+
+    private void reallyGenerateSwitch(MethodVisitor mv, Label[] labels, Map<Label, ReportingPlace> label2Place) {
+      Label afterSwitch = new Label();
+      mv.visitVarInsn(ILOAD, 0);
+      mv.visitTableSwitchInsn(0, labels.length - 1 , labels[0], labels);
+
+      for (Map.Entry<Label, ReportingPlace> entry : label2Place.entrySet()) {
+        mv.visitLabel(entry.getKey());
+        generateCaseBody(getSwitchedValue(entry.getValue()));
+        mv.visitJumpInsn(GOTO, afterSwitch);
+      }
+
+      mv.visitLabel(afterSwitch);
+    }
+
+    private Map<Label, ReportingPlace> deduplicateLabels(Label[] labels) {
+      Map<Label, ReportingPlace> label2Place = new LinkedHashMap<>();
+      for (int i = 0; i < labels.length; i++) {
+        if (!label2Place.containsKey(labels[i])) {
+          label2Place.put(labels[i], myReportingPlaces.get(i));
+        }
+      }
+      return label2Place;
+    }
+
+    private Label[] getCaseLabels() {
+      Map<T, Label> labelsByValue = new HashMap<>();
+      Label[] labels = new Label[myReportingPlaces.size()];
+      for (int i = 0; i < myReportingPlaces.size(); i++) {
+        labels[i] = getOrCreateLabel(labelsByValue, getSwitchedValue(myReportingPlaces.get(i)));
+      }
+      return labelsByValue.size() == 1 ? null : labels;
+    }
+
+    private Label getOrCreateLabel(Map<T, Label> labelsByValue, T key) {
+      Label label = labelsByValue.get(key);
+      if (label == null) {
+        labelsByValue.put(key, label = new Label());
+      }
+      return label;
+    }
+
+    abstract void generateCaseBody(T switchedValue);
+
+    abstract T getSwitchedValue(ReportingPlace place);
+  }
+}

ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeClassReader.java

@@ -0,0 +1,43 @@
+// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package com.intellij.compiler.instrumentation.
+// The only changes are the package name and the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm).
+package org.jetbrains.sbtidea.instrumentation.notNullVerification;
+
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.Label;
+
+import java.io.IOException;
+import java.io.InputStream;
+
+public final class FailSafeClassReader extends ClassReader {
+  private static final Label INVALID = new Label();
+
+  public FailSafeClassReader(byte[] b) {
+    super(b);
+  }
+
+  public FailSafeClassReader(byte[] b, int off, int len) {
+    super(b, off, len);
+  }
+
+  public FailSafeClassReader(InputStream is) throws IOException {
+    super(is);
+  }
+
+  public FailSafeClassReader(String name) throws IOException {
+    super(name);
+  }
+
+  @Override
+  protected Label readLabel(int offset, Label[] labels) {
+    // attempt to workaround javac bug:
+    // annotation table from original method is duplicated for synthetic bridge methods.
+    // All offsets in the duplicated table is taken from original annotations table and obviously are not relevant for the bridge method
+    if (offset >= 0 && offset < labels.length) {
+      return super.readLabel(offset, labels);
+    }
+    else {
+      return INVALID;
+    }
+  }
+}

ideaSupport/src/main/java/org/jetbrains/sbtidea/instrumentation/notNullVerification/FailSafeMethodVisitor.java

@@ -0,0 +1,35 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+// Copied from IntelliJ IDEA Community Edition (commit 7a35a7d7fe64), original package com.intellij.compiler.instrumentation.
+// The only changes are the package name and the ASM package (org.jetbrains.org.objectweb.asm -> org.objectweb.asm).
+package org.jetbrains.sbtidea.instrumentation.notNullVerification;
+
+import org.objectweb.asm.AnnotationVisitor;
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.TypePath;
+
+/**
+ * To be used together with FailSafeClassReader: adds null checks for labels describing annotation visibility range.
+ * For incorrectly generated annotations FailSafeClassReader returns null labels. Local variables annotations with null labels
+ * will be ignored by this visitor.
+ */
+public class FailSafeMethodVisitor extends MethodVisitor {
+  public FailSafeMethodVisitor(int api, MethodVisitor mv) {
+    super(api, mv);
+  }
+
+  @Override
+  public AnnotationVisitor visitLocalVariableAnnotation(int typeRef, TypePath typePath, Label[] start, Label[] end, int[] index, String desc, boolean visible) {
+    for (Label aStart : start) {
+      if (aStart == null) {
+        return null;
+      }
+    }
+    for (Label anEnd : end) {
+      if (anEnd == null) {
+        return null;
+      }
+    }
+    return super.visitLocalVariableAnnotation(typeRef, typePath, start, end, index, desc, visible);
+  }
+}