diff --git a/src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesCheckConsumer.java b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCheckConsumer.java
similarity index 90%
rename from src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesCheckConsumer.java
rename to src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCheckConsumer.java
index 649efde..6912b15 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/internal/TraceReferencesCheckConsumer.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCheckConsumer.java
@@ -2,10 +2,12 @@
 // 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;
+package com.android.tools.r8.tracereferences;
 
 import com.android.tools.r8.DiagnosticsHandler;
+import com.android.tools.r8.Keep;
 import com.android.tools.r8.diagnostic.DefinitionContext;
+import com.android.tools.r8.diagnostic.MissingDefinitionsDiagnostic;
 import com.android.tools.r8.diagnostic.internal.DefinitionContextUtils;
 import com.android.tools.r8.diagnostic.internal.MissingClassInfoImpl;
 import com.android.tools.r8.diagnostic.internal.MissingDefinitionsDiagnosticImpl;
@@ -15,15 +17,18 @@
 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.
+ * A {@link TraceReferencesConsumer.ForwardingConsumer}, which forwards all callbacks to the wrapped
+ * {@link TraceReferencesConsumer}.
+ *
+ * <p>This consumer collects the set of missing definitions and reports a {@link
+ * com.android.tools.r8.diagnostic.MissingDefinitionsDiagnostic} as an error, if any missing
+ * definitions were found.
  */
+@Keep
 public class TraceReferencesCheckConsumer extends TraceReferencesConsumer.ForwardingConsumer {
 
   private final Map<ClassReference, Map<Object, DefinitionContext>> missingClassesContexts =
@@ -33,10 +38,6 @@
   private final Map<MethodReference, Map<Object, DefinitionContext>> missingMethodsContexts =
       new ConcurrentHashMap<>();
 
-  public TraceReferencesCheckConsumer() {
-    this(TraceReferencesConsumer.emptyConsumer());
-  }
-
   public TraceReferencesCheckConsumer(TraceReferencesConsumer consumer) {
     super(consumer);
   }
@@ -97,9 +98,18 @@
   @Override
   public void finished(DiagnosticsHandler handler) {
     super.finished(handler);
-    if (isEmpty()) {
-      return;
+    if (!isEmpty()) {
+      handler.error(buildDiagnostic());
     }
+  }
+
+  private boolean isEmpty() {
+    return missingClassesContexts.isEmpty()
+        && missingFieldsContexts.isEmpty()
+        && missingMethodsContexts.isEmpty();
+  }
+
+  private MissingDefinitionsDiagnostic buildDiagnostic() {
     MissingDefinitionsDiagnosticImpl.Builder diagnosticBuilder =
         MissingDefinitionsDiagnosticImpl.builder();
     missingClassesContexts.forEach(
@@ -123,12 +133,6 @@
                     .setMethod(reference)
                     .addReferencedFromContexts(referencedFrom.values())
                     .build()));
-    handler.error(diagnosticBuilder.build());
-  }
-
-  private boolean isEmpty() {
-    return missingClassesContexts.isEmpty()
-        && missingFieldsContexts.isEmpty()
-        && missingMethodsContexts.isEmpty();
+    return diagnosticBuilder.build();
   }
 }
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 82915ee..1c98765 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommand.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommand.java
@@ -20,7 +20,6 @@
 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;
@@ -331,7 +330,7 @@
     }
 
     public Builder setConsumer(TraceReferencesConsumer consumer) {
-      this.consumer = new TraceReferencesCheckConsumer(consumer);
+      this.consumer = consumer;
       return this;
     }
 
