Rewrite holder and arguments for DexCallSites in Writer

Bug: 176885013
Bug: 176486859
Change-Id: I1e34dabf3c142ebaba04cc3209c51e0e8aba363f
diff --git a/src/main/java/com/android/tools/r8/dex/VirtualFile.java b/src/main/java/com/android/tools/r8/dex/VirtualFile.java
index bce34ff..983e4eb 100644
--- a/src/main/java/com/android/tools/r8/dex/VirtualFile.java
+++ b/src/main/java/com/android/tools/r8/dex/VirtualFile.java
@@ -221,6 +221,7 @@
         graphLens,
         namingLens,
         initClassLens,
+        transaction.rewriter,
         indexedItems.classes,
         indexedItems.protos,
         indexedItems.types,
diff --git a/src/main/java/com/android/tools/r8/graph/ObjectToOffsetMapping.java b/src/main/java/com/android/tools/r8/graph/ObjectToOffsetMapping.java
index f315247..196bc31 100644
--- a/src/main/java/com/android/tools/r8/graph/ObjectToOffsetMapping.java
+++ b/src/main/java/com/android/tools/r8/graph/ObjectToOffsetMapping.java
@@ -53,6 +53,7 @@
       GraphLens graphLens,
       NamingLens namingLens,
       InitClassLens initClassLens,
+      LensCodeRewriterUtils lensCodeRewriter,
       Collection<DexProgramClass> classes,
       Collection<DexProto> protos,
       Collection<DexType> types,
@@ -76,7 +77,7 @@
     this.graphLens = graphLens;
     this.namingLens = namingLens;
     this.initClassLens = initClassLens;
-    this.lensCodeRewriter = new LensCodeRewriterUtils(appView);
+    this.lensCodeRewriter = lensCodeRewriter;
     timing.begin("Sort strings");
     this.strings = createSortedMap(strings, DexString::compareTo, this::setFirstJumboString);
     CompareToVisitor visitor =
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriterUtils.java b/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriterUtils.java
index ae5e3f9..39f7dd1 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriterUtils.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriterUtils.java
@@ -39,6 +39,10 @@
 
   private final Map<DexProto, DexProto> protoFixupCache = new ConcurrentHashMap<>();
 
