core: fix 'break' in complex 'if' in loop (fix #67)

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

@@ -3,6 +3,7 @@ package jadx.core.dex.attributes;
 import jadx.core.dex.attributes.annotations.AnnotationsList;
 import jadx.core.dex.attributes.annotations.MethodParameters;
 import jadx.core.dex.attributes.nodes.DeclareVariablesAttr;
+import jadx.core.dex.attributes.nodes.EdgeInsnAttr;
 import jadx.core.dex.attributes.nodes.EnumClassAttr;
 import jadx.core.dex.attributes.nodes.EnumMapAttr;
 import jadx.core.dex.attributes.nodes.FieldReplaceAttr;
@@ -30,6 +31,7 @@ public class AType<T extends IAttribute> {
 
 	public static final AType<AttrList<JumpInfo>> JUMP = new AType<AttrList<JumpInfo>>();
 	public static final AType<AttrList<LoopInfo>> LOOP = new AType<AttrList<LoopInfo>>();
+	public static final AType<AttrList<EdgeInsnAttr>> EDGE_INSN = new AType<AttrList<EdgeInsnAttr>>();
 
 	public static final AType<ExcHandlerAttr> EXC_HANDLER = new AType<ExcHandlerAttr>();
 	public static final AType<CatchAttr> CATCH_BLOCK = new AType<CatchAttr>();

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

+package jadx.core.dex.attributes.nodes;
+
+import jadx.core.dex.attributes.AType;
+import jadx.core.dex.attributes.AttrList;
+import jadx.core.dex.attributes.IAttribute;
+import jadx.core.dex.nodes.BlockNode;
+import jadx.core.dex.nodes.InsnNode;
+
+public class EdgeInsnAttr implements IAttribute {
+
+	private final BlockNode start;
+	private final BlockNode end;
+	private final InsnNode insn;
+
+	public static void addEdgeInsn(BlockNode start, BlockNode end, InsnNode insn) {
+		EdgeInsnAttr edgeInsnAttr = new EdgeInsnAttr(start, end, insn);
+		start.addAttr(AType.EDGE_INSN, edgeInsnAttr);
+		end.addAttr(AType.EDGE_INSN, edgeInsnAttr);
+	}
+
+	public EdgeInsnAttr(BlockNode start, BlockNode end, InsnNode insn) {
+		this.start = start;
+		this.end = end;
+		this.insn = insn;
+	}
+
+	@Override
+	public AType<AttrList<EdgeInsnAttr>> getType() {
+		return AType.EDGE_INSN;
+	}
+
+	public BlockNode getStart() {
+		return start;
+	}
+
+	public BlockNode getEnd() {
+		return end;
+	}
+
+	public InsnNode getInsn() {
+		return insn;
+	}
+
+	@Override
+	public String toString() {
+		return "EDGE_INSN: " + start + "->" + end + " " + insn;
+	}
+}

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

@@ -3,6 +3,7 @@ package jadx.core.dex.visitors.regions;
 import jadx.core.Consts;
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.attributes.AType;
+import jadx.core.dex.attributes.nodes.EdgeInsnAttr;
 import jadx.core.dex.attributes.nodes.LoopInfo;
 import jadx.core.dex.attributes.nodes.LoopLabelAttr;
 import jadx.core.dex.instructions.IfNode;
@@ -409,7 +410,7 @@ public class RegionMaker {
 			return false;
 		}
 		InsnNode breakInsn = new InsnNode(InsnType.BREAK, 0);
-		insertBlock.getInstructions().add(breakInsn);
+		EdgeInsnAttr.addEdgeInsn(insertBlock, insertBlock.getSuccessors().get(0), breakInsn);
 		stack.addExit(exit);
 		// add label to 'break' if needed
 		addBreakLabel(exitEdge, exit, breakInsn);
@@ -619,8 +620,9 @@ public class RegionMaker {
 		ifRegion.setCondition(currentIf.getCondition());
 		currentRegion.getSubBlocks().add(ifRegion);
 
+		BlockNode outBlock = currentIf.getOutBlock();
 		stack.push(ifRegion);
-		stack.addExit(currentIf.getOutBlock());
+		stack.addExit(outBlock);
 
 		ifRegion.setThenRegion(makeRegion(currentIf.getThenBlock(), stack));
 		BlockNode elseBlock = currentIf.getElseBlock();
@@ -630,8 +632,41 @@ public class RegionMaker {
 			ifRegion.setElseRegion(makeRegion(elseBlock, stack));
 		}
 
+		// insert edge insns in new 'else' branch
+		// TODO: make more common algorithm
+		if (ifRegion.getElseRegion() == null && outBlock != null) {
+			List<EdgeInsnAttr> edgeInsnAttrs = outBlock.getAll(AType.EDGE_INSN);
+			if (!edgeInsnAttrs.isEmpty()) {
+				Region elseRegion = new Region(ifRegion);
+				for (EdgeInsnAttr edgeInsnAttr : edgeInsnAttrs) {
+					if (edgeInsnAttr.getEnd().equals(outBlock)) {
+						addEdgeInsn(currentIf, elseRegion, edgeInsnAttr);
+					}
+				}
+				ifRegion.setElseRegion(elseRegion);
+			}
+		}
+
 		stack.pop();
-		return currentIf.getOutBlock();
+		return outBlock;
+	}
+
+	private void addEdgeInsn(IfInfo ifInfo, Region region, EdgeInsnAttr edgeInsnAttr) {
+		BlockNode start = edgeInsnAttr.getStart();
+		if (start.contains(AFlag.SKIP)) {
+			return;
+		}
+		boolean fromThisIf = false;
+		for (BlockNode ifBlock : ifInfo.getMergedBlocks()) {
+			if (ifBlock.getSuccessors().contains(start)) {
+				fromThisIf = true;
+				break;
+			}
+		}
+		if (!fromThisIf) {
+			return;
+		}
+		region.add(start);
 	}
 
 	private BlockNode processSwitch(IRegion currentRegion, BlockNode block, SwitchNode insn, RegionStack stack) {

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

 package jadx.core.dex.visitors.regions;
 
 import jadx.core.dex.attributes.AFlag;
+import jadx.core.dex.attributes.AType;
+import jadx.core.dex.attributes.nodes.EdgeInsnAttr;
 import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.IBlock;
@@ -19,6 +21,7 @@ import jadx.core.utils.RegionUtils;
 import jadx.core.utils.exceptions.JadxException;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -75,10 +78,33 @@ public class RegionMakerVisitor extends AbstractVisitor {
 			} else if (region instanceof SwitchRegion) {
 				// insert 'break' in switch cases (run after try/catch insertion)
 				processSwitch(mth, (SwitchRegion) region);
+			} else if (region instanceof Region) {
+				insertEdgeInsn((Region) region);
 			}
 		}
 	}
 
+	/**
+	 * Insert insn block from edge insn attribute.
+	 */
+	private static void insertEdgeInsn(Region region) {
+		List<IContainer> subBlocks = region.getSubBlocks();
+		if (subBlocks.isEmpty()) {
+			return;
+		}
+		IContainer last = subBlocks.get(subBlocks.size() - 1);
+		List<EdgeInsnAttr> edgeInsnAttrs = last.getAll(AType.EDGE_INSN);
+		if (edgeInsnAttrs.isEmpty()) {
+			return;
+		}
+		EdgeInsnAttr insnAttr = edgeInsnAttrs.get(0);
+		if (!insnAttr.getStart().equals(last)) {
+			return;
+		}
+		List<InsnNode> insns = Collections.singletonList(insnAttr.getInsn());
+		region.add(new InsnContainer(insns));
+	}
+
 	private static void processSwitch(MethodNode mth, SwitchRegion sw) {
 		for (IContainer c : sw.getBranches()) {
 			if (!(c instanceof Region)) {

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

@@ -64,7 +64,7 @@ public class DebugUtils {
 	}
 
 	public static void printRegions(MethodNode mth, boolean printInsns) {
-		LOG.debug("|{}", mth.toString());
+		LOG.debug("|{}", mth);
 		printRegion(mth, mth.getRegion(), "|  ", printInsns);
 	}
 
@@ -75,7 +75,7 @@ public class DebugUtils {
 			if (container instanceof IRegion) {
 				printRegion(mth, (IRegion) container, indent, printInsns);
 			} else {
-				LOG.debug("{}{}", indent, container);
+				LOG.debug("{}{} {}", indent, container, container.getAttributesString());
 				if (printInsns && container instanceof IBlock) {
 					IBlock block = (IBlock) container;
 					printInsns(mth, indent, block);

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

+package jadx.tests.integration.loops;
+
+import jadx.core.dex.nodes.ClassNode;
+import jadx.tests.api.IntegrationTest;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.junit.Test;
+
+import static jadx.tests.api.utils.JadxMatchers.containsOne;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+
+public class TestBreakInComplexIf extends IntegrationTest {
+
+	public static class TestCls {
+
+		private int test(Map<String, Point> map, int mapX) {
+			int length = 1;
+			for (int x = mapX + 1; x < 100; x++) {
+				Point tile = map.get(x + "");
+				if (tile == null || tile.y != 100) {
+					break;
+				}
+				length++;
+			}
+			return length;
+		}
+
+		class Point {
+			public final int x;
+			public final int y;
+
+			Point(int x, int y) {
+				this.x = x;
+				this.y = y;
+			}
+		}
+
+		public void check() {
+			Map<String, Point> map = new HashMap<String, Point>();
+			map.put("3", new Point(100, 100));
+			map.put("4", new Point(60, 100));
+			assertThat(test(map, 2), is(3));
+		}
+	}
+
+	@Test
+	public void test() {
+		ClassNode cls = getClassNode(TestCls.class);
+		String code = cls.getCode().toString();
+
+		assertThat(code, containsOne("if (tile == null || tile.y != 100) {"));
+		assertThat(code, containsOne("break;"));
+	}
+
+	@Test
+	public void testNoDebug() {
+		noDebugInfo();
+		ClassNode cls = getClassNode(TestCls.class);
+		String code = cls.getCode().toString();
+		assertThat(code, containsOne("break;"));
+	}
+}

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

+package jadx.tests.integration.loops;
+
+import jadx.core.dex.nodes.ClassNode;
+import jadx.tests.api.IntegrationTest;
+
+import java.util.Arrays;
+import java.util.List;
+
+import org.junit.Test;
+
+import static jadx.tests.api.utils.JadxMatchers.countString;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+
+public class TestBreakInComplexIf2 extends IntegrationTest {
+
+	public static class TestCls {
+
+		private int test(List<String> list) {
+			int length = 0;
+			for (String str : list) {
+				if (str.isEmpty() || str.length() > 4) {
+					break;
+				}
+				if (str.equals("skip")) {
+					continue;
+				}
+				if (str.equals("a")) {
+					break;
+				}
+				length++;
+			}
+			return length;
+		}
+
+		public void check() {
+			assertThat(test(Arrays.asList("x", "y", "skip", "z", "a")), is(3));
+			assertThat(test(Arrays.asList("x", "skip", "")), is(1));
+			assertThat(test(Arrays.asList("skip", "y", "12345")), is(1));
+		}
+	}
+
+	@Test
+	public void test() {
+		ClassNode cls = getClassNode(TestCls.class);
+		String code = cls.getCode().toString();
+
+		assertThat(code, countString(2, "break;"));
+	}
+
+	@Test
+	public void testNoDebug() {
+		noDebugInfo();
+		ClassNode cls = getClassNode(TestCls.class);
+		String code = cls.getCode().toString();
+
+		assertThat(code, countString(2, "break;"));
+	}
+}