Rewrite CONSTANT_Dynamic bootstrap method signature (part 1)

Rewrite the first argument of CONSTANT_Dynamic bootstrap methods
to java.lang.Object when desugaring with API levels below 26.

This is for first step for D8 only.

Bug: 210485236
Change-Id: I1c8669810dc72814c6833519f45168c53138a84e
diff --git a/src/main/java/com/android/tools/r8/cf/code/CfFrame.java b/src/main/java/com/android/tools/r8/cf/code/CfFrame.java
index f8b10e0..47124c0 100644
--- a/src/main/java/com/android/tools/r8/cf/code/CfFrame.java
+++ b/src/main/java/com/android/tools/r8/cf/code/CfFrame.java
@@ -532,6 +532,27 @@
   }
 
   public CfFrame map(java.util.function.Function<DexType, DexType> func) {
+    boolean mapped = false;
+    for (int var : locals.keySet()) {
+      CfFrame.FrameType originalType = locals.get(var);
+      CfFrame.FrameType mappedType = originalType.map(func);
+      mapped = originalType != mappedType;
+      if (mapped) {
+        break;
+      }
+    }
+    if (!mapped) {
+      for (FrameType frameType : stack) {
+        CfFrame.FrameType mappedType = frameType.map(func);
+        mapped = frameType != mappedType;
+        if (mapped) {
+          break;
+        }
+      }
+    }
+    if (!mapped) {
+      return this;
+    }
     Int2ReferenceSortedMap<FrameType> newLocals = new Int2ReferenceAVLTreeMap<>();
     for (int var : locals.keySet()) {
       newLocals.put(var, locals.get(var).map(func));
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/CfInstructionDesugaringEventConsumer.java b/src/main/java/com/android/tools/r8/ir/desugar/CfInstructionDesugaringEventConsumer.java
index de940f8..71db541 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/CfInstructionDesugaringEventConsumer.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/CfInstructionDesugaringEventConsumer.java
@@ -245,6 +245,7 @@
 
     private void finalizeConstantDynamicDesugaring(Consumer<ProgramMethod> needsProcessing) {
       for (ConstantDynamicClass constantDynamicClass : synthesizedConstantDynamicClasses) {
+        constantDynamicClass.rewriteBootstrapMethodSignatureIfNeeded();
         constantDynamicClass.getConstantDynamicProgramClass().forEachProgramMethod(needsProcessing);
       }
       synthesizedConstantDynamicClasses.clear();
@@ -408,6 +409,7 @@
     public void finalizeDesugaring() {
       finalizeInvokeSpecialDesugaring();
       finalizeLambdaDesugaring();
+      // TODO(b/210485236): Finalize constant dynamic desugaring by rewriting signature if needed.
     }
 
     private void finalizeInvokeSpecialDesugaring() {
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/constantdynamic/ConstantDynamicClass.java b/src/main/java/com/android/tools/r8/ir/desugar/constantdynamic/ConstantDynamicClass.java
index fb848df..1eb43e3 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/constantdynamic/ConstantDynamicClass.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/constantdynamic/ConstantDynamicClass.java
@@ -34,6 +34,7 @@
 import com.android.tools.r8.contexts.CompilationContext.MethodProcessingContext;
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.CfCode;
+import com.android.tools.r8.graph.Code;
 import com.android.tools.r8.graph.DexEncodedField;
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexField;
@@ -41,6 +42,7 @@
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexMethodHandle;
 import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.FieldAccessFlags;
 import com.android.tools.r8.graph.MethodAccessFlags;
 import com.android.tools.r8.graph.MethodResolutionResult;
@@ -55,6 +57,8 @@
 import com.android.tools.r8.ir.optimize.UtilityMethodsForCodeOptimizations.MethodSynthesizerConsumer;
 import com.android.tools.r8.ir.optimize.UtilityMethodsForCodeOptimizations.UtilityMethodForCodeOptimizations;
 import com.android.tools.r8.synthesis.SyntheticProgramClassBuilder;
+import com.android.tools.r8.utils.AndroidApiLevel;
+import com.android.tools.r8.utils.ListUtils;
 import com.android.tools.r8.utils.collections.ImmutableDeque;
 import com.android.tools.r8.utils.collections.ImmutableInt2ReferenceSortedMap;
 import com.google.common.collect.ImmutableList;
@@ -80,6 +84,8 @@
   public final DexField constantValueField;
   private final DexMethod getConstMethod;
   private final Behaviour behaviour;
+  private DexEncodedMethod bootstrapMethodImpl;
+  private DexMethod finalBootstrapMethodReference;
 
   // Considered final but is set after due to circularity in allocation.
   private DexProgramClass clazz = null;
@@ -115,12 +121,30 @@
             .resolveMethod(bootstrapMethodReference, bootstrapMethodHandle.isInterface);
     if (resolution.isSingleResolution()
         && resolution.asSingleResolution().getResolvedMethod().isStatic()) {
-      // Ensure that the bootstrap method is accessible from the generated class.
       SingleResolutionResult result = resolution.asSingleResolution();
-      MethodAccessFlags flags = result.getResolvedMethod().getAccessFlags();
-      flags.unsetPrivate();
-      flags.setPublic();
+      bootstrapMethodImpl = result.getResolvedMethod();
+      if (shouldRewriteBootstrapMethodSignature()) {
+        // The bootstrap method will have its signature modified to have type Object as its first
+        // argument.
+        this.finalBootstrapMethodReference =
+            factory.createMethod(
+                result.getResolvedHolder().getType(),
+                factory.createProto(
+                    bootstrapMethodReference.getReturnType(),
+                    factory.objectType,
+                    factory.stringType,
+                    factory.classType),
+                bootstrapMethodImpl.getName());
+      } else {
+        this.finalBootstrapMethodReference = bootstrapMethodReference;
+        // Ensure that the bootstrap method is accessible from the generated class.
+        MethodAccessFlags flags = bootstrapMethodImpl.getAccessFlags();
+        flags.unsetPrivate();
+        flags.setPublic();
+      }
+
       behaviour = CACHE_CONSTANT;
+
       synthesizeConstantDynamicClass(builder);
     } else {
       // Unconditionally throw as the RI.
@@ -128,6 +152,12 @@
     }
   }
 
+  private boolean shouldRewriteBootstrapMethodSignature() {
+    // TODO(b/210485236): Check for references to the bootstrap method outside of dynamic constant.
+    return !appView.enableWholeProgramOptimizations()
+        && appView.options().getMinApiLevel().isLessThan(AndroidApiLevel.O);
+  }
+
   public Collection<CfInstruction> desugarConstDynamicInstruction(
       CfConstDynamic invoke,
       FreshLocalProvider freshLocalProvider,
@@ -214,13 +244,11 @@
 
   private void invokeBootstrapMethod(ImmutableList.Builder<CfInstruction> instructions) {
     assert reference.getBootstrapMethod().type.isInvokeStatic();
-    DexMethodHandle bootstrapMethodHandle = reference.getBootstrapMethod();
-    DexMethod bootstrapMethodReference = bootstrapMethodHandle.asMethod();
     // TODO(b/178172809): Use MethodHandle.invokeWithArguments if supported.
     instructions.add(new CfConstNull());
     instructions.add(new CfConstString(reference.getName()));
     instructions.add(new CfConstClass(reference.getType()));
-    instructions.add(new CfInvoke(INVOKESTATIC, bootstrapMethodReference, false));
+    instructions.add(new CfInvoke(INVOKESTATIC, finalBootstrapMethodReference, false));
     instructions.add(new CfCheckCast(reference.getType()));
   }
 
@@ -317,4 +345,76 @@
     assert clazz != null;
     this.clazz = clazz;
   }
+
+  public void rewriteBootstrapMethodSignatureIfNeeded() {
+    if (!shouldRewriteBootstrapMethodSignature() || behaviour != CACHE_CONSTANT) {
+      return;
+    }
+    DexProgramClass bootstrapMethodHolder =
+        appView.definitionFor(bootstrapMethodImpl.getHolderType()).asProgramClass();
+    DexEncodedMethod replacement =
+        bootstrapMethodHolder
+            .getMethodCollection()
+            .replaceDirectMethod(
+                bootstrapMethodImpl.getReference(),
+                encodedMethod -> {
+                  MethodAccessFlags newAccessFlags = encodedMethod.accessFlags.copy();
+                  // Ensure that the bootstrap method is accessible from the generated class.
+                  newAccessFlags.unsetPrivate();
+                  newAccessFlags.setPublic();
+                  DexEncodedMethod newMethod =
+                      DexEncodedMethod.syntheticBuilder()
+                          .setMethod(finalBootstrapMethodReference)
+                          .setAccessFlags(newAccessFlags)
+                          .setGenericSignature(encodedMethod.getGenericSignature())
+                          .setAnnotations(encodedMethod.annotations())
+                          .setParameterAnnotations(encodedMethod.parameterAnnotationsList)
+                          .setCode(adaptCode(encodedMethod))
+                          .setApiLevelForDefinition(encodedMethod.getApiLevelForDefinition())
+                          .setApiLevelForCode(encodedMethod.getApiLevelForCode())
+                          .build();
+                  newMethod.copyMetadata(appView, encodedMethod);
+                  return newMethod;
+                });
+    if (replacement != null) {
+      // Since we've copied the code object from an existing method, the code should already be
+      // processed, and thus we don't need to schedule it for processing in D8.
+      assert !appView.options().isGeneratingClassFiles() || replacement.getCode().isCfCode();
+      assert !appView.options().isGeneratingDex() || replacement.getCode().isDexCode();
+    }
+    // The method might already have been moved by another dynamic constant targeting it.
+    // If so, it must be defined on the holder.
+    ProgramMethod modified =
+        bootstrapMethodHolder.lookupProgramMethod(finalBootstrapMethodReference);
+    assert modified != null;
+    assert modified.getDefinition().isPublicMethod();
+  }
+
+  private DexType mapLookupTypeToObject(DexType type) {
+    return type == appView.dexItemFactory().lookupType ? appView.dexItemFactory().objectType : type;
+  }
+
+  private Code adaptCode(DexEncodedMethod method) {
+    assert behaviour == CACHE_CONSTANT;
+    if (method.getCode().isDexCode()) {
+      return method.getCode();
+    }
+    CfCode code = method.getCode().asCfCode();
+    List<CfInstruction> newInstructions =
+        ListUtils.mapOrElse(
+            code.getInstructions(),
+            instruction ->
+                instruction.isFrame()
+                    ? instruction.asFrame().map(this::mapLookupTypeToObject)
+                    : instruction);
+    return code.getInstructions() != newInstructions
+        ? new CfCode(
+            method.getHolderType(),
+            code.getMaxStack(),
+            code.getMaxLocals(),
+            newInstructions,
+            code.getTryCatchRanges(),
+            code.getLocalVariables())
+        : code;
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/desugar/constantdynamic/ConstantDynamicGetDeclaredMethodsTest.java b/src/test/java/com/android/tools/r8/desugar/constantdynamic/ConstantDynamicGetDeclaredMethodsTest.java
index de9685e..4810b1d 100644
--- a/src/test/java/com/android/tools/r8/desugar/constantdynamic/ConstantDynamicGetDeclaredMethodsTest.java
+++ b/src/test/java/com/android/tools/r8/desugar/constantdynamic/ConstantDynamicGetDeclaredMethodsTest.java
@@ -34,7 +34,7 @@
         getTestParameters().withAllRuntimes().withAllApiLevelsAlsoForCf().build());
   }
 
-  private static final String EXPECTED_OUTPUT =
+  private static final String EXPECTED_OUTPUT_WITH_METHOD_HANDLES =
       StringUtils.lines(
           "Hello, world!",
           "myConstant",
@@ -42,6 +42,14 @@
           "java.lang.invoke.MethodHandles$Lookup",
           "java.lang.String",
           "java.lang.Class");
+  private static final String EXPECTED_OUTPUT_WITHOUT_METHOD_HANDLES =
+      StringUtils.lines(
+          "Hello, world!",
+          "myConstant",
+          "3",
+          "java.lang.Object",
+          "java.lang.String",
+          "java.lang.Class");
   private static final String EXPECTED_OUTPUT_R8 =
       StringUtils.lines("Hello, world!", "No myConstant method");
 
@@ -56,7 +64,7 @@
     testForJvm()
         .addProgramClassFileData(getTransformedClasses())
         .run(parameters.getRuntime(), MAIN_CLASS)
-        .assertSuccessWithOutput(EXPECTED_OUTPUT);
+        .assertSuccessWithOutput(EXPECTED_OUTPUT_WITH_METHOD_HANDLES);
   }
 
   @Test
@@ -64,24 +72,22 @@
     testForDesugaring(parameters)
         .addProgramClassFileData(getTransformedClasses())
         .run(parameters.getRuntime(), MAIN_CLASS)
-        // TODO(b/210485236): This should never fail.
         .applyIf(
             // When not desugaring the CF code requires JDK 11.
             DesugarTestConfiguration::isNotDesugared,
             r -> {
               if (parameters.isCfRuntime()
                   && parameters.getRuntime().asCf().isNewerThanOrEqual(CfVm.JDK11)) {
-                r.assertSuccessWithOutput(EXPECTED_OUTPUT);
+                r.assertSuccessWithOutput(EXPECTED_OUTPUT_WITH_METHOD_HANDLES);
               } else {
                 r.assertFailureWithErrorThatThrows(UnsupportedClassVersionError.class);
               }
             },
             c ->
                 DesugarTestConfiguration.isDesugared(c)
-                    && parameters.isDexRuntime()
-                    && parameters.asDexRuntime().getVersion().isOlderThan(Version.V8_1_0),
-            r -> r.assertFailureWithErrorThatThrows(ClassNotFoundException.class),
-            r -> r.assertSuccessWithOutput(EXPECTED_OUTPUT));
+                    && parameters.getApiLevel().isLessThan(AndroidApiLevel.O),
+            r -> r.assertSuccessWithOutput(EXPECTED_OUTPUT_WITHOUT_METHOD_HANDLES),
+            r -> r.assertSuccessWithOutput(EXPECTED_OUTPUT_WITH_METHOD_HANDLES));
   }
 
   @Test
