Minor cleanup prior to definitionFor(DexField) removal

Change-Id: I845054a98954c96ca5db315af5d476c4c2dcf32b
Bug: 157616970
diff --git a/src/main/java/com/android/tools/r8/graph/DexField.java b/src/main/java/com/android/tools/r8/graph/DexField.java
index d10a2c8..31017af 100644
--- a/src/main/java/com/android/tools/r8/graph/DexField.java
+++ b/src/main/java/com/android/tools/r8/graph/DexField.java
@@ -6,6 +6,7 @@
 import com.android.tools.r8.dex.IndexedItemCollection;
 import com.android.tools.r8.errors.CompilationError;
 import com.android.tools.r8.naming.NamingLens;
+import java.util.function.BiConsumer;
 
 public class DexField extends DexMember<DexEncodedField, DexField> {
 
@@ -23,6 +24,15 @@
   }
 
   @Override
+  public <T> void apply(
+      BiConsumer<DexType, T> classConsumer,
+      BiConsumer<DexField, T> fieldConsumer,
+      BiConsumer<DexMethod, T> methodConsumer,
+      T arg) {
+    fieldConsumer.accept(this, arg);
+  }
+
+  @Override
   public int computeHashCode() {
     return holder.hashCode()
         + type.hashCode() * 7
diff --git a/src/main/java/com/android/tools/r8/graph/DexMethod.java b/src/main/java/com/android/tools/r8/graph/DexMethod.java
index d3dd5fc..89c8f05 100644
--- a/src/main/java/com/android/tools/r8/graph/DexMethod.java
+++ b/src/main/java/com/android/tools/r8/graph/DexMethod.java
@@ -11,6 +11,7 @@
 import com.android.tools.r8.references.TypeReference;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.BiConsumer;
 
 public class DexMethod extends DexMember<DexEncodedMethod, DexMethod> {
 
@@ -29,6 +30,15 @@
   }
 
   @Override
+  public <T> void apply(
+      BiConsumer<DexType, T> classConsumer,
+      BiConsumer<DexField, T> fieldConsumer,
+      BiConsumer<DexMethod, T> methodConsumer,
+      T arg) {
+    methodConsumer.accept(this, arg);
+  }
+
+  @Override
   public String toString() {
     return "Method " + holder + "." + name + " " + proto.toString();
   }
diff --git a/src/main/java/com/android/tools/r8/graph/DexReference.java b/src/main/java/com/android/tools/r8/graph/DexReference.java
index 42286db..3735378 100644
--- a/src/main/java/com/android/tools/r8/graph/DexReference.java
+++ b/src/main/java/com/android/tools/r8/graph/DexReference.java
@@ -3,7 +3,7 @@
 // BSD-style license that can be found in the LICENSE file.
 package com.android.tools.r8.graph;
 
-import java.util.Objects;
+import java.util.function.BiConsumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
@@ -13,6 +13,12 @@
  */
 public abstract class DexReference extends IndexedDexItem {
 
+  public abstract <T> void apply(
+      BiConsumer<DexType, T> classConsumer,
+      BiConsumer<DexField, T> fieldConsumer,
+      BiConsumer<DexMethod, T> methodConsumer,
+      T arg);
+
   public boolean isDexType() {
     return false;
   }
@@ -49,22 +55,6 @@
     return DexItem.filter(stream, DexReference.class);
   }
 
-  public DexDefinition toDefinition(AppInfo appInfo) {
-    if (isDexType()) {
-      return appInfo.definitionFor(asDexType());
-    } else if (isDexField()) {
-      return appInfo.definitionFor(asDexField());
-    } else {
-      assert isDexMethod();
-      return appInfo.definitionFor(asDexMethod());
-    }
-  }
-
-  public static Stream<DexDefinition> mapToDefinition(
-      Stream<DexReference> references, AppInfo appInfo) {
-    return references.map(r -> r.toDefinition(appInfo)).filter(Objects::nonNull);
-  }
-
   private static <T extends DexReference> Stream<T> filter(
       Stream<DexReference> stream,
       Predicate<DexReference> pred,
diff --git a/src/main/java/com/android/tools/r8/graph/DexType.java b/src/main/java/com/android/tools/r8/graph/DexType.java
index 01c5298..dfe07f6 100644
--- a/src/main/java/com/android/tools/r8/graph/DexType.java
+++ b/src/main/java/com/android/tools/r8/graph/DexType.java
@@ -24,14 +24,13 @@
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.InternalOptions.OutlineOptions;
-import com.android.tools.r8.utils.Pair;
 import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Sets;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
+import java.util.function.BiConsumer;
 import java.util.function.Predicate;
 
 public class DexType extends DexReference implements PresortedComparable<DexType> {
@@ -117,6 +116,15 @@
   }
 
   @Override
+  public <T> void apply(
+      BiConsumer<DexType, T> classConsumer,
+      BiConsumer<DexField, T> fieldConsumer,
+      BiConsumer<DexMethod, T> methodConsumer,
+      T arg) {
+    classConsumer.accept(this, arg);
+  }
+
+  @Override
   public String toSourceString() {
     if (toStringCache == null) {
       // TODO(ager): Pass in a ProguardMapReader to map names back to original names.
@@ -409,15 +417,4 @@
   public String getPackageName() {
     return DescriptorUtils.getPackageNameFromBinaryName(toBinaryName());
   }
-
-  public Pair<String, String> rewritingPrefixIn(Map<String, String> map) {
-    // TODO(b/134732760): Rewrite this to use descriptors and not Strings.
-    String javaClassName = this.toString();
-    for (String rewritePrefix : map.keySet()) {
-      if (javaClassName.startsWith(rewritePrefix)) {
-        return new Pair<>(rewritePrefix, map.get(rewritePrefix));
-      }
-    }
-    return null;
-  }
 }
diff --git a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
index 1bc26f2..faed477 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -94,7 +94,7 @@
 import com.android.tools.r8.shaking.GraphReporter.KeepReasonWitness;
 import com.android.tools.r8.shaking.KeepInfoCollection.MutableKeepInfoCollection;
 import com.android.tools.r8.shaking.RootSetBuilder.ConsequentRootSet;
-import com.android.tools.r8.shaking.RootSetBuilder.DependentItems;
+import com.android.tools.r8.shaking.RootSetBuilder.ItemsWithRules;
 import com.android.tools.r8.shaking.RootSetBuilder.RootSet;
 import com.android.tools.r8.shaking.ScopedDexMethodSet.AddMethodIfMoreVisibleResult;
 import com.android.tools.r8.utils.Action;
@@ -608,7 +608,7 @@
     items.entrySet().forEach(this::enqueueRootItem);
   }
 
-  private void enqueueRootItems(DependentItems items) {
+  private void enqueueRootItems(ItemsWithRules items) {
     items.forEachField(this::enqueueRootField);
     items.forEachMethod(this::enqueueRootMethod);
     items.forEachClass(this::enqueueRootClass);
diff --git a/src/main/java/com/android/tools/r8/shaking/RootSetBuilder.java b/src/main/java/com/android/tools/r8/shaking/RootSetBuilder.java
index f542984..1d85ff2 100644
--- a/src/main/java/com/android/tools/r8/shaking/RootSetBuilder.java
+++ b/src/main/java/com/android/tools/r8/shaking/RootSetBuilder.java
@@ -99,7 +99,7 @@
   private final Set<DexType> neverClassInline = Sets.newIdentityHashSet();
   private final Set<DexType> neverMerge = Sets.newIdentityHashSet();
   private final Set<DexReference> neverPropagateValue = Sets.newIdentityHashSet();
-  private final Map<DexReference, MutableDependentItems> dependentNoShrinking =
+  private final Map<DexReference, MutableItemsWithRules> dependentNoShrinking =
       new IdentityHashMap<>();
   private final Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule =
       new IdentityHashMap<>();
@@ -1009,8 +1009,8 @@
     }
     // Keep the type if the item is also kept.
     dependentNoShrinking
-        .computeIfAbsent(item.toReference(), x -> new MutableDependentItems())
-        .addDependentClass(type, context);
+        .computeIfAbsent(item.toReference(), x -> new MutableItemsWithRules())
+        .addClassWithRule(type, context);
     // Unconditionally add to no-obfuscation, as that is only checked for surviving items.
     noObfuscation.add(type);
   }
@@ -1097,8 +1097,8 @@
       if (!modifiers.allowsShrinking) {
         if (precondition != null) {
           dependentNoShrinking
-              .computeIfAbsent(precondition.toReference(), x -> new MutableDependentItems())
-              .addDependentItem(item.toReference(), keepRule);
+              .computeIfAbsent(precondition.toReference(), x -> new MutableItemsWithRules())
+              .addReferenceWithRule(item.toReference(), keepRule);
         } else {
           noShrinking.computeIfAbsent(item.toReference(), i -> new HashSet<>()).add(keepRule);
         }
@@ -1270,7 +1270,7 @@
     final Set<DexType> neverClassInline;
     final Map<DexReference, Set<ProguardKeepRuleBase>> noShrinking;
     final Set<DexReference> noObfuscation;
-    final Map<DexReference, MutableDependentItems> dependentNoShrinking;
+    final Map<DexReference, MutableItemsWithRules> dependentNoShrinking;
     final Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule;
     final List<DelayedRootSetActionItem> delayedRootSetActionItems;
 
@@ -1279,7 +1279,7 @@
         Set<DexType> neverClassInline,
         Map<DexReference, Set<ProguardKeepRuleBase>> noShrinking,
         Set<DexReference> noObfuscation,
-        Map<DexReference, MutableDependentItems> dependentNoShrinking,
+        Map<DexReference, MutableItemsWithRules> dependentNoShrinking,
         Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule,
         List<DelayedRootSetActionItem> delayedRootSetActionItems) {
       this.neverInline = neverInline;
@@ -1306,7 +1306,7 @@
 
     public void forEachMemberWithDependentItems(
         DexDefinitionSupplier definitions,
-        BiConsumer<DexEncodedMember<?, ?>, DependentItems> consumer) {
+        BiConsumer<DexEncodedMember<?, ?>, ItemsWithRules> consumer) {
       dependentNoShrinking.forEach(
           (reference, dependentItems) -> {
             if (reference.isDexMember()) {
@@ -1387,9 +1387,9 @@
           });
     }
 
-    DependentItems getDependentItems(DexDefinition item) {
-      DependentItems found = dependentNoShrinking.get(item.toReference());
-      return found != null ? found : DependentItems.empty();
+    ItemsWithRules getDependentItems(DexDefinition item) {
+      ItemsWithRules found = dependentNoShrinking.get(item.toReference());
+      return found != null ? found : ItemsWithRules.empty();
     }
 
     Set<ProguardKeepRuleBase> getDependentKeepClassCompatRule(DexType type) {
@@ -1397,93 +1397,89 @@
     }
   }
 
-  abstract static class DependentItems {
+  abstract static class ItemsWithRules {
 
-    public static DependentItems empty() {
-      return MutableDependentItems.EMPTY;
+    public static ItemsWithRules empty() {
+      return MutableItemsWithRules.EMPTY;
     }
 
     public abstract void forEachClass(BiConsumer<DexType, Set<ProguardKeepRuleBase>> consumer);
 
-    public abstract void forEachField(BiConsumer<DexField, Set<ProguardKeepRuleBase>> consumer);
+    public abstract void forEachField(
+        BiConsumer<? super DexField, Set<ProguardKeepRuleBase>> consumer);
 
     public abstract void forEachMember(
         BiConsumer<DexMember<?, ?>, Set<ProguardKeepRuleBase>> consumer);
 
-    public abstract void forEachMethod(BiConsumer<DexMethod, Set<ProguardKeepRuleBase>> consumer);
+    public abstract void forEachMethod(
+        BiConsumer<? super DexMethod, Set<ProguardKeepRuleBase>> consumer);
   }
 
-  static class MutableDependentItems extends DependentItems {
+  static class MutableItemsWithRules extends ItemsWithRules {
 
-    private static final DependentItems EMPTY =
-        new MutableDependentItems(
+    private static final ItemsWithRules EMPTY =
+        new MutableItemsWithRules(
             Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
 
-    final Map<DexType, Set<ProguardKeepRuleBase>> dependentClasses;
-    final Map<DexField, Set<ProguardKeepRuleBase>> dependentFields;
-    final Map<DexMethod, Set<ProguardKeepRuleBase>> dependentMethods;
+    final Map<DexType, Set<ProguardKeepRuleBase>> classesWithRules;
+    final Map<DexField, Set<ProguardKeepRuleBase>> fieldsWithRules;
+    final Map<DexMethod, Set<ProguardKeepRuleBase>> methodsWithRules;
 
-    MutableDependentItems() {
+    MutableItemsWithRules() {
       this(new IdentityHashMap<>(), new IdentityHashMap<>(), new IdentityHashMap<>());
     }
 
-    private MutableDependentItems(
-        Map<DexType, Set<ProguardKeepRuleBase>> dependentClasses,
-        Map<DexField, Set<ProguardKeepRuleBase>> dependentFields,
-        Map<DexMethod, Set<ProguardKeepRuleBase>> dependentMethods) {
-      this.dependentClasses = dependentClasses;
-      this.dependentFields = dependentFields;
-      this.dependentMethods = dependentMethods;
+    private MutableItemsWithRules(
+        Map<DexType, Set<ProguardKeepRuleBase>> classesWithRules,
+        Map<DexField, Set<ProguardKeepRuleBase>> fieldsWithRules,
+        Map<DexMethod, Set<ProguardKeepRuleBase>> methodsWithRules) {
+      this.classesWithRules = classesWithRules;
+      this.fieldsWithRules = fieldsWithRules;
+      this.methodsWithRules = methodsWithRules;
     }
 
-    public void addAll(DependentItems items) {
-      items.forEachClass(dependentClasses::put);
-      items.forEachField(dependentFields::put);
-      items.forEachMethod(dependentMethods::put);
+    public void addAll(ItemsWithRules items) {
+      items.forEachClass(classesWithRules::put);
+      items.forEachField(fieldsWithRules::put);
+      items.forEachMethod(methodsWithRules::put);
     }
 
-    public void addDependentItem(DexReference reference, ProguardKeepRuleBase rule) {
-      if (reference.isDexField()) {
-        addDependentField(reference.asDexField(), rule);
-      } else if (reference.isDexMethod()) {
-        addDependentMethod(reference.asDexMethod(), rule);
-      } else {
-        assert reference.isDexType();
-        addDependentClass(reference.asDexType(), rule);
-      }
+    public void addClassWithRule(DexType type, ProguardKeepRuleBase rule) {
+      classesWithRules.computeIfAbsent(type, ignore -> new HashSet<>()).add(rule);
     }
 
-    public void addDependentClass(DexType type, ProguardKeepRuleBase rule) {
-      dependentClasses.computeIfAbsent(type, ignore -> new HashSet<>()).add(rule);
+    public void addFieldWithRule(DexField field, ProguardKeepRuleBase rule) {
+      fieldsWithRules.computeIfAbsent(field, ignore -> new HashSet<>()).add(rule);
     }
 
-    public void addDependentField(DexField field, ProguardKeepRuleBase rule) {
-      dependentFields.computeIfAbsent(field, ignore -> new HashSet<>()).add(rule);
+    public void addMethodWithRule(DexMethod method, ProguardKeepRuleBase rule) {
+      methodsWithRules.computeIfAbsent(method, ignore -> new HashSet<>()).add(rule);
     }
 
-    public void addDependentMethod(DexMethod method, ProguardKeepRuleBase rule) {
-      dependentMethods.computeIfAbsent(method, ignore -> new HashSet<>()).add(rule);
+    public void addReferenceWithRule(DexReference reference, ProguardKeepRuleBase rule) {
+      reference.apply(
+          this::addClassWithRule, this::addFieldWithRule, this::addMethodWithRule, rule);
     }
 
     @Override
     public void forEachClass(BiConsumer<DexType, Set<ProguardKeepRuleBase>> consumer) {
-      dependentClasses.forEach(consumer);
+      classesWithRules.forEach(consumer);
     }
 
     @Override
-    public void forEachField(BiConsumer<DexField, Set<ProguardKeepRuleBase>> consumer) {
-      dependentFields.forEach(consumer);
+    public void forEachField(BiConsumer<? super DexField, Set<ProguardKeepRuleBase>> consumer) {
+      fieldsWithRules.forEach(consumer);
     }
 
     @Override
     public void forEachMember(BiConsumer<DexMember<?, ?>, Set<ProguardKeepRuleBase>> consumer) {
-      dependentFields.forEach(consumer);
-      dependentMethods.forEach(consumer);
+      forEachField(consumer);
+      forEachMethod(consumer);
     }
 
     @Override
-    public void forEachMethod(BiConsumer<DexMethod, Set<ProguardKeepRuleBase>> consumer) {
-      dependentMethods.forEach(consumer);
+    public void forEachMethod(BiConsumer<? super DexMethod, Set<ProguardKeepRuleBase>> consumer) {
+      methodsWithRules.forEach(consumer);
     }
   }
 
@@ -1529,7 +1525,7 @@
         Map<DexReference, ProguardMemberRule> mayHaveSideEffects,
         Map<DexReference, ProguardMemberRule> noSideEffects,
         Map<DexReference, ProguardMemberRule> assumedValues,
-        Map<DexReference, MutableDependentItems> dependentNoShrinking,
+        Map<DexReference, MutableItemsWithRules> dependentNoShrinking,
         Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule,
         Set<DexReference> identifierNameStrings,
         Set<ProguardIfRule> ifRules,
@@ -1595,11 +1591,11 @@
     }
 
     // Add dependent items that depend on -if rules.
-    private void addDependentItems(Map<DexReference, ? extends DependentItems> dependentItems) {
+    private void addDependentItems(Map<DexReference, ? extends ItemsWithRules> dependentItems) {
       dependentItems.forEach(
           (reference, dependence) ->
               dependentNoShrinking
-                  .computeIfAbsent(reference, x -> new MutableDependentItems())
+                  .computeIfAbsent(reference, x -> new MutableItemsWithRules())
                   .addAll(dependence));
     }
 
@@ -1863,7 +1859,7 @@
         Set<DexType> neverClassInline,
         Map<DexReference, Set<ProguardKeepRuleBase>> noShrinking,
         Set<DexReference> noObfuscation,
-        Map<DexReference, MutableDependentItems> dependentNoShrinking,
+        Map<DexReference, MutableItemsWithRules> dependentNoShrinking,
         Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule,
         List<DelayedRootSetActionItem> delayedRootSetActionItems) {
       super(