[Retrace] Add support for retracing stacktraces with outlines

Bug: 201397823
Change-Id: I298e29bc4597403dde1c849beca72d27eb3acbae
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 9ce44ff..5bb9ca1 100644
--- a/src/main/java/com/android/tools/r8/retrace/Retrace.java
+++ b/src/main/java/com/android/tools/r8/retrace/Retrace.java
@@ -199,8 +199,9 @@
               ST parsedLine = stackTraceLineParser.parse(stackTraceLine);
               List<RetracedNodeState<T, ST>> newLeaves = new ArrayList<>();
               for (RetracedNodeState<T, ST> previousNode : acc) {
-                proxyRetracer
-                    .retrace(parsedLine, previousNode.context)
+                RetraceStackTraceElementProxyResult<T, ST> result =
+                    proxyRetracer.retrace(parsedLine, previousNode.context);
+                result.stream()
                     .forEach(
                         retracedElement -> {
                           if (retracedElement.isTopFrame() || !retracedElement.hasRetracedClass()) {
@@ -212,7 +213,7 @@
                 if (!previousNode.hasChildren()) {
                   // This happens when there is nothing to retrace. Add the node to newLeaves to
                   // ensure we keep retracing this path.
-                  previousNode.addChild(null, RetraceStackTraceContext.empty());
+                  previousNode.addChild(null, result.getResultContext());
                 }
                 newLeaves.addAll(previousNode.getChildren());
               }
@@ -232,8 +233,7 @@
     Map<RetraceStackTraceElementProxy<T, ST>, List<T>> ambiguousBlocks = new HashMap<>();
     List<RetraceStackTraceElementProxy<T, ST>> ambiguousKeys = new ArrayList<>();
     ST parsedLine = stackTraceLineParser.parse(stackTraceFrame);
-    proxyRetracer
-        .retrace(parsedLine, RetraceStackTraceContext.empty())
+    proxyRetracer.retrace(parsedLine, RetraceStackTraceContext.empty()).stream()
         .forEach(
             retracedElement -> {
               if (retracedElement.isTopFrame() || !retracedElement.hasRetracedClass()) {
@@ -259,8 +259,7 @@
    */
   public List<T> retraceLine(T stackTraceLine) {
     ST parsedLine = stackTraceLineParser.parse(stackTraceLine);
-    return proxyRetracer
-        .retrace(parsedLine, RetraceStackTraceContext.empty())
+    return proxyRetracer.retrace(parsedLine, RetraceStackTraceContext.empty()).stream()
         .map(
             retraceFrame -> {
               retraceFrame.getOriginalItem().toRetracedItem(retraceFrame, isVerbose);
diff --git a/src/main/java/com/android/tools/r8/retrace/RetraceStackTraceElementProxyResult.java b/src/main/java/com/android/tools/r8/retrace/RetraceStackTraceElementProxyResult.java
new file mode 100644
index 0000000..1db1f08
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/retrace/RetraceStackTraceElementProxyResult.java
@@ -0,0 +1,22 @@
+// 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.retrace;
+
+import com.android.tools.r8.Keep;
+import java.util.stream.Stream;
+
+@Keep
+public interface RetraceStackTraceElementProxyResult<T, ST extends StackTraceElementProxy<T, ST>> {
+
+  Stream<? extends RetraceStackTraceElementProxy<T, ST>> stream();
+
+  /**
+   * If the stream is empty, use getResultContext to obtain the resulting stack trace context. Due
+   * to the lazyness of streams the result is only populated after querying the stream.
+   *
+   * @return the resulting stack trace context.
+   */
+  RetraceStackTraceContext getResultContext();
+}
diff --git a/src/main/java/com/android/tools/r8/retrace/StackTraceElementProxyRetracer.java b/src/main/java/com/android/tools/r8/retrace/StackTraceElementProxyRetracer.java
index 851bbc2..e792ef0 100644
--- a/src/main/java/com/android/tools/r8/retrace/StackTraceElementProxyRetracer.java
+++ b/src/main/java/com/android/tools/r8/retrace/StackTraceElementProxyRetracer.java
@@ -6,13 +6,11 @@
 
 import com.android.tools.r8.Keep;
 import com.android.tools.r8.retrace.internal.StackTraceElementProxyRetracerImpl;
-import java.util.stream.Stream;
 
 @Keep
 public interface StackTraceElementProxyRetracer<T, ST extends StackTraceElementProxy<T, ST>> {
 
-  Stream<? extends RetraceStackTraceElementProxy<T, ST>> retrace(
-      ST element, RetraceStackTraceContext context);
+  RetraceStackTraceElementProxyResult<T, ST> retrace(ST element, RetraceStackTraceContext context);
 
   static <T, ST extends StackTraceElementProxy<T, ST>>
       StackTraceElementProxyRetracer<T, ST> createDefault(Retracer retracer) {
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/RetraceStackTraceElementProxyResultImpl.java b/src/main/java/com/android/tools/r8/retrace/internal/RetraceStackTraceElementProxyResultImpl.java
new file mode 100644
index 0000000..1317d9f
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/retrace/internal/RetraceStackTraceElementProxyResultImpl.java
@@ -0,0 +1,67 @@
+// 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.retrace.internal;
+
+import com.android.tools.r8.retrace.RetraceStackTraceContext;
+import com.android.tools.r8.retrace.RetraceStackTraceElementProxyResult;
+import com.android.tools.r8.retrace.StackTraceElementProxy;
+import com.android.tools.r8.retrace.internal.StackTraceElementProxyRetracerImpl.RetraceStackTraceElementProxyImpl;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
+
+public class RetraceStackTraceElementProxyResultImpl<T, ST extends StackTraceElementProxy<T, ST>>
+    implements RetraceStackTraceElementProxyResult<T, ST> {
+
+  private final Stream<? extends RetraceStackTraceElementProxyImpl<T, ST>> resultStream;
+  private final Supplier<RetraceStackTraceContext> resultContext;
+
+  private RetraceStackTraceElementProxyResultImpl(
+      Stream<? extends RetraceStackTraceElementProxyImpl<T, ST>> resultStream,
+      Supplier<RetraceStackTraceContext> resultContext) {
+    this.resultStream = resultStream;
+    this.resultContext = resultContext;
+  }
+
+  @Override
+  public Stream<? extends RetraceStackTraceElementProxyImpl<T, ST>> stream() {
+    return resultStream;
+  }
+
+  @Override
+  public RetraceStackTraceContext getResultContext() {
+    return resultContext.get();
+  }
+
+  Builder<T, ST> builder() {
+    return Builder.<T, ST>create().setResultStream(resultStream).setResultContext(resultContext);
+  }
+
+  static class Builder<T, ST extends StackTraceElementProxy<T, ST>> {
+
+    Stream<? extends RetraceStackTraceElementProxyImpl<T, ST>> resultStream;
+    Supplier<RetraceStackTraceContext> resultContext;
+
+    private Builder() {}
+
+    Builder<T, ST> setResultStream(
+        Stream<? extends RetraceStackTraceElementProxyImpl<T, ST>> resultStream) {
+      this.resultStream = resultStream;
+      return this;
+    }
+
+    Builder<T, ST> setResultContext(Supplier<RetraceStackTraceContext> resultContext) {
+      this.resultContext = resultContext;
+      return this;
+    }
+
+    RetraceStackTraceElementProxyResultImpl<T, ST> build() {
+      return new RetraceStackTraceElementProxyResultImpl<>(resultStream, resultContext);
+    }
+
+    static <T, ST extends StackTraceElementProxy<T, ST>> Builder<T, ST> create() {
+      return new Builder<>();
+    }
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/retrace/internal/StackTraceElementProxyRetracerImpl.java b/src/main/java/com/android/tools/r8/retrace/internal/StackTraceElementProxyRetracerImpl.java
index 317eb15..5623bf2 100644
--- a/src/main/java/com/android/tools/r8/retrace/internal/StackTraceElementProxyRetracerImpl.java
+++ b/src/main/java/com/android/tools/r8/retrace/internal/StackTraceElementProxyRetracerImpl.java
@@ -7,26 +7,33 @@
 import com.android.tools.r8.references.Reference;
 import com.android.tools.r8.references.TypeReference;
 import com.android.tools.r8.retrace.RetraceClassResult;
+import com.android.tools.r8.retrace.RetraceFieldElement;
 import com.android.tools.r8.retrace.RetraceFieldResult;
+import com.android.tools.r8.retrace.RetraceFrameElement;
 import com.android.tools.r8.retrace.RetraceFrameResult;
 import com.android.tools.r8.retrace.RetraceStackTraceContext;
 import com.android.tools.r8.retrace.RetraceStackTraceElementProxy;
+import com.android.tools.r8.retrace.RetraceStackTraceElementProxyResult;
+import com.android.tools.r8.retrace.RetraceThrownExceptionElement;
 import com.android.tools.r8.retrace.RetraceTypeResult;
 import com.android.tools.r8.retrace.RetraceTypeResult.Element;
 import com.android.tools.r8.retrace.RetracedClassReference;
 import com.android.tools.r8.retrace.RetracedFieldReference;
 import com.android.tools.r8.retrace.RetracedMethodReference;
+import com.android.tools.r8.retrace.RetracedSingleFrame;
 import com.android.tools.r8.retrace.RetracedSourceFile;
 import com.android.tools.r8.retrace.RetracedTypeReference;
 import com.android.tools.r8.retrace.Retracer;
 import com.android.tools.r8.retrace.StackTraceElementProxy;
 import com.android.tools.r8.retrace.StackTraceElementProxyRetracer;
+import com.android.tools.r8.utils.Box;
 import com.android.tools.r8.utils.ListUtils;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.OptionalInt;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -40,185 +47,242 @@
   }
 
   @Override
-  public Stream<? extends RetraceStackTraceElementProxy<T, ST>> retrace(
+  public RetraceStackTraceElementProxyResult<T, ST> retrace(
       ST element, RetraceStackTraceContext context) {
-    Stream<RetraceStackTraceElementProxyImpl<T, ST>> currentResults =
-        Stream.of(RetraceStackTraceElementProxyImpl.create(element, context));
+    RetraceStackTraceElementProxyResultImpl<T, ST> currentResult =
+        RetraceStackTraceElementProxyResultImpl.Builder.<T, ST>create()
+            .setResultStream(Stream.of(RetraceStackTraceElementProxyImpl.create(element, context)))
+            .setResultContext(RetraceStackTraceContext::empty)
+            .build();
     if (!element.hasClassName()
         && !element.hasFieldOrReturnType()
         && !element.hasMethodArguments()) {
-      return currentResults;
+      return currentResult;
     }
-    currentResults = retraceFieldOrReturnType(currentResults, element);
-    currentResults = retracedMethodArguments(currentResults, element);
+    currentResult = retraceFieldOrReturnType(currentResult, element);
+    currentResult = retracedMethodArguments(currentResult, element);
     if (element.hasClassName()) {
       RetraceClassResult classResult = retracer.retraceClass(element.getClassReference());
       if (element.hasMethodName()) {
-        currentResults = retraceMethod(currentResults, element, classResult);
+        currentResult = retraceMethod(currentResult, element, classResult, context);
       } else if (element.hasFieldName()) {
-        currentResults = retraceField(currentResults, element, classResult);
+        currentResult = retraceField(currentResult, element, classResult);
       } else {
-        currentResults = retraceClassOrType(currentResults, element, classResult);
+        currentResult = retraceClassOrType(currentResult, classResult);
       }
     }
-    return currentResults;
+    return currentResult;
   }
 
-  private Stream<RetraceStackTraceElementProxyImpl<T, ST>> retraceClassOrType(
-      Stream<RetraceStackTraceElementProxyImpl<T, ST>> currentResults,
-      ST element,
+  private RetraceStackTraceElementProxyResultImpl<T, ST> retraceClassOrType(
+      RetraceStackTraceElementProxyResultImpl<T, ST> currentResult,
       RetraceClassResult classResult) {
-    return currentResults.flatMap(
-        proxy ->
-            // We assume, since no method was defined for this stack trace element, that this was a
-            // thrown exception.
-            classResult.lookupThrownException(proxy.getContext()).stream()
-                .map(
-                    thrownExceptionElement ->
-                        proxy
-                            .builder()
-                            .setRetracedClass(thrownExceptionElement.getRetracedClass())
-                            .joinAmbiguous(classResult.isAmbiguous())
-                            .setTopFrame(true)
-                            .setContext(thrownExceptionElement.getContext())
-                            .applyIf(
-                                element.hasSourceFile(),
-                                builder -> {
-                                  RetracedSourceFile sourceFile =
-                                      thrownExceptionElement.getSourceFile();
-                                  builder.setSourceFile(
-                                      sourceFile.hasRetraceResult()
-                                          ? sourceFile.getSourceFile()
-                                          : RetraceUtils.inferSourceFile(
-                                              thrownExceptionElement
-                                                  .getRetracedClass()
-                                                  .getTypeName(),
-                                              element.getSourceFile(),
-                                              classResult.hasRetraceResult()));
-                                })
-                            .build()));
+    return currentResult
+        .builder()
+        .setResultStream(
+            currentResult.stream()
+                .flatMap(
+                    proxy ->
+                        // We assume, since no method was defined for this stack trace element, that
+                        // this was a thrown exception.
+                        classResult.lookupThrownException(proxy.getContext()).stream()
+                            .map(
+                                thrownExceptionElement ->
+                                    buildProxyForRewrittenThrownExceptionElement(
+                                        classResult, proxy, thrownExceptionElement))))
+        .build();
   }
 
-  private Stream<RetraceStackTraceElementProxyImpl<T, ST>> retraceMethod(
-      Stream<RetraceStackTraceElementProxyImpl<T, ST>> currentResults,
+  private RetraceStackTraceElementProxyImpl<T, ST> buildProxyForRewrittenThrownExceptionElement(
+      RetraceClassResult classResult,
+      RetraceStackTraceElementProxyImpl<T, ST> proxy,
+      RetraceThrownExceptionElement thrownExceptionElement) {
+    return proxy
+        .builder()
+        .setRetracedClass(thrownExceptionElement.getRetracedClass())
+        .joinAmbiguous(classResult.isAmbiguous())
+        .setTopFrame(true)
+        .setContext(thrownExceptionElement.getContext())
+        .apply(
+            setSourceFileOnProxy(
+                thrownExceptionElement::getSourceFile,
+                thrownExceptionElement.getRetracedClass(),
+                classResult))
+        .build();
+  }
+
+  private RetraceStackTraceElementProxyResultImpl<T, ST> retraceMethod(
+      RetraceStackTraceElementProxyResultImpl<T, ST> currentResult,
+      ST element,
+      RetraceClassResult classResult,
+      RetraceStackTraceContext context) {
+    Box<RetraceStackTraceContext> resultingContext = new Box<>(RetraceStackTraceContext.empty());
+    RetraceStackTraceElementProxyResultImpl.Builder<T, ST> resultBuilder =
+        currentResult.builder().setResultContext(resultingContext::get);
+    return resultBuilder
+        .setResultStream(
+            currentResult.stream()
+                .flatMap(
+                    proxy -> {
+                      RetraceFrameResult frameResult =
+                          classResult.lookupFrame(
+                              context,
+                              element.hasLineNumber()
+                                  ? OptionalInt.of(element.getLineNumber())
+                                  : OptionalInt.empty(),
+                              element.getMethodName());
+                      return frameResult.stream()
+                          .flatMap(
+                              frameElement -> {
+                                resultingContext.set(frameElement.getRetraceStackTraceContext());
+                                return frameElement
+                                    .streamRewritten(context)
+                                    .map(
+                                        singleFrame ->
+                                            buildProxyForRewrittenFrameElement(
+                                                element,
+                                                classResult,
+                                                proxy,
+                                                frameResult,
+                                                frameElement,
+                                                singleFrame));
+                              });
+                    }))
+        .build();
+  }
+
+  private RetraceStackTraceElementProxyImpl<T, ST> buildProxyForRewrittenFrameElement(
+      ST element,
+      RetraceClassResult classResult,
+      RetraceStackTraceElementProxyImpl<T, ST> proxy,
+      RetraceFrameResult frameResult,
+      RetraceFrameElement frameElement,
+      RetracedSingleFrame singleFrame) {
+    boolean isTopFrame = singleFrame.getIndex() == 0;
+    RetracedMethodReference method = singleFrame.getMethodReference();
+    return proxy
+        .builder()
+        .setRetracedClass(method.getHolderClass())
+        .setRetracedMethod(method)
+        .joinAmbiguous(frameResult.isAmbiguous())
+        .setTopFrame(isTopFrame)
+        .setContext(frameElement.getRetraceStackTraceContext())
+        .applyIf(
+            element.hasLineNumber(),
+            builder -> {
+              builder.setLineNumber(method.getOriginalPositionOrDefault(element.getLineNumber()));
+            })
+        .apply(
+            setSourceFileOnProxy(
+                () -> frameElement.getSourceFile(method), method.getHolderClass(), classResult))
+        .build();
+  }
+
+  private RetraceStackTraceElementProxyResultImpl<T, ST> retraceField(
+      RetraceStackTraceElementProxyResultImpl<T, ST> currentResult,
       ST element,
       RetraceClassResult classResult) {
-    return currentResults.flatMap(
-        proxy -> {
-          RetraceFrameResult frameResult =
-              classResult.lookupFrame(
-                  proxy.context,
-                  element.hasLineNumber()
-                      ? OptionalInt.of(element.getLineNumber())
-                      : OptionalInt.empty(),
-                  element.getMethodName());
-          return frameResult.stream()
-              .flatMap(
-                  frameElement ->
-                      frameElement
-                          .streamRewritten(proxy.getContext())
+    return currentResult
+        .builder()
+        .setResultStream(
+            currentResult.stream()
+                .flatMap(
+                    proxy -> {
+                      RetraceFieldResult retraceFieldResult =
+                          classResult.lookupField(element.getFieldName());
+                      return retraceFieldResult.stream()
                           .map(
-                              singleFrame -> {
-                                boolean isTopFrame = singleFrame.getIndex() == 0;
-                                RetracedMethodReference method = singleFrame.getMethodReference();
-                                return proxy
-                                    .builder()
-                                    .setRetracedClass(method.getHolderClass())
-                                    .setRetracedMethod(method)
-                                    .joinAmbiguous(frameResult.isAmbiguous())
-                                    .setTopFrame(isTopFrame)
-                                    .setContext(frameElement.getRetraceStackTraceContext())
-                                    .applyIf(
-                                        element.hasLineNumber(),
-                                        builder -> {
-                                          builder.setLineNumber(
-                                              method.getOriginalPositionOrDefault(
-                                                  element.getLineNumber()));
-                                        })
-                                    .applyIf(
-                                        element.hasSourceFile(),
-                                        builder -> {
-                                          RetracedSourceFile sourceFileResult =
-                                              frameElement.getSourceFile(method);
-                                          builder.setSourceFile(
-                                              sourceFileResult.hasRetraceResult()
-                                                  ? sourceFileResult.getSourceFile()
-                                                  : RetraceUtils.inferSourceFile(
-                                                      method.getHolderClass().getTypeName(),
-                                                      element.getSourceFile(),
-                                                      classResult.hasRetraceResult()));
-                                        })
-                                    .build();
-                              }));
-        });
+                              fieldElement ->
+                                  buildProxyForRewrittenFieldElement(
+                                      classResult, proxy, retraceFieldResult, fieldElement));
+                    }))
+        .build();
   }
 
-  private Stream<RetraceStackTraceElementProxyImpl<T, ST>> retraceField(
-      Stream<RetraceStackTraceElementProxyImpl<T, ST>> currentResults,
-      ST element,
+  private RetraceStackTraceElementProxyImpl<T, ST> buildProxyForRewrittenFieldElement(
+      RetraceClassResult classResult,
+      RetraceStackTraceElementProxyImpl<T, ST> proxy,
+      RetraceFieldResult retraceFieldResult,
+      RetraceFieldElement fieldElement) {
+    return proxy
+        .builder()
+        .setRetracedClass(fieldElement.getField().getHolderClass())
+        .setRetracedField(fieldElement.getField())
+        .joinAmbiguous(retraceFieldResult.isAmbiguous())
+        .setTopFrame(true)
+        .apply(
+            setSourceFileOnProxy(
+                fieldElement::getSourceFile, fieldElement.getField().getHolderClass(), classResult))
+        .build();
+  }
+
+  private Consumer<RetraceStackTraceElementProxyImpl.Builder<T, ST>> setSourceFileOnProxy(
+      Supplier<RetracedSourceFile> sourceFile,
+      RetracedClassReference classReference,
       RetraceClassResult classResult) {
-    return currentResults.flatMap(
-        proxy -> {
-          RetraceFieldResult retraceFieldResult = classResult.lookupField(element.getFieldName());
-          return retraceFieldResult.stream()
-              .map(
-                  fieldElement ->
-                      proxy
-                          .builder()
-                          .setRetracedClass(fieldElement.getField().getHolderClass())
-                          .setRetracedField(fieldElement.getField())
-                          .joinAmbiguous(retraceFieldResult.isAmbiguous())
-                          .setTopFrame(true)
-                          .applyIf(
-                              element.hasSourceFile(),
-                              builder -> {
-                                RetracedSourceFile sourceFile = fieldElement.getSourceFile();
-                                builder.setSourceFile(
-                                    sourceFile.hasRetraceResult()
-                                        ? sourceFile.getSourceFile()
-                                        : RetraceUtils.inferSourceFile(
-                                            fieldElement.getField().getHolderClass().getTypeName(),
-                                            element.getSourceFile(),
-                                            classResult.hasRetraceResult()));
-                              })
-                          .build());
-        });
+    return proxy -> {
+      ST original = proxy.originalElement;
+      if (!original.hasSourceFile()) {
+        return;
+      }
+      RetracedSourceFile retracedSourceFile = sourceFile.get();
+      proxy.setSourceFile(
+          retracedSourceFile.hasRetraceResult()
+              ? retracedSourceFile.getSourceFile()
+              : RetraceUtils.inferSourceFile(
+                  classReference.getTypeName(),
+                  original.getSourceFile(),
+                  classResult.hasRetraceResult()));
+    };
   }
 
-  private Stream<RetraceStackTraceElementProxyImpl<T, ST>> retraceFieldOrReturnType(
-      Stream<RetraceStackTraceElementProxyImpl<T, ST>> currentResults, ST element) {
+  private RetraceStackTraceElementProxyResultImpl<T, ST> retraceFieldOrReturnType(
+      RetraceStackTraceElementProxyResultImpl<T, ST> currentResult, ST element) {
     if (!element.hasFieldOrReturnType()) {
-      return currentResults;
+      return currentResult;
     }
+    RetraceStackTraceElementProxyResultImpl.Builder<T, ST> resultBuilder = currentResult.builder();
     String elementOrReturnType = element.getFieldOrReturnType();
     if (elementOrReturnType.equals("void")) {
-      return currentResults.map(
-          proxy ->
-              proxy
-                  .builder()
-                  .setRetracedFieldOrReturnType(RetracedTypeReferenceImpl.createVoid())
-                  .build());
+      return resultBuilder
+          .setResultStream(
+              currentResult.stream()
+                  .map(
+                      proxy ->
+                          buildProxyForRewrittenReturnType(
+                              proxy, RetracedTypeReferenceImpl.createVoid(), proxy.isAmbiguous())))
+          .build();
     } else {
       TypeReference typeReference = Reference.typeFromTypeName(elementOrReturnType);
       RetraceTypeResult retraceTypeResult = retracer.retraceType(typeReference);
       List<Element> retracedElements = retraceTypeResult.stream().collect(Collectors.toList());
-      return currentResults.flatMap(
-          proxy ->
-              retracedElements.stream()
-                  .map(
-                      retracedResult ->
-                          proxy
-                              .builder()
-                              .setRetracedFieldOrReturnType(retracedResult.getType())
-                              .joinAmbiguous(retraceTypeResult.isAmbiguous())
-                              .build()));
+      return resultBuilder
+          .setResultStream(
+              currentResult.stream()
+                  .flatMap(
+                      proxy ->
+                          retracedElements.stream()
+                              .map(
+                                  retracedResult ->
+                                      buildProxyForRewrittenReturnType(
+                                          proxy,
+                                          retracedResult.getType(),
+                                          retraceTypeResult.isAmbiguous()))))
+          .build();
     }
   }
 
-  private Stream<RetraceStackTraceElementProxyImpl<T, ST>> retracedMethodArguments(
-      Stream<RetraceStackTraceElementProxyImpl<T, ST>> currentResults, ST element) {
+  private RetraceStackTraceElementProxyImpl<T, ST> buildProxyForRewrittenReturnType(
+      RetraceStackTraceElementProxyImpl<T, ST> proxy,
+      RetracedTypeReference type,
+      boolean isAmbiguous) {
+    return proxy.builder().setRetracedFieldOrReturnType(type).joinAmbiguous(isAmbiguous).build();
+  }
+
+  private RetraceStackTraceElementProxyResultImpl<T, ST> retracedMethodArguments(
+      RetraceStackTraceElementProxyResultImpl<T, ST> currentResult, ST element) {
     if (!element.hasMethodArguments()) {
-      return currentResults;
+      return currentResult;
     }
     List<RetraceTypeResult> retracedResults =
         Arrays.stream(element.getMethodArguments().split(","))
@@ -243,19 +307,24 @@
               return newResult;
             });
     boolean isAmbiguous = allRetracedArguments.size() > 1;
-    return currentResults.flatMap(
-        proxy ->
-            allRetracedArguments.stream()
-                .map(
-                    retracedArguments ->
-                        proxy
-                            .builder()
-                            .setRetracedMethodArguments(retracedArguments)
-                            .joinAmbiguous(isAmbiguous)
-                            .build()));
+    return currentResult
+        .builder()
+        .setResultStream(
+            currentResult.stream()
+                .flatMap(
+                    proxy ->
+                        allRetracedArguments.stream()
+                            .map(
+                                retracedArguments ->
+                                    proxy
+                                        .builder()
+                                        .setRetracedMethodArguments(retracedArguments)
+                                        .joinAmbiguous(isAmbiguous)
+                                        .build())))
+        .build();
   }
 
-  public static class RetraceStackTraceElementProxyImpl<T, ST extends StackTraceElementProxy<T, ST>>
+  static class RetraceStackTraceElementProxyImpl<T, ST extends StackTraceElementProxy<T, ST>>
       implements RetraceStackTraceElementProxy<T, ST> {
 
     private final ST originalItem;
@@ -383,7 +452,7 @@
           originalItem, null, null, null, null, null, null, -1, false, false, context);
     }
 
-    private Builder<T, ST> builder() {
+    Builder<T, ST> builder() {
       Builder<T, ST> builder = new Builder<>(originalItem);
       builder.classContext = retracedClass;
       builder.methodContext = retracedMethod;
@@ -522,6 +591,11 @@
         return this;
       }
 
+      private Builder<T, ST> apply(Consumer<Builder<T, ST>> consumer) {
+        consumer.accept(this);
+        return this;
+      }
+
       private Builder<T, ST> applyIf(boolean condition, Consumer<Builder<T, ST>> consumer) {
         if (condition) {
           consumer.accept(this);
diff --git a/src/test/java/com/android/tools/r8/retrace/RetraceTests.java b/src/test/java/com/android/tools/r8/retrace/RetraceTests.java
index fdc9e96..e4332b9 100644
--- a/src/test/java/com/android/tools/r8/retrace/RetraceTests.java
+++ b/src/test/java/com/android/tools/r8/retrace/RetraceTests.java
@@ -32,6 +32,7 @@
 import com.android.tools.r8.retrace.stacktraces.FoundMethodVerboseStackTrace;
 import com.android.tools.r8.retrace.stacktraces.InlineFileNameStackTrace;
 import com.android.tools.r8.retrace.stacktraces.InlineFileNameWithInnerClassesStackTrace;
+import com.android.tools.r8.retrace.stacktraces.InlineInOutlineStackTrace;
 import com.android.tools.r8.retrace.stacktraces.InlineNoLineNumberStackTrace;
 import com.android.tools.r8.retrace.stacktraces.InlineSourceFileContextStackTrace;
 import com.android.tools.r8.retrace.stacktraces.InlineWithLineNumbersStackTrace;
@@ -47,6 +48,9 @@
 import com.android.tools.r8.retrace.stacktraces.NullStackTrace;
 import com.android.tools.r8.retrace.stacktraces.ObfucatedExceptionClassStackTrace;
 import com.android.tools.r8.retrace.stacktraces.ObfuscatedRangeToSingleLineStackTrace;
+import com.android.tools.r8.retrace.stacktraces.OutlineInOutlineStackTrace;
+import com.android.tools.r8.retrace.stacktraces.OutlineSimpleStackTrace;
+import com.android.tools.r8.retrace.stacktraces.OutlineWithInliningStackTrace;
 import com.android.tools.r8.retrace.stacktraces.OverloadSameLineTest;
 import com.android.tools.r8.retrace.stacktraces.RetraceAssertionErrorStackTrace;
 import com.android.tools.r8.retrace.stacktraces.SingleLineNoLineNumberStackTrace;
@@ -321,6 +325,26 @@
     runRetraceTest(new DifferentLineNumberSpanStackTrace());
   }
 
+  @Test
+  public void testOutlineSimpleStackTrace() throws Exception {
+    runExperimentalRetraceTest(new OutlineSimpleStackTrace());
+  }
+
+  @Test
+  public void testOutlineWithInliningStackTrace() throws Exception {
+    runExperimentalRetraceTest(new OutlineWithInliningStackTrace());
+  }
+
+  @Test
+  public void testOutlineInOutlineStackTrace() throws Exception {
+    runExperimentalRetraceTest(new OutlineInOutlineStackTrace());
+  }
+
+  @Test
+  public void testInlineInOutlineStackTrace() throws Exception {
+    runExperimentalRetraceTest(new InlineInOutlineStackTrace());
+  }
+
   private void inspectRetraceTest(
       StackTraceForTest stackTraceForTest, Consumer<Retracer> inspection) {
     inspection.accept(
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineOutlineTest.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInOutlineStackTrace.java
similarity index 96%
rename from src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineOutlineTest.java
rename to src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInOutlineStackTrace.java
index ac0d27a..c2b5c47 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineOutlineTest.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiOutlineInOutlineStackTrace.java
@@ -25,9 +25,9 @@
 import org.junit.runners.Parameterized;
 
 @RunWith(Parameterized.class)
-public class RetraceApiOutlineOutlineTest extends RetraceApiTestBase {
+public class RetraceApiOutlineInOutlineStackTrace extends RetraceApiTestBase {
 
-  public RetraceApiOutlineOutlineTest(TestParameters parameters) {
+  public RetraceApiOutlineInOutlineStackTrace(TestParameters parameters) {
     super(parameters);
   }
 
diff --git a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiTestCollection.java b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiTestCollection.java
index 5b37834..25d25b2 100644
--- a/src/test/java/com/android/tools/r8/retrace/api/RetraceApiTestCollection.java
+++ b/src/test/java/com/android/tools/r8/retrace/api/RetraceApiTestCollection.java
@@ -38,7 +38,7 @@
           RetraceApiRewriteFrameInlineNpeResidualTest.ApiTest.class,
           RetraceApiOutlineNoInlineTest.ApiTest.class,
           RetraceApiOutlineInlineTest.ApiTest.class,
-          RetraceApiOutlineOutlineTest.ApiTest.class,
+          RetraceApiOutlineInOutlineStackTrace.ApiTest.class,
           RetraceApiInlineInOutlineTest.ApiTest.class,
           RetraceApiSingleFrameTest.ApiTest.class);
 
diff --git a/src/test/java/com/android/tools/r8/retrace/stacktraces/InlineInOutlineStackTrace.java b/src/test/java/com/android/tools/r8/retrace/stacktraces/InlineInOutlineStackTrace.java
new file mode 100644
index 0000000..a231a0f
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/stacktraces/InlineInOutlineStackTrace.java
@@ -0,0 +1,56 @@
+// 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.retrace.stacktraces;
+
+import com.android.tools.r8.utils.StringUtils;
+import java.util.Arrays;
+import java.util.List;
+
+public class InlineInOutlineStackTrace implements StackTraceForTest {
+
+  @Override
+  public List<String> obfuscatedStackTrace() {
+    return Arrays.asList("java.io.IOException: INVALID_SENDER", "\tat a.a(:2)", "\tat b.s(:27)");
+  }
+
+  @Override
+  public String mapping() {
+    return StringUtils.lines(
+        "# { id: 'com.android.tools.r8.mapping', version: 'experimental' }",
+        "outline.Class -> a:",
+        "  1:2:int some.inlinee():75:76 -> a",
+        "  1:2:int outline():0 -> a",
+        "# { 'id':'com.android.tools.r8.outline' }",
+        "some.Class -> b:",
+        "  1:1:void foo.bar.Baz.qux():42:42 -> s",
+        "  5:5:int foo.bar.baz.outlineCaller(int):98:98 -> s",
+        "  5:5:int outlineCaller(int):24 -> s",
+        "  27:27:int outlineCaller(int):0:0 -> s",
+        "# { 'id':'com.android.tools.r8.outlineCallsite', 'positions': { '1': 4, '2': 5 } }");
+  }
+
+  @Override
+  public List<String> retracedStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER",
+        "\tat some.inlinee(some.java:76)",
+        "\tat foo.bar.baz.outlineCaller(baz.java:98)",
+        "\tat some.Class.outlineCaller(Class.java:24)");
+  }
+
+  @Override
+  public List<String> retraceVerboseStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER",
+        "\tat some.int inlinee()(some.java:76)",
+        "\tat foo.bar.baz.int outlineCaller(int)(baz.java:98)",
+        "\tat some.Class.int outlineCaller(int)(Class.java:24)");
+  }
+
+  @Override
+  public int expectedWarnings() {
+    return 0;
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineInOutlineStackTrace.java b/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineInOutlineStackTrace.java
new file mode 100644
index 0000000..ae8b52b
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineInOutlineStackTrace.java
@@ -0,0 +1,55 @@
+// 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.retrace.stacktraces;
+
+import com.android.tools.r8.utils.StringUtils;
+import java.util.Arrays;
+import java.util.List;
+
+public class OutlineInOutlineStackTrace implements StackTraceForTest {
+
+  @Override
+  public List<String> obfuscatedStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER", "\tat a.a(:4)", "\tat b.a(:6)", "\tat c.a(:28)");
+  }
+
+  @Override
+  public String mapping() {
+    return StringUtils.lines(
+        "# { id: 'com.android.tools.r8.mapping', version: 'experimental' }",
+        "outline1.Class -> a:",
+        "  4:4:int outline():0:0 -> a",
+        "# { 'id':'com.android.tools.r8.outline' }",
+        "outline2.Class -> b:",
+        "  6:6:int outline():0:0 -> a",
+        "# { 'id':'com.android.tools.r8.outlineCallsite', 'positions': { '4': 43 } }",
+        "  42:43:int outline():0:0 -> a",
+        "# { 'id':'com.android.tools.r8.outline' }",
+        "some.Class -> c:",
+        "  1:1:void foo.bar.Baz.qux():42:42 -> a",
+        "  10:11:int outlineCaller(int):98:98 -> a",
+        "  28:28:int outlineCaller(int):0:0 -> a",
+        "# { 'id':'com.android.tools.r8.outlineCallsite', 'positions': { '42': 10, '43': 11 } }");
+  }
+
+  @Override
+  public List<String> retracedStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER", "\tat some.Class.outlineCaller(Class.java:98)");
+  }
+
+  @Override
+  public List<String> retraceVerboseStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER",
+        "\tat some.Class.int outlineCaller(int)(Class.java:98)");
+  }
+
+  @Override
+  public int expectedWarnings() {
+    return 0;
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineSimpleStackTrace.java b/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineSimpleStackTrace.java
new file mode 100644
index 0000000..ad0dddc
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineSimpleStackTrace.java
@@ -0,0 +1,49 @@
+// 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.retrace.stacktraces;
+
+import com.android.tools.r8.utils.StringUtils;
+import java.util.Arrays;
+import java.util.List;
+
+public class OutlineSimpleStackTrace implements StackTraceForTest {
+
+  @Override
+  public List<String> obfuscatedStackTrace() {
+    return Arrays.asList("java.io.IOException: INVALID_SENDER", "\tat a.a(:1)", "\tat b.s(:27)");
+  }
+
+  @Override
+  public String mapping() {
+    return StringUtils.lines(
+        "# { id: 'com.android.tools.r8.mapping', version: 'experimental' }",
+        "outline.Class -> a:",
+        "  1:2:int outline() -> a",
+        "# { 'id':'com.android.tools.r8.outline' }",
+        "some.Class -> b:",
+        "  1:1:void foo.bar.Baz.qux():42:42 -> s",
+        "  4:4:int outlineCaller(int):98:98 -> s",
+        "  27:27:int outlineCaller(int):0:0 -> s",
+        "# { 'id':'com.android.tools.r8.outlineCallsite', 'positions': { '1': 4, '2': 5 } }");
+  }
+
+  @Override
+  public List<String> retracedStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER", "\tat some.Class.outlineCaller(Class.java:98)");
+  }
+
+  @Override
+  public List<String> retraceVerboseStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER",
+        "\tat some.Class.int outlineCaller(int)(Class.java:98)");
+  }
+
+  @Override
+  public int expectedWarnings() {
+    return 0;
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineWithInliningStackTrace.java b/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineWithInliningStackTrace.java
new file mode 100644
index 0000000..a8198a2
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/retrace/stacktraces/OutlineWithInliningStackTrace.java
@@ -0,0 +1,53 @@
+// 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.retrace.stacktraces;
+
+import com.android.tools.r8.utils.StringUtils;
+import java.util.Arrays;
+import java.util.List;
+
+public class OutlineWithInliningStackTrace implements StackTraceForTest {
+
+  @Override
+  public List<String> obfuscatedStackTrace() {
+    return Arrays.asList("java.io.IOException: INVALID_SENDER", "\tat a.a(:2)", "\tat b.s(:27)");
+  }
+
+  @Override
+  public String mapping() {
+    return StringUtils.lines(
+        "# { id: 'com.android.tools.r8.mapping', version: 'experimental' }",
+        "outline.Class -> a:",
+        "  1:2:int outline():0 -> a",
+        "# { 'id':'com.android.tools.r8.outline' }",
+        "some.Class -> b:",
+        "  1:1:void foo.bar.Baz.qux():42:42 -> s",
+        "  5:5:int foo.bar.baz.outlineCaller(int):98:98 -> s",
+        "  5:5:int outlineCaller(int):24 -> s",
+        "  27:27:int outlineCaller(int):0:0 -> s",
+        "# { 'id':'com.android.tools.r8.outlineCallsite', 'positions': { '1': 4, '2': 5 } }");
+  }
+
+  @Override
+  public List<String> retracedStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER",
+        "\tat foo.bar.baz.outlineCaller(baz.java:98)",
+        "\tat some.Class.outlineCaller(Class.java:24)");
+  }
+
+  @Override
+  public List<String> retraceVerboseStackTrace() {
+    return Arrays.asList(
+        "java.io.IOException: INVALID_SENDER",
+        "\tat foo.bar.baz.int outlineCaller(int)(baz.java:98)",
+        "\tat some.Class.int outlineCaller(int)(Class.java:24)");
+  }
+
+  @Override
+  public int expectedWarnings() {
+    return 0;
+  }
+}
diff --git a/third_party/retrace/binary_compatibility.tar.gz.sha1 b/third_party/retrace/binary_compatibility.tar.gz.sha1
index 2aa51b9..531812a 100644
--- a/third_party/retrace/binary_compatibility.tar.gz.sha1
+++ b/third_party/retrace/binary_compatibility.tar.gz.sha1
@@ -1 +1 @@
-c2fb502a7f826443d6e41083e9d2f06ef27a5aa7
\ No newline at end of file
+865741441b0a433306aff59af5f3825041411cbe
\ No newline at end of file
diff --git a/tools/utils.py b/tools/utils.py
index 641e352..53aac35 100644
--- a/tools/utils.py
+++ b/tools/utils.py
@@ -41,8 +41,6 @@
 R8 = 'r8'
 R8LIB = 'r8lib'
 R8LIB_NO_DEPS = 'r8LibNoDeps'
-R8RETRACE = 'R8Retrace'
-R8RETRACE_NO_DEPS = 'R8RetraceNoDeps'
 R8_SRC = 'sourceJar'
 LIBRARY_DESUGAR_CONVERSIONS = 'buildLibraryDesugarConversions'
 
@@ -54,8 +52,6 @@
 R8_SRC_JAR = os.path.join(LIBS, 'r8-src.jar')
 R8LIB_EXCLUDE_DEPS_JAR = os.path.join(LIBS, 'r8lib-exclude-deps.jar')
 R8_FULL_EXCLUDE_DEPS_JAR = os.path.join(LIBS, 'r8-full-exclude-deps.jar')
-R8RETRACE_JAR = os.path.join(LIBS, 'r8retrace.jar')
-R8RETRACE_EXCLUDE_DEPS_JAR = os.path.join(LIBS, 'r8retrace-exclude-deps.jar')
 MAVEN_ZIP = os.path.join(LIBS, 'r8.zip')
 MAVEN_ZIP_LIB = os.path.join(LIBS, 'r8lib.zip')
 LIBRARY_DESUGAR_CONVERSIONS_ZIP = os.path.join(LIBS, 'library_desugar_conversions.zip')