[DEX] Switch and payload accessing

Author: REAndroid

src/main/java/com/reandroid/dex/ins/InsSwitch.java

@@ -21,42 +21,42 @@
 
 public abstract class InsSwitch extends Ins31t {
 
-    private InsSwitchPayload switchPayload;
+    private InsSwitchPayload<?> switchPayload;
 
     public InsSwitch(Opcode<?> opcode) {
         super(opcode);
     }
 
-    public InsSwitchPayload getPayload() {
-        InsSwitchPayload switchPayload = this.switchPayload;
-        if(switchPayload == null) {
+    public InsSwitchPayload<? extends SwitchEntry> getPayload() {
+        InsSwitchPayload<?> switchPayload = this.switchPayload;
+        if (switchPayload == null) {
             switchPayload = findByAddress();
-            if(switchPayload != null) {
+            if (switchPayload != null) {
                 this.switchPayload = switchPayload;
                 switchPayload.setSwitch(this);
             }
         }
         return switchPayload;
     }
-    public void setPayload(InsSwitchPayload switchPayload) {
+    public void setPayload(InsSwitchPayload<?> switchPayload) {
         this.switchPayload = switchPayload;
         setTargetAddress(switchPayload.getAddress());
     }
 
-    private InsSwitchPayload findByAddress() {
+    private InsSwitchPayload<?> findByAddress() {
         InstructionList instructionList = getInstructionList();
-        if(instructionList != null){
-            Iterator<? extends InsSwitchPayload> iterator = instructionList
+        if (instructionList != null) {
+            Iterator<? extends InsSwitchPayload<?>> iterator = instructionList
                     .iterator(getPayloadOpcode());
             int address = getTargetAddress();
-            while (iterator.hasNext()){
-                InsSwitchPayload switchPayload = iterator.next();
-                if(switchPayload.getAddress() == address){
+            while (iterator.hasNext()) {
+                InsSwitchPayload<?> switchPayload = iterator.next();
+                if (switchPayload.getAddress() == address) {
                     return switchPayload;
                 }
             }
         }
         return null;
     }
-    public abstract Opcode<? extends InsSwitchPayload> getPayloadOpcode();
+    public abstract Opcode<? extends InsSwitchPayload<?>> getPayloadOpcode();
 }

src/main/java/com/reandroid/dex/ins/SwitchEntry.java

@@ -21,6 +21,13 @@ public interface SwitchEntry extends PayloadEntry, Label {
     @Override
     InsSwitchPayload<?> getPayload();
 
+    default InsSwitch getInsSwitch() {
+        InsSwitchPayload<?> payload = getPayload();
+        if (payload != null) {
+            return payload.getSwitch();
+        }
+        return null;
+    }
     @Override
     default boolean isRemoved() {
         return getPayload().isRemoved();
@@ -32,7 +39,7 @@ default void updateTargetAddress() {
     }
 
     default void addEquivalentIfEq(int constRegister) {
-        InsSwitch insSwitch = getPayload().getSwitch();
+        InsSwitch insSwitch = getInsSwitch();
         InstructionList instructionList = insSwitch.getInstructionList();
 
         Ins targetIns = getTargetIns();

src/main/java/com/reandroid/dex/model/DexInstruction.java

@@ -24,14 +24,7 @@
 import com.reandroid.dex.id.IdItem;
 import com.reandroid.dex.id.MethodId;
 import com.reandroid.dex.id.StringId;
-import com.reandroid.dex.ins.ConstNumber;
-import com.reandroid.dex.ins.ConstNumberLong;
-import com.reandroid.dex.ins.Ins;
-import com.reandroid.dex.ins.InsConstStringJumbo;
-import com.reandroid.dex.ins.Label;
-import com.reandroid.dex.ins.Opcode;
-import com.reandroid.dex.ins.RegistersSet;
-import com.reandroid.dex.ins.SizeXIns;
+import com.reandroid.dex.ins.*;
 import com.reandroid.dex.key.FieldKey;
 import com.reandroid.dex.key.Key;
 import com.reandroid.dex.key.MethodKey;
@@ -311,6 +304,15 @@ public boolean isMethodInvokeInterface() {
     public boolean isMethodInvoke() {
         return getOpcode().isMethodInvoke();
     }
+    public boolean isMove() {
+        return getOpcode().isMove();
+    }
+    public boolean isMoveResult() {
+        return getOpcode().isMoveResult();
+    }
+    public boolean hasOutRegisters() {
+        return getOpcode().hasOutRegisters();
+    }
     public int getTargetAddress() {
         Ins ins = getIns();
         if (ins instanceof Label) {
@@ -536,6 +538,36 @@ public Iterator<DexInstruction> getTargetingInstructionsIfOpcode(Predicate<Opcod
         iterator = FilterIterator.of(iterator, ins -> predicate.test(ins.getOpcode()));
         return DexInstruction.createAll(getDexMethod(), iterator);
     }
+    // for switch payload instruction
+    public Iterator<DexInstruction> getTargetSwitchCases() {
+        Ins ins = getIns();
+        if (!(ins instanceof InsSwitchPayload)) {
+            return EmptyIterator.of();
+        }
+        InsSwitchPayload<? extends SwitchEntry> payload = (InsSwitchPayload<? extends SwitchEntry>) ins;
+        Iterator<? extends SwitchEntry> switchEntryIterator = payload.iterator();
+        if (!switchEntryIterator.hasNext()) {
+            return EmptyIterator.of();
+        }
+        Iterator<Ins> iterator = ComputeIterator.of(switchEntryIterator, SwitchEntry::getTargetIns);
+        iterator = CollectionUtil.copyOfUniqueOf(iterator);
+        return DexInstruction.createAll(getDexMethod(), iterator);
+    }
+    public DexInstruction getTargetingSwitch() {
+        return CollectionUtil.getFirst(getTargetingSwitches(null));
+    }
+    public Iterator<DexInstruction> getTargetingSwitches(Opcode<? extends InsSwitch> switchOpcode) {
+        Iterator<SwitchEntry> switchEntryIterator = getIns().getForcedExtraLines(SwitchEntry.class);
+        if (!switchEntryIterator.hasNext()) {
+            return EmptyIterator.of();
+        }
+        Iterator<InsSwitch> iterator = ComputeIterator.of(switchEntryIterator, SwitchEntry::getInsSwitch);
+        if (switchOpcode != null) {
+            iterator = FilterIterator.of(iterator, ins -> ins.getOpcode() == switchOpcode);
+        }
+        iterator = CollectionUtil.copyOfUniqueOf(iterator);
+        return DexInstruction.createAll(getDexMethod(), iterator);
+    }
     public DexInstruction getNext() {
         return getDexMethod().getInstruction(getIndex() + 1);
     }
@@ -548,6 +580,16 @@ public DexInstruction getPrevious() {
     public Iterator<DexInstruction> getPreviousInstructions() {
         return LinkedIterator.of(this, DexInstruction::getPrevious);
     }
+    public DexInstruction getMoveResult() {
+        if (!hasOutRegisters()) {
+            return null;
+        }
+        DexInstruction next = getNext();
+        if (next != null && next.isMoveResult()) {
+            return next;
+        }
+        return null;
+    }
 
     public DexInstruction getPreviousReader(int register) {
         return getPreviousReader(register, CollectionUtil.getAcceptAll());
@@ -643,7 +685,7 @@ public String toString() {
         return getIns().toString();
     }
 
-    public static Iterator<DexInstruction> createAll(DexMethod dexMethod, Iterator<Ins> iterator) {
+    public static Iterator<DexInstruction> createAll(DexMethod dexMethod, Iterator<? extends Ins> iterator) {
         if (dexMethod == null || !iterator.hasNext()) {
             return EmptyIterator.of();
         }