diff --git a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommandParser.java b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommandParser.java
index 96956b0..c3c8e30 100644
--- a/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommandParser.java
+++ b/src/main/java/com/android/tools/r8/tracereferences/TraceReferencesCommandParser.java
@@ -186,17 +186,19 @@
 
     switch (command) {
       case CHECK:
-        builder.setConsumer(TraceReferencesConsumer.emptyConsumer());
+        builder.setConsumer(
+            new TraceReferencesCheckConsumer(TraceReferencesConsumer.emptyConsumer()));
         break;
       case KEEP_RULES:
         builder.setConsumer(
-            TraceReferencesKeepRules.builder()
-                .setAllowObfuscation(allowObfuscation)
-                .setOutputConsumer(
-                    output != null
-                        ? new FileConsumer(output)
-                        : new WriterConsumer(null, new PrintWriter(System.out)))
-                .build());
+            new TraceReferencesCheckConsumer(
+                TraceReferencesKeepRules.builder()
+                    .setAllowObfuscation(allowObfuscation)
+                    .setOutputConsumer(
+                        output != null
+                            ? new FileConsumer(output)
+                            : new WriterConsumer(null, new PrintWriter(System.out)))
+                    .build()));
         break;
       default:
         throw new Unreachable();
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 fac4e79..d905d4c 100644
--- a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesCommandTest.java
+++ b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesCommandTest.java
@@ -241,7 +241,7 @@
               .addLibraryFiles(ToolHelper.getAndroidJar(AndroidApiLevel.P))
               .addTargetFiles(targetJar)
               .addSourceFiles(sourceJar)
-              .setConsumer(consumer)
+              .setConsumer(new TraceReferencesCheckConsumer(consumer))
               .build());
       assertEquals(expected, stringConsumer.get());
       if (diagnosticsCheckerConsumer != null) {
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 a0436a9..c12e681 100644
--- a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesDiagnosticTest.java
+++ b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesDiagnosticTest.java
@@ -74,7 +74,8 @@
               .addLibraryFiles(ToolHelper.getAndroidJar(AndroidApiLevel.P))
               .addSourceFiles(sourceJar)
               .addTargetFiles(targetJar)
-              .setConsumer(TraceReferencesConsumer.emptyConsumer())
+              .setConsumer(
+                  new TraceReferencesCheckConsumer(TraceReferencesConsumer.emptyConsumer()))
               .build());
       fail("Unexpected success");
     } catch (CompilationFailedException e) {
@@ -165,7 +166,8 @@
               .addLibraryFiles(ToolHelper.getAndroidJar(AndroidApiLevel.P))
               .addSourceFiles(sourceJar)
               .addTargetFiles(targetJar)
-              .setConsumer(TraceReferencesConsumer.emptyConsumer())
+              .setConsumer(
+                  new TraceReferencesCheckConsumer(TraceReferencesConsumer.emptyConsumer()))
               .build());
       fail("Unexpected success");
     } catch (CompilationFailedException e) {
@@ -235,7 +237,8 @@
               .addLibraryFiles(ToolHelper.getAndroidJar(AndroidApiLevel.P))
               .addSourceFiles(sourceJar)
               .addTargetFiles(targetJar)
-              .setConsumer(TraceReferencesConsumer.emptyConsumer())
+              .setConsumer(
+                  new TraceReferencesCheckConsumer(TraceReferencesConsumer.emptyConsumer()))
               .build());
       fail("Unexpected success");
     } catch (CompilationFailedException e) {
diff --git a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesMissingReferencesInDexTest.java b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesMissingReferencesInDexTest.java
index 61f5dab..b1e75c2 100644
--- a/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesMissingReferencesInDexTest.java
+++ b/src/test/java/com/android/tools/r8/tracereferences/TraceReferencesMissingReferencesInDexTest.java
@@ -76,7 +76,7 @@
           TraceReferencesCommand.builder(diagnosticsChecker)
               .addLibraryFiles(ToolHelper.getAndroidJar(AndroidApiLevel.P))
               .addSourceFiles(sourceDex)
-              .setConsumer(consumer)
+              .setConsumer(new TraceReferencesCheckConsumer(consumer))
               .build());
       fail("Expected compilation to fail");
     } catch (CompilationFailedException e) {
@@ -113,7 +113,7 @@
           TraceReferencesCommand.builder(diagnosticsChecker)
               .addLibraryFiles(ToolHelper.getAndroidJar(AndroidApiLevel.P))
               .addSourceFiles(sourceDex)
-              .setConsumer(consumer)
+              .setConsumer(new TraceReferencesCheckConsumer(consumer))
               .build());
       fail("Expected compilation to fail");
     } catch (CompilationFailedException e) {