Reland "Report contextual information in trace references"

This reverts commit 4c655b76f9d0c7ca1383d5af10bd6ea71d4fb0c2.

Bug: 186729231
Change-Id: I4f005ebddc46d801ebd765304f5c3127f91bfda9
diff --git a/src/main/java/com/android/tools/r8/desugar/desugaredlibrary/DesugaredLibraryKeepRuleGenerator.java b/src/main/java/com/android/tools/r8/desugar/desugaredlibrary/DesugaredLibraryKeepRuleGenerator.java
index fb4b86e..6386b61 100644
--- a/src/main/java/com/android/tools/r8/desugar/desugaredlibrary/DesugaredLibraryKeepRuleGenerator.java
+++ b/src/main/java/com/android/tools/r8/desugar/desugaredlibrary/DesugaredLibraryKeepRuleGenerator.java
@@ -131,7 +131,10 @@
       ClassReference rewrittenReference = rewrittenWithLens(tracedClass.getReference());
       super.acceptType(
           rewrittenReference != tracedClass.getReference()
-              ? new TracedClassImpl(rewrittenReference, tracedClass.getAccessFlags())
+              ? new TracedClassImpl(
+                  rewrittenReference,
+                  tracedClass.getReferencedFromContext(),
+                  tracedClass.getAccessFlags())
               : tracedClass,
           handler);
     }
@@ -141,7 +144,10 @@
       FieldReference rewrittenReference = rewrittenWithLens(tracedField.getReference());
       super.acceptField(
           rewrittenReference != tracedField.getReference()
-              ? new TracedFieldImpl(rewrittenReference, tracedField.getAccessFlags())
+              ? new TracedFieldImpl(
+                  rewrittenReference,
+                  tracedField.getReferencedFromContext(),
+                  tracedField.getAccessFlags())
               : tracedField,
           handler);
     }
@@ -151,7 +157,10 @@
       MethodReference rewrittenReference = rewrittenWithLens(tracedMethod.getReference());
       super.acceptMethod(
           rewrittenReference != tracedMethod.getReference()
-              ? new TracedMethodImpl(rewrittenReference, tracedMethod.getAccessFlags())
+              ? new TracedMethodImpl(
+                  rewrittenReference,
+                  tracedMethod.getReferencedFromContext(),
+                  tracedMethod.getAccessFlags())
               : tracedMethod,
           handler);
     }
diff --git a/src/main/java/com/android/tools/r8/diagnostic/internal/MissingClassInfoImpl.java b/src/main/java/com/android/tools/r8/diagnostic/internal/MissingClassInfoImpl.java
index 82d5a35..011bfdd 100644
--- a/src/main/java/com/android/tools/r8/diagnostic/internal/MissingClassInfoImpl.java
+++ b/src/main/java/com/android/tools/r8/diagnostic/internal/MissingClassInfoImpl.java
@@ -29,7 +29,7 @@
     return classReference;
   }
 
