Remove all uses of lookupMethodInAllContexts()

Change-Id: I38d3c259aba1b19b30c65d91b1cec364afe45f8f
Bug: 157446971
diff --git a/src/main/java/com/android/tools/r8/R8.java b/src/main/java/com/android/tools/r8/R8.java
index 8406706..7984ff2 100644
--- a/src/main/java/com/android/tools/r8/R8.java
+++ b/src/main/java/com/android/tools/r8/R8.java
@@ -515,10 +515,8 @@
             assert changed;
             appView.setVerticallyMergedClasses(verticalClassMerger.getMergedClasses());
             application = application.asDirect().rewrittenWithLens(lens);
-            lens.initializeCacheForLookupMethodInAllContexts();
             appViewWithLiveness.setAppInfo(
                 appViewWithLiveness.appInfo().rewrittenWithLens(application.asDirect(), lens));
-            lens.unsetCacheForLookupMethodInAllContexts();
           }
           timing.end();
         }
diff --git a/src/main/java/com/android/tools/r8/graph/GraphLense.java b/src/main/java/com/android/tools/r8/graph/GraphLense.java
index ea55ec4..95d4085 100644
--- a/src/main/java/com/android/tools/r8/graph/GraphLense.java
+++ b/src/main/java/com/android/tools/r8/graph/GraphLense.java
@@ -4,6 +4,7 @@
 package com.android.tools.r8.graph;
 
 import com.android.tools.r8.ir.code.Invoke.Type;
