Allowing multiple GeneratedExtensionRegistrys

This is to allow Chrome to ship its multiple GeneratedExtensionRegistry
changes, where each feature module gets a GeneratedExtensionRegistry,
which is loaded when the module is.

Bug: 178224903
Change-Id: I80330799756a6a0a4562d7647989bf31ab7d14df
diff --git a/src/main/java/com/android/tools/r8/graph/FieldAccessInfo.java b/src/main/java/com/android/tools/r8/graph/FieldAccessInfo.java
index 19ac497..bea5b47 100644
--- a/src/main/java/com/android/tools/r8/graph/FieldAccessInfo.java
+++ b/src/main/java/com/android/tools/r8/graph/FieldAccessInfo.java
@@ -52,5 +52,7 @@
 
   boolean isWrittenOnlyInMethodSatisfying(Predicate<ProgramMethod> predicate);
 
+  boolean isReadOnlyInMethodSatisfying(Predicate<ProgramMethod> predicate);
+
   boolean isWrittenOutside(DexEncodedMethod method);
 }
diff --git a/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java b/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java
index 871ed69..4edec8a 100644
--- a/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java
+++ b/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java
@@ -243,6 +243,14 @@
   }
 
   /**
+   * Returns true if this field is only read by methods for which {@param predicate} returns true.
+   */
+  @Override
+  public boolean isReadOnlyInMethodSatisfying(Predicate<ProgramMethod> predicate) {
+    return readsWithContexts.isAccessedOnlyInMethodSatisfying(predicate);
+  }
+
+  /**
    * Returns true if this field is written by a method in the program other than {@param method}.
    */
   @Override
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedExtensionRegistryShrinker.java b/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedExtensionRegistryShrinker.java
index 3cb884a..bf81f7e 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedExtensionRegistryShrinker.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedExtensionRegistryShrinker.java
@@ -219,9 +219,9 @@
       return false;
     }
 
