API to configure partial optimization

This is an API which is supposed to replace the experimental API
enableExperimentalPartialShrinking.

The initial API supports including packages and classes. Both packages
and classes are specified using the R8 reference API.

Mixing the experimental API (including through system properties) and
this new API is not allowed.

Change-Id: I0d99989d972ab1da017dc8a6a55d5e66a46dd9e6
Bug: b/309743298
diff --git a/src/main/java/com/android/tools/r8/PartialOptimizationConfigurationBuilder.java b/src/main/java/com/android/tools/r8/PartialOptimizationConfigurationBuilder.java
new file mode 100644
index 0000000..63604f1
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/PartialOptimizationConfigurationBuilder.java
@@ -0,0 +1,31 @@
+// Copyright (c) 2025, 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;
+
+import com.android.tools.r8.keepanno.annotations.KeepForApi;
+import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.references.PackageReference;
+
+/** API for building partial optimization configuration for the compiler. */
+@KeepForApi
+public interface PartialOptimizationConfigurationBuilder {
+
+  /**
+   * Add a class to be optimized with partial optimization. Note that this does *not* include inner
+   * classes.
+   *
+   * @param classReference class to be optimized.
+   * @return instance which received the call for chaining of calls.
+   */
+  PartialOptimizationConfigurationBuilder addClass(ClassReference classReference);
+
+  /**
+   * Add a complete package to be optimized with partial optimization. Note that this does *not*
+   * include sub-packages.
+   *
+   * @param packageReference package to be optimized.
+   * @return instance which received the call for chaining of calls.
+   */
+  PartialOptimizationConfigurationBuilder addPackage(PackageReference packageReference);
+}
diff --git a/src/main/java/com/android/tools/r8/PartialOptimizationConfigurationProvider.java b/src/main/java/com/android/tools/r8/PartialOptimizationConfigurationProvider.java
new file mode 100644
index 0000000..f54ba47
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/PartialOptimizationConfigurationProvider.java
@@ -0,0 +1,13 @@
+// Copyright (c) 2025, 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;
+
+import com.android.tools.r8.keepanno.annotations.KeepForApi;
+
+/** API for providing partial optimization configuration to the compiler. */
+@KeepForApi
+public interface PartialOptimizationConfigurationProvider {
+
+  void getPartialOptimizationConfiguration(PartialOptimizationConfigurationBuilder builder);
+}
diff --git a/src/main/java/com/android/tools/r8/R8Command.java b/src/main/java/com/android/tools/r8/R8Command.java
index 0afb85b..b6790e9 100644
--- a/src/main/java/com/android/tools/r8/R8Command.java
+++ b/src/main/java/com/android/tools/r8/R8Command.java
@@ -161,6 +161,8 @@
         ResourceShrinkerConfiguration.DEFAULT_CONFIGURATION;
     private R8PartialCompilationConfiguration partialCompilationConfiguration =
         R8PartialCompilationConfiguration.fromSystemProperties();
+    private final List<PartialOptimizationConfigurationProvider>
+        partialOptimizationConfigurationProviders = new ArrayList<>();
 
     private final ProguardConfigurationParserOptions.Builder parserOptionsBuilder =
         ProguardConfigurationParserOptions.builder().readEnvironment();
