Fix enum unboxing bug in presence of subclasses

Bug: b/500074331
Change-Id: Ibf9343de9e98c44d7cc5f8f57559f4a07759538d

src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingLens.java

diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingLens.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingLens.java
index 33e9763..7a01842 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingLens.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxingLens.java
@@ -108,7 +108,10 @@
       AbstractValue unboxedEnumValue,
       DexType enumType) {
     DexMethod enumMethod = method.withHolder(enumType, dexItemFactory());
-    DexMethod rewrittenEnumMethod = newMethodSignatures.getRepresentativeValue(enumMethod);
+    DexMethod rewrittenEnumMethod = methodMap.apply(enumMethod);
+    if (rewrittenEnumMethod == null) {
+      rewrittenEnumMethod = newMethodSignatures.getRepresentativeValue(enumMethod);
+    }
     if (!unboxedEnumValue.isSingleNumberValue()) {
       return rewrittenEnumMethod;
     }

src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java

diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java b/src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java
index 791093e..666b867 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/MutableMethodOptimizationInfo.java
@@ -784,8 +784,15 @@
           dynamicType.getDynamicUpperBoundType(staticReturnType);
       TypeElement newDynamicUpperBoundType =
           newDynamicType.getDynamicUpperBoundType(staticReturnType);
+      // Check that the new dynamic return type is compatible with the old dynamic return type.
+      // It should generally always be stronger, but there are a few cases where this is not
+      // guaranteed, for example, due to the fact that we can reason precisely about virtual
+      // dispatch when we have a specific receiver type, but not about the behavior of _dispatch_
+      // methods inserted by enum unboxing, even when the dispatch argument is constant.
       assert newDynamicUpperBoundType.lessThanOrEqualUpToNullability(
-              previousDynamicUpperBoundType, appView)
+                  previousDynamicUpperBoundType, appView)
+              || previousDynamicUpperBoundType.lessThanOrEqualUpToNullability(
+                  newDynamicUpperBoundType, appView)
           : "upper bound type changed from "
               + previousDynamicUpperBoundType
               + " to "

src/test/java/com/android/tools/r8/enumunboxing/EnumWithSubclassOverrideTest.java

diff --git a/src/test/java/com/android/tools/r8/enumunboxing/EnumWithSubclassOverrideTest.java b/src/test/java/com/android/tools/r8/enumunboxing/EnumWithSubclassOverrideTest.java
index 4dcb4b3..9127324 100644
--- a/src/test/java/com/android/tools/r8/enumunboxing/EnumWithSubclassOverrideTest.java
+++ b/src/test/java/com/android/tools/r8/enumunboxing/EnumWithSubclassOverrideTest.java
@@ -49,8 +49,7 @@
         .addEnumUnboxingInspector(inspector -> inspector.assertUnboxed(MyEnum.class))
         .compile()
         .run(parameters.getRuntime(), Main.class)
-        // TODO(b/500074331): Should print "Sub B side effect".
-        .assertSuccessWithOutputLines("Base copy A", "Base copy B");
+        .assertSuccessWithOutputLines("Sub B side effect", "Base copy A", "Base copy B");
   }
 
   static class Main {