Emit feature split metadata

Bug: b/367558466
Change-Id: If95422f098f8eb4702903b56ccd6cd33dc9db15a
diff --git a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
index c1b8ee6..55274e8 100644
--- a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
@@ -213,8 +213,7 @@
     return desugaredLibraryCodeToKeep;
   }
 
-  private List<VirtualFile> distribute(ExecutorService executorService)
-      throws ExecutionException, IOException {
+  private List<VirtualFile> distribute(ExecutorService executorService) {
     Collection<DexProgramClass> classes = appView.appInfo().classes();
     Collection<DexProgramClass> globalSynthetics = new ArrayList<>();
     if (appView.options().intermediate && appView.options().hasGlobalSyntheticsConsumer()) {
diff --git a/src/main/java/com/android/tools/r8/dex/VirtualFile.java b/src/main/java/com/android/tools/r8/dex/VirtualFile.java
index 2d96710..c8f16fc 100644
--- a/src/main/java/com/android/tools/r8/dex/VirtualFile.java
+++ b/src/main/java/com/android/tools/r8/dex/VirtualFile.java
@@ -161,6 +161,10 @@
     return featureSplit;
   }
 
+  public FeatureSplit getFeatureSplitOrBase() {
+    return featureSplit != null ? featureSplit : FeatureSplit.BASE;
+  }
+
   public StartupProfile getStartupProfile() {
     return startupProfile;
   }
diff --git a/src/main/java/com/android/tools/r8/metadata/R8BuildMetadata.java b/src/main/java/com/android/tools/r8/metadata/R8BuildMetadata.java
index 096f514..6d118ab 100644
--- a/src/main/java/com/android/tools/r8/metadata/R8BuildMetadata.java
+++ b/src/main/java/com/android/tools/r8/metadata/R8BuildMetadata.java
@@ -8,6 +8,9 @@
 import com.android.tools.r8.metadata.impl.R8BaselineProfileRewritingOptionsImpl;
 import com.android.tools.r8.metadata.impl.R8BuildMetadataImpl;
 import com.android.tools.r8.metadata.impl.R8CompilationInfoImpl;
+import com.android.tools.r8.metadata.impl.R8DexFileMetadataImpl;
+import com.android.tools.r8.metadata.impl.R8FeatureSplitMetadataImpl;
+import com.android.tools.r8.metadata.impl.R8FeatureSplitsMetadataImpl;
 import com.android.tools.r8.metadata.impl.R8KeepAttributesOptionsImpl;
 import com.android.tools.r8.metadata.impl.R8LibraryDesugaringOptionsImpl;
 import com.android.tools.r8.metadata.impl.R8OptionsImpl;
@@ -30,6 +33,11 @@
             R8BaselineProfileRewritingOptions.class,
             deserializeTo(R8BaselineProfileRewritingOptionsImpl.class))
         .registerTypeAdapter(R8CompilationInfo.class, deserializeTo(R8CompilationInfoImpl.class))
+        .registerTypeAdapter(R8DexFileMetadata.class, deserializeTo(R8DexFileMetadataImpl.class))
+        .registerTypeAdapter(
+            R8FeatureSplitMetadata.class, deserializeTo(R8FeatureSplitMetadataImpl.class))
+        .registerTypeAdapter(
+            R8FeatureSplitsMetadata.class, deserializeTo(R8FeatureSplitsMetadataImpl.class))
         .registerTypeAdapter(
             R8KeepAttributesOptions.class, deserializeTo(R8KeepAttributesOptionsImpl.class))
         .registerTypeAdapter(
@@ -60,7 +68,12 @@
   /**
    * @return null if not compiling to dex.
    */
-  List<String> getDexChecksums();
+  List<R8DexFileMetadata> getDexFilesMetadata();
+
+  /**
+   * @return null if not using feature splits.
+   */
+  R8FeatureSplitsMetadata getFeatureSplitsMetadata();
 
   /**
    * @return null if resource optimization is disabled.
diff --git a/src/main/java/com/android/tools/r8/metadata/R8DexFileMetadata.java b/src/main/java/com/android/tools/r8/metadata/R8DexFileMetadata.java
new file mode 100644
index 0000000..fc841d0
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/metadata/R8DexFileMetadata.java
@@ -0,0 +1,21 @@
+// Copyright (c) 2024, 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.metadata;
+
+import com.android.tools.r8.keepanno.annotations.KeepForApi;
+
+@KeepForApi
+public interface R8DexFileMetadata {
+
+  /**
+   * Returns the SHA-256 checksum of the entire dex file.
+   *
+   * <p>This can be used to check if the given dex file has been tampered with after compilation.
+   *
+   * <p>Note: This differs from the checksum in the dex format, as the checksum embedded in the dex
+   * is the adler32 checksum of the dex file excluding the magic value and the checksum itself. See
+   * also https://source.android.com/docs/core/runtime/dex-format.
+   */
+  String getChecksum();
+}
diff --git a/src/main/java/com/android/tools/r8/metadata/R8FeatureSplitMetadata.java b/src/main/java/com/android/tools/r8/metadata/R8FeatureSplitMetadata.java
new file mode 100644
index 0000000..f9c5d47
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/metadata/R8FeatureSplitMetadata.java
@@ -0,0 +1,13 @@
+// Copyright (c) 2024, 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.metadata;
+
+import com.android.tools.r8.keepanno.annotations.KeepForApi;
+import java.util.List;
+
+@KeepForApi
+public interface R8FeatureSplitMetadata {
+
+  List<R8DexFileMetadata> getDexFilesMetadata();
+}
diff --git a/src/main/java/com/android/tools/r8/metadata/R8FeatureSplitsMetadata.java b/src/main/java/com/android/tools/r8/metadata/R8FeatureSplitsMetadata.java
new file mode 100644
index 0000000..66fd1e2
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/metadata/R8FeatureSplitsMetadata.java
@@ -0,0 +1,15 @@
+// Copyright (c) 2024, 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.metadata;
+
+import com.android.tools.r8.keepanno.annotations.KeepForApi;
+import java.util.List;
+
+@KeepForApi
+public interface R8FeatureSplitsMetadata {
+
+  List<R8FeatureSplitMetadata> getFeatureSplits();
+
+  boolean isIsolatedSplitsEnabled();
+}
diff --git a/src/main/java/com/android/tools/r8/metadata/impl/BuildMetadataFactory.java b/src/main/java/com/android/tools/r8/metadata/impl/BuildMetadataFactory.java
index 31a76a8..1bb4869 100644
--- a/src/main/java/com/android/tools/r8/metadata/impl/BuildMetadataFactory.java
+++ b/src/main/java/com/android/tools/r8/metadata/impl/BuildMetadataFactory.java
@@ -3,6 +3,9 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.metadata.impl;
 
+import static com.android.tools.r8.utils.MapUtils.ignoreKey;
+
+import com.android.tools.r8.FeatureSplit;
 import com.android.tools.r8.Version;
 import com.android.tools.r8.dex.VirtualFile;
 import com.android.tools.r8.graph.AppInfo;
@@ -11,7 +14,11 @@
 import com.android.tools.r8.metadata.D8BuildMetadata;
 import com.android.tools.r8.metadata.R8BuildMetadata;
 import com.android.tools.r8.utils.InternalOptions;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ExecutorService;
 
 public class BuildMetadataFactory {
@@ -27,15 +34,30 @@
       AppView<? extends AppInfoWithClassHierarchy> appView,
       ExecutorService executorService,
       List<VirtualFile> virtualFiles) {
+    Map<FeatureSplit, List<VirtualFile>> virtualFilesForFeatureSplit = new IdentityHashMap<>();
+    for (VirtualFile virtualFile : virtualFiles) {
+      FeatureSplit featureSplit = virtualFile.getFeatureSplitOrBase();
+      virtualFilesForFeatureSplit
+          .computeIfAbsent(featureSplit, ignoreKey(ArrayList::new))
+          .add(virtualFile);
+    }
+    List<VirtualFile> baseVirtualFiles =
+        virtualFilesForFeatureSplit.getOrDefault(FeatureSplit.BASE, Collections.emptyList());
     InternalOptions options = appView.options();
     return R8BuildMetadataImpl.builder()
         .setOptions(new R8OptionsImpl(options))
         .setBaselineProfileRewritingOptions(R8BaselineProfileRewritingOptionsImpl.create(options))
         .setCompilationInfo(R8CompilationInfoImpl.create(executorService, options))
-        .applyIf(options.isGeneratingDex(), builder -> builder.setDexChecksums(virtualFiles))
+        .applyIf(
+            options.isGeneratingDex(), builder -> builder.setDexFilesMetadata(baseVirtualFiles))
+        .applyIf(
+            options.hasFeatureSplitConfiguration(),
+            builder ->
+                builder.setFeatureSplitsMetadata(
+                    R8FeatureSplitsMetadataImpl.create(appView, virtualFilesForFeatureSplit)))
         .setResourceOptimizationOptions(R8ResourceOptimizationOptionsImpl.create(options))
         .setStartupOptimizationOptions(
-            R8StartupOptimizationOptionsImpl.create(options, virtualFiles))
+            R8StartupOptimizationOptionsImpl.create(options, baseVirtualFiles))
         .setVersion(Version.LABEL)
         .build();
   }