@@ -653,6 +655,20 @@
     }
 
     /**
+     * Add {@link PartialOptimizationConfigurationProvider}s to enable R8 optimizations on a part of
+     * the program only. The {@link PartialOptimizationConfigurationProvider}s specifies exactly
+     * which classes R8 can optimize.
+     *
+     * @param providers instances implementing {@link PartialOptimizationConfigurationProvider} to
+     *     provide the partial compilation configuration.
+     */
+    public Builder addPartialOptimizationConfigurationProviders(
+        PartialOptimizationConfigurationProvider... providers) {
+      partialOptimizationConfigurationProviders.addAll(Arrays.asList(providers));
+      return this;
+    }
+
+    /**
      * Add a collection of startup profile providers that should be used for distributing the
      * program classes in DEX. The given startup profiles are also used to disallow optimizations
      * across the startup and post-startup boundary.
@@ -767,14 +783,28 @@
       if (hasDesugaredLibraryConfiguration() && getDisableDesugaring()) {
         reporter.error("Using desugared library configuration requires desugaring to be enabled");
       }
-      if (partialCompilationConfiguration.isEnabled()) {
-        validateR8Partial();
-      }
+      buildAndValidateR8Partial();
       super.validate();
     }
 
-    private void validateR8Partial() {
+    private void buildAndValidateR8Partial() {
+      if (!partialCompilationConfiguration.isEnabled()
+          && partialOptimizationConfigurationProviders.isEmpty()) {
+        return;
+      }
       Reporter reporter = getReporter();
+      if (partialCompilationConfiguration.isEnabled()
+          && !partialOptimizationConfigurationProviders.isEmpty()) {
+        reporter.error("Cannot mix experimental partial compilation with partial optimization");
+      }
+      if (!partialOptimizationConfigurationProviders.isEmpty()) {
+        assert !partialCompilationConfiguration.isEnabled();
+        R8PartialCompilationConfiguration.Builder configurationBuilder =
+            R8PartialCompilationConfiguration.builder();
+        partialOptimizationConfigurationProviders.forEach(
+            provider -> provider.getPartialOptimizationConfiguration(configurationBuilder));
+        partialCompilationConfiguration = configurationBuilder.build();
+      }
       if (!(getProgramConsumer() instanceof DexIndexedConsumer)) {
         reporter.error("Partial shrinking does not support generating class files");
       }
diff --git a/src/main/java/com/android/tools/r8/partial/R8PartialCompilationConfiguration.java b/src/main/java/com/android/tools/r8/partial/R8PartialCompilationConfiguration.java
index bb66350..486fe2f 100644
--- a/src/main/java/com/android/tools/r8/partial/R8PartialCompilationConfiguration.java
+++ b/src/main/java/com/android/tools/r8/partial/R8PartialCompilationConfiguration.java
@@ -3,6 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.partial;
 
+import com.android.tools.r8.PartialOptimizationConfigurationBuilder;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DirectMappedDexApplication;
 import com.android.tools.r8.partial.predicate.AllClassesMatcher;
@@ -13,6 +14,8 @@
 import com.android.tools.r8.partial.predicate.R8PartialPredicate;
 import com.android.tools.r8.partial.predicate.R8PartialPredicateCollection;
 import com.android.tools.r8.partial.predicate.UnnamedPackageMatcher;
+import com.android.tools.r8.references.ClassReference;
+import com.android.tools.r8.references.PackageReference;
 import com.android.tools.r8.utils.ConsumerUtils;
 import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.InternalOptions;
@@ -154,15 +157,13 @@
     return randomizeForTesting != null;
   }
 
-  public static class Builder {
+  public static class Builder implements PartialOptimizationConfigurationBuilder {
     private final R8PartialPredicateCollection includePredicates =
         new R8PartialPredicateCollection();
     private final R8PartialPredicateCollection excludePredicates =
         new R8PartialPredicateCollection();
     private Random randomizeForTesting;
 
-    private Builder() {}
-
     public R8PartialCompilationConfiguration build() {
       return new R8PartialCompilationConfiguration(
           !includePredicates.isEmpty() || randomizeForTesting != null,
@@ -265,5 +266,17 @@
       }
       return descriptorPrefixWithoutWildcards;
     }
+
+    @Override
+    public PartialOptimizationConfigurationBuilder addClass(ClassReference classReference) {
+      includePredicates.add(createPredicate("L" + classReference.getBinaryName()));
+      return this;
+    }
+
+    @Override
+    public PartialOptimizationConfigurationBuilder addPackage(PackageReference packageReference) {
+      includePredicates.add(createPredicate("L" + packageReference.getPackageBinaryName() + "/*"));
+      return this;
+    }
   }
 }
diff --git a/src/main/java/com/android/tools/r8/utils/ThrowingTriConsumer.java b/src/main/java/com/android/tools/r8/utils/ThrowingTriConsumer.java
new file mode 100644
index 0000000..818a612
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/utils/ThrowingTriConsumer.java
@@ -0,0 +1,16 @@
+// Copyright (c) 2025, 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;
+
+/**
+ * Similar to a {@link TriConsumer} but throws a single {@link Throwable}.
+ *
+ * @param <T> the type of the first argument
+ * @param <U> the type of the second argument
+ * @param <V> the type of the third argument * @param <E> the type of the {@link Throwable}
+ */
+@FunctionalInterface
+public interface ThrowingTriConsumer<T, U, V, E extends Throwable> {
+  void accept(T t, U u, V v) throws E;
+}
diff --git a/src/test/java/com/android/tools/r8/compilerapi/CompilerApiTestCollection.java b/src/test/java/com/android/tools/r8/compilerapi/CompilerApiTestCollection.java
index d32d194..14f1af6 100644
--- a/src/test/java/com/android/tools/r8/compilerapi/CompilerApiTestCollection.java
+++ b/src/test/java/com/android/tools/r8/compilerapi/CompilerApiTestCollection.java
@@ -38,6 +38,7 @@
 import com.android.tools.r8.compilerapi.testsetup.ApiTestingSetUpTest;
 import com.android.tools.r8.compilerapi.wrappers.CommandLineParserTest;
 import com.android.tools.r8.compilerapi.wrappers.EnableMissingLibraryApiModelingTest;
