Introduce a new NestedGraphLens with newMethodSignatures

Bug: 182099754
Change-Id: Ic6990f9361b8218889f3a440fd7805921455dd23
diff --git a/src/main/java/com/android/tools/r8/graph/NestedGraphLens.java b/src/main/java/com/android/tools/r8/graph/NestedGraphLens.java
new file mode 100644
index 0000000..e367004
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/graph/NestedGraphLens.java
@@ -0,0 +1,308 @@
+// Copyright (c) 2021, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.graph;
+
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
+import com.android.tools.r8.ir.code.Invoke.Type;
+import com.android.tools.r8.utils.IterableUtils;
+import com.android.tools.r8.utils.collections.BidirectionalManyToManyRepresentativeMap;
+import com.android.tools.r8.utils.collections.BidirectionalManyToOneRepresentativeMap;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * GraphLens implementation with a parent lens using a simple mapping for type, method and field
+ * mapping.
+ *
+ * <p>Subclasses can override the lookup methods.
+ *
+ * <p>For method mapping where invocation type can change just override {@link
+ * #mapInvocationType(DexMethod, DexMethod, Type)} if the default name mapping applies, and only
+ * invocation type might need to change.
+ */
+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;
+
+  // Map that stores the new signature of methods that have been affected by class merging, unused
+  // argument removal, repackaging, synthetic finalization, etc. This is needed to generate a
+  // correct mapping file in the end.
+  //
+  // The difference to `methodMap` is that `methodMap` specifies how invoke-method instructions
+  // should be rewritten, whereas this map specifies where methods have been moved. In most cases
+  // the two maps are the same, but in a few cases we move a method m1 to m2 and rewrite invokes to
+  // m1 to m3.
+  //
+  // The static type of this map is a many-to-many map to facilitate both one-to-many mappings and
+  // many-to-one mappings. Currently, the concrete map is always *either* a one-to-many or a
+  // many-to-one map, and not a many-to-many map.
+  //
+  // One-to-many mappings are generally found when methods are split into two. This could be due to
+  // the code object being moved elsewhere (e.g., interface desugaring).
+  //
+  // Many-to-one mappings are generally found when methods are shared, e.g., due to horizontal class
+  // merging of synthetic finalization.
+  protected BidirectionalManyToManyRepresentativeMap<DexMethod, DexMethod> newMethodSignatures;
+
+  // Overrides this if the sub type needs to be a nested lens while it doesn't have any mappings
+  // at all, e.g., publicizer lens that changes invocation type only.
+  protected boolean isLegitimateToHaveEmptyMappings() {
+    return false;
+  }
+
+  public NestedGraphLens(
+      Map<DexType, DexType> typeMap,
+      Map<DexMethod, DexMethod> methodMap,
+      BidirectionalManyToOneRepresentativeMap<DexField, DexField> fieldMap,
+      BidirectionalManyToManyRepresentativeMap<DexMethod, DexMethod> newMethodSignatures,
+      GraphLens previousLens,
+      DexItemFactory dexItemFactory) {
+    super(dexItemFactory, previousLens);
+    assert !typeMap.isEmpty()
+        || !methodMap.isEmpty()
+        || !fieldMap.isEmpty()
+        || isLegitimateToHaveEmptyMappings();
+    this.typeMap = typeMap.isEmpty() ? null : typeMap;
+    this.methodMap = methodMap;
+    this.fieldMap = fieldMap;
+    this.newMethodSignatures = newMethodSignatures;
+    this.dexItemFactory = dexItemFactory;
+  }
+
+  public static Builder builder() {
+    return new Builder();
+  }
+
+  protected DexType internalGetOriginalType(DexType previous) {
+    return previous;
+  }
+
+  protected Iterable<DexType> internalGetOriginalTypes(DexType previous) {
+    return IterableUtils.singleton(internalGetOriginalType(previous));
+  }
+
+  @Override
+  public DexType getOriginalType(DexType type) {
+    return getPrevious().getOriginalType(internalGetOriginalType(type));
+  }
+
+  @Override
+  public Iterable<DexType> getOriginalTypes(DexType type) {
+    return IterableUtils.flatMap(internalGetOriginalTypes(type), getPrevious()::getOriginalTypes);
+  }
+
+  @Override
+  public DexField getOriginalFieldSignature(DexField field) {
+    DexField originalField = fieldMap.getRepresentativeKeyOrDefault(field, field);
+    return getPrevious().getOriginalFieldSignature(originalField);
+  }
+
+  @Override
+  public DexMethod getOriginalMethodSignature(DexMethod method) {
+    DexMethod originalMethod = internalGetPreviousMethodSignature(method);
+    return getPrevious().getOriginalMethodSignature(originalMethod);
+  }
+
+  @Override
+  public DexField getRenamedFieldSignature(DexField originalField) {
+    DexField renamedField = getPrevious().getRenamedFieldSignature(originalField);
+    return internalGetNextFieldSignature(renamedField);
+  }
+
+  @Override
+  public DexMethod getRenamedMethodSignature(DexMethod originalMethod, GraphLens applied) {
+    if (this == applied) {
+      return originalMethod;
+    }
+    DexMethod renamedMethod = getPrevious().getRenamedMethodSignature(originalMethod, applied);
+    return internalGetNextMethodSignature(renamedMethod);
+  }
+
+  @Override
+  protected DexType internalDescribeLookupClassType(DexType previous) {
+    return typeMap != null ? typeMap.getOrDefault(previous, previous) : previous;
+  }
+
+  @Override
+  protected FieldLookupResult internalDescribeLookupField(FieldLookupResult previous) {
+    if (previous.hasReboundReference()) {
+      // Rewrite the rebound reference and then "fixup" the non-rebound reference.
+      DexField rewrittenReboundReference = previous.getRewrittenReboundReference(fieldMap);
+      DexField rewrittenNonReboundReference =
+          previous.getReference() == previous.getReboundReference()
+              ? rewrittenReboundReference
+              : rewrittenReboundReference.withHolder(
+                  internalDescribeLookupClassType(previous.getReference().getHolderType()),
+                  dexItemFactory);
+      return FieldLookupResult.builder(this)
+          .setReboundReference(rewrittenReboundReference)
+          .setReference(rewrittenNonReboundReference)
+          .setCastType(previous.getRewrittenCastType(this::internalDescribeLookupClassType))
+          .build();
+    } else {
+      // TODO(b/168282032): We should always have the rebound reference, so this should become
+      //  unreachable.
+      DexField rewrittenReference = previous.getRewrittenReference(fieldMap);
+      return FieldLookupResult.builder(this)
+          .setReference(rewrittenReference)
+          .setCastType(previous.getRewrittenCastType(this::internalDescribeLookupClassType))
+          .build();
+    }
+  }
+
+  @Override
+  public MethodLookupResult internalDescribeLookupMethod(
+      MethodLookupResult previous, DexMethod context) {
+    if (previous.hasReboundReference()) {
+      // TODO(sgjesse): Should we always do interface to virtual mapping? Is it a performance win
+      //  that only subclasses which are known to need it actually do it?
+      DexMethod rewrittenReboundReference = previous.getRewrittenReboundReference(methodMap);
+      DexMethod rewrittenReference =
+          previous.getReference() == previous.getReboundReference()
+              ? rewrittenReboundReference
+              : // This assumes that the holder will always be moved in lock-step with the method!
+              rewrittenReboundReference.withHolder(
+                  internalDescribeLookupClassType(previous.getReference().getHolderType()),
+                  dexItemFactory);
+      return MethodLookupResult.builder(this)
+          .setReference(rewrittenReference)
+          .setReboundReference(rewrittenReboundReference)
+          .setPrototypeChanges(
+              internalDescribePrototypeChanges(
+                  previous.getPrototypeChanges(), rewrittenReboundReference))
+          .setType(
+              mapInvocationType(
+                  rewrittenReboundReference, previous.getReference(), previous.getType()))
+          .build();
+    } else {
+      // TODO(b/168282032): We should always have the rebound reference, so this should become
+      //  unreachable.
+      DexMethod newMethod = methodMap.get(previous.getReference());
+      if (newMethod == null) {
+        return previous;
+      }
+      // TODO(sgjesse): Should we always do interface to virtual mapping? Is it a performance win
+      //  that only subclasses which are known to need it actually do it?
+      return MethodLookupResult.builder(this)
+          .setReference(newMethod)
+          .setPrototypeChanges(
+              internalDescribePrototypeChanges(previous.getPrototypeChanges(), newMethod))
+          .setType(mapInvocationType(newMethod, previous.getReference(), previous.getType()))
+          .build();
+    }
+  }
+
+  @Override
+  public RewrittenPrototypeDescription lookupPrototypeChangesForMethodDefinition(DexMethod method) {
+    DexMethod previous = internalGetPreviousMethodSignature(method);
+    RewrittenPrototypeDescription lookup =
+        getPrevious().lookupPrototypeChangesForMethodDefinition(previous);
+    return internalDescribePrototypeChanges(lookup, method);
+  }
+
+  protected RewrittenPrototypeDescription internalDescribePrototypeChanges(
+      RewrittenPrototypeDescription prototypeChanges, DexMethod method) {
+    return prototypeChanges;
+  }
+
+  protected DexField internalGetNextFieldSignature(DexField field) {
+    return fieldMap.getOrDefault(field, field);
+  }
+
+  @Override
+  protected DexMethod internalGetPreviousMethodSignature(DexMethod method) {
+    return newMethodSignatures.getRepresentativeKeyOrDefault(method, method);
+  }
+
+  protected DexMethod internalGetNextMethodSignature(DexMethod method) {
+    return newMethodSignatures.getRepresentativeValueOrDefault(method, method);
+  }
+
+  @Override
+  public DexMethod lookupGetFieldForMethod(DexField field, DexMethod context) {
+    return getPrevious().lookupGetFieldForMethod(field, context);
+  }
+
+  @Override
+  public DexMethod lookupPutFieldForMethod(DexField field, DexMethod context) {
+    return getPrevious().lookupPutFieldForMethod(field, context);
+  }
+
+  /**
+   * Default invocation type mapping.
+   *
+   * <p>This is an identity mapping. If a subclass need invocation type mapping either override this
+   * method or {@link #lookupMethod(DexMethod, DexMethod, Type)}
+   */
+  protected Type mapInvocationType(DexMethod newMethod, DexMethod originalMethod, Type type) {
+    return type;
+  }
+
+  /**
+   * Standard mapping between interface and virtual invoke type.
+   *
+   * <p>Handle methods moved from interface to class or class to interface.
+   */
+  public static Type mapVirtualInterfaceInvocationTypes(
+      DexDefinitionSupplier definitions, DexMethod newMethod, DexMethod originalMethod, Type type) {
+    if (type == Type.VIRTUAL || type == Type.INTERFACE) {
+      // Get the invoke type of the actual definition.
+      DexClass newTargetClass = definitions.definitionFor(newMethod.getHolderType());
+      if (newTargetClass == null) {
+        return type;
+      }
+      DexClass originalTargetClass = definitions.definitionFor(originalMethod.getHolderType());
+      if (originalTargetClass != null
+          && (originalTargetClass.isInterface() ^ (type == Type.INTERFACE))) {
+        // The invoke was wrong to start with, so we keep it wrong. This is to ensure we get
+        // the IncompatibleClassChangeError the original invoke would have triggered.
+        return newTargetClass.accessFlags.isInterface() ? Type.VIRTUAL : Type.INTERFACE;
+      }
+      return newTargetClass.accessFlags.isInterface() ? Type.INTERFACE : Type.VIRTUAL;
+    }
+    return type;
+  }
+
+  @Override
+  public boolean isContextFreeForMethods() {
+    return getPrevious().isContextFreeForMethods();
+  }
+
+  @Override
+  public boolean verifyIsContextFreeForMethod(DexMethod method) {
+    assert getPrevious().verifyIsContextFreeForMethod(method);
+    return true;
+  }
+
+  @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());
+    }
+    fieldMap.forEachManyToOneMapping(
+        (keys, value) -> {
+          builder.append(
+              keys.stream()
+                  .map(DexField::toSourceString)
+                  .collect(Collectors.joining("," + System.lineSeparator())));
+          builder.append(" -> ");
+          builder.append(value.toSourceString()).append(System.lineSeparator());
+        });
+    builder.append(getPrevious().toString());
+    return builder.toString();
+  }
+}