diff --git a/src/main/java/com/android/tools/r8/metadata/impl/R8BuildMetadataImpl.java b/src/main/java/com/android/tools/r8/metadata/impl/R8BuildMetadataImpl.java
index daf5193..edabf58 100644
--- a/src/main/java/com/android/tools/r8/metadata/impl/R8BuildMetadataImpl.java
+++ b/src/main/java/com/android/tools/r8/metadata/impl/R8BuildMetadataImpl.java
@@ -14,9 +14,12 @@
 import com.android.tools.r8.metadata.R8BaselineProfileRewritingOptions;
 import com.android.tools.r8.metadata.R8BuildMetadata;
 import com.android.tools.r8.metadata.R8CompilationInfo;
+import com.android.tools.r8.metadata.R8DexFileMetadata;
+import com.android.tools.r8.metadata.R8FeatureSplitsMetadata;
 import com.android.tools.r8.metadata.R8Options;
 import com.android.tools.r8.metadata.R8ResourceOptimizationOptions;
 import com.android.tools.r8.metadata.R8StartupOptimizationOptions;
+import com.android.tools.r8.utils.ListUtils;
 import com.google.gson.Gson;
 import com.google.gson.annotations.Expose;
 import com.google.gson.annotations.SerializedName;
@@ -46,8 +49,12 @@
   private final R8CompilationInfo compilationInfo;
 
   @Expose
