Update handling of annotations in presence of kept graph consumer

Change-Id: Ib30f6f71a04afa0d9bc7a3467830d570f7231cd8
diff --git a/src/main/java/com/android/tools/r8/experimental/graphinfo/AnnotationGraphNode.java b/src/main/java/com/android/tools/r8/experimental/graphinfo/AnnotationGraphNode.java
index 5242510..7baa5fa 100644
--- a/src/main/java/com/android/tools/r8/experimental/graphinfo/AnnotationGraphNode.java
+++ b/src/main/java/com/android/tools/r8/experimental/graphinfo/AnnotationGraphNode.java
@@ -4,31 +4,44 @@
 package com.android.tools.r8.experimental.graphinfo;
 
 import com.android.tools.r8.Keep;
+import java.util.Objects;
 
 @Keep
 public final class AnnotationGraphNode extends GraphNode {
 
   private final GraphNode annotatedNode;
+  private final ClassGraphNode annotationClassNode;
 
-  public AnnotationGraphNode(GraphNode annotatedNode) {
+  public AnnotationGraphNode(GraphNode annotatedNode, ClassGraphNode annotationClassNode) {
     super(annotatedNode.isLibraryNode());
     this.annotatedNode = annotatedNode;
+    this.annotationClassNode = annotationClassNode;
   }
 
   public GraphNode getAnnotatedNode() {
     return annotatedNode;
   }
 
+  public ClassGraphNode getAnnotationClassNode() {
+    return annotationClassNode;
+  }
+
   @Override
   public boolean equals(Object o) {
-    return this == o
-        || (o instanceof AnnotationGraphNode
-            && ((AnnotationGraphNode) o).annotatedNode.equals(annotatedNode));
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof AnnotationGraphNode)) {
+      return false;
+    }
+    AnnotationGraphNode node = (AnnotationGraphNode) o;
+    return annotatedNode.equals(node.annotatedNode)
+        && annotationClassNode.equals(node.annotationClassNode);
   }
 
   @Override
   public int hashCode() {
-    return 7 * annotatedNode.hashCode();
+    return Objects.hash(annotatedNode, annotationClassNode);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
index edb24b0..8b7c63a 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -42,7 +42,6 @@
 import com.android.tools.r8.graph.DexEncodedMember;
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexField;
-import com.android.tools.r8.graph.DexItem;
 import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexItemFactory.ClassMethods;
 import com.android.tools.r8.graph.DexLibraryClass;
@@ -1698,7 +1697,7 @@
       return;
     }
     if (!type.isClassType()) {
-      // Ignore primitive types.
+      // Ignore primitive types and void.
       return;
     }
     DexProgramClass clazz = getProgramClassOrNull(type, context);
@@ -1961,11 +1960,14 @@
       }
       return;
     }
-    KeepReason reason = KeepReason.annotatedOn(annotatedItem.getDefinition());
-    graphReporter.registerAnnotation(annotation, reason);
+
+    // Report that the annotation is retained due to the annotated item.
+    graphReporter.registerAnnotation(annotation, annotatedItem);
+
+    // Report that the items referenced from inside the annotation are retained due to the
+    // annotation.
     AnnotationReferenceMarker referenceMarker =
-        new AnnotationReferenceMarker(
-            annotation.getAnnotationType(), annotatedItem, appView.dexItemFactory(), reason);
+        new AnnotationReferenceMarker(annotation, annotatedItem);
     annotation.annotation.collectIndexedItems(referenceMarker);
   }
 
