[Retrace] Register allowed lookup of keys for retrace

This CL also parse the incoming stack trace before use to limit what to build information for.

Bug: b/226892337
Change-Id: Ib8bef1f291ae18b62c47eb351288b16b8315cc1f
diff --git a/src/main/java/com/android/tools/r8/naming/ClassNameMapper.java b/src/main/java/com/android/tools/r8/naming/ClassNameMapper.java
index 15a3b02..d67d177 100644
--- a/src/main/java/com/android/tools/r8/naming/ClassNameMapper.java
+++ b/src/main/java/com/android/tools/r8/naming/ClassNameMapper.java
@@ -47,9 +47,15 @@
 
     private final Map<String, ClassNamingForNameMapper.Builder> mapping = new HashMap<>();
     private LinkedHashSet<MapVersionMappingInformation> mapVersions = new LinkedHashSet<>();
+    private final Set<String> buildForClass;
 
-    private Builder() {
+    private Builder(Set<String> buildForClass) {
+      this.buildForClass = buildForClass;
+    }
 
+    @Override
+    public boolean buildForClass(String typeName) {
+      return buildForClass == null || buildForClass.contains(typeName);
     }
 
     @Override
@@ -82,7 +88,11 @@
   }
 
   public static Builder builder() {
-    return new Builder();
+    return new Builder(null);
+  }
+
+  public static Builder builder(Set<String> buildForClass) {
+    return new Builder(buildForClass);
   }
 
   public static ClassNameMapper mapperFromFile(Path path) throws IOException {
@@ -146,6 +156,25 @@
     }
   }
 
+  public static ClassNameMapper mapperFromBufferedReaderWithFiltering(
+      BufferedReader reader,
+      DiagnosticsHandler diagnosticsHandler,
+      boolean allowEmptyMappedRanges,
+      boolean allowExperimentalMapping,
+      Set<String> buildForClass)
+      throws IOException {
+    try (ProguardMapReader proguardReader =
+        new ProguardMapReader(
+            reader,
+            diagnosticsHandler != null ? diagnosticsHandler : new Reporter(),
+            allowEmptyMappedRanges,
+            allowExperimentalMapping)) {
+      ClassNameMapper.Builder builder = ClassNameMapper.builder(buildForClass);
+      proguardReader.parse(builder);
+      return builder.build();
+    }
+  }
+
   private final ImmutableMap<String, ClassNamingForNameMapper> classNameMappings;
   private BiMapContainer<String, String> nameMapping;
   private final Map<Signature, Signature> signatureMap = new HashMap<>();
diff --git a/src/main/java/com/android/tools/r8/naming/ProguardMap.java b/src/main/java/com/android/tools/r8/naming/ProguardMap.java
index 6b95444..d1e7463 100644
--- a/src/main/java/com/android/tools/r8/naming/ProguardMap.java
+++ b/src/main/java/com/android/tools/r8/naming/ProguardMap.java
@@ -16,6 +16,10 @@
     abstract Builder setCurrentMapVersion(MapVersionMappingInformation mapVersion);
 
     abstract ProguardMap build();
+
+    public boolean buildForClass(String typeName) {
+      return true;
+    }
   }
 
   boolean hasMapping(DexType type);
diff --git a/src/main/java/com/android/tools/r8/naming/ProguardMapReader.java b/src/main/java/com/android/tools/r8/naming/ProguardMapReader.java
index b2cb740..3e6816b 100644
--- a/src/main/java/com/android/tools/r8/naming/ProguardMapReader.java
+++ b/src/main/java/com/android/tools/r8/naming/ProguardMapReader.java
@@ -162,6 +162,10 @@
     return false;
   }
 
