Refactor repackaging to use new NestedGraphLens

Bug: 182099754
Change-Id: I2c9dbcb50705f107f4fc124b9560cbe6aab0cfb5

2 files changed, 52 insertions, 50 deletions

diff --git a/src/main/java/com/android/tools/r8/graph/NestedGraphLens.java b/src/main/java/com/android/tools/r8/graph/NestedGraphLens.java
index e36700494..cf2e16313 100644
--- a/src/main/java/com/android/tools/r8/graph/NestedGraphLens.java
+++ b/src/main/java/com/android/tools/r8/graph/NestedGraphLens.java
@@ -24,8 +24,6 @@ import java.util.stream.Collectors;
  */
 public class NestedGraphLens extends NonIdentityGraphLens {
 
-  protected final DexItemFactory dexItemFactory;
-
   protected final BidirectionalManyToOneRepresentativeMap<DexField, DexField> fieldMap;
   protected final Map<DexMethod, DexMethod> methodMap;
   protected final Map<DexType, DexType> typeMap;
@@ -57,22 +55,28 @@ public class NestedGraphLens extends NonIdentityGraphLens {
   }
 
   public NestedGraphLens(
-      Map<DexType, DexType> typeMap,
-      Map<DexMethod, DexMethod> methodMap,
+      AppView<?> appView,
+      BidirectionalManyToOneRepresentativeMap<DexField, DexField> fieldMap,
+      BidirectionalManyToOneRepresentativeMap<DexMethod, DexMethod> methodMap,
+      Map<DexType, DexType> typeMap) {
+    this(appView, fieldMap, methodMap.getForwardMap(), typeMap, methodMap);
+  }
+
+  public NestedGraphLens(
+      AppView<?> appView,
       BidirectionalManyToOneRepresentativeMap<DexField, DexField> fieldMap,
-      BidirectionalManyToManyRepresentativeMap<DexMethod, DexMethod> newMethodSignatures,
-      GraphLens previousLens,
-      DexItemFactory dexItemFactory) {
-    super(dexItemFactory, previousLens);
+      Map<DexMethod, DexMethod> methodMap,
+      Map<DexType, DexType> typeMap,
+      BidirectionalManyToManyRepresentativeMap<DexMethod, DexMethod> newMethodSignatures) {
+    super(appView);
     assert !typeMap.isEmpty()
         || !methodMap.isEmpty()
         || !fieldMap.isEmpty()
         || isLegitimateToHaveEmptyMappings();
-    this.typeMap = typeMap.isEmpty() ? null : typeMap;
-    this.methodMap = methodMap;
     this.fieldMap = fieldMap;
+    this.methodMap = methodMap;
+    this.typeMap = typeMap;
     this.newMethodSignatures = newMethodSignatures;
-    this.dexItemFactory = dexItemFactory;
   }
 
   public static Builder builder() {
@@ -126,7 +130,7 @@ public class NestedGraphLens extends NonIdentityGraphLens {
 
   @Override
   protected DexType internalDescribeLookupClassType(DexType previous) {
-    return typeMap != null ? typeMap.getOrDefault(previous, previous) : previous;
+    return typeMap.getOrDefault(previous, previous);
   }
 
   @Override
@@ -139,7 +143,7 @@ public class NestedGraphLens extends NonIdentityGraphLens {
               ? rewrittenReboundReference
               : rewrittenReboundReference.withHolder(
                   internalDescribeLookupClassType(previous.getReference().getHolderType()),
-                  dexItemFactory);
+                  dexItemFactory());
       return FieldLookupResult.builder(this)
           .setReboundReference(rewrittenReboundReference)
           .setReference(rewrittenNonReboundReference)
@@ -169,7 +173,7 @@ public class NestedGraphLens extends NonIdentityGraphLens {
               : // This assumes that the holder will always be moved in lock-step with the method!
               rewrittenReboundReference.withHolder(
                   internalDescribeLookupClassType(previous.getReference().getHolderType()),
-                  dexItemFactory);
+                  dexItemFactory());
       return MethodLookupResult.builder(this)
           .setReference(rewrittenReference)
           .setReboundReference(rewrittenReboundReference)
@@ -283,24 +287,28 @@ public class NestedGraphLens extends NonIdentityGraphLens {
   @Override
   public String toString() {
     StringBuilder builder = new StringBuilder();
-    if (typeMap != null) {
-      for (Map.Entry<DexType, DexType> entry : typeMap.entrySet()) {
-        builder.append(entry.getKey().toSourceString()).append(" -> ");
-        builder.append(entry.getValue().toSourceString()).append(System.lineSeparator());
-      }
-    }
-    for (Map.Entry<DexMethod, DexMethod> entry : methodMap.entrySet()) {
-      builder.append(entry.getKey().toSourceString()).append(" -> ");
-      builder.append(entry.getValue().toSourceString()).append(System.lineSeparator());
-    }
+    typeMap.forEach(
+        (from, to) ->
+            builder
+                .append(from.getTypeName())
+                .append(" -> ")
+                .append(to.getTypeName())
+                .append(System.lineSeparator()));
+    methodMap.forEach(
+        (from, to) ->
+            builder
+                .append(from.toSourceString())
+                .append(" -> ")
+                .append(to.toSourceString())
+                .append(System.lineSeparator()));
     fieldMap.forEachManyToOneMapping(
-        (keys, value) -> {
+        (fromSet, to) -> {
           builder.append(
-              keys.stream()
+              fromSet.stream()
                   .map(DexField::toSourceString)
                   .collect(Collectors.joining("," + System.lineSeparator())));
           builder.append(" -> ");
-          builder.append(value.toSourceString()).append(System.lineSeparator());
+          builder.append(to.toSourceString()).append(System.lineSeparator());
         });
     builder.append(getPrevious().toString());
     return builder.toString();
diff --git a/src/main/java/com/android/tools/r8/repackaging/RepackagingLens.java b/src/main/java/com/android/tools/r8/repackaging/RepackagingLens.java
index d0f6c2f1d..64fe0ebed 100644
--- a/src/main/java/com/android/tools/r8/repackaging/RepackagingLens.java
+++ b/src/main/java/com/android/tools/r8/repackaging/RepackagingLens.java
@@ -10,7 +10,7 @@ import com.android.tools.r8.graph.DexMember;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexReference;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.LegacyNestedGraphLens;
+import com.android.tools.r8.graph.NestedGraphLens;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.IterableUtils;
 import com.android.tools.r8.utils.collections.BidirectionalOneToOneHashMap;
@@ -20,25 +20,19 @@ import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import java.util.Map;
 
-public class RepackagingLens extends LegacyNestedGraphLens {
+public class RepackagingLens extends NestedGraphLens {
 
-  private final BiMap<DexType, DexType> originalTypes;
+  private final BiMap<DexType, DexType> newTypes;
   private final Map<String, String> packageRenamings;
 
   private RepackagingLens(
       AppView<AppInfoWithLiveness> appView,
       BidirectionalOneToOneMap<DexField, DexField> newFieldSignatures,
-      BidirectionalOneToOneMap<DexMethod, DexMethod> originalMethodSignatures,
-      BiMap<DexType, DexType> originalTypes,
+      BidirectionalOneToOneMap<DexMethod, DexMethod> newMethodSignatures,
+      BiMap<DexType, DexType> newTypes,
       Map<String, String> packageRenamings) {
-    super(
-        originalTypes.inverse(),
-        originalMethodSignatures.getInverseOneToOneMap().getForwardMap(),
-        newFieldSignatures,
-        originalMethodSignatures,
-        appView.graphLens(),
-        appView.dexItemFactory());
-    this.originalTypes = originalTypes;
+    super(appView, newFieldSignatures, newMethodSignatures, newTypes);
+    this.newTypes = newTypes;
     this.packageRenamings = packageRenamings;
   }
 
@@ -49,7 +43,7 @@ public class RepackagingLens extends LegacyNestedGraphLens {
 
   @Override
   public DexType getOriginalType(DexType type) {
-    DexType previous = originalTypes.getOrDefault(type, type);
+    DexType previous = newTypes.inverse().getOrDefault(type, type);
     return getPrevious().getOriginalType(previous);
   }
 
@@ -73,7 +67,7 @@ public class RepackagingLens extends LegacyNestedGraphLens {
   }
 
   private boolean isSimpleTypeRenamingOrEqual(DexType from, DexType to) {
-    return from == to || originalTypes.get(to) == from;
+    return from == to || newTypes.get(from) == to;
   }
 
   private boolean isSimpleTypeRenamingOrEqual(DexMember<?, ?> from, DexMember<?, ?> to) {
@@ -82,16 +76,16 @@ public class RepackagingLens extends LegacyNestedGraphLens {
     }
     return IterableUtils.testPairs(
         this::isSimpleTypeRenamingOrEqual,
-        from.getReferencedBaseTypes(dexItemFactory),
-        to.getReferencedBaseTypes(dexItemFactory));
+        from.getReferencedBaseTypes(dexItemFactory()),
+        to.getReferencedBaseTypes(dexItemFactory()));
   }
 
   public static class Builder {
 
-    protected final BiMap<DexType, DexType> originalTypes = HashBiMap.create();
+    protected final BiMap<DexType, DexType> newTypes = HashBiMap.create();
     protected final MutableBidirectionalOneToOneMap<DexField, DexField> newFieldSignatures =
         new BidirectionalOneToOneHashMap<>();
-    protected final MutableBidirectionalOneToOneMap<DexMethod, DexMethod> originalMethodSignatures =
+    protected final MutableBidirectionalOneToOneMap<DexMethod, DexMethod> newMethodSignatures =
         new BidirectionalOneToOneHashMap<>();
 
     public void recordMove(DexField from, DexField to) {
@@ -99,18 +93,18 @@ public class RepackagingLens extends LegacyNestedGraphLens {
     }
 
     public void recordMove(DexMethod from, DexMethod to) {
-      originalMethodSignatures.put(to, from);
+      newMethodSignatures.put(from, to);
     }
 
     public void recordMove(DexType from, DexType to) {
-      originalTypes.put(to, from);
+      newTypes.put(from, to);
     }
 
     public RepackagingLens build(
         AppView<AppInfoWithLiveness> appView, Map<String, String> packageRenamings) {
-      assert !originalTypes.isEmpty();
+      assert !newTypes.isEmpty();
       return new RepackagingLens(
-          appView, newFieldSignatures, originalMethodSignatures, originalTypes, packageRenamings);
+          appView, newFieldSignatures, newMethodSignatures, newTypes, packageRenamings);
     }
   }
 }