Handle default interface methods in vertical class merging

Change-Id: Ib3f061ffa0326395be47f200bc01331fe0456cff
diff --git a/src/main/java/com/android/tools/r8/R8.java b/src/main/java/com/android/tools/r8/R8.java
index fd547cf..c03d928 100644
--- a/src/main/java/com/android/tools/r8/R8.java
+++ b/src/main/java/com/android/tools/r8/R8.java
@@ -369,7 +369,7 @@
         if (options.enableClassMerging && options.enableInlining) {
           timing.begin("ClassMerger");
           VerticalClassMerger classMerger =
-              new VerticalClassMerger(application, appViewWithLiveness, timing);
+              new VerticalClassMerger(application, appViewWithLiveness, executorService, timing);
           appView.setGraphLense(classMerger.run());
           timing.end();
           application = application.asDirect().rewrittenWithLense(appView.getGraphLense());
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/MethodPoolCollection.java b/src/main/java/com/android/tools/r8/ir/optimize/MethodPoolCollection.java
index 9af3967..bd3acb6 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/MethodPoolCollection.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/MethodPoolCollection.java
@@ -7,13 +7,16 @@
 import com.android.tools.r8.graph.DexApplication;
 import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexMethod;
-import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.utils.MethodSignatureEquivalence;
 import com.android.tools.r8.utils.ThreadUtils;
 import com.android.tools.r8.utils.Timing;
 import com.google.common.base.Equivalence;
 import com.google.common.base.Equivalence.Wrapper;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Deque;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -22,6 +25,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
+import java.util.function.Predicate;
 
 // Per-class collection of method signatures.
 //
@@ -41,13 +45,30 @@
     timing.begin("Building method pool collection");
     try {
       List<Future<?>> futures = new ArrayList<>();
-      submitAll(application.classes(), futures, executorService);
+      @SuppressWarnings("unchecked")
+      List<DexClass> classes = (List) application.classes();
+      submitAll(classes, futures, executorService);
       ThreadUtils.awaitFutures(futures);
     } finally {
       timing.end();
     }
   }
 
+  public MethodPool buildForHierarchy(
+      DexClass clazz, ExecutorService executorService, Timing timing) throws ExecutionException {
+    timing.begin("Building method pool collection");
+    try {
+      List<Future<?>> futures = new ArrayList<>();
+      submitAll(
+          getAllSuperTypesInclusive(clazz, methodPools::containsKey), futures, executorService);
+      submitAll(getAllSubTypesExclusive(clazz, methodPools::containsKey), futures, executorService);
+      ThreadUtils.awaitFutures(futures);
+    } finally {
+      timing.end();
+    }
+    return get(clazz);
+  }
+
   public MethodPool get(DexClass clazz) {
     assert methodPools.containsKey(clazz);
     return methodPools.get(clazz);
@@ -64,8 +85,8 @@
   }
 
   private void submitAll(
-      Iterable<DexProgramClass> classes, List<Future<?>> futures, ExecutorService executorService) {
-    for (DexProgramClass clazz : classes) {
+      Iterable<DexClass> classes, List<Future<?>> futures, ExecutorService executorService) {
+    for (DexClass clazz : classes) {
       futures.add(executorService.submit(computeMethodPoolPerClass(clazz)));
     }
   }
@@ -101,6 +122,51 @@
     };
   }
 