-  @SerializedName("dexChecksums")
-  private final List<String> dexChecksums;
+  @SerializedName("dexFilesMetadata")
+  private final List<R8DexFileMetadata> dexFilesMetadata;
+
+  @Expose
+  @SerializedName("featureSplitsMetadata")
+  private final R8FeatureSplitsMetadata featureSplitsMetadata;
 
   @Expose
   @SerializedName("resourceOptimizationOptions")
@@ -65,14 +72,16 @@
       R8Options options,
       R8BaselineProfileRewritingOptions baselineProfileRewritingOptions,
       R8CompilationInfo compilationInfo,
-      List<String> dexChecksums,
+      List<R8DexFileMetadata> dexFilesMetadata,
+      R8FeatureSplitsMetadata featureSplitsMetadata,
       R8ResourceOptimizationOptions resourceOptimizationOptions,
       R8StartupOptimizationOptions startupOptimizationOptions,
       String version) {
     this.options = options;
     this.baselineProfileRewritingOptions = baselineProfileRewritingOptions;
     this.compilationInfo = compilationInfo;
-    this.dexChecksums = dexChecksums;
+    this.dexFilesMetadata = dexFilesMetadata;
+    this.featureSplitsMetadata = featureSplitsMetadata;
     this.resourceOptimizationOptions = resourceOptimizationOptions;
     this.startupOptimizationOptions = startupOptimizationOptions;
     this.version = version;
@@ -98,8 +107,13 @@
   }
 
   @Override
