Cleanup class initializer defaults optimization

Change-Id: Icda7e56061ebc3e089eb0e6d7190525461039ce9
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/fieldaccess/FieldAssignmentTracker.java b/src/main/java/com/android/tools/r8/ir/analysis/fieldaccess/FieldAssignmentTracker.java
index 62460e5..0bef6b1 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/fieldaccess/FieldAssignmentTracker.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/fieldaccess/FieldAssignmentTracker.java
@@ -200,7 +200,7 @@
       AbstractValue abstractValue =
           abstractInstanceFieldValuesForClass.getOrDefault(field, UnknownValue.getInstance());
       if (abstractValue.isBottom()) {
-        feedback.modifyAppInfoWithLiveness().removeInstantiatedType(clazz);
+        feedback.modifyAppInfoWithLiveness(modifier -> modifier.removeInstantiatedType(clazz));
         break;
       }
       if (abstractValue.isUnknown()) {
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
index a4d036c..084b1f6 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
@@ -1344,7 +1344,7 @@
 
     timing.begin("Optimize class initializers");
     ClassInitializerDefaultsResult classInitializerDefaultsResult =
-        classInitializerDefaultsOptimization.optimize(method, code);
+        classInitializerDefaultsOptimization.optimize(method, code, feedback);
     timing.end();
 
     if (Log.ENABLED) {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/ClassInitializerDefaultsOptimization.java b/src/main/java/com/android/tools/r8/ir/optimize/ClassInitializerDefaultsOptimization.java
index e35087a..1e02c2b 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/ClassInitializerDefaultsOptimization.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/ClassInitializerDefaultsOptimization.java
@@ -45,6 +45,7 @@
 import com.android.tools.r8.ir.code.StaticPut;
 import com.android.tools.r8.ir.code.Value;
 import com.android.tools.r8.ir.conversion.IRConverter;
+import com.android.tools.r8.ir.optimize.info.OptimizationFeedback;
 import com.android.tools.r8.naming.dexitembasedstring.ClassNameComputationInfo;
 import com.android.tools.r8.naming.dexitembasedstring.ClassNameComputationInfo.ClassNameMapping;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
@@ -91,32 +92,20 @@
     }
   }
 
-  private class WaveDoneAction implements Action {
+  private static class WaveDoneAction implements Action {
 
     private final Map<DexEncodedField, DexValue> fieldsWithStaticValues = new IdentityHashMap<>();
-    private final Set<DexField> noLongerWrittenFields = Sets.newIdentityHashSet();
 
-    public WaveDoneAction(
-        Map<DexEncodedField, DexValue> fieldsWithStaticValues,
-        Set<DexField> noLongerWrittenFields) {
+    WaveDoneAction(Map<DexEncodedField, DexValue> fieldsWithStaticValues) {
       this.fieldsWithStaticValues.putAll(fieldsWithStaticValues);
-      this.noLongerWrittenFields.addAll(noLongerWrittenFields);
     }
 
-    public synchronized void join(
-        Map<DexEncodedField, DexValue> fieldsWithStaticValues,
-        Set<DexField> noLongerWrittenFields) {
+    public synchronized void join(Map<DexEncodedField, DexValue> fieldsWithStaticValues) {
       this.fieldsWithStaticValues.putAll(fieldsWithStaticValues);
-      this.noLongerWrittenFields.addAll(noLongerWrittenFields);
     }
 
     @Override
     public void execute() {
-      // Update AppInfo.
-      AppView<AppInfoWithLiveness> appViewWithLiveness = appView.withLiveness();
-      appViewWithLiveness.setAppInfo(
-          appViewWithLiveness.appInfo().withoutStaticFieldsWrites(noLongerWrittenFields));
-
       // Update static field values of classes.
       fieldsWithStaticValues.forEach(DexEncodedField::setStaticValue);
     }
@@ -134,7 +123,8 @@
     this.dexItemFactory = appView.dexItemFactory();
   }
 
-  public ClassInitializerDefaultsResult optimize(DexEncodedMethod method, IRCode code) {
+  public ClassInitializerDefaultsResult optimize(
+      DexEncodedMethod method, IRCode code, OptimizationFeedback feedback) {
     if (appView.options().debug || method.getOptimizationInfo().isReachabilitySensitive()) {
       return ClassInitializerDefaultsResult.empty();
     }
@@ -269,17 +259,21 @@
           }
         }
 
-        // Finally, remove these fields from the set of assigned static fields.
+        // Remove these fields from the set of assigned static fields.
+        feedback.modifyAppInfoWithLiveness(
+            modifier -> candidates.forEach(modifier::removeWrittenField));
+
+        // Update the static value of the fields when the wave ends.
         synchronized (this) {
           if (waveDoneAction == null) {
-            waveDoneAction = new WaveDoneAction(fieldsWithStaticValues, candidates);
+            waveDoneAction = new WaveDoneAction(fieldsWithStaticValues);
             converter.addWaveDoneAction(
                 () -> {
                   waveDoneAction.execute();
                   waveDoneAction = null;
                 });
           } else {
-            waveDoneAction.join(fieldsWithStaticValues, candidates);
+            waveDoneAction.join(fieldsWithStaticValues);
           }
         }
       } else {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedback.java b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedback.java
index b65a579..46cff7b 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedback.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedback.java
@@ -9,9 +9,11 @@
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.ir.conversion.FieldOptimizationFeedback;
 import com.android.tools.r8.ir.conversion.MethodOptimizationFeedback;
+import com.android.tools.r8.shaking.AppInfoWithLivenessModifier;
 import com.android.tools.r8.utils.ThreadUtils;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
+import java.util.function.Consumer;
 
 public abstract class OptimizationFeedback
     implements FieldOptimizationFeedback, MethodOptimizationFeedback {
@@ -34,4 +36,8 @@
         },
         executorService);
   }
+
+  public void modifyAppInfoWithLiveness(Consumer<AppInfoWithLivenessModifier> consumer) {
+    // Intentionally empty.
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackDelayed.java b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackDelayed.java
index 61f5a16..c241c0a 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackDelayed.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/OptimizationFeedbackDelayed.java
@@ -24,6 +24,7 @@
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
+import java.util.function.Consumer;
 
 public class OptimizationFeedbackDelayed extends OptimizationFeedback {
 
@@ -58,10 +59,6 @@
     return info;
   }
 
-  public AppInfoWithLivenessModifier modifyAppInfoWithLiveness() {
-    return appInfoWithLivenessModifier;
-  }
-
   @Override
   public void fixupOptimizationInfos(
       AppView<?> appView, ExecutorService executorService, OptimizationInfoFixer fixer)
@@ -70,6 +67,11 @@
     super.fixupOptimizationInfos(appView, executorService, fixer);
   }
 
+  @Override
+  public void modifyAppInfoWithLiveness(Consumer<AppInfoWithLivenessModifier> consumer) {
+    consumer.accept(appInfoWithLivenessModifier);
+  }
+
   public void refineAppInfoWithLiveness(AppInfoWithLiveness appInfo) {
     appInfoWithLivenessModifier.modify(appInfo);
   }
diff --git a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
index 0b9cfdf..509a832 100644
--- a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
+++ b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
@@ -719,6 +719,10 @@
     return fieldAccessInfoCollection;
   }
 