+  private Set<DexClass> getAllSuperTypesInclusive(
+      DexClass subject, Predicate<DexClass> stoppingCriterion) {
+    Set<DexClass> superTypes = new HashSet<>();
+    Deque<DexClass> worklist = new ArrayDeque<>();
+    worklist.add(subject);
+    while (!worklist.isEmpty()) {
+      DexClass clazz = worklist.pop();
+      if (stoppingCriterion.test(clazz)) {
+        continue;
+      }
+      if (superTypes.add(clazz)) {
+        if (clazz.superType != null) {
+          addNonNull(worklist, application.definitionFor(clazz.superType));
+        }
+        for (DexType interfaceType : clazz.interfaces.values) {
+          addNonNull(worklist, application.definitionFor(interfaceType));
+        }
+      }
+    }
+    return superTypes;
+  }
+
+  private Set<DexClass> getAllSubTypesExclusive(
+      DexClass subject, Predicate<DexClass> stoppingCriterion) {
+    Set<DexClass> subTypes = new HashSet<>();
+    Deque<DexClass> worklist = new ArrayDeque<>();
+    subject.type.forAllExtendsSubtypes(
+        type -> addNonNull(worklist, application.definitionFor(type)));
+    subject.type.forAllImplementsSubtypes(
+        type -> addNonNull(worklist, application.definitionFor(type)));
+    while (!worklist.isEmpty()) {
+      DexClass clazz = worklist.pop();
+      if (stoppingCriterion.test(clazz)) {
+        continue;
+      }
+      if (subTypes.add(clazz)) {
+        clazz.type.forAllExtendsSubtypes(
+            type -> addNonNull(worklist, application.definitionFor(type)));
+        clazz.type.forAllImplementsSubtypes(
+            type -> addNonNull(worklist, application.definitionFor(type)));
+      }
+    }
+    return subTypes;
+  }
+
   public static class MethodPool {
     private MethodPool superType;
     private final Set<MethodPool> interfaces = new HashSet<>();
@@ -124,6 +190,10 @@
       assert added;
     }
 
+    public void seen(DexMethod method) {
+      seen(MethodSignatureEquivalence.get().wrap(method));
+    }
+
     public synchronized void seen(Wrapper<DexMethod> method) {
       boolean added = methodPool.add(method);
       assert added;
@@ -144,4 +214,10 @@
           || subTypes.stream().anyMatch(subType -> subType.hasSeenDownwardRecursive(method));
     }
   }