-  public List<String> getDexChecksums() {
-    return dexChecksums;
+  public List<R8DexFileMetadata> getDexFilesMetadata() {
+    return dexFilesMetadata;
+  }
+
+  @Override
+  public R8FeatureSplitsMetadata getFeatureSplitsMetadata() {
+    return featureSplitsMetadata;
   }
 
   @Override
@@ -127,7 +141,8 @@
     private R8Options options;
     private R8BaselineProfileRewritingOptions baselineProfileRewritingOptions;
     private R8CompilationInfo compilationInfo;
-    private List<String> dexChecksums;
+    private List<R8DexFileMetadata> dexFilesMetadata;
+    private R8FeatureSplitsMetadata featureSplitsMetadata;
     private R8ResourceOptimizationOptions resourceOptimizationOptions;
     private R8StartupOptimizationOptions startupOptimizationOptions;
     private String version;
@@ -155,12 +170,18 @@
       return this;
     }
 
-    public Builder setDexChecksums(List<VirtualFile> virtualFiles) {
-      this.dexChecksums =
+    public Builder setDexFilesMetadata(List<VirtualFile> virtualFiles) {
+      List<String> checksums =
           virtualFiles.stream()
               .filter(not(VirtualFile::isEmpty))
               .map(virtualFile -> virtualFile.getChecksumForBuildMetadata().toString())
               .collect(Collectors.toList());
+      this.dexFilesMetadata = ListUtils.map(checksums, R8DexFileMetadataImpl::new);
+      return this;
+    }
+
+    public Builder setFeatureSplitsMetadata(R8FeatureSplitsMetadata featureSplitsMetadata) {
+      this.featureSplitsMetadata = featureSplitsMetadata;
       return this;
     }
 
@@ -186,7 +207,8 @@
           options,
           baselineProfileRewritingOptions,
           compilationInfo,
-          dexChecksums,
+          dexFilesMetadata,
+          featureSplitsMetadata,
           resourceOptimizationOptions,
           startupOptimizationOptions,
           version);