-  public static class Builder extends MissingDefinitionInfoBase.Builder {
+  public static class Builder extends MissingDefinitionInfoBase.Builder<Builder> {
 
     private ClassReference classReference;
 
@@ -43,5 +43,10 @@
     public MissingDefinitionInfo build() {
       return new MissingClassInfoImpl(classReference, referencedFromContextsBuilder.build());
     }
+
+    @Override
+    Builder self() {
+      return this;
+    }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/diagnostic/internal/MissingDefinitionInfoBase.java b/src/main/java/com/android/tools/r8/diagnostic/internal/MissingDefinitionInfoBase.java
index 8d90a92..57a6de3 100644
--- a/src/main/java/com/android/tools/r8/diagnostic/internal/MissingDefinitionInfoBase.java
+++ b/src/main/java/com/android/tools/r8/diagnostic/internal/MissingDefinitionInfoBase.java
@@ -29,16 +29,25 @@
     return referencedFromContexts;
   }
 
-  public abstract static class Builder {
+  public abstract static class Builder<B extends Builder<B>> {
 
     final ImmutableList.Builder<DefinitionContext> referencedFromContextsBuilder =
         ImmutableList.builder();
 
     Builder() {}
 
-    public Builder addReferencedFromContext(DefinitionContext missingDefinitionContext) {
+    public B addReferencedFromContext(DefinitionContext missingDefinitionContext) {
       referencedFromContextsBuilder.add(missingDefinitionContext);
-      return this;
+      return self();
     }
+
+    public B addReferencedFromContexts(Iterable<DefinitionContext> missingDefinitionContexts) {
+      for (DefinitionContext missingDefinitionContext : missingDefinitionContexts) {
+        addReferencedFromContext(missingDefinitionContext);
+      }
+      return self();
+    }
+
+    abstract B self();
   }
 }
diff --git a/src/main/java/com/android/tools/r8/diagnostic/internal/MissingFieldInfoImpl.java b/src/main/java/com/android/tools/r8/diagnostic/internal/MissingFieldInfoImpl.java
index d7acc21..679158d 100644
--- a/src/main/java/com/android/tools/r8/diagnostic/internal/MissingFieldInfoImpl.java
+++ b/src/main/java/com/android/tools/r8/diagnostic/internal/MissingFieldInfoImpl.java
@@ -29,7 +29,7 @@
     return fieldReference;
   }
 
-  public static class Builder extends MissingDefinitionInfoBase.Builder {
+  public static class Builder extends MissingDefinitionInfoBase.Builder<Builder> {
 
     private FieldReference fieldReference;
 
@@ -43,5 +43,10 @@
     public MissingDefinitionInfo build() {
       return new MissingFieldInfoImpl(fieldReference, referencedFromContextsBuilder.build());
     }
+
+    @Override
+    Builder self() {
+      return this;
+    }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/diagnostic/internal/MissingMethodInfoImpl.java b/src/main/java/com/android/tools/r8/diagnostic/internal/MissingMethodInfoImpl.java
index 8ab85d9..39f7e20 100644
--- a/src/main/java/com/android/tools/r8/diagnostic/internal/MissingMethodInfoImpl.java
+++ b/src/main/java/com/android/tools/r8/diagnostic/internal/MissingMethodInfoImpl.java
@@ -29,7 +29,7 @@
     return methodReference;
   }
 
-  public static class Builder extends MissingDefinitionInfoBase.Builder {
+  public static class Builder extends MissingDefinitionInfoBase.Builder<Builder> {
 
     private MethodReference methodReference;
 
@@ -43,5 +43,10 @@
     public MissingDefinitionInfo build() {
       return new MissingMethodInfoImpl(methodReference, referencedFromContextsBuilder.build());
     }
+
+    @Override
+    Builder self() {
+      return this;
+    }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/tracereferences/Formatter.java b/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
index fd188c0..2f19681 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/Formatter.java
@@ -10,6 +10,7 @@
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedClass;
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedField;
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedMethod;
+import com.android.tools.r8.tracereferences.internal.TraceReferencesResult;
 import com.android.tools.r8.utils.ListUtils;
 import com.android.tools.r8.utils.StringUtils;
 import com.android.tools.r8.utils.StringUtils.BraceType;
@@ -76,7 +77,11 @@
   protected abstract void printTypeFooter();
 
   void format(TraceReferencesResult result) {
-    print(result.types, result.keepPackageNames, result.fields, result.methods);
+    print(
+        result.getTracedClasses(),
+        result.getTracedPackageNames(),
+        result.getTracedFields(),
+        result.getTracedMethods());
   }
 
   private void print(
diff --git a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommand.java b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommand.java
index 1c98765..82915ee 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommand.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommand.java
@@ -20,6 +20,7 @@
 import com.android.tools.r8.ResourceException;
 import com.android.tools.r8.origin.Origin;
 import com.android.tools.r8.origin.PathOrigin;
+import com.android.tools.r8.tracereferences.internal.TraceReferencesCheckConsumer;
 import com.android.tools.r8.utils.ArchiveResourceProvider;
 import com.android.tools.r8.utils.Box;
 import com.android.tools.r8.utils.ExceptionDiagnostic;
@@ -330,7 +331,7 @@
     }
 
     public Builder setConsumer(TraceReferencesConsumer consumer) {
-      this.consumer = consumer;
+      this.consumer = new TraceReferencesCheckConsumer(consumer);
       return this;
     }
 
diff --git a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesConsumer.java b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesConsumer.java
index f364c72..e5251f7 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesConsumer.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesConsumer.java
@@ -6,6 +6,7 @@
 import com.android.tools.r8.DiagnosticsHandler;
 import com.android.tools.r8.Keep;
 import com.android.tools.r8.KeepForSubclassing;
+import com.android.tools.r8.diagnostic.DefinitionContext;
 import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.references.FieldReference;
 import com.android.tools.r8.references.MethodReference;
@@ -55,6 +56,9 @@
     /** Returns the reference traced. */
     T getReference();
 
+    /** Returns the context from which this was referenced. */
+    DefinitionContext getReferencedFromContext();
+
     /**
      * Returns the access flags for the reference traced. If the definition is not found (<code>
      * isMissingDefinition()</code> returns <code>true</code>) the access flags are not known and
diff --git a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesKeepRules.java b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesKeepRules.java
index d1926db..f648da4 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesKeepRules.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesKeepRules.java
@@ -7,6 +7,7 @@
 import com.android.tools.r8.Keep;
 import com.android.tools.r8.StringConsumer;
 import com.android.tools.r8.StringConsumer.FileConsumer;
+import com.android.tools.r8.tracereferences.internal.TraceReferencesResult;
 import java.nio.file.Path;
 
 /**
diff --git a/src/main/java/com/android/tools/r8/tracereferences/Tracer.java b/src/main/java/com/android/tools/r8/tracereferences/Tracer.java
index 3e8de35..21bb4be 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/Tracer.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/Tracer.java
@@ -6,10 +6,8 @@
 
 import com.android.tools.r8.DiagnosticsHandler;
 import com.android.tools.r8.dex.ApplicationReader;
-import com.android.tools.r8.diagnostic.internal.MissingClassInfoImpl;
-import com.android.tools.r8.diagnostic.internal.MissingDefinitionsDiagnosticImpl;
-import com.android.tools.r8.diagnostic.internal.MissingFieldInfoImpl;
-import com.android.tools.r8.diagnostic.internal.MissingMethodInfoImpl;
+import com.android.tools.r8.diagnostic.DefinitionContext;
+import com.android.tools.r8.diagnostic.internal.DefinitionContextUtils;
 import com.android.tools.r8.features.ClassToFeatureSplitMap;
 import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
 import com.android.tools.r8.graph.AppView;
@@ -17,7 +15,6 @@
 import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexClassAndField;
 import com.android.tools.r8.graph.DexClassAndMethod;
-import com.android.tools.r8.graph.DexEncodedField;
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexMethod;
@@ -30,6 +27,7 @@
 import com.android.tools.r8.graph.GraphLens.FieldLookupResult;
 import com.android.tools.r8.graph.GraphLens.MethodLookupResult;
 import com.android.tools.r8.graph.InitClassLens;
+import com.android.tools.r8.graph.ProgramField;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.graph.ResolutionResult;
 import com.android.tools.r8.graph.UseRegistry;
@@ -101,11 +99,12 @@
   public void run(TraceReferencesConsumer consumer) {
     UseCollector useCollector = new UseCollector(appInfo, consumer, diagnostics, targetPredicate);
     for (DexProgramClass clazz : appInfo.classes()) {
-      useCollector.registerSuperType(clazz, clazz.superType);
+      DefinitionContext classContext = DefinitionContextUtils.create(clazz);
+      useCollector.registerSuperType(clazz, clazz.superType, classContext);
       for (DexType implementsType : clazz.getInterfaces()) {
-        useCollector.registerSuperType(clazz, implementsType);
+        useCollector.registerSuperType(clazz, implementsType, classContext);
       }
-      clazz.forEachField(useCollector::registerField);
+      clazz.forEachProgramField(useCollector::registerField);
       clazz.forEachProgramMethod(
           method -> {
             useCollector.registerMethod(method);
@@ -113,7 +112,6 @@
           });
     }
     consumer.finished(diagnostics);
-    useCollector.reportMissingDefinitions();
   }
 
   // The graph lens is intentionally only made accessible to the MethodUseCollector, since the
@@ -146,37 +144,39 @@
       return targetPredicate.test(type);
     }
 
-    private void addType(DexType type) {
+    private void addType(DexType type, DefinitionContext referencedFrom) {
       if (type.isArrayType()) {
-        addType(type.toBaseType(factory));
+        addType(type.toBaseType(factory), referencedFrom);
         return;
       }
       if (type.isPrimitiveType() || type.isVoidType()) {
         return;
       }
       assert type.isClassType();
-      addClassType(type);
+      addClassType(type, referencedFrom);
     }
 
-    private void addTypes(DexTypeList types) {
-      types.forEach(this::addType);
+    private void addTypes(DexTypeList types, DefinitionContext referencedFrom) {
+      for (DexType type : types) {
+        addType(type, referencedFrom);
+      }
     }
 
-    private void addClassType(DexType type) {
+    private void addClassType(DexType type, DefinitionContext referencedFrom) {
       assert type.isClassType();
       DexClass clazz = appInfo.definitionFor(type);
       if (clazz != null) {
-        addClass(clazz);
+        addClass(clazz, referencedFrom);
       } else {
-        TracedClassImpl tracedClass = new TracedClassImpl(type);
+        TracedClassImpl tracedClass = new TracedClassImpl(type, referencedFrom);
         collectMissingClass(tracedClass);
         consumer.acceptType(tracedClass, diagnostics);
       }
     }
 
-    private void addClass(DexClass clazz) {
+    private void addClass(DexClass clazz, DefinitionContext referencedFrom) {
       if (isTargetType(clazz.getType())) {
-        TracedClassImpl tracedClass = new TracedClassImpl(clazz);
+        TracedClassImpl tracedClass = new TracedClassImpl(clazz, referencedFrom);
         consumer.acceptType(tracedClass, diagnostics);
         if (clazz.getAccessFlags().isVisibilityDependingOnPackage()) {
           consumer.acceptPackage(
@@ -185,7 +185,8 @@
       }
     }
 
-    private void addSuperMethodFromTarget(DexClassAndMethod method) {
+    private void addSuperMethodFromTarget(
+        DexClassAndMethod method, DefinitionContext referencedFrom) {
       assert !method.isProgramMethod();
       assert isTargetType(method.getHolderType());
 
@@ -195,7 +196,7 @@
       // - The holder type is registered from visiting the extends/implements clause of the sub
       //   class.
 
-      TracedMethodImpl tracedMethod = new TracedMethodImpl(method);
+      TracedMethodImpl tracedMethod = new TracedMethodImpl(method, referencedFrom);
       if (isTargetType(method.getHolderType())) {
         consumer.acceptMethod(tracedMethod, diagnostics);
         if (method.getAccessFlags().isVisibilityDependingOnPackage()) {
@@ -227,38 +228,20 @@
       collectMissing(tracedMethod, missingMethods);
     }
 
-    private void reportMissingDefinitions() {
-      if (missingClasses.size() > 0 || missingFields.size() > 0 || missingMethods.size() > 0) {
-        MissingDefinitionsDiagnosticImpl.Builder diagnosticBuilder =
-            MissingDefinitionsDiagnosticImpl.builder();
-        missingClasses.forEach(
-            classReference ->
-                diagnosticBuilder.addMissingDefinitionInfo(
-                    MissingClassInfoImpl.builder().setClass(classReference).build()));
-        missingFields.forEach(
-            fieldReference ->
-                diagnosticBuilder.addMissingDefinitionInfo(
-                    MissingFieldInfoImpl.builder().setField(fieldReference).build()));
-        missingMethods.forEach(
-            methodReference ->
-                diagnosticBuilder.addMissingDefinitionInfo(
-                    MissingMethodInfoImpl.builder().setMethod(methodReference).build()));
-        diagnostics.error(diagnosticBuilder.build());
-      }
-    }
-
-    private void registerField(DexEncodedField field) {
-      addType(field.getType());
+    private void registerField(ProgramField field) {
+      DefinitionContext referencedFrom = DefinitionContextUtils.create(field);
+      addType(field.getType(), referencedFrom);
     }
 
     private void registerMethod(ProgramMethod method) {
-      addTypes(method.getParameters());
-      addType(method.getReturnType());
+      DefinitionContext referencedFrom = DefinitionContextUtils.create(method);
+      addTypes(method.getParameters(), referencedFrom);
+      addType(method.getReturnType(), referencedFrom);
       for (DexAnnotation annotation : method.getDefinition().annotations().annotations) {
         if (annotation.getAnnotationType() == appInfo.dexItemFactory().annotationThrows) {
           DexValueArray dexValues = annotation.annotation.elements[0].value.asDexValueArray();
           for (DexValue dexValType : dexValues.getValues()) {
-            addType(dexValType.asDexValueType().value);
+            addType(dexValType.asDexValueType().value, referencedFrom);
           }
         }
       }
@@ -270,7 +253,7 @@
       if (superTarget != null
           && !superTarget.isProgramMethod()
           && isTargetType(superTarget.getHolderType())) {
-        addSuperMethodFromTarget(superTarget);
+        addSuperMethodFromTarget(superTarget, referencedFrom);
       }
     }
 
@@ -278,8 +261,9 @@
       method.registerCodeReferences(new MethodUseCollector(method, graphLens, initClassLens));
     }
 
-    private void registerSuperType(DexProgramClass clazz, DexType superType) {
-      addType(superType);
+    private void registerSuperType(
+        DexProgramClass clazz, DexType superType, DefinitionContext referencedFrom) {
+      addType(superType, referencedFrom);
       // If clazz overrides any methods in superType, we should keep those as well.
       clazz.forEachMethod(
           method -> {
@@ -290,7 +274,7 @@
             if (resolvedMethod != null
                 && !resolvedMethod.isProgramMethod()
                 && isTargetType(resolvedMethod.getHolderType())) {
-              addSuperMethodFromTarget(resolvedMethod);
+              addSuperMethodFromTarget(resolvedMethod, referencedFrom);
             }
           });
     }
@@ -300,6 +284,7 @@
       private final ProgramMethod context;
       private final GraphLens graphLens;
       private final InitClassLens initClassLens;
+      private final DefinitionContext referencedFrom;
 
       public MethodUseCollector(
           ProgramMethod context, GraphLens graphLens, InitClassLens initClassLens) {
@@ -307,6 +292,7 @@
         this.context = context;
         this.graphLens = graphLens;
         this.initClassLens = initClassLens;
+        this.referencedFrom = DefinitionContextUtils.create(context);
       }
 
       // Method references.
@@ -358,7 +344,7 @@
         DexMethod method = lookupResult.getReference();
         if (method.getHolderType().isArrayType()) {
           assert lookupResult.getType().isVirtual();
-          addType(method.getHolderType());
+          addType(method.getHolderType(), referencedFrom);
           return;
         }
         assert lookupResult.getType().isInterface() || lookupResult.getType().isVirtual();
@@ -374,15 +360,15 @@
       private void handleRewrittenMethodReference(
           DexMethod method, DexClassAndMethod resolvedMethod) {
         assert resolvedMethod == null || resolvedMethod.getReference().match(method);
-        addType(method.getHolderType());
-        addTypes(method.getParameters());
-        addType(method.getReturnType());
+        addType(method.getHolderType(), referencedFrom);
+        addTypes(method.getParameters(), referencedFrom);
+        addType(method.getReturnType(), referencedFrom);
         if (resolvedMethod != null) {
           if (isTargetType(resolvedMethod.getHolderType())) {
             if (resolvedMethod.getHolderType() != method.getHolderType()) {
-              addType(resolvedMethod.getHolderType());
+              addType(resolvedMethod.getHolderType(), referencedFrom);
             }
-            TracedMethodImpl tracedMethod = new TracedMethodImpl(resolvedMethod);
+            TracedMethodImpl tracedMethod = new TracedMethodImpl(resolvedMethod, referencedFrom);
             consumer.acceptMethod(tracedMethod, diagnostics);
             if (resolvedMethod.getAccessFlags().isVisibilityDependingOnPackage()) {
               consumer.acceptPackage(
@@ -391,7 +377,7 @@
             }
           }
         } else {
-          TracedMethodImpl tracedMethod = new TracedMethodImpl(method);
+          TracedMethodImpl tracedMethod = new TracedMethodImpl(method, referencedFrom);
           collectMissingMethod(tracedMethod);
           consumer.acceptMethod(tracedMethod, diagnostics);
         }
@@ -432,16 +418,16 @@
       }
 
       private void handleRewrittenFieldReference(DexField field) {
-        addType(field.getHolderType());
-        addType(field.getType());
+        addType(field.getHolderType(), referencedFrom);
+        addType(field.getType(), referencedFrom);
 
         DexClassAndField resolvedField = appInfo.resolveField(field).getResolutionPair();
         if (resolvedField != null) {
           if (isTargetType(resolvedField.getHolderType())) {
             if (resolvedField.getHolderType() != field.getHolderType()) {
-              addClass(resolvedField.getHolder());
+              addClass(resolvedField.getHolder(), referencedFrom);
             }
-            TracedFieldImpl tracedField = new TracedFieldImpl(resolvedField);
+            TracedFieldImpl tracedField = new TracedFieldImpl(resolvedField, referencedFrom);
             consumer.acceptField(tracedField, diagnostics);
             if (resolvedField.getAccessFlags().isVisibilityDependingOnPackage()) {
               consumer.acceptPackage(
@@ -450,7 +436,7 @@
             }
           }
         } else {
-          TracedFieldImpl tracedField = new TracedFieldImpl(field);
+          TracedFieldImpl tracedField = new TracedFieldImpl(field, referencedFrom);
           collectMissingField(tracedField);
           consumer.acceptField(tracedField, diagnostics);
         }
@@ -460,7 +446,7 @@
 
       @Override
       public void registerTypeReference(DexType type) {
-        addType(graphLens.lookupType(type));
+        addType(graphLens.lookupType(type), referencedFrom);
       }
     }
   }
diff --git a/src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesCheckConsumer.java b/src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesCheckConsumer.java
new file mode 100644
index 0000000..649efde
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesCheckConsumer.java
@@ -0,0 +1,134 @@
+// Copyright (c) 2021, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.tracereferences.internal;
+
+import com.android.tools.r8.DiagnosticsHandler;
+import com.android.tools.r8.diagnostic.DefinitionContext;
+import com.android.tools.r8.diagnostic.internal.DefinitionContextUtils;
+import com.android.tools.r8.diagnostic.internal.MissingClassInfoImpl;
+import com.android.tools.r8.diagnostic.internal.MissingDefinitionsDiagnosticImpl;
+import com.android.tools.r8.diagnostic.internal.MissingFieldInfoImpl;
+import com.android.tools.r8.diagnostic.internal.MissingMethodInfoImpl;
+import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.references.FieldReference;
+import com.android.tools.r8.references.MethodReference;
+import com.android.tools.r8.references.PackageReference;
+import com.android.tools.r8.tracereferences.TraceReferencesConsumer;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Collects the set of missing definitions and reports a {@link
+ * com.android.tools.r8.diagnostic.MissingDefinitionsDiagnostic}, if any missing definitions were
+ * found.
+ */
+public class TraceReferencesCheckConsumer extends TraceReferencesConsumer.ForwardingConsumer {
+
+  private final Map<ClassReference, Map<Object, DefinitionContext>> missingClassesContexts =
+      new ConcurrentHashMap<>();
+  private final Map<FieldReference, Map<Object, DefinitionContext>> missingFieldsContexts =
+      new ConcurrentHashMap<>();
+  private final Map<MethodReference, Map<Object, DefinitionContext>> missingMethodsContexts =
+      new ConcurrentHashMap<>();
+
+  public TraceReferencesCheckConsumer() {
+    this(TraceReferencesConsumer.emptyConsumer());
+  }
+
+  public TraceReferencesCheckConsumer(TraceReferencesConsumer consumer) {
+    super(consumer);
+  }
+
+  @Override
+  public void acceptType(TracedClass tracedClass, DiagnosticsHandler handler) {
+    super.acceptType(tracedClass, handler);
+    if (tracedClass.isMissingDefinition()) {
+      Map<Object, DefinitionContext> missingClassContexts =
+          missingClassesContexts.computeIfAbsent(
+              tracedClass.getReference(), ignore -> new ConcurrentHashMap<>());
+      DefinitionContextUtils.accept(
+          tracedClass.getReferencedFromContext(),
+          classContext -> missingClassContexts.put(classContext.getClassReference(), classContext),
+          fieldContext -> missingClassContexts.put(fieldContext.getFieldReference(), fieldContext),
+          methodContext ->
+              missingClassContexts.put(methodContext.getMethodReference(), methodContext));
+    }
+  }
+
+  @Override
+  public void acceptField(TracedField tracedField, DiagnosticsHandler handler) {
+    super.acceptField(tracedField, handler);
+    if (tracedField.isMissingDefinition()) {
+      Map<Object, DefinitionContext> missingFieldContexts =
+          missingFieldsContexts.computeIfAbsent(
+              tracedField.getReference(), ignore -> new ConcurrentHashMap<>());
+      DefinitionContextUtils.accept(
+          tracedField.getReferencedFromContext(),
+          classContext -> missingFieldContexts.put(classContext.getClassReference(), classContext),
+          fieldContext -> missingFieldContexts.put(fieldContext.getFieldReference(), fieldContext),
+          methodContext ->
+              missingFieldContexts.put(methodContext.getMethodReference(), methodContext));
+    }
+  }
+
+  @Override
+  public void acceptMethod(TracedMethod tracedMethod, DiagnosticsHandler handler) {
+    super.acceptMethod(tracedMethod, handler);
+    if (tracedMethod.isMissingDefinition()) {
+      Map<Object, DefinitionContext> missingMethodContexts =
+          missingMethodsContexts.computeIfAbsent(
+              tracedMethod.getReference(), ignore -> new ConcurrentHashMap<>());
+      DefinitionContextUtils.accept(
+          tracedMethod.getReferencedFromContext(),
+          classContext -> missingMethodContexts.put(classContext.getClassReference(), classContext),
+          fieldContext -> missingMethodContexts.put(fieldContext.getFieldReference(), fieldContext),
+          methodContext ->
+              missingMethodContexts.put(methodContext.getMethodReference(), methodContext));
+    }
+  }
+
+  @Override
+  public void acceptPackage(PackageReference pkg, DiagnosticsHandler handler) {
+    super.acceptPackage(pkg, handler);
+  }
+
+  @Override
+  public void finished(DiagnosticsHandler handler) {
+    super.finished(handler);
+    if (isEmpty()) {
+      return;
+    }
+    MissingDefinitionsDiagnosticImpl.Builder diagnosticBuilder =
+        MissingDefinitionsDiagnosticImpl.builder();
+    missingClassesContexts.forEach(
+        (reference, referencedFrom) ->
+            diagnosticBuilder.addMissingDefinitionInfo(
+                MissingClassInfoImpl.builder()
+                    .setClass(reference)
+                    .addReferencedFromContexts(referencedFrom.values())
+                    .build()));
+    missingFieldsContexts.forEach(
+        (reference, referencedFrom) ->
+            diagnosticBuilder.addMissingDefinitionInfo(
+                MissingFieldInfoImpl.builder()
+                    .setField(reference)
+                    .addReferencedFromContexts(referencedFrom.values())
+                    .build()));
+    missingMethodsContexts.forEach(
+        (reference, referencedFrom) ->
+            diagnosticBuilder.addMissingDefinitionInfo(
+                MissingMethodInfoImpl.builder()
+                    .setMethod(reference)
+                    .addReferencedFromContexts(referencedFrom.values())
+                    .build()));
+    handler.error(diagnosticBuilder.build());
+  }
+
+  private boolean isEmpty() {
+    return missingClassesContexts.isEmpty()
+        && missingFieldsContexts.isEmpty()
+        && missingMethodsContexts.isEmpty();
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesResult.java b/src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesResult.java
similarity index 73%
rename from src/main/java/com/android/tools/r8/tracereferences/TraceReferencesResult.java
rename to src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesResult.java
index 1b16d54..94ffff7 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesResult.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesResult.java
@@ -1,13 +1,14 @@
-// Copyright (c) 2020, the R8 project authors. Please see the AUTHORS file
+// Copyright (c) 2021, the R8 project authors. Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
-package com.android.tools.r8.tracereferences;
+package com.android.tools.r8.tracereferences.internal;
 
 import com.android.tools.r8.DiagnosticsHandler;
 import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.references.FieldReference;
 import com.android.tools.r8.references.MethodReference;
 import com.android.tools.r8.references.PackageReference;
+import com.android.tools.r8.tracereferences.TraceReferencesConsumer;
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedClass;
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedField;
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedMethod;
@@ -18,10 +19,10 @@
 
 public class TraceReferencesResult {
 
-  final Set<TracedClass> types;
-  final Map<ClassReference, Set<TracedField>> fields;
-  final Map<ClassReference, Set<TracedMethod>> methods;
-  final Set<PackageReference> keepPackageNames;
+  private final Set<TracedClass> types;
+  private final Map<ClassReference, Set<TracedField>> fields;
+  private final Map<ClassReference, Set<TracedMethod>> methods;
+  private final Set<PackageReference> keepPackageNames;
 
   TraceReferencesResult(
       Set<TracedClass> types,
@@ -34,11 +35,27 @@
     this.keepPackageNames = keepPackageNames;
   }
 
+  public Set<TracedClass> getTracedClasses() {
+    return types;
+  }
+
+  public Map<ClassReference, Set<TracedField>> getTracedFields() {
+    return fields;
+  }
+
+  public Map<ClassReference, Set<TracedMethod>> getTracedMethods() {
+    return methods;
+  }
+
+  public Set<PackageReference> getTracedPackageNames() {
+    return keepPackageNames;
+  }
+
   public static Builder builder() {
     return new Builder();
   }
 
-  static class Builder implements TraceReferencesConsumer {
+  public static class Builder implements TraceReferencesConsumer {
     private final Set<TracedClass> types = new HashSet<>();
     private final Map<ClassReference, Set<TracedField>> fields = new HashMap<>();
     private final Map<ClassReference, Set<TracedMethod>> methods = new HashMap<>();
@@ -69,7 +86,7 @@
     @Override
     public void finished(DiagnosticsHandler handler) {}
 
-    TraceReferencesResult build() {
+    public TraceReferencesResult build() {
       return new TraceReferencesResult(types, fields, methods, keepPackageNames);
     }
   }
diff --git a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedClassImpl.java b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedClassImpl.java
index fd9a353..c876442 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedClassImpl.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedClassImpl.java
@@ -4,6 +4,7 @@
 
 package com.android.tools.r8.tracereferences.internal;
 
+import com.android.tools.r8.diagnostic.DefinitionContext;
 import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.references.ClassReference;
@@ -12,16 +13,22 @@
 
 public class TracedClassImpl extends TracedReferenceBase<ClassReference, ClassAccessFlags>
     implements TracedClass {
-  public TracedClassImpl(DexType type) {
-    this(type.asClassReference(), null);
+  public TracedClassImpl(DexType type, DefinitionContext referencedFrom) {
+    this(type.asClassReference(), referencedFrom, null);
   }
 
-  public TracedClassImpl(DexClass clazz) {
-    this(clazz.getClassReference(), new ClassAccessFlagsImpl(clazz.getAccessFlags()));
+  public TracedClassImpl(DexClass clazz, DefinitionContext referencedFrom) {
+    this(
+        clazz.getClassReference(),
+        referencedFrom,
+        new ClassAccessFlagsImpl(clazz.getAccessFlags()));
   }
 
-  public TracedClassImpl(ClassReference classReference, ClassAccessFlags accessFlags) {
-    super(classReference, accessFlags, accessFlags == null);
+  public TracedClassImpl(
+      ClassReference classReference,
+      DefinitionContext referencedFrom,
+      ClassAccessFlags accessFlags) {
+    super(classReference, referencedFrom, accessFlags, accessFlags == null);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedFieldImpl.java b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedFieldImpl.java
index 9c6ac2d..59e58db 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedFieldImpl.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedFieldImpl.java
@@ -4,6 +4,7 @@
 
 package com.android.tools.r8.tracereferences.internal;
 
+import com.android.tools.r8.diagnostic.DefinitionContext;
 import com.android.tools.r8.graph.DexClassAndField;
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.references.FieldReference;
@@ -12,16 +13,22 @@
 
 public class TracedFieldImpl extends TracedReferenceBase<FieldReference, FieldAccessFlags>
     implements TracedField {
-  public TracedFieldImpl(DexField field) {
-    this(field.asFieldReference(), null);
+  public TracedFieldImpl(DexField field, DefinitionContext referencedFrom) {
+    this(field.asFieldReference(), referencedFrom, null);
   }
 
-  public TracedFieldImpl(DexClassAndField field) {
-    this(field.getFieldReference(), new FieldAccessFlagsImpl(field.getAccessFlags()));
+  public TracedFieldImpl(DexClassAndField field, DefinitionContext referencedFrom) {
+    this(
+        field.getFieldReference(),
+        referencedFrom,
+        new FieldAccessFlagsImpl(field.getAccessFlags()));
   }
 
-  public TracedFieldImpl(FieldReference fieldReference, FieldAccessFlags accessFlags) {
-    super(fieldReference, accessFlags, accessFlags == null);
+  public TracedFieldImpl(
+      FieldReference fieldReference,
+      DefinitionContext referencedFrom,
+      FieldAccessFlags accessFlags) {
+    super(fieldReference, referencedFrom, accessFlags, accessFlags == null);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java
index c12cbd4..1ea760a 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedMethodImpl.java
@@ -4,6 +4,7 @@
 
 package com.android.tools.r8.tracereferences.internal;
 
+import com.android.tools.r8.diagnostic.DefinitionContext;
 import com.android.tools.r8.graph.DexClassAndMethod;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.references.MethodReference;
@@ -12,16 +13,22 @@
 
 public class TracedMethodImpl extends TracedReferenceBase<MethodReference, MethodAccessFlags>
     implements TracedMethod {
-  public TracedMethodImpl(DexMethod method) {
-    this(method.asMethodReference(), null);
+  public TracedMethodImpl(DexMethod method, DefinitionContext referencedFrom) {
+    this(method.asMethodReference(), referencedFrom, null);
   }
 
-  public TracedMethodImpl(DexClassAndMethod method) {
-    this(method.getMethodReference(), new MethodAccessFlagsImpl(method.getAccessFlags()));
+  public TracedMethodImpl(DexClassAndMethod method, DefinitionContext referencedFrom) {
+    this(
+        method.getMethodReference(),
+        referencedFrom,
+        new MethodAccessFlagsImpl(method.getAccessFlags()));
   }
 
-  public TracedMethodImpl(MethodReference methodReference, MethodAccessFlags accessFlags) {
-    super(methodReference, accessFlags, accessFlags == null);
+  public TracedMethodImpl(
+      MethodReference methodReference,
+      DefinitionContext referencedFrom,
+      MethodAccessFlags accessFlags) {
+    super(methodReference, referencedFrom, accessFlags, accessFlags == null);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedReferenceBase.java b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedReferenceBase.java
index 254048b..d426b96 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/internal/TracedReferenceBase.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/internal/TracedReferenceBase.java
@@ -4,16 +4,21 @@
 
 package com.android.tools.r8.tracereferences.internal;
 
+import com.android.tools.r8.diagnostic.DefinitionContext;
 import com.android.tools.r8.tracereferences.TraceReferencesConsumer.TracedReference;
 
 abstract class TracedReferenceBase<T, F> implements TracedReference<T, F> {
+
   private final T reference;
+  private final DefinitionContext referencedFrom;
   private final F accessFlags;
   private final boolean missingDefinition;
 
-  TracedReferenceBase(T reference, F accessFlags, boolean missingDefinition) {
+  TracedReferenceBase(
+      T reference, DefinitionContext referencedFrom, F accessFlags, boolean missingDefinition) {
     assert accessFlags != null || missingDefinition;
     this.reference = reference;
+    this.referencedFrom = referencedFrom;
     this.accessFlags = accessFlags;
     this.missingDefinition = missingDefinition;
   }
@@ -24,6 +29,11 @@
   }
 
   @Override
+  public DefinitionContext getReferencedFromContext() {
+    return referencedFrom;
+  }
+
+  @Override
   public boolean isMissingDefinition() {
     return missingDefinition;
   }
diff --git a/src/test/java/com/android/tools/r8/missingclasses/MissingClassesTestBase.java b/src/test/java/com/android/tools/r8/missingclasses/MissingClassesTestBase.java
index 4dfeca1..2e6ffd2 100644
--- a/src/test/java/com/android/tools/r8/missingclasses/MissingClassesTestBase.java
+++ b/src/test/java/com/android/tools/r8/missingclasses/MissingClassesTestBase.java
@@ -14,9 +14,9 @@
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.ThrowableConsumer;
 import com.android.tools.r8.diagnostic.DefinitionContext;
-import com.android.tools.r8.diagnostic.internal.DefinitionContextUtils;
 import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.references.Reference;
+import com.android.tools.r8.utils.MissingDefinitionsDiagnosticTestUtils;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.List;
@@ -106,8 +106,8 @@
     inspectDiagnosticsWithIgnoreWarnings(
         diagnostics,
         referencedFrom,
-        getExpectedDiagnosticMessage(
-            DefinitionContextUtils.toSourceString(referencedFrom[0]), referencedFrom.length));
+        MissingDefinitionsDiagnosticTestUtils.getMissingClassMessage(
+            getMissingClassReference(), referencedFrom));
   }
 
   void inspectDiagnosticsWithIgnoreWarnings(
@@ -136,8 +136,8 @@
     inspectDiagnosticsWithNoRules(
         diagnostics,
         referencedFrom,
-        getExpectedDiagnosticMessage(
-            DefinitionContextUtils.toSourceString(referencedFrom[0]), referencedFrom.length));
+        MissingDefinitionsDiagnosticTestUtils.getMissingClassMessage(
+            getMissingClassReference(), referencedFrom));
   }
 
   void inspectDiagnosticsWithNoRules(
@@ -159,19 +159,4 @@
                     .assertHasMessage(expectedDiagnosticMessage)
                     .assertNumberOfMissingClasses(1));
   }
-
-  private String getExpectedDiagnosticMessage(String referencedFrom, int numberOfContexts) {
-    StringBuilder builder =
-        new StringBuilder("Missing class ")
-            .append(getMissingClassReference().getTypeName())
-            .append(" (referenced from: ")
-            .append(referencedFrom);
-    if (numberOfContexts > 1) {
-      builder.append(" and ").append(numberOfContexts - 1).append(" other context");
-      if (numberOfContexts > 2) {
-        builder.append("s");
-      }
-    }
-    return builder.append(")").toString();
-  }
 }
diff --git a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesCommandTest.java b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesCommandTest.java
index 02deedd..fac4e79 100644
--- a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesCommandTest.java
+++ b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesCommandTest.java
@@ -3,6 +3,9 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.tracereferences;
 
+import static com.android.tools.r8.utils.MissingDefinitionsDiagnosticTestUtils.getMissingClassMessage;
+import static com.android.tools.r8.utils.MissingDefinitionsDiagnosticTestUtils.getMissingFieldMessage;
+import static com.android.tools.r8.utils.MissingDefinitionsDiagnosticTestUtils.getMissingMethodMessage;
 import static org.hamcrest.CoreMatchers.containsString;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
@@ -17,6 +20,8 @@
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.TestParametersCollection;
 import com.android.tools.r8.ToolHelper;
+import com.android.tools.r8.diagnostic.DefinitionContext;
+import com.android.tools.r8.diagnostic.internal.DefinitionMethodContextImpl;
 import com.android.tools.r8.origin.Origin;
 import com.android.tools.r8.references.Reference;
 import com.android.tools.r8.utils.AndroidApiLevel;
@@ -571,18 +576,23 @@
       System.setOut(originalOut);
     }
 
+    DefinitionContext referencedFrom =
+        DefinitionMethodContextImpl.builder()
+            .setMethodContext(Reference.methodFromMethod(Source.class.getDeclaredMethod("source")))
+            .setOrigin(getOrigin(Source.class))
+            .build();
     assertThat(
         baosErr.toString(Charsets.UTF_8.name()),
         containsString(
             StringUtils.lines(
-                "Warning: Missing class " + Target.class.getTypeName(),
-                "Missing field "
-                    + FieldReferenceUtils.toSourceString(
-                        FieldReferenceUtils.fieldFromField(Target.class, "field")),
-                "Missing method "
-                    + MethodReferenceUtils.toSourceString(
-                        MethodReferenceUtils.methodFromMethod(
-                            Target.class, "method", int.class)))));
+                "Warning: "
+                    + getMissingClassMessage(
+                        Reference.classFromClass(Target.class), referencedFrom),
+                getMissingFieldMessage(
+                    FieldReferenceUtils.fieldFromField(Target.class, "field"), referencedFrom),
+                getMissingMethodMessage(
+                    MethodReferenceUtils.methodFromMethod(Target.class, "method", int.class),
+                    referencedFrom))));
     assertEquals(0, baosOut.size());
   }
 
@@ -621,18 +631,24 @@
     }
 
     assertEquals(0, baosErr.size());
+
+    DefinitionContext referencedFrom =
+        DefinitionMethodContextImpl.builder()
+            .setMethodContext(Reference.methodFromMethod(Source.class.getDeclaredMethod("source")))
+            .setOrigin(getOrigin(Source.class))
+            .build();
     assertThat(
         baosOut.toString(Charsets.UTF_8.name()),
         containsString(
             StringUtils.lines(
-                "Info: Missing class " + Target.class.getTypeName(),
-                "Missing field "
-                    + FieldReferenceUtils.toSourceString(
-                        FieldReferenceUtils.fieldFromField(Target.class, "field")),
-                "Missing method "
-                    + MethodReferenceUtils.toSourceString(
-                        MethodReferenceUtils.methodFromMethod(
-                            Target.class, "method", int.class)))));
+                "Info: "
+                    + getMissingClassMessage(
+                        Reference.classFromClass(Target.class), referencedFrom),
+                getMissingFieldMessage(
+                    FieldReferenceUtils.fieldFromField(Target.class, "field"), referencedFrom),
+                getMissingMethodMessage(
+                    MethodReferenceUtils.methodFromMethod(Target.class, "method", int.class),
+                    referencedFrom))));
   }
 
   private void checkTargetPartlyMissing(DiagnosticsChecker diagnosticsChecker) {
diff --git a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesDiagnosticTest.java b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesDiagnosticTest.java
index 292271d..a0436a9 100644
--- a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesDiagnosticTest.java
+++ b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesDiagnosticTest.java
@@ -3,6 +3,9 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.tracereferences;
 
+import static com.android.tools.r8.utils.MissingDefinitionsDiagnosticTestUtils.getMissingClassMessage;
+import static com.android.tools.r8.utils.MissingDefinitionsDiagnosticTestUtils.getMissingFieldMessage;
+import static com.android.tools.r8.utils.MissingDefinitionsDiagnosticTestUtils.getMissingMethodMessage;
 import static org.junit.Assert.fail;
 
 import com.android.tools.r8.CompilationFailedException;
@@ -12,7 +15,10 @@
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.TestParametersCollection;
 import com.android.tools.r8.ToolHelper;
+import com.android.tools.r8.diagnostic.DefinitionContext;
+import com.android.tools.r8.diagnostic.internal.DefinitionMethodContextImpl;
 import com.android.tools.r8.references.Reference;
+import com.android.tools.r8.tracereferences.TraceReferencesCommandTest.Source;
 import com.android.tools.r8.utils.AndroidApiLevel;
 import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.StringDiagnostic;
@@ -75,21 +81,41 @@
       // Expected.
     }
 
+    DefinitionContext referencedFrom =
+        DefinitionMethodContextImpl.builder()
+            .setMethodContext(Reference.methodFromMethod(Source.class.getDeclaredMethod("source")))
+            .setOrigin(getOrigin(Source.class))
+            .build();
     testDiagnosticMessages.inspectErrors(
         diagnostic ->
             diagnostic
                 .assertIsMissingDefinitionsDiagnostic()
                 .assertHasMessage(
                     StringUtils.joinLines(
-                        "Missing class " + Target1.class.getTypeName(),
-                        "Missing method void " + Target1.class.getTypeName() + ".<init>()",
-                        "Missing class " + Target2.class.getTypeName(),
-                        "Missing method void " + Target2.class.getTypeName() + ".<init>()",
-                        "Missing class " + Target3.class.getTypeName(),
-                        "Missing method void " + Target3.class.getTypeName() + ".<init>()",
-                        "Missing field int " + Target.class.getTypeName() + ".missingField1",
-                        "Missing field int " + Target.class.getTypeName() + ".missingField2",
-                        "Missing method void " + Target.class.getTypeName() + ".missingMethod()"))
+                        getMissingClassMessage(Target1.class, referencedFrom),
+                        getMissingMethodMessage(
+                            Reference.methodFromMethod(Target1.class.getDeclaredConstructor()),
+                            referencedFrom),
+                        getMissingClassMessage(Target2.class, referencedFrom),
+                        getMissingMethodMessage(
+                            Reference.methodFromMethod(Target2.class.getDeclaredConstructor()),
+                            referencedFrom),
+                        getMissingClassMessage(Target3.class, referencedFrom),
+                        getMissingMethodMessage(
+                            Reference.methodFromMethod(Target3.class.getDeclaredConstructor()),
+                            referencedFrom),
+                        getMissingFieldMessage(
+                            Reference.fieldFromField(
+                                Target.class.getDeclaredField("missingField1")),
+                            referencedFrom),
+                        getMissingFieldMessage(
+                            Reference.fieldFromField(
+                                Target.class.getDeclaredField("missingField2")),
+                            referencedFrom),
+                        getMissingMethodMessage(
+                            Reference.methodFromMethod(
+                                Target.class.getDeclaredMethod("missingMethod")),
+                            referencedFrom)))
                 .assertIsAllMissingClasses(Target1.class, Target2.class, Target3.class)
                 .assertIsAllMissingFields(
                     Reference.fieldFromField(Target.class.getField("missingField1")),
@@ -146,15 +172,29 @@
       // Expected.
     }
 
+    DefinitionContext referencedFrom =
+        DefinitionMethodContextImpl.builder()
+            .setMethodContext(Reference.methodFromMethod(Source.class.getDeclaredMethod("source")))
+            .setOrigin(getOrigin(Source.class))
+            .build();
     testDiagnosticMessages.inspectErrors(
         diagnostic ->
             diagnostic
                 .assertIsMissingDefinitionsDiagnostic()
                 .assertHasMessage(
                     StringUtils.joinLines(
-                        "Missing field int " + Target.class.getTypeName() + ".missingField1",
-                        "Missing field int " + Target.class.getTypeName() + ".missingField2",
-                        "Missing method void " + Target.class.getTypeName() + ".missingMethod()"))
+                        getMissingFieldMessage(
+                            Reference.fieldFromField(
+                                Target.class.getDeclaredField("missingField1")),
+                            referencedFrom),
+                        getMissingFieldMessage(
+                            Reference.fieldFromField(
+                                Target.class.getDeclaredField("missingField2")),
+                            referencedFrom),
+                        getMissingMethodMessage(
+                            Reference.methodFromMethod(
+                                Target.class.getDeclaredMethod("missingMethod")),
+                            referencedFrom)))
                 .assertNoMissingClasses()
                 .assertIsAllMissingFields(
                     Reference.fieldFromField(Target.class.getField("missingField1")),
@@ -202,12 +242,19 @@
       // Expected.
     }
 
+    DefinitionContext referencedFrom =
+        DefinitionMethodContextImpl.builder()
+            .setMethodContext(Reference.methodFromMethod(Source.class.getDeclaredMethod("source")))
+            .setOrigin(getOrigin(Source.class))
+            .build();
     testDiagnosticMessages.inspectErrors(
         diagnostic ->
             diagnostic
                 .assertIsMissingDefinitionsDiagnostic()
                 .assertHasMessage(
-                    "Missing method void " + Target.class.getTypeName() + ".missingMethod()")
+                    getMissingMethodMessage(
+                        Reference.methodFromMethod(Target.class.getDeclaredMethod("missingMethod")),
+                        referencedFrom))
                 .assertNoMissingClasses()
                 .assertNoMissingFields()
                 .assertIsAllMissingMethods(
diff --git a/src/test/java/com/android/tools/r8/utils/MissingDefinitionsDiagnosticTestUtils.java b/src/test/java/com/android/tools/r8/utils/MissingDefinitionsDiagnosticTestUtils.java
new file mode 100644
index 0000000..18814e1
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/utils/MissingDefinitionsDiagnosticTestUtils.java
@@ -0,0 +1,61 @@
+// Copyright (c) 2021, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.utils;
+
+import com.android.tools.r8.diagnostic.DefinitionContext;
+import com.android.tools.r8.diagnostic.internal.DefinitionContextUtils;
+import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.references.FieldReference;
+import com.android.tools.r8.references.MethodReference;
+import com.android.tools.r8.references.Reference;
+
+public class MissingDefinitionsDiagnosticTestUtils {
+
+  public static String getMissingClassMessage(
+      Class<?> missingClass, DefinitionContext... referencedFrom) {
+    return getMissingClassMessage(Reference.classFromClass(missingClass), referencedFrom);
+  }
+
+  public static String getMissingClassMessage(
+      ClassReference missingReference, DefinitionContext... referencedFrom) {
+    StringBuilder builder =
+        new StringBuilder("Missing class ").append(missingReference.getTypeName());
+    appendReferencedFromSuffix(builder, referencedFrom);
+    return builder.toString();
+  }
+
+  public static String getMissingFieldMessage(
+      FieldReference missingReference, DefinitionContext... referencedFrom) {
+    StringBuilder builder =
+        new StringBuilder("Missing field ")
+            .append(FieldReferenceUtils.toSourceString(missingReference));
+    appendReferencedFromSuffix(builder, referencedFrom);
+    return builder.toString();
+  }
+
+  public static String getMissingMethodMessage(
+      MethodReference missingReference, DefinitionContext... referencedFrom) {
+    StringBuilder builder =
+        new StringBuilder("Missing method ")
+            .append(MethodReferenceUtils.toSourceString(missingReference));
+    appendReferencedFromSuffix(builder, referencedFrom);
+    return builder.toString();
+  }
+
+  private static void appendReferencedFromSuffix(
+      StringBuilder builder, DefinitionContext... referencedFrom) {
+    builder
+        .append(" (referenced from: ")
+        .append(DefinitionContextUtils.toSourceString(referencedFrom[0]));
+    int numberOfContexts = referencedFrom.length;
+    if (numberOfContexts > 1) {
+      builder.append(" and ").append(numberOfContexts - 1).append(" other context");
+      if (numberOfContexts > 2) {
+        builder.append("s");
+      }
+    }
+    builder.append(")");
+  }
+}