Fix JSpecify support on JDK 21 (#869)

Our JSpecify mode would not run on JDK 21 due some some internal javac
API changes. This PR adapts to those changes using code specific to JDK
21.

Fixes #827

---------

Co-authored-by: Manu Sridharan <msridhar@gmail.com>

4 files changed, 180 insertions, 33 deletions

diff --git a/buildSrc/src/main/groovy/nullaway.java-test-conventions.gradle b/buildSrc/src/main/groovy/nullaway.java-test-conventions.gradle
index 07f8f26..6a0e641 100644
--- a/buildSrc/src/main/groovy/nullaway.java-test-conventions.gradle
+++ b/buildSrc/src/main/groovy/nullaway.java-test-conventions.gradle
@@ -44,22 +44,6 @@ configurations.create('transitiveSourcesElements') {
     }
 }
 
-// Share the coverage data to be aggregated for the whole product
-configurations.create('coverageDataElements') {
-    visible = false
-    canBeResolved = false
-    canBeConsumed = true
-    extendsFrom(configurations.implementation)
-    attributes {
-        attribute(Usage.USAGE_ATTRIBUTE, objects.named(Usage, Usage.JAVA_RUNTIME))
-        attribute(Category.CATEGORY_ATTRIBUTE, objects.named(Category, Category.DOCUMENTATION))
-        attribute(DocsType.DOCS_TYPE_ATTRIBUTE, objects.named(DocsType, 'jacoco-coverage-data'))
-    }
-    // This will cause the test task to run if the coverage data is requested by the aggregation task
-    outgoing.artifact(tasks.named("test").map { task ->
-        task.extensions.getByType(JacocoTaskExtension).destinationFile
-    })
-}
 
 test {
     maxHeapSize = "1024m"
@@ -118,3 +102,23 @@ def testJdk21 = tasks.register("testJdk21", Test) {
 tasks.named('check').configure {
     dependsOn testJdk21
 }
+
+
+// Share the coverage data to be aggregated for the whole product
+configurations.create('coverageDataElements') {
+    visible = false
+    canBeResolved = false
+    canBeConsumed = true
+    extendsFrom(configurations.implementation)
+    attributes {
+        attribute(Usage.USAGE_ATTRIBUTE, objects.named(Usage, Usage.JAVA_RUNTIME))
+        attribute(Category.CATEGORY_ATTRIBUTE, objects.named(Category, Category.DOCUMENTATION))
+        attribute(DocsType.DOCS_TYPE_ATTRIBUTE, objects.named(DocsType, 'jacoco-coverage-data'))
+    }
+    // This will cause the test tasks to run if the coverage data is requested by the aggregation task
+    tasks.withType(Test).forEach {task ->
+        // tasks.named(task.name) is a hack to introduce the right task dependence in Gradle.  We will fix this
+        // correctly when addressing https://github.com/uber/NullAway/issues/871
+        outgoing.artifact(tasks.named(task.name).map { t -> t.extensions.getByType(JacocoTaskExtension).destinationFile })
+    }
+}
diff --git a/nullaway/build.gradle b/nullaway/build.gradle
index 90ac062..829917f 100644
--- a/nullaway/build.gradle
+++ b/nullaway/build.gradle
@@ -124,14 +124,6 @@ tasks.named('check').configure {
     dependsOn(jdk8Test)
 }
 
-tasks.named('testJdk21', Test).configure {
-    filter {
-        // JSpecify Generics tests do not yet pass on JDK 21
-        // See https://github.com/uber/NullAway/issues/827
-        excludeTestsMatching "com.uber.nullaway.NullAwayJSpecifyGenericsTests"
-    }
-}
-
 // Create a task to build NullAway with NullAway checking enabled
 tasks.register('buildWithNullAway', JavaCompile) {
     onlyIf {
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index 925fab0..1ab3b12 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -476,7 +476,7 @@ public class NullAway extends BugChecker
       doUnboxingCheck(state, tree.getExpression());
     }
     // generics check
-    if (lhsType != null && lhsType.getTypeArguments().length() > 0) {
+    if (lhsType != null && lhsType.getTypeArguments().length() > 0 && config.isJSpecifyMode()) {
       GenericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
     }
 
@@ -695,7 +695,9 @@ public class NullAway extends BugChecker
     if (!withinAnnotatedCode(state)) {
       return Description.NO_MATCH;
     }
-    GenericsChecks.checkInstantiationForParameterizedTypedTree(tree, state, this, config);
+    if (config.isJSpecifyMode()) {
+      GenericsChecks.checkInstantiationForParameterizedTypedTree(tree, state, this, config);
+    }
     return Description.NO_MATCH;
   }
 
@@ -1424,7 +1426,7 @@ public class NullAway extends BugChecker
       return Description.NO_MATCH;
     }
     VarSymbol symbol = ASTHelpers.getSymbol(tree);
-    if (tree.getInitializer() != null) {
+    if (tree.getInitializer() != null && config.isJSpecifyMode()) {
       GenericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
     }
 
@@ -1577,7 +1579,9 @@ public class NullAway extends BugChecker
   public Description matchConditionalExpression(
       ConditionalExpressionTree tree, VisitorState state) {
     if (withinAnnotatedCode(state)) {
-      GenericsChecks.checkTypeParameterNullnessForConditionalExpression(tree, this, state);
+      if (config.isJSpecifyMode()) {
+        GenericsChecks.checkTypeParameterNullnessForConditionalExpression(tree, this, state);
+      }
       doUnboxingCheck(state, tree.getCondition());
     }
     return Description.NO_MATCH;
@@ -1708,8 +1712,10 @@ public class NullAway extends BugChecker
                       : Nullness.NONNULL);
         }
       }
-      GenericsChecks.compareGenericTypeParameterNullabilityForCall(
-          formalParams, actualParams, methodSymbol.isVarArgs(), this, state);
+      if (config.isJSpecifyMode()) {
+        GenericsChecks.compareGenericTypeParameterNullabilityForCall(
+            formalParams, actualParams, methodSymbol.isVarArgs(), this, state);
+      }
     }
 
     // Allow handlers to override the list of non-null argument positions
diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java b/nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java
index 2cde64f..822fcf1 100644
--- a/nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java
+++ b/nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java
@@ -14,6 +14,10 @@ import com.sun.source.util.SimpleTreeVisitor;
 import com.sun.tools.javac.code.Attribute;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.TypeMetadata;
+import com.sun.tools.javac.util.ListBuffer;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -74,10 +78,10 @@ public class PreservedAnnotationTreeVisitor extends SimpleTreeVisitor<Type, Void
                       new Attribute.TypeCompound(
                           nullableType, com.sun.tools.javac.util.List.nil(), null)))
               : com.sun.tools.javac.util.List.nil();
-      TypeMetadata typeMetadata =
-          new TypeMetadata(new TypeMetadata.Annotations(nullableAnnotationCompound));
+      TypeMetadata typeMetadata = TYPE_METADATA_BUILDER.create(nullableAnnotationCompound);
       Type currentTypeArgType = curTypeArg.accept(this, null);
-      Type newTypeArgType = currentTypeArgType.cloneWithMetadata(typeMetadata);
+      Type newTypeArgType =
+          TYPE_METADATA_BUILDER.cloneTypeWithMetadata(currentTypeArgType, typeMetadata);
       newTypeArgs.add(newTypeArgType);
     }
     Type.ClassType finalType =
@@ -91,4 +95,145 @@ public class PreservedAnnotationTreeVisitor extends SimpleTreeVisitor<Type, Void
   protected Type defaultAction(Tree node, Void unused) {
     return castToNonNull(ASTHelpers.getType(node));
   }