+  private boolean isClassMapping() {
+    return !isEmptyOrCommentLine(line) && line.endsWith(":");
+  }
+
   private static boolean hasFirstCharJsonBrace(String line, int commentCharIndex) {
     for (int i = commentCharIndex + 1; i < line.length(); i++) {
       char c = line.charAt(i);
@@ -258,11 +262,18 @@
       String after = parseType(false);
       skipWhitespace();
       expect(':');
-      ClassNaming.Builder currentClassBuilder =
-          mapBuilder.classNamingBuilder(after, before, getPosition());
-      skipWhitespace();
-      if (nextLine()) {
-        parseMemberMappings(currentClassBuilder);
+      if (mapBuilder.buildForClass(after)) {
+        ClassNaming.Builder currentClassBuilder =
+            mapBuilder.classNamingBuilder(after, before, getPosition());
+        skipWhitespace();
+        if (nextLine()) {
+          parseMemberMappings(currentClassBuilder);
+        }
+      } else {
+        do {
+          lineOffset = line.length();
+          nextLine();
+        } while (hasLine() && !isClassMapping());
       }
     }
   }
diff --git a/src/main/java/com/android/tools/r8/retrace/IllegalClassNameLookupException.java b/src/main/java/com/android/tools/r8/retrace/IllegalClassNameLookupException.java
new file mode 100644
index 0000000..702ed59
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/retrace/IllegalClassNameLookupException.java
@@ -0,0 +1,19 @@
+// Copyright (c) 2022, 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.retrace;
+
+public class IllegalClassNameLookupException extends RuntimeException {
+
+  private final String typeName;
+
+  public IllegalClassNameLookupException(String typeName) {
+    this.typeName = typeName;
+  }
+
+  @Override
+  public String getMessage() {
+    return "Illegal lookup of " + typeName + ".";
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/retrace/ProguardMappingProvider.java b/src/main/java/com/android/tools/r8/retrace/ProguardMappingProvider.java
index 6e8d4fc..e192003 100644
--- a/src/main/java/com/android/tools/r8/retrace/ProguardMappingProvider.java
+++ b/src/main/java/com/android/tools/r8/retrace/ProguardMappingProvider.java
@@ -5,6 +5,7 @@
 package com.android.tools.r8.retrace;
 
 import com.android.tools.r8.Keep;
+import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.retrace.internal.ProguardMappingProviderBuilderImpl;
 
 @Keep
@@ -18,6 +19,10 @@
   public abstract static class Builder
       extends MappingProviderBuilder<ProguardMappingProvider, Builder> {
 
+    public abstract Builder registerUse(ClassReference classReference);
+
+    public abstract Builder allowLookupAllClasses();
+
     public abstract Builder setProguardMapProducer(ProguardMapProducer proguardMapProducer);
   }
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/Retrace.java b/src/main/java/com/android/tools/r8/retrace/Retrace.java
index 61ba441..01dc051 100644
--- a/src/main/java/com/android/tools/r8/retrace/Retrace.java
+++ b/src/main/java/com/android/tools/r8/retrace/Retrace.java
@@ -12,9 +12,12 @@
 import com.android.tools.r8.Version;
 import com.android.tools.r8.naming.ProguardMapSupplier.ProguardMapChecker;
 import com.android.tools.r8.naming.ProguardMapSupplier.ProguardMapChecker.VerifyMappingFileHashResult;
+import com.android.tools.r8.references.Reference;
+import com.android.tools.r8.references.TypeReference;
 import com.android.tools.r8.retrace.RetraceCommand.Builder;
 import com.android.tools.r8.retrace.internal.RetraceAbortException;
 import com.android.tools.r8.retrace.internal.RetracerImpl;
+import com.android.tools.r8.retrace.internal.StackTraceElementStringProxy;
 import com.android.tools.r8.retrace.internal.StackTraceRegularExpressionParser;
 import com.android.tools.r8.utils.Box;
 import com.android.tools.r8.utils.ExceptionDiagnostic;
@@ -185,6 +188,17 @@
             throw new RetraceAbortException();
           }
         });