+import com.android.tools.r8.utils.SetUtils;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import com.google.common.collect.ImmutableMap;
@@ -18,7 +19,6 @@
 import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.Function;
 import java.util.function.Supplier;
 
 /**
@@ -189,13 +189,6 @@
 
   public abstract RewrittenPrototypeDescription lookupPrototypeChanges(DexMethod method);
 
-  // Context sensitive graph lenses should override this method.
-  public Set<DexMethod> lookupMethodInAllContexts(DexMethod method) {
-    DexMethod result = lookupMethod(method);
-    assert result != null;
-    return ImmutableSet.of(result);
-  }
-
   public abstract DexField lookupField(DexField field);
 
   public DexMethod lookupGetFieldForMethod(DexField field, DexMethod context) {
@@ -280,37 +273,40 @@
     return true;
   }
 
-  public ImmutableSet<DexReference> rewriteReferencesConservatively(Set<DexReference> original) {
-    ImmutableSet.Builder<DexReference> builder = ImmutableSet.builder();
-    for (DexReference item : original) {
-      if (item.isDexMethod()) {
-        DexMethod method = item.asDexMethod();
-        builder.addAll(lookupMethodInAllContexts(method));
-      } else {
-        builder.add(lookupReference(item));
-      }
+  public DexReference rewriteReference(DexReference reference) {
+    if (reference.isDexField()) {
+      return getRenamedFieldSignature(reference.asDexField());
     }
-    return builder.build();
+    if (reference.isDexMethod()) {
+      return getRenamedMethodSignature(reference.asDexMethod());
+    }
+    assert reference.isDexType();
+    return lookupType(reference.asDexType());
   }
 
-  public Object2BooleanMap<DexReference> rewriteReferencesConservatively(
-      Object2BooleanMap<DexReference> original) {
-    Object2BooleanMap<DexReference> result = new Object2BooleanArrayMap<>();
-    for (Object2BooleanMap.Entry<DexReference> entry : original.object2BooleanEntrySet()) {
-      DexReference item = entry.getKey();
-      if (item.isDexMethod()) {
-        DexMethod method = item.asDexMethod();
-        for (DexMethod candidate : lookupMethodInAllContexts(method)) {
-          result.put(candidate, entry.getBooleanValue());
-        }
-      } else {
-        result.put(lookupReference(item), entry.getBooleanValue());
-      }
+  public Set<DexReference> rewriteReferences(Set<DexReference> references) {
+    Set<DexReference> result = SetUtils.newIdentityHashSet(references.size());
+    for (DexReference reference : references) {
+      result.add(rewriteReference(reference));
     }
     return result;
   }
 
-  public ImmutableSortedSet<DexMethod> rewriteMethodsWithRenamedSignature(Set<DexMethod> methods) {
+  public <T> ImmutableMap<DexReference, T> rewriteReferenceKeys(Map<DexReference, T> map) {
+    ImmutableMap.Builder<DexReference, T> builder = ImmutableMap.builder();
+    map.forEach((reference, value) -> builder.put(rewriteReference(reference), value));
+    return builder.build();
+  }
+
+  public Object2BooleanMap<DexReference> rewriteReferenceKeys(Object2BooleanMap<DexReference> map) {
+    Object2BooleanMap<DexReference> result = new Object2BooleanArrayMap<>();
+    for (Object2BooleanMap.Entry<DexReference> entry : map.object2BooleanEntrySet()) {
+      result.put(rewriteReference(entry.getKey()), entry.getBooleanValue());
+    }
+    return result;
+  }
+
+  public ImmutableSortedSet<DexMethod> rewriteMethods(Set<DexMethod> methods) {
     ImmutableSortedSet.Builder<DexMethod> builder =
         new ImmutableSortedSet.Builder<>(PresortedComparable::slowCompare);
     for (DexMethod method : methods) {
@@ -319,25 +315,24 @@
     return builder.build();
   }
 
-  public ImmutableSortedSet<DexMethod> rewriteMethodsConservatively(Set<DexMethod> original) {
-    ImmutableSortedSet.Builder<DexMethod> builder =
+  public <T> ImmutableMap<DexField, T> rewriteFieldKeys(Map<DexField, T> map) {
+    ImmutableMap.Builder<DexField, T> builder = ImmutableMap.builder();
+    map.forEach((field, value) -> builder.put(getRenamedFieldSignature(field), value));
+    return builder.build();
+  }
+
+  public ImmutableSet<DexType> rewriteTypes(Set<DexType> types) {
+    ImmutableSortedSet.Builder<DexType> builder =
         new ImmutableSortedSet.Builder<>(PresortedComparable::slowCompare);
-    if (isContextFreeForMethods()) {
-      for (DexMethod item : original) {
-        builder.add(lookupMethod(item));
-      }
-    } else {
-      for (DexMethod item : original) {
-        builder.addAll(lookupMethodInAllContexts(item));
-      }
+    for (DexType type : types) {
+      builder.add(lookupType(type));
     }
     return builder.build();
   }
 
-  public static <T extends DexReference, S> ImmutableMap<T, S> rewriteReferenceKeys(
-      Map<T, S> original, Function<T, T> rewrite) {
-    ImmutableMap.Builder<T, S> builder = ImmutableMap.builder();
-    original.forEach((item, value) -> builder.put(rewrite.apply(item), value));
+  public <T> ImmutableMap<DexType, T> rewriteTypeKeys(Map<DexType, T> map) {
+    ImmutableMap.Builder<DexType, T> builder = ImmutableMap.builder();
+    map.forEach((type, value) -> builder.put(lookupType(type), value));
     return builder.build();
   }
 
@@ -728,15 +723,6 @@
     }
 
     @Override
-    public Set<DexMethod> lookupMethodInAllContexts(DexMethod method) {
-      Set<DexMethod> result = Sets.newIdentityHashSet();
-      for (DexMethod previous : previousLense.lookupMethodInAllContexts(method)) {
-        result.add(methodMap.getOrDefault(previous, previous));
-      }
-      return result;
-    }
-
-    @Override
     public DexField lookupField(DexField field) {
       DexField previous = previousLense.lookupField(field);
       return fieldMap.getOrDefault(previous, previous);
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/NestedPrivateMethodLense.java b/src/main/java/com/android/tools/r8/ir/desugar/NestedPrivateMethodLense.java
index a6118a8..ded1538 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/NestedPrivateMethodLense.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/NestedPrivateMethodLense.java
@@ -13,10 +13,8 @@
 import com.android.tools.r8.graph.RewrittenPrototypeDescription;
 import com.android.tools.r8.ir.code.Invoke;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Sets;
 import java.util.IdentityHashMap;
 import java.util.Map;
-import java.util.Set;
 
 public class NestedPrivateMethodLense extends NestedGraphLense {
 
@@ -109,19 +107,6 @@
   }
 
   @Override
-  public Set<DexMethod> lookupMethodInAllContexts(DexMethod method) {
-    Set<DexMethod> result = Sets.newIdentityHashSet();
-    for (DexMethod previous : previousLense.lookupMethodInAllContexts(method)) {
-      result.add(previous);
-      DexMethod bridge = methodMap.get(previous);
-      if (bridge != null) {
-        result.add(bridge);
-      }
-    }
-    return result;
-  }
-
-  @Override
   public GraphLenseLookupResult lookupMethod(
       DexMethod method, DexMethod context, Invoke.Type type) {
     assert originalMethodSignatures == null;
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxer.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxer.java
index 3239c8c..dcbf623 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxer.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/EnumUnboxer.java
@@ -4,6 +4,7 @@
 
 package com.android.tools.r8.ir.optimize.enums;
 
+import static com.android.tools.r8.graph.DexProgramClass.asProgramClassOrNull;
 import static com.android.tools.r8.ir.analysis.type.Nullability.definitelyNotNull;
 
 import com.android.tools.r8.graph.AppView;
@@ -53,6 +54,7 @@
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedback.OptimizationInfoFixer;
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedbackDelayed;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.shaking.AppInfoWithLivenessModifier;
 import com.android.tools.r8.utils.BooleanUtils;
 import com.android.tools.r8.utils.Reporter;
 import com.android.tools.r8.utils.StringDiagnostic;
@@ -325,6 +327,8 @@
       return;
     }
     ImmutableSet<DexType> enumsToUnbox = ImmutableSet.copyOf(this.enumsUnboxingCandidates.keySet());
+    // Update keep info on any of the enum methods of the removed classes.
+    updatePinnedItems(enumsToUnbox);
     enumUnboxerRewriter = new EnumUnboxingRewriter(appView, enumsToUnbox);
     NestedGraphLense enumUnboxingLens = new TreeFixer(enumsToUnbox).fixupTypeReferences();
     appView.setUnboxedEnums(enumUnboxerRewriter.getEnumsToUnbox());
@@ -368,6 +372,16 @@
     postBuilder.rewrittenWithLens(appView, previousLens);
   }
 
+  private void updatePinnedItems(Set<DexType> enumsToUnbox) {
+    AppInfoWithLivenessModifier modifier = AppInfoWithLiveness.modifier();
+    for (DexType type : enumsToUnbox) {
+      DexProgramClass clazz = asProgramClassOrNull(appView.definitionFor(type));
+      assert !appView.appInfo().isPinned(clazz.type);
+      modifier.removePinnedClassMembers(clazz);
+    }
+    modifier.modify(appView.appInfo());
+  }
+
   public void finishAnalysis() {
     for (DexType toUnbox : enumsUnboxingCandidates.keySet()) {
       DexProgramClass enumClass = appView.definitionForProgramType(toUnbox);
diff --git a/src/main/java/com/android/tools/r8/optimize/MemberRebindingLense.java b/src/main/java/com/android/tools/r8/optimize/MemberRebindingLense.java
index 95e2772..01c923d 100644
--- a/src/main/java/com/android/tools/r8/optimize/MemberRebindingLense.java
+++ b/src/main/java/com/android/tools/r8/optimize/MemberRebindingLense.java
@@ -15,7 +15,6 @@
 import java.util.Collections;
 import java.util.IdentityHashMap;
 import java.util.Map;
-import java.util.Set;
 
 public class MemberRebindingLense extends NestedGraphLense {
 
@@ -99,11 +98,6 @@
   }
 
   @Override
-  public Set<DexMethod> lookupMethodInAllContexts(DexMethod method) {
-    return previousLense.lookupMethodInAllContexts(method);
-  }
-
-  @Override
   protected Type mapInvocationType(DexMethod newMethod, DexMethod originalMethod, Type type) {
     return super.mapVirtualInterfaceInvocationTypes(appView, newMethod, originalMethod, type);
   }
diff --git a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
index 989ae9a..34fdfe7 100644
--- a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
+++ b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
@@ -5,7 +5,6 @@
 
 import static com.android.tools.r8.graph.DexEncodedMethod.asProgramMethodOrNull;
 import static com.android.tools.r8.graph.DexProgramClass.asProgramClassOrNull;
-import static com.android.tools.r8.graph.GraphLense.rewriteReferenceKeys;
 import static com.android.tools.r8.graph.ResolutionResult.SingleResolutionResult.isOverriding;
 
 import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
@@ -48,12 +47,11 @@
 import com.android.tools.r8.utils.InternalOptions;
 import com.android.tools.r8.utils.ListUtils;
 import com.android.tools.r8.utils.PredicateSet;
+import com.android.tools.r8.utils.SetUtils;
 import com.android.tools.r8.utils.TraversalContinuation;
 import com.android.tools.r8.utils.Visibility;
 import com.android.tools.r8.utils.collections.ProgramMethodSet;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedSet;
-import com.google.common.collect.ImmutableSortedSet.Builder;
 import com.google.common.collect.Sets;
 import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import it.unimi.dsi.fastutil.objects.Object2BooleanMap;
@@ -69,7 +67,6 @@
 import java.util.SortedSet;
 import java.util.TreeMap;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /** Encapsulates liveness and reachability information for an application. */
