fix: correct arg replace in PHI instruction (#462)

jadx-core/src/main/java/jadx/core/dex/instructions/PhiInsn.java

@@ -1,9 +1,10 @@
 package jadx.core.dex.instructions;
 
-import java.util.LinkedHashMap;
-import java.util.Map;
+import java.util.ArrayList;
+import java.util.List;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.instructions.args.ArgType;
@@ -17,11 +18,12 @@
 
 public final class PhiInsn extends InsnNode {
 
-	private final Map<RegisterArg, BlockNode> blockBinds;
+	// map arguments to blocks (in same order as in arguments list)
+	private final List<BlockNode> blockBinds;
 
 	public PhiInsn(int regNum, int predecessors) {
 		super(InsnType.PHI, predecessors);
-		this.blockBinds = new LinkedHashMap<>(predecessors);
+		this.blockBinds = new ArrayList<>(predecessors);
 		setResult(InsnArg.reg(regNum, ArgType.UNKNOWN));
 		add(AFlag.DONT_INLINE);
 		add(AFlag.DONT_GENERATE);
@@ -34,19 +36,24 @@ public RegisterArg bindArg(BlockNode pred) {
 	}
 
 	public void bindArg(RegisterArg arg, BlockNode pred) {
-		if (blockBinds.containsValue(pred)) {
+		if (blockBinds.contains(pred)) {
 			throw new JadxRuntimeException("Duplicate predecessors in PHI insn: " + pred + ", " + this);
 		}
-		addArg(arg);
-		blockBinds.put(arg, pred);
+		super.addArg(arg);
+		blockBinds.add(pred);
 	}
 
+	@Nullable
 	public BlockNode getBlockByArg(RegisterArg arg) {
-		return blockBinds.get(arg);
+		int index = getArgIndex(arg);
+		if (index == -1) {
+			return null;
+		}
+		return blockBinds.get(index);
 	}
 
-	public Map<RegisterArg, BlockNode> getBlockBinds() {
-		return blockBinds;
+	public BlockNode getBlockByArgIndex(int argIndex) {
+		return blockBinds.get(argIndex);
 	}
 
 	@Override
@@ -57,39 +64,52 @@ public RegisterArg getArg(int n) {
 
 	@Override
 	public boolean removeArg(InsnArg arg) {
-		if (!(arg instanceof RegisterArg)) {
+		int index = getArgIndex(arg);
+		if (index == -1) {
 			return false;
 		}
-		RegisterArg reg = (RegisterArg) arg;
-		if (super.removeArg(reg)) {
-			blockBinds.remove(reg);
-			reg.getSVar().removeUse(reg);
-			InsnRemover.fixUsedInPhiFlag(reg);
-			return true;
-		}
-		return false;
+		removeArg(index);
+		return true;
+	}
+
+	@Override
+	protected RegisterArg removeArg(int index) {
+		RegisterArg reg = (RegisterArg) super.removeArg(index);
+		blockBinds.remove(index);
+		InsnRemover.fixUsedInPhiFlag(reg);
+		return reg;
 	}
 
 	@Override
 	public boolean replaceArg(InsnArg from, InsnArg to) {
 		if (!(from instanceof RegisterArg) || !(to instanceof RegisterArg)) {
 			return false;
 		}
-		BlockNode pred = getBlockByArg((RegisterArg) from);
+
+		int argIndex = getArgIndex(from);
+		if (argIndex == -1) {
+			return false;
+		}
+		BlockNode pred = getBlockByArgIndex(argIndex);
 		if (pred == null) {
 			throw new JadxRuntimeException("Unknown predecessor block by arg " + from + " in PHI: " + this);
 		}
-		if (removeArg(from)) {
-			RegisterArg reg = (RegisterArg) to;
-			bindArg(reg, pred);
-			reg.getSVar().setUsedInPhi(this);
-		}
+		removeArg(argIndex);
+
+		RegisterArg reg = (RegisterArg) to;
+		bindArg(reg, pred);
+		reg.getSVar().setUsedInPhi(this);
 		return true;
 	}
 
+	@Override
+	public void addArg(InsnArg arg) {
+		throw new JadxRuntimeException("Direct addArg is forbidden for PHI insn, bindArg must be used");
+	}
+
 	@Override
 	public void setArg(int n, InsnArg arg) {
-		throw new JadxRuntimeException("Unsupported operation for PHI node");
+		throw new JadxRuntimeException("Direct setArg is forbidden for PHI insn, bindArg must be used");
 	}
 
 	@Override

jadx-core/src/main/java/jadx/core/dex/nodes/InsnNode.java

@@ -139,15 +139,29 @@ public boolean replaceArg(InsnArg from, InsnArg to) {
 	}
 
 	protected boolean removeArg(InsnArg arg) {
+		int index = getArgIndex(arg);
+		if (index == -1) {
+			return false;
+		}
+		removeArg(index);
+		return true;
+	}
+
+	protected InsnArg removeArg(int index) {
+		InsnArg arg = arguments.get(index);
+		arguments.remove(index);
+		InsnRemover.unbindArgUsage(null, arg);
+		return arg;
+	}
+
+	protected int getArgIndex(InsnArg arg) {
 		int count = getArgsCount();
 		for (int i = 0; i < count; i++) {
 			if (arg == arguments.get(i)) {
-				arguments.remove(i);
-				InsnRemover.unbindArgUsage(null, arg);
-				return true;
+				return i;
 			}
 		}
-		return false;
+		return -1;
 	}
 
 	protected void addReg(DecodedInstruction insn, int i, ArgType type) {

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

@@ -3,7 +3,6 @@
 import java.util.ArrayList;
 import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -318,10 +317,12 @@ private boolean tryInsertAdditionalInsn(MethodNode mth, SSAVar var) {
 				return false;
 			}
 		}
-		for (Map.Entry<RegisterArg, BlockNode> entry : phiInsn.getBlockBinds().entrySet()) {
-			RegisterArg reg = entry.getKey();
+
+		int argsCount = phiInsn.getArgsCount();
+		for (int argIndex = 0; argIndex < argsCount; argIndex++) {
+			RegisterArg reg = phiInsn.getArg(argIndex);
 			if (reg.getSVar() == var) {
-				BlockNode blockNode = entry.getValue();
+				BlockNode blockNode = phiInsn.getBlockByArgIndex(argIndex);
 				InsnNode lastInsn = BlockUtils.getLastInsn(blockNode);
 				if (lastInsn != null && BlockSplitter.makeSeparate(lastInsn.getType())) {
 					if (Consts.DEBUG) {

jadx-core/src/main/java/jadx/core/utils/DebugUtils.java

@@ -162,9 +162,6 @@ private static void checkPHI(MethodNode mth) {
 				if (insn.getType() == InsnType.PHI) {
 					PhiInsn phi = (PhiInsn) insn;
 					phis.add(phi);
-					if (phi.getArgsCount() != phi.getBlockBinds().size()) {
-						throw new JadxRuntimeException("Incorrect args and binds in PHI");
-					}
 					if (phi.getArgsCount() == 0) {
 						throw new JadxRuntimeException("No args and binds in PHI");
 					}