@@ -4492,20 +4494,12 @@
 
   private class AnnotationReferenceMarker implements IndexedItemCollection {
 
-    private final DexItem annotationHolder;
     private final ProgramDefinition context;
-    private final DexItemFactory dexItemFactory;
     private final KeepReason reason;
 
-    private AnnotationReferenceMarker(
-        DexItem annotationHolder,
-        ProgramDefinition context,
-        DexItemFactory dexItemFactory,
-        KeepReason reason) {
-      this.annotationHolder = annotationHolder;
+    private AnnotationReferenceMarker(DexAnnotation annotation, ProgramDefinition context) {
       this.context = context;
-      this.dexItemFactory = dexItemFactory;
-      this.reason = reason;
+      this.reason = KeepReason.referencedInAnnotation(annotation, context);
     }
 
     @Override
@@ -4535,17 +4529,15 @@
                 : fieldAccessInfoCollection.extend(
                     fieldReference, new FieldAccessInfoImpl(fieldReference));
         fieldAccessInfo.setReadFromAnnotation();
-        markFieldAsLive(field, context, KeepReason.referencedInAnnotation(annotationHolder));
+        markFieldAsLive(field, context, reason);
         // When an annotation has a field of an enum type the JVM will use the values() method on
         // that enum class if the field is referenced.
         if (options.isGeneratingClassFiles() && field.getHolder().isEnum()) {
-          markEnumValuesAsReachable(
-              field.getHolder(), KeepReason.referencedInAnnotation(annotationHolder));
+          markEnumValuesAsReachable(field.getHolder(), reason);
         }
       } else {
         // There is no dispatch on annotations, so only keep what is directly referenced.
-        workList.enqueueMarkFieldAsReachableAction(
-            field, context, KeepReason.referencedInAnnotation(annotationHolder));
+        workList.enqueueMarkFieldAsReachableAction(field, context, reason);
       }
       return false;
     }
@@ -4562,17 +4554,13 @@
       if (target != null) {
         // There is no dispatch on annotations, so only keep what is directly referenced.
         if (target.getReference() == method) {
-          markDirectStaticOrConstructorMethodAsLive(
-              new ProgramMethod(holder, target),
-              KeepReason.referencedInAnnotation(annotationHolder));
+          markDirectStaticOrConstructorMethodAsLive(new ProgramMethod(holder, target), reason);
         }
       } else {
         target = holder.lookupVirtualMethod(method);
         // There is no dispatch on annotations, so only keep what is directly referenced.
         if (target != null && target.getReference() == method) {
-          markMethodAsTargeted(
-              new ProgramMethod(holder, target),
-              KeepReason.referencedInAnnotation(annotationHolder));
+          markMethodAsTargeted(new ProgramMethod(holder, target), reason);
         }
       }
       return false;
@@ -4600,11 +4588,7 @@
 
     @Override
     public boolean addType(DexType type) {
-      // Annotations can also contain the void type, which is not a class type, so filter it out
-      // here.
-      if (type != dexItemFactory.voidType) {
-        markTypeAsLive(type, context, reason);
-      }
+      markTypeAsLive(type, context, reason);
       return false;
     }
   }
diff --git a/src/main/java/com/android/tools/r8/shaking/GraphReporter.java b/src/main/java/com/android/tools/r8/shaking/GraphReporter.java
index 963e036..6146c2c 100644
--- a/src/main/java/com/android/tools/r8/shaking/GraphReporter.java
+++ b/src/main/java/com/android/tools/r8/shaking/GraphReporter.java
@@ -3,7 +3,6 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.shaking;
 
-import com.android.tools.r8.errors.Unimplemented;
 import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.experimental.graphinfo.AnnotationGraphNode;
 import com.android.tools.r8.experimental.graphinfo.ClassGraphNode;
@@ -21,7 +20,6 @@
 import com.android.tools.r8.graph.DexEncodedField;
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexField;
-import com.android.tools.r8.graph.DexItem;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexReference;
@@ -51,7 +49,7 @@
   private final CollectingGraphConsumer verificationGraphConsumer;
 
   // Canonicalization of external graph-nodes and edge info.
-  private final Map<DexItem, AnnotationGraphNode> annotationNodes = new IdentityHashMap<>();
+  private final Map<DexAnnotation, AnnotationGraphNode> annotationNodes = new IdentityHashMap<>();
   private final Map<DexType, ClassGraphNode> classNodes = new IdentityHashMap<>();
   private final Map<DexMethod, MethodGraphNode> methodNodes = new IdentityHashMap<>();
   private final Map<DexField, FieldGraphNode> fieldNodes = new IdentityHashMap<>();
