Fix JSpecify support on JDK 21

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,21 @@ 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 ->
+        outgoing.artifact(tasks.named(task.name).map { t -> t.extensions.getByType(JacocoTaskExtension).destinationFile })
+    }
+}

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 {

nullaway/src/main/java/com/uber/nullaway/NullAway.java

@@ -476,7 +476,7 @@ public Description matchAssignment(AssignmentTree tree, VisitorState state) {
       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);
     }
 
@@ -694,7 +694,9 @@ public Description matchParameterizedType(ParameterizedTypeTree tree, VisitorSta
     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;
   }
 
@@ -1423,7 +1425,7 @@ public Description matchVariable(VariableTree tree, VisitorState state) {
       return Description.NO_MATCH;
     }
     VarSymbol symbol = ASTHelpers.getSymbol(tree);
-    if (tree.getInitializer() != null) {
+    if (tree.getInitializer() != null && config.isJSpecifyMode()) {
       GenericsChecks.checkTypeParameterNullnessForAssignability(tree, this, state);
     }
 
@@ -1576,7 +1578,9 @@ public Description matchUnary(UnaryTree tree, VisitorState state) {
   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;
@@ -1707,8 +1711,10 @@ private Description handleInvocation(
                       : 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

nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java

@@ -14,6 +14,10 @@
 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,9 @@
                       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 = tmBuilder.create(nullableAnnotationCompound);
       Type currentTypeArgType = curTypeArg.accept(this, null);
-      Type newTypeArgType = currentTypeArgType.cloneWithMetadata(typeMetadata);
+      Type newTypeArgType = tmBuilder.cloneTypeWithMetadata(currentTypeArgType, typeMetadata);
       newTypeArgs.add(newTypeArgType);
     }
     Type.ClassType finalType =
@@ -91,4 +94,99 @@
   protected Type defaultAction(Tree node, Void unused) {
     return castToNonNull(ASTHelpers.getType(node));
   }
+
+  private interface TypeMetadataBuilder {
+    TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound> attrs);
+
+    Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metaData);
+  }
+
+  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));
+    }
+
+    @Override
+    public Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metadata) {
+      return typeToBeCloned.cloneWithMetadata(metadata);
+    }
+  }
+
+  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);
+      }
+    }
+
+    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);
+      }
+    }
+
+    @Override
+    public Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metadata) {
+      try {
+        // In Jdk21 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);
+      }
+    }
+  }
+
+  private static final TypeMetadataBuilder tmBuilder =
+      getVersion() >= 21
+          ? new JDK21TypeMetadataBuilder()
+          : new JDK17AndEarlierTypeMetadataBuilder();
+
+  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);
+  }
 }