core: fix processing 'if' at loop end

jadx-core/src/main/java/jadx/core/dex/attributes/AttributeStorage.java

@@ -72,7 +72,7 @@ public class AttributeStorage {
 		if (attrList == null) {
 			return Collections.emptyList();
 		}
-		return attrList.getList();
+		return Collections.unmodifiableList(attrList.getList());
 	}
 
 	public void remove(AFlag flag) {

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

@@ -2,6 +2,7 @@ package jadx.core.dex.visitors;
 
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.instructions.InsnType;
+import jadx.core.dex.instructions.PhiInsn;
 import jadx.core.dex.instructions.args.InsnArg;
 import jadx.core.dex.instructions.args.InsnWrapArg;
 import jadx.core.dex.instructions.args.RegisterArg;
@@ -194,15 +195,12 @@ public class CodeShrinker extends AbstractVisitor {
 //					continue;
 //				}
 				SSAVar sVar = arg.getSVar();
-				if (sVar.getAssign() == null) {
-					continue;
-				}
 				// allow inline only one use arg or 'this'
 				if (sVar.getVariableUseCount() != 1 && !arg.isThis()) {
 					continue;
 				}
 				InsnNode assignInsn = sVar.getAssign().getParentInsn();
-				if (assignInsn == null) {
+				if (assignInsn == null || assignInsn instanceof PhiInsn) {
 					continue;
 				}
 				int assignPos = insnList.getIndex(assignInsn);

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

@@ -2,6 +2,7 @@ package jadx.core.dex.visitors.regions;
 
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.attributes.AType;
+import jadx.core.dex.attributes.nodes.LoopInfo;
 import jadx.core.dex.instructions.IfNode;
 import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.instructions.args.InsnArg;
@@ -53,8 +54,8 @@ public class IfMakerHelper {
 			info.setOutBlock(null);
 			return info;
 		}
-		boolean badThen = thenBlock.contains(AFlag.LOOP_START) || !allPathsFromIf(thenBlock, info);
-		boolean badElse = elseBlock.contains(AFlag.LOOP_START) || !allPathsFromIf(elseBlock, info);
+		boolean badThen = isBadBranchBlock(info, thenBlock);
+		boolean badElse = isBadBranchBlock(info, elseBlock);
 		if (badThen && badElse) {
 			LOG.debug("Stop processing blocks after 'if': {}, method: {}", info.getIfBlock(), mth);
 			return null;
@@ -92,6 +93,26 @@ public class IfMakerHelper {
 		return info;
 	}
 
+	private static boolean isBadBranchBlock(IfInfo info, BlockNode block) {
+		// check if block at end of loop edge
+		if (block.contains(AFlag.LOOP_START) && block.getPredecessors().size() == 1) {
+			BlockNode pred = block.getPredecessors().get(0);
+			if (pred.contains(AFlag.LOOP_END)) {
+				List<LoopInfo> startLoops = block.getAll(AType.LOOP);
+				List<LoopInfo> endLoops = pred.getAll(AType.LOOP);
+				// search for same loop
+				for (LoopInfo startLoop : startLoops) {
+					for (LoopInfo endLoop : endLoops) {
+						if (startLoop == endLoop) {
+							return true;
+						}
+					}
+				}
+			}
+		}
+		return !allPathsFromIf(block, info);
+	}
+
 	private static boolean allPathsFromIf(BlockNode block, IfInfo info) {
 		List<BlockNode> preds = block.getPredecessors();
 		Set<BlockNode> ifBlocks = info.getMergedBlocks();

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

@@ -84,7 +84,8 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 			return false;
 		}
 		PhiInsn phiInsn = incrArg.getSVar().getUsedInPhi();
-		if (phiInsn.getArgsCount() != 2
+		if (phiInsn == null
+				|| phiInsn.getArgsCount() != 2
 				|| !phiInsn.getArg(1).equals(incrArg)
 				|| incrArg.getSVar().getUseCount() != 1) {
 			return false;
@@ -102,6 +103,12 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 		if (!usedOnlyInLoop(mth, loopRegion, arg)) {
 			return false;
 		}
+		// can't make loop if argument from increment instruction is assign in loop
+		for (InsnArg iArg : incrInsn.getArguments()) {
+			if (iArg.isRegister() && assignOnlyInLoop(mth, loopRegion, (RegisterArg) iArg)) {
+				return false;
+			}
+		}
 
 		// all checks passed
 		initInsn.add(AFlag.SKIP);
@@ -188,7 +195,7 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 			if (wrapArg != null) {
 				wrapArg.getParentInsn().replaceArg(wrapArg, iterVar);
 			} else {
-				LOG.debug(" Wrapped insn not found: {}, mth: {}", arrGetInsn, mth);
+				LOG.debug(" checkArrayForEach: Wrapped insn not found: {}, mth: {}", arrGetInsn, mth);
 			}
 		}
 		return new ForEachLoop(iterVar, len.getArg(0));
@@ -237,7 +244,7 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 					}
 				}
 			} else {
-				LOG.warn(" Wrapped insn not found: {}, mth: {}", nextCall, mth);
+				LOG.warn(" checkIterableForEach: Wrapped insn not found: {}, mth: {}", nextCall, mth);
 				return false;
 			}
 		} else {
@@ -295,6 +302,25 @@ public class LoopRegionVisitor extends AbstractVisitor implements IRegionVisitor
 		return false;
 	}
 
+	private static boolean assignOnlyInLoop(MethodNode mth, LoopRegion loopRegion, RegisterArg arg) {
+		InsnNode assignInsn = arg.getAssignInsn();
+		if (assignInsn == null) {
+			return true;
+		}
+		if (!argInLoop(mth, loopRegion, assignInsn.getResult())) {
+			return false;
+		}
+		if (assignInsn instanceof PhiInsn) {
+			PhiInsn phiInsn = (PhiInsn) assignInsn;
+			for (InsnArg phiArg : phiInsn.getArguments()) {
+				if (!assignOnlyInLoop(mth, loopRegion, (RegisterArg) phiArg)) {
+					return false;
+				}
+			}
+		}
+		return true;
+	}
+
 	private static boolean usedOnlyInLoop(MethodNode mth, LoopRegion loopRegion, RegisterArg arg) {
 		List<RegisterArg> useList = arg.getSVar().getUseList();
 		for (RegisterArg useArg : useList) {

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

@@ -2,8 +2,10 @@ package jadx.core.utils;
 
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.attributes.AType;
+import jadx.core.dex.attributes.nodes.PhiListAttr;
 import jadx.core.dex.instructions.IfNode;
 import jadx.core.dex.instructions.InsnType;
+import jadx.core.dex.instructions.PhiInsn;
 import jadx.core.dex.instructions.args.InsnArg;
 import jadx.core.dex.instructions.args.InsnWrapArg;
 import jadx.core.dex.instructions.mods.TernaryInsn;
@@ -119,6 +121,9 @@ public class BlockUtils {
 	}
 
 	public static BlockNode getBlockByInsn(MethodNode mth, InsnNode insn) {
+		if (insn instanceof PhiInsn) {
+			return searchBlockWithPhi(mth, (PhiInsn) insn);
+		}
 		if (insn.contains(AFlag.WRAPPED)) {
 			return getBlockByWrappedInsn(mth, insn);
 		}
@@ -130,6 +135,20 @@ public class BlockUtils {
 		return null;
 	}
 
+	private static BlockNode searchBlockWithPhi(MethodNode mth, PhiInsn insn) {
+		for (BlockNode block : mth.getBasicBlocks()) {
+			PhiListAttr phiListAttr = block.get(AType.PHI_LIST);
+			if (phiListAttr != null) {
+				for (PhiInsn phiInsn : phiListAttr.getList()) {
+					if (phiInsn == insn) {
+						return block;
+					}
+				}
+			}
+		}
+		return null;
+	}
+
 	private static BlockNode getBlockByWrappedInsn(MethodNode mth, InsnNode insn) {
 		for (BlockNode bn : mth.getBasicBlocks()) {
 			for (InsnNode bi : bn.getInstructions()) {

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

+package jadx.tests.integration.loops;
+
+import jadx.core.dex.nodes.ClassNode;
+import jadx.tests.api.IntegrationTest;
+
+import org.junit.Test;
+
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.not;
+import static org.junit.Assert.assertThat;
+
+public class TestIfInLoop2 extends IntegrationTest {
+
+	public static class TestCls {
+		public static void test(String str) {
+			int len = str.length();
+			int at = 0;
+			while (at < len) {
+				char c = str.charAt(at);
+				int endAt = at + 1;
+				if (c == 'A') {
+					while (endAt < len) {
+						c = str.charAt(endAt);
+						if (c == 'B') {
+							break;
+						}
+						endAt++;
+					}
+				}
+				at = endAt;
+			}
+		}
+	}
+
+	@Test
+	public void test() {
+		ClassNode cls = getClassNode(TestCls.class);
+		String code = cls.getCode().toString();
+
+		assertThat(code, not(containsString("for (int at = 0; at < len; at = endAt) {")));
+	}
+}