commit 07a2ae7eda08623baeca57285c9bdbcb41e28070
author Christoffer Quist Adamsen <christofferqa@google.com> Tue Apr 13 19:33:06 2021 +0200
committer Christoffer Quist Adamsen <christofferqa@google.com> Tue Apr 13 19:33:06 2021 +0200
tree 9cccfa72804278d366835fb04e03a2827290ba5b
parent 49d3649076e6a71ea92684d4d4f11f1238cbdbe3

Fix lens rewriting of invoke signatures

Change-Id: I8744720c1a3180e4a784e2d20beff3b2c56d0bc7

src/main/java/com/android/tools/r8/graph/MethodAccessInfoCollection.java

diff --git a/src/main/java/com/android/tools/r8/graph/MethodAccessInfoCollection.java b/src/main/java/com/android/tools/r8/graph/MethodAccessInfoCollection.java
index 1af8be2..f56cc34 100644
--- a/src/main/java/com/android/tools/r8/graph/MethodAccessInfoCollection.java
+++ b/src/main/java/com/android/tools/r8/graph/MethodAccessInfoCollection.java
@@ -4,8 +4,9 @@
 
 package com.android.tools.r8.graph;
 
+import com.android.tools.r8.graph.GraphLens.MethodLookupResult;
+import com.android.tools.r8.ir.code.Invoke.Type;
 import com.android.tools.r8.utils.ConsumerUtils;
-import com.android.tools.r8.utils.MapUtils;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
 import com.google.common.collect.Sets;
 import java.util.IdentityHashMap;
@@ -89,24 +90,31 @@
 
   public MethodAccessInfoCollection rewrittenWithLens(
       DexDefinitionSupplier definitions, GraphLens lens) {
-    return new MethodAccessInfoCollection(
-        rewriteInvokesWithLens(directInvokes, definitions, lens),
-        rewriteInvokesWithLens(interfaceInvokes, definitions, lens),
-        rewriteInvokesWithLens(staticInvokes, definitions, lens),
-        rewriteInvokesWithLens(superInvokes, definitions, lens),
-        rewriteInvokesWithLens(virtualInvokes, definitions, lens));
+    MethodAccessInfoCollection.Builder<?> builder = identityBuilder();
+    rewriteInvokesWithLens(builder, directInvokes, definitions, lens, Type.DIRECT);
+    rewriteInvokesWithLens(builder, interfaceInvokes, definitions, lens, Type.INTERFACE);
+    rewriteInvokesWithLens(builder, staticInvokes, definitions, lens, Type.STATIC);
+    rewriteInvokesWithLens(builder, superInvokes, definitions, lens, Type.SUPER);
+    rewriteInvokesWithLens(builder, virtualInvokes, definitions, lens, Type.VIRTUAL);
+    return builder.build();
   }
 
-  private static Map<DexMethod, ProgramMethodSet> rewriteInvokesWithLens(
-      Map<DexMethod, ProgramMethodSet> invokes, DexDefinitionSupplier definitions, GraphLens lens) {
-    return MapUtils.map(
-        invokes,
-        IdentityHashMap::new,
-        lens::getRenamedMethodSignature,
-        methods -> methods.rewrittenWithLens(definitions, lens),
-        (methods, other) -> {
-          methods.addAll(other);
-          return methods;
+  private static void rewriteInvokesWithLens(
+      MethodAccessInfoCollection.Builder<?> builder,
+      Map<DexMethod, ProgramMethodSet> invokes,
+      DexDefinitionSupplier definitions,
+      GraphLens lens,
+      Type type) {
+    invokes.forEach(
+        (reference, contexts) -> {
+          ProgramMethodSet newContexts = contexts.rewrittenWithLens(definitions, lens);
+          for (ProgramMethod newContext : newContexts) {
+            MethodLookupResult methodLookupResult =
+                lens.lookupMethod(reference, newContext.getReference(), type);
+            DexMethod newReference = methodLookupResult.getReference();
+            Type newType = methodLookupResult.getType();
+            builder.registerInvokeInContext(newReference, newContext, newType);
+          }
         });
   }
 
@@ -151,6 +159,25 @@
       return virtualInvokes;
     }
 
+    public boolean registerInvokeInContext(
+        DexMethod invokedMethod, ProgramMethod context, Type type) {
+      switch (type) {
+        case DIRECT:
+          return registerInvokeDirectInContext(invokedMethod, context);
+        case INTERFACE:
+          return registerInvokeInterfaceInContext(invokedMethod, context);
+        case STATIC:
+          return registerInvokeStaticInContext(invokedMethod, context);
+        case SUPER:
+          return registerInvokeSuperInContext(invokedMethod, context);
+        case VIRTUAL:
+          return registerInvokeVirtualInContext(invokedMethod, context);
+        default:
+          assert false;
+          return false;
+      }
+    }
+
     public boolean registerInvokeDirectInContext(DexMethod invokedMethod, ProgramMethod context) {
       return registerInvokeMethodInContext(invokedMethod, context, directInvokes);
     }