refactor: improve variables handling in instruction wrapping

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

 package jadx.core.dex.instructions.args;
 
-import java.util.ArrayList;
-import java.util.List;
-
 import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -12,7 +9,10 @@ import com.android.dx.io.instructions.DecodedInstruction;
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.nodes.InsnNode;
+import jadx.core.dex.nodes.MethodNode;
+import jadx.core.utils.InsnRemover;
 import jadx.core.utils.InsnUtils;
+import jadx.core.utils.exceptions.JadxRuntimeException;
 
 /**
  * Instruction argument,
@@ -93,7 +93,7 @@ public abstract class InsnArg extends Typed {
 		this.parentInsn = parentInsn;
 	}
 
-	public InsnArg wrapInstruction(InsnNode insn) {
+	public InsnArg wrapInstruction(MethodNode mth, InsnNode insn) {
 		InsnNode parent = parentInsn;
 		if (parent == null) {
 			return null;
@@ -118,20 +118,13 @@ public abstract class InsnArg extends Typed {
 				}
 			}
 		}
-		insn.add(AFlag.WRAPPED);
-		InsnArg arg = wrapArg(insn);
+		InsnArg arg = wrapInsnIntoArg(insn);
 		parent.setArg(i, arg);
+		InsnRemover.unbindArgUsage(mth, this);
+		InsnRemover.unbindResult(mth, insn);
 		return arg;
 	}
 
-	public static void updateParentInsn(InsnNode fromInsn, InsnNode toInsn) {
-		List<RegisterArg> args = new ArrayList<>();
-		fromInsn.getRegisterArgs(args);
-		for (RegisterArg reg : args) {
-			reg.setParentInsn(toInsn);
-		}
-	}
-
 	private static int getArgIndex(InsnNode parent, InsnArg arg) {
 		int count = parent.getArgsCount();
 		for (int i = 0; i < count; i++) {
@@ -142,23 +135,43 @@ public abstract class InsnArg extends Typed {
 		return -1;
 	}
 
-	public static InsnArg wrapArg(InsnNode insn) {
+	public static InsnArg wrapInsnIntoArg(InsnNode insn) {
 		InsnArg arg;
+		InsnType type = insn.getType();
+		if (type == InsnType.CONST || type == InsnType.MOVE) {
+			arg = insn.getArg(0);
+			insn.add(AFlag.REMOVE);
+			insn.add(AFlag.DONT_GENERATE);
+		} else {
+			arg = wrapArg(insn);
+		}
+		return arg;
+	}
+
+	/**
+	 * Prefer {@link InsnArg#wrapInsnIntoArg}.
+	 * This method don't support MOVE and CONST insns!
+	 */
+	public static InsnArg wrapArg(InsnNode insn) {
+		InsnArg arg = wrap(insn);
+		insn.add(AFlag.WRAPPED);
 		switch (insn.getType()) {
-			case MOVE:
 			case CONST:
-				arg = insn.getArg(0);
-				break;
+			case MOVE:
+				throw new JadxRuntimeException("Don't wrap MOVE or CONST insns: " + insn);
+
 			case CONST_STR:
-				arg = wrap(insn);
 				arg.setType(ArgType.STRING);
 				break;
 			case CONST_CLASS:
-				arg = wrap(insn);
 				arg.setType(ArgType.CLASS);
 				break;
+
 			default:
-				arg = wrap(insn);
+				RegisterArg resArg = insn.getResult();
+				if (resArg != null) {
+					arg.setType(resArg.getType());
+				}
 				break;
 		}
 		return arg;

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

 package jadx.core.dex.visitors;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 
@@ -130,8 +131,8 @@ public class ClassModifier extends AbstractVisitor {
 		if (arg.getSVar().getUseCount() != 0) {
 			InsnNode iget = new IndexInsnNode(InsnType.IGET, fieldInfo, 1);
 			iget.addArg(insn.getArg(1));
-			for (InsnArg insnArg : arg.getSVar().getUseList()) {
-				insnArg.wrapInstruction(iget);
+			for (InsnArg insnArg : new ArrayList<>(arg.getSVar().getUseList())) {
+				insnArg.wrapInstruction(mth, iget);
 			}
 		}
 		return true;

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

@@ -23,7 +23,6 @@ import jadx.core.dex.visitors.ssa.SSATransform;
 import jadx.core.dex.visitors.typeinference.TypeInferenceVisitor;
 import jadx.core.utils.InsnRemover;
 import jadx.core.utils.exceptions.JadxException;
-import jadx.core.utils.exceptions.JadxOverflowException;
 
 @JadxVisitor(
 		name = "Constants Inline",
@@ -195,7 +194,7 @@ public class ConstInlineVisitor extends AbstractVisitor {
 				fieldNode = mth.getParentClass().getConstField((int) literal, false);
 			}
 			if (fieldNode != null) {
-				litArg.wrapInstruction(new IndexInsnNode(InsnType.SGET, fieldNode.getFieldInfo(), 0));
+				litArg.wrapInstruction(mth, new IndexInsnNode(InsnType.SGET, fieldNode.getFieldInfo(), 0));
 			}
 		} else {
 			if (!useInsn.replaceArg(arg, constArg.duplicate())) {
@@ -204,12 +203,6 @@ public class ConstInlineVisitor extends AbstractVisitor {
 		}
 		if (insnType == InsnType.RETURN) {
 			useInsn.setSourceLine(constInsn.getSourceLine());
-		} else if (insnType == InsnType.MOVE) {
-			try {
-				replaceConst(mth, useInsn, constArg, toRemove);
-			} catch (StackOverflowError e) {
-				throw new JadxOverflowException("Stack overflow at const inline visitor");
-			}
 		}
 		return true;
 	}

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

@@ -120,7 +120,7 @@ public class ModVisitor extends AbstractVisitor {
 						break;
 
 					case ARITH:
-						processArith(parentClass, (ArithNode) insn);
+						processArith(mth, parentClass, (ArithNode) insn);
 						break;
 
 					case CHECK_CAST:
@@ -187,7 +187,7 @@ public class ModVisitor extends AbstractVisitor {
 		}
 	}
 
-	private static void processArith(ClassNode parentClass, ArithNode arithNode) {
+	private static void processArith(MethodNode mth, ClassNode parentClass, ArithNode arithNode) {
 		if (arithNode.getArgsCount() != 2) {
 			throw new JadxRuntimeException("Invalid args count in insn: " + arithNode);
 		}

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

@@ -136,14 +136,18 @@ public class ReSugarCode extends AbstractVisitor {
 		ArgType arrType = newArrayInsn.getArrayType();
 		InsnNode filledArr = new FilledNewArrayNode(arrType.getArrayElement(), len);
 		filledArr.setResult(arrArg);
+
+		InsnNode lastPut = Utils.last(arrPuts);
 		for (InsnNode put : arrPuts) {
 			filledArr.addArg(put.getArg(2));
-			remover.addWithoutUnbind(put);
+			if (put != lastPut) {
+				remover.addWithoutUnbind(put);
+			}
 			InsnRemover.unbindArgUsage(mth, put.getArg(0));
 		}
 		remover.addWithoutUnbind(newArrayInsn);
 
-		int replaceIndex = InsnList.getIndex(instructions, Utils.last(arrPuts));
+		int replaceIndex = InsnList.getIndex(instructions, lastPut);
 		instructions.set(replaceIndex, filledArr);
 	}
 

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

@@ -97,7 +97,7 @@ public class SimplifyVisitor extends AbstractVisitor {
 			if (arg.isInsnWrap()) {
 				InsnNode ni = simplifyInsn(mth, block, ((InsnWrapArg) arg).getWrapInsn());
 				if (ni != null) {
-					arg.wrapInstruction(ni);
+					arg.wrapInstruction(mth, ni);
 				}
 			}
 		}
@@ -332,16 +332,11 @@ public class SimplifyVisitor extends AbstractVisitor {
 				}
 				args.add(arg);
 			}
+			removeStringBuilderInsns(mth, toStrInsn, chain);
+
 			InsnNode concatInsn = new InsnNode(InsnType.STR_CONCAT, args);
 			concatInsn.setResult(toStrInsn.getResult());
 			concatInsn.copyAttributesFrom(toStrInsn);
-
-			InsnRemover insnRemover = new InsnRemover(mth);
-			for (InsnNode insnNode : chain) {
-				insnRemover.addAndUnbind(insnNode);
-			}
-			insnRemover.perform();
-
 			return concatInsn;
 		} catch (Exception e) {
 			LOG.warn("Can't convert string concatenation: {} insn: {}", mth, toStrInsn, e);
@@ -349,6 +344,23 @@ public class SimplifyVisitor extends AbstractVisitor {
 		return null;
 	}
 
+	/**
+	 * Remove and unbind all instructions with StringBuilder
+	 */
+	private static void removeStringBuilderInsns(MethodNode mth, InvokeNode toStrInsn, List<InsnNode> chain) {
+		InsnRemover.unbindAllArgs(mth, toStrInsn);
+		for (InsnNode insnNode : chain) {
+			InsnRemover.unbindAllArgs(mth, insnNode);
+		}
+		InsnRemover insnRemover = new InsnRemover(mth);
+		for (InsnNode insnNode : chain) {
+			if (insnNode != toStrInsn) {
+				insnRemover.addAndUnbind(insnNode);
+			}
+		}
+		insnRemover.perform();
+	}
+
 	private static List<InsnNode> flattenInsnChainUntil(InsnNode insn, InsnType insnType) {
 		List<InsnNode> chain = new ArrayList<>();
 		InsnArg arg = insn.getArg(0);

jadx-core/src/main/java/jadx/core/dex/visitors/regions/TernaryMod.java

@@ -104,8 +104,9 @@ public class TernaryMod implements IRegionIterativeVisitor {
 				resArg = thenResArg;
 				thenPhi.removeArg(elseResArg);
 			}
-			TernaryInsn ternInsn = new TernaryInsn(ifRegion.getCondition(),
-					resArg, InsnArg.wrapArg(thenInsn), InsnArg.wrapArg(elseInsn));
+			InsnArg thenArg = InsnArg.wrapInsnIntoArg(thenInsn);
+			InsnArg elseArg = InsnArg.wrapInsnIntoArg(elseInsn);
+			TernaryInsn ternInsn = new TernaryInsn(ifRegion.getCondition(), resArg, thenArg, elseArg);
 			ternInsn.setSourceLine(thenInsn.getSourceLine());
 
 			InsnRemover.unbindResult(mth, elseInsn);
@@ -142,7 +143,9 @@ public class TernaryMod implements IRegionIterativeVisitor {
 			TernaryInsn ternInsn = new TernaryInsn(ifRegion.getCondition(), null, thenArg, elseArg);
 			ternInsn.setSourceLine(thenInsn.getSourceLine());
 			InsnNode retInsn = new InsnNode(InsnType.RETURN, 1);
-			retInsn.addArg(InsnArg.wrapArg(ternInsn));
+			InsnArg arg = InsnArg.wrapInsnIntoArg(ternInsn);
+			arg.setType(thenArg.getType());
+			retInsn.addArg(arg);
 
 			header.getInstructions().clear();
 			header.getInstructions().add(retInsn);
@@ -234,7 +237,7 @@ public class TernaryMod implements IRegionIterativeVisitor {
 	/**
 	 * Convert one variable change with only 'then' branch:
 	 * 'if (c) {r = a;}' to 'r = c ? a : r'
-	 * Also convert only if 'r' used only once
+	 * Convert if 'r' used only once
 	 */
 	private static boolean processOneBranchTernary(MethodNode mth, IfRegion ifRegion) {
 		IContainer thenRegion = ifRegion.getThenRegion();
@@ -276,10 +279,10 @@ public class TernaryMod implements IRegionIterativeVisitor {
 		// all checks passed
 		InsnList.remove(block, insn);
 		TernaryInsn ternInsn = new TernaryInsn(ifRegion.getCondition(),
-				phiInsn.getResult(), InsnArg.wrapArg(insn), otherArg);
+				phiInsn.getResult(), InsnArg.wrapInsnIntoArg(insn), otherArg);
 		ternInsn.setSourceLine(insn.getSourceLine());
 
-		InsnRemover.unbindInsn(mth, phiInsn);
+		InsnRemover.unbindAllArgs(mth, phiInsn);
 		header.getInstructions().clear();
 		header.getInstructions().add(ternInsn);
 

jadx-core/src/main/java/jadx/core/dex/visitors/shrink/CodeShrinkVisitor.java

@@ -20,7 +20,6 @@ import jadx.core.dex.visitors.JadxVisitor;
 import jadx.core.dex.visitors.ModVisitor;
 import jadx.core.utils.BlockUtils;
 import jadx.core.utils.InsnList;
-import jadx.core.utils.InsnRemover;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
 @JadxVisitor(
@@ -117,10 +116,9 @@ public class CodeShrinkVisitor extends AbstractVisitor {
 		if (parentInsn != null && parentInsn.getType() == InsnType.RETURN) {
 			parentInsn.setSourceLine(insn.getSourceLine());
 		}
-		boolean replaced = arg.wrapInstruction(insn) != null;
+		boolean replaced = arg.wrapInstruction(mth, insn) != null;
 		if (replaced) {
 			InsnList.remove(block, insn);
-			InsnRemover.unbindResult(mth, insn);
 		}
 		return replaced;
 	}

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

@@ -3,9 +3,11 @@ package jadx.core.utils;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import org.jetbrains.annotations.Nullable;
 
+import jadx.core.Consts;
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.instructions.PhiInsn;
@@ -16,6 +18,7 @@ import jadx.core.dex.instructions.args.SSAVar;
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.InsnNode;
 import jadx.core.dex.nodes.MethodNode;
+import jadx.core.utils.exceptions.JadxRuntimeException;
 
 /**
  * Helper class for correct instructions removing,
@@ -80,6 +83,7 @@ public class InsnRemover {
 		}
 		unbindResult(mth, insn);
 		insn.add(AFlag.REMOVE);
+		insn.add(AFlag.DONT_GENERATE);
 	}
 
 	public static void unbindResult(@Nullable MethodNode mth, InsnNode insn) {
@@ -113,6 +117,19 @@ public class InsnRemover {
 				}
 			}
 			mth.removeSVar(ssaVar);
+			return;
+		}
+		if (Consts.DEBUG) { // TODO: enable this
+			throw new JadxRuntimeException("Can't remove SSA var, still in use, count: " + useCount
+					+ ", list:\n  " + ssaVar.getUseList().stream()
+							.map(arg -> arg + " from " + arg.getParentInsn())
+							.collect(Collectors.joining("\n  ")));
+		}
+	}
+
+	public static void unbindAllArgs(@Nullable MethodNode mth, InsnNode insn) {
+		for (InsnArg arg : insn.getArguments()) {
+			unbindArgUsage(mth, arg);
 		}
 	}
 
@@ -137,12 +154,18 @@ public class InsnRemover {
 		}
 		for (InsnNode rem : toRemove) {
 			int insnsCount = insns.size();
+			boolean found = false;
 			for (int i = 0; i < insnsCount; i++) {
 				if (insns.get(i) == rem) {
 					insns.remove(i);
+					found = true;
 					break;
 				}
 			}
+			if (!found && Consts.DEBUG) { // TODO: enable this
+				throw new JadxRuntimeException("Can't remove insn:\n " + rem
+						+ "\nnot found in list:\n " + Utils.listToString(insns, "\n "));
+			}
 		}
 	}