diff --git a/src/main/java/com/android/tools/r8/DexSegments.java b/src/main/java/com/android/tools/r8/DexSegments.java
index 117006b..81b904c 100644
--- a/src/main/java/com/android/tools/r8/DexSegments.java
+++ b/src/main/java/com/android/tools/r8/DexSegments.java
@@ -154,6 +154,11 @@
     return run(command.getInputApp());
   }
 
+  public static Map<Integer, SegmentInfo> runForTesting(AndroidApp app)
+      throws IOException, ResourceException {
+    return run(app);
+  }
+
   public static Int2ReferenceMap<SegmentInfo> run(AndroidApp app)
       throws IOException, ResourceException {
     Int2ReferenceMap<SegmentInfo> result = new Int2ReferenceLinkedOpenHashMap<>();
diff --git a/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java b/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
index 6bd6774..e9956ad 100644
--- a/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
+++ b/src/test/testbase/java/com/android/tools/r8/R8TestCompileResult.java
@@ -8,6 +8,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 
+import com.android.tools.r8.DexSegments.SegmentInfo;
 import com.android.tools.r8.TestRuntime.DexRuntime;
 import com.android.tools.r8.ToolHelper.DexVm;
 import com.android.tools.r8.ToolHelper.ProcessResult;
@@ -34,6 +35,8 @@
 import com.android.tools.r8.utils.codeinspector.CodeInspector;
 import com.android.tools.r8.utils.codeinspector.FoundClassSubject;
 import com.android.tools.r8.utils.graphinspector.GraphInspector;
+import it.unimi.dsi.fastutil.ints.Int2ReferenceLinkedOpenHashMap;
+import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
@@ -363,11 +366,19 @@
     if (results.isBenchmarkingCodeSize()) {
       AndroidApp appWithFeatures =
           features.isEmpty() ? app : AndroidApp.builder(app).addProgramFiles(features).build();
-      results.addDexSegmentsSizeResult(DexSegments.run(appWithFeatures));
+      results.addDexSegmentsSizeResult(runDexSegments(appWithFeatures));
     }
     return self();
   }
 
+  private Int2ReferenceMap<SegmentInfo> runDexSegments(AndroidApp app)
+      throws IOException, ResourceException {
+    Map<Integer, SegmentInfo> result = DexSegments.runForTesting(app);
+    Int2ReferenceMap<SegmentInfo> rewrittenResult = new Int2ReferenceLinkedOpenHashMap<>();
+    rewrittenResult.putAll(result);
+    return rewrittenResult;
+  }
+
   @Override
   public R8TestCompileResult benchmarkDex2OatCodeSize(BenchmarkResults results) throws IOException {
     if (results.isBenchmarkingCodeSize()) {