fix: improve restructure of nested try/catch/finally blocks (#2033)

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

@@ -33,6 +33,7 @@ public enum AFlag {
 	EXC_TOP_SPLITTER,
 	EXC_BOTTOM_SPLITTER,
 	FINALLY_INSNS,
+	IGNORE_THROW_SPLIT,
 
 	SKIP_FIRST_ARG,
 	SKIP_ARG, // skip argument in invoke call

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

@@ -6,6 +6,7 @@
 import java.util.Comparator;
 import java.util.Deque;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -283,14 +284,7 @@ private static boolean checkTryCatchRelation(List<TryCatchBlockAttr> tryBlocks,
 		boolean catchInTry = innerTryBlock.getBlocks().stream().anyMatch(isHandlersIntersects(outerTryBlock));
 		boolean blocksOutsideHandler = outerTryBlock.getBlocks().stream().anyMatch(b -> !handlerBlocks.contains(b));
 
-		boolean makeInner = catchInHandler && (catchInTry || blocksOutsideHandler);
-		if (makeInner && innerTryBlock.isAllHandler()) {
-			// inner try block can't have catch-all handler
-			outerTryBlock.setBlocks(Utils.concatDistinct(outerTryBlock.getBlocks(), innerTryBlock.getBlocks()));
-			innerTryBlock.clear();
-			return false;
-		}
-		if (makeInner) {
+		if (catchInHandler && (catchInTry || blocksOutsideHandler)) {
 			// convert to inner
 			List<BlockNode> mergedBlocks = Utils.concatDistinct(outerTryBlock.getBlocks(), innerTryBlock.getBlocks());
 			innerTryBlock.getHandlers().removeAll(outerTryBlock.getHandlers());
@@ -299,7 +293,8 @@ private static boolean checkTryCatchRelation(List<TryCatchBlockAttr> tryBlocks,
 			outerTryBlock.setBlocks(mergedBlocks);
 			return false;
 		}
-		if (innerTryBlock.getHandlers().containsAll(outerTryBlock.getHandlers())) {
+		Set<ExceptionHandler> innerHandlerSet = new HashSet<>(innerTryBlock.getHandlers());
+		if (innerHandlerSet.containsAll(outerTryBlock.getHandlers())) {
 			// merge
 			List<BlockNode> mergedBlocks = Utils.concatDistinct(outerTryBlock.getBlocks(), innerTryBlock.getBlocks());
 			List<ExceptionHandler> handlers = Utils.concatDistinct(outerTryBlock.getHandlers(), innerTryBlock.getHandlers());

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

@@ -1,8 +1,12 @@
 package jadx.core.dex.visitors.blocks;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import org.jetbrains.annotations.Nullable;
@@ -20,6 +24,7 @@
 import jadx.core.dex.nodes.Edge;
 import jadx.core.dex.nodes.InsnNode;
 import jadx.core.dex.nodes.MethodNode;
+import jadx.core.dex.trycatch.ExcHandlerAttr;
 import jadx.core.dex.trycatch.TryCatchBlockAttr;
 import jadx.core.dex.visitors.AbstractVisitor;
 import jadx.core.utils.BlockUtils;
@@ -264,7 +269,7 @@ private static boolean modifyBlocksTree(MethodNode mth) {
 		if (mergeConstReturn(mth)) {
 			return true;
 		}
-		return splitReturnBlocks(mth);
+		return splitExitBlocks(mth);
 	}
 
 	private static boolean mergeConstReturn(MethodNode mth) {
@@ -478,11 +483,13 @@ private static boolean splitLoops(MethodNode mth, BlockNode block, List<LoopInfo
 		return false;
 	}
 
-	private static boolean splitReturnBlocks(MethodNode mth) {
+	private static boolean splitExitBlocks(MethodNode mth) {
 		boolean changed = false;
 		for (BlockNode preExitBlock : mth.getPreExitBlocks()) {
 			if (splitReturn(mth, preExitBlock)) {
 				changed = true;
+			} else if (splitThrow(mth, preExitBlock)) {
+				changed = true;
 			}
 		}
 		if (changed) {
@@ -522,7 +529,7 @@ private static boolean splitReturn(MethodNode mth, BlockNode returnBlock) {
 		}
 		if (returnInsn.getArgsCount() == 1
 				&& returnBlock.getInstructions().size() == 1
-				&& !isReturnArgAssignInPred(preds, returnInsn)) {
+				&& !isArgAssignInPred(preds, returnInsn.getArg(0))) {
 			return false;
 		}
 
@@ -546,11 +553,67 @@ private static boolean splitReturn(MethodNode mth, BlockNode returnBlock) {
 		return true;
 	}
 
-	private static boolean isReturnArgAssignInPred(List<BlockNode> preds, InsnNode returnInsn) {
-		InsnArg retArg = returnInsn.getArg(0);
-		if (retArg.isRegister()) {
-			RegisterArg arg = (RegisterArg) retArg;
-			int regNum = arg.getRegNum();
+	private static boolean splitThrow(MethodNode mth, BlockNode exitBlock) {
+		if (exitBlock.contains(AFlag.IGNORE_THROW_SPLIT)) {
+			return false;
+		}
+		List<BlockNode> preds = exitBlock.getPredecessors();
+		if (preds.size() < 2) {
+			return false;
+		}
+		InsnNode throwInsn = BlockUtils.getLastInsn(exitBlock);
+		if (throwInsn == null || throwInsn.getType() != InsnType.THROW) {
+			return false;
+		}
+		// split only for several exception handlers
+		// traverse predecessors to exception handler
+		Map<BlockNode, ExcHandlerAttr> handlersMap = new HashMap<>(preds.size());
+		Set<BlockNode> handlers = new HashSet<>(preds.size());
+		for (BlockNode pred : preds) {
+			BlockUtils.visitPredecessorsUntil(mth, pred, block -> {
+				ExcHandlerAttr excHandlerAttr = block.get(AType.EXC_HANDLER);
+				if (excHandlerAttr == null) {
+					return false;
+				}
+				boolean correctHandler = excHandlerAttr.getHandler().getBlocks().contains(block);
+				if (correctHandler && isArgAssignInPred(Collections.singletonList(block), throwInsn.getArg(0))) {
+					handlersMap.put(pred, excHandlerAttr);
+					handlers.add(block);
+				}
+				return correctHandler;
+			});
+		}
+		if (handlers.size() == 1) {
+			exitBlock.add(AFlag.IGNORE_THROW_SPLIT);
+			return false;
+		}
+
+		boolean first = true;
+		for (BlockNode pred : new ArrayList<>(preds)) {
+			if (first) {
+				first = false;
+			} else {
+				BlockNode newThrowBlock = BlockSplitter.startNewBlock(mth, -1);
+				newThrowBlock.add(AFlag.SYNTHETIC);
+				for (InsnNode oldInsn : exitBlock.getInstructions()) {
+					InsnNode copyInsn = oldInsn.copyWithoutSsa();
+					copyInsn.add(AFlag.SYNTHETIC);
+					newThrowBlock.getInstructions().add(copyInsn);
+				}
+				newThrowBlock.copyAttributesFrom(exitBlock);
+				ExcHandlerAttr excHandlerAttr = handlersMap.get(pred);
+				if (excHandlerAttr != null) {
+					excHandlerAttr.getHandler().addBlock(newThrowBlock);
+				}
+				BlockSplitter.replaceConnection(pred, exitBlock, newThrowBlock);
+			}
+		}
+		return true;
+	}
+
+	private static boolean isArgAssignInPred(List<BlockNode> preds, InsnArg arg) {
+		if (arg.isRegister()) {
+			int regNum = ((RegisterArg) arg).getRegNum();
 			for (BlockNode pred : preds) {
 				for (InsnNode insnNode : pred.getInstructions()) {
 					RegisterArg result = insnNode.getResult();

jadx-core/src/main/java/jadx/core/dex/visitors/finaly/MarkFinallyVisitor.java

@@ -44,7 +44,6 @@
 )
 public class MarkFinallyVisitor extends AbstractVisitor {
 	private static final Logger LOG = LoggerFactory.getLogger(MarkFinallyVisitor.class);
-	// private static final Logger LOG = LoggerFactory.getLogger(MarkFinallyVisitor.class);
 
 	@Override
 	public void visit(MethodNode mth) {
@@ -164,15 +163,22 @@ private static boolean extractFinally(MethodNode mth, TryCatchBlockAttr tryBlock
 			mergeInnerTryBlocks = false;
 		}
 
-		// remove 'finally' from 'try' blocks, check all up paths on each exit (connected with finally exit)
+		// remove 'finally' from 'try' blocks,
+		// check all up paths on each exit (connected with 'finally' exit)
 		List<BlockNode> tryBlocks = allHandler.getTryBlock().getBlocks();
 		BlockNode bottomBlock = BlockUtils.getBottomBlock(allHandler.getBlocks());
 		if (bottomBlock == null) {
+			if (Consts.DEBUG_FINALLY) {
+				LOG.warn("No bottom block for handler: {} and blocks: {}", allHandler, allHandler.getBlocks());
+			}
 			return false;
 		}
 		BlockNode bottomFinallyBlock = BlockUtils.followEmptyPath(bottomBlock);
 		BlockNode bottom = BlockUtils.getNextBlock(bottomFinallyBlock);
 		if (bottom == null) {
+			if (Consts.DEBUG_FINALLY) {
+				LOG.warn("Finally bottom block not found for: {} and: {}", bottomBlock, bottomFinallyBlock);
+			}
 			return false;
 		}
 		boolean found = false;
@@ -197,12 +203,15 @@ private static boolean extractFinally(MethodNode mth, TryCatchBlockAttr tryBlock
 				}
 			}
 		}
-		if (Consts.DEBUG_FINALLY) {
-			LOG.debug("Result slices:\n{}", extractInfo);
-		}
 		if (!found) {
+			if (Consts.DEBUG_FINALLY) {
+				LOG.info("Dup not found for all handler: {}", allHandler);
+			}
 			return false;
 		}
+		if (Consts.DEBUG_FINALLY) {
+			LOG.debug("Result slices:\n{}", extractInfo);
+		}
 		if (!checkSlices(extractInfo)) {
 			mth.addWarnComment("Finally extract failed");
 			return false;

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

@@ -512,22 +512,29 @@ public static List<BlockNode> collectPredecessors(MethodNode mth, BlockNode star
 			bs.or(blocksToBitSet(mth, stopBlocks));
 		}
 		List<BlockNode> list = new ArrayList<>();
-		traversePredecessors(start, bs, list::add);
+		traversePredecessors(start, bs, block -> {
+			list.add(block);
+			return false;
+		});
 		return list;
 	}
 
+	public static void visitPredecessorsUntil(MethodNode mth, BlockNode start, Predicate<BlockNode> visitor) {
+		traversePredecessors(start, newBlocksBitSet(mth), visitor);
+	}
+
 	/**
-	 * Up BFS
+	 * Up BFS.
+	 * To stop return true from predicate
 	 */
-	private static void traversePredecessors(BlockNode start, BitSet visited, Consumer<BlockNode> visitor) {
+	private static void traversePredecessors(BlockNode start, BitSet visited, Predicate<BlockNode> visitor) {
 		Queue<BlockNode> queue = new ArrayDeque<>();
 		queue.add(start);
 		while (true) {
 			BlockNode current = queue.poll();
-			if (current == null) {
+			if (current == null || visitor.test(current)) {
 				return;
 			}
-			visitor.accept(current);
 			for (BlockNode next : current.getPredecessors()) {
 				int id = next.getId();
 				if (!visited.get(id)) {

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

@@ -0,0 +1,19 @@
+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 TestNestedTryCatch5 extends SmaliTest {
+
+	@Test
+	public void test() {
+		disableCompilation();
+		assertThat(getClassNodeFromSmali())
+				.code()
+				.doesNotContain("?? ")
+				.countString(3, "throw ");
+	}
+}