-    ProgramMethod uniqueReadContext = fieldAccessInfo.getUniqueReadContext();
-    return uniqueReadContext != null
-        && references.isFindLiteExtensionByNumberMethod(uniqueReadContext);
+    // Multiple GeneratedExtensionRegistries exist in Chrome; 1 per feature split.
+    return fieldAccessInfo.isReadOnlyInMethodSatisfying(
+        method -> references.isFindLiteExtensionByNumberMethod(method));
   }
 
   private void forEachDeadProtoExtensionField(Consumer<DexField> consumer) {
diff --git a/src/test/java/com/android/tools/r8/internal/proto/Proto2ShrinkingTest.java b/src/test/java/com/android/tools/r8/internal/proto/Proto2ShrinkingTest.java
index 2d2e79f..348a7f5 100644
--- a/src/test/java/com/android/tools/r8/internal/proto/Proto2ShrinkingTest.java
+++ b/src/test/java/com/android/tools/r8/internal/proto/Proto2ShrinkingTest.java
@@ -209,10 +209,17 @@
 
   private void verifyUnusedExtensionsAreRemoved(
       CodeInspector inputInspector, CodeInspector outputInspector) {
+    verifyUnusedExtensionsAreRemoved(
+        inputInspector,
+        outputInspector,
+        "com.google.protobuf.proto2_registryGeneratedExtensionRegistryLite");
+  }
+
+  private void verifyUnusedExtensionsAreRemoved(
+      CodeInspector inputInspector, CodeInspector outputInspector, String extensionRegistryName) {
     // Verify that the registry was split across multiple methods in the input.
     {
-      ClassSubject generatedExtensionRegistryLoader =
-          inputInspector.clazz("com.google.protobuf.proto2_registryGeneratedExtensionRegistryLite");
+      ClassSubject generatedExtensionRegistryLoader = inputInspector.clazz(extensionRegistryName);
       assertThat(generatedExtensionRegistryLoader, isPresent());
       assertThat(
           generatedExtensionRegistryLoader.uniqueMethodWithName("findLiteExtensionByNumber"),
@@ -228,9 +235,7 @@
     // Verify that the registry methods are still present in the output.
     // TODO(b/112437944): Should they be optimized into a single findLiteExtensionByNumber() method?
     {
-      ClassSubject generatedExtensionRegistryLoader =
-          outputInspector.clazz(
-              "com.google.protobuf.proto2_registryGeneratedExtensionRegistryLite");
+      ClassSubject generatedExtensionRegistryLoader = outputInspector.clazz(extensionRegistryName);
       assertThat(generatedExtensionRegistryLoader, isPresent());
       assertThat(
           generatedExtensionRegistryLoader.uniqueMethodWithName("findLiteExtensionByNumber"),
@@ -373,4 +378,40 @@
             inspector ->
                 assertRewrittenProtoSchemasMatch(new CodeInspector(PROGRAM_FILES), inspector));
   }
+
+  @Test
+  public void testTwoExtensionRegistrys() throws Exception {
+    CodeInspector inputInspector = new CodeInspector(PROGRAM_FILES);
+    testForR8(parameters.getBackend())
+        .addProgramFiles(PROGRAM_FILES)
+        .addKeepMainRule("proto2.TestClass")
+        .addKeepRuleFiles(PROTOBUF_LITE_PROGUARD_RULES)
+        .addKeepRules(findLiteExtensionByNumberInDuplicateCalledRule())
+        .addKeepRules(allGeneratedMessageLiteSubtypesAreInstantiatedRule())
+        // TODO(b/173340579): This rule should not be needed to allow shrinking of
+        //  PartiallyUsed$Enum.
+        .addNoHorizontalClassMergingRule(PARTIALLY_USED + "$Enum$1")
+        .allowAccessModification(allowAccessModification)
+        .allowDiagnosticMessages()
+        .allowUnusedDontWarnPatterns()
+        .allowUnusedProguardConfigurationRules()
+        .enableProtoShrinking()
+        .minification(enableMinification)
+        .setMinApi(parameters.getApiLevel())
+        .compile()
+        .assertAllInfoMessagesMatch(
+            containsString("Proguard configuration rule does not match anything"))
+        .assertAllWarningMessagesMatch(
+            anyOf(
+                equalTo("Resource 'META-INF/MANIFEST.MF' already exists."),
+                containsString("required for default or static interface methods desugaring")))
+        .inspect(
+            outputInspector -> {
+              verifyUnusedExtensionsAreRemoved(inputInspector, outputInspector);
+              verifyUnusedExtensionsAreRemoved(
+                  inputInspector,
+                  outputInspector,
+                  "com.google.protobuf.proto2_registryGeneratedExtensionRegistryLiteDuplicate");
+            });
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/internal/proto/ProtoShrinkingTestBase.java b/src/test/java/com/android/tools/r8/internal/proto/ProtoShrinkingTestBase.java
index 23d12f6..478a51a 100644
--- a/src/test/java/com/android/tools/r8/internal/proto/ProtoShrinkingTestBase.java
+++ b/src/test/java/com/android/tools/r8/internal/proto/ProtoShrinkingTestBase.java
@@ -77,6 +77,13 @@
         "}");
   }
 
+  static String findLiteExtensionByNumberInDuplicateCalledRule() {
+    return StringUtils.lines(
+        "-keep class com.google.protobuf.proto2_registryGeneratedExtensionRegistryLiteDuplicate {",
+        "  findLiteExtensionByNumber(...);",
+        "}");
+  }
+
   static String keepAllProtosRule() {
     return StringUtils.lines(
         "-if class * extends com.google.protobuf.GeneratedMessageLite",
diff --git a/third_party/proto.tar.gz.sha1 b/third_party/proto.tar.gz.sha1
index dae5c12..d18398f 100644
--- a/third_party/proto.tar.gz.sha1
+++ b/third_party/proto.tar.gz.sha1
@@ -1 +1 @@
-18151ef2484c03b5d1f8fc0084202cb9482f664d
\ No newline at end of file
+25a84e6afdc9906e7795b58ac4dda5d8526e209f
\ No newline at end of file