diff --git a/src/main/java/com/android/tools/r8/metadata/impl/R8DexFileMetadataImpl.java b/src/main/java/com/android/tools/r8/metadata/impl/R8DexFileMetadataImpl.java
new file mode 100644
index 0000000..e0680dc
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/metadata/impl/R8DexFileMetadataImpl.java
@@ -0,0 +1,36 @@
+// Copyright (c) 2024, 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.metadata.impl;
+
+import com.android.tools.r8.keepanno.annotations.AnnotationPattern;
+import com.android.tools.r8.keepanno.annotations.FieldAccessFlags;
+import com.android.tools.r8.keepanno.annotations.KeepConstraint;
+import com.android.tools.r8.keepanno.annotations.KeepItemKind;
+import com.android.tools.r8.keepanno.annotations.UsedByReflection;
+import com.android.tools.r8.metadata.R8DexFileMetadata;
+import com.google.gson.annotations.Expose;
+import com.google.gson.annotations.SerializedName;
+
+@UsedByReflection(
+    description = "Keep and preserve @SerializedName for correct (de)serialization",
+    constraints = {KeepConstraint.LOOKUP},
+    constrainAnnotations = @AnnotationPattern(constant = SerializedName.class),
+    kind = KeepItemKind.CLASS_AND_FIELDS,
+    fieldAccess = {FieldAccessFlags.PRIVATE},
+    fieldAnnotatedByClassConstant = SerializedName.class)
+public class R8DexFileMetadataImpl implements R8DexFileMetadata {
+
+  @Expose
+  @SerializedName("checksum")
+  private final String checksum;
+
+  public R8DexFileMetadataImpl(String checksum) {
+    this.checksum = checksum;
+  }
+
+  @Override
+  public String getChecksum() {
+    return checksum;
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/metadata/impl/R8FeatureSplitMetadataImpl.java b/src/main/java/com/android/tools/r8/metadata/impl/R8FeatureSplitMetadataImpl.java
new file mode 100644
index 0000000..ae567d7
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/metadata/impl/R8FeatureSplitMetadataImpl.java
@@ -0,0 +1,38 @@
+// Copyright (c) 2024, 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.metadata.impl;
+
+import com.android.tools.r8.keepanno.annotations.AnnotationPattern;
+import com.android.tools.r8.keepanno.annotations.FieldAccessFlags;
+import com.android.tools.r8.keepanno.annotations.KeepConstraint;
+import com.android.tools.r8.keepanno.annotations.KeepItemKind;
+import com.android.tools.r8.keepanno.annotations.UsedByReflection;
+import com.android.tools.r8.metadata.R8DexFileMetadata;
+import com.android.tools.r8.metadata.R8FeatureSplitMetadata;
+import com.google.gson.annotations.Expose;
+import com.google.gson.annotations.SerializedName;
+import java.util.List;
+
+@UsedByReflection(
+    description = "Keep and preserve @SerializedName for correct (de)serialization",
+    constraints = {KeepConstraint.LOOKUP},
+    constrainAnnotations = @AnnotationPattern(constant = SerializedName.class),
+    kind = KeepItemKind.CLASS_AND_FIELDS,
+    fieldAccess = {FieldAccessFlags.PRIVATE},
+    fieldAnnotatedByClassConstant = SerializedName.class)
+public class R8FeatureSplitMetadataImpl implements R8FeatureSplitMetadata {
+
+  @Expose
+  @SerializedName("dexFilesMetadata")
+  private final List<R8DexFileMetadata> dexFilesMetadata;
+
+  public R8FeatureSplitMetadataImpl(List<R8DexFileMetadata> dexFilesMetadata) {
+    this.dexFilesMetadata = dexFilesMetadata;
+  }
+
+  @Override
+  public List<R8DexFileMetadata> getDexFilesMetadata() {
+    return dexFilesMetadata;
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/metadata/impl/R8FeatureSplitsMetadataImpl.java b/src/main/java/com/android/tools/r8/metadata/impl/R8FeatureSplitsMetadataImpl.java
new file mode 100644
index 0000000..5d558b2
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/metadata/impl/R8FeatureSplitsMetadataImpl.java
@@ -0,0 +1,82 @@
+// Copyright (c) 2024, 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.metadata.impl;
+
+import com.android.tools.r8.FeatureSplit;
+import com.android.tools.r8.dex.VirtualFile;
+import com.android.tools.r8.features.FeatureSplitConfiguration;
+import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.keepanno.annotations.AnnotationPattern;
+import com.android.tools.r8.keepanno.annotations.FieldAccessFlags;
+import com.android.tools.r8.keepanno.annotations.KeepConstraint;
+import com.android.tools.r8.keepanno.annotations.KeepItemKind;
+import com.android.tools.r8.keepanno.annotations.UsedByReflection;
+import com.android.tools.r8.metadata.R8DexFileMetadata;
+import com.android.tools.r8.metadata.R8FeatureSplitMetadata;
+import com.android.tools.r8.metadata.R8FeatureSplitsMetadata;
+import com.android.tools.r8.utils.InternalOptions;
+import com.android.tools.r8.utils.ListUtils;
+import com.google.gson.annotations.Expose;
+import com.google.gson.annotations.SerializedName;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+@UsedByReflection(
+    description = "Keep and preserve @SerializedName for correct (de)serialization",
+    constraints = {KeepConstraint.LOOKUP},
+    constrainAnnotations = @AnnotationPattern(constant = SerializedName.class),
+    kind = KeepItemKind.CLASS_AND_FIELDS,
+    fieldAccess = {FieldAccessFlags.PRIVATE},
+    fieldAnnotatedByClassConstant = SerializedName.class)
+public class R8FeatureSplitsMetadataImpl implements R8FeatureSplitsMetadata {
+
+  @Expose
+  @SerializedName("featureSplitsMetadata")
+  private final List<R8FeatureSplitMetadata> featureSplitsMetadata;
+
+  @Expose
+  @SerializedName("isolatedSplits")
+  private final boolean isolatedSplits;
+
+  public R8FeatureSplitsMetadataImpl(
+      FeatureSplitConfiguration configuration, List<R8FeatureSplitMetadata> featureSplitsMetadata) {
+    this.featureSplitsMetadata = featureSplitsMetadata;
+    this.isolatedSplits = configuration.isIsolatedSplitsEnabled();
+  }
+
+  public static R8FeatureSplitsMetadataImpl create(
+      AppView<? extends AppInfoWithClassHierarchy> appView,
+      Map<FeatureSplit, List<VirtualFile>> virtualFilesForFeatureSplit) {
+    InternalOptions options = appView.options();
+    assert options.hasFeatureSplitConfiguration();
+    FeatureSplitConfiguration configuration = options.getFeatureSplitConfiguration();
+    List<FeatureSplit> featureSplits = configuration.getFeatureSplits();
+    List<R8FeatureSplitMetadata> featureSplitsMetadata = new ArrayList<>();
+    for (FeatureSplit featureSplit : featureSplits) {
+      List<VirtualFile> featureSplitVirtualFiles =
+          virtualFilesForFeatureSplit.getOrDefault(featureSplit, Collections.emptyList());
+      List<R8DexFileMetadata> dexFilesMetadata =
+          ListUtils.map(
+              featureSplitVirtualFiles,
+              featureSplitVirtualFile ->
+                  new R8DexFileMetadataImpl(
+                      featureSplitVirtualFile.getChecksumForBuildMetadata().toString()));
+      featureSplitsMetadata.add(new R8FeatureSplitMetadataImpl(dexFilesMetadata));
+    }
+    return new R8FeatureSplitsMetadataImpl(configuration, featureSplitsMetadata);
+  }
+
+  @Override
+  public List<R8FeatureSplitMetadata> getFeatureSplits() {
+    return featureSplitsMetadata;
+  }
+
+  @Override
+  public boolean isIsolatedSplitsEnabled() {
+    return isolatedSplits;
+  }
+}
diff --git a/src/test/java/com/android/tools/r8/metadata/R8BuildMetadataTest.java b/src/test/java/com/android/tools/r8/metadata/R8BuildMetadataTest.java
index 04b401c..1365c88 100644
--- a/src/test/java/com/android/tools/r8/metadata/R8BuildMetadataTest.java
+++ b/src/test/java/com/android/tools/r8/metadata/R8BuildMetadataTest.java
@@ -55,11 +55,13 @@
             .addKeepMainRule(Main.class)
             .addArtProfileForRewriting(
                 ExternalArtProfile.builder().addClassRule(mainReference).build())
-            .apply(StartupTestingUtils.addStartupProfile(startupProfile))
             .applyIf(
                 parameters.isDexRuntime(),
                 testBuilder ->
-                    testBuilder.addAndroidResources(getTestResources()).enableOptimizedShrinking())
+                    testBuilder
+                        .addAndroidResources(getTestResources())
+                        .apply(StartupTestingUtils.addStartupProfile(startupProfile))
+                        .enableOptimizedShrinking())
             .allowDiagnosticInfoMessages(parameters.canUseNativeMultidex())
             .collectBuildMetadata()
             .setMinApi(parameters)
@@ -103,10 +105,14 @@
     }
     R8StartupOptimizationOptions startupOptimizationOptions =
         buildMetadata.getStartupOptizationOptions();
-    assertNotNull(startupOptimizationOptions);
-    assertEquals(
-        parameters.isDexRuntime() && parameters.canUseNativeMultidex() ? 1 : 0,
-        startupOptimizationOptions.getNumberOfStartupDexFiles());
+    if (parameters.isDexRuntime()) {
+      assertNotNull(startupOptimizationOptions);
+      assertEquals(
+          parameters.isDexRuntime() && parameters.canUseNativeMultidex() ? 1 : 0,
+          startupOptimizationOptions.getNumberOfStartupDexFiles());
+    } else {
+      assertNull(startupOptimizationOptions);
+    }
     assertEquals(Version.LABEL, buildMetadata.getVersion());
   }