Enable devirtualization in ProgramRewritingTest

Bug: b/391572031
Fixes: b/140233505
Change-Id: I225e9543aebb449ec5f16a497aacccff5a8d7d4f
diff --git a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/ProgramRewritingTest.java b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/ProgramRewritingTest.java
index 2faea77..ca43f85 100644
--- a/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/ProgramRewritingTest.java
+++ b/src/test/java/com/android/tools/r8/desugar/desugaredlibrary/ProgramRewritingTest.java
@@ -94,14 +94,6 @@
                 parameters, libraryDesugaringSpecification, compilationSpecification)
             .addProgramFiles(Paths.get(ToolHelper.EXAMPLES_JAVA9_BUILD_DIR + "stream.jar"))
             .addKeepMainRule(TEST_CLASS)
-            .applyIf(
-                compilationSpecification.isProgramShrink(),
-                b ->
-                    b.addOptionsModification(
-                        options -> {
-                          // TODO(b/140233505): Allow devirtualization once fixed.
-                          options.enableDevirtualization = false;
-                        }))
             .compile()
             .inspect(this::checkRewrittenInvokes)
             .inspectKeepRules(
@@ -146,11 +138,25 @@
         classSubject
             .uniqueMethodWithOriginalName("main")
             .streamInstructions()
-            .filter(instr -> instr.isInvokeInterface() || instr.isInvokeStatic())
+            .filter(
+                instr ->
+                    instr.isInvokeInterface() || instr.isInvokeStatic() || instr.isInvokeVirtual())
+            .filter(
+                instr -> {
+                  String holder = instr.getMethod().getHolderType().getTypeName();
+                  return !holder.equals("java.lang.Class")
+                      && !holder.equals("java.lang.Object")
+                      && !holder.equals("java.io.PrintStream");
+                })
             .collect(Collectors.toList());
+    assertEquals(22, invokes.size());
     assertInvokeStaticMatching(invokes, 0, "Set$-EL;spliterator");
     assertInvokeStaticMatching(invokes, 1, "Collection$-EL;stream");
-    assertInvokeInterfaceMatching(invokes, 2, "Set;iterator");
+    if (compilationSpecification.isProgramShrink()) {
+      assertInvokeVirtualMatching(invokes, 2, "HashSet;iterator");
+    } else {
+      assertInvokeInterfaceMatching(invokes, 2, "Set;iterator");
+    }
     assertInvokeStaticMatching(invokes, 3, "Collection$-EL;stream");
     assertInvokeStaticMatching(invokes, 4, "Set$-EL;spliterator");
     assertInvokeInterfaceMatching(invokes, 8, "Iterator;remove");
@@ -165,7 +171,6 @@
     assertInvokeStaticMatching(invokes, 19, "Comparator$-CC;comparingInt");
     assertInvokeStaticMatching(invokes, 20, "List$-EL;sort");
     assertInvokeStaticMatching(invokes, 21, "Collection$-EL;stream");
-    assertEquals(22, invokes.size());
   }
 
   private void assertInvokeInterfaceMatching(List<InstructionSubject> invokes, int i, String s) {
@@ -178,6 +183,11 @@
     assertTrue(invokes.get(i).toString().contains(s));
   }
 
+  private void assertInvokeVirtualMatching(List<InstructionSubject> invokes, int i, String s) {
+    assertTrue(invokes.get(i).isInvokeVirtual());
+    assertTrue(invokes.get(i).toString().contains(s));
+  }
+
   private void assertGeneratedKeepRulesAreCorrect(String keepRules) {
 
     String prefix = libraryDesugaringSpecification.functionPrefix(parameters);