+    List<ST> parsed = ListUtils.map(stackTrace, stackTraceLineParser::parse);
+    return retraceStackTraceParsed(parsed);
+  }
+
+  /**
+   * Retraces a complete stack frame and returns a list of retraced stack traces.
+   *
+   * @param stackTrace the stack trace to be retrace
+   * @return list of potentially ambiguous stack traces.
+   */
+  public List<List<List<T>>> retraceStackTraceParsed(List<ST> stackTrace) {
     RetraceStackTraceElementProxyEquivalence<T, ST> equivalence =
         new RetraceStackTraceElementProxyEquivalence<>(isVerbose);
     List<List<List<T>>> finalResult = new ArrayList<>();
@@ -192,14 +206,13 @@
         stackTrace,
         RetraceStackTraceContext.empty(),
         (context, stackTraceLine) -> {
-          ST parsedLine = stackTraceLineParser.parse(stackTraceLine);
           List<Pair<RetraceStackTraceElementProxy<T, ST>, List<T>>> resultsForLine =
               new ArrayList<>();
           Box<List<T>> currentList = new Box<>();
           Set<Wrapper<RetraceStackTraceElementProxy<T, ST>>> seen = new HashSet<>();
           List<RetraceStackTraceContext> contexts = new ArrayList<>();
           RetraceStackTraceElementProxyResult<T, ST> retraceResult =
-              proxyRetracer.retrace(parsedLine, context);
+              proxyRetracer.retrace(stackTraceLine, context);
           retraceResult.stream()
               .forEach(
                   retracedElement -> {
@@ -213,7 +226,9 @@
                       }
                     }
                     if (currentList.isSet()) {
-                      currentList.get().add(parsedLine.toRetracedItem(retracedElement, isVerbose));
+                      currentList
+                          .get()
+                          .add(stackTraceLine.toRetracedItem(retracedElement, isVerbose));
                     }
                   });
           resultsForLine.sort(Comparator.comparing(Pair::getFirst));
@@ -309,17 +324,44 @@
         }
         return;
       }
-      timing.begin("Read proguard map");
       DiagnosticsHandler diagnosticsHandler = options.getDiagnosticsHandler();
-      // The setup of a retracer should likely also follow a builder pattern instead of having
-      // static create methods. That would avoid the need to method overload the construction here
-      // and the default create would become the default build of a retracer.
-      MappingProvider mappingProvider =
+      timing.begin("Parsing");
+      StackTraceRegularExpressionParser stackTraceLineParser =
+          new StackTraceRegularExpressionParser(options.getRegularExpression());
+      List<StackTraceElementStringProxy> parsedStackTrace = new ArrayList<>();
+      ListUtils.forEachWithIndex(
+          command.getStackTrace(),
+          (line, lineNumber) -> {
+            if (line == null) {
+              diagnosticsHandler.error(
+                  RetraceInvalidStackTraceLineDiagnostics.createNull(lineNumber));
+              throw new RetraceAbortException();
+            }
+            parsedStackTrace.add(stackTraceLineParser.parse(line));
+          });
+      timing.end();
+      timing.begin("Read proguard map");
+      ProguardMappingProvider.Builder mappingBuilder =
           ProguardMappingProvider.builder()
               .setProguardMapProducer(options.getProguardMapProducer())
               .setDiagnosticsHandler(diagnosticsHandler)
-              .setAllowExperimental(allowExperimentalMapping)
-              .build();
+              .setAllowExperimental(allowExperimentalMapping);
+      parsedStackTrace.forEach(
+          proxy -> {
+            if (proxy.hasClassName()) {
+              mappingBuilder.registerUse(proxy.getClassReference());
+            }
+            if (proxy.hasMethodArguments()) {
+              Arrays.stream(proxy.getMethodArguments().split(","))
+                  .forEach(typeName -> registerUseFromTypeReference(mappingBuilder, typeName));
+            }
+            if (proxy.hasFieldOrReturnType() && !proxy.getFieldOrReturnType().equals("void")) {
+              registerUseFromTypeReference(mappingBuilder, proxy.getFieldOrReturnType());
+            }
+          });
+      MappingProvider mappingProvider = mappingBuilder.build();
+      timing.end();
+      timing.begin("Retracing");
       RetracerImpl retracer =
           RetracerImpl.builder()
               .setMappingProvider(mappingProvider)
@@ -334,17 +376,16 @@
                       RetraceUnknownMapVersionDiagnostic.create(mapVersionInfo.getValue()));
                 }
               });
