adding logic for using addMetadata and dropMetadata through methodHandle

uber · Nov 25, 2023 · ef844c8 · ef844c8
1 parent 89474f7
commit ef844c8
Showing 1 changed file with 51 additions and 1 deletion.
diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java b/nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java
@@ -80,7 +80,7 @@ public Type visitParameterizedType(ParameterizedTypeTree tree, Void p) {
               : com.sun.tools.javac.util.List.nil();
       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 =
@@ -97,6 +97,8 @@ protected Type defaultAction(Tree node, Void unused) {
 
   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 {
@@ -105,11 +107,18 @@ private static class JDK17AndEarlierTypeMetadataBuilder implements TypeMetadataB
     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 cloneAddMetadataHandle = createAddMethodHandle();
+    private static final MethodHandle cloneDropMetadataHandle = createDropMethodHandle();
 
     private static MethodHandle createHandle() {
       MethodHandles.Lookup lookup = MethodHandles.lookup();
@@ -123,6 +132,30 @@ private static MethodHandle createHandle() {
       }
     }
 
+    private static MethodHandle createAddMethodHandle() {
+      MethodHandles.Lookup lookup = MethodHandles.lookup();
+      MethodType mt = MethodType.methodType(Type.class, TypeMetadata.class);
+      try {
+        return lookup.findVirtual(com.sun.tools.javac.code.Type.class, "addMetadata", mt);
+      } catch (NoSuchMethodException e) {
+        throw new RuntimeException(e);
+      } catch (IllegalAccessException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    private static MethodHandle createDropMethodHandle() {
+      MethodHandles.Lookup lookup = MethodHandles.lookup();
+      MethodType mt = MethodType.methodType(Type.class, Class.class);
+      try {
+        return lookup.findVirtual(com.sun.tools.javac.code.Type.class, "dropMetadata", 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<>();
@@ -133,6 +166,23 @@ public TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound>
         throw new RuntimeException(e);
       }
     }
+
+    @Override
+    public Type cloneTypeWithMetaData(Type typeToBeCloned, TypeMetadata metaData) {
+      try {
+        Type clonedType = null;
+        if (metaData.getClass().getComponentType() != null) {
+          clonedType = (Type) cloneAddMetadataHandle.invoke(typeToBeCloned, metaData);
+        } else {
+          Type clonedTypeDrop =
+              (Type) cloneDropMetadataHandle.invoke(typeToBeCloned, metaData.getClass());
+          clonedType = (Type) cloneAddMetadataHandle.invoke(clonedTypeDrop, metaData);
+        }
+        return clonedType;
+      } catch (Throwable e) {
+        throw new RuntimeException(e);
+      }
+    }
   }
 
   private static final TypeMetadataBuilder tmBuilder =