+
+  private static <T> void addNonNull(Collection<T> collection, T item) {
+    if (item != null) {
+      collection.add(item);
+    }
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/shaking/VerticalClassMerger.java b/src/main/java/com/android/tools/r8/shaking/VerticalClassMerger.java
index 4d072c4..451b2cb 100644
--- a/src/main/java/com/android/tools/r8/shaking/VerticalClassMerger.java
+++ b/src/main/java/com/android/tools/r8/shaking/VerticalClassMerger.java
@@ -7,6 +7,7 @@
 import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.graph.AppInfo.ResolutionResult;
 import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.Code;
 import com.android.tools.r8.graph.DexAnnotationSet;
 import com.android.tools.r8.graph.DexApplication;
 import com.android.tools.r8.graph.DexClass;
@@ -22,6 +23,7 @@
 import com.android.tools.r8.graph.DexTypeList;
 import com.android.tools.r8.graph.GraphLense;
 import com.android.tools.r8.graph.GraphLense.Builder;
+import com.android.tools.r8.graph.GraphLense.GraphLenseLookupResult;
 import com.android.tools.r8.graph.JarCode;
 import com.android.tools.r8.graph.KeyedDexItem;
 import com.android.tools.r8.graph.MethodAccessFlags;
@@ -30,6 +32,8 @@
 import com.android.tools.r8.graph.UseRegistry;
 import com.android.tools.r8.ir.code.Invoke.Type;
 import com.android.tools.r8.ir.optimize.Inliner.Constraint;
+import com.android.tools.r8.ir.optimize.MethodPoolCollection;
+import com.android.tools.r8.ir.optimize.MethodPoolCollection.MethodPool;
 import com.android.tools.r8.ir.synthetic.ForwardMethodSourceCode;
 import com.android.tools.r8.ir.synthetic.SynthesizedCode;
 import com.android.tools.r8.logging.Log;
@@ -58,7 +62,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
 import java.util.function.Predicate;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.MethodNode;
 
 /**
  * Merges Supertypes with a single implementation into their single subtype.
@@ -150,7 +158,9 @@
 
   private final DexApplication application;
   private final AppInfoWithLiveness appInfo;
+  private final ExecutorService executorService;
   private final GraphLense graphLense;
+  private final MethodPoolCollection methodPoolCollection;
   private final Timing timing;
   private Collection<DexMethod> invokes;
 
@@ -170,10 +180,15 @@
   private final VerticalClassMergerGraphLense.Builder renamedMembersLense;
 
   public VerticalClassMerger(
-      DexApplication application, AppView<AppInfoWithLiveness> appView, Timing timing) {
+      DexApplication application,
+      AppView<AppInfoWithLiveness> appView,
+      ExecutorService executorService,
+      Timing timing) {
     this.application = application;
     this.appInfo = appView.getAppInfo();
+    this.executorService = executorService;
     this.graphLense = appView.getGraphLense();
+    this.methodPoolCollection = new MethodPoolCollection(application);
     this.renamedMembersLense = VerticalClassMergerGraphLense.builder(appInfo);
     this.timing = timing;
 
@@ -520,7 +535,7 @@
     }
   }
 
-  public GraphLense run() {
+  public GraphLense run() throws ExecutionException {
     timing.begin("merge");
     GraphLense mergingGraphLense = mergeClasses(graphLense);
     timing.end();
@@ -561,7 +576,7 @@
     }
   }
 
-  private GraphLense mergeClasses(GraphLense graphLense) {
+  private GraphLense mergeClasses(GraphLense graphLense) throws ExecutionException {
     Deque<DexProgramClass> worklist = new ArrayDeque<>();
     Set<DexProgramClass> seenBefore = new HashSet<>();
 
@@ -726,7 +741,7 @@
       this.target = target;
     }
 
-    public boolean merge() {
+    public boolean merge() throws ExecutionException {
       // Merge the class [clazz] into [targetClass] by adding all methods to
       // targetClass that are not currently contained.
       // Step 1: Merge methods
@@ -792,12 +807,37 @@
           }
         }
 
-        // This virtual method could be called directly from a sub class via an invoke-super
-        // instruction. Therefore, we translate this virtual method into a direct method, such that
-        // relevant invoke-super instructions can be rewritten into invoke-direct instructions.
-        DexEncodedMethod resultingDirectMethod =
-            renameMethod(virtualMethod, availableMethodSignatures, Rename.ALWAYS);
-        makePrivate(resultingDirectMethod);
+        DexEncodedMethod resultingDirectMethod;
+        if (source.accessFlags.isInterface()) {
+          // Moving a default interface method into its subtype. This method could be hit directly
+          // via an invoke-super instruction from any of the transitive subtypes of this interface,
+          // due to the way invoke-super works on default interface methods. In order to be able
+          // to hit this method directly after the merge, we need to make it public, and find a
+          // method name that does not collide with one in the hierarchy of this class.
+          MethodPool methodPoolForTarget =
+              methodPoolCollection.buildForHierarchy(target, executorService, timing);
+          resultingDirectMethod =
+              renameMethod(
+                  virtualMethod,
+                  method ->
+                      availableMethodSignatures.test(method)
+                          && !methodPoolForTarget.hasSeen(
+                              MethodSignatureEquivalence.get().wrap(method)),
+                  Rename.ALWAYS,
+                  getStaticProto(virtualMethod.method.holder, virtualMethod.method.proto));
+          makeStatic(resultingDirectMethod);
+
+          // Update method pool collection now that we are adding a new public method.
+          methodPoolForTarget.seen(resultingDirectMethod.method);
+        } else {
+          // This virtual method could be called directly from a sub class via an invoke-super in-
+          // struction. Therefore, we translate this virtual method into a direct method, such that
+          // relevant invoke-super instructions can be rewritten into invoke-direct instructions.
+          resultingDirectMethod =
+              renameMethod(virtualMethod, availableMethodSignatures, Rename.ALWAYS);
+          makePrivate(resultingDirectMethod);
+        }
+
         add(directMethods, resultingDirectMethod, MethodSignatureEquivalence.get());
 
         // Record that invoke-super instructions in the target class should be redirected to the
@@ -811,7 +851,7 @@
           // Note that this method is added independently of whether it will actually be used. If
           // it turns out that the method is never used, it will be removed by the final round
           // of tree shaking.
-          shadowedBy = buildBridgeMethod(virtualMethod, resultingDirectMethod.method);
+          shadowedBy = buildBridgeMethod(virtualMethod, resultingDirectMethod);
           add(virtualMethods, shadowedBy, MethodSignatureEquivalence.get());
         }
 
@@ -903,7 +943,8 @@
         // rewrite any invocations on the form "invoke-super J.m()" to "invoke-direct C.m$I()",
         // if I has a supertype J. This is due to the fact that invoke-super instructions that
         // resolve to a method on an interface never hit an implementation below that interface.
-        deferredRenamings.mapVirtualMethodToDirectInType(oldTarget, newTarget, target.type);
+        deferredRenamings.mapVirtualMethodToDirectInType(
+            oldTarget, new GraphLenseLookupResult(newTarget, Type.STATIC), target.type);
       } else {
         // If we merge class B into class C, and class C contains an invocation super.m(), then it
         // is insufficient to rewrite "invoke-super B.m()" to "invoke-direct C.m$B()" (the method
@@ -922,7 +963,7 @@
                   || appInfo.lookupSuperTarget(signatureInHolder, holder.type) != null;
           if (resolutionSucceeds) {
             deferredRenamings.mapVirtualMethodToDirectInType(
-                signatureInHolder, newTarget, target.type);
+                signatureInHolder, new GraphLenseLookupResult(newTarget, Type.DIRECT), target.type);
           } else {
             break;
           }
@@ -944,7 +985,9 @@
                       || appInfo.lookupSuperTarget(signatureInHolder, holder.type) != null;
               if (resolutionSucceededBeforeMerge) {
                 deferredRenamings.mapVirtualMethodToDirectInType(
-                    signatureInType, newTarget, target.type);
+                    signatureInType,
+                    new GraphLenseLookupResult(newTarget, Type.DIRECT),
+                    target.type);
               }
             }
           }
@@ -974,23 +1017,38 @@
     }
 
     private DexEncodedMethod buildBridgeMethod(
-        DexEncodedMethod method, DexMethod invocationTarget) {
+        DexEncodedMethod method, DexEncodedMethod invocationTarget) {
       DexType holder = target.type;
-      DexProto proto = invocationTarget.proto;
+      DexProto proto = method.method.proto;
       DexString name = method.method.name;
       MethodAccessFlags accessFlags = method.accessFlags.copy();
       accessFlags.setBridge();
       accessFlags.setSynthetic();
       accessFlags.unsetAbstract();
-      DexEncodedMethod bridge = new DexEncodedMethod(
-          application.dexItemFactory.createMethod(holder, proto, name),
-          accessFlags,
-          DexAnnotationSet.empty(),
-          ParameterAnnotationsList.empty(),
-          new SynthesizedCode(
-              new ForwardMethodSourceCode(holder, proto, holder, invocationTarget, Type.DIRECT),
-              registry -> registry.registerInvokeDirect(invocationTarget)),
-          method.hasClassFileVersion() ? method.getClassFileVersion() : -1);
+      SynthesizedCode code;
+      if (invocationTarget.isPrivateMethod()) {
+        assert !invocationTarget.isStaticMethod();
+        code =
+            new SynthesizedCode(
+                new ForwardMethodSourceCode(
+                    holder, proto, holder, invocationTarget.method, Type.DIRECT),
+                registry -> registry.registerInvokeDirect(invocationTarget.method));
+      } else {
+        assert invocationTarget.isStaticMethod();
+        code =
+            new SynthesizedCode(
+                new ForwardMethodSourceCode(
+                    holder, proto, null, invocationTarget.method, Type.STATIC),
+                registry -> registry.registerInvokeStatic(invocationTarget.method));
+      }
+      DexEncodedMethod bridge =
+          new DexEncodedMethod(
+              application.dexItemFactory.createMethod(holder, proto, name),
+              accessFlags,
+              DexAnnotationSet.empty(),
+              ParameterAnnotationsList.empty(),
+              code,
+              method.hasClassFileVersion() ? method.getClassFileVersion() : -1);
       if (method.getOptimizationInfo().isPublicized()) {
         // The bridge is now the public method serving the role of the original method, and should
         // reflect that this method was publicized.
@@ -1106,6 +1164,14 @@
 
     private DexEncodedMethod renameMethod(
         DexEncodedMethod method, Predicate<DexMethod> availableMethodSignatures, Rename strategy) {
+      return renameMethod(method, availableMethodSignatures, strategy, method.method.proto);
+    }
+
+    private DexEncodedMethod renameMethod(
+        DexEncodedMethod method,
+        Predicate<DexMethod> availableMethodSignatures,
+        Rename strategy,
+        DexProto newProto) {
       // We cannot handle renaming static initializers yet and constructors should have been
       // renamed already.
       assert !method.accessFlags.isConstructor() || strategy == Rename.NEVER;
@@ -1115,8 +1181,7 @@
       DexMethod newSignature;
       switch (strategy) {
         case IF_NEEDED:
-          newSignature =
-              application.dexItemFactory.createMethod(target.type, method.method.proto, oldName);
+          newSignature = application.dexItemFactory.createMethod(target.type, newProto, oldName);
           if (availableMethodSignatures.test(newSignature)) {
             break;
           }
@@ -1126,15 +1191,13 @@
           int count = 1;
           do {
             DexString newName = getFreshName(oldName.toSourceString(), count, oldHolder);
-            newSignature =
-                application.dexItemFactory.createMethod(target.type, method.method.proto, newName);
+            newSignature = application.dexItemFactory.createMethod(target.type, newProto, newName);
             count++;
           } while (!availableMethodSignatures.test(newSignature));
           break;
 
         case NEVER:
-          newSignature =
-              application.dexItemFactory.createMethod(target.type, method.method.proto, oldName);
+          newSignature = application.dexItemFactory.createMethod(target.type, newProto, oldName);
           assert availableMethodSignatures.test(newSignature);
           break;
 
@@ -1173,6 +1236,24 @@
     method.accessFlags.setPrivate();
   }
 
+  private static void makeStatic(DexEncodedMethod method) {
+    method.accessFlags.setStatic();
+
+    Code code = method.getCode();
+    if (code.isJarCode()) {
+      MethodNode node = code.asJarCode().getNode();
+      node.access |= Opcodes.ACC_STATIC;
+      node.desc = method.method.proto.toDescriptorString();
+    }
+  }
+
+  private DexProto getStaticProto(DexType receiverType, DexProto proto) {
+    DexType[] parameterTypes = new DexType[proto.parameters.size() + 1];
+    parameterTypes[0] = receiverType;
+    System.arraycopy(proto.parameters.values, 0, parameterTypes, 1, proto.parameters.size());
+    return appInfo.dexItemFactory.createProto(proto.returnType, parameterTypes);
+  }
+
   private class TreeFixer {
 
     private final Builder lense = GraphLense.builder();
diff --git a/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java b/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java
index fa10215..5974e60 100644
--- a/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java
+++ b/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java
@@ -48,14 +48,15 @@
   private final AppInfo appInfo;
 
   private final Set<DexMethod> mergedMethods;
-  private final Map<DexType, Map<DexMethod, DexMethod>> contextualVirtualToDirectMethodMaps;
+  private final Map<DexType, Map<DexMethod, GraphLenseLookupResult>>
+      contextualVirtualToDirectMethodMaps;
 
   public VerticalClassMergerGraphLense(
       AppInfo appInfo,
       Map<DexField, DexField> fieldMap,
       Map<DexMethod, DexMethod> methodMap,
       Set<DexMethod> mergedMethods,
-      Map<DexType, Map<DexMethod, DexMethod>> contextualVirtualToDirectMethodMaps,
+      Map<DexType, Map<DexMethod, GraphLenseLookupResult>> contextualVirtualToDirectMethodMaps,
       GraphLense previousLense) {
     super(ImmutableMap.of(), methodMap, fieldMap, previousLense, appInfo.dexItemFactory);
     this.appInfo = appInfo;
@@ -73,17 +74,18 @@
     assert isContextFreeForMethod(method) || (context != null && type != null);
     GraphLenseLookupResult previous = previousLense.lookupMethod(method, context, type);
     if (previous.getType() == Type.SUPER && !mergedMethods.contains(context.method)) {
-      Map<DexMethod, DexMethod> virtualToDirectMethodMap =
+      Map<DexMethod, GraphLenseLookupResult> virtualToDirectMethodMap =
           contextualVirtualToDirectMethodMaps.get(context.method.holder);
       if (virtualToDirectMethodMap != null) {
-        DexMethod directMethod = virtualToDirectMethodMap.get(previous.getMethod());
-        if (directMethod != null) {
+        GraphLenseLookupResult lookup = virtualToDirectMethodMap.get(previous.getMethod());
+        if (lookup != null) {
           // If the super class A of the enclosing class B (i.e., context.method.holder)
           // has been merged into B during vertical class merging, and this invoke-super instruction
           // was resolving to a method in A, then the target method has been changed to a direct
           // method and moved into B, so that we need to use an invoke-direct instruction instead of
-          // invoke-super.
-          return new GraphLenseLookupResult(directMethod, Type.DIRECT);
+          // invoke-super (or invoke-static, if the method was originally a default interface
+          // method).
+          return lookup;
         }
       }
     }
@@ -102,11 +104,11 @@
     ImmutableSet.Builder<DexMethod> builder = ImmutableSet.builder();
     for (DexMethod previous : previousLense.lookupMethodInAllContexts(method)) {
       builder.add(methodMap.getOrDefault(previous, previous));
-      for (Map<DexMethod, DexMethod> virtualToDirectMethodMap :
+      for (Map<DexMethod, GraphLenseLookupResult> virtualToDirectMethodMap :
           contextualVirtualToDirectMethodMaps.values()) {
-        DexMethod directMethod = virtualToDirectMethodMap.get(previous);
-        if (directMethod != null) {
-          builder.add(directMethod);
+        GraphLenseLookupResult lookup = virtualToDirectMethodMap.get(previous);
+        if (lookup != null) {
+          builder.add(lookup.getMethod());
         }
       }
     }
@@ -124,7 +126,7 @@
       return false;
     }
     DexMethod previous = previousLense.lookupMethod(method);
-    for (Map<DexMethod, DexMethod> virtualToDirectMethodMap :
+    for (Map<DexMethod, GraphLenseLookupResult> virtualToDirectMethodMap :
         contextualVirtualToDirectMethodMaps.values()) {
       if (virtualToDirectMethodMap.containsKey(previous)) {
         return false;
@@ -139,8 +141,8 @@
     protected final Map<DexField, DexField> fieldMap = new HashMap<>();
     protected final Map<DexMethod, DexMethod> methodMap = new HashMap<>();
     private final ImmutableSet.Builder<DexMethod> mergedMethodsBuilder = ImmutableSet.builder();
-    private final Map<DexType, Map<DexMethod, DexMethod>> contextualVirtualToDirectMethodMaps =
-        new HashMap<>();
+    private final Map<DexType, Map<DexMethod, GraphLenseLookupResult>>
+        contextualVirtualToDirectMethodMaps = new HashMap<>();
 
     private Builder(AppInfo appInfo) {
       this.appInfo = appInfo;
@@ -189,7 +191,7 @@
     }
 
     public boolean hasMappingForSignatureInContext(DexType context, DexMethod signature) {
-      Map<DexMethod, DexMethod> virtualToDirectMethodMap =
+      Map<DexMethod, GraphLenseLookupResult> virtualToDirectMethodMap =
           contextualVirtualToDirectMethodMaps.get(context);
       if (virtualToDirectMethodMap != null) {
         return virtualToDirectMethodMap.containsKey(signature);
@@ -209,8 +211,9 @@
       methodMap.put(from, to);
     }
 
-    public void mapVirtualMethodToDirectInType(DexMethod from, DexMethod to, DexType type) {
-      Map<DexMethod, DexMethod> virtualToDirectMethodMap =
+    public void mapVirtualMethodToDirectInType(
+        DexMethod from, GraphLenseLookupResult to, DexType type) {
+      Map<DexMethod, GraphLenseLookupResult> virtualToDirectMethodMap =
           contextualVirtualToDirectMethodMaps.computeIfAbsent(type, key -> new HashMap<>());
       virtualToDirectMethodMap.put(from, to);
     }
@@ -220,8 +223,10 @@
       methodMap.putAll(builder.methodMap);
       mergedMethodsBuilder.addAll(builder.mergedMethodsBuilder.build());
       for (DexType context : builder.contextualVirtualToDirectMethodMaps.keySet()) {
-        Map<DexMethod, DexMethod> current = contextualVirtualToDirectMethodMaps.get(context);
-        Map<DexMethod, DexMethod> other = builder.contextualVirtualToDirectMethodMaps.get(context);
+        Map<DexMethod, GraphLenseLookupResult> current =
+            contextualVirtualToDirectMethodMaps.get(context);
+        Map<DexMethod, GraphLenseLookupResult> other =
+            builder.contextualVirtualToDirectMethodMaps.get(context);
         if (current != null) {
           current.putAll(other);
         } else {