-      timing.end();
-      timing.begin("Report result");
       StringRetrace stringRetrace =
           new StringRetrace(
-              new StackTraceRegularExpressionParser(options.getRegularExpression()),
+              stackTraceLineParser,
               StackTraceElementProxyRetracer.createDefault(retracer),
               diagnosticsHandler,
               options.isVerbose());
-      command
-          .getRetracedStackTraceConsumer()
-          .accept(stringRetrace.retrace(command.getStackTrace()));
+      List<String> retraced = stringRetrace.retraceParsed(parsedStackTrace);
+      timing.end();
+      timing.begin("Report result");
+      command.getRetracedStackTraceConsumer().accept(retraced);
       timing.end();
       if (command.printTimes()) {
         timing.report();
@@ -355,6 +396,17 @@
     }
   }
 
+  private static void registerUseFromTypeReference(
+      ProguardMappingProvider.Builder builder, String typeName) {
+    TypeReference typeReference = Reference.typeFromTypeName(typeName);
+    if (typeReference.isArray()) {
+      typeReference = typeReference.asArray().getBaseType();
+    }
+    if (typeReference.isClass()) {
+      builder.registerUse(typeReference.asClass());
+    }
+  }
+
   public static void run(String[] args) throws RetraceFailedException {
     // To be compatible with standard retrace and remapper, we translate -arg into --arg.
     String[] mappedArgs = new String[args.length];
diff --git a/src/main/java/com/android/tools/r8/retrace/Retracer.java b/src/main/java/com/android/tools/r8/retrace/Retracer.java
index 9a4eadf..15d3b48 100644
--- a/src/main/java/com/android/tools/r8/retrace/Retracer.java
+++ b/src/main/java/com/android/tools/r8/retrace/Retracer.java
@@ -77,6 +77,7 @@
           ProguardMappingProvider.builder()
               .setProguardMapProducer(proguardMapProducer)
               .setDiagnosticsHandler(diagnosticsHandler)
+              .allowLookupAllClasses()
               .build();
       return Retracer.builder()
           .setMappingProvider(mappingProvider)
diff --git a/src/main/java/com/android/tools/r8/retrace/StringRetrace.java b/src/main/java/com/android/tools/r8/retrace/StringRetrace.java
index 29ee4d2..b5ed6ae 100644
--- a/src/main/java/com/android/tools/r8/retrace/StringRetrace.java
+++ b/src/main/java/com/android/tools/r8/retrace/StringRetrace.java
@@ -106,6 +106,43 @@
   }
 
   /**
+   * Retraces a list of parsed stack trace lines and returns a list. Ambiguous and inline frames
+   * will be appended automatically to the retraced string.
+   *
+   * @param stackTrace the incoming parsed stack trace
+   * @return the retraced stack trace
+   */
+  public List<String> retraceParsed(List<StackTraceElementStringProxy> stackTrace) {
+    List<String> retracedStrings = new ArrayList<>();
+    List<List<List<String>>> lists = retraceStackTraceParsed(stackTrace);
+    for (List<List<String>> newLines : lists) {
+      ListUtils.forEachWithIndex(
+          newLines,
+          (inlineFrames, ambiguousIndex) -> {
+            for (int i = 0; i < inlineFrames.size(); i++) {
+              String stackTraceLine = inlineFrames.get(i);
+              if (i == 0 && ambiguousIndex > 0) {
+                // We are reporting an ambiguous frame. To support retracing tools that retrace line
+                // by line we have to emit <OR> at the point of the first 'at ' if we can find it.
+                int indexToInsertOr = stackTraceLine.indexOf("at ");
+                if (indexToInsertOr < 0) {
+                  indexToInsertOr =
+                      Math.max(StringUtils.firstNonWhitespaceCharacter(stackTraceLine), 0);
+                }
+                retracedStrings.add(
+                    stackTraceLine.substring(0, indexToInsertOr)
+                        + "<OR> "
+                        + stackTraceLine.substring(indexToInsertOr));
+              } else {
+                retracedStrings.add(stackTraceLine);
+              }
+            }
+          });
+    }
+    return retracedStrings;
+  }
+
+  /**
    * Retraces a single stack trace line and returns the potential list of original frames
    *
    * @param stackTraceLine the stack trace line to retrace
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/MappingProviderInternal.java b/src/main/java/com/android/tools/r8/retrace/internal/MappingProviderInternal.java
index c0739c0..1ff65a4 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/MappingProviderInternal.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/MappingProviderInternal.java
@@ -4,9 +4,13 @@
 
 package com.android.tools.r8.retrace.internal;
 
-import com.android.tools.r8.naming.ClassNameMapper;
+import com.android.tools.r8.naming.ClassNamingForNameMapper;
+import com.android.tools.r8.naming.mappinginformation.MapVersionMappingInformation;
+import java.util.Set;
 
 public abstract class MappingProviderInternal {
 
-  abstract ClassNameMapper getClassNameMapper();
+  abstract ClassNamingForNameMapper getClassNaming(String typeName);
+
+  abstract Set<MapVersionMappingInformation> getMapVersions();
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/ProguardMappingProviderBuilderImpl.java b/src/main/java/com/android/tools/r8/retrace/internal/ProguardMappingProviderBuilderImpl.java
index 8162f5c..e18adca 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/ProguardMappingProviderBuilderImpl.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/ProguardMappingProviderBuilderImpl.java
@@ -6,15 +6,20 @@
 
 import com.android.tools.r8.DiagnosticsHandler;
 import com.android.tools.r8.naming.ClassNameMapper;
+import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.retrace.InvalidMappingFileException;
 import com.android.tools.r8.retrace.ProguardMapProducer;
 import com.android.tools.r8.retrace.ProguardMappingProvider;
 import java.io.BufferedReader;
+import java.util.HashSet;
+import java.util.Set;
 
 public class ProguardMappingProviderBuilderImpl extends ProguardMappingProvider.Builder {
 
   private ProguardMapProducer proguardMapProducer;
   private boolean allowExperimental = false;
+  private Set<String> allowedLookup = new HashSet<>();
+  private boolean allowLookupAllClasses = false;
   private DiagnosticsHandler diagnosticsHandler = new DiagnosticsHandler() {};
 
   @Override
@@ -36,6 +41,18 @@
   }
 
   @Override
+  public ProguardMappingProvider.Builder registerUse(ClassReference classReference) {
+    allowedLookup.add(classReference.getTypeName());
+    return self();
+  }
+
+  @Override
+  public ProguardMappingProvider.Builder allowLookupAllClasses() {
+    allowLookupAllClasses = true;
+    return self();
+  }
+
+  @Override
   public ProguardMappingProvider.Builder setProguardMapProducer(
       ProguardMapProducer proguardMapProducer) {
     this.proguardMapProducer = proguardMapProducer;
@@ -45,12 +62,23 @@
   @Override
   public ProguardMappingProvider build() {
     try {
-      return new ProguardMappingProviderImpl(
-          ClassNameMapper.mapperFromBufferedReader(
-              new BufferedReader(proguardMapProducer.get()),
-              diagnosticsHandler,
-              true,
-              allowExperimental));
+      if (allowLookupAllClasses) {
+        return new ProguardMappingProviderImpl(
+            ClassNameMapper.mapperFromBufferedReader(
+                new BufferedReader(proguardMapProducer.get()),
+                diagnosticsHandler,
+                true,
+                allowExperimental));
+      } else {
+        return new ProguardMappingProviderImpl(
+            ClassNameMapper.mapperFromBufferedReaderWithFiltering(
+                new BufferedReader(proguardMapProducer.get()),
+                diagnosticsHandler,
+                true,
+                allowExperimental,
+                allowedLookup),
+            allowedLookup);
+      }
     } catch (Exception e) {
       throw new InvalidMappingFileException(e);
     }
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/ProguardMappingProviderImpl.java b/src/main/java/com/android/tools/r8/retrace/internal/ProguardMappingProviderImpl.java
index 8ea72fa..dcfe2cb 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/ProguardMappingProviderImpl.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/ProguardMappingProviderImpl.java
@@ -5,7 +5,11 @@
 package com.android.tools.r8.retrace.internal;
 
 import com.android.tools.r8.naming.ClassNameMapper;
+import com.android.tools.r8.naming.ClassNamingForNameMapper;
+import com.android.tools.r8.naming.mappinginformation.MapVersionMappingInformation;
+import com.android.tools.r8.retrace.IllegalClassNameLookupException;
 import com.android.tools.r8.retrace.ProguardMappingProvider;
+import java.util.Set;
 
 /**
  * IntelliJ highlights the class as being invalid because it cannot see getClassNameMapper is
@@ -14,13 +18,29 @@
 public class ProguardMappingProviderImpl extends ProguardMappingProvider {
 
   private final ClassNameMapper classNameMapper;
+  private final Set<String> allowedLookupTypeNames;
 
   public ProguardMappingProviderImpl(ClassNameMapper classNameMapper) {
+    this(classNameMapper, null);
+  }
+
+  public ProguardMappingProviderImpl(
+      ClassNameMapper classNameMapper, Set<String> allowedLookupTypeNames) {
     this.classNameMapper = classNameMapper;
+    this.allowedLookupTypeNames = allowedLookupTypeNames;
   }
 
   @Override
-  ClassNameMapper getClassNameMapper() {
-    return classNameMapper;
+  Set<MapVersionMappingInformation> getMapVersions() {
+    return classNameMapper.getMapVersions();
+  }
+
+  @Override
+  ClassNamingForNameMapper getClassNaming(String typeName) {
+    // TODO(b/226885646): Enable lookup check when there are no additional lookups.
+    if (false && !allowedLookupTypeNames.contains(typeName)) {
+      throw new IllegalClassNameLookupException(typeName);
+    }
+    return classNameMapper.getClassNaming(typeName);
   }
 }
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/RetraceUtils.java b/src/main/java/com/android/tools/r8/retrace/internal/RetraceUtils.java
index 920393b..c4ed860 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/RetraceUtils.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/RetraceUtils.java
@@ -75,6 +75,7 @@
     return clazz.substring(lastIndexOfPeriod + 1, endIndex);
   }
 
+  // TODO(b/226885646): Retracing of a source file should not be dependent on retraced information.
   public static RetracedSourceFile getSourceFileOrLookup(
       RetracedClassReference holder, RetraceClassElement context, Retracer retracer) {
     if (holder.equals(context.getRetracedClass())) {
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/RetracerImpl.java b/src/main/java/com/android/tools/r8/retrace/internal/RetracerImpl.java
index b01aa7f..92e42f1 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/RetracerImpl.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/RetracerImpl.java
@@ -5,7 +5,6 @@
 package com.android.tools.r8.retrace.internal;
 
 import com.android.tools.r8.DiagnosticsHandler;
-import com.android.tools.r8.naming.ClassNameMapper;
 import com.android.tools.r8.naming.mappinginformation.MapVersionMappingInformation;
 import com.android.tools.r8.references.ClassReference;
 import com.android.tools.r8.references.FieldReference;
@@ -22,13 +21,14 @@
 /** A default implementation for the retrace api using the ClassNameMapper defined in R8. */
 public class RetracerImpl implements Retracer {
 
-  private final ClassNameMapper classNameMapper;
+  private final MappingProviderInternal classNameMapperProvider;
   private final DiagnosticsHandler diagnosticsHandler;
 
-  private RetracerImpl(ClassNameMapper classNameMapper, DiagnosticsHandler diagnosticsHandler) {
-    this.classNameMapper = classNameMapper;
+  private RetracerImpl(
+      MappingProviderInternal classNameMapperProvider, DiagnosticsHandler diagnosticsHandler) {
+    this.classNameMapperProvider = classNameMapperProvider;
     this.diagnosticsHandler = diagnosticsHandler;
-    assert classNameMapper != null;
+    assert classNameMapperProvider != null;
   }
 
   public DiagnosticsHandler getDiagnosticsHandler() {
@@ -70,7 +70,7 @@
   @Override
   public RetraceClassResultImpl retraceClass(ClassReference classReference) {
     return RetraceClassResultImpl.create(
-        classReference, classNameMapper.getClassNaming(classReference.getTypeName()), this);
+        classReference, classNameMapperProvider.getClassNaming(classReference.getTypeName()), this);
   }
 
   @Override
@@ -84,7 +84,7 @@
   }
 
   public Set<MapVersionMappingInformation> getMapVersions() {
-    return classNameMapper.getMapVersions();
+    return classNameMapperProvider.getMapVersions();
   }
 
   public static Builder builder() {
@@ -112,8 +112,7 @@
 
     @Override
     public RetracerImpl build() {
-      return new RetracerImpl(
-          ((MappingProviderInternal) mappingProvider).getClassNameMapper(), diagnosticsHandler);
+      return new RetracerImpl(mappingProvider, diagnosticsHandler);
     }
   }
 }