+  // Map from original call sites to rewritten call sites to lookup call-sites when writing.
+  // TODO(b/176885013): This is redundant if we canonicalize call sites.
+  private final Map<DexCallSite, DexCallSite> rewrittenCallSiteCache = new ConcurrentHashMap<>();
+
   public LensCodeRewriterUtils(AppView<?> appView) {
     this.appView = appView;
     this.definitions = appView;
@@ -56,24 +60,30 @@
   }
 
   public DexCallSite rewriteCallSite(DexCallSite callSite, ProgramMethod context) {
-    DexItemFactory dexItemFactory = definitions.dexItemFactory();
-    DexProto newMethodProto = rewriteProto(callSite.methodProto);
-    DexMethodHandle newBootstrapMethod =
-        rewriteDexMethodHandle(
-            callSite.bootstrapMethod, NOT_ARGUMENT_TO_LAMBDA_METAFACTORY, context);
-    boolean isLambdaMetaFactory =
-        dexItemFactory.isLambdaMetafactoryMethod(callSite.bootstrapMethod.asMethod());
-    MethodHandleUse methodHandleUse =
-        isLambdaMetaFactory ? ARGUMENT_TO_LAMBDA_METAFACTORY : NOT_ARGUMENT_TO_LAMBDA_METAFACTORY;
-    List<DexValue> newArgs =
-        rewriteBootstrapArguments(callSite.bootstrapArgs, methodHandleUse, context);
-    if (!newMethodProto.equals(callSite.methodProto)
-        || newBootstrapMethod != callSite.bootstrapMethod
-        || !newArgs.equals(callSite.bootstrapArgs)) {
-      return dexItemFactory.createCallSite(
-          callSite.methodName, newMethodProto, newBootstrapMethod, newArgs);
-    }
-    return callSite;
+    return rewrittenCallSiteCache.computeIfAbsent(
+        callSite,
+        ignored -> {
+          DexItemFactory dexItemFactory = definitions.dexItemFactory();
+          DexProto newMethodProto = rewriteProto(callSite.methodProto);
+          DexMethodHandle newBootstrapMethod =
+              rewriteDexMethodHandle(
+                  callSite.bootstrapMethod, NOT_ARGUMENT_TO_LAMBDA_METAFACTORY, context);
+          boolean isLambdaMetaFactory =
+              dexItemFactory.isLambdaMetafactoryMethod(callSite.bootstrapMethod.asMethod());
+          MethodHandleUse methodHandleUse =
+              isLambdaMetaFactory
+                  ? ARGUMENT_TO_LAMBDA_METAFACTORY
+                  : NOT_ARGUMENT_TO_LAMBDA_METAFACTORY;
+          List<DexValue> newArgs =
+              rewriteBootstrapArguments(callSite.bootstrapArgs, methodHandleUse, context);
+          if (!newMethodProto.equals(callSite.methodProto)
+              || newBootstrapMethod != callSite.bootstrapMethod
+              || !newArgs.equals(callSite.bootstrapArgs)) {
+            return dexItemFactory.createCallSite(
+                callSite.methodName, newMethodProto, newBootstrapMethod, newArgs);
+          }
+          return callSite;
+        });
   }
 
   public DexMethodHandle rewriteDexMethodHandle(
@@ -96,11 +106,15 @@
         // MethodHandles that are not arguments to a lambda metafactory will not be desugared
         // away. Therefore they could flow to a MethodHandle.invokeExact call which means that
         // we cannot member rebind. We therefore keep the receiver and also pin the receiver
-        // with a keep rule (see Enqueuer.registerMethodHandle).
+        // with a keep rule (see Enqueuer.traceMethodHandle).
+        // Note that the member can be repackaged or minified.
         actualTarget =
             definitions
                 .dexItemFactory()
-                .createMethod(invokedMethod.holder, rewrittenTarget.proto, rewrittenTarget.name);
+                .createMethod(
+                    graphLens().lookupType(invokedMethod.holder),
+                    rewrittenTarget.proto,
+                    rewrittenTarget.name);
         newType = oldType;
         if (oldType.isInvokeDirect()) {
           // For an invoke direct, the rewritten target must have the same holder as the original.
diff --git a/src/test/java/com/android/tools/r8/dex/DebugByteCodeWriterTest.java b/src/test/java/com/android/tools/r8/dex/DebugByteCodeWriterTest.java
index d2a5729..330535e 100644
--- a/src/test/java/com/android/tools/r8/dex/DebugByteCodeWriterTest.java
+++ b/src/test/java/com/android/tools/r8/dex/DebugByteCodeWriterTest.java
@@ -17,6 +17,7 @@
 import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.InitClassLens;
 import com.android.tools.r8.graph.ObjectToOffsetMapping;
+import com.android.tools.r8.ir.conversion.LensCodeRewriterUtils;
 import com.android.tools.r8.naming.NamingLens;
 import com.android.tools.r8.utils.InternalOptions;
 import com.android.tools.r8.utils.Reporter;
@@ -41,15 +42,18 @@
   }
 
   private ObjectToOffsetMapping emptyObjectTObjectMapping() {
-    return new ObjectToOffsetMapping(
+    AppView<AppInfo> appView =
         AppView.createForD8(
             AppInfo.createInitialAppInfo(
                 DexApplication.builder(
                         new InternalOptions(new DexItemFactory(), new Reporter()), null)
-                    .build())),
+                    .build()));
+    return new ObjectToOffsetMapping(
+        appView,
         GraphLens.getIdentityLens(),
         NamingLens.getIdentityLens(),
         InitClassLens.getDefault(),
+        new LensCodeRewriterUtils(appView),
         Collections.emptyList(),
         Collections.emptyList(),
         Collections.emptyList(),
diff --git a/src/test/java/com/android/tools/r8/repackage/RepackageCustomMethodHandleTest.java b/src/test/java/com/android/tools/r8/repackage/RepackageCustomMethodHandleTest.java
index a97ecf1..f20603d 100644
--- a/src/test/java/com/android/tools/r8/repackage/RepackageCustomMethodHandleTest.java
+++ b/src/test/java/com/android/tools/r8/repackage/RepackageCustomMethodHandleTest.java
@@ -4,12 +4,10 @@
 
 package com.android.tools.r8.repackage;
 
-import static com.android.tools.r8.DiagnosticsMatcher.diagnosticMessage;
 import static com.android.tools.r8.shaking.ProguardConfigurationParser.FLATTEN_PACKAGE_HIERARCHY;
 import static com.android.tools.r8.shaking.ProguardConfigurationParser.REPACKAGE_CLASSES;
-import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.MatcherAssert.assertThat;
 
-import com.android.tools.r8.CompilationFailedException;
 import com.android.tools.r8.NeverClassInline;
 import com.android.tools.r8.NeverInline;
 import com.android.tools.r8.TestParameters;
@@ -36,6 +34,8 @@
 @RunWith(Parameterized.class)
 public class RepackageCustomMethodHandleTest extends RepackageTestBase {
 
+  private static final String EXPECTED = "InvokeCustom::foo";
+
   @Parameters(name = "{1}, kind: {0}")
   public static List<Object[]> data() {
     return buildParameters(
@@ -59,29 +59,25 @@
         .addProgramClassFileData(
             transformer(Main.class).addClassTransformer(generateCallSiteInvoke()).transform())
         .run(parameters.getRuntime(), Main.class)
-        .assertSuccessWithOutputLines("InvokeCustom::foo");
+        .assertSuccessWithOutputLines(EXPECTED);
   }
 
-  @Test(expected = CompilationFailedException.class)
+  @Test()
   public void testR8() throws Exception {
     testForR8(parameters.getBackend())
         .addProgramClasses(InvokeCustom.class)
         .addProgramClassFileData(
             transformer(Main.class).addClassTransformer(generateCallSiteInvoke()).transform())
         .addKeepMainRule(Main.class)
+        .addKeepClassAndMembersRules(Main.class)
         .apply(this::configureRepackaging)
         .enableInliningAnnotations()
         .enableNeverClassInliningAnnotations()
         .setMinApi(parameters.getApiLevel())
-        .compileWithExpectedDiagnostics(
-            diagnostics -> {
-              // TODO(b/176486859): This should not lookup the original packaged name.
-              diagnostics.assertErrorsMatch(
-                  diagnosticMessage(
-                      containsString(
-                          "Failed lookup of non-missing type: "
-                              + InvokeCustom.class.getTypeName())));
-            });
+        .compile()
+        .inspect(inspector -> assertThat(InvokeCustom.class, isRepackaged(inspector)))
+        .run(parameters.getRuntime(), Main.class)
+        .assertSuccessWithOutputLines(EXPECTED);
   }
 
   private ClassTransformer generateCallSiteInvoke() {