Fix argument offsets in GeneratedMessageLiteShrinker

Change-Id: I33e082eca17f58a764b6f967c052cfa0b09ce890
diff --git a/src/main/java/com/android/tools/r8/graph/Descriptor.java b/src/main/java/com/android/tools/r8/graph/Descriptor.java
index 09bef34..f0ba1b6 100644
--- a/src/main/java/com/android/tools/r8/graph/Descriptor.java
+++ b/src/main/java/com/android/tools/r8/graph/Descriptor.java
@@ -6,6 +6,8 @@
 public abstract class Descriptor<T extends DexItem, S extends Descriptor<T,S>>
     extends DexReference implements PresortedComparable<S> {
 
+  public abstract boolean match(S entry);
+
   public abstract boolean match(T entry);
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/graph/DexField.java b/src/main/java/com/android/tools/r8/graph/DexField.java
index 1f106b7..dfd3203 100644
--- a/src/main/java/com/android/tools/r8/graph/DexField.java
+++ b/src/main/java/com/android/tools/r8/graph/DexField.java
@@ -118,8 +118,13 @@
   }
 
   @Override
-  public boolean match(DexEncodedField entry) {
-    return entry.field.name == name && entry.field.type == type;
+  public boolean match(DexField field) {
+    return field.name == name && field.type == type;
+  }
+
+  @Override
+  public boolean match(DexEncodedField encodedField) {
+    return match(encodedField.field);
   }
 
   public String qualifiedName() {
diff --git a/src/main/java/com/android/tools/r8/graph/DexMethod.java b/src/main/java/com/android/tools/r8/graph/DexMethod.java
index 332d0ea..4974dbf 100644
--- a/src/main/java/com/android/tools/r8/graph/DexMethod.java
+++ b/src/main/java/com/android/tools/r8/graph/DexMethod.java
@@ -148,8 +148,13 @@
   }
 
   @Override
-  public boolean match(DexEncodedMethod entry) {
-    return entry.method.name == name && entry.method.proto == proto;
+  public boolean match(DexMethod method) {
+    return method.name == name && method.proto == proto;
+  }
+
+  @Override
+  public boolean match(DexEncodedMethod encodedMethod) {
+    return match(encodedMethod.method);
   }
 
   public String qualifiedName() {
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedMessageLiteShrinker.java b/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedMessageLiteShrinker.java
index e447e7a..4429aa0 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedMessageLiteShrinker.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/proto/GeneratedMessageLiteShrinker.java
@@ -20,11 +20,11 @@
 import com.android.tools.r8.ir.code.Instruction;
 import com.android.tools.r8.ir.code.InstructionListIterator;
 import com.android.tools.r8.ir.code.InvokeMethod;
-import com.android.tools.r8.ir.code.InvokeStatic;
 import com.android.tools.r8.ir.code.MemberType;
 import com.android.tools.r8.ir.code.NewArrayEmpty;
 import com.android.tools.r8.ir.code.Value;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.BooleanUtils;
 import java.util.List;
 
 public class GeneratedMessageLiteShrinker {
@@ -74,21 +74,18 @@
       return;
     }
 
-    InvokeMethod newMessageInfoInvoke = null;
-    for (Instruction instruction : code.instructions()) {
-      if (instruction.isInvokeStatic()) {
-        InvokeStatic invoke = instruction.asInvokeStatic();
-        if (invoke.getInvokedMethod() == references.newMessageInfoMethod
-            || invoke.getInvokedMethod() == references.rawMessageInfoConstructor) {
-          newMessageInfoInvoke = invoke;
-          break;
-        }
-      }
-    }
-
+    InvokeMethod newMessageInfoInvoke = getNewMessageInfoInvoke(code);
     if (newMessageInfoInvoke != null) {
-      Value infoValue = newMessageInfoInvoke.inValues().get(1).getAliasedValue();
-      Value objectsValue = newMessageInfoInvoke.inValues().get(2).getAliasedValue();
+      // If this invoke is targeting RawMessageInfo.<init>(...) then `info` and `objects` is at
+      // positions 2 and 3, respectively, and not position 1 and 2 as when calling the static method
+      // GeneratedMessageLite.newMessageInfo().
+      int adjustment = BooleanUtils.intValue(newMessageInfoInvoke.isInvokeDirect());
+      assert adjustment == 0
+          ? newMessageInfoInvoke.getInvokedMethod().match(references.newMessageInfoMethod)
+          : newMessageInfoInvoke.getInvokedMethod() == references.rawMessageInfoConstructor;
+
+      Value infoValue = newMessageInfoInvoke.inValues().get(1 + adjustment).getAliasedValue();
+      Value objectsValue = newMessageInfoInvoke.inValues().get(2 + adjustment).getAliasedValue();
 
       // Decode the arguments passed to newMessageInfo().
       ProtoMessageInfo protoMessageInfo = decoder.run(infoValue, objectsValue, context);
@@ -159,11 +156,25 @@
       }
     }
 
-    // Pass the newly created `objects` array to newMessageInfo().
-    newMessageInfoInvoke.replaceValue(2, newObjectsValue);
+    // Pass the newly created `objects` array to RawMessageInfo.<init>(...) or
+    // GeneratedMessageLite.newMessageInfo().
+    int adjustment = BooleanUtils.intValue(newMessageInfoInvoke.isInvokeDirect());
+    newMessageInfoInvoke.replaceValue(2 + adjustment, newObjectsValue);
 
     if (hasIntroducedIdentifierNameString) {
       method.getMutableOptimizationInfo().markUseIdentifierNameString();
     }
   }
+
+  private InvokeMethod getNewMessageInfoInvoke(IRCode code) {
+    for (Instruction instruction : code.instructions()) {
+      if (instruction.isInvokeMethod()) {
+        InvokeMethod invoke = instruction.asInvokeMethod();
+        if (references.isMessageInfoConstructionMethod(invoke.getInvokedMethod())) {
+          return invoke;
+        }
+      }
+    }
+    return null;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/proto/ProtoReferences.java b/src/main/java/com/android/tools/r8/ir/analysis/proto/ProtoReferences.java
index 85b1590..da66929 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/proto/ProtoReferences.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/proto/ProtoReferences.java
@@ -82,4 +82,8 @@
     return method.proto == findLiteExtensionByNumberProto
         && method.name.startsWith(findLiteExtensionByNumberName);
   }
+
+  public boolean isMessageInfoConstructionMethod(DexMethod method) {
+    return method.match(newMessageInfoMethod) || method == rawMessageInfoConstructor;
+  }
 }