Mark hazzer and one-of proto fields as read from dynamicMethod() if they are written in the app

Bug: 144142878
Change-Id: I533dff931e1b47a0a6b8bd2ab3fe90e9ad41c603

src/main/java/com/android/tools/r8/graph/FieldAccessInfo.java

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 079b8bb..f573b76 100644
--- a/src/main/java/com/android/tools/r8/graph/FieldAccessInfo.java
+++ b/src/main/java/com/android/tools/r8/graph/FieldAccessInfo.java
@@ -7,6 +7,7 @@
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
+import java.util.function.Predicate;
 
 /** Provides immutable access to {@link FieldAccessInfoImpl}. */
 public interface FieldAccessInfo {
@@ -29,5 +30,7 @@
 
   boolean isWritten();
 
+  boolean isWrittenInMethodSatisfying(Predicate<DexEncodedMethod> predicate);
+
   boolean isWrittenOutside(DexEncodedMethod method);
 }

src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java

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 a263cb6..5924435 100644
--- a/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java
+++ b/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java
@@ -147,6 +147,23 @@
   }
 
   /**
+   * Returns true if this field is written by a method for which {@param predicate} returns true.
+   */
+  @Override
+  public boolean isWrittenInMethodSatisfying(Predicate<DexEncodedMethod> predicate) {
+    if (writesWithContexts != null) {
+      for (Set<DexEncodedMethod> encodedWriteContexts : writesWithContexts.values()) {
+        for (DexEncodedMethod encodedWriteContext : encodedWriteContexts) {
+          if (predicate.test(encodedWriteContext)) {
+            return true;
+          }
+        }
+      }
+    }
+    return false;
+  }
+
+  /**
    * Returns true if this field is written by a method in the program other than {@param method}.
    */
   @Override

src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java

diff --git a/src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java b/src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java
index fce7b09..92954e4 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/proto/schema/ProtoEnqueuerExtension.java
@@ -30,6 +30,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Predicate;
 
 // TODO(b/112437944): Handle cycles in the graph + add a test that fails with the current
 //  implementation. The current caching mechanism is unsafe, because we may mark a message as not
@@ -222,6 +223,19 @@
         }
 
         if (newlyLiveField != null) {
+          // Mark hazzer and one-of proto fields as read from dynamicMethod() if they are written in
+          // the app. This is needed to ensure that field writes are not removed from the app.
+          DexEncodedMethod defaultInitializer = clazz.getDefaultInitializer();
+          assert defaultInitializer != null;
+          Predicate<DexEncodedMethod> neitherDefaultConstructorNorDynamicMethod =
+              writer -> writer != defaultInitializer && writer != dynamicMethod;
+          if (enqueuer.isFieldWrittenInMethodSatisfying(
+              newlyLiveField, neitherDefaultConstructorNorDynamicMethod)) {
+            enqueuer.registerFieldRead(newlyLiveField.field, dynamicMethod);
+          }
+
+          // Unconditionally register the hazzer and one-of proto fields as written from
+          // dynamicMethod().
           if (enqueuer.registerFieldWrite(newlyLiveField.field, dynamicMethod)) {
             worklist.enqueueMarkReachableFieldAction(
                 clazz, newlyLiveField, KeepReason.reflectiveUseIn(dynamicMethod));

src/main/java/com/android/tools/r8/shaking/Enqueuer.java

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 245f3c6..a9f60ec 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -1907,6 +1907,12 @@
     return info != null && info.isRead();
   }
 
+  public boolean isFieldWrittenInMethodSatisfying(
+      DexEncodedField field, Predicate<DexEncodedMethod> predicate) {
+    FieldAccessInfoImpl info = fieldAccessInfoCollection.get(field.field);
+    return info != null && info.isWrittenInMethodSatisfying(predicate);
+  }
+
   public boolean isFieldWrittenOutsideDefaultConstructor(DexEncodedField field) {
     FieldAccessInfoImpl info = fieldAccessInfoCollection.get(field.field);
     if (info == null) {