Move invocation type mapping to GraphLense

Method lookup in a GraphLense now returns both the method to invoke and
the invocation type to use as its result.

Before this happened in the LenseCodeRewriter after mapping the
method. This was somewhat fragile and disconnected from the method
mappping. Each lense providing the method mapping should also have the
information to determine the invocation type to use for the mapped
method.

Change-Id: I4dc65f5df8bfd5a392d8f660b83c61f7acd35282
diff --git a/src/main/java/com/android/tools/r8/graph/GraphLense.java b/src/main/java/com/android/tools/r8/graph/GraphLense.java
index 096de5a..4126383 100644
--- a/src/main/java/com/android/tools/r8/graph/GraphLense.java
+++ b/src/main/java/com/android/tools/r8/graph/GraphLense.java
@@ -20,11 +20,35 @@
  * <li>Renaming private methods/fields.</li>
  * <li>Moving methods/fields to a super/subclass.</li>
  * <li>Replacing method/field references by the same method/field on a super/subtype</li>
+ * <li>Moved methods might require changed invocation type at the call site</li>
  * </ul>
  * Note that the latter two have to take visibility into account.
  */
 public abstract class GraphLense {
 
+  /**
+   * Result of a method lookup in a GraphLense.
+   *
+   * This provide the new target and the invoke type to use.
+   */
+  public static class GraphLenseLookupResult {
+    private final DexMethod method;
+    private final Type type;
+
+    public GraphLenseLookupResult(DexMethod method, Type type) {
+      this.method = method;
+      this.type = type;
+    }
+
+    public DexMethod getMethod() {
+      return method;
+    }
+
+    public Type getType() {
+      return type;
+    }
+  }
+
   public static class Builder {
 
     protected Builder() {
@@ -68,10 +92,11 @@
   // This overload can be used when the graph lense is known to be context insensitive.
   public DexMethod lookupMethod(DexMethod method) {
     assert isContextFreeForMethod(method);
-    return lookupMethod(method, null, null);
+    return lookupMethod(method, null, null).getMethod();
   }
 
-  public abstract DexMethod lookupMethod(DexMethod method, DexEncodedMethod context, Type type);
+  public abstract GraphLenseLookupResult lookupMethod(
+      DexMethod method, DexEncodedMethod context, Type type);
 
   // Context sensitive graph lenses should override this method.
   public Set<DexMethod> lookupMethodInAllContexts(DexMethod method) {
@@ -107,8 +132,9 @@
     }
 
     @Override
-    public DexMethod lookupMethod(DexMethod method, DexEncodedMethod context, Type type) {
-      return method;
+    public GraphLenseLookupResult lookupMethod(
+        DexMethod method, DexEncodedMethod context, Type type) {
+      return new GraphLenseLookupResult(method, type);
     }
 
     @Override
@@ -122,19 +148,29 @@
     }
   }
 
+  /**
+   * GraphLense implementation with a parent lense using a simple mapping for type, method and
+   * field mapping.
+   *
+   * Subclasses can override the lookup methods.
+   *
+   * For method mapping where invocation type can change just override
+   * {@link #mapInvocationType(DexMethod, DexMethod, DexEncodedMethod, Type)} if
+   * the default name mapping applies, and only invocation type might need to change.
+   */
   public static class NestedGraphLense extends GraphLense {
 
-    private final GraphLense previousLense;
+    protected final GraphLense previousLense;
     protected final DexItemFactory dexItemFactory;
 
-    private final Map<DexType, DexType> typeMap;
+    protected final Map<DexType, DexType> typeMap;
     private final Map<DexType, DexType> arrayTypeCache = new IdentityHashMap<>();
-    private final Map<DexMethod, DexMethod> methodMap;
-    private final Map<DexField, DexField> fieldMap;
+    protected final Map<DexMethod, DexMethod> methodMap;
+    protected final Map<DexField, DexField> fieldMap;
 
     public NestedGraphLense(Map<DexType, DexType> typeMap, Map<DexMethod, DexMethod> methodMap,
         Map<DexField, DexField> fieldMap, GraphLense previousLense, DexItemFactory dexItemFactory) {
-      this.typeMap = typeMap;
+      this.typeMap = typeMap.isEmpty() ? null : typeMap;
       this.methodMap = methodMap;
       this.fieldMap = fieldMap;
       this.previousLense = previousLense;
@@ -161,13 +197,58 @@
         }
       }
       DexType previous = previousLense.lookupType(type);
-      return typeMap.getOrDefault(previous, previous);
+      return typeMap != null ? typeMap.getOrDefault(previous, previous) : previous;
     }
 
     @Override
-    public DexMethod lookupMethod(DexMethod method, DexEncodedMethod context, Type type) {
-      DexMethod previous = previousLense.lookupMethod(method, context, type);
-      return methodMap.getOrDefault(previous, previous);
+    public GraphLenseLookupResult lookupMethod(
+        DexMethod method, DexEncodedMethod context, Type type) {
+      GraphLenseLookupResult previous = previousLense.lookupMethod(method, context, type);
+      DexMethod newMethod = methodMap.get(previous.getMethod());
+      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 new GraphLenseLookupResult(
+          newMethod, mapInvocationType(newMethod, method, context, type));
+    }
+
+    /**
+     * Default invocation type mapping.
+     *
+     * This is an identity mapping. If a subclass need invocation type mapping either override
+     * this method or {@link #lookupMethod(DexMethod, DexEncodedMethod, Type)}
+     */
+    protected Type mapInvocationType(
+        DexMethod newMethod, DexMethod originalMethod, DexEncodedMethod context, Type type) {
+      return type;
+    }
+
+    /**
+     * Standard mapping between interface and virtual invoke type.
+     *
+     * Handle methods moved from interface to class or class to interface.
+     */
+    final protected Type mapVirtualInterfaceInvocationTypes(
+        AppInfo appInfo, DexMethod newMethod, DexMethod originalMethod,
+        DexEncodedMethod context, Type type) {
+      if (type == Type.VIRTUAL || type == Type.INTERFACE) {
+        // Get the invoke type of the actual definition.
+        DexClass newTargetClass = appInfo.definitionFor(newMethod.holder);
+        if (newTargetClass == null) {
+          return type;
+        }
+        DexClass originalTargetClass = appInfo.definitionFor(originalMethod.holder);
+        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
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/CallGraph.java b/src/main/java/com/android/tools/r8/ir/conversion/CallGraph.java
index 21b8dc4..85506ad 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/CallGraph.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/CallGraph.java
@@ -363,7 +363,7 @@
 
     private void processInvoke(Type type, DexMethod method) {
       DexEncodedMethod source = caller.method;
-      method = graphLense.lookupMethod(method, source, type);
+      method = graphLense.lookupMethod(method, source, type).getMethod();
       DexEncodedMethod definition = appInfo.lookup(type, method, source.method.holder);
       if (definition != null) {
         assert !source.accessFlags.isBridge() || definition != caller.method;
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
index 7d7e5f3..52919c6 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
@@ -141,7 +141,7 @@
       this.memberValuePropagation =
           options.enableValuePropagation ?
               new MemberValuePropagation(appInfo.withLiveness()) : null;
-      this.lensCodeRewriter = new LensCodeRewriter(graphLense, appInfo.withSubtyping(), options);
+      this.lensCodeRewriter = new LensCodeRewriter(graphLense, appInfo.withSubtyping());
       if (appInfo.hasLiveness()) {
         // When disabling the pruner here, also disable the ProtoLiteExtension in R8.java.
         this.protoLiteRewriter =
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java b/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
index 0e4be8b..3c9850a 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/LensCodeRewriter.java
@@ -5,7 +5,6 @@
 
 import com.android.tools.r8.graph.AppInfoWithSubtyping;
 import com.android.tools.r8.graph.DexCallSite;
-import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexMethod;
@@ -15,6 +14,7 @@
 import com.android.tools.r8.graph.DexValue;
 import com.android.tools.r8.graph.DexValue.DexValueMethodHandle;
 import com.android.tools.r8.graph.GraphLense;
+import com.android.tools.r8.graph.GraphLense.GraphLenseLookupResult;
 import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.CheckCast;
 import com.android.tools.r8.ir.code.ConstClass;
@@ -34,7 +34,6 @@
 import com.android.tools.r8.ir.code.StaticGet;
 import com.android.tools.r8.ir.code.StaticPut;
 import com.android.tools.r8.ir.code.Value;
-import com.android.tools.r8.utils.InternalOptions;
 import java.util.List;
 import java.util.ListIterator;
 import java.util.stream.Collectors;
@@ -43,13 +42,10 @@
 
   private final GraphLense graphLense;
   private final AppInfoWithSubtyping appInfo;
-  private final InternalOptions options;
 
-  public LensCodeRewriter(
-      GraphLense graphLense, AppInfoWithSubtyping appInfo, InternalOptions options) {
+  public LensCodeRewriter(GraphLense graphLense, AppInfoWithSubtyping appInfo) {
     this.graphLense = graphLense;
     this.appInfo = appInfo;
-    this.options = options;
   }
 
   private Value makeOutValue(Instruction insn, IRCode code) {
@@ -100,8 +96,10 @@
           if (!invokedHolder.isClassType()) {
             continue;
           }
-          DexMethod actualTarget = graphLense.lookupMethod(invokedMethod, method, invoke.getType());
-          Invoke.Type invokeType = getInvokeType(invoke, actualTarget, invokedMethod, method);
+          GraphLenseLookupResult lenseLookup =
+              graphLense.lookupMethod(invokedMethod, method, invoke.getType());
+          DexMethod actualTarget = lenseLookup.getMethod();
+          Invoke.Type invokeType = lenseLookup.getType();
           if (actualTarget != invokedMethod || invoke.getType() != invokeType) {
             Invoke newInvoke = Invoke.create(invokeType, actualTarget, null,
                     invoke.outValue(), invoke.inValues());
@@ -220,17 +218,13 @@
       DexEncodedMethod method, DexMethodHandle methodHandle) {
     if (methodHandle.isMethodHandle()) {
       DexMethod invokedMethod = methodHandle.asMethod();
-      DexMethod actualTarget =
-          graphLense.lookupMethod(invokedMethod, method, methodHandle.type.toInvokeType());
+      GraphLenseLookupResult lenseLookup = graphLense.lookupMethod(
+          invokedMethod, method, methodHandle.type.toInvokeType());
+      DexMethod actualTarget = lenseLookup.getMethod();
       if (actualTarget != invokedMethod) {
-        DexClass clazz = appInfo.definitionFor(actualTarget.holder);
-        MethodHandleType newType = methodHandle.type;
-        if (clazz != null
-            && (newType.isInvokeInterface() || newType.isInvokeInstance())) {
-          newType = clazz.accessFlags.isInterface()
-              ? MethodHandleType.INVOKE_INTERFACE
-              : MethodHandleType.INVOKE_INSTANCE;
-        }
+        MethodHandleType newType = lenseLookup.getType() == Type.INTERFACE
+            ? MethodHandleType.INVOKE_INTERFACE
+            : MethodHandleType.INVOKE_INSTANCE;
         return new DexMethodHandle(newType, actualTarget);
       }
     } else {
@@ -242,55 +236,4 @@
     }
     return methodHandle;
   }
-
-  private Type getInvokeType(
-      InvokeMethod invoke,
-      DexMethod actualTarget,
-      DexMethod originalTarget,
-      DexEncodedMethod invocationContext) {
-    // We might move methods from interfaces to classes and vice versa. So we have to support
-    // fixing the invoke kind, yet only if it was correct to start with.
-    if (invoke.isInvokeVirtual() || invoke.isInvokeInterface()) {
-      // Get the invoke type of the actual definition.
-      DexClass newTargetClass = appInfo.definitionFor(actualTarget.holder);
-      if (newTargetClass == null) {
-        return invoke.getType();
-      }
-      DexClass originalTargetClass = appInfo.definitionFor(originalTarget.holder);
-      if (originalTargetClass != null
-          && (originalTargetClass.isInterface() ^ (invoke.getType() == 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;
-    }
-    if (options.enableClassMerging && invoke.isInvokeSuper()) {
-      if (actualTarget.getHolder() == invocationContext.method.getHolder()) {
-        DexClass targetClass = appInfo.definitionFor(actualTarget.holder);
-        if (targetClass == null) {
-          return invoke.getType();
-        }
-
-        // If the super class A of the enclosing class B (i.e., invocationContext.method.holder)
-        // has been merged into B during vertical class merging, and this invoke-super instruction
-        // was resolving to a method in A, then the target method has been changed to a direct
-        // method and moved into B, so that we need to use an invoke-direct instruction instead of
-        // invoke-super.
-        //
-        // At this point, we have an invoke-super instruction where the static target is the
-        // enclosing class. However, such an instruction could occur even if a subclass has never
-        // been merged into the enclosing class. Therefore, to determine if vertical class merging
-        // has been applied, we look if there is a direct method with the right signature, and only
-        // return Type.DIRECT in that case.
-        DexEncodedMethod method = targetClass.lookupDirectMethod(actualTarget);
-        if (method != null) {
-          // The target method has been moved from the super class into the sub class during class
-          // merging such that we now need to use an invoke-direct instruction.
-          return Type.DIRECT;
-        }
-      }
-    }
-    return invoke.getType();
-  }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/Inliner.java b/src/main/java/com/android/tools/r8/ir/optimize/Inliner.java
index 029b169..25cd969 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/Inliner.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/Inliner.java
@@ -271,7 +271,7 @@
       Origin origin = appInfo.originFor(target.method.holder);
       IRCode code = target.buildInliningIR(appInfo, options, generator, callerPosition, origin);
       if (!target.isProcessed()) {
-        new LensCodeRewriter(graphLense, appInfo, options).rewrite(code, target);
+        new LensCodeRewriter(graphLense, appInfo).rewrite(code, target);
       }
       return code;
     }
diff --git a/src/main/java/com/android/tools/r8/optimize/BridgeMethodAnalysis.java b/src/main/java/com/android/tools/r8/optimize/BridgeMethodAnalysis.java
index 56089f4..492bcd0 100644
--- a/src/main/java/com/android/tools/r8/optimize/BridgeMethodAnalysis.java
+++ b/src/main/java/com/android/tools/r8/optimize/BridgeMethodAnalysis.java
@@ -46,14 +46,14 @@
         assert !method.accessFlags.isPrivate() && !method.accessFlags.isConstructor();
         if (kind == InvokeKind.STATIC) {
           assert method.accessFlags.isStatic();
-          DexMethod actualTarget = lense.lookupMethod(target, method, Type.STATIC);
+          DexMethod actualTarget = lense.lookupMethod(target, method, Type.STATIC).getMethod();
           DexEncodedMethod targetMethod = appInfo.lookupStaticTarget(actualTarget);
           if (targetMethod != null) {
             addForwarding(method, targetMethod);
           }
         } else if (kind == InvokeKind.VIRTUAL) {
           // TODO(herhut): Add support for bridges with multiple targets.
-          DexMethod actualTarget = lense.lookupMethod(target, method, Type.VIRTUAL);
+          DexMethod actualTarget = lense.lookupMethod(target, method, Type.VIRTUAL).getMethod();
           DexEncodedMethod targetMethod = appInfo.lookupSingleVirtualTarget(actualTarget);
           if (targetMethod != null) {
             addForwarding(method, targetMethod);
@@ -93,15 +93,16 @@
     }
 
     @Override
-    public DexMethod lookupMethod(DexMethod method, DexEncodedMethod context, Type type) {
-      DexMethod previous = previousLense.lookupMethod(method, context, type);
-      DexMethod bridge = bridgeTargetToBridgeMap.get(previous);
+    public GraphLenseLookupResult lookupMethod(
+        DexMethod method, DexEncodedMethod context, Type type) {
+      GraphLenseLookupResult previous = previousLense.lookupMethod(method, context, type);
+      DexMethod bridge = bridgeTargetToBridgeMap.get(previous.getMethod());
       // Do not forward calls from a bridge method to itself while the bridge method is still
       // a bridge.
       if (bridge == null || (context.accessFlags.isBridge() && bridge == context.method)) {
         return previous;
       }
-      return bridge;
+      return new GraphLenseLookupResult(bridge, type);
     }
 
     @Override
diff --git a/src/main/java/com/android/tools/r8/optimize/MemberRebindingAnalysis.java b/src/main/java/com/android/tools/r8/optimize/MemberRebindingAnalysis.java
index fc408e3..da7a787 100644
--- a/src/main/java/com/android/tools/r8/optimize/MemberRebindingAnalysis.java
+++ b/src/main/java/com/android/tools/r8/optimize/MemberRebindingAnalysis.java
@@ -26,12 +26,13 @@
 
   private final AppInfoWithLiveness appInfo;
   private final GraphLense lense;
-  private final GraphLense.Builder builder = GraphLense.builder();
+  private final MemberRebindingLense.Builder builder;
 
   public MemberRebindingAnalysis(AppInfoWithLiveness appInfo, GraphLense lense) {
     assert lense.isContextFreeForMethods();
     this.appInfo = appInfo;
     this.lense = lense;
+    this.builder = MemberRebindingLense.builder(appInfo);
   }
 
   private DexMethod validTargetFor(DexMethod target, DexMethod original) {
@@ -250,6 +251,6 @@
         mergeFieldAccessContexts(appInfo.instanceFieldReads, appInfo.instanceFieldWrites),
         appInfo::resolveFieldOn, DexClass::lookupField);
 
-    return builder.build(appInfo.dexItemFactory, lense);
+    return builder.build(lense);
   }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/MemberRebindingLense.java b/src/main/java/com/android/tools/r8/optimize/MemberRebindingLense.java
new file mode 100644
index 0000000..4ee0cee
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/MemberRebindingLense.java
@@ -0,0 +1,56 @@
+// Copyright (c) 2018, 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.optimize;
+
+import com.android.tools.r8.graph.AppInfo;
+import com.android.tools.r8.graph.DexEncodedMethod;
+import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.graph.DexMethod;
+import com.android.tools.r8.graph.GraphLense;
+import com.android.tools.r8.graph.GraphLense.NestedGraphLense;
+import com.android.tools.r8.ir.code.Invoke.Type;
+import com.google.common.collect.ImmutableMap;
+import java.util.Map;
+
+public class MemberRebindingLense extends NestedGraphLense {
+  public static class Builder extends NestedGraphLense.Builder {
+    private final AppInfo appInfo;
+
+    protected Builder(AppInfo appInfo) {
+      this.appInfo = appInfo;
+    }
+
+    public GraphLense build(GraphLense previousLense) {
+      assert typeMap.isEmpty();
+      if (methodMap.isEmpty() && fieldMap.isEmpty()) {
+        return previousLense;
+      }
+      return new MemberRebindingLense(appInfo, methodMap, fieldMap, previousLense);
+    }
+  }
+
+  private final AppInfo appInfo;
+
+  public MemberRebindingLense(
+      AppInfo appInfo,
+      Map<DexMethod, DexMethod> methodMap,
+      Map<DexField, DexField> fieldMap,
+      GraphLense previousLense) {
+    super(ImmutableMap.of(), methodMap, fieldMap, previousLense, appInfo.dexItemFactory);
+    this.appInfo = appInfo;
+  }
+
+  public static Builder builder(AppInfo appInfo) {
+    return new Builder(appInfo);
+  }
+
+
+  @Override
+  protected Type mapInvocationType(
+      DexMethod newMethod, DexMethod originalMethod, DexEncodedMethod context, Type type) {
+    return super.mapVirtualInterfaceInvocationTypes(
+        appInfo, newMethod, originalMethod, context, type);
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/shaking/VerticalClassMerger.java b/src/main/java/com/android/tools/r8/shaking/VerticalClassMerger.java
index c737055..ddb8efa 100644
--- a/src/main/java/com/android/tools/r8/shaking/VerticalClassMerger.java
+++ b/src/main/java/com/android/tools/r8/shaking/VerticalClassMerger.java
@@ -308,7 +308,7 @@
 
     // The resulting graph lense that should be used after class merging.
     VerticalClassMergerGraphLense.Builder renamedMembersLense =
-        new VerticalClassMergerGraphLense.Builder();
+        VerticalClassMergerGraphLense.builder(appInfo);
 
     Iterator<DexProgramClass> classIterator = classes.iterator();
 
@@ -431,7 +431,7 @@
     private final DexClass source;
     private final DexClass target;
     private final VerticalClassMergerGraphLense.Builder deferredRenamings =
-        new VerticalClassMergerGraphLense.Builder();
+        VerticalClassMergerGraphLense.builder(appInfo);
     private boolean abortMerge = false;
 
     private ClassMerger(DexClass source, DexClass target) {
diff --git a/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java b/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java
index 942a857..99fb63a 100644
--- a/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java
+++ b/src/main/java/com/android/tools/r8/shaking/VerticalClassMergerGraphLense.java
@@ -4,11 +4,13 @@
 
 package com.android.tools.r8.shaking;
 
+import com.android.tools.r8.graph.AppInfo;
 import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.GraphLense;
+import com.android.tools.r8.graph.GraphLense.NestedGraphLense;
 import com.android.tools.r8.ir.code.Invoke.Type;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -40,47 +42,57 @@
 // invocation will hit the same implementation as the original super.m() call.
 //
 // For the invocation "invoke-virtual A.m()" in B.m2, this graph lense will return the method B.m.
-public class VerticalClassMergerGraphLense extends GraphLense {
-  private final GraphLense previousLense;
+public class VerticalClassMergerGraphLense extends NestedGraphLense {
+  private final AppInfo appInfo;
 
-  private final Map<DexField, DexField> fieldMap;
-  private final Map<DexMethod, DexMethod> methodMap;
   private final Set<DexMethod> mergedMethods;
   private final Map<DexType, Map<DexMethod, DexMethod>> contextualVirtualToDirectMethodMaps;
 
   public VerticalClassMergerGraphLense(
+      AppInfo appInfo,
       Map<DexField, DexField> fieldMap,
       Map<DexMethod, DexMethod> methodMap,
       Set<DexMethod> mergedMethods,
       Map<DexType, Map<DexMethod, DexMethod>> contextualVirtualToDirectMethodMaps,
       GraphLense previousLense) {
-    this.previousLense = previousLense;
-    this.fieldMap = fieldMap;
-    this.methodMap = methodMap;
+    super(ImmutableMap.of(), methodMap, fieldMap, previousLense, appInfo.dexItemFactory);
+    this.appInfo = appInfo;
     this.mergedMethods = mergedMethods;
     this.contextualVirtualToDirectMethodMaps = contextualVirtualToDirectMethodMaps;
   }
 
-  @Override
-  public DexType lookupType(DexType type) {
-    return previousLense.lookupType(type);
+  public static Builder builder(AppInfo appInfo) {
+    return new Builder(appInfo);
   }
 
   @Override
-  public DexMethod lookupMethod(DexMethod method, DexEncodedMethod context, Type type) {
+  public GraphLenseLookupResult lookupMethod(
+      DexMethod method, DexEncodedMethod context, Type type) {
     assert isContextFreeForMethod(method) || (context != null && type != null);
-    DexMethod previous = previousLense.lookupMethod(method, context, type);
-    if (type == Type.SUPER && !mergedMethods.contains(context.method)) {
+    GraphLenseLookupResult previous = previousLense.lookupMethod(method, context, type);
+    if (previous.getType() == Type.SUPER && !mergedMethods.contains(context.method)) {
       Map<DexMethod, DexMethod> virtualToDirectMethodMap =
           contextualVirtualToDirectMethodMaps.get(context.method.holder);
       if (virtualToDirectMethodMap != null) {
-        DexMethod directMethod = virtualToDirectMethodMap.get(previous);
+        DexMethod directMethod = virtualToDirectMethodMap.get(previous.getMethod());
         if (directMethod != null) {
-          return directMethod;
+          // If the super class A of the enclosing class B (i.e., context.method.holder)
+          // has been merged into B during vertical class merging, and this invoke-super instruction
+          // was resolving to a method in A, then the target method has been changed to a direct
+          // method and moved into B, so that we need to use an invoke-direct instruction instead of
+          // invoke-super.
+          return new GraphLenseLookupResult(directMethod, Type.DIRECT);
         }
       }
     }
-    return methodMap.getOrDefault(previous, previous);
+    return super.lookupMethod(previous.getMethod(), context, previous.getType());
+  }
+
+  @Override
+  protected Type mapInvocationType(
+      DexMethod newMethod, DexMethod originalMethod, DexEncodedMethod context, Type type) {
+    return super.mapVirtualInterfaceInvocationTypes(
+        appInfo, newMethod, originalMethod, context, type);
   }
 
   @Override
@@ -100,12 +112,6 @@
   }
 
   @Override
-  public DexField lookupField(DexField field) {
-    DexField previous = previousLense.lookupField(field);
-    return fieldMap.getOrDefault(previous, previous);
-  }
-
-  @Override
   public boolean isContextFreeForMethods() {
     return contextualVirtualToDirectMethodMaps.isEmpty() && previousLense.isContextFreeForMethods();
   }
@@ -126,6 +132,7 @@
   }
 
   public static class Builder {
+    private final AppInfo appInfo;
 
     private final ImmutableMap.Builder<DexField, DexField> fieldMapBuilder = ImmutableMap.builder();
     private final ImmutableMap.Builder<DexMethod, DexMethod> methodMapBuilder =
@@ -134,7 +141,9 @@
     private final Map<DexType, Map<DexMethod, DexMethod>> contextualVirtualToDirectMethodMaps =
         new HashMap<>();
 
-    public Builder() {}
+    private Builder(AppInfo appInfo) {
+      this.appInfo = appInfo;
+    }
 
     public GraphLense build(GraphLense previousLense) {
       Map<DexField, DexField> fieldMap = fieldMapBuilder.build();
@@ -145,6 +154,7 @@
         return previousLense;
       }
       return new VerticalClassMergerGraphLense(
+          appInfo,
           fieldMap,
           methodMap,
           mergedMethodsBuilder.build(),