fix: support multi-entry loops (simple case) (#1320)

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

@@ -25,6 +25,7 @@ import jadx.core.dex.attributes.nodes.PhiListAttr;
 import jadx.core.dex.attributes.nodes.RegDebugInfoAttr;
 import jadx.core.dex.attributes.nodes.RenameReasonAttr;
 import jadx.core.dex.attributes.nodes.SkipMethodArgsAttr;
+import jadx.core.dex.attributes.nodes.SpecialEdgeAttr;
 import jadx.core.dex.attributes.nodes.TmpEdgeAttr;
 import jadx.core.dex.nodes.IMethodDetails;
 import jadx.core.dex.trycatch.CatchAttr;
@@ -76,6 +77,7 @@ public final class AType<T extends IJadxAttribute> implements IJadxAttrType<T> {
 	public static final AType<ForceReturnAttr> FORCE_RETURN = new AType<>();
 	public static final AType<AttrList<LoopInfo>> LOOP = new AType<>();
 	public static final AType<AttrList<EdgeInsnAttr>> EDGE_INSN = new AType<>();
+	public static final AType<AttrList<SpecialEdgeAttr>> SPECIAL_EDGE = new AType<>();
 	public static final AType<TmpEdgeAttr> TMP_EDGE = new AType<>();
 	public static final AType<TryCatchBlockAttr> TRY_BLOCK = new AType<>();
 

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

 package jadx.core.dex.attributes.nodes;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -20,10 +19,10 @@ public class LoopInfo {
 	private int id;
 	private LoopInfo parentLoop;
 
-	public LoopInfo(BlockNode start, BlockNode end) {
+	public LoopInfo(BlockNode start, BlockNode end, Set<BlockNode> loopBlocks) {
 		this.start = start;
 		this.end = end;
-		this.loopBlocks = Collections.unmodifiableSet(BlockUtils.getAllPathsBlocks(start, end));
+		this.loopBlocks = loopBlocks;
 	}
 
 	public BlockNode getStart() {

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

+package jadx.core.dex.attributes.nodes;
+
+import jadx.api.plugins.input.data.attributes.IJadxAttribute;
+import jadx.core.dex.attributes.AType;
+import jadx.core.dex.attributes.AttrList;
+import jadx.core.dex.nodes.BlockNode;
+
+public class SpecialEdgeAttr implements IJadxAttribute {
+
+	public enum SpecialEdgeType {
+		BACK_EDGE,
+		CROSS_EDGE
+	}
+
+	private final SpecialEdgeType type;
+	private final BlockNode start;
+	private final BlockNode end;
+
+	public SpecialEdgeAttr(SpecialEdgeType type, BlockNode start, BlockNode end) {
+		this.type = type;
+		this.start = start;
+		this.end = end;
+	}
+
+	public SpecialEdgeType getType() {
+		return type;
+	}
+
+	public BlockNode getStart() {
+		return start;
+	}
+
+	public BlockNode getEnd() {
+		return end;
+	}
+
+	@Override
+	public AType<AttrList<SpecialEdgeAttr>> getAttrType() {
+		return AType.SPECIAL_EDGE;
+	}
+
+	@Override
+	public String toString() {
+		return type + ": " + start + " -> " + end;
+	}
+}

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

@@ -53,6 +53,10 @@ public class BlockProcessor extends AbstractVisitor {
 			clearBlocksState(mth);
 			computeDominators(mth);
 		}
+		if (FixMultiEntryLoops.process(mth)) {
+			clearBlocksState(mth);
+			computeDominators(mth);
+		}
 		updateCleanSuccessors(mth);
 
 		int i = 0;
@@ -347,7 +351,8 @@ public class BlockProcessor extends AbstractVisitor {
 					successor.add(AFlag.LOOP_START);
 					block.add(AFlag.LOOP_END);
 
-					LoopInfo loop = new LoopInfo(successor, block);
+					Set<BlockNode> loopBlocks = BlockUtils.getAllPathsBlocks(successor, block);
+					LoopInfo loop = new LoopInfo(successor, block, loopBlocks);
 					successor.addAttr(AType.LOOP, loop);
 					block.addAttr(AType.LOOP, loop);
 				}

jadx-core/src/main/java/jadx/core/dex/visitors/blocks/BlockSplitter.java

@@ -192,6 +192,14 @@ public class BlockSplitter extends AbstractVisitor {
 		return newBlock;
 	}
 
+	static void copyBlockData(BlockNode from, BlockNode to) {
+		List<InsnNode> toInsns = to.getInstructions();
+		for (InsnNode insn : from.getInstructions()) {
+			toInsns.add(insn.copyWithoutSsa());
+		}
+		to.copyAttributesFrom(from);
+	}
+
 	static void replaceTarget(BlockNode source, BlockNode oldTarget, BlockNode newTarget) {
 		InsnNode lastInsn = BlockUtils.getLastInsn(source);
 		if (lastInsn instanceof TargetInsnNode) {

jadx-core/src/main/java/jadx/core/dex/visitors/blocks/FixMultiEntryLoops.java

+package jadx.core.dex.visitors.blocks;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import jadx.core.dex.attributes.AType;
+import jadx.core.dex.attributes.nodes.SpecialEdgeAttr;
+import jadx.core.dex.attributes.nodes.SpecialEdgeAttr.SpecialEdgeType;
+import jadx.core.dex.nodes.BlockNode;
+import jadx.core.dex.nodes.MethodNode;
+import jadx.core.utils.ListUtils;
+
+public class FixMultiEntryLoops {
+
+	public static boolean process(MethodNode mth) {
+		try {
+			detectSpecialEdges(mth);
+		} catch (Exception e) {
+			mth.addWarnComment("Failed to detect multi-entry loops", e);
+			return false;
+		}
+		List<SpecialEdgeAttr> specialEdges = mth.getAll(AType.SPECIAL_EDGE);
+		List<SpecialEdgeAttr> multiEntryLoops = specialEdges.stream()
+				.filter(e -> e.getType() == SpecialEdgeType.BACK_EDGE)
+				.filter(e -> !isSingleEntryLoop(e))
+				.collect(Collectors.toList());
+		if (multiEntryLoops.isEmpty()) {
+			return false;
+		}
+		try {
+			List<SpecialEdgeAttr> crossEdges = ListUtils.filter(specialEdges, e -> e.getType() == SpecialEdgeType.CROSS_EDGE);
+			boolean changed = false;
+			for (SpecialEdgeAttr backEdge : multiEntryLoops) {
+				changed |= fixLoop(mth, backEdge, crossEdges);
+			}
+			return changed;
+		} catch (Exception e) {
+			mth.addWarnComment("Failed to fix multi-entry loops", e);
+			return false;
+		}
+	}
+
+	private static boolean fixLoop(MethodNode mth, SpecialEdgeAttr backEdge, List<SpecialEdgeAttr> crossEdges) {
+		BlockNode header = backEdge.getEnd();
+		BlockNode headerIDom = header.getIDom();
+		SpecialEdgeAttr subEntry = ListUtils.filterOnlyOne(crossEdges, e -> e.getStart() == headerIDom);
+		if (subEntry == null || !isSupportedPattern(header, subEntry)) {
+			// TODO: for now only sub entry in header successor is supported
+			mth.addWarnComment("Unsupported multi-entry loop pattern (" + backEdge + "). Please submit an issue!!!");
+			return false;
+		}
+		BlockNode loopEnd = backEdge.getStart();
+		BlockNode subEntryBlock = subEntry.getEnd();
+		BlockNode copyHeader = BlockSplitter.insertBlockBetween(mth, loopEnd, header);
+		BlockSplitter.copyBlockData(header, copyHeader);
+		BlockSplitter.replaceConnection(copyHeader, header, subEntryBlock);
+		mth.addDebugComment("Duplicate block to fix multi-entry loop: " + backEdge);
+		return true;
+	}
+
+	private static boolean isSupportedPattern(BlockNode header, SpecialEdgeAttr subEntry) {
+		return ListUtils.isSingleElement(header.getSuccessors(), subEntry.getEnd());
+	}
+
+	private static boolean isSingleEntryLoop(SpecialEdgeAttr e) {
+		BlockNode header = e.getEnd();
+		BlockNode loopEnd = e.getStart();
+		return header == loopEnd
+				|| loopEnd.getDoms().get(header.getId()); // header dominates loop end
+	}
+
+	private enum BlockColor {
+		WHITE, GRAY, BLACK
+	}
+
+	private static void detectSpecialEdges(MethodNode mth) {
+		List<BlockNode> blocks = mth.getBasicBlocks();
+		BlockColor[] colors = new BlockColor[blocks.size()];
+		Arrays.fill(colors, BlockColor.WHITE);
+		colorDFS(mth, blocks, colors, mth.getEnterBlock().getId());
+	}
+
+	// TODO: transform to non-recursive form
+	private static void colorDFS(MethodNode mth, List<BlockNode> blocks, BlockColor[] colors, int cur) {
+		colors[cur] = BlockColor.GRAY;
+		BlockNode block = blocks.get(cur);
+		for (BlockNode v : block.getSuccessors()) {
+			int vId = v.getId();
+			switch (colors[vId]) {
+				case WHITE:
+					colorDFS(mth, blocks, colors, vId);
+					break;
+				case GRAY:
+					mth.addAttr(AType.SPECIAL_EDGE, new SpecialEdgeAttr(SpecialEdgeType.BACK_EDGE, block, v));
+					break;
+				case BLACK:
+					mth.addAttr(AType.SPECIAL_EDGE, new SpecialEdgeAttr(SpecialEdgeType.CROSS_EDGE, block, v));
+					break;
+			}
+		}
+		colors[cur] = BlockColor.BLACK;
+	}
+}

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

@@ -100,6 +100,19 @@ public class ListUtils {
 		return list;
 	}
 
+	public static <T> List<T> filter(List<T> list, Predicate<T> filter) {
+		if (list == null || list.isEmpty()) {
+			return Collections.emptyList();
+		}
+		List<T> result = new ArrayList<>();
+		for (T element : list) {
+			if (filter.test(element)) {
+				result.add(element);
+			}
+		}
+		return result;
+	}
+
 	/**
 	 * Search exactly one element in list by filter
 	 *
@@ -134,4 +147,16 @@ public class ListUtils {
 		}
 		return true;
 	}
+
+	public static <T> boolean anyMatch(List<T> list, Predicate<T> test) {
+		if (list == null || list.isEmpty()) {
+			return false;
+		}
+		for (T element : list) {
+			if (test.test(element)) {
+				return true;
+			}
+		}
+		return false;
+	}
 }

jadx-core/src/test/java/jadx/tests/external/BaseExternalTest.java

@@ -10,6 +10,7 @@ import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import jadx.api.CommentsLevel;
 import jadx.api.ICodeWriter;
 import jadx.api.JadxArgs;
 import jadx.api.JadxDecompiler;
@@ -42,6 +43,7 @@ public abstract class BaseExternalTest extends IntegrationTest {
 		args.setSkipFilesSave(true);
 		args.setSkipResources(true);
 		args.setShowInconsistentCode(true);
+		args.setCommentsLevel(CommentsLevel.DEBUG);
 		return args;
 	}
 

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

+package jadx.tests.integration.loops;
+
+import org.junit.jupiter.api.Test;
+
+import jadx.tests.api.SmaliTest;
+
+import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;
+
+public class TestMultiEntryLoop extends SmaliTest {
+
+	@Test
+	public void test() {
+		assertThat(getClassNodeFromSmali())
+				.code()
+				.containsOne("while (true) {");
+	}
+}

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

+.class public Lloops/TestMultiEntryLoop;
+.super Ljava/lang/Object;
+
+.field private static arr:[B
+
+.method private static test(III)Ljava/lang/String;
+    .registers 9
+
+    mul-int/lit8 p1, p1, 0x2
+
+    rsub-int/lit8 p1, p1, 0x6f
+
+    mul-int/lit8 p0, p0, 0x2
+
+    add-int/lit8 p0, p0, 0x1c
+
+    mul-int/lit8 p2, p2, 0x2
+
+    add-int/lit8 p2, p2, 0x4
+
+    new-instance v0, Ljava/lang/String;
+
+    const/4 v5, -0x1
+
+    sget-object v4, Lloops/TestMultiEntryLoop;->arr:[B
+
+    new-array v1, p0, [B
+
+    add-int/lit8 p0, p0, -0x1
+
+    if-nez v4, :cond_1e
+
+    move v2, p1
+
+    move v3, p2
+
+    :goto_19
+    add-int/2addr v2, v3
+
+    add-int/lit8 p1, v2, -0x8
+
+    add-int/lit8 p2, p2, 0x1
+
+    :cond_1e
+    add-int/lit8 v5, v5, 0x1
+
+    int-to-byte v2, p1
+
+    aput-byte v2, v1, v5
+
+    if-ne v5, p0, :cond_2a
+
+    invoke-direct {v0, v1}, Ljava/lang/String;-><init>([B)V
+
+    return-object v0
+
+    :cond_2a
+    move v2, p1
+
+    aget-byte v3, v4, p2
+
+    goto :goto_19
+.end method