+  FieldAccessInfoCollectionImpl getMutableFieldAccessInfoCollection() {
+    return fieldAccessInfoCollection;
+  }
+
   /** This method provides immutable access to `objectAllocationInfoCollection`. */
   public ObjectAllocationInfoCollection getObjectAllocationInfoCollection() {
     return objectAllocationInfoCollection;
diff --git a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLivenessModifier.java b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLivenessModifier.java
index bd982f1..75c27fd 100644
--- a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLivenessModifier.java
+++ b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLivenessModifier.java
@@ -4,7 +4,10 @@
 
 package com.android.tools.r8.shaking;
 
+import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.FieldAccessInfoCollectionImpl;
+import com.android.tools.r8.graph.FieldAccessInfoImpl;
 import com.android.tools.r8.graph.ObjectAllocationInfoCollectionImpl;
 import com.google.common.collect.Sets;
 import java.util.Set;
@@ -12,7 +15,8 @@
 /** Used to mutate AppInfoWithLiveness between waves. */
 public class AppInfoWithLivenessModifier {
 
-  private final Set<DexProgramClass> noLongerInstantiatedClasses = Sets.newIdentityHashSet();
+  private final Set<DexProgramClass> noLongerInstantiatedClasses = Sets.newConcurrentHashSet();
+  private final Set<DexField> noLongerWrittenFields = Sets.newConcurrentHashSet();
 
   AppInfoWithLivenessModifier() {}
 
@@ -24,10 +28,27 @@
     noLongerInstantiatedClasses.add(clazz);
   }
 
+  public void removeWrittenField(DexField field) {
+    noLongerWrittenFields.add(field);
+  }
+
   public void modify(AppInfoWithLiveness appInfo) {
+    // Instantiated classes.
     ObjectAllocationInfoCollectionImpl objectAllocationInfoCollection =
         appInfo.getMutableObjectAllocationInfoCollection();
     noLongerInstantiatedClasses.forEach(objectAllocationInfoCollection::markNoLongerInstantiated);
+
+    // Written fields.
+    FieldAccessInfoCollectionImpl fieldAccessInfoCollection =
+        appInfo.getMutableFieldAccessInfoCollection();
+    noLongerWrittenFields.forEach(
+        field -> {
+          FieldAccessInfoImpl fieldAccessInfo = fieldAccessInfoCollection.get(field);
+          if (fieldAccessInfo != null) {
+            fieldAccessInfo.clearWrites();
+          }
+        });
+
     clear();
   }