+
+  /**
+   * Abstracts over the different APIs for building {@link TypeMetadata} objects in different JDK
+   * versions.
+   */
+  private interface TypeMetadataBuilder {
+    TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound> attrs);
+
+    Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metaData);
+  }
+
+  /**
+   * Provides implementations for methods under TypeMetadataBuilder compatible with JDK 17 and
+   * earlier versions.
+   */
+  private static class JDK17AndEarlierTypeMetadataBuilder implements TypeMetadataBuilder {
+
+    @Override
+    public TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound> attrs) {
+      return new TypeMetadata(new TypeMetadata.Annotations(attrs));
+    }
+
+    /**
+     * Clones the given type with the specified Metadata for getting the right nullability
+     * annotations.
+     *
+     * @param typeToBeCloned The Type we want to clone with the required Nullability Metadata
+     * @param metadata The required Nullability metadata which is lost from the type
+     * @return Type after it has been cloned by applying the required Nullability metadata
+     */
+    @Override
+    public Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metadata) {
+      return typeToBeCloned.cloneWithMetadata(metadata);
+    }
+  }
+
+  /**
+   * Provides implementations for methods under TypeMetadataBuilder compatible with the updates made
+   * to the library methods for Jdk 21. The implementation calls the logic specific to JDK 21
+   * indirectly using MethodHandles since we still need the code to compile on earlier versions.
+   */
+  private static class JDK21TypeMetadataBuilder implements TypeMetadataBuilder {
+
+    private static final MethodHandle typeMetadataHandle = createHandle();
+    private static final MethodHandle addMetadataHandle =
+        createVirtualMethodHandle(Type.class, TypeMetadata.class, Type.class, "addMetadata");
+    private static final MethodHandle dropMetadataHandle =
+        createVirtualMethodHandle(Type.class, Class.class, Type.class, "dropMetadata");
+
+    private static MethodHandle createHandle() {
+      MethodHandles.Lookup lookup = MethodHandles.lookup();
+      MethodType mt = MethodType.methodType(void.class, com.sun.tools.javac.util.ListBuffer.class);
+      try {
+        return lookup.findConstructor(TypeMetadata.Annotations.class, mt);
+      } catch (NoSuchMethodException e) {
+        throw new RuntimeException(e);
+      } catch (IllegalAccessException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    /**
+     * Used to get a MethodHandle for a virtual method from the specified class
+     *
+     * @param retTypeClass Class to indicate the return type of the desired method
+     * @param paramTypeClass Class to indicate the parameter type of the desired method
+     * @param refClass Class within which the desired method is contained
+     * @param methodName Name of the desired method
+     * @return The appropriate MethodHandle for the virtual method
+     */
+    private static MethodHandle createVirtualMethodHandle(
+        Class<?> retTypeClass, Class<?> paramTypeClass, Class<?> refClass, String methodName) {
+      MethodHandles.Lookup lookup = MethodHandles.lookup();
+      MethodType mt = MethodType.methodType(retTypeClass, paramTypeClass);
+      try {
+        return lookup.findVirtual(refClass, methodName, mt);
+      } catch (NoSuchMethodException e) {
+        throw new RuntimeException(e);
+      } catch (IllegalAccessException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    @Override
+    public TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound> attrs) {
+      ListBuffer<Attribute.TypeCompound> b = new ListBuffer<>();
+      b.appendList(attrs);
+      try {
+        return (TypeMetadata) typeMetadataHandle.invoke(b);
+      } catch (Throwable e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    /**
+     * Calls dropMetadata and addMetadata using MethodHandles for JDK 21, which removed the previous
+     * cloneWithMetadata method.
+     *
+     * @param typeToBeCloned The Type we want to clone with the required Nullability metadata
+     * @param metadata The required Nullability metadata
+     * @return Cloned Type with the necessary Nullability metadata
+     */
+    @Override
+    public Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metadata) {
+      try {
+        // In JDK 21 addMetadata works if there is no metadata associated with the type, so we
+        // create a copy without the existing metadata first and then add it
+        Type clonedTypeWithoutMetadata =
+            (Type) dropMetadataHandle.invoke(typeToBeCloned, metadata.getClass());
+        return (Type) addMetadataHandle.invoke(clonedTypeWithoutMetadata, metadata);
+      } catch (Throwable e) {
+        throw new RuntimeException(e);
+      }
+    }
+  }
+
+  /** The TypeMetadataBuilder to be used for the current JDK version. */
+  private static final TypeMetadataBuilder TYPE_METADATA_BUILDER =
+      getVersion() >= 21
+          ? new JDK21TypeMetadataBuilder()
+          : new JDK17AndEarlierTypeMetadataBuilder();
+
+  /**
+   * Utility method to get the current JDK version, that works on Java 8 and above.
+   *
+   * <p>TODO remove this method once we drop support for Java 8
+   *
+   * @return the current JDK version
+   */
+  private static int getVersion() {
+    String version = System.getProperty("java.version");
+    if (version.startsWith("1.")) {
+      version = version.substring(2, 3);
+    } else {
+      int dot = version.indexOf(".");
+      if (dot != -1) {
+        version = version.substring(0, dot);
+      }
+    }
+    return Integer.parseInt(version);
+  }
 }