@@ -445,7 +442,7 @@
         previous.callSites,
         additionalPinnedItems == null
             ? previous.pinnedItems
-            : CollectionUtils.mergeSets(previous.pinnedItems, additionalPinnedItems),
+            : SetUtils.newIdentityHashSet(previous.pinnedItems, additionalPinnedItems),
         previous.allowAccessModification,
         previous.mayHaveSideEffects,
         previous.noSideEffects,
@@ -752,6 +749,10 @@
     singleTargetLookupCache.removeInstantiatedType(clazz.type, this);
   }
 
+  void removePinnedItem(DexReference item) {
+    pinnedItems.remove(item);
+  }
+
   private boolean assertNoItemRemoved(Collection<DexReference> items, Collection<DexType> types) {
     Set<DexType> typeSet = ImmutableSet.copyOf(types);
     for (DexReference item : items) {
@@ -882,27 +883,15 @@
     return holder == null || holder.isLibraryClass() || holder.isClasspathClass();
   }
 
-  private static <T extends PresortedComparable<T>> ImmutableSortedSet<T> rewriteItems(
-      Set<T> original, Function<T, T> rewrite) {
-    Builder<T> builder = new Builder<>(PresortedComparable::slowCompare);
-    for (T item : original) {
-      builder.add(rewrite.apply(item));
-    }
-    return builder.build();
-  }
-
-  private static <T extends PresortedComparable<T>>
-      SortedMap<T, ProgramMethodSet> rewriteKeysConservativelyWhileMergingValues(
-          Map<T, ProgramMethodSet> original, Function<T, Set<T>> rewrite) {
-    SortedMap<T, ProgramMethodSet> result = new TreeMap<>(PresortedComparable::slowCompare);
-    for (T item : original.keySet()) {
-      Set<T> rewrittenKeys = rewrite.apply(item);
-      for (T rewrittenKey : rewrittenKeys) {
-        result
-            .computeIfAbsent(rewrittenKey, k -> ProgramMethodSet.create())
-            .addAll(original.get(item));
-      }
-    }
+  private static SortedMap<DexMethod, ProgramMethodSet> rewriteInvokesWithContexts(
+      Map<DexMethod, ProgramMethodSet> invokes, GraphLense lens) {
+    SortedMap<DexMethod, ProgramMethodSet> result = new TreeMap<>(PresortedComparable::slowCompare);
+    invokes.forEach(
+        (method, contexts) ->
+            result
+                .computeIfAbsent(
+                    lens.getRenamedMethodSignature(method), ignore -> ProgramMethodSet.create())
+                .addAll(contexts));
     return Collections.unmodifiableSortedMap(result);
   }
 
@@ -998,54 +987,54 @@
             .filter(Objects::nonNull)
             .collect(Collectors.toList()));
 
+    assert lens.assertReferencesNotModified(pinnedItems);
+
     return new AppInfoWithLiveness(
         application,
         deadProtoTypes,
         missingTypes,
-        rewriteItems(liveTypes, lens::lookupType),
-        rewriteItems(instantiatedAppServices, lens::lookupType),
-        lens.rewriteMethodsConservatively(targetedMethods),
-        lens.rewriteMethodsConservatively(failedResolutionTargets),
-        lens.rewriteMethodsConservatively(bootstrapMethods),
-        lens.rewriteMethodsConservatively(methodsTargetedByInvokeDynamic),
-        lens.rewriteMethodsConservatively(virtualMethodsTargetedByInvokeDirect),
-        lens.rewriteMethodsConservatively(liveMethods),
+        lens.rewriteTypes(liveTypes),
+        lens.rewriteTypes(instantiatedAppServices),
+        lens.rewriteMethods(targetedMethods),
+        lens.rewriteMethods(failedResolutionTargets),
+        lens.rewriteMethods(bootstrapMethods),
+        lens.rewriteMethods(methodsTargetedByInvokeDynamic),
+        lens.rewriteMethods(virtualMethodsTargetedByInvokeDirect),
+        lens.rewriteMethods(liveMethods),
         fieldAccessInfoCollection.rewrittenWithLens(application, lens),
         objectAllocationInfoCollection.rewrittenWithLens(application, lens),
-        rewriteKeysConservativelyWhileMergingValues(
-            virtualInvokes, lens::lookupMethodInAllContexts),
-        rewriteKeysConservativelyWhileMergingValues(
-            interfaceInvokes, lens::lookupMethodInAllContexts),
-        rewriteKeysConservativelyWhileMergingValues(superInvokes, lens::lookupMethodInAllContexts),
-        rewriteKeysConservativelyWhileMergingValues(directInvokes, lens::lookupMethodInAllContexts),
-        rewriteKeysConservativelyWhileMergingValues(staticInvokes, lens::lookupMethodInAllContexts),
+        rewriteInvokesWithContexts(virtualInvokes, lens),
+        rewriteInvokesWithContexts(interfaceInvokes, lens),
+        rewriteInvokesWithContexts(superInvokes, lens),
+        rewriteInvokesWithContexts(directInvokes, lens),
+        rewriteInvokesWithContexts(staticInvokes, lens),
         // TODO(sgjesse): Rewrite call sites as well? Right now they are only used by minification
         //   after second tree shaking.
         callSites,
-        lens.rewriteReferencesConservatively(pinnedItems),
-        lens.rewriteReferencesConservatively(allowAccessModification),
-        rewriteReferenceKeys(mayHaveSideEffects, lens::lookupReference),
-        rewriteReferenceKeys(noSideEffects, lens::lookupReference),
-        rewriteReferenceKeys(assumedValues, lens::lookupReference),
-        lens.rewriteMethodsWithRenamedSignature(alwaysInline),
-        lens.rewriteMethodsWithRenamedSignature(forceInline),
-        lens.rewriteMethodsWithRenamedSignature(neverInline),
-        lens.rewriteMethodsWithRenamedSignature(whyAreYouNotInlining),
-        lens.rewriteMethodsWithRenamedSignature(keepConstantArguments),
-        lens.rewriteMethodsWithRenamedSignature(keepUnusedArguments),
-        lens.rewriteMethodsWithRenamedSignature(reprocess),
-        lens.rewriteMethodsWithRenamedSignature(neverReprocess),
+        pinnedItems,
+        lens.rewriteReferences(allowAccessModification),
+        lens.rewriteReferenceKeys(mayHaveSideEffects),
+        lens.rewriteReferenceKeys(noSideEffects),
+        lens.rewriteReferenceKeys(assumedValues),
+        lens.rewriteMethods(alwaysInline),
+        lens.rewriteMethods(forceInline),
+        lens.rewriteMethods(neverInline),
+        lens.rewriteMethods(whyAreYouNotInlining),
+        lens.rewriteMethods(keepConstantArguments),
+        lens.rewriteMethods(keepUnusedArguments),
+        lens.rewriteMethods(reprocess),
+        lens.rewriteMethods(neverReprocess),
         alwaysClassInline.rewriteItems(lens::lookupType),
-        rewriteItems(neverClassInline, lens::lookupType),
-        rewriteItems(neverMerge, lens::lookupType),
-        lens.rewriteReferencesConservatively(neverPropagateValue),
-        lens.rewriteReferencesConservatively(identifierNameStrings),
+        lens.rewriteTypes(neverClassInline),
+        lens.rewriteTypes(neverMerge),
+        lens.rewriteReferences(neverPropagateValue),
+        lens.rewriteReferenceKeys(identifierNameStrings),
         // Don't rewrite pruned types - the removed types are identified by their original name.
         prunedTypes,
-        rewriteReferenceKeys(switchMaps, lens::lookupField),
+        lens.rewriteFieldKeys(switchMaps),
         enumValueInfoMaps.rewrittenWithLens(lens),
-        rewriteItems(constClassReferences, lens::lookupType),
-        rewriteReferenceKeys(initClassReferences, lens::lookupType));
+        lens.rewriteTypes(constClassReferences),
+        lens.rewriteTypeKeys(initClassReferences));
   }
 
   /**
diff --git a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLivenessModifier.java b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLivenessModifier.java
index 32a8501..ae59f8b 100644
--- a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLivenessModifier.java
+++ b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLivenessModifier.java
@@ -6,6 +6,7 @@
 
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.DexReference;
 import com.android.tools.r8.graph.FieldAccessInfoCollectionImpl;
 import com.android.tools.r8.graph.FieldAccessInfoImpl;
 import com.google.common.collect.Sets;
@@ -15,6 +16,7 @@
 public class AppInfoWithLivenessModifier {
 
   private final Set<DexProgramClass> noLongerInstantiatedClasses = Sets.newConcurrentHashSet();
+  private final Set<DexReference> noLongerPinnedItems = Sets.newConcurrentHashSet();
   private final Set<DexField> noLongerWrittenFields = Sets.newConcurrentHashSet();
 
   AppInfoWithLivenessModifier() {}
@@ -27,6 +29,10 @@
     noLongerInstantiatedClasses.add(clazz);
   }
 
+  public void removePinnedClassMembers(DexProgramClass clazz) {
+    clazz.members().forEach(member -> noLongerPinnedItems.add(member.toReference()));
+  }
+
   public void removeWrittenField(DexField field) {
     noLongerWrittenFields.add(field);
   }
@@ -36,6 +42,8 @@
     noLongerInstantiatedClasses.forEach(appInfo::removeFromSingleTargetLookupCache);
     appInfo.mutateObjectAllocationInfoCollection(
         mutator -> noLongerInstantiatedClasses.forEach(mutator::markNoLongerInstantiated));
+    // Pinned items.
+    noLongerPinnedItems.forEach(appInfo::removePinnedItem);
     // Written fields.
     FieldAccessInfoCollectionImpl fieldAccessInfoCollection =
         appInfo.getMutableFieldAccessInfoCollection();
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 1202384..2a0a843 100644
--- a/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java
+++ b/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java
@@ -17,7 +17,6 @@
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Sets;
 import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.Set;
@@ -55,8 +54,6 @@
   private Set<DexMethod> mergedMethods;
   private final Map<DexMethod, DexMethod> originalMethodSignaturesForBridges;
 
-  private Map<DexMethod, Set<DexType>> contextsForContextSensitiveMethods;
-
   private VerticalClassMergerGraphLense(
       AppView<?> appView,
       Map<DexType, DexType> typeMap,
@@ -82,24 +79,6 @@
     this.originalMethodSignaturesForBridges = originalMethodSignaturesForBridges;
   }
 
-  public void initializeCacheForLookupMethodInAllContexts() {
-    assert contextsForContextSensitiveMethods == null;
-    contextsForContextSensitiveMethods = new IdentityHashMap<>();
-    contextualVirtualToDirectMethodMaps.forEach(
-        (type, virtualToDirectMethodMap) -> {
-          for (DexMethod method : virtualToDirectMethodMap.keySet()) {
-            contextsForContextSensitiveMethods
-                .computeIfAbsent(method, ignore -> Sets.newIdentityHashSet())
-                .add(type);
-          }
-        });
-  }
-
-  public void unsetCacheForLookupMethodInAllContexts() {
-    assert contextsForContextSensitiveMethods != null;
-    contextsForContextSensitiveMethods = null;
-  }
-
   @Override
   public DexType getOriginalType(DexType type) {
     return previousLense.getOriginalType(type);
@@ -145,22 +124,6 @@
   }
 
   @Override
-  public Set<DexMethod> lookupMethodInAllContexts(DexMethod method) {
-    assert contextsForContextSensitiveMethods != null;
-    ImmutableSet.Builder<DexMethod> builder = ImmutableSet.builder();
-    for (DexMethod previous : previousLense.lookupMethodInAllContexts(method)) {
-      builder.add(methodMap.getOrDefault(previous, previous));
-      Set<DexType> contexts = contextsForContextSensitiveMethods.get(previous);
-      if (contexts != null) {
-        for (DexType context : contexts) {
-          builder.add(contextualVirtualToDirectMethodMaps.get(context).get(previous).getMethod());
-        }
-      }
-    }
-    return builder.build();
-  }
-
-  @Override
   public boolean isContextFreeForMethods() {
     return contextualVirtualToDirectMethodMaps.isEmpty() && previousLense.isContextFreeForMethods();
   }
diff --git a/src/main/java/com/android/tools/r8/utils/SetUtils.java b/src/main/java/com/android/tools/r8/utils/SetUtils.java
index c961003..89880cc 100644
--- a/src/main/java/com/android/tools/r8/utils/SetUtils.java
+++ b/src/main/java/com/android/tools/r8/utils/SetUtils.java
@@ -24,6 +24,13 @@
     return result;
   }
 
+  public static <T> Set<T> newIdentityHashSet(Iterable<T> c1, Iterable<T> c2) {
+    Set<T> result = Sets.newIdentityHashSet();
+    c1.forEach(result::add);
+    c2.forEach(result::add);
+    return result;
+  }
+
   public static <T> Set<T> newIdentityHashSet(int capacity) {
     return Collections.newSetFromMap(new IdentityHashMap<>(capacity));
   }