@@ -110,11 +116,10 @@
             parameters.getApiLevel().isLessThan(AndroidApiLevel.O),
             b -> b.addDontWarn(MethodHandles.Lookup.class))
         .run(parameters.getRuntime(), MAIN_CLASS)
-        // TODO(b/210485236): This should never fail.
         .applyIf(
             parameters.getDexRuntimeVersion().isOlderThan(Version.V8_1_0),
             b -> b.assertFailureWithErrorThatThrows(ClassNotFoundException.class),
-            b -> b.assertSuccessWithOutput(EXPECTED_OUTPUT));
+            b -> b.assertSuccessWithOutput(EXPECTED_OUTPUT_WITH_METHOD_HANDLES));
   }
 
   private byte[] getTransformedClasses() throws IOException {
diff --git a/src/test/java/com/android/tools/r8/desugar/constantdynamic/JacocoConstantDynamicGetDeclaredMethods.java b/src/test/java/com/android/tools/r8/desugar/constantdynamic/JacocoConstantDynamicGetDeclaredMethods.java
index 79fc834..8524f5d 100644
--- a/src/test/java/com/android/tools/r8/desugar/constantdynamic/JacocoConstantDynamicGetDeclaredMethods.java
+++ b/src/test/java/com/android/tools/r8/desugar/constantdynamic/JacocoConstantDynamicGetDeclaredMethods.java
@@ -49,13 +49,20 @@
   public static JacocoClasses testClasses;
 
   private static final String MAIN_CLASS = TestRunner.class.getTypeName();
-  private static final String EXPECTED_OUTPUT =
+  private static final String EXPECTED_OUTPUT_WITH_METHOD_HANDLES =
       StringUtils.lines(
           jacocoBootstrapMethodName,
           "3",
           "java.lang.invoke.MethodHandles$Lookup",
           "java.lang.String",
           "java.lang.Class");
+  private static final String EXPECTED_OUTPUT_WITHOUT_METHOD_HANDLES =
+      StringUtils.lines(
+          jacocoBootstrapMethodName,
+          "3",
+          "java.lang.Object",
+          "java.lang.String",
+          "java.lang.Class");
 
   @BeforeClass
   public static void setUpInput() throws IOException {
@@ -85,7 +92,7 @@
         .addProgramFiles(testClasses.getOriginal())
         .enableJaCoCoAgent(ToolHelper.JACOCO_AGENT, agentOutputOnTheFly)
         .run(parameters.getRuntime(), MAIN_CLASS)
-        .assertSuccessWithOutput(EXPECTED_OUTPUT);
+        .assertSuccessWithOutput(EXPECTED_OUTPUT_WITH_METHOD_HANDLES);
     checkJacocoReport(agentOutputOnTheFly);
 
     // Run the instrumented code.
@@ -94,7 +101,7 @@
         .addProgramFiles(testClasses.getInstrumented())
         .configureJaCoCoAgentForOfflineInstrumentedCode(ToolHelper.JACOCO_AGENT, agentOutputOffline)
         .run(parameters.getRuntime(), MAIN_CLASS)
-        .assertSuccessWithOutput(EXPECTED_OUTPUT);
+        .assertSuccessWithOutput(EXPECTED_OUTPUT_WITH_METHOD_HANDLES);
     checkJacocoReport(agentOutputOffline);
   }
 
@@ -108,11 +115,10 @@
         .setMinApi(parameters.getApiLevel())
         .compile()
         .runWithJaCoCo(agentOutput, parameters.getRuntime(), MAIN_CLASS)
-        // TODO(b/210485236): This should never fail.
         .applyIf(
-            parameters.getDexRuntimeVersion().isOlderThan(Version.V8_1_0),
-            b -> b.assertFailureWithErrorThatThrows(ClassNotFoundException.class),
-            b -> b.assertSuccessWithOutput(EXPECTED_OUTPUT));
+            parameters.getApiLevel().isLessThan(AndroidApiLevel.O),
+            b -> b.assertSuccessWithOutput(EXPECTED_OUTPUT_WITHOUT_METHOD_HANDLES),
+            b -> b.assertSuccessWithOutput(EXPECTED_OUTPUT_WITH_METHOD_HANDLES));
     checkJacocoReport(agentOutput);
   }
 
@@ -160,11 +166,10 @@
             "javax.management.**")
         .compile()
         .runWithJaCoCo(agentOutput, parameters.getRuntime(), MAIN_CLASS)
-        // TODO(b/210485236): This should never fail.
         .applyIf(
             parameters.getDexRuntimeVersion().isOlderThan(Version.V8_1_0),
             b -> b.assertFailureWithErrorThatThrows(ClassNotFoundException.class),
-            b -> b.assertSuccessWithOutput(EXPECTED_OUTPUT));
+            b -> b.assertSuccessWithOutput(EXPECTED_OUTPUT_WITH_METHOD_HANDLES));
     checkJacocoReport(agentOutput);
   }