fix: improve top block search for try/catch (#1633)

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

@@ -19,29 +19,59 @@
 
 public final class BlockNode extends AttrNode implements IBlock, Comparable<BlockNode> {
 
+	/**
+	 * Const ID
+	 */
+	private final int cid;
+
+	/**
+	 * ID linked to position in blocks list (easier to use BitSet)
+	 * TODO: rename to avoid confusion
+	 */
 	private int id;
+
+	/**
+	 * Offset in methods bytecode
+	 */
 	private final int startOffset;
+
 	private final List<InsnNode> instructions = new ArrayList<>(2);
 
 	private List<BlockNode> predecessors = new ArrayList<>(1);
 	private List<BlockNode> successors = new ArrayList<>(1);
 	private List<BlockNode> cleanSuccessors;
 
-	// all dominators
+	/**
+	 * All dominators, excluding self
+	 */
 	private BitSet doms = EmptyBitSet.EMPTY;
-	// dominance frontier
+
+	/**
+	 * Dominance frontier
+	 */
 	private BitSet domFrontier;
-	// immediate dominator
+
+	/**
+	 * Immediate dominator
+	 */
 	private BlockNode idom;
-	// blocks on which dominates this block
+
+	/**
+	 * Blocks on which dominates this block
+	 */
 	private List<BlockNode> dominatesOn = new ArrayList<>(3);
 
-	public BlockNode(int id, int offset) {
+	public BlockNode(int cid, int id, int offset) {
+		this.cid = cid;
 		this.id = id;
 		this.startOffset = offset;
 	}
 
-	public void setId(int id) {
+	public int getCId() {
+		return cid;
+	}
+
+	void setId(int id) {
 		this.id = id;
 	}
 
@@ -188,12 +218,12 @@ public boolean equals(Object obj) {
 			return false;
 		}
 		BlockNode other = (BlockNode) obj;
-		return id == other.id && startOffset == other.startOffset;
+		return cid == other.cid && startOffset == other.startOffset;
 	}
 
 	@Override
 	public int compareTo(@NotNull BlockNode o) {
-		return Integer.compare(id, o.id);
+		return Integer.compare(cid, o.cid);
 	}
 
 	@Override
@@ -203,6 +233,6 @@ public String baseString() {
 
 	@Override
 	public String toString() {
-		return "B:" + id + ':' + InsnUtils.formatOffset(startOffset);
+		return "B:" + cid + ':' + InsnUtils.formatOffset(startOffset);
 	}
 }

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

@@ -61,6 +61,7 @@ public class MethodNode extends NotificationAttrNode implements IMethodDetails,
 	private List<RegisterArg> argsList;
 	private InsnNode[] instructions;
 	private List<BlockNode> blocks;
+	private int blocksMaxCId;
 	private BlockNode enterBlock;
 	private BlockNode exitBlock;
 	private List<SSAVar> sVars;
@@ -318,6 +319,15 @@ public List<BlockNode> getBasicBlocks() {
 
 	public void setBasicBlocks(List<BlockNode> blocks) {
 		this.blocks = blocks;
+		int i = 0;
+		for (BlockNode block : blocks) {
+			block.setId(i);
+			i++;
+		}
+	}
+
+	public int getNextBlockCId() {
+		return blocksMaxCId++;
 	}
 
 	public BlockNode getEnterBlock() {

jadx-core/src/main/java/jadx/core/dex/trycatch/TryCatchBlockAttr.java

@@ -60,6 +60,10 @@ public boolean isThrowOnly() {
 		return throwFound;
 	}
 
+	public int getId() {
+		return id;
+	}
+
 	public List<ExceptionHandler> getHandlers() {
 		return handlers;
 	}

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

@@ -100,7 +100,7 @@ public void process(MethodNode mth) {
 					if (insnArr == null) {
 						return;
 					}
-					BlockNode block = new BlockNode(0, 0);
+					BlockNode block = new BlockNode(0, 0, 0);
 					List<InsnNode> insnList = block.getInstructions();
 					for (InsnNode insn : insnArr) {
 						if (insn != null) {
@@ -199,7 +199,7 @@ private void processBlock(MethodNode mth, BlockNode block, boolean error) {
 				dot.add("color=red,");
 			}
 			dot.add("label=\"{");
-			dot.add(String.valueOf(block.getId())).add("\\:\\ ");
+			dot.add(String.valueOf(block.getCId())).add("\\:\\ ");
 			dot.add(InsnUtils.formatOffset(block.getStartOffset()));
 			if (!attrs.isEmpty()) {
 				dot.add('|').add(attrs);
@@ -230,10 +230,10 @@ private void processBlock(MethodNode mth, BlockNode block, boolean error) {
 
 			if (PRINT_DOMINATORS) {
 				for (BlockNode c : block.getDominatesOn()) {
-					conn.startLine(block.getId() + " -> " + c.getId() + "[color=green];");
+					conn.startLine(block.getCId() + " -> " + c.getCId() + "[color=green];");
 				}
 				for (BlockNode dom : BlockUtils.bitSetToBlocks(mth, block.getDomFrontier())) {
-					conn.startLine("f_" + block.getId() + " -> f_" + dom.getId() + "[color=blue];");
+					conn.startLine("f_" + block.getCId() + " -> f_" + dom.getCId() + "[color=blue];");
 				}
 			}
 		}
@@ -273,7 +273,7 @@ private String attributesString(IAttributeNode block) {
 		private String makeName(IContainer c) {
 			String name;
 			if (c instanceof BlockNode) {
-				name = "Node_" + ((BlockNode) c).getId();
+				name = "Node_" + ((BlockNode) c).getCId();
 			} else if (c instanceof IBlock) {
 				name = "Node_" + c.getClass().getSimpleName() + '_' + c.hashCode();
 			} else {

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

@@ -171,10 +171,6 @@ private static void initExcHandlers(MethodNode mth) {
 		}
 	}
 
-	protected static void removeTmpConnections(MethodNode mth) {
-		mth.getBasicBlocks().forEach(BlockExceptionHandler::removeTmpConnection);
-	}
-
 	private static void removeTmpConnection(BlockNode block) {
 		TmpEdgeAttr tmpEdgeAttr = block.get(AType.TMP_EDGE);
 		if (tmpEdgeAttr != null) {
@@ -402,6 +398,13 @@ private static BlockNode searchTopBlock(MethodNode mth, List<BlockNode> blocks)
 		}
 		BlockNode topDom = BlockUtils.getCommonDominator(mth, blocks);
 		if (topDom != null) {
+			// dominator always return one up block if blocks already contains dominator, use successor instead
+			if (topDom.getSuccessors().size() == 1) {
+				BlockNode upBlock = topDom.getSuccessors().get(0);
+				if (blocks.contains(upBlock)) {
+					return upBlock;
+				}
+			}
 			return adjustTopBlock(topDom);
 		}
 		throw new JadxRuntimeException("Failed to find top block for try-catch from: " + blocks);

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

@@ -137,8 +137,9 @@ static BlockNode connectNewBlock(MethodNode mth, BlockNode block, int offset) {
 	}
 
 	static BlockNode startNewBlock(MethodNode mth, int offset) {
-		BlockNode block = new BlockNode(mth.getBasicBlocks().size(), offset);
-		mth.getBasicBlocks().add(block);
+		List<BlockNode> blocks = mth.getBasicBlocks();
+		BlockNode block = new BlockNode(mth.getNextBlockCId(), blocks.size(), offset);
+		blocks.add(block);
 		return block;
 	}
 

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

@@ -25,23 +25,16 @@ public static void compute(MethodNode mth) {
 		List<BlockNode> sorted = sortBlocks(mth);
 		BlockNode[] doms = build(sorted);
 		apply(sorted, doms);
-		mth.setBasicBlocks(sorted);
 	}
 
 	private static List<BlockNode> sortBlocks(MethodNode mth) {
 		int blocksCount = mth.getBasicBlocks().size();
-		BitSet reachSet = new BitSet(blocksCount);
 		List<BlockNode> sorted = new ArrayList<>(blocksCount);
-		BlockUtils.dfsVisit(mth, b -> {
-			sorted.add(b);
-			reachSet.set(b.getId());
-		});
-		if (reachSet.cardinality() != blocksCount) {
+		BlockUtils.dfsVisit(mth, sorted::add);
+		if (sorted.size() != blocksCount) {
 			throw new JadxRuntimeException("Found unreachable blocks");
 		}
-		for (int i = 0; i < blocksCount; i++) {
-			sorted.get(i).setId(i);
-		}
+		mth.setBasicBlocks(sorted);
 		return sorted;
 	}
 

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

@@ -697,7 +697,7 @@ public static BlockNode traverseWhileDominates(BlockNode dom, BlockNode start) {
 	}
 
 	/**
-	 * Search lowest common ancestor in dominator tree for input set.
+	 * Search the lowest common ancestor in dominator tree for input set.
 	 */
 	@Nullable
 	public static BlockNode getCommonDominator(MethodNode mth, List<BlockNode> blocks) {

jadx-core/src/test/java/jadx/tests/integration/trycatch/TestNestedTryCatch4.java

@@ -0,0 +1,18 @@
+package jadx.tests.integration.trycatch;
+
+import org.junit.jupiter.api.Test;
+
+import jadx.tests.api.SmaliTest;
+
+import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;
+
+public class TestNestedTryCatch4 extends SmaliTest {
+
+	@Test
+	public void test() {
+		disableCompilation();
+		assertThat(getClassNodeFromSmali())
+				.code()
+				.doesNotContain("?? ");
+	}
+}