Merge "Rewrite method handles via lense lookup results."
diff --git a/src/main/java/com/android/tools/r8/graph/DexMethodHandle.java b/src/main/java/com/android/tools/r8/graph/DexMethodHandle.java
index cd3c093..bfdc1b7 100644
--- a/src/main/java/com/android/tools/r8/graph/DexMethodHandle.java
+++ b/src/main/java/com/android/tools/r8/graph/DexMethodHandle.java
@@ -177,7 +177,8 @@
         case INVOKE_SUPER:
           return Type.SUPER;
         default:
-          throw new Unreachable("DexMethodHandle with unexpected type: " + this);
+          throw new Unreachable(
+              "Conversion to invoke type with unexpected method handle: " + this);
       }
     }
   }
diff --git a/src/main/java/com/android/tools/r8/ir/code/Invoke.java b/src/main/java/com/android/tools/r8/ir/code/Invoke.java
index 819bb4b..fc5ee05 100644
--- a/src/main/java/com/android/tools/r8/ir/code/Invoke.java
+++ b/src/main/java/com/android/tools/r8/ir/code/Invoke.java
@@ -12,6 +12,7 @@
 import com.android.tools.r8.graph.AppInfo;
 import com.android.tools.r8.graph.DexItem;
 import com.android.tools.r8.graph.DexMethod;
+import com.android.tools.r8.graph.DexMethodHandle.MethodHandleType;
 import com.android.tools.r8.graph.DexProto;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.ir.analysis.type.TypeLatticeElement;
@@ -30,7 +31,29 @@
     NEW_ARRAY,
     MULTI_NEW_ARRAY,
     CUSTOM,
-    POLYMORPHIC
+    POLYMORPHIC;
+
+    public MethodHandleType toMethodHandle(DexMethod targetMethod) {
+      switch (this) {
+        case STATIC:
+          return MethodHandleType.INVOKE_STATIC;
+        case VIRTUAL:
+          return MethodHandleType.INVOKE_INSTANCE;
+        case DIRECT:
+          if (targetMethod.name.toString().equals("<init>")) {
+            return MethodHandleType.INVOKE_CONSTRUCTOR;
+          } else {
+            return MethodHandleType.INVOKE_DIRECT;
+          }
+        case INTERFACE:
+          return MethodHandleType.INVOKE_INTERFACE;
+        case SUPER:
+          return MethodHandleType.INVOKE_SUPER;
+        default:
+          throw new Unreachable(
+              "Conversion to method handle with unexpected invoke type: " + this);
+      }
+    }
   }
 
   public Invoke(Value result, List<Value> arguments) {
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 6b016b9..9a88674 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
@@ -78,13 +78,13 @@
           DexProto newMethodProto =
               appInfo.dexItemFactory.applyClassMappingToProto(
                   callSite.methodProto, graphLense::lookupType, protoFixupCache);
-          DexMethodHandle newBootstrapMethod = rewriteDexMethodHandle(method,
-              callSite.bootstrapMethod);
+          DexMethodHandle newBootstrapMethod =
+              rewriteDexMethodHandle(callSite.bootstrapMethod, method);
           List<DexValue> newArgs = callSite.bootstrapArgs.stream().map(
               (arg) -> {
                 if (arg instanceof DexValueMethodHandle) {
                   return new DexValueMethodHandle(
-                      rewriteDexMethodHandle(method, ((DexValueMethodHandle) arg).value));
+                      rewriteDexMethodHandle(((DexValueMethodHandle) arg).value, method));
                 }
                 return arg;
               })
@@ -226,16 +226,17 @@
   }
 
   private DexMethodHandle rewriteDexMethodHandle(
-      DexEncodedMethod method, DexMethodHandle methodHandle) {
+      DexMethodHandle methodHandle, DexEncodedMethod context) {
     if (methodHandle.isMethodHandle()) {
       DexMethod invokedMethod = methodHandle.asMethod();
+      MethodHandleType oldType = methodHandle.type;
       GraphLenseLookupResult lenseLookup =
-          graphLense.lookupMethod(invokedMethod, method, methodHandle.type.toInvokeType());
+          graphLense.lookupMethod(invokedMethod, context, oldType.toInvokeType());
       DexMethod actualTarget = lenseLookup.getMethod();
+      MethodHandleType newType = lenseLookup.getType().toMethodHandle(actualTarget);
       if (actualTarget != invokedMethod) {
-        MethodHandleType newType = methodHandle.type;
         DexClass clazz = appInfo.definitionFor(actualTarget.holder);
-        if (clazz != null && (newType.isInvokeInterface() || newType.isInvokeInstance())) {
+        if (clazz != null && (oldType.isInvokeInterface() || oldType.isInvokeInstance())) {
           newType =
               lenseLookup.getType() == Type.INTERFACE
                   ? MethodHandleType.INVOKE_INTERFACE
@@ -243,6 +244,9 @@
         }
         return new DexMethodHandle(newType, actualTarget);
       }
+      if (oldType != newType) {
+        return new DexMethodHandle(newType, actualTarget);
+      }
     } else {
       DexField field = methodHandle.asField();
       DexField actualField = graphLense.lookupField(field);
diff --git a/src/test/java/com/android/tools/r8/classmerging/ClassMergingTest.java b/src/test/java/com/android/tools/r8/classmerging/ClassMergingTest.java
index 9ca9290..59accb3 100644
--- a/src/test/java/com/android/tools/r8/classmerging/ClassMergingTest.java
+++ b/src/test/java/com/android/tools/r8/classmerging/ClassMergingTest.java
@@ -45,8 +45,7 @@
 import org.junit.runner.RunWith;
 
 // TODO(christofferqa): Add tests to check that statically typed invocations on method handles
-// continue to work after class merging. Rewriting of method handles should be carried out by
-// LensCodeRewriter.rewriteDexMethodHandle.
+// continue to work after class merging.
 @RunWith(VmTestRunner.class)
 public class ClassMergingTest extends TestBase {