fix: preserve code semantics on array-for-each transform (#893)

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

@@ -225,6 +225,9 @@ public class InsnGen {
 	private static final Set<Flags> BODY_ONLY_NOWRAP_FLAGS = EnumSet.of(Flags.BODY_ONLY_NOWRAP);
 
 	protected void makeInsn(InsnNode insn, CodeWriter code, Flags flag) throws CodegenException {
+		if (insn.getType() == InsnType.REGION_ARG) {
+			return;
+		}
 		try {
 			if (flag == Flags.BODY_ONLY || flag == Flags.BODY_ONLY_NOWRAP) {
 				makeInsnBody(code, insn, flag == Flags.BODY_ONLY ? BODY_ONLY_FLAG : BODY_ONLY_NOWRAP_FLAGS);

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

@@ -181,18 +181,6 @@ public class RegionGen extends InsnGen {
 	}
 
 	private CodeWriter makeLoop(LoopRegion region, CodeWriter code) throws CodegenException {
-		BlockNode header = region.getHeader();
-		if (header != null) {
-			List<InsnNode> headerInsns = header.getInstructions();
-			if (headerInsns.size() > 1) {
-				mth.addWarn("Found not inlined instructions from loop header");
-				int last = headerInsns.size() - 1;
-				for (int i = 0; i < last; i++) {
-					InsnNode insn = headerInsns.get(i);
-					makeInsn(insn, code);
-				}
-			}
-		}
 		LoopLabelAttr labelAttr = region.getInfo().getStart().get(AType.LOOP_LABEL);
 		if (labelAttr != null) {
 			code.startLine(mgen.getNameGen().getLoopLabel(labelAttr)).add(':');

jadx-core/src/main/java/jadx/core/dex/attributes/nodes/LoopInfo.java

@@ -72,6 +72,10 @@ public class LoopInfo {
 		return edges;
 	}
 
+	public BlockNode getPreHeader() {
+		return BlockUtils.selectOther(end, start.getPredecessors());
+	}
+
 	public int getId() {
 		return id;
 	}

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

@@ -1,6 +1,7 @@
 package jadx.core.dex.instructions;
 
 import jadx.core.dex.nodes.InsnNode;
+import jadx.core.utils.StringUtils;
 
 public final class ConstStringNode extends InsnNode {
 
@@ -34,6 +35,6 @@ public final class ConstStringNode extends InsnNode {
 
 	@Override
 	public String toString() {
-		return super.toString() + " \"" + str + '"';
+		return super.toString() + ' ' + StringUtils.getInstance().unescapeString(str);
 	}
 }

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

@@ -66,6 +66,9 @@ public enum InsnType {
 	ONE_ARG,
 	PHI,
 
+	// fake insn to keep arguments which will be used in regions codegen
+	REGION_ARG,
+
 	// TODO: now multidimensional arrays created using Array.newInstance function
 	NEW_MULTIDIM_ARRAY
 }

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

@@ -1,6 +1,5 @@
 package jadx.core.dex.instructions.args;
 
-import jadx.api.JadxArgs;
 import jadx.core.codegen.TypeGen;
 import jadx.core.utils.StringUtils;
 import jadx.core.utils.exceptions.JadxRuntimeException;
@@ -72,12 +71,10 @@ public final class LiteralArg extends InsnArg {
 		return literal == that.literal && getType().equals(that.getType());
 	}
 
-	private static final StringUtils DEF_STRING_UTILS = new StringUtils(new JadxArgs());
-
 	@Override
 	public String toString() {
 		try {
-			String value = TypeGen.literalToString(literal, getType(), DEF_STRING_UTILS, true, false);
+			String value = TypeGen.literalToString(literal, getType(), StringUtils.getInstance(), true, false);
 			if (getType().equals(ArgType.BOOLEAN) && (value.equals("true") || value.equals("false"))) {
 				return value;
 			}

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

@@ -272,6 +272,15 @@ public class InsnNode extends LineAttrNode {
 		return true;
 	}
 
+	public boolean containsWrappedInsn() {
+		for (InsnArg arg : this.getArguments()) {
+			if (arg.isInsnWrap()) {
+				return true;
+			}
+		}
+		return false;
+	}
+
 	/**
 	 * 'Soft' equals, don't compare arguments, only instruction specific parameters.
 	 */

jadx-core/src/main/java/jadx/core/dex/regions/loops/ForEachLoop.java

@@ -1,25 +1,40 @@
 package jadx.core.dex.regions.loops;
 
+import jadx.core.dex.attributes.AFlag;
+import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.instructions.args.InsnArg;
 import jadx.core.dex.instructions.args.RegisterArg;
+import jadx.core.dex.nodes.InsnNode;
 
 public final class ForEachLoop extends LoopType {
-	private final RegisterArg varArg;
-	private final InsnArg iterableArg;
+	private final InsnNode varArgInsn;
+	private final InsnNode iterableArgInsn;
 
 	public ForEachLoop(RegisterArg varArg, InsnArg iterableArg) {
-		this.varArg = varArg;
-		this.iterableArg = iterableArg;
+		// store for-each args in fake instructions to
+		// save code semantics and allow args manipulations like args inlining
+		varArgInsn = new InsnNode(InsnType.REGION_ARG, 1);
+		varArgInsn.add(AFlag.DONT_INLINE);
+		varArgInsn.setResult(varArg.duplicate());
+
+		iterableArgInsn = new InsnNode(InsnType.REGION_ARG, 1);
+		iterableArgInsn.add(AFlag.DONT_INLINE);
+		iterableArgInsn.addArg(iterableArg.duplicate());
 
 		// will be declared at codegen
-		varArg.getSVar().getCodeVar().setDeclared(true);
+		getVarArg().getSVar().getCodeVar().setDeclared(true);
+	}
+
+	public void injectFakeInsns(LoopRegion loopRegion) {
+		loopRegion.getInfo().getPreHeader().getInstructions().add(iterableArgInsn);
+		loopRegion.getHeader().getInstructions().add(0, varArgInsn);
 	}
 
 	public RegisterArg getVarArg() {
-		return varArg;
+		return varArgInsn.getResult();
 	}
 
 	public InsnArg getIterableArg() {
-		return iterableArg;
+		return iterableArgInsn.getArg(0);
 	}
 }

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

@@ -511,11 +511,44 @@ public class BlockProcessor extends AbstractVisitor {
 			LoopInfo loop = loops.get(0);
 			return insertBlocksForBreak(mth, loop)
 					|| insertBlocksForContinue(mth, loop)
-					|| insertBlockForProdecessors(mth, loop);
+					|| insertBlockForProdecessors(mth, loop)
+					|| insertPreHeader(mth, loop);
 		}
 		return false;
 	}
 
+	/**
+	 * Insert simple path block before loop header
+	 */
+	private static boolean insertPreHeader(MethodNode mth, LoopInfo loop) {
+		BlockNode start = loop.getStart();
+		List<BlockNode> preds = start.getPredecessors();
+		int predsCount = preds.size() - 1; // don't count back edge
+		if (predsCount == 1) {
+			return false;
+		}
+		if (predsCount == 0) {
+			if (!start.contains(AFlag.MTH_ENTER_BLOCK)) {
+				mth.addWarnComment("Unexpected block without predecessors: " + start);
+			}
+			BlockNode newEnterBlock = BlockSplitter.startNewBlock(mth, -1);
+			newEnterBlock.add(AFlag.SYNTHETIC);
+			newEnterBlock.add(AFlag.MTH_ENTER_BLOCK);
+			mth.setEnterBlock(newEnterBlock);
+			start.remove(AFlag.MTH_ENTER_BLOCK);
+			BlockSplitter.connect(newEnterBlock, start);
+			return true;
+		}
+		// multiple predecessors
+		BlockNode preHeader = BlockSplitter.startNewBlock(mth, -1);
+		preHeader.add(AFlag.SYNTHETIC);
+		for (BlockNode pred : new ArrayList<>(preds)) {
+			BlockSplitter.replaceConnection(pred, start, preHeader);
+		}
+		BlockSplitter.connect(preHeader, start);
+		return true;
+	}
+
 	/**
 	 * Insert additional blocks for possible 'break' insertion
 	 */

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

@@ -37,6 +37,8 @@ import jadx.core.dex.visitors.JadxVisitor;
 import jadx.core.dex.visitors.regions.variables.ProcessVariables;
 import jadx.core.dex.visitors.shrink.CodeShrinkVisitor;
 import jadx.core.utils.BlockUtils;
+import jadx.core.utils.InsnRemover;
+import jadx.core.utils.InsnUtils;
 import jadx.core.utils.RegionUtils;
 import jadx.core.utils.exceptions.JadxOverflowException;
 
@@ -50,7 +52,6 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 
 	@Override
 	public void visit(MethodNode mth) {
-		// DebugUtils.checkMethod(mth);
 		DepthRegionTraversal.traverse(mth, this);
 	}
 
@@ -133,11 +134,7 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 		incrInsn.add(AFlag.DONT_GENERATE);
 
 		LoopType arrForEach = checkArrayForEach(mth, loopRegion, initInsn, incrInsn, condition);
-		if (arrForEach != null) {
-			loopRegion.setType(arrForEach);
-		} else {
-			loopRegion.setType(new ForLoop(initInsn, incrInsn));
-		}
+		loopRegion.setType(arrForEach != null ? arrForEach : new ForLoop(initInsn, incrInsn));
 		return true;
 	}
 
@@ -172,7 +169,7 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 		condArg = args.get(0);
 		RegisterArg arrIndex = args.get(1);
 		InsnNode arrGetInsn = arrIndex.getParentInsn();
-		if (arrGetInsn == null || arrGetInsn.getType() != InsnType.AGET) {
+		if (arrGetInsn == null || arrGetInsn.getType() != InsnType.AGET || arrGetInsn.containsWrappedInsn()) {
 			return null;
 		}
 		if (!condition.isCompare()) {
@@ -211,23 +208,25 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 		condArg.add(AFlag.DONT_GENERATE);
 		bCondArg.add(AFlag.DONT_GENERATE);
 		arrGetInsn.add(AFlag.DONT_GENERATE);
-
-		// inline array variable
-		if (arrayArg.isRegister()) {
-			((RegisterArg) arrayArg).getSVar().removeUse((RegisterArg) arrGetInsn.getArg(0));
-		}
-		CodeShrinkVisitor.shrinkMethod(mth);
-		len.add(AFlag.DONT_GENERATE);
+		compare.getInsn().add(AFlag.DONT_GENERATE);
 
 		if (arrGetInsn.contains(AFlag.WRAPPED)) {
 			InsnArg wrapArg = BlockUtils.searchWrappedInsnParent(mth, arrGetInsn);
 			if (wrapArg != null && wrapArg.getParentInsn() != null) {
-				wrapArg.getParentInsn().replaceArg(wrapArg, iterVar);
+				InsnNode parentInsn = wrapArg.getParentInsn();
+				parentInsn.replaceArg(wrapArg, iterVar.duplicate());
+				parentInsn.rebindArgs();
 			} else {
 				LOG.debug(" checkArrayForEach: Wrapped insn not found: {}, mth: {}", arrGetInsn, mth);
 			}
 		}
-		return new ForEachLoop(iterVar, len.getArg(0));
+		ForEachLoop forEachLoop = new ForEachLoop(iterVar, len.getArg(0));
+		forEachLoop.injectFakeInsns(loopRegion);
+		if (InsnUtils.dontGenerateIfNotUsed(len)) {
+			InsnRemover.remove(mth, len);
+		}
+		CodeShrinkVisitor.shrinkMethod(mth);
+		return forEachLoop;
 	}
 
 	private static boolean checkIterableForEach(MethodNode mth, LoopRegion loopRegion, IfCondition condition) {
@@ -306,7 +305,9 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 		for (RegisterArg itArg : itUseList) {
 			itArg.add(AFlag.DONT_GENERATE);
 		}
-		loopRegion.setType(new ForEachLoop(iterVar, iterableArg));
+		ForEachLoop forEachLoop = new ForEachLoop(iterVar, iterableArg);
+		forEachLoop.injectFakeInsns(loopRegion);
+		loopRegion.setType(forEachLoop);
 		return true;
 	}
 

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

@@ -182,6 +182,8 @@ public class InsnRemover {
 		BlockNode block = BlockUtils.getBlockByInsn(mth, insn);
 		if (block != null) {
 			remove(mth, block, insn);
+		} else {
+			mth.addWarnComment("Not found block with instruction: " + insn);
 		}
 	}
 

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

@@ -8,6 +8,7 @@ import org.slf4j.LoggerFactory;
 
 import com.android.dx.io.instructions.DecodedInstruction;
 
+import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.attributes.AType;
 import jadx.core.dex.info.FieldInfo;
 import jadx.core.dex.instructions.ConstClassNode;
@@ -17,6 +18,7 @@ import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.instructions.args.InsnArg;
 import jadx.core.dex.instructions.args.InsnWrapArg;
 import jadx.core.dex.instructions.args.RegisterArg;
+import jadx.core.dex.instructions.args.SSAVar;
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.DexNode;
 import jadx.core.dex.nodes.FieldNode;
@@ -193,4 +195,20 @@ public class InsnUtils {
 		}
 		return null;
 	}
+
+	public static boolean dontGenerateIfNotUsed(InsnNode insn) {
+		RegisterArg resArg = insn.getResult();
+		if (resArg != null) {
+			SSAVar ssaVar = resArg.getSVar();
+			for (RegisterArg arg : ssaVar.getUseList()) {
+				InsnNode parentInsn = arg.getParentInsn();
+				if (parentInsn != null
+						&& !parentInsn.contains(AFlag.DONT_GENERATE)) {
+					return false;
+				}
+			}
+		}
+		insn.add(AFlag.DONT_GENERATE);
+		return true;
+	}
 }

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

@@ -3,6 +3,11 @@ package jadx.core.utils;
 import jadx.api.JadxArgs;
 
 public class StringUtils {
+	private static final StringUtils DEFAULT_INSTANCE = new StringUtils(new JadxArgs());
+
+	public static StringUtils getInstance() {
+		return DEFAULT_INSTANCE;
+	}
 
 	private final boolean escapeUnicode;
 

jadx-core/src/test/java/jadx/tests/integration/loops/TestBreakWithLabel.java

@@ -36,7 +36,7 @@ public class TestBreakWithLabel extends IntegrationTest {
 	}
 
 	@Test
-	public void test() throws Exception {
+	public void test() {
 		ClassNode cls = getClassNode(TestCls.class);
 		String code = cls.getCode().toString();
 

jadx-core/src/test/java/jadx/tests/integration/loops/TestLoopRestore.java

@@ -0,0 +1,19 @@
+package jadx.tests.integration.loops;
+
+import org.junit.jupiter.api.Test;
+
+import jadx.core.dex.nodes.ClassNode;
+import jadx.tests.api.SmaliTest;
+
+import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;
+
+public class TestLoopRestore extends SmaliTest {
+
+	@Test
+	public void test() {
+		ClassNode cls = getClassNodeFromSmali();
+		assertThat(cls).code()
+				.containsOne("try {")
+				.containsOne("for (byte b : digest) {");
+	}
+}

jadx-core/src/test/java/jadx/tests/integration/types/TestGenerics2.java

@@ -15,11 +15,9 @@ public class TestGenerics2 extends SmaliTest {
 		public void test() {
 			Map<Integer, String> map = this.field;
 			useInt(map.size());
-			Iterator<Map.Entry<Integer, String>> it = map.entrySet().iterator();
-			while (it.hasNext()) {
-				Map.Entry<Integer, String> next = it.next();
-				useInt(next.getKey().intValue());
-				next.getValue().trim();
+			for (Map.Entry<Integer, String> entry : map.entrySet()) {
+				useInt(entry.getKey().intValue());
+				entry.getValue().trim();
 			}
 		}
 	*/
@@ -30,8 +28,8 @@ public class TestGenerics2 extends SmaliTest {
 		ClassNode cls = getClassNodeFromSmali();
 		String code = cls.getCode().toString();
 
-		assertThat(code, containsOne("Entry<Integer, String> next"));
-		assertThat(code, containsOne("useInt(next.getKey().intValue());")); // no Integer cast
-		assertThat(code, containsOne("next.getValue().trim();")); // no String cast
+		assertThat(code, containsOne("for (Map.Entry<Integer, String> entry : map.entrySet()) {"));
+		assertThat(code, containsOne("useInt(entry.getKey().intValue());")); // no Integer cast
+		assertThat(code, containsOne("entry.getValue().trim();")); // no String cast
 	}
 }

jadx-core/src/test/smali/loops/TestLoopRestore.smali

@@ -0,0 +1,74 @@
+.class public Lloops/TestLoopRestore;
+.super Ljava/lang/Object;
+.source "SourceFile.java"
+
+.method private test([B)Ljava/lang/String;
+    .registers 10
+
+    const/16 v0, 0x10
+    new-array v0, v0, [C
+    fill-array-data v0, :array_3c
+
+    :try_start_7
+    const-string v1, "MD5"
+    invoke-static {v1}, Ljava/security/MessageDigest;->getInstance(Ljava/lang/String;)Ljava/security/MessageDigest;
+    move-result-object v1
+
+    invoke-virtual {v1, p1}, Ljava/security/MessageDigest;->update([B)V
+    invoke-virtual {v1}, Ljava/security/MessageDigest;->digest()[B
+    move-result-object p1
+
+    array-length v1, p1
+    mul-int/lit8 v2, v1, 0x2
+    new-array v2, v2, [C
+    :try_end_19
+    .catch Ljava/lang/Exception; {:try_start_7 .. :try_end_19} :catch_3a
+
+    const/4 v3, 0x0
+    const/4 v4, 0x0
+
+    :goto_1b
+    if-ge v3, v1, :cond_34
+
+    aget-byte v5, p1, v3
+    add-int/lit8 v6, v4, 0x1
+    ushr-int/lit8 v7, v5, 0x4
+    and-int/lit8 v7, v7, 0xf
+    aget-char v7, v0, v7
+    aput-char v7, v2, v4
+    add-int/lit8 v4, v6, 0x1
+    and-int/lit8 v5, v5, 0xf
+    aget-char v5, v0, v5
+    aput-char v5, v2, v6
+    add-int/lit8 v3, v3, 0x1
+    goto :goto_1b
+
+    :cond_34
+    new-instance p1, Ljava/lang/String;
+    invoke-direct {p1, v2}, Ljava/lang/String;-><init>([C)V
+    return-object p1
+
+    :catch_3a
+    const/4 p1, 0x0
+    return-object p1
+
+    :array_3c
+    .array-data 2
+        0x30s
+        0x31s
+        0x32s
+        0x33s
+        0x34s
+        0x35s
+        0x36s
+        0x37s
+        0x38s
+        0x39s
+        0x61s
+        0x62s
+        0x63s
+        0x64s
+        0x65s
+        0x66s
+    .end array-data
+.end method