fix: several improvements for multi-variable type search (#720)

jadx-core/src/main/java/jadx/core/codegen/NameGen.java

@@ -132,6 +132,9 @@ private String makeArgName(CodeVar var) {
 		if (!NameMapper.isValidAndPrintable(varName)) {
 			varName = getFallbackName(var);
 		}
+		if (Consts.DEBUG) {
+			varName += '_' + getFallbackName(var);
+		}
 		return varName;
 	}
 

jadx-core/src/main/java/jadx/core/dex/instructions/args/InsnWrapArg.java

@@ -1,17 +1,24 @@
 package jadx.core.dex.instructions.args;
 
+import java.util.Objects;
+
 import org.jetbrains.annotations.NotNull;
 
+import jadx.core.dex.instructions.ConstStringNode;
+import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.nodes.InsnNode;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
 public final class InsnWrapArg extends InsnArg {
 
 	private final InsnNode wrappedInsn;
 
-	public InsnWrapArg(@NotNull InsnNode insn) {
+	/**
+	 * Use {@link InsnArg#wrapInsnIntoArg(InsnNode)} instead this constructor
+	 */
+	InsnWrapArg(@NotNull InsnNode insn) {
 		RegisterArg result = insn.getResult();
-		this.type = result != null ? result.getType() : ArgType.VOID;
+		this.type = result != null ? result.getType() : ArgType.UNKNOWN;
 		this.wrappedInsn = insn;
 	}
 
@@ -29,7 +36,9 @@ public void setParentInsn(InsnNode parentInsn) {
 
 	@Override
 	public InsnArg duplicate() {
-		return copyCommonParams(new InsnWrapArg(wrappedInsn.copy()));
+		InsnWrapArg copy = new InsnWrapArg(wrappedInsn.copy());
+		copy.setType(type);
+		return copyCommonParams(copy);
 	}
 
 	@Override
@@ -67,6 +76,10 @@ public boolean equals(Object o) {
 
 	@Override
 	public String toString() {
+		if (wrappedInsn.getType() == InsnType.CONST_STR
+				&& Objects.equals(type, ArgType.STRING)) {
+			return "(\"" + ((ConstStringNode) wrappedInsn).getString() + "\")";
+		}
 		return "(wrap: " + type + "\n  " + wrappedInsn + ')';
 	}
 }

jadx-core/src/main/java/jadx/core/dex/instructions/args/SSAVar.java

@@ -175,6 +175,11 @@ public void setCodeVar(@NotNull CodeVar codeVar) {
 		codeVar.addSsaVar(this);
 	}
 
+	public void resetTypeAndCodeVar() {
+		this.typeInfo.reset();
+		this.codeVar = null;
+	}
+
 	public boolean isCodeVarSet() {
 		return codeVar != null;
 	}

jadx-core/src/main/java/jadx/core/dex/visitors/ConstInlineVisitor.java

@@ -87,11 +87,13 @@ private static void checkInsn(MethodNode mth, InsnNode insn, List<InsnNode> toRe
 			String s = ((ConstStringNode) insn).getString();
 			FieldNode f = mth.getParentClass().getConstField(s);
 			if (f == null) {
-				constArg = InsnArg.wrapArg(insn.copy());
+				InsnNode copy = insn.copy();
+				copy.setResult(null);
+				constArg = InsnArg.wrapArg(copy);
 			} else {
 				InsnNode constGet = new IndexInsnNode(InsnType.SGET, f.getFieldInfo(), 0);
-				constGet.setResult(insn.getResult().duplicate());
 				constArg = InsnArg.wrapArg(constGet);
+				constArg.setType(ArgType.STRING);
 			}
 		} else {
 			return;

jadx-core/src/main/java/jadx/core/dex/visitors/InitCodeVariables.java

@@ -31,6 +31,13 @@ public void visit(MethodNode mth) throws JadxException {
 		initCodeVars(mth);
 	}
 
+	public static void rerun(MethodNode mth) {
+		for (SSAVar sVar : mth.getSVars()) {
+			sVar.resetTypeAndCodeVar();
+		}
+		initCodeVars(mth);
+	}
+
 	private static void initCodeVars(MethodNode mth) {
 		for (RegisterArg mthArg : mth.getArguments(true)) {
 			initCodeVar(mthArg.getSVar());
@@ -40,7 +47,7 @@ private static void initCodeVars(MethodNode mth) {
 		}
 	}
 
-	public static void initCodeVar(SSAVar ssaVar) {
+	private static void initCodeVar(SSAVar ssaVar) {
 		if (ssaVar.isCodeVarSet()) {
 			return;
 		}

jadx-core/src/main/java/jadx/core/dex/visitors/SimplifyVisitor.java

@@ -165,7 +165,7 @@ private static void simplifyStringConstructor(RootNode root, ConstructorInsn ins
 						}
 					}
 					if (printable >= arr.length - printable) {
-						InsnWrapArg wa = new InsnWrapArg(new ConstStringNode(new String(arr)));
+						InsnArg wa = InsnArg.wrapArg(new ConstStringNode(new String(arr)));
 						if (insn.getArgsCount() == 1) {
 							insn.setArg(0, wa);
 						} else {
@@ -176,7 +176,7 @@ private static void simplifyStringConstructor(RootNode root, ConstructorInsn ins
 									ArgType.array(ArgType.BYTE));
 							InvokeNode in = new InvokeNode(mi, InvokeType.VIRTUAL, 1);
 							in.addArg(wa);
-							insn.setArg(0, new InsnWrapArg(in));
+							insn.setArg(0, InsnArg.wrapArg(in));
 						}
 					}
 				}

jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockSplitter.java

@@ -37,7 +37,7 @@ public class BlockSplitter extends AbstractVisitor {
 			InsnType.MONITOR_EXIT,
 			InsnType.THROW);
 
-	public static boolean makeSeparate(InsnType insnType) {
+	public static boolean isSeparate(InsnType insnType) {
 		return SEPARATE_INSNS.contains(insnType);
 	}
 
@@ -89,7 +89,7 @@ private static void splitBasicBlocks(MethodNode mth) {
 				InsnType type = prevInsn.getType();
 				if (type == InsnType.GOTO
 						|| type == InsnType.THROW
-						|| makeSeparate(type)) {
+						|| isSeparate(type)) {
 
 					if (type == InsnType.RETURN || type == InsnType.THROW) {
 						mth.addExitBlock(curBlock);
@@ -102,7 +102,7 @@ private static void splitBasicBlocks(MethodNode mth) {
 					startNew = true;
 				} else {
 					startNew = isSplitByJump(prevInsn, insn)
-							|| makeSeparate(insn.getType())
+							|| isSeparate(insn.getType())
 							|| isDoWhile(blocksMap, curBlock, insn)
 							|| insn.contains(AType.EXC_HANDLER)
 							|| prevInsn.contains(AFlag.TRY_LEAVE)

jadx-core/src/main/java/jadx/core/dex/visitors/regions/variables/ProcessVariables.java

@@ -77,7 +77,7 @@ private void checkCodeVars(MethodNode mth, List<CodeVar> codeVars) {
 						.forEach(ssaVar -> {
 							ArgType ssaType = ssaVar.getAssign().getInitType();
 							if (ssaType.isTypeKnown()) {
-								TypeCompare comparator = mth.root().getTypeUpdate().getComparator();
+								TypeCompare comparator = mth.root().getTypeUpdate().getTypeCompare();
 								TypeCompareEnum result = comparator.compareTypes(ssaType, codeVarType);
 								if (result == TypeCompareEnum.CONFLICT || result.isNarrow()) {
 									mth.addWarn("Incorrect type for immutable var: ssa=" + ssaType

jadx-core/src/main/java/jadx/core/dex/visitors/typeinference/TypeCompare.java

@@ -23,11 +23,13 @@ public class TypeCompare {
 	private static final Logger LOG = LoggerFactory.getLogger(TypeCompare.class);
 
 	private final RootNode root;
-	private final ArgTypeComparator comparator;
+	private final Comparator<ArgType> comparator;
+	private final Comparator<ArgType> reversedComparator;
 
 	public TypeCompare(RootNode root) {
 		this.root = root;
 		this.comparator = new ArgTypeComparator();
+		this.reversedComparator = comparator.reversed();
 	}
 
 	/**
@@ -192,10 +194,14 @@ private TypeCompareEnum compareGenericTypeWithObject(ArgType genericType, ArgTyp
 		return NARROW_BY_GENERIC;
 	}
 
-	public ArgTypeComparator getComparator() {
+	public Comparator<ArgType> getComparator() {
 		return comparator;
 	}
 
+	public Comparator<ArgType> getReversedComparator() {
+		return reversedComparator;
+	}
+
 	private final class ArgTypeComparator implements Comparator<ArgType> {
 		@Override
 		public int compare(ArgType a, ArgType b) {

jadx-core/src/main/java/jadx/core/dex/visitors/typeinference/TypeCompareEnum.java

@@ -31,6 +31,10 @@ public TypeCompareEnum invert() {
 		}
 	}
 
+	public boolean isEqual() {
+		return this == EQUAL;
+	}
+
 	public boolean isWider() {
 		return this == WIDER || this == WIDER_BY_GENERIC;
 	}

jadx-core/src/main/java/jadx/core/dex/visitors/typeinference/TypeInferenceVisitor.java

@@ -68,7 +68,33 @@ public void visit(MethodNode mth) {
 		if (mth.isNoCode()) {
 			return;
 		}
-		// collect initial type bounds from assign and usages
+		boolean resolved = runTypePropagation(mth);
+		if (!resolved) {
+			boolean moveAdded = false;
+			for (SSAVar var : new ArrayList<>(mth.getSVars())) {
+				moveAdded |= tryInsertAdditionalInsn(mth, var);
+			}
+			if (moveAdded) {
+				InitCodeVariables.rerun(mth);
+				resolved = runTypePropagation(mth);
+			}
+			if (!resolved) {
+				resolved = runMultiVariableSearch(mth);
+			}
+		}
+		if (resolved) {
+			for (SSAVar var : new ArrayList<>(mth.getSVars())) {
+				processIncompatiblePrimitives(mth, var);
+			}
+		}
+	}
+
+	/**
+	 * Guess type from usage and try to set it to current variable
+	 * and all connected instructions with {@link TypeUpdate#apply(SSAVar, ArgType)}
+	 */
+	private boolean runTypePropagation(MethodNode mth) {
+		// collect initial type bounds from assign and usages`
 		mth.getSVars().forEach(this::attachBounds);
 		mth.getSVars().forEach(this::mergePhiBounds);
 
@@ -86,27 +112,20 @@ public void visit(MethodNode mth) {
 				resolved = false;
 			}
 		}
-		if (resolved) {
-			for (SSAVar var : new ArrayList<>(mth.getSVars())) {
-				processIncompatiblePrimitives(mth, var);
-			}
-		} else {
-			for (SSAVar var : new ArrayList<>(mth.getSVars())) {
-				tryInsertAdditionalInsn(mth, var);
-			}
-			runMultiVariableSearch(mth);
-		}
+		return resolved;
 	}
 
-	private void runMultiVariableSearch(MethodNode mth) {
+	private boolean runMultiVariableSearch(MethodNode mth) {
 		TypeSearch typeSearch = new TypeSearch(mth);
 		try {
 			boolean success = typeSearch.run();
 			if (!success) {
 				mth.addWarn("Multi-variable type inference failed");
 			}
+			return success;
 		} catch (Exception e) {
 			mth.addWarn("Multi-variable type inference failed. Error: " + Utils.getStackTrace(e));
+			return false;
 		}
 	}
 
@@ -186,7 +205,7 @@ private Optional<ArgType> selectBestTypeFromBounds(Set<ITypeBound> bounds) {
 		return bounds.stream()
 				.map(ITypeBound::getType)
 				.filter(Objects::nonNull)
-				.max(typeUpdate.getArgTypeComparator());
+				.max(typeUpdate.getTypeCompare().getComparator());
 	}
 
 	private void attachBounds(SSAVar var) {
@@ -337,9 +356,13 @@ private boolean tryInsertAdditionalInsn(MethodNode mth, SSAVar var) {
 		if (usedInPhiList.isEmpty()) {
 			return false;
 		}
-		if (var.getUseCount() == 1) {
-			InsnNode assignInsn = var.getAssign().getAssignInsn();
-			if (assignInsn != null && assignInsn.getType() == InsnType.MOVE) {
+		InsnNode assignInsn = var.getAssign().getAssignInsn();
+		if (assignInsn != null) {
+			InsnType assignType = assignInsn.getType();
+			if (assignType == InsnType.CONST) {
+				return false;
+			}
+			if (assignType == InsnType.MOVE && var.getUseCount() == 1) {
 				return false;
 			}
 		}
@@ -358,11 +381,16 @@ private boolean insertMoveForPhi(MethodNode mth, PhiInsn phiInsn, SSAVar var) {
 			if (reg.getSVar() == var) {
 				BlockNode blockNode = phiInsn.getBlockByArgIndex(argIndex);
 				InsnNode lastInsn = BlockUtils.getLastInsn(blockNode);
-				if (lastInsn != null && BlockSplitter.makeSeparate(lastInsn.getType())) {
-					if (Consts.DEBUG) {
-						LOG.warn("Can't insert move for PHI in block with separate insn: {}", lastInsn);
+				if (lastInsn != null && BlockSplitter.isSeparate(lastInsn.getType())) {
+					// can't insert move in block with separate instruction
+					// trying previous block
+					List<BlockNode> preds = blockNode.getPredecessors();
+					if (preds.size() == 1) {
+						blockNode = preds.get(0);
+					} else {
+						mth.addWarn("Failed to insert additional move for type inference");
+						return false;
 					}
-					return false;
 				}
 
 				int regNum = reg.getRegNum();
@@ -377,15 +405,6 @@ private boolean insertMoveForPhi(MethodNode mth, PhiInsn phiInsn, SSAVar var) {
 				blockNode.getInstructions().add(moveInsn);
 
 				phiInsn.replaceArg(reg, reg.duplicate(regNum, newSsaVar));
-
-				attachBounds(var);
-				for (InsnArg phiArg : phiInsn.getArguments()) {
-					attachBounds(((RegisterArg) phiArg).getSVar());
-				}
-				for (InsnArg phiArg : phiInsn.getArguments()) {
-					mergePhiBounds(((RegisterArg) phiArg).getSVar());
-				}
-				InitCodeVariables.initCodeVar(newSsaVar);
 				return true;
 			}
 		}

jadx-core/src/main/java/jadx/core/dex/visitors/typeinference/TypeInfo.java

@@ -25,6 +25,11 @@ public Set<ITypeBound> getBounds() {
 		return bounds;
 	}
 
+	public void reset() {
+		this.type = ArgType.UNKNOWN;
+		this.bounds.clear();
+	}
+
 	@Override
 	public String toString() {
 		return "TypeInfo{type=" + type + ", bounds=" + bounds + '}';