@@ -359,11 +357,13 @@
     return registerEdge(getClassGraphNode(clazz.type), reason);
   }
 
-  public KeepReasonWitness registerAnnotation(DexAnnotation annotation, KeepReason reason) {
+  public KeepReasonWitness registerAnnotation(
+      DexAnnotation annotation, ProgramDefinition annotatedItem) {
+    KeepReason reason = KeepReason.annotatedOn(annotatedItem.getDefinition());
     if (skipReporting(reason)) {
       return KeepReasonWitness.INSTANCE;
     }
-    return registerEdge(getAnnotationGraphNode(annotation.annotation.type), reason);
+    return registerEdge(getAnnotationGraphNode(annotation, annotatedItem), reason);
   }
 
   public KeepReasonWitness registerMethod(DexEncodedMethod method, KeepReason reason) {
@@ -432,15 +432,18 @@
     return appView.appInfo().definitionForWithoutExistenceAssert(type);
   }
 
-  AnnotationGraphNode getAnnotationGraphNode(DexItem type) {
+  AnnotationGraphNode getAnnotationGraphNode(
+      DexAnnotation annotation, ProgramDefinition annotatedItem) {
     return annotationNodes.computeIfAbsent(
-        type,
-        t -> {
-          if (t instanceof DexType) {
-            return new AnnotationGraphNode(getClassGraphNode(((DexType) t)));
-          }
-          throw new Unimplemented(
-              "Incomplete support for annotation node on item: " + type.getClass());
+        annotation,
+        key -> {
+          GraphNode annotatedNode =
+              annotatedItem
+                  .getReference()
+                  .apply(
+                      this::getClassGraphNode, this::getFieldGraphNode, this::getMethodGraphNode);
+          ClassGraphNode annotationClassNode = getClassGraphNode(annotation.getAnnotationType());
+          return new AnnotationGraphNode(annotatedNode, annotationClassNode);
         });
   }
 
diff --git a/src/main/java/com/android/tools/r8/shaking/KeepReason.java b/src/main/java/com/android/tools/r8/shaking/KeepReason.java
index 5e6b3fd..2493dd2 100644
--- a/src/main/java/com/android/tools/r8/shaking/KeepReason.java
+++ b/src/main/java/com/android/tools/r8/shaking/KeepReason.java
@@ -6,12 +6,13 @@
 import com.android.tools.r8.experimental.graphinfo.GraphEdgeInfo;
 import com.android.tools.r8.experimental.graphinfo.GraphEdgeInfo.EdgeKind;
 import com.android.tools.r8.experimental.graphinfo.GraphNode;
+import com.android.tools.r8.graph.DexAnnotation;
 import com.android.tools.r8.graph.DexDefinition;
 import com.android.tools.r8.graph.DexEncodedMethod;
-import com.android.tools.r8.graph.DexItem;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.ProgramDefinition;
 import com.android.tools.r8.graph.ProgramMethod;
 
 // TODO(herhut): Canonicalize reason objects.
@@ -61,8 +62,9 @@
     return new ReferencedFrom(method.getDefinition());
   }
 
-  public static KeepReason referencedInAnnotation(DexItem holder) {
-    return new ReferencedInAnnotation(holder);
+  public static KeepReason referencedInAnnotation(
+      DexAnnotation annotation, ProgramDefinition annotatedItem) {
+    return new ReferencedInAnnotation(annotation, annotatedItem);
   }
 
   public boolean isDueToKeepRule() {
@@ -229,10 +231,12 @@
 
   private static class ReferencedInAnnotation extends KeepReason {
 
-    private final DexItem holder;
+    private final DexAnnotation annotation;
+    private final ProgramDefinition annotatedItem;
 
-    private ReferencedInAnnotation(DexItem holder) {
-      this.holder = holder;
+    private ReferencedInAnnotation(DexAnnotation annotation, ProgramDefinition annotatedItem) {
+      this.annotation = annotation;
+      this.annotatedItem = annotatedItem;
     }
 
     @Override
@@ -242,7 +246,7 @@
 
     @Override
     public GraphNode getSourceNode(GraphReporter graphReporter) {
-      return graphReporter.getAnnotationGraphNode(holder);
+      return graphReporter.getAnnotationGraphNode(annotation, annotatedItem);
     }
   }
 
diff --git a/src/test/java/com/android/tools/r8/utils/graphinspector/GraphInspector.java b/src/test/java/com/android/tools/r8/utils/graphinspector/GraphInspector.java
index b08651f..e51c162 100644
--- a/src/test/java/com/android/tools/r8/utils/graphinspector/GraphInspector.java
+++ b/src/test/java/com/android/tools/r8/utils/graphinspector/GraphInspector.java
@@ -11,6 +11,7 @@
 
 import com.android.tools.r8.errors.Unimplemented;
 import com.android.tools.r8.errors.Unreachable;
+import com.android.tools.r8.experimental.graphinfo.AnnotationGraphNode;
 import com.android.tools.r8.experimental.graphinfo.ClassGraphNode;
 import com.android.tools.r8.experimental.graphinfo.FieldGraphNode;
 import com.android.tools.r8.experimental.graphinfo.GraphEdgeInfo;
@@ -624,6 +625,14 @@
   private final Map<MethodReference, MethodGraphNode> methods;
   private final Map<FieldReference, FieldGraphNode> fields;
 
+  // Maps (annotated item, annotation type) to annotation node.
+  private final Map<ClassReference, Map<ClassReference, AnnotationGraphNode>> classAnnotations =
+      new HashMap<>();
+  private final Map<FieldReference, Map<ClassReference, AnnotationGraphNode>> fieldAnnotations =
+      new HashMap<>();
+  private final Map<MethodReference, Map<ClassReference, AnnotationGraphNode>> methodAnnotations =
+      new HashMap<>();
+
   public GraphInspector(CollectingGraphConsumer consumer, CodeInspector inspector) {
     this.consumer = consumer;
     this.inspector = inspector;
@@ -646,8 +655,30 @@
       } else if (target instanceof KeepRuleGraphNode) {
         KeepRuleGraphNode node = (KeepRuleGraphNode) target;
         rules.add(node);
+      } else if (target instanceof AnnotationGraphNode) {
+        AnnotationGraphNode node = (AnnotationGraphNode) target;
+        GraphNode annotatedNode = node.getAnnotatedNode();
+        Map<ClassReference, AnnotationGraphNode> annotationsOnAnnotatedNode;
+        if (annotatedNode instanceof ClassGraphNode) {
+          annotationsOnAnnotatedNode =
+              classAnnotations.computeIfAbsent(
+                  ((ClassGraphNode) annotatedNode).getReference(), key -> new HashMap<>());
+        } else if (annotatedNode instanceof FieldGraphNode) {
+          annotationsOnAnnotatedNode =
+              fieldAnnotations.computeIfAbsent(
+                  ((FieldGraphNode) annotatedNode).getReference(), key -> new HashMap<>());
+        } else if (annotatedNode instanceof MethodGraphNode) {
+          annotationsOnAnnotatedNode =
+              methodAnnotations.computeIfAbsent(
+                  ((MethodGraphNode) annotatedNode).getReference(), key -> new HashMap<>());
+        } else {
+          throw new Unreachable(
+              "Incomplete support for annotations on non-class, non-field, non-method items: "
+                  + annotatedNode.getClass().getTypeName());
+        }
+        annotationsOnAnnotatedNode.put(node.getAnnotationClassNode().getReference(), node);
       } else {
-        throw new Unimplemented("Incomplet support for graph node type: " + target.getClass());
+        throw new Unimplemented("Incomplete support for graph node type: " + target.getClass());
       }
       Map<GraphNode, Set<GraphEdgeInfo>> sources = consumer.getSourcesTargeting(target);
       for (GraphNode source : sources.keySet()) {