Remove all uses of definitionFor(DexReference)

Bug: 157616970
Change-Id: I9e225f441bf1fb605d5b5ce48ca83a37ba7c1318
diff --git a/src/main/java/com/android/tools/r8/graph/AppInfo.java b/src/main/java/com/android/tools/r8/graph/AppInfo.java
index 7931476..e49f177 100644
--- a/src/main/java/com/android/tools/r8/graph/AppInfo.java
+++ b/src/main/java/com/android/tools/r8/graph/AppInfo.java
@@ -130,20 +130,6 @@
     return app.classesWithDeterministicOrder();
   }
 
-  @Deprecated
-  @Override
-  public DexDefinition definitionFor(DexReference reference) {
-    assert checkIfObsolete();
-    if (reference.isDexType()) {
-      return definitionFor(reference.asDexType());
-    }
-    if (reference.isDexMethod()) {
-      return definitionFor(reference.asDexMethod());
-    }
-    assert reference.isDexField();
-    return definitionFor(reference.asDexField());
-  }
-
   @Override
   public DexClass definitionFor(DexType type) {
     return definitionForWithoutExistenceAssert(type);
diff --git a/src/main/java/com/android/tools/r8/graph/AppView.java b/src/main/java/com/android/tools/r8/graph/AppView.java
index 99a1803..33c63dc 100644
--- a/src/main/java/com/android/tools/r8/graph/AppView.java
+++ b/src/main/java/com/android/tools/r8/graph/AppView.java
@@ -228,12 +228,6 @@
 
   @Deprecated
   @Override
-  public final DexDefinition definitionFor(DexReference reference) {
-    return appInfo().definitionFor(reference);
-  }
-
-  @Deprecated
-  @Override
   public final DexEncodedField definitionFor(DexField field) {
     return appInfo().definitionFor(field);
   }
diff --git a/src/main/java/com/android/tools/r8/graph/DexClass.java b/src/main/java/com/android/tools/r8/graph/DexClass.java
index c2cdc32..5e18e05 100644
--- a/src/main/java/com/android/tools/r8/graph/DexClass.java
+++ b/src/main/java/com/android/tools/r8/graph/DexClass.java
@@ -406,6 +406,15 @@
     return methodCollection.getVirtualMethod(predicate);
   }
 
