Remove Proto assertion and fix determinism

Bug: 179146128
Change-Id: I9c3a86bd34913589669062b59c05cdc7d0747cfe
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/proto/ProtoEnumSwitchMapRemover.java b/src/main/java/com/android/tools/r8/ir/analysis/proto/ProtoEnumSwitchMapRemover.java
index 00e824f..5c3692c 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/proto/ProtoEnumSwitchMapRemover.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/proto/ProtoEnumSwitchMapRemover.java
@@ -4,7 +4,6 @@
 
 package com.android.tools.r8.ir.analysis.proto;
 
-import com.android.tools.r8.errors.CompilationError;
 import com.android.tools.r8.graph.DexEncodedField;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexType;
@@ -12,6 +11,7 @@
 import com.android.tools.r8.ir.analysis.fieldvalueanalysis.StaticFieldValues.EnumStaticFieldValues;
 import com.android.tools.r8.ir.analysis.value.ObjectState;
 import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
+import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -20,7 +20,8 @@
 
   private final ProtoReferences references;
 
-  private final Map<DexType, EnumStaticFieldValues> staticFieldValuesMap =
+  private final Map<DexType, EnumStaticFieldValues> staticFieldValuesMap = new IdentityHashMap<>();
+  private final Map<DexType, EnumStaticFieldValues> staticFieldValuesMapDelayed =
       new ConcurrentHashMap<>();
 
   public ProtoEnumSwitchMapRemover(ProtoReferences references) {
@@ -34,10 +35,15 @@
     assert clazz.isEnum();
     EnumStaticFieldValues enumStaticFieldValues = staticFieldValues.asEnumStaticFieldValues();
     if (isProtoEnum(clazz)) {
-      staticFieldValuesMap.put(clazz.type, enumStaticFieldValues);
+      staticFieldValuesMapDelayed.put(clazz.type, enumStaticFieldValues);
     }
   }
 
+  public void updateVisibleStaticFieldValues() {
+    staticFieldValuesMap.putAll(staticFieldValuesMapDelayed);
+    staticFieldValuesMapDelayed.clear();
+  }
+
   private boolean isProtoEnum(DexProgramClass clazz) {
     assert clazz.isEnum();
     if (clazz.type == references.methodToInvokeType) {
@@ -53,9 +59,8 @@
     }
     EnumStaticFieldValues enumStaticFieldValues = staticFieldValuesMap.get(enumClass.type);
     if (enumStaticFieldValues == null) {
-      if (enumClass.type == references.methodToInvokeType) {
-        throw new CompilationError("Proto optimizations: missing information for MethodToInvoke.");
-      }
+      // If the switch map is found in a wave previous to the wave containing the enum clinit,
+      // then bail out. This can happen but is extremely uncommon.
       return null;
     }
     ObjectState state =
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 b1e68ace..d7d0344 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
@@ -872,6 +872,9 @@
     if (options.enableFieldAssignmentTracker) {
       fieldAccessAnalysis.fieldAssignmentTracker().waveDone(wave, delayedOptimizationFeedback);
     }
+    if (appView.options().protoShrinking().enableRemoveProtoEnumSwitchMap()) {
+      appView.protoShrinker().protoEnumSwitchMapRemover.updateVisibleStaticFieldValues();
+    }
     assert delayedOptimizationFeedback.noUpdatesLeft();
     onWaveDoneActions.forEach(com.android.tools.r8.utils.Action::execute);
     onWaveDoneActions = null;