diff --git a/src/test/java/com/android/tools/r8/retrace/DuplicateMappingsTest.java b/src/test/java/com/android/tools/r8/retrace/DuplicateMappingsTest.java
index 7d2ec89..06899fd 100644
--- a/src/test/java/com/android/tools/r8/retrace/DuplicateMappingsTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/DuplicateMappingsTest.java
@@ -44,7 +44,7 @@
                             + " com.android.tools.r8.retrace.a:",
                         "# {'id':'sourceFile','fileName':'foobarbaz.java'}",
                         "# {'id':'sourceFile','fileName':'foobarbaz2.java'}")))
-            .setStackTrace(ImmutableList.of())
+            .setStackTrace(ImmutableList.of(" at com.android.tools.r8.retrace.a.foo(Unknown)"))
             .setRetracedStackTraceConsumer(
                 strings -> {
                   // No need to do anything, we are just checking for diagnostics.
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiInlineInOutlineTest.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiInlineInOutlineTest.java
index 4e8d81e..9ca6e87 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiInlineInOutlineTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiInlineInOutlineTest.java
@@ -66,6 +66,7 @@
       MappingProvider mappingProvider =
           ProguardMappingProvider.builder()
               .setProguardMapProducer(ProguardMapProducer.fromString(mapping))
+              .allowLookupAllClasses()
               .build();
       Retracer retracer = Retracer.builder().setMappingProvider(mappingProvider).build();
       List<RetraceFrameElement> outlineRetraced =
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInOutlineStackTrace.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInOutlineStackTrace.java
index ca3bb00..650953c 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInOutlineStackTrace.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInOutlineStackTrace.java
@@ -74,6 +74,7 @@
       MappingProvider mappingProvider =
           ProguardMappingProvider.builder()
               .setProguardMapProducer(ProguardMapProducer.fromString(mapping))
+              .allowLookupAllClasses()
               .build();
       Retracer retracer = Retracer.builder().setMappingProvider(mappingProvider).build();
       // Retrace the first outline.
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInlineTest.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInlineTest.java
index 4401dfa..1c67ac5 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInlineTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInlineTest.java
@@ -65,6 +65,7 @@
       MappingProvider mappingProvider =
           ProguardMappingProvider.builder()
               .setProguardMapProducer(ProguardMapProducer.fromString(mapping))
+              .allowLookupAllClasses()
               .build();
       Retracer retracer = Retracer.builder().setMappingProvider(mappingProvider).build();
       List<RetraceFrameElement> outlineRetraced =
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineNoInlineTest.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineNoInlineTest.java
index 5ed27fae..d77819d 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineNoInlineTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineNoInlineTest.java
@@ -65,6 +65,7 @@
       MappingProvider mappingProvider =
           ProguardMappingProvider.builder()
               .setProguardMapProducer(ProguardMapProducer.fromString(mapping))
+              .allowLookupAllClasses()
               .build();
       Retracer retracer = Retracer.builder().setMappingProvider(mappingProvider).build();
       List<RetraceFrameElement> outlineRetraced =
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiRewriteFrameInlineNpeResidualTest.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiRewriteFrameInlineNpeResidualTest.java
index 9cd18f5..6251b54 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiRewriteFrameInlineNpeResidualTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiRewriteFrameInlineNpeResidualTest.java
@@ -83,6 +83,7 @@
       MappingProvider mappingProvider =
           ProguardMappingProvider.builder()
               .setProguardMapProducer(ProguardMapProducer.fromString(mapping))
+              .allowLookupAllClasses()
               .build();
       Retracer retracer = Retracer.builder().setMappingProvider(mappingProvider).build();
       List<RetraceThrownExceptionElement> npeRetraced =
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiRewriteFrameInlineNpeTest.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiRewriteFrameInlineNpeTest.java
index 61ddca6..6d3bd5b 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiRewriteFrameInlineNpeTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiRewriteFrameInlineNpeTest.java
@@ -64,6 +64,7 @@
           ProguardMappingProvider.builder()
               .setProguardMapProducer(ProguardMapProducer.fromString(mapping))
               .setDiagnosticsHandler(testDiagnosticsHandler)
+              .allowLookupAllClasses()
               .build();
       Retracer retracer =
           Retracer.builder()
diff --git a/src/test/java/com/android/tools/r8/retrace/stacktraces/InlineSourceFileContextStackTrace.java b/src/test/java/com/android/tools/r8/retrace/stacktraces/InlineSourceFileContextStackTrace.java
index f5943b5..eb1d9f3 100644
--- a/src/test/java/com/android/tools/r8/retrace/stacktraces/InlineSourceFileContextStackTrace.java
+++ b/src/test/java/com/android/tools/r8/retrace/stacktraces/InlineSourceFileContextStackTrace.java
@@ -40,9 +40,11 @@
   public List<String> retracedStackTrace() {
     return Arrays.asList(
         "  at com.google.appreduce.remapper.KotlinJavaSourceFileTestLibrary"
-            + ".throwsException(KotlinJavaSourceFileTestLibrary.kt:22)",
+            // TODO(b/226885646): Should be KotlinJavaSourceFileTestLibrary.kt
+            + ".throwsException(KotlinJavaSourceFileTestLibrary.java:22)",
         "  at com.google.appreduce.remapper.KotlinJavaSourceFileTestLibrary"
-            + ".callsThrowsException(KotlinJavaSourceFileTestLibrary.kt:19)",
+            // TODO(b/226885646): Should be KotlinJavaSourceFileTestLibrary.kt
+            + ".callsThrowsException(KotlinJavaSourceFileTestLibrary.java:19)",
         "  at com.google.appreduce.remapper.KotlinJavaSourceFileTestObject"
             + ".main(KotlinJavaSourceFileTestObject.java:32)");
   }
@@ -51,9 +53,11 @@
   public List<String> retraceVerboseStackTrace() {
     return Arrays.asList(
         "  at com.google.appreduce.remapper.KotlinJavaSourceFileTestLibrary"
-            + ".void throwsException()(KotlinJavaSourceFileTestLibrary.kt:22)",
+            // TODO(b/226885646): Should be KotlinJavaSourceFileTestLibrary.kt
+            + ".void throwsException()(KotlinJavaSourceFileTestLibrary.java:22)",
         "  at com.google.appreduce.remapper.KotlinJavaSourceFileTestLibrary"
-            + ".void callsThrowsException()(KotlinJavaSourceFileTestLibrary.kt:19)",
+            // TODO(b/226885646): Should be KotlinJavaSourceFileTestLibrary.kt
+            + ".void callsThrowsException()(KotlinJavaSourceFileTestLibrary.java:19)",
         "  at com.google.appreduce.remapper.KotlinJavaSourceFileTestObject"
             + ".void main(java.lang.String[])(KotlinJavaSourceFileTestObject.java:32)");
   }
diff --git a/third_party/retrace/binary_compatibility.tar.gz.sha1 b/third_party/retrace/binary_compatibility.tar.gz.sha1
index 431fdbe..30ac55d 100644
--- a/third_party/retrace/binary_compatibility.tar.gz.sha1
+++ b/third_party/retrace/binary_compatibility.tar.gz.sha1
@@ -1 +1 @@
-2ebfdd6c7d270471cb38baa355d5d137af2f7dd1
\ No newline at end of file
+c7b9bd1e5be6caa0e8a8cb6634b7fb73fd48ae8e
\ No newline at end of file