Untangle optimized resource shrinking from legacy

This moves all optimized shrinking logic to R8ResourceShrinkerState, only initializing this once based on the resource inputs. The parsing of resource tables is now only done at setup time before starting the enqueuer.

This prepares us for tracing the xml files directly.

Also, add support for feature splits in optimized shrinking.

Recursive tracing, i.e., marking of transitively reachable resources
are now done through the trace method when code gives rise to a new
resource.

Bug: b/287398085
Change-Id: I85b20442f472bb02203fdc88ea42c6df47eba584
diff --git a/src/main/java/com/android/tools/r8/R8.java b/src/main/java/com/android/tools/r8/R8.java
index eb4825c..b00b3af 100644
--- a/src/main/java/com/android/tools/r8/R8.java
+++ b/src/main/java/com/android/tools/r8/R8.java
@@ -918,7 +918,7 @@
     dexFileContent.forEach(resourceShrinkerBuilder::addDexInput);
     try {
       addResourcesToBuilder(
-          resourceShrinkerBuilder, reporter, options.androidResourceProvider, null);
+          resourceShrinkerBuilder, reporter, options.androidResourceProvider, FeatureSplit.BASE);
       if (options.featureSplitConfiguration != null) {
         for (FeatureSplit featureSplit : options.featureSplitConfiguration.getFeatureSplits()) {
           if (featureSplit.getAndroidResourceProvider() != null) {
@@ -936,9 +936,7 @@
       LegacyResourceShrinker shrinker = resourceShrinkerBuilder.build();
       ShrinkerResult shrinkerResult;
       if (options.resourceShrinkerConfiguration.isOptimizedShrinking()) {
-        shrinkerResult =
-            shrinker.shrinkModel(
-                appView.getResourceShrinkerState().getR8ResourceShrinkerModel(), true);
+        shrinkerResult = appView.getResourceShrinkerState().shrinkModel();
       } else {
         shrinkerResult = shrinker.run();
       }
@@ -949,7 +947,7 @@
           toKeep,
           options.androidResourceProvider,
           options.androidResourceConsumer,
-          null);
+          FeatureSplit.BASE);
       if (options.featureSplitConfiguration != null) {
         for (FeatureSplit featureSplit : options.featureSplitConfiguration.getFeatureSplits()) {
           if (featureSplit.getAndroidResourceProvider() != null) {
diff --git a/src/main/java/com/android/tools/r8/graph/AppView.java b/src/main/java/com/android/tools/r8/graph/AppView.java
index e3c04a8..abb62bf 100644
--- a/src/main/java/com/android/tools/r8/graph/AppView.java
+++ b/src/main/java/com/android/tools/r8/graph/AppView.java
@@ -64,6 +64,7 @@
 import com.android.tools.r8.utils.InternalOptions.TestingOptions;
 import com.android.tools.r8.utils.OptionalBool;
 import com.android.tools.r8.utils.Reporter;
+import com.android.tools.r8.utils.ResourceShrinkerUtils;
 import com.android.tools.r8.utils.ThrowingConsumer;
 import com.android.tools.r8.utils.Timing;
 import com.android.tools.r8.utils.threads.ThreadTask;
@@ -218,6 +219,9 @@
     } else {
       this.argumentPropagator = null;
     }
+    if (enableWholeProgramOptimizations() && options().isOptimizedResourceShrinking()) {
+      resourceShrinkerState = ResourceShrinkerUtils.createResourceShrinkerState(this);
+    }
     timing.end();
     this.libraryMethodSideEffectModelCollection =
         timing.time("Library side-effects", () -> new LibraryMethodSideEffectModelCollection(this));
diff --git a/src/main/java/com/android/tools/r8/graph/analysis/ResourceAccessAnalysis.java b/src/main/java/com/android/tools/r8/graph/analysis/ResourceAccessAnalysis.java
index 6b64f53..5670096 100644
--- a/src/main/java/com/android/tools/r8/graph/analysis/ResourceAccessAnalysis.java
+++ b/src/main/java/com/android/tools/r8/graph/analysis/ResourceAccessAnalysis.java
@@ -5,9 +5,6 @@
 package com.android.tools.r8.graph.analysis;
 
 import com.android.build.shrinker.r8integration.R8ResourceShrinkerState;
-import com.android.tools.r8.AndroidResourceInput;
-import com.android.tools.r8.AndroidResourceInput.Kind;
-import com.android.tools.r8.ResourceException;
 import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
 import com.android.tools.r8.graph.AppView;
@@ -44,18 +41,7 @@
       AppView<? extends AppInfoWithClassHierarchy> appView, Enqueuer enqueuer) {
     this.appView = appView;
     this.enqueuer = enqueuer;
-    this.resourceShrinkerState = new R8ResourceShrinkerState();
-    try {
-      for (AndroidResourceInput androidResource :
-          appView.options().androidResourceProvider.getAndroidResources()) {
-        if (androidResource.getKind() == Kind.RESOURCE_TABLE) {
-          resourceShrinkerState.setResourceTableInput(androidResource.getByteStream());
-          break;
-        }
-      }
-    } catch (ResourceException e) {
-      throw appView.reporter().fatalError("Failed initializing resource table");
-    }
+    this.resourceShrinkerState = appView.getResourceShrinkerState();
   }
 
   public static void register(
@@ -68,6 +54,8 @@
   @Override
   public void done(Enqueuer enqueuer) {
     EnqueuerFieldAccessAnalysis.super.done(enqueuer);
+    // We clear the bits here, since we will trace the final reachable entries in the second round.
+    resourceShrinkerState.clearReachableBits();
   }
 
   private static boolean enabled(
diff --git a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
index 3fa01e0..2a36daf 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -17,11 +17,7 @@
 import static com.android.tools.r8.utils.MapUtils.ignoreKey;
 import static java.util.Collections.emptySet;
 
-import com.android.build.shrinker.r8integration.R8ResourceShrinkerState;
-import com.android.tools.r8.AndroidResourceInput;
-import com.android.tools.r8.AndroidResourceInput.Kind;
 import com.android.tools.r8.Diagnostic;
-import com.android.tools.r8.ResourceException;
 import com.android.tools.r8.cf.code.CfInstruction;
 import com.android.tools.r8.cf.code.CfInvoke;
 import com.android.tools.r8.contexts.CompilationContext.MethodProcessingContext;
@@ -487,8 +483,6 @@
 
   private final ProfileCollectionAdditions profileCollectionAdditions;
 
-  private final R8ResourceShrinkerState r8ResourceShrinkerState;
-
   Enqueuer(
       AppView<? extends AppInfoWithClassHierarchy> appView,
       ProfileCollectionAdditions profileCollectionAdditions,
@@ -547,27 +541,6 @@
 
     objectAllocationInfoCollection =
         ObjectAllocationInfoCollectionImpl.builder(mode.isInitialTreeShaking(), graphReporter);
-    r8ResourceShrinkerState = setupResourceShrinkerState(appView);
-  }
-
-  private R8ResourceShrinkerState setupResourceShrinkerState(
-      AppView<? extends AppInfoWithClassHierarchy> appView) {
-    R8ResourceShrinkerState r8ResourceShrinkerState = new R8ResourceShrinkerState();
-    if (options.resourceShrinkerConfiguration.isOptimizedShrinking()
-        && options.androidResourceProvider != null) {
-      try {
-        for (AndroidResourceInput androidResource :
-            options.androidResourceProvider.getAndroidResources()) {
-          if (androidResource.getKind() == Kind.RESOURCE_TABLE) {
-            r8ResourceShrinkerState.setResourceTableInput(androidResource.getByteStream());
-            break;
-          }
-        }
-      } catch (ResourceException e) {
-        throw appView.reporter().fatalError("Failed initializing resource table");
-      }
-    }
-    return r8ResourceShrinkerState;
   }
 
   private AppInfoWithClassHierarchy appInfo() {
@@ -1146,7 +1119,7 @@
   }
 
   public void traceResourceValue(int value) {
-    r8ResourceShrinkerState.trace(value);
+    appView.getResourceShrinkerState().trace(value);
   }
 
   public void traceReflectiveFieldWrite(ProgramField field, ProgramMethod context) {
@@ -3792,7 +3765,6 @@
     EnqueuerResult result = createEnqueuerResult(appInfo, timing);
     profileCollectionAdditions.commit(appView);
     timing.end();
-    appView.setResourceShrinkerState(r8ResourceShrinkerState);
     return result;
   }
 
diff --git a/src/main/java/com/android/tools/r8/utils/ResourceShrinkerUtils.java b/src/main/java/com/android/tools/r8/utils/ResourceShrinkerUtils.java
new file mode 100644
index 0000000..2d4e961
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/utils/ResourceShrinkerUtils.java
@@ -0,0 +1,88 @@
+// 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.utils;
+
+import com.android.build.shrinker.r8integration.R8ResourceShrinkerState;
+import com.android.tools.r8.AndroidResourceInput;
+import com.android.tools.r8.FeatureSplit;
+import com.android.tools.r8.ResourceException;
+import com.android.tools.r8.graph.AppView;
+import java.io.InputStream;
+import java.util.Collection;
+
+public class ResourceShrinkerUtils {
+
+  public static R8ResourceShrinkerState createResourceShrinkerState(AppView<?> appView) {
+    R8ResourceShrinkerState state =
+        new R8ResourceShrinkerState(
+            exception -> appView.reporter().fatalError(new ExceptionDiagnostic(exception)));
+    InternalOptions options = appView.options();
+    if (options.resourceShrinkerConfiguration.isOptimizedShrinking()
+        && options.androidResourceProvider != null) {
+      try {
+        addResources(
+            appView,
+            state,
+            options.androidResourceProvider.getAndroidResources(),
+            FeatureSplit.BASE);
+        if (options.hasFeatureSplitConfiguration()) {
+          for (FeatureSplit featureSplit :
+              options.getFeatureSplitConfiguration().getFeatureSplits()) {
+            if (featureSplit.getAndroidResourceProvider() != null) {
+              addResources(
+                  appView,
+                  state,
+                  featureSplit.getAndroidResourceProvider().getAndroidResources(),
+                  featureSplit);
+            }
+          }
+        }
+      } catch (ResourceException e) {
+        throw appView.reporter().fatalError("Failed initializing resource table");
+      }
+      state.setupReferences();
+    }
+    return state;
+  }
+
+  private static void addResources(
+      AppView<?> appView,
+      R8ResourceShrinkerState state,
+      Collection<AndroidResourceInput> androidResources,
+      FeatureSplit featureSplit)
+      throws ResourceException {
+    for (AndroidResourceInput androidResource : androidResources) {
+      switch (androidResource.getKind()) {
+        case MANIFEST:
+          state.setManifestProvider(
+              () -> wrapThrowingInputStreamResource(appView, androidResource));
+          break;
+        case RESOURCE_TABLE:
+          state.addResourceTable(androidResource.getByteStream(), featureSplit);
+          break;
+        case XML_FILE:
+          state.addXmlFileProvider(
+              () -> wrapThrowingInputStreamResource(appView, androidResource),
+              androidResource.getPath().location());
+          break;
+        case RES_FOLDER_FILE:
+          state.addResFileProvider(
+              () -> wrapThrowingInputStreamResource(appView, androidResource),
+              androidResource.getPath().location());
+          break;
+        case UNKNOWN:
+          break;
+      }
+    }
+  }
+
+  private static InputStream wrapThrowingInputStreamResource(
+      AppView<?> appView, AndroidResourceInput androidResource) {
+    try {
+      return androidResource.getByteStream();
+    } catch (ResourceException ex) {
+      throw appView.reporter().fatalError("Failed reading " + androidResource.getPath().location());
+    }
+  }
+}
diff --git a/src/resourceshrinker/java/com/android/build/shrinker/r8integration/LegacyResourceShrinker.java b/src/resourceshrinker/java/com/android/build/shrinker/r8integration/LegacyResourceShrinker.java
index f37d381..7874b5c 100644
--- a/src/resourceshrinker/java/com/android/build/shrinker/r8integration/LegacyResourceShrinker.java
+++ b/src/resourceshrinker/java/com/android/build/shrinker/r8integration/LegacyResourceShrinker.java
@@ -14,7 +14,6 @@
 import com.android.build.shrinker.ResourceShrinkerImplKt;
 import com.android.build.shrinker.ResourceTableUtilKt;
 import com.android.build.shrinker.graph.ProtoResourcesGraphBuilder;
-import com.android.build.shrinker.graph.ResFolderFileTree;
 import com.android.build.shrinker.obfuscation.ProguardMappingsRecorder;
 import com.android.build.shrinker.r8integration.R8ResourceShrinkerState.R8ResourceShrinkerModel;
 import com.android.build.shrinker.usages.DexFileAnalysisCallback;
@@ -134,11 +133,6 @@
       ResourceTable loadedResourceTable = ResourceTable.parseFrom(pathAndBytes.bytes);
       model.instantiateFromResourceTable(loadedResourceTable, false);
     }
-    return shrinkModel(model, false);
-  }
-
-  public ShrinkerResult shrinkModel(
-      R8ResourceShrinkerModel model, boolean exactMatchingOfStyleablesAndAttr) throws IOException {
     if (proguardMapStrings != null) {
       new ProguardMappingsRecorder(proguardMapStrings).recordObfuscationMappings(model);
       proguardMapStrings = null;
@@ -171,13 +165,7 @@
     for (PathAndBytes pathAndBytes : resourceTables.keySet()) {
       ResourceTable resourceTable = ResourceTable.parseFrom(pathAndBytes.bytes);
       new ProtoResourcesGraphBuilder(
-              new ResFolderFileTree() {
-                @Override
-                public byte[] getEntryByName(String pathInRes) {
-                  return resFolderMappings.get(pathInRes).getBytes();
-                }
-              },
-              unused -> resourceTable)
+              pathInRes -> resFolderMappings.get(pathInRes).getBytes(), unused -> resourceTable)
           .buildGraph(model);
     }
     ResourceStore resourceStore = model.getResourceStore();
@@ -194,8 +182,7 @@
         resEntriesToKeep.add(xmlInput.path.toString());
       }
     }
-    List<Integer> resourceIdsToRemove =
-        getResourceIdsToRemove(unusedResources, model, exactMatchingOfStyleablesAndAttr);
+    List<Integer> resourceIdsToRemove = getResourceIdsToRemove(unusedResources);
     Map<FeatureSplit, ResourceTable> shrunkenTables = new HashMap<>();
     for (Entry<PathAndBytes, FeatureSplit> entry : resourceTables.entrySet()) {
       ResourceTable shrunkenResourceTable =
@@ -206,21 +193,11 @@
     return new ShrinkerResult(resEntriesToKeep.build(), shrunkenTables);
   }
 
-  private static List<Integer> getResourceIdsToRemove(
-      List<Resource> unusedResources,
-      R8ResourceShrinkerModel model,
-      boolean exactMatchingOfStyleablesAndAttr) {
-    if (!exactMatchingOfStyleablesAndAttr) {
+  private static List<Integer> getResourceIdsToRemove(List<Resource> unusedResources) {
       return unusedResources.stream()
           .filter(s -> s.type != ResourceType.ID)
           .map(resource -> resource.value)
           .collect(Collectors.toList());
-    }
-    return model.getResourceStore().getResources().stream()
-        .filter(r -> !r.isReachable())
-        .filter(r -> r.type != ResourceType.ID)
-        .map(r -> r.value)
-        .collect(Collectors.toList());
   }
 
   // Lifted from com/android/utils/XmlUtils.java which we can't easily update internal dependency
diff --git a/src/resourceshrinker/java/com/android/build/shrinker/r8integration/R8ResourceShrinkerState.java b/src/resourceshrinker/java/com/android/build/shrinker/r8integration/R8ResourceShrinkerState.java
index 5e0a8ae..2392218 100644
--- a/src/resourceshrinker/java/com/android/build/shrinker/r8integration/R8ResourceShrinkerState.java
+++ b/src/resourceshrinker/java/com/android/build/shrinker/r8integration/R8ResourceShrinkerState.java
@@ -1,48 +1,182 @@
 // Copyright (c) 2023, 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.build.shrinker.r8integration;
 
+import static com.android.build.shrinker.r8integration.LegacyResourceShrinker.getUtfReader;
+
 import com.android.aapt.Resources.ConfigValue;
 import com.android.aapt.Resources.Entry;
 import com.android.aapt.Resources.Item;
 import com.android.aapt.Resources.ResourceTable;
 import com.android.aapt.Resources.Value;
+import com.android.aapt.Resources.XmlNode;
 import com.android.build.shrinker.NoDebugReporter;
+import com.android.build.shrinker.ResourceShrinkerImplKt;
 import com.android.build.shrinker.ResourceShrinkerModel;
 import com.android.build.shrinker.ResourceTableUtilKt;
 import com.android.build.shrinker.ShrinkerDebugReporter;
+import com.android.build.shrinker.graph.ProtoResourcesGraphBuilder;
+import com.android.build.shrinker.r8integration.LegacyResourceShrinker.ShrinkerResult;
+import com.android.build.shrinker.usages.ProtoAndroidManifestUsageRecorderKt;
+import com.android.build.shrinker.usages.ToolsAttributeUsageRecorderKt;
 import com.android.ide.common.resources.ResourcesUtil;
+import com.android.ide.common.resources.usage.ResourceStore;
 import com.android.ide.common.resources.usage.ResourceUsageModel;
 import com.android.ide.common.resources.usage.ResourceUsageModel.Resource;
 import com.android.resources.ResourceType;
+import com.android.tools.r8.FeatureSplit;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 public class R8ResourceShrinkerState {
 
-  private R8ResourceShrinkerModel r8ResourceShrinkerModel;
+  private final Function<Exception, RuntimeException> errorHandler;
+  private final R8ResourceShrinkerModel r8ResourceShrinkerModel;
+  private final Map<String, Supplier<InputStream>> xmlFileProviders = new HashMap<>();
+
+  private Supplier<InputStream> manifestProvider;
+  private final Map<String, Supplier<InputStream>> resfileProviders = new HashMap<>();
+  private final Map<ResourceTable, FeatureSplit> resourceTables = new HashMap<>();
+
+  public R8ResourceShrinkerState(Function<Exception, RuntimeException> errorHandler) {
+    r8ResourceShrinkerModel = new R8ResourceShrinkerModel(NoDebugReporter.INSTANCE, true);
+    this.errorHandler = errorHandler;
+  }
 
   public List<String> trace(int id) {
     Resource resource = r8ResourceShrinkerModel.getResourceStore().getResource(id);
+    if (resource == null) {
+      return Collections.emptyList();
+    }
     ResourceUsageModel.markReachable(resource);
+    if (resource.references != null) {
+      for (Resource reference : resource.references) {
+        if (!reference.isReachable()) {
+          trace(reference.value);
+        }
+      }
+    }
     return Collections.emptyList();
   }
 
-  public void setResourceTableInput(InputStream inputStream) {
-    r8ResourceShrinkerModel = new R8ResourceShrinkerModel(NoDebugReporter.INSTANCE, true);
-    r8ResourceShrinkerModel.instantiateFromResourceTable(inputStream, true);
+  public void setManifestProvider(Supplier<InputStream> manifestProvider) {
+    this.manifestProvider = manifestProvider;
+  }
+
+  public void addXmlFileProvider(Supplier<InputStream> inputStreamSupplier, String location) {
+    this.xmlFileProviders.put(location, inputStreamSupplier);
+  }
+
+  public void addResFileProvider(Supplier<InputStream> inputStreamSupplier, String location) {
+    this.resfileProviders.put(location, inputStreamSupplier);
+  }
+
+  public void addResourceTable(InputStream inputStream, FeatureSplit featureSplit) {
+    this.resourceTables.put(
+        r8ResourceShrinkerModel.instantiateFromResourceTable(inputStream, true), featureSplit);
   }
 
   public R8ResourceShrinkerModel getR8ResourceShrinkerModel() {
     return r8ResourceShrinkerModel;
   }
 
+  private byte[] getXmlOrResFileBytes(String path) {
+    assert !path.startsWith("res/");
+    String pathWithRes = "res/" + path;
+    Supplier<InputStream> inputStreamSupplier = xmlFileProviders.get(pathWithRes);
+    if (inputStreamSupplier == null) {
+      inputStreamSupplier = resfileProviders.get(pathWithRes);
+    }
+    if (inputStreamSupplier == null) {
+      // Ill formed resource table with file references inside res/ that does not exist.
+      return null;
+    }
+    try {
+      return inputStreamSupplier.get().readAllBytes();
+    } catch (IOException ex) {
+      throw errorHandler.apply(ex);
+    }
+  }
+
+  public void setupReferences() {
+    for (ResourceTable resourceTable : resourceTables.keySet()) {
+      new ProtoResourcesGraphBuilder(this::getXmlOrResFileBytes, unused -> resourceTable)
+          .buildGraph(r8ResourceShrinkerModel);
+    }
+  }
+
+  public ShrinkerResult shrinkModel() throws IOException {
+    updateModelWithManifestReferences();
+    updateModelWithKeepXmlReferences();
+    ResourceStore resourceStore = r8ResourceShrinkerModel.getResourceStore();
+    resourceStore.processToolsAttributes();
+    ImmutableSet<String> resEntriesToKeep = getResEntriesToKeep(resourceStore);
+    List<Integer> resourceIdsToRemove = getResourcesToRemove();
+
+    Map<FeatureSplit, ResourceTable> shrunkenTables = new IdentityHashMap<>();
+    resourceTables.forEach(
+        (resourceTable, featureSplit) ->
+            shrunkenTables.put(
+                featureSplit,
+                ResourceTableUtilKt.nullOutEntriesWithIds(resourceTable, resourceIdsToRemove)));
+
+    return new ShrinkerResult(resEntriesToKeep, shrunkenTables);
+  }
+
+  private ImmutableSet<String> getResEntriesToKeep(ResourceStore resourceStore) {
+    ImmutableSet.Builder<String> resEntriesToKeep = new ImmutableSet.Builder<>();
+    for (String path : Iterables.concat(xmlFileProviders.keySet(), resfileProviders.keySet())) {
+      if (ResourceShrinkerImplKt.isJarPathReachable(resourceStore, path)) {
+        resEntriesToKeep.add(path);
+      }
+    }
+    return resEntriesToKeep.build();
+  }
+
+  private List<Integer> getResourcesToRemove() {
+    return r8ResourceShrinkerModel.getResourceStore().getResources().stream()
+        .filter(r -> !r.isReachable() && !r.isPublic())
+        .filter(r -> r.type != ResourceType.ID)
+        .map(r -> r.value)
+        .collect(Collectors.toList());
+  }
+
+  // Temporary to support updating the reachable entries from the manifest, we need to instead
+  // trace these in the enqueuer.
+  public void updateModelWithManifestReferences() throws IOException {
+    if (manifestProvider == null) {
+      return;
+    }
+    ProtoAndroidManifestUsageRecorderKt.recordUsagesFromNode(
+        XmlNode.parseFrom(manifestProvider.get()), r8ResourceShrinkerModel);
+  }
+
+  public void updateModelWithKeepXmlReferences() throws IOException {
+    for (Map.Entry<String, Supplier<InputStream>> entry : xmlFileProviders.entrySet()) {
+      if (entry.getKey().startsWith("res/raw")) {
+        ToolsAttributeUsageRecorderKt.processRawXml(
+            getUtfReader(entry.getValue().get().readAllBytes()), r8ResourceShrinkerModel);
+      }
+    }
+  }
+
+  public void clearReachableBits() {
+    for (Resource resource : r8ResourceShrinkerModel.getResourceStore().getResources()) {
+      resource.setReachable(false);
+    }
+  }
+
   public static class R8ResourceShrinkerModel extends ResourceShrinkerModel {
     private final Map<Integer, String> stringResourcesWithSingleValue = new HashMap<>();
 
@@ -56,10 +190,11 @@
     }
 
     // Similar to instantiation in ProtoResourceTableGatherer, but using an inputstream.
-    void instantiateFromResourceTable(InputStream inputStream, boolean includeStyleables) {
+    ResourceTable instantiateFromResourceTable(InputStream inputStream, boolean includeStyleables) {
       try {
         ResourceTable resourceTable = ResourceTable.parseFrom(inputStream);
         instantiateFromResourceTable(resourceTable, includeStyleables);
+        return resourceTable;
       } catch (IOException ex) {
         throw new RuntimeException(ex);
       }
diff --git a/src/test/java/com/android/tools/r8/R8TestBuilder.java b/src/test/java/com/android/tools/r8/R8TestBuilder.java
index 2c3ea07..bcd5227 100644
--- a/src/test/java/com/android/tools/r8/R8TestBuilder.java
+++ b/src/test/java/com/android/tools/r8/R8TestBuilder.java
@@ -920,16 +920,19 @@
   public T addFeatureSplitAndroidResources(AndroidTestResource testResource, String featureName)
       throws IOException {
     Path outputFile = getState().getNewTempFile("resourceshrinkeroutput_" + featureName + ".zip");
+    Path programOut = getState().getNewTempFile("feature_output" + featureName + ".jar");
     resourceShrinkerOutputForFeatures.put(featureName, outputFile);
     getBuilder()
         .addFeatureSplit(
             featureSplitGenerator -> {
-              featureSplitGenerator.setAndroidResourceConsumer(
-                  new ArchiveProtoAndroidResourceConsumer(outputFile));
               Path resourceZip = testResource.getResourceZip();
-              featureSplitGenerator.setAndroidResourceProvider(
-                  new ArchiveProtoAndroidResourceProvider(
-                      resourceZip, new PathOrigin(resourceZip)));
+              featureSplitGenerator
+                  .setAndroidResourceConsumer(new ArchiveProtoAndroidResourceConsumer(outputFile))
+                  .setAndroidResourceProvider(
+                      new ArchiveProtoAndroidResourceProvider(
+                          resourceZip, new PathOrigin(resourceZip)))
+                  .setProgramConsumer(DexIndexedConsumer.emptyConsumer());
+
               return featureSplitGenerator.build();
             });
     return addProgramClassFileData(testResource.getRClass().getClassFileData());
diff --git a/src/test/java/com/android/tools/r8/androidresources/RClassResourceGeneration.java b/src/test/java/com/android/tools/r8/androidresources/RClassResourceGeneration.java
index 7cb981b..9f5f557 100644
--- a/src/test/java/com/android/tools/r8/androidresources/RClassResourceGeneration.java
+++ b/src/test/java/com/android/tools/r8/androidresources/RClassResourceGeneration.java
@@ -31,7 +31,7 @@
 
   @Parameters(name = "{0}")
   public static TestParametersCollection parameters() {
-    return getTestParameters().withDexRuntimes().withAllApiLevels().build();
+    return getTestParameters().withDefaultDexRuntime().withAllApiLevels().build();
   }
 
   @BeforeClass
diff --git a/src/test/java/com/android/tools/r8/androidresources/ResourceShrinkingWithFeatures.java b/src/test/java/com/android/tools/r8/androidresources/ResourceShrinkingWithFeatures.java
index 2b257b4..221f011 100644
--- a/src/test/java/com/android/tools/r8/androidresources/ResourceShrinkingWithFeatures.java
+++ b/src/test/java/com/android/tools/r8/androidresources/ResourceShrinkingWithFeatures.java
@@ -8,15 +8,17 @@
 
 import com.android.tools.r8.CompilationFailedException;
 import com.android.tools.r8.R8FullTestBuilder;
+import com.android.tools.r8.R8TestBuilder;
 import com.android.tools.r8.R8TestCompileResult;
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.TestParameters;
-import com.android.tools.r8.TestParametersCollection;
 import com.android.tools.r8.ToolHelper;
 import com.android.tools.r8.androidresources.AndroidResourceTestingUtils.AndroidTestResource;
 import com.android.tools.r8.androidresources.AndroidResourceTestingUtils.AndroidTestResourceBuilder;
 import com.android.tools.r8.androidresources.ResourceShrinkingWithFeatures.FeatureSplit.FeatureSplitMain;
+import com.android.tools.r8.utils.BooleanUtils;
 import java.io.IOException;
+import java.util.List;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
@@ -30,9 +32,14 @@
   @Parameter(0)
   public TestParameters parameters;
 
-  @Parameters(name = "{0}")
-  public static TestParametersCollection parameters() {
-    return getTestParameters().withDexRuntimes().withAllApiLevels().build();
+  @Parameter(1)
+  public boolean optimized;
+
+  @Parameters(name = "{0}, optimized: {1}")
+  public static List<Object[]> data() {
+    return buildParameters(
+        getTestParameters().withDefaultDexRuntime().withAllApiLevels().build(),
+        BooleanUtils.values());
   }
 
   public static AndroidTestResource getTestResources(TemporaryFolder temp) throws Exception {
@@ -57,6 +64,7 @@
       testForR8(parameters.getBackend())
           .setMinApi(parameters)
           .addProgramClasses(Base.class)
+          .applyIf(optimized, R8TestBuilder::enableOptimizedShrinking)
           .addFeatureSplit(builder -> builder.build())
           .compileWithExpectedDiagnostics(
               diagnostics -> {
@@ -95,6 +103,7 @@
             .addAndroidResources(getTestResources(temp))
             .addFeatureSplitAndroidResources(
                 getFeatureSplitTestResources(featureSplitTemp), FeatureSplit.class.getName())
+            .applyIf(optimized, R8FullTestBuilder::enableOptimizedShrinking)
             .addKeepMainRule(Base.class)
             .addKeepMainRule(FeatureSplitMain.class)
             .compile();
diff --git a/src/test/java/com/android/tools/r8/androidresources/SimpleNoCodeReferenceAndroidResourceTest.java b/src/test/java/com/android/tools/r8/androidresources/SimpleNoCodeReferenceAndroidResourceTest.java
index 34b2b9a..030af1c 100644
--- a/src/test/java/com/android/tools/r8/androidresources/SimpleNoCodeReferenceAndroidResourceTest.java
+++ b/src/test/java/com/android/tools/r8/androidresources/SimpleNoCodeReferenceAndroidResourceTest.java
@@ -35,7 +35,7 @@
 
   @Parameters(name = "{0}")
   public static TestParametersCollection parameters() {
-    return getTestParameters().withDexRuntimes().withAllApiLevels().build();
+    return getTestParameters().withDefaultDexRuntime().withAllApiLevels().build();
   }
 
   @Test
diff --git a/src/test/java/com/android/tools/r8/androidresources/TestShrinkingWithCodeReferences.java b/src/test/java/com/android/tools/r8/androidresources/TestShrinkingWithCodeReferences.java
index 00b000f..5bd64e3 100644
--- a/src/test/java/com/android/tools/r8/androidresources/TestShrinkingWithCodeReferences.java
+++ b/src/test/java/com/android/tools/r8/androidresources/TestShrinkingWithCodeReferences.java
@@ -23,7 +23,7 @@
 
   @Parameters(name = "{0}")
   public static TestParametersCollection parameters() {
-    return getTestParameters().withDexRuntimes().withAllApiLevels().build();
+    return getTestParameters().withDefaultDexRuntime().withAllApiLevels().build();
   }
 
   public static AndroidTestResource getTestResources(TemporaryFolder temp) throws Exception {
diff --git a/src/test/java/com/android/tools/r8/androidresources/XmlFilesWithReferences.java b/src/test/java/com/android/tools/r8/androidresources/XmlFilesWithReferences.java
index fb43ba7..9f80fe1 100644
--- a/src/test/java/com/android/tools/r8/androidresources/XmlFilesWithReferences.java
+++ b/src/test/java/com/android/tools/r8/androidresources/XmlFilesWithReferences.java
@@ -23,7 +23,7 @@
 
   @Parameters(name = "{0}")
   public static TestParametersCollection parameters() {
-    return getTestParameters().withDexRuntimes().withAllApiLevels().build();
+    return getTestParameters().withDefaultDexRuntime().withAllApiLevels().build();
   }
 
   public static AndroidTestResource getTestResources(TemporaryFolder temp) throws Exception {