Fixing METHOD_PROTO and DualReferenceInstruction writing via class interning in DexPool

1 files changed, 33 insertions, 21 deletions

diff --git a/dexlib2/src/main/java/org/jf/dexlib2/writer/pool/ClassPool.java b/dexlib2/src/main/java/org/jf/dexlib2/writer/pool/ClassPool.java
index bf51f129..08a60946 100644
--- a/dexlib2/src/main/java/org/jf/dexlib2/writer/pool/ClassPool.java
+++ b/dexlib2/src/main/java/org/jf/dexlib2/writer/pool/ClassPool.java
@@ -44,6 +44,7 @@ import org.jf.dexlib2.builder.MutableMethodImplementation;
 import org.jf.dexlib2.formatter.DexFormatter;
 import org.jf.dexlib2.iface.*;
 import org.jf.dexlib2.iface.debug.*;
+import org.jf.dexlib2.iface.instruction.DualReferenceInstruction;
 import org.jf.dexlib2.iface.instruction.Instruction;
 import org.jf.dexlib2.iface.instruction.ReferenceInstruction;
 import org.jf.dexlib2.iface.reference.*;
@@ -133,27 +134,12 @@ public class ClassPool extends BasePool<String, PoolClassDef> implements ClassSe
             for (Instruction instruction: methodImpl.getInstructions()) {
                 hasInstruction = true;
                 if (instruction instanceof ReferenceInstruction) {
-                    Reference reference = ((ReferenceInstruction)instruction).getReference();
-                    switch (instruction.getOpcode().referenceType) {
-                        case ReferenceType.STRING:
-                            dexPool.stringSection.intern((StringReference)reference);
-                            break;
-                        case ReferenceType.TYPE:
-                            dexPool.typeSection.intern(((TypeReference)reference).getType());
-                            break;
-                        case ReferenceType.FIELD:
-                            dexPool.fieldSection.intern((FieldReference) reference);
-                            break;
-                        case ReferenceType.METHOD:
-                            dexPool.methodSection.intern((MethodReference)reference);
-                            break;
-                        case ReferenceType.CALL_SITE:
-                            dexPool.callSiteSection.intern((CallSiteReference) reference);
-                            break;
-                        default:
-                            throw new ExceptionWithContext("Unrecognized reference type: %d",
-                                    instruction.getOpcode().referenceType);
-                    }
+                    ReferenceInstruction refInst = (ReferenceInstruction)instruction;
+                    internReference(refInst.getReference(), refInst.getReferenceType());
+                }
+                if (instruction instanceof DualReferenceInstruction) {
+                    DualReferenceInstruction dualRefInst = (DualReferenceInstruction)instruction;
+                    internReference(dualRefInst.getReference2(), dualRefInst.getReferenceType2());
                 }
             }
 
@@ -170,6 +156,32 @@ public class ClassPool extends BasePool<String, PoolClassDef> implements ClassSe
         }
     }
 
+    private void internReference(@Nonnull Reference reference, int referenceType) {
+        switch (referenceType) {
+            case ReferenceType.STRING:
+                dexPool.stringSection.intern((StringReference)reference);
+                break;
+            case ReferenceType.TYPE:
+                dexPool.typeSection.intern(((TypeReference)reference).getType());
+                break;
+            case ReferenceType.FIELD:
+                dexPool.fieldSection.intern((FieldReference) reference);
+                break;
+            case ReferenceType.METHOD:
+                dexPool.methodSection.intern((MethodReference)reference);
+                break;
+            case ReferenceType.METHOD_PROTO:
+                dexPool.protoSection.intern((MethodProtoReference)reference);
+                break;
+            case ReferenceType.CALL_SITE:
+                dexPool.callSiteSection.intern((CallSiteReference) reference);
+                break;
+            default:
+                throw new ExceptionWithContext("Unrecognized reference type: %d",
+                        referenceType);
+        }
+    }
+
     private void internDebug(@Nonnull Method method) {
         for (MethodParameter param: method.getParameters()) {
             String paramName = param.getName();