Remove rewritten collection building in prefix lens.

This CL lets the prefix mapper define the canonical map of rewritten
types and avoids the code that would iterate the item factory strings.

Bug: 150736225

Change-Id: I81bc3841fbc8752766aee7e498aabdea61ac56d3
diff --git a/src/main/java/com/android/tools/r8/ir/desugar/PrefixRewritingMapper.java b/src/main/java/com/android/tools/r8/ir/desugar/PrefixRewritingMapper.java
index 1c3b6a2..6f5c3f2 100644
--- a/src/main/java/com/android/tools/r8/ir/desugar/PrefixRewritingMapper.java
+++ b/src/main/java/com/android/tools/r8/ir/desugar/PrefixRewritingMapper.java
@@ -18,6 +18,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Consumer;
 
 public abstract class PrefixRewritingMapper {
 
@@ -47,6 +48,8 @@
 
   public abstract boolean isRewriting();
 
+  public abstract void forAllRewrittenTypes(Consumer<DexType> consumer);
+
   public static class DesugarPrefixRewritingMapper extends PrefixRewritingMapper {
 
     private final Set<DexType> notRewritten = Sets.newConcurrentHashSet();
@@ -70,6 +73,11 @@
       return factory.createString("L" + DescriptorUtils.getBinaryNameFromJavaType(prefix));
     }
 
+    @Override
+    public void forAllRewrittenTypes(Consumer<DexType> consumer) {
+      rewritten.keySet().forEach(consumer);
+    }
+
     private void validatePrefixes(Map<String, String> initialPrefixes) {
       String[] prefixes = initialPrefixes.keySet().toArray(new String[0]);
       for (int i = 0; i < prefixes.length; i++) {
@@ -192,5 +200,8 @@
     public boolean isRewriting() {
       return false;
     }
+
+    @Override
+    public void forAllRewrittenTypes(Consumer<DexType> consumer) {}
   }
 }
diff --git a/src/main/java/com/android/tools/r8/naming/PrefixRewritingNamingLens.java b/src/main/java/com/android/tools/r8/naming/PrefixRewritingNamingLens.java
index 777d006..6c78e0b 100644
--- a/src/main/java/com/android/tools/r8/naming/PrefixRewritingNamingLens.java
+++ b/src/main/java/com/android/tools/r8/naming/PrefixRewritingNamingLens.java
@@ -9,15 +9,13 @@
 import com.android.tools.r8.graph.DexCallSite;
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexItem;
-import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexString;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.InnerClassAttribute;
-import com.android.tools.r8.ir.desugar.PrefixRewritingMapper;
 import com.android.tools.r8.utils.InternalOptions;
-import com.google.common.collect.ImmutableMap;
-import java.util.IdentityHashMap;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -28,9 +26,9 @@
 // Naming lens for rewriting type prefixes.
 public class PrefixRewritingNamingLens extends NamingLens {
 
-  final Map<DexType, DexString> classRenaming = new IdentityHashMap<>();
   final NamingLens namingLens;
   final InternalOptions options;
+  final AppView<?> appView;
 
   public static NamingLens createPrefixRewritingNamingLens(AppView<?> appView) {
     return createPrefixRewritingNamingLens(appView, NamingLens.getIdentityLens());
@@ -45,18 +43,21 @@
   }
 
   public PrefixRewritingNamingLens(NamingLens namingLens, AppView<?> appView) {
+    this.appView = appView;
     this.namingLens = namingLens;
     this.options = appView.options();
-    DexItemFactory itemFactory = options.itemFactory;
-    PrefixRewritingMapper rewritePrefix = appView.rewritePrefix;
-    itemFactory.forAllTypes(
-        type -> {
-          if (rewritePrefix.hasRewrittenType(type, appView)) {
-            classRenaming.put(type, rewritePrefix.rewrittenType(type, appView).descriptor);
-          }
-        });
-    // Verify that no type would have been renamed by both lenses.
-    assert namingLens.verifyNoOverlap(classRenaming);
+  }
+
+  private boolean isRenamed(DexType type) {
+    return getRenaming(type) != null;
+  }
+
+  private DexString getRenaming(DexType type) {
+    DexString descriptor = null;
+    if (appView.rewritePrefix.hasRewrittenType(type, appView)) {
+      descriptor = appView.rewritePrefix.rewrittenType(type, appView).descriptor;
+    }
+    return descriptor;
   }
 
   @Override
@@ -66,17 +67,18 @@
 
   @Override
   public DexString prefixRewrittenType(DexType type) {
-    return classRenaming.get(type);
+    return getRenaming(type);
   }
 
   @Override
   public DexString lookupDescriptor(DexType type) {
-    return classRenaming.getOrDefault(type, namingLens.lookupDescriptor(type));
+    DexString renaming = getRenaming(type);
+    return renaming != null ? renaming : namingLens.lookupDescriptor(type);
   }
 
   @Override
   public DexString lookupInnerName(InnerClassAttribute attribute, InternalOptions options) {
-    if (classRenaming.containsKey(attribute.getInner())) {
+    if (isRenamed(attribute.getInner())) {
       // Prefix rewriting does not influence the inner name.
       return attribute.getInnerName();
     }
@@ -85,7 +87,7 @@
 
   @Override
   public DexString lookupName(DexMethod method) {
-    if (classRenaming.containsKey(method.holder)) {
+    if (isRenamed(method.holder)) {
       // Prefix rewriting does not influence the method name.
       return method.name;
     }
@@ -94,7 +96,7 @@
 
   @Override
   public DexString lookupMethodName(DexCallSite callSite) {
-    if (classRenaming.containsKey(callSite.bootstrapMethod.rewrittenTarget.holder)) {
+    if (isRenamed(callSite.bootstrapMethod.rewrittenTarget.holder)) {
       // Prefix rewriting does not influence the inner name.
       return callSite.methodName;
     }
@@ -103,7 +105,7 @@
 
   @Override
   public DexString lookupName(DexField field) {
-    if (classRenaming.containsKey(field.holder)) {
+    if (isRenamed(field.holder)) {
       // Prefix rewriting does not influence the field name.
       return field.name;
     }
@@ -127,9 +129,10 @@
   }
 
   private boolean verifyNotPrefixRewrittenPackage(String packageName) {
-    for (DexType dexType : classRenaming.keySet()) {
-      assert !dexType.getPackageDescriptor().equals(packageName);
-    }
+    appView.rewritePrefix.forAllRewrittenTypes(
+        dexType -> {
+          assert !dexType.getPackageDescriptor().equals(packageName);
+        });
     return true;
   }
 
@@ -140,7 +143,7 @@
     // If compiling the desugared library, the mapping needs to be printed.
     // When debugging the program, both mapping files need to be merged.
     if (options.isDesugaredLibraryCompilation()) {
-      classRenaming.keySet().forEach(consumer);
+      appView.rewritePrefix.forAllRewrittenTypes(consumer);
     }
     namingLens.forAllRenamedTypes(consumer);
   }
@@ -150,13 +153,16 @@
       Class<T> clazz, Predicate<T> predicate, Function<T, String> namer) {
     Map<String, T> renamedItemsPrefixRewritting;
     if (clazz == DexType.class) {
-      renamedItemsPrefixRewritting =
-          classRenaming.keySet().stream()
-              .filter(item -> predicate.test(clazz.cast(item)))
-              .map(clazz::cast)
-              .collect(ImmutableMap.toImmutableMap(namer, i -> i));
+      renamedItemsPrefixRewritting = new HashMap<>();
+      appView.rewritePrefix.forAllRewrittenTypes(
+          item -> {
+            T cast = clazz.cast(item);
+            if (predicate.test(cast)) {
+              renamedItemsPrefixRewritting.put(namer.apply(cast), cast);
+            }
+          });
     } else {
-      renamedItemsPrefixRewritting = ImmutableMap.of();
+      renamedItemsPrefixRewritting = Collections.emptyMap();
     }
     Map<String, T> renamedItemsMinifier = namingLens.getRenamedItems(clazz, predicate, namer);
     // The Collector throws an exception for duplicated keys.