+  /** Find member in this class matching {@param member}. */
+  @SuppressWarnings("unchecked")
+  public <D extends DexEncodedMember<D, R>, R extends DexMember<D, R>> D lookupMember(
+      DexMember<D, R> member) {
+    DexEncodedMember<?, ?> definition =
+        member.isDexField() ? lookupField(member.asDexField()) : lookupMethod(member.asDexMethod());
+    return (D) definition;
+  }
+
   /** Find method in this class matching {@param method}. */
   public DexEncodedMethod lookupMethod(DexMethod method) {
     return methodCollection.getMethod(method);
diff --git a/src/main/java/com/android/tools/r8/graph/DexDefinitionSupplier.java b/src/main/java/com/android/tools/r8/graph/DexDefinitionSupplier.java
index a7a1b1e..18d1f77 100644
--- a/src/main/java/com/android/tools/r8/graph/DexDefinitionSupplier.java
+++ b/src/main/java/com/android/tools/r8/graph/DexDefinitionSupplier.java
@@ -7,9 +7,6 @@
 public interface DexDefinitionSupplier {
 
   @Deprecated
-  DexDefinition definitionFor(DexReference reference);
-
-  @Deprecated
   DexEncodedField definitionFor(DexField field);
 
   @Deprecated
@@ -19,7 +16,11 @@
   @SuppressWarnings("unchecked")
   default <D extends DexEncodedMember<D, R>, R extends DexMember<D, R>>
       DexEncodedMember<D, R> definitionFor(DexMember<D, R> member) {
-    return (DexEncodedMember<D, R>) definitionFor((DexReference) member);
+    if (member.isDexField()) {
+      return (DexEncodedMember<D, R>) definitionFor(member.asDexField());
+    }
+    assert member.isDexMethod();
+    return (DexEncodedMember<D, R>) definitionFor(member.asDexMethod());
   }
 
   DexClass definitionFor(DexType type);
diff --git a/src/main/java/com/android/tools/r8/graph/DirectMappedDexApplication.java b/src/main/java/com/android/tools/r8/graph/DirectMappedDexApplication.java
index 60a1c49..d66ef7e 100644
--- a/src/main/java/com/android/tools/r8/graph/DirectMappedDexApplication.java
+++ b/src/main/java/com/android/tools/r8/graph/DirectMappedDexApplication.java
@@ -75,19 +75,6 @@
 
   @Deprecated
   @Override
-  public DexDefinition definitionFor(DexReference reference) {
-    if (reference.isDexType()) {
-      return definitionFor(reference.asDexType());
-    }
-    if (reference.isDexMethod()) {
-      return definitionFor(reference.asDexMethod());
-    }
-    assert reference.isDexField();
-    return definitionFor(reference.asDexField());
-  }
-
-  @Deprecated
-  @Override
   public DexEncodedField definitionFor(DexField field) {
     DexClass clazz = definitionFor(field.holder);
     return clazz != null ? clazz.lookupField(field) : 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 ffb45b6..26ed5d0 100644
--- a/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
+++ b/src/main/java/com/android/tools/r8/shaking/Enqueuer.java
@@ -94,6 +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.RootSet;
 import com.android.tools.r8.shaking.ScopedDexMethodSet.AddMethodIfMoreVisibleResult;
 import com.android.tools.r8.utils.Action;
@@ -503,17 +504,6 @@
     recordTypeReference(field.type);
   }
 
-  public DexDefinition definitionFor(DexReference reference) {
-    if (reference.isDexType()) {
-      return definitionFor(reference.asDexType());
-    } else if (reference.isDexMethod()) {
-      return definitionFor(reference.asDexMethod());
-    } else {
-      assert reference.isDexField();
-      return definitionFor(reference.asDexField());
-    }
-  }
-
   public DexEncodedField definitionFor(DexField field) {
     DexClass clazz = definitionFor(field.holder);
     if (clazz == null) {
@@ -614,24 +604,98 @@
     }
   }
 
-  private static <T> SetWithReason<T> newSetWithoutReasonReporter() {
-    return new SetWithReason<>((f, r) -> {});
-  }
-
   private void enqueueRootItems(Map<DexReference, Set<ProguardKeepRuleBase>> items) {
     items.entrySet().forEach(this::enqueueRootItem);
   }
 
+  private void enqueueRootItems(DependentItems items) {
+    items.forEachField(this::enqueueRootField);
+    items.forEachMethod(this::enqueueRootMethod);
+    items.forEachClass(this::enqueueRootClass);
+  }
+
   private void enqueueRootItem(Entry<DexReference, Set<ProguardKeepRuleBase>> root) {
-    DexDefinition item = appView.definitionFor(root.getKey());
-    if (item != null) {
-      enqueueRootItem(item, root.getValue());
+    DexReference reference = root.getKey();
+    Set<ProguardKeepRuleBase> rules = root.getValue();
+    if (reference.isDexField()) {
+      enqueueRootField(reference.asDexField(), rules);
+    } else if (reference.isDexMethod()) {
+      enqueueRootMethod(reference.asDexMethod(), rules);
+    } else if (reference.isDexType()) {
+      enqueueRootClass(reference.asDexType(), rules);
     } else {
-      // TODO(b/123923324): Verify that root items are present.
-      // assert false : "Expected root item `" + root.getKey().toSourceString() + "` to be present";
+      throw new Unreachable();
     }
   }
 
+  // TODO(b/123923324): Verify that root items are present.
+  private void enqueueRootClass(DexType type, Set<ProguardKeepRuleBase> rules) {
+    DexProgramClass clazz = asProgramClassOrNull(appView.definitionFor(type));
+    if (clazz != null) {
+      enqueueRootClass(clazz, rules, null);
+    }
+  }
+
+  private void enqueueRootClass(
+      DexProgramClass clazz, Set<ProguardKeepRuleBase> rules, DexDefinition precondition) {
+    KeepReasonWitness witness = graphReporter.reportKeepClass(precondition, rules, clazz);
+    keepClassWithRules(clazz, rules);
+    if (clazz.isAnnotation()) {
+      workList.enqueueMarkAnnotationInstantiatedAction(clazz, witness);
+    } else if (clazz.isInterface()) {
+      workList.enqueueMarkInterfaceInstantiatedAction(clazz, witness);
+    } else {
+      workList.enqueueMarkInstantiatedAction(clazz, null, InstantiationReason.KEEP_RULE, witness);
+      if (clazz.hasDefaultInitializer()) {
+        ProgramMethod defaultInitializer = clazz.getProgramDefaultInitializer();
+        if (forceProguardCompatibility) {
+          workList.enqueueMarkMethodKeptAction(
+              defaultInitializer,
+              graphReporter.reportCompatKeepDefaultInitializer(defaultInitializer));
+        }
+        if (clazz.isExternalizable(appView)) {
+          enqueueMarkMethodLiveAction(defaultInitializer, witness);
+        }
+      }
+    }
+  }
+
+  // TODO(b/123923324): Verify that root items are present.
+  private void enqueueRootField(DexField reference, Set<ProguardKeepRuleBase> rules) {
+    DexProgramClass holder = getProgramClassOrNull(reference.holder);
+    if (holder != null) {
+      ProgramField field = holder.lookupProgramField(reference);
+      if (field != null) {
+        enqueueRootField(field, rules, null);
+      }
+    }
+  }
+
+  private void enqueueRootField(
+      ProgramField field, Set<ProguardKeepRuleBase> rules, DexDefinition precondition) {
+    keepFieldWithRules(field.getHolder(), field.getDefinition(), rules);
+    workList.enqueueMarkFieldKeptAction(
+        field, graphReporter.reportKeepField(precondition, rules, field.getDefinition()));
+  }
+
+  // TODO(b/123923324): Verify that root items are present.
+  private void enqueueRootMethod(DexMethod reference, Set<ProguardKeepRuleBase> rules) {
+    DexProgramClass holder = getProgramClassOrNull(reference.holder);
+    if (holder != null) {
+      ProgramMethod method = holder.lookupProgramMethod(reference);
+      if (method != null) {
+        enqueueRootMethod(method, rules, null);
+      }
+    }
+  }
+
+  private void enqueueRootMethod(
+      ProgramMethod method, Set<ProguardKeepRuleBase> rules, DexDefinition precondition) {
+    keepMethodWithRules(method.getHolder(), method.getDefinition(), rules);
+    workList.enqueueMarkMethodKeptAction(
+        method, graphReporter.reportKeepMethod(precondition, rules, method.getDefinition()));
+  }
+
   private void enqueueRootItem(DexDefinition item, Set<ProguardKeepRuleBase> rules) {
     internalEnqueueRootItem(item, rules, null);
   }
@@ -640,43 +704,20 @@
       DexDefinition item, Set<ProguardKeepRuleBase> rules, DexDefinition precondition) {
     if (item.isDexClass()) {
       DexProgramClass clazz = item.asDexClass().asProgramClass();
-      KeepReasonWitness witness = graphReporter.reportKeepClass(precondition, rules, clazz);
-      keepClassWithRules(clazz, rules);
-      if (clazz.isAnnotation()) {
-        workList.enqueueMarkAnnotationInstantiatedAction(clazz, witness);
-      } else if (clazz.isInterface()) {
-        workList.enqueueMarkInterfaceInstantiatedAction(clazz, witness);
-      } else {
-        workList.enqueueMarkInstantiatedAction(clazz, null, InstantiationReason.KEEP_RULE, witness);
-        if (clazz.hasDefaultInitializer()) {
-          ProgramMethod defaultInitializer = clazz.getProgramDefaultInitializer();
-          if (forceProguardCompatibility) {
-            workList.enqueueMarkMethodKeptAction(
-                defaultInitializer,
-                graphReporter.reportCompatKeepDefaultInitializer(defaultInitializer));
-          }
-          if (clazz.isExternalizable(appView)) {
-            enqueueMarkMethodLiveAction(defaultInitializer, witness);
-          }
-        }
+      if (clazz != null) {
+        enqueueRootClass(clazz, rules, precondition);
       }
     } else if (item.isDexEncodedField()) {
       DexEncodedField field = item.asDexEncodedField();
       DexProgramClass holder = getProgramClassOrNull(field.holder());
       if (holder != null) {
-        keepFieldWithRules(holder, field, rules);
-        workList.enqueueMarkFieldKeptAction(
-            new ProgramField(holder, field),
-            graphReporter.reportKeepField(precondition, rules, field));
+        enqueueRootField(new ProgramField(holder, field), rules, precondition);
       }
     } else if (item.isDexEncodedMethod()) {
-      DexEncodedMethod encodedMethod = item.asDexEncodedMethod();
-      DexProgramClass holder = getProgramClassOrNull(encodedMethod.holder());
+      DexEncodedMethod method = item.asDexEncodedMethod();
+      DexProgramClass holder = getProgramClassOrNull(method.holder());
       if (holder != null) {
-        keepMethodWithRules(holder, encodedMethod, rules);
-        workList.enqueueMarkMethodKeptAction(
-            new ProgramMethod(holder, encodedMethod),
-            graphReporter.reportKeepMethod(precondition, rules, encodedMethod));
+        enqueueRootMethod(new ProgramMethod(holder, method), rules, precondition);
       }
     } else {
       throw new IllegalArgumentException(item.toString());
@@ -1621,8 +1662,8 @@
   }
 
   private void enqueueHolderWithDependentInstanceConstructor(
-      DexProgramClass clazz, ProgramMethod instanceInitializer, Set<ProguardKeepRuleBase> reasons) {
-    enqueueRootItem(clazz, reasons);
+      ProgramMethod instanceInitializer, Set<ProguardKeepRuleBase> reasons) {
+    enqueueRootItem(instanceInitializer.getHolder(), reasons);
   }
 
   private void processAnnotations(DexProgramClass holder, DexDefinition annotatedItem) {
@@ -4129,12 +4170,6 @@
 
     @Deprecated
     @Override
-    public DexDefinition definitionFor(DexReference reference) {
-      return enqueuer.definitionFor(reference);
-    }
-
-    @Deprecated
-    @Override
     public DexEncodedField definitionFor(DexField field) {
       return enqueuer.definitionFor(field);
     }
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 859e313..8be43f4 100644
--- a/src/main/java/com/android/tools/r8/shaking/RootSetBuilder.java
+++ b/src/main/java/com/android/tools/r8/shaking/RootSetBuilder.java
@@ -71,6 +71,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -98,8 +99,8 @@
   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, Map<DexReference, Set<ProguardKeepRuleBase>>>
-      dependentNoShrinking = new IdentityHashMap<>();
+  private final Map<DexReference, MutableDependentItems> dependentNoShrinking =
+      new IdentityHashMap<>();
   private final Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule =
       new IdentityHashMap<>();
   private final Map<DexReference, ProguardMemberRule> mayHaveSideEffects = new IdentityHashMap<>();
@@ -1008,9 +1009,8 @@
     }
     // Keep the type if the item is also kept.
     dependentNoShrinking
-        .computeIfAbsent(item.toReference(), x -> new IdentityHashMap<>())
-        .computeIfAbsent(type, k -> new HashSet<>())
-        .add(context);
+        .computeIfAbsent(item.toReference(), x -> new MutableDependentItems())
+        .addDependentClass(type, context);
     // Unconditionally add to no-obfuscation, as that is only checked for surviving items.
     noObfuscation.add(type);
   }
@@ -1097,9 +1097,8 @@
       if (!modifiers.allowsShrinking) {
         if (precondition != null) {
           dependentNoShrinking
-              .computeIfAbsent(precondition.toReference(), x -> new IdentityHashMap<>())
-              .computeIfAbsent(item.toReference(), i -> new HashSet<>())
-              .add(keepRule);
+              .computeIfAbsent(precondition.toReference(), x -> new MutableDependentItems())
+              .addDependentItem(item.toReference(), keepRule);
         } else {
           noShrinking.computeIfAbsent(item.toReference(), i -> new HashSet<>()).add(keepRule);
         }
@@ -1271,7 +1270,7 @@
     final Set<DexType> neverClassInline;
     final Map<DexReference, Set<ProguardKeepRuleBase>> noShrinking;
     final Set<DexReference> noObfuscation;
-    final Map<DexReference, Map<DexReference, Set<ProguardKeepRuleBase>>> dependentNoShrinking;
+    final Map<DexReference, MutableDependentItems> dependentNoShrinking;
     final Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule;
     final List<DelayedRootSetActionItem> delayedRootSetActionItems;
 
@@ -1280,7 +1279,7 @@
         Set<DexType> neverClassInline,
         Map<DexReference, Set<ProguardKeepRuleBase>> noShrinking,
         Set<DexReference> noObfuscation,
-        Map<DexReference, Map<DexReference, Set<ProguardKeepRuleBase>>> dependentNoShrinking,
+        Map<DexReference, MutableDependentItems> dependentNoShrinking,
         Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule,
         List<DelayedRootSetActionItem> delayedRootSetActionItems) {
       this.neverInline = neverInline;
@@ -1292,10 +1291,6 @@
       this.delayedRootSetActionItems = delayedRootSetActionItems;
     }
 
-    public boolean noShrinking(DexReference reference) {
-      return noShrinking.containsKey(reference);
-    }
-
     public void forEachClassWithDependentItems(
         DexDefinitionSupplier definitions, Consumer<DexProgramClass> consumer) {
       for (DexReference reference : dependentNoShrinking.keySet()) {
@@ -1324,19 +1319,34 @@
     public void forEachDependentInstanceConstructor(
         DexProgramClass clazz,
         AppView<?> appView,
-        Consumer3<DexProgramClass, ProgramMethod, Set<ProguardKeepRuleBase>> fn) {
+        BiConsumer<ProgramMethod, Set<ProguardKeepRuleBase>> fn) {
       getDependentItems(clazz)
-          .forEach(
+          .forEachMethod(
               (reference, reasons) -> {
-                if (reference.isDexMethod()) {
-                  DexMethod methodReference = reference.asDexMethod();
-                  DexProgramClass holder =
-                      asProgramClassOrNull(appView.definitionForHolder(methodReference));
-                  if (holder != null) {
-                    ProgramMethod method = holder.lookupProgramMethod(methodReference);
-                    if (method != null && method.getDefinition().isInstanceInitializer()) {
-                      fn.accept(clazz, method, reasons);
-                    }
+                DexProgramClass holder =
+                    asProgramClassOrNull(appView.definitionForHolder(reference));
+                if (holder != null) {
+                  ProgramMethod method = holder.lookupProgramMethod(reference);
+                  if (method != null && method.getDefinition().isInstanceInitializer()) {
+                    fn.accept(method, reasons);
+                  }
+                }
+              });
+    }
+
+    public void forEachDependentMember(
+        DexDefinition item,
+        AppView<?> appView,
+        Consumer3<DexDefinition, DexDefinition, Set<ProguardKeepRuleBase>> fn) {
+      getDependentItems(item)
+          .forEachMember(
+              (reference, reasons) -> {
+                DexProgramClass holder =
+                    asProgramClassOrNull(appView.definitionForHolder(reference));
+                if (holder != null) {
+                  DexEncodedMember<?, ?> member = holder.lookupMember(reference);
+                  if (member != null) {
+                    fn.accept(item, member, reasons);
                   }
                 }
               });
@@ -1346,35 +1356,33 @@
         DexDefinition item,
         AppView<?> appView,
         Consumer3<DexDefinition, DexDefinition, Set<ProguardKeepRuleBase>> fn) {
-      getDependentItems(item)
-          .forEach(
-              (reference, reasons) -> {
-                DexDefinition definition = appView.definitionFor(reference);
-                if (definition != null
-                    && !definition.isDexClass()
-                    && !definition.isStaticMember()) {
-                  fn.accept(item, definition, reasons);
-                }
-              });
+      forEachDependentMember(
+          item,
+          appView,
+          (precondition, member, reasons) -> {
+            if (!member.isStatic()) {
+              fn.accept(precondition, member, reasons);
+            }
+          });
     }
 
     public void forEachDependentStaticMember(
         DexDefinition item,
         AppView<?> appView,
         Consumer3<DexDefinition, DexDefinition, Set<ProguardKeepRuleBase>> fn) {
-      getDependentItems(item)
-          .forEach(
-              (reference, reasons) -> {
-                DexDefinition definition = appView.definitionFor(reference);
-                if (definition != null && !definition.isDexClass() && definition.isStaticMember()) {
-                  fn.accept(item, definition, reasons);
-                }
-              });
+      forEachDependentMember(
+          item,
+          appView,
+          (precondition, member, reasons) -> {
+            if (member.isStatic()) {
+              fn.accept(precondition, member, reasons);
+            }
+          });
     }
 
-    Map<DexReference, Set<ProguardKeepRuleBase>> getDependentItems(DexDefinition item) {
-      return Collections.unmodifiableMap(
-          dependentNoShrinking.getOrDefault(item.toReference(), Collections.emptyMap()));
+    DependentItems getDependentItems(DexDefinition item) {
+      DependentItems found = dependentNoShrinking.get(item.toReference());
+      return found != null ? found : DependentItems.empty();
     }
 
     Set<ProguardKeepRuleBase> getDependentKeepClassCompatRule(DexType type) {
@@ -1382,6 +1390,96 @@
     }
   }
 
+  abstract static class DependentItems {
+
+    public static DependentItems empty() {
+      return MutableDependentItems.EMPTY;
+    }
+
+    public abstract void forEachClass(BiConsumer<DexType, Set<ProguardKeepRuleBase>> consumer);
+
+    public abstract void forEachField(BiConsumer<DexField, Set<ProguardKeepRuleBase>> consumer);
+
+    public abstract void forEachMember(
+        BiConsumer<DexMember<?, ?>, Set<ProguardKeepRuleBase>> consumer);
+
+    public abstract void forEachMethod(BiConsumer<DexMethod, Set<ProguardKeepRuleBase>> consumer);
+  }
+
+  static class MutableDependentItems extends DependentItems {
+
+    private static final DependentItems EMPTY =
+        new MutableDependentItems(
+            Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
+
+    final Map<DexType, Set<ProguardKeepRuleBase>> dependentClasses;
+    final Map<DexField, Set<ProguardKeepRuleBase>> dependentFields;
+    final Map<DexMethod, Set<ProguardKeepRuleBase>> dependentMethods;
+
+    MutableDependentItems() {
+      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;
+    }
+
+    public void addAll(DependentItems items) {
+      items.forEachClass(dependentClasses::put);
+      items.forEachField(dependentFields::put);
+      items.forEachMethod(dependentMethods::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 addDependentClass(DexType type, ProguardKeepRuleBase rule) {
+      dependentClasses.computeIfAbsent(type, ignore -> new HashSet<>()).add(rule);
+    }
+
+    public void addDependentField(DexField field, ProguardKeepRuleBase rule) {
+      dependentFields.computeIfAbsent(field, ignore -> new HashSet<>()).add(rule);
+    }
+
+    public void addDependentMethod(DexMethod method, ProguardKeepRuleBase rule) {
+      dependentMethods.computeIfAbsent(method, ignore -> new HashSet<>()).add(rule);
+    }
+
+    @Override
+    public void forEachClass(BiConsumer<DexType, Set<ProguardKeepRuleBase>> consumer) {
+      dependentClasses.forEach(consumer);
+    }
+
+    @Override
+    public void forEachField(BiConsumer<DexField, Set<ProguardKeepRuleBase>> consumer) {
+      dependentFields.forEach(consumer);
+    }
+
+    @Override
+    public void forEachMember(BiConsumer<DexMember<?, ?>, Set<ProguardKeepRuleBase>> consumer) {
+      dependentFields.forEach(consumer);
+      dependentMethods.forEach(consumer);
+    }
+
+    @Override
+    public void forEachMethod(BiConsumer<DexMethod, Set<ProguardKeepRuleBase>> consumer) {
+      dependentMethods.forEach(consumer);
+    }
+  }
+
   public static class RootSet extends RootSetBase {
 
     public final ImmutableList<DexReference> reasonAsked;
@@ -1424,7 +1522,7 @@
         Map<DexReference, ProguardMemberRule> mayHaveSideEffects,
         Map<DexReference, ProguardMemberRule> noSideEffects,
         Map<DexReference, ProguardMemberRule> assumedValues,
-        Map<DexReference, Map<DexReference, Set<ProguardKeepRuleBase>>> dependentNoShrinking,
+        Map<DexReference, MutableDependentItems> dependentNoShrinking,
         Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule,
         Set<DexReference> identifierNameStrings,
         Set<ProguardIfRule> ifRules,
@@ -1490,13 +1588,12 @@
     }
 
     // Add dependent items that depend on -if rules.
-    private void addDependentItems(
-        Map<DexReference, Map<DexReference, Set<ProguardKeepRuleBase>>> dependentItems) {
+    private void addDependentItems(Map<DexReference, ? extends DependentItems> dependentItems) {
       dependentItems.forEach(
           (reference, dependence) ->
               dependentNoShrinking
-                  .computeIfAbsent(reference, x -> new IdentityHashMap<>())
-                  .putAll(dependence));
+                  .computeIfAbsent(reference, x -> new MutableDependentItems())
+                  .addAll(dependence));
     }
 
     public void copy(DexReference original, DexReference rewritten) {
@@ -1759,7 +1856,7 @@
         Set<DexType> neverClassInline,
         Map<DexReference, Set<ProguardKeepRuleBase>> noShrinking,
         Set<DexReference> noObfuscation,
-        Map<DexReference, Map<DexReference, Set<ProguardKeepRuleBase>>> dependentNoShrinking,
+        Map<DexReference, MutableDependentItems> dependentNoShrinking,
         Map<DexType, Set<ProguardKeepRuleBase>> dependentKeepClassCompatRule,
         List<DelayedRootSetActionItem> delayedRootSetActionItems) {
       super(