+import com.android.tools.r8.partial.PartialOptimizationApiTest;
 import com.android.tools.r8.partial.PartialShrinkingPreviewApiTest;
 import com.google.common.collect.ImmutableList;
 import java.nio.file.Path;
@@ -83,7 +84,8 @@
           ProtectApiSurfaceTest.ApiTest.class);
 
   private static final List<Class<? extends CompilerApiTest>> CLASSES_PENDING_BINARY_COMPATIBILITY =
-      ImmutableList.of(PartialShrinkingPreviewApiTest.ApiTest.class);
+      ImmutableList.of(
+          PartialOptimizationApiTest.ApiTest.class, PartialShrinkingPreviewApiTest.ApiTest.class);
 
   private final TemporaryFolder temp;
 
@@ -118,7 +120,9 @@
 
   @Override
   public List<Class<?>> getPendingAdditionalClassesForTests() {
-    return ImmutableList.of();
+    return ImmutableList.of(
+        PartialOptimizationApiTest.ProviderUsingClass.class,
+        PartialOptimizationApiTest.ProviderUsingPackage.class);
   }
 
   @Override
diff --git a/src/test/java/com/android/tools/r8/partial/PartialOptimizationApiTest.java b/src/test/java/com/android/tools/r8/partial/PartialOptimizationApiTest.java
new file mode 100644
index 0000000..87f8faa
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/partial/PartialOptimizationApiTest.java
@@ -0,0 +1,135 @@
+// Copyright (c) 2025, 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.partial;
+
+import static com.android.tools.r8.MarkerMatcher.markerBackend;
+import static com.android.tools.r8.MarkerMatcher.markerCompilationMode;
+import static com.android.tools.r8.MarkerMatcher.markerMinApi;
+import static com.android.tools.r8.MarkerMatcher.markerTool;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import com.android.tools.r8.CompilationMode;
+import com.android.tools.r8.DexIndexedConsumer;
+import com.android.tools.r8.PartialOptimizationConfigurationBuilder;
+import com.android.tools.r8.PartialOptimizationConfigurationProvider;
+import com.android.tools.r8.ProgramConsumer;
+import com.android.tools.r8.R8;
+import com.android.tools.r8.R8Command;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.compilerapi.CompilerApiTest;
+import com.android.tools.r8.compilerapi.CompilerApiTestRunner;
+import com.android.tools.r8.dex.Marker;
+import com.android.tools.r8.dex.Marker.Tool;
+import com.android.tools.r8.origin.Origin;
+import com.android.tools.r8.references.Reference;
+import com.android.tools.r8.utils.AndroidApiLevel;
+import com.android.tools.r8.utils.ThrowingTriConsumer;
+import com.android.tools.r8.utils.codeinspector.CodeInspector;
+import java.nio.file.Path;
+import java.util.Collection;
+import org.hamcrest.CoreMatchers;
+import org.junit.Test;
+
+public class PartialOptimizationApiTest extends CompilerApiTestRunner {
+
+  public static final int MIN_API_LEVEL = 31;
+
+  public PartialOptimizationApiTest(TestParameters parameters) {
+    super(parameters);
+  }
+
+  @Override
+  public Class<? extends CompilerApiTest> binaryTestClass() {
+    return ApiTest.class;
+  }
+
+  @Test
+  public void testR8() throws Exception {
+    ApiTest test = new ApiTest(ApiTest.PARAMETERS);
+    runTest(test::runR8);
+  }
+
+  private void runTest(
+      ThrowingTriConsumer<
+              ProgramConsumer,
+              PartialOptimizationConfigurationProvider,
+              PartialOptimizationConfigurationProvider,
+              Exception>
+          test)
+      throws Exception {
+    Path output = temp.newFolder().toPath().resolve("out.jar");
+    ProviderUsingClass providerUsingClass = new ProviderUsingClass();
+    ProviderUsingPackage providerUsingPackage = new ProviderUsingPackage();
+    test.accept(
+        new DexIndexedConsumer.ArchiveConsumer(output), providerUsingClass, providerUsingPackage);
+
+    assertTrue(providerUsingClass.called);
+    assertTrue(providerUsingPackage.called);
+    Collection<Marker> markers = new CodeInspector(output).getMarkers();
+    assertEquals(1, markers.size());
+    assertThat(
+        markers,
+        CoreMatchers.everyItem(
+            CoreMatchers.allOf(
+                markerBackend(Backend.DEX),
+                markerCompilationMode(CompilationMode.RELEASE),
+                markerMinApi(AndroidApiLevel.getAndroidApiLevel(MIN_API_LEVEL)),
+                markerTool(Tool.R8Partial))));
+  }
+
+  public static class ApiTest extends CompilerApiTest {
+
+    public ApiTest(Object parameters) {
+      super(parameters);
+    }
+
+    public void runR8(
+        ProgramConsumer programConsumer,
+        PartialOptimizationConfigurationProvider provider1,
+        PartialOptimizationConfigurationProvider provider2)
+        throws Exception {
+      R8.run(
+          R8Command.builder()
+              .addClassProgramData(getBytesForClass(getMockClass()), Origin.unknown())
+              .addProguardConfiguration(getKeepMainRules(getMockClass()), Origin.unknown())
+              .addLibraryFiles(getJava8RuntimeJar())
+              .setProgramConsumer(programConsumer)
+              .addPartialOptimizationConfigurationProviders(provider1, provider2)
+              .setMinApiLevel(MIN_API_LEVEL)
+              .build());
+    }
+
+    @Test
+    public void testR8() throws Exception {
+      runR8(
+          DexIndexedConsumer.emptyConsumer(), new ProviderUsingClass(), new ProviderUsingPackage());
+    }
+  }
+
+  public static class ProviderUsingClass implements PartialOptimizationConfigurationProvider {
+    private boolean called = false;
+
+    @Override
+    public void getPartialOptimizationConfiguration(
+        PartialOptimizationConfigurationBuilder builder) {
+      called = true;
+      builder.addClass(
+          Reference.classFromTypeName("com.android.tools.r8.compilerapi.mockdata.MockClass"));
+    }
+  }
+
+  public static class ProviderUsingPackage implements PartialOptimizationConfigurationProvider {
+    private boolean called = false;
+
+    @Override
+    public void getPartialOptimizationConfiguration(
+        PartialOptimizationConfigurationBuilder builder) {
+      called = true;
+      builder.addPackage(
+          Reference.packageFromString("com.android.tools.r8.compilerapi.mockpackage"));
+    }
+  }
+}