Version 2.1.56

Cherry-pick: Add phi check when checking in value definition in service rewriter
CL: https://r8-review.googlesource.com/c/r8/+/52665

Cherry-pick: Reland "Ensure no service loader rewriting if it crosses a feature split"
CL: https://r8-review.googlesource.com/c/r8/+/52542

bug: 162568140
Change-Id: I05cb372595d3755e81690cb8caab06c01b4c9c0b
diff --git a/src/main/java/com/android/tools/r8/FeatureSplit.java b/src/main/java/com/android/tools/r8/FeatureSplit.java
index 5947786..9bb637e 100644
--- a/src/main/java/com/android/tools/r8/FeatureSplit.java
+++ b/src/main/java/com/android/tools/r8/FeatureSplit.java
@@ -29,6 +29,9 @@
  */
 @Keep
 public final class FeatureSplit {
+
+  public static final FeatureSplit BASE = new FeatureSplit(null, null);
+
   private final ProgramConsumer programConsumer;
   private final List<ProgramResourceProvider> programResourceProviders;
 
diff --git a/src/main/java/com/android/tools/r8/Version.java b/src/main/java/com/android/tools/r8/Version.java
index ae9866f..4037c85 100644
--- a/src/main/java/com/android/tools/r8/Version.java
+++ b/src/main/java/com/android/tools/r8/Version.java
@@ -11,7 +11,7 @@
 
   // This field is accessed from release scripts using simple pattern matching.
   // Therefore, changing this field could break our release scripts.
-  public static final String LABEL = "2.1.55";
+  public static final String LABEL = "2.1.56";
 
   private Version() {
   }
diff --git a/src/main/java/com/android/tools/r8/graph/AppServices.java b/src/main/java/com/android/tools/r8/graph/AppServices.java
index 5c1c7ad..0dbaa30 100644
--- a/src/main/java/com/android/tools/r8/graph/AppServices.java
+++ b/src/main/java/com/android/tools/r8/graph/AppServices.java
@@ -8,6 +8,8 @@
 import com.android.tools.r8.DataEntryResource;
 import com.android.tools.r8.DataResourceProvider;
 import com.android.tools.r8.DataResourceProvider.Visitor;
+import com.android.tools.r8.FeatureSplit;
+import com.android.tools.r8.ProgramResourceProvider;
 import com.android.tools.r8.ResourceException;
 import com.android.tools.r8.errors.CompilationError;
 import com.android.tools.r8.origin.Origin;
@@ -21,7 +23,7 @@
 import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.IdentityHashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -36,9 +38,9 @@
   private final AppView<?> appView;
 
   // Mapping from service types to service implementation types.
-  private final Map<DexType, List<DexType>> services;
+  private final Map<DexType, Map<FeatureSplit, List<DexType>>> services;
 
-  private AppServices(AppView<?> appView, Map<DexType, List<DexType>> services) {
+  private AppServices(AppView<?> appView, Map<DexType, Map<FeatureSplit, List<DexType>>> services) {
     this.appView = appView;
     this.services = services;
   }
@@ -54,66 +56,124 @@
 
   public List<DexType> serviceImplementationsFor(DexType serviceType) {
     assert verifyRewrittenWithLens();
-    assert services.containsKey(serviceType);
-    List<DexType> serviceImplementationTypes = services.get(serviceType);
-    if (serviceImplementationTypes == null) {
+    Map<FeatureSplit, List<DexType>> featureSplitListMap = services.get(serviceType);
+    if (featureSplitListMap == null) {
       assert false
           : "Unexpected attempt to get service implementations for non-service type `"
               + serviceType.toSourceString()
               + "`";
       return ImmutableList.of();
     }
-    return serviceImplementationTypes;
+    ImmutableList.Builder<DexType> builder = ImmutableList.builder();
+    for (List<DexType> implementations : featureSplitListMap.values()) {
+      builder.addAll(implementations);
+    }
+    return builder.build();
+  }
+
+  public boolean hasServiceImplementationsInFeature(DexType serviceType) {
+    if (appView.options().featureSplitConfiguration == null) {
+      return false;
+    }
+    Map<FeatureSplit, List<DexType>> featureImplementations = services.get(serviceType);
+    if (featureImplementations == null || featureImplementations.isEmpty()) {
+      assert false
+          : "Unexpected attempt to get service implementations for non-service type `"
+              + serviceType.toSourceString()
+              + "`";
+      return true;
+    }
+    if (featureImplementations.size() > 1
+        || !featureImplementations.containsKey(FeatureSplit.BASE)) {
+      return true;
+    }
+    // Check if service is defined feature
+    DexProgramClass serviceClass = appView.definitionForProgramType(serviceType);
+    if (appView.options().featureSplitConfiguration.isInFeature(serviceClass)) {
+      return true;
+    }
+    for (DexType dexType : featureImplementations.get(FeatureSplit.BASE)) {
+      DexProgramClass implementationClass = appView.definitionForProgramType(dexType);
+      if (appView.options().featureSplitConfiguration.isInFeature(implementationClass)) {
+        return true;
+      }
+    }
+    return false;
   }
 
   public AppServices rewrittenWithLens(GraphLense graphLens) {
-    ImmutableMap.Builder<DexType, List<DexType>> rewrittenServices = ImmutableMap.builder();
-    for (Entry<DexType, List<DexType>> entry : services.entrySet()) {
+    ImmutableMap.Builder<DexType, Map<FeatureSplit, List<DexType>>> rewrittenFeatureMappings =
+        ImmutableMap.builder();
+    for (Entry<DexType, Map<FeatureSplit, List<DexType>>> entry : services.entrySet()) {
       DexType rewrittenServiceType = graphLens.lookupType(entry.getKey());
-      ImmutableList.Builder<DexType> rewrittenServiceImplementationTypes = ImmutableList.builder();
-      for (DexType serviceImplementationType : entry.getValue()) {
-        rewrittenServiceImplementationTypes.add(graphLens.lookupType(serviceImplementationType));
+      ImmutableMap.Builder<FeatureSplit, List<DexType>> rewrittenFeatureImplementations =
+          ImmutableMap.builder();
+      for (Entry<FeatureSplit, List<DexType>> featureSplitImpls : entry.getValue().entrySet()) {
+        ImmutableList.Builder<DexType> rewrittenServiceImplementationTypes =
+            ImmutableList.builder();
+        for (DexType serviceImplementationType : featureSplitImpls.getValue()) {
+          rewrittenServiceImplementationTypes.add(graphLens.lookupType(serviceImplementationType));
+        }
+        rewrittenFeatureImplementations.put(
+            featureSplitImpls.getKey(), rewrittenServiceImplementationTypes.build());
       }
-      rewrittenServices.put(rewrittenServiceType, rewrittenServiceImplementationTypes.build());
+      rewrittenFeatureMappings.put(rewrittenServiceType, rewrittenFeatureImplementations.build());
     }
-    return new AppServices(appView, rewrittenServices.build());
+    return new AppServices(appView, rewrittenFeatureMappings.build());
   }
 
   public AppServices prunedCopy(Collection<DexType> removedClasses) {
-    ImmutableMap.Builder<DexType, List<DexType>> rewrittenServicesBuilder = ImmutableMap.builder();
-    for (Entry<DexType, List<DexType>> entry : services.entrySet()) {
-      if (!removedClasses.contains(entry.getKey())) {
-        DexType serviceType = entry.getKey();
+    ImmutableMap.Builder<DexType, Map<FeatureSplit, List<DexType>>> rewrittenServicesBuilder =
+        ImmutableMap.builder();
+    for (Entry<DexType, Map<FeatureSplit, List<DexType>>> entry : services.entrySet()) {
+      if (removedClasses.contains(entry.getKey())) {
+        continue;
+      }
+      ImmutableMap.Builder<FeatureSplit, List<DexType>> prunedFeatureSplitImpls =
+          ImmutableMap.builder();
+      for (Entry<FeatureSplit, List<DexType>> featureSplitEntry : entry.getValue().entrySet()) {
         ImmutableList.Builder<DexType> rewrittenServiceImplementationTypesBuilder =
             ImmutableList.builder();
-        for (DexType serviceImplementationType : entry.getValue()) {
+        for (DexType serviceImplementationType : featureSplitEntry.getValue()) {
           if (!removedClasses.contains(serviceImplementationType)) {
             rewrittenServiceImplementationTypesBuilder.add(serviceImplementationType);
           }
         }
-        List<DexType> rewrittenServiceImplementationTypes =
+        List<DexType> prunedFeatureSplitImplementations =
             rewrittenServiceImplementationTypesBuilder.build();
-        if (rewrittenServiceImplementationTypes.size() > 0) {
-          rewrittenServicesBuilder.put(
-              serviceType, rewrittenServiceImplementationTypesBuilder.build());
+        if (prunedFeatureSplitImplementations.size() > 0) {
+          prunedFeatureSplitImpls.put(
+              featureSplitEntry.getKey(), rewrittenServiceImplementationTypesBuilder.build());
         }
       }
+      ImmutableMap<FeatureSplit, List<DexType>> prunedServiceImplementations =
+          prunedFeatureSplitImpls.build();
+      if (prunedServiceImplementations.size() > 0) {
+        rewrittenServicesBuilder.put(entry.getKey(), prunedServiceImplementations);
+      }
     }
     return new AppServices(appView, rewrittenServicesBuilder.build());
   }
 
   private boolean verifyRewrittenWithLens() {
-    for (Entry<DexType, List<DexType>> entry : services.entrySet()) {
+    for (Entry<DexType, Map<FeatureSplit, List<DexType>>> entry : services.entrySet()) {
       assert entry.getKey() == appView.graphLense().lookupType(entry.getKey());
-      for (DexType type : entry.getValue()) {
-        assert type == appView.graphLense().lookupType(type);
+      for (Entry<FeatureSplit, List<DexType>> featureEntry : entry.getValue().entrySet()) {
+        for (DexType type : featureEntry.getValue()) {
+          assert type == appView.graphLense().lookupType(type);
+        }
       }
     }
     return true;
   }
 
   public void visit(BiConsumer<DexType, List<DexType>> consumer) {
-    services.forEach(consumer);
+    services.forEach(
+        (type, featureImpls) -> {
+          ImmutableList.Builder<DexType> builder = ImmutableList.builder();
+          featureImpls.values().forEach(builder::addAll);
+          consumer.accept(type, builder.build());
+        });
   }
 
   public static Builder builder(AppView<?> appView) {
@@ -123,7 +183,7 @@
   public static class Builder {
 
     private final AppView<?> appView;
-    private final Map<DexType, List<DexType>> services = new IdentityHashMap<>();
+    private final Map<DexType, Map<FeatureSplit, List<DexType>>> services = new LinkedHashMap<>();
 
     private Builder(AppView<?> appView) {
       this.appView = appView;
@@ -131,14 +191,27 @@
 
     public AppServices build() {
       for (DataResourceProvider provider : appView.appInfo().app().dataResourceProviders) {
-        readServices(provider);
+        readServices(provider, FeatureSplit.BASE);
+      }
+      if (appView.options().featureSplitConfiguration != null) {
+        List<FeatureSplit> featureSplits =
+            appView.options().featureSplitConfiguration.getFeatureSplits();
+        for (FeatureSplit featureSplit : featureSplits) {
+          for (ProgramResourceProvider provider : featureSplit.getProgramResourceProviders()) {
+            DataResourceProvider dataResourceProvider = provider.getDataResourceProvider();
+            if (dataResourceProvider != null) {
+              readServices(dataResourceProvider, featureSplit);
+            }
+          }
+        }
       }
       return new AppServices(appView, services);
     }
 
-    private void readServices(DataResourceProvider dataResourceProvider) {
+    private void readServices(
+        DataResourceProvider dataResourceProvider, FeatureSplit featureSplit) {
       try {
-        dataResourceProvider.accept(new DataResourceProviderVisitor());
+        dataResourceProvider.accept(new DataResourceProviderVisitor(featureSplit));
       } catch (ResourceException e) {
         throw new CompilationError(e.getMessage(), e);
       }
@@ -146,6 +219,12 @@
 
     private class DataResourceProviderVisitor implements Visitor {
 
+      private final FeatureSplit featureSplit;
+
+      public DataResourceProviderVisitor(FeatureSplit featureSplit) {
+        this.featureSplit = featureSplit;
+      }
+
       @Override
       public void visit(DataDirectoryResource directory) {
         // Ignore.
@@ -162,8 +241,10 @@
               DexType serviceType = appView.dexItemFactory().createType(serviceDescriptor);
               byte[] bytes = ByteStreams.toByteArray(file.getByteStream());
               String contents = new String(bytes, Charset.defaultCharset());
+              Map<FeatureSplit, List<DexType>> featureSplitImplementations =
+                  services.computeIfAbsent(serviceType, k -> new LinkedHashMap<>());
               List<DexType> serviceImplementations =
-                  services.computeIfAbsent(serviceType, (key) -> new ArrayList<>());
+                  featureSplitImplementations.computeIfAbsent(featureSplit, f -> new ArrayList<>());
               readServiceImplementationsForService(
                   contents, file.getOrigin(), serviceImplementations);
             }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java b/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java
index d0ad73a..235ede5 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/ServiceLoaderRewriter.java
@@ -130,8 +130,16 @@
         continue;
       }
 
-      // Check that ClassLoader used is the ClassLoader defined for the the service configuration
+      // Check that we are not service loading anything from a feature into base.
+      if (appView.appServices().hasServiceImplementationsInFeature(constClass.getValue())) {
+        continue;
+      }
+
+      // Check that ClassLoader used is the ClassLoader defined for the service configuration
       // that we are instantiating or NULL.
+      if (serviceLoaderLoad.inValues().get(1).isPhi()) {
+        continue;
+      }
       InvokeVirtual classLoaderInvoke =
           serviceLoaderLoad.inValues().get(1).definition.asInvokeVirtual();
       boolean isGetClassLoaderOnConstClassOrNull =
diff --git a/src/test/java/com/android/tools/r8/rewrite/ServiceLoaderRewritingTest.java b/src/test/java/com/android/tools/r8/rewrite/ServiceLoaderRewritingTest.java
index 5f46a45..1ccedb0 100644
--- a/src/test/java/com/android/tools/r8/rewrite/ServiceLoaderRewritingTest.java
+++ b/src/test/java/com/android/tools/r8/rewrite/ServiceLoaderRewritingTest.java
@@ -135,6 +135,20 @@
     }
   }
 
+  public static class LoadWhereClassLoaderIsPhi {
+
+    public static void main(String[] args) {
+      ServiceLoader.load(
+              Service.class,
+              System.currentTimeMillis() > 0
+                  ? Thread.currentThread().getContextClassLoader()
+                  : null)
+          .iterator()
+          .next()
+          .print();
+    }
+  }
+
   @Parameterized.Parameters(name = "{0}")
   public static TestParametersCollection data() {
     return getTestParameters().withAllRuntimesAndApiLevels().build();
@@ -207,7 +221,7 @@
     testForR8(parameters.getBackend())
         .addInnerClasses(ServiceLoaderRewritingTest.class)
         .addKeepMainRule(MainWithTryCatchRunner.class)
-        .setMinApi(parameters.getRuntime())
+        .setMinApi(parameters.getApiLevel())
         .addDataEntryResources(
             DataEntryResource.fromBytes(
                 StringUtils.lines(ServiceImpl.class.getTypeName(), ServiceImpl2.class.getTypeName())
@@ -267,7 +281,7 @@
             .addInnerClasses(ServiceLoaderRewritingTest.class)
             .addKeepMainRule(EscapingRunner.class)
             .enableInliningAnnotations()
-            .setMinApi(parameters.getRuntime())
+            .setMinApi(parameters.getApiLevel())
             .noMinification()
             .addDataEntryResources(
                 DataEntryResource.fromBytes(
@@ -291,6 +305,37 @@
   }
 
   @Test
+  public void testDoNoRewriteWhenClassLoaderIsPhi()
+      throws IOException, CompilationFailedException, ExecutionException {
+    Path path = temp.newFile("out.zip").toPath();
+    CodeInspector inspector =
+        testForR8(parameters.getBackend())
+            .addInnerClasses(ServiceLoaderRewritingTest.class)
+            .addKeepMainRule(LoadWhereClassLoaderIsPhi.class)
+            .enableInliningAnnotations()
+            .setMinApi(parameters.getApiLevel())
+            .addDataEntryResources(
+                DataEntryResource.fromBytes(
+                    StringUtils.lines(ServiceImpl.class.getTypeName()).getBytes(),
+                    "META-INF/services/" + Service.class.getTypeName(),
+                    Origin.unknown()))
+            .compile()
+            .writeToZip(path)
+            .run(parameters.getRuntime(), LoadWhereClassLoaderIsPhi.class)
+            .assertSuccessWithOutputLines("Hello World!")
+            .inspector();
+
+    // Check that we have not rewritten the calls to ServiceLoader.load.
+    assertEquals(1, getServiceLoaderLoads(inspector, LoadWhereClassLoaderIsPhi.class));
+
+    // Check that we have not removed the service configuration from META-INF/services.
+    ZipFile zip = new ZipFile(path.toFile());
+    ClassSubject serviceImpl = inspector.clazz(ServiceImpl.class);
+    assertTrue(serviceImpl.isPresent());
+    assertNotNull(zip.getEntry("META-INF/services/" + serviceImpl.getFinalName()));
+  }
+
+  @Test
   public void testKeepAsOriginal()
       throws IOException, CompilationFailedException, ExecutionException {
     // The CL that changed behaviour after Nougat is: