From ef844c80a1ad82d0699e4082748dfe995c49f6a2 Mon Sep 17 00:00:00 2001
From: Abhijit Kulkarni <akulk022@ucr.edu>
Date: Sat, 25 Nov 2023 14:12:57 -0800
Subject: [PATCH] adding logic for using addMetadata and dropMetadata through
 methodHandle

---
 .../PreservedAnnotationTreeVisitor.java       | 52 ++++++++++++++++++-
 1 file changed, 51 insertions(+), 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
index 84dc4e7b84..dfb3fe6774 100644
--- 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 =