Infer the register type based on the result of an instance-of instruction

In oat files, it's possible that a check-cast is removed based on the
register being provably of the type being cast to, based on the result
of an instance-of check.

Supporting this in general would require more sophisticated static analysis
than is currently done. Instead, this adds a special-case specifically
for the case of an instance-of followed immediately by an if-eqz

2 files changed, 169 insertions, 39 deletions

diff --git a/dexlib2/src/main/java/org/jf/dexlib2/analysis/AnalyzedInstruction.java b/dexlib2/src/main/java/org/jf/dexlib2/analysis/AnalyzedInstruction.java
index 30cc9067..1f600835 100644
--- a/dexlib2/src/main/java/org/jf/dexlib2/analysis/AnalyzedInstruction.java
+++ b/dexlib2/src/main/java/org/jf/dexlib2/analysis/AnalyzedInstruction.java
@@ -31,12 +31,15 @@
 
 package org.jf.dexlib2.analysis;
 
+import com.google.common.base.Objects;
+import com.google.common.collect.Maps;
 import org.jf.dexlib2.iface.instruction.*;
 import org.jf.dexlib2.iface.reference.MethodReference;
 import org.jf.dexlib2.iface.reference.Reference;
 import org.jf.util.ExceptionWithContext;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import java.util.*;
 
 public class AnalyzedInstruction implements Comparable<AnalyzedInstruction> {
@@ -71,6 +74,12 @@ public class AnalyzedInstruction implements Comparable<AnalyzedInstruction> {
     protected final RegisterType[] postRegisterMap;
 
     /**
+     * This contains optional register type overrides for register types from predecessors
+     */
+    @Nullable
+    protected Map<PredecessorOverrideKey, RegisterType> predecessorRegisterOverrides = null;
+
+    /**
      * When deodexing, we might need to deodex this instruction multiple times, when we merge in new register
      * information. When this happens, we need to restore the original (odexed) instruction, so we can deodex it again
      */
@@ -101,6 +110,17 @@ public class AnalyzedInstruction implements Comparable<AnalyzedInstruction> {
         return Collections.unmodifiableSortedSet(predecessors);
     }
 
+    public RegisterType getPredecessorRegisterType(@Nonnull AnalyzedInstruction predecessor, int registerNumber) {
+        if (predecessorRegisterOverrides != null) {
+            RegisterType override = predecessorRegisterOverrides.get(
+                    new PredecessorOverrideKey(predecessor, registerNumber));
+            if (override != null) {
+                return override;
+            }
+        }
+        return predecessor.postRegisterMap[registerNumber];
+    }
+
     protected boolean addPredecessor(AnalyzedInstruction predecessor) {
         return predecessors.add(predecessor);
     }
@@ -193,39 +213,82 @@ public class AnalyzedInstruction implements Comparable<AnalyzedInstruction> {
 
     /**
      * Iterates over the predecessors of this instruction, and merges all the post-instruction register types for the
-     * given register. Any dead, unreachable, or odexed predecessor is ignored
+     * given register. Any dead, unreachable, or odexed predecessor is ignored. This takes into account any overridden
+     * predecessor register types
+     *
      * @param registerNumber the register number
      * @return The register type resulting from merging the post-instruction register types from all predecessors
      */
-    protected RegisterType mergePreRegisterTypeFromPredecessors(int registerNumber) {
+    protected RegisterType getMergedPreRegisterTypeFromPredecessors(int registerNumber) {
         RegisterType mergedRegisterType = null;
         for (AnalyzedInstruction predecessor: predecessors) {
-            RegisterType predecessorRegisterType = predecessor.postRegisterMap[registerNumber];
-            assert predecessorRegisterType != null;
-            mergedRegisterType = predecessorRegisterType.merge(mergedRegisterType);
+            RegisterType predecessorRegisterType = getPredecessorRegisterType(predecessor, registerNumber);
+            if (predecessorRegisterType != null) {
+                if (mergedRegisterType == null) {
+                    mergedRegisterType = predecessorRegisterType;
+                } else {
+                    mergedRegisterType = predecessorRegisterType.merge(mergedRegisterType);
+                }
+            }
         }
         return mergedRegisterType;
     }
+    /**
+     * Sets the "post-instruction" register type as indicated.
+     * @param registerNumber Which register to set
+     * @param registerType The "post-instruction" register type
+     * @returns true if the given register type is different than the existing post-instruction register type
+     */
+    protected boolean setPostRegisterType(int registerNumber, RegisterType registerType) {
+        assert registerNumber >= 0 && registerNumber < postRegisterMap.length;
+        assert registerType != null;
 
-    /*
-      * Sets the "post-instruction" register type as indicated.
-      * @param registerNumber Which register to set
-      * @param registerType The "post-instruction" register type
-      * @returns true if the given register type is different than the existing post-instruction register type
-      */
-     protected boolean setPostRegisterType(int registerNumber, RegisterType registerType) {
-         assert registerNumber >= 0 && registerNumber < postRegisterMap.length;
-         assert registerType != null;
+        RegisterType oldRegisterType = postRegisterMap[registerNumber];
+        if (oldRegisterType.equals(registerType)) {
+            return false;
+        }
 
-         RegisterType oldRegisterType = postRegisterMap[registerNumber];
-         if (oldRegisterType.equals(registerType)) {
-             return false;
-         }
+        postRegisterMap[registerNumber] = registerType;
+        return true;
+    }
 
-         postRegisterMap[registerNumber] = registerType;
-         return true;
-     }
+    /**
+     * Adds an override for a register type from a predecessor.
+     *
+     * This is used to set the register type for only one branch from a conditional jump.
+     *
+     * @param predecessor Which predecessor is being overriden
+     * @param registerNumber The register number of the register being overriden
+     * @param registerType The overridden register type
+     * @param verifiedInstructions
+     *
+     * @return true if the post-instruction register type for this instruction changed as a result of this override
+     */
+    protected boolean overridePredecessorRegisterType(@Nonnull AnalyzedInstruction predecessor, int registerNumber,
+                                                      @Nonnull RegisterType registerType, BitSet verifiedInstructions) {
+        if (predecessorRegisterOverrides == null) {
+            predecessorRegisterOverrides = Maps.newHashMap();
+        }
+        predecessorRegisterOverrides.put(new PredecessorOverrideKey(predecessor, registerNumber), registerType);
+
+        RegisterType mergedType = getMergedPreRegisterTypeFromPredecessors(registerNumber);
+
+        if (preRegisterMap[registerNumber].equals(mergedType)) {
+            return false;
+        }
+
+        preRegisterMap[registerNumber] = mergedType;
+        verifiedInstructions.clear(instructionIndex);
+
+        if (!setsRegister(registerNumber)) {
+            if (!postRegisterMap[registerNumber].equals(mergedType)) {
+                postRegisterMap[registerNumber] = mergedType;
+                return true;
+            }
+        }
 
+        return false;
+    }
 
     protected boolean isInvokeInit() {
         if (instruction == null || !instruction.getOpcode().canInitializeReference()) {
@@ -328,5 +391,27 @@ public class AnalyzedInstruction implements Comparable<AnalyzedInstruction> {
             return 1;
         }
     }
+
+    private static class PredecessorOverrideKey {
+        public final AnalyzedInstruction analyzedInstruction;
+        public final int registerNumber;
+
+        public PredecessorOverrideKey(AnalyzedInstruction analyzedInstruction, int registerNumber) {
+            this.analyzedInstruction = analyzedInstruction;
+            this.registerNumber = registerNumber;
+        }
+
+        @Override public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            PredecessorOverrideKey that = (PredecessorOverrideKey)o;
+            return com.google.common.base.Objects.equal(registerNumber, that.registerNumber) &&
+                    Objects.equal(analyzedInstruction, that.analyzedInstruction);
+        }
+
+        @Override public int hashCode() {
+            return Objects.hashCode(analyzedInstruction, registerNumber);
+        }
+    }
 }
 
diff --git a/dexlib2/src/main/java/org/jf/dexlib2/analysis/MethodAnalyzer.java b/dexlib2/src/main/java/org/jf/dexlib2/analysis/MethodAnalyzer.java
index 3d0318c4..e634297f 100644
--- a/dexlib2/src/main/java/org/jf/dexlib2/analysis/MethodAnalyzer.java
+++ b/dexlib2/src/main/java/org/jf/dexlib2/analysis/MethodAnalyzer.java
@@ -336,27 +336,16 @@ public class MethodAnalyzer {
                 registerType);
     }
 
-    private void setPostRegisterTypeAndPropagateChanges(@Nonnull AnalyzedInstruction analyzedInstruction,
-                                                        int registerNumber, @Nonnull RegisterType registerType) {
-
-        BitSet changedInstructions = new BitSet(analyzedInstructions.size());
-
-        if (!analyzedInstruction.setPostRegisterType(registerNumber, registerType)) {
-            return;
-        }
-
-        propagateRegisterToSuccessors(analyzedInstruction, registerNumber, changedInstructions);
-
+    private void propagateChanges(@Nonnull BitSet changedInstructions, int registerNumber,
+                                  @Nonnull RegisterType registerType) {
         //Using a for loop inside the while loop optimizes for the common case of the successors of an instruction
         //occurring after the instruction. Any successors that occur prior to the instruction will be picked up on
         //the next iteration of the while loop.
-        //This could also be done recursively, but in large methods it would likely cause very deep recursion,
-        //which requires the user to specify a larger stack size. This isn't really a problem, but it is slightly
-        //annoying.
+        //This could also be done recursively, but in large methods it would likely cause very deep recursion.
         while (!changedInstructions.isEmpty()) {
             for (int instructionIndex=changedInstructions.nextSetBit(0);
-                     instructionIndex>=0;
-                     instructionIndex=changedInstructions.nextSetBit(instructionIndex+1)) {
+                 instructionIndex>=0;
+                 instructionIndex=changedInstructions.nextSetBit(instructionIndex+1)) {
 
                 changedInstructions.clear(instructionIndex);
 
@@ -364,6 +353,44 @@ public class MethodAnalyzer {
                         changedInstructions);
             }
         }
+    }
+
+    private void overridePredecessorRegisterTypeAndPropagateChanges(
+            @Nonnull AnalyzedInstruction analyzedInstruction, @Nonnull AnalyzedInstruction predecessor,
+            int registerNumber, @Nonnull RegisterType registerType) {
+        BitSet changedInstructions = new BitSet(analyzedInstructions.size());
+
+        if (!analyzedInstruction.overridePredecessorRegisterType(
+                predecessor, registerNumber, registerType, analyzedState)) {
+            return;
+        }
+        changedInstructions.set(analyzedInstruction.instructionIndex);
+
+        propagateChanges(changedInstructions, registerNumber, registerType);
+
+        if (registerType.category == RegisterType.LONG_LO) {
+            checkWidePair(registerNumber, analyzedInstruction);
+            overridePredecessorRegisterTypeAndPropagateChanges(analyzedInstruction, predecessor, registerNumber + 1,
+                    RegisterType.LONG_HI_TYPE);
+        } else if (registerType.category == RegisterType.DOUBLE_LO) {
+            checkWidePair(registerNumber, analyzedInstruction);
+            overridePredecessorRegisterTypeAndPropagateChanges(analyzedInstruction, predecessor, registerNumber + 1,
+                    RegisterType.DOUBLE_HI_TYPE);
+        }
+    }
+
+    private void setPostRegisterTypeAndPropagateChanges(@Nonnull AnalyzedInstruction analyzedInstruction,
+                                                        int registerNumber, @Nonnull RegisterType registerType) {
+
+        BitSet changedInstructions = new BitSet(analyzedInstructions.size());
+
+        if (!analyzedInstruction.setPostRegisterType(registerNumber, registerType)) {
+            return;
+        }
+
+        propagateRegisterToSuccessors(analyzedInstruction, registerNumber, changedInstructions);
+
+        propagateChanges(changedInstructions, registerNumber, registerType);
 
         if (registerType.category == RegisterType.LONG_LO) {
             checkWidePair(registerNumber, analyzedInstruction);
@@ -1037,8 +1064,11 @@ public class MethodAnalyzer {
     }
 
     private void analyzeMoveResult(@Nonnull AnalyzedInstruction analyzedInstruction) {
-        AnalyzedInstruction previousInstruction = analyzedInstructions.valueAt(analyzedInstruction.instructionIndex-1);
-        if (!previousInstruction.instruction.getOpcode().setsResult()) {
+        AnalyzedInstruction previousInstruction = null;
+        if (analyzedInstruction.instructionIndex > 0) {
+            previousInstruction = analyzedInstructions.valueAt(analyzedInstruction.instructionIndex-1);
+        }
+        if (previousInstruction == null || !previousInstruction.instruction.getOpcode().setsResult()) {
             throw new AnalysisException(analyzedInstruction.instruction.getOpcode().name + " must occur after an " +
                     "invoke-*/fill-new-array instruction");
         }
@@ -1132,6 +1162,21 @@ public class MethodAnalyzer {
 
     private void analyzeInstanceOf(@Nonnull AnalyzedInstruction analyzedInstruction) {
         setDestinationRegisterTypeAndPropagateChanges(analyzedInstruction, RegisterType.BOOLEAN_TYPE);
+
+        int instructionIndex = analyzedInstruction.getInstructionIndex();
+        AnalyzedInstruction nextAnalyzedInstruction = analyzedInstructions.valueAt(instructionIndex + 1);
+
+        Instruction nextInstruction = nextAnalyzedInstruction.instruction;
+        if (nextInstruction.getOpcode() == Opcode.IF_EQZ) {
+            if (((Instruction21t)nextInstruction).getRegisterA() == analyzedInstruction.getDestinationRegister()) {
+                Reference reference = ((Instruction22c)analyzedInstruction.getInstruction()).getReference();
+                RegisterType registerType = RegisterType.getRegisterType(classPath, (TypeReference)reference);
+                int objectRegister = ((TwoRegisterInstruction)analyzedInstruction.getInstruction()).getRegisterB();
+
+                overridePredecessorRegisterTypeAndPropagateChanges(analyzedInstructions.valueAt(instructionIndex + 2),
+                        nextAnalyzedInstruction, objectRegister, registerType);
+            }
+        }
     }
 
     private void analyzeArrayLength(@Nonnull AnalyzedInstruction analyzedInstruction) {