core: bind blocks for target instructions at early stage

jadx-core/src/main/java/jadx/core/dex/instructions/GotoNode.java

 package jadx.core.dex.instructions;
 
-import jadx.core.dex.nodes.InsnNode;
+import jadx.core.dex.nodes.BlockNode;
 import jadx.core.utils.InsnUtils;
 
-public class GotoNode extends InsnNode {
+public class GotoNode extends TargetInsnNode {
 
 	protected int target;
 
@@ -20,6 +20,15 @@ public class GotoNode extends InsnNode {
 		return target;
 	}
 
+	@Override
+	public boolean replaceTargetBlock(BlockNode origin, BlockNode replace) {
+		return false;
+	}
+
+	@Override
+	public void initBlocks(BlockNode curBlock) {
+	}
+
 	@Override
 	public String toString() {
 		return super.toString() + "-> " + InsnUtils.formatOffset(target);

jadx-core/src/main/java/jadx/core/dex/instructions/IfNode.java

@@ -56,6 +56,7 @@ public class IfNode extends GotoNode {
 		setArg(1, arg2);
 	}
 
+	@Override
 	public void initBlocks(BlockNode curBlock) {
 		thenBlock = getBlockByOffset(target, curBlock.getSuccessors());
 		if (curBlock.getSuccessors().size() == 1) {
@@ -65,6 +66,19 @@ public class IfNode extends GotoNode {
 		}
 	}
 
+	@Override
+	public boolean replaceTargetBlock(BlockNode origin, BlockNode replace) {
+		if (thenBlock == origin) {
+			thenBlock = replace;
+			return true;
+		}
+		if (elseBlock == origin) {
+			elseBlock = replace;
+			return true;
+		}
+		return false;
+	}
+
 	public BlockNode getThenBlock() {
 		return thenBlock;
 	}

jadx-core/src/main/java/jadx/core/dex/instructions/SwitchNode.java

 package jadx.core.dex.instructions;
 
 import java.util.Arrays;
+import java.util.List;
 
 import jadx.core.dex.instructions.args.InsnArg;
+import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.InsnNode;
 import jadx.core.utils.InsnUtils;
 
-public class SwitchNode extends InsnNode {
+import static jadx.core.utils.BlockUtils.getBlockByOffset;
+
+public class SwitchNode extends TargetInsnNode {
 
 	private final Object[] keys;
 	private final int[] targets;
 	private final int def; // next instruction
 
+	private BlockNode[] targetBlocks;
+	private BlockNode defTargetBlock;
+
 	public SwitchNode(InsnArg arg, Object[] keys, int[] targets, int def) {
 		super(InsnType.SWITCH, 1);
 		this.keys = keys;
@@ -36,6 +43,42 @@ public class SwitchNode extends InsnNode {
 		return def;
 	}
 
+	public BlockNode[] getTargetBlocks() {
+		return targetBlocks;
+	}
+
+	public BlockNode getDefTargetBlock() {
+		return defTargetBlock;
+	}
+
+	@Override
+	public void initBlocks(BlockNode curBlock) {
+		List<BlockNode> successors = curBlock.getSuccessors();
+		int len = targets.length;
+		targetBlocks = new BlockNode[len];
+		for (int i = 0; i < len; i++) {
+			targetBlocks[i] = getBlockByOffset(targets[i], successors);
+		}
+		defTargetBlock = getBlockByOffset(def, successors);
+	}
+
+	@Override
+	public boolean replaceTargetBlock(BlockNode origin, BlockNode replace) {
+		int count = 0;
+		int len = targetBlocks.length;
+		for (int i = 0; i < len; i++) {
+			if (targetBlocks[i] == origin) {
+				targetBlocks[i] = replace;
+				count++;
+			}
+		}
+		if (defTargetBlock == origin) {
+			defTargetBlock = replace;
+			count++;
+		}
+		return count > 0;
+	}
+
 	@Override
 	public boolean isSame(InsnNode obj) {
 		if (this == obj) {

jadx-core/src/main/java/jadx/core/dex/instructions/TargetInsnNode.java

+package jadx.core.dex.instructions;
+
+import jadx.core.dex.nodes.BlockNode;
+import jadx.core.dex.nodes.InsnNode;
+
+public abstract class TargetInsnNode extends InsnNode {
+
+	public TargetInsnNode(InsnType type, int argsCount) {
+		super(type, argsCount);
+	}
+
+	public abstract void initBlocks(BlockNode curBlock);
+
+	public abstract boolean replaceTargetBlock(BlockNode origin, BlockNode replace);
+}

jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinallyExtract.java

@@ -577,8 +577,7 @@ public class BlockFinallyExtract extends AbstractVisitor {
 			BlockNode newPred = BlockSplitter.insertBlockBetween(mth, pred, sOut);
 			for (BlockNode predBlock : new ArrayList<>(sOut.getPredecessors())) {
 				if (predBlock != newPred) {
-					removeConnection(predBlock, sOut);
-					connect(predBlock, newPred);
+					BlockSplitter.replaceConnection(predBlock, sOut, newPred);
 				}
 			}
 			rOut.getPredecessors().clear();
@@ -605,6 +604,7 @@ public class BlockFinallyExtract extends AbstractVisitor {
 			connect(middle, startBlock);
 			addIgnoredEdge(middle, startBlock);
 			connect(middle, rOut);
+			BlockSplitter.replaceTarget(middle, remBlock, rOut);
 		}
 
 		// mark blocks for remove

jadx-core/src/main/java/jadx/core/dex/visitors/blocksmaker/BlockFinish.java

 package jadx.core.dex.visitors.blocksmaker;
 
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import jadx.core.dex.attributes.AType;
-import jadx.core.dex.instructions.IfNode;
-import jadx.core.dex.instructions.InsnType;
 import jadx.core.dex.nodes.BlockNode;
-import jadx.core.dex.nodes.InsnNode;
 import jadx.core.dex.nodes.MethodNode;
 import jadx.core.dex.trycatch.ExcHandlerAttr;
 import jadx.core.dex.trycatch.SplitterBlockAttr;
@@ -30,26 +26,12 @@ public class BlockFinish extends AbstractVisitor {
 
 		for (BlockNode block : mth.getBasicBlocks()) {
 			block.updateCleanSuccessors();
-			initBlocksInIfNodes(block);
 			fixSplitterBlock(block);
 		}
 
 		mth.finishBasicBlocks();
 	}
 
-	/**
-	 * Init 'then' and 'else' blocks for 'if' instruction.
-	 */
-	private static void initBlocksInIfNodes(BlockNode block) {
-		List<InsnNode> instructions = block.getInstructions();
-		if (instructions.size() == 1) {
-			InsnNode insn = instructions.get(0);
-			if (insn.getType() == InsnType.IF) {
-				((IfNode) insn).initBlocks(block);
-			}
-		}
-	}
-
 	/**
 	 * For evey exception handler must be only one splitter block,
 	 * select correct one and remove others if necessary.

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

@@ -24,7 +24,6 @@ import jadx.core.utils.BlockUtils;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
 import static jadx.core.dex.visitors.blocksmaker.BlockSplitter.connect;
-import static jadx.core.dex.visitors.blocksmaker.BlockSplitter.removeConnection;
 import static jadx.core.utils.EmptyBitSet.EMPTY;
 
 public class BlockProcessor extends AbstractVisitor {
@@ -39,7 +38,7 @@ public class BlockProcessor extends AbstractVisitor {
 	}
 
 	public static void rerun(MethodNode mth) {
-		removeBlocks(mth);
+		removeMarkedBlocks(mth);
 		clearBlocksState(mth);
 		processBlocksTree(mth);
 	}
@@ -247,65 +246,69 @@ public class BlockProcessor extends AbstractVisitor {
 			if (block.getPredecessors().isEmpty() && block != mth.getEnterBlock()) {
 				throw new JadxRuntimeException("Unreachable block: " + block);
 			}
+			if (checkLoops(mth, block)) {
+				return true;
+			}
+		}
+		return splitReturn(mth);
+	}
 
-			// check loops
-			List<LoopInfo> loops = block.getAll(AType.LOOP);
-			if (loops.size() > 1) {
-				boolean oneHeader = true;
-				for (LoopInfo loop : loops) {
-					if (loop.getStart() != block) {
-						oneHeader = false;
-						break;
-					}
+	private static boolean checkLoops(MethodNode mth, BlockNode block) {
+		// check loops
+		List<LoopInfo> loops = block.getAll(AType.LOOP);
+		if (loops.size() > 1) {
+			boolean oneHeader = true;
+			for (LoopInfo loop : loops) {
+				if (loop.getStart() != block) {
+					oneHeader = false;
+					break;
+				}
+			}
+			if (oneHeader) {
+				// several back edges connected to one loop header => make additional block
+				BlockNode newLoopEnd = BlockSplitter.startNewBlock(mth, block.getStartOffset());
+				newLoopEnd.add(AFlag.SYNTHETIC);
+				connect(newLoopEnd, block);
+				for (LoopInfo la : loops) {
+					BlockSplitter.replaceConnection(la.getEnd(), block, newLoopEnd);
 				}
-				if (oneHeader) {
-					// several back edges connected to one loop header => make additional block
-					BlockNode newLoopHeader = BlockSplitter.startNewBlock(mth, block.getStartOffset());
-					newLoopHeader.add(AFlag.SYNTHETIC);
-					connect(newLoopHeader, block);
-					for (LoopInfo la : loops) {
-						BlockNode node = la.getEnd();
-						removeConnection(node, block);
-						connect(node, newLoopHeader);
+				return true;
+			}
+		}
+		if (loops.size() == 1) {
+			LoopInfo loop = loops.get(0);
+			// insert additional blocks for possible 'break' insertion
+			List<Edge> edges = loop.getExitEdges();
+			if (!edges.isEmpty()) {
+				boolean change = false;
+				for (Edge edge : edges) {
+					BlockNode target = edge.getTarget();
+					if (!target.contains(AFlag.SYNTHETIC)) {
+						BlockSplitter.insertBlockBetween(mth, edge.getSource(), target);
+						change = true;
 					}
+				}
+				if (change) {
 					return true;
 				}
 			}
-			if (loops.size() == 1) {
-				LoopInfo loop = loops.get(0);
-				// insert additional blocks for possible 'break' insertion
-				List<Edge> edges = loop.getExitEdges();
-				if (!edges.isEmpty()) {
-					boolean change = false;
-					for (Edge edge : edges) {
-						BlockNode target = edge.getTarget();
-						if (!target.contains(AFlag.SYNTHETIC)) {
-							BlockSplitter.insertBlockBetween(mth, edge.getSource(), target);
-							change = true;
-						}
-					}
-					if (change) {
-						return true;
+			// insert additional blocks for possible 'continue' insertion
+			BlockNode loopEnd = loop.getEnd();
+			if (loopEnd.getPredecessors().size() > 1) {
+				boolean change = false;
+				List<BlockNode> nodes = new ArrayList<>(loopEnd.getPredecessors());
+				for (BlockNode pred : nodes) {
+					if (!pred.contains(AFlag.SYNTHETIC)) {
+						BlockSplitter.insertBlockBetween(mth, pred, loopEnd);
+						change = true;
 					}
 				}
-				// insert additional blocks for possible 'continue' insertion
-				BlockNode loopEnd = loop.getEnd();
-				if (loopEnd.getPredecessors().size() > 1) {
-					boolean change = false;
-					List<BlockNode> nodes = new ArrayList<>(loopEnd.getPredecessors());
-					for (BlockNode pred : nodes) {
-						if (!pred.contains(AFlag.SYNTHETIC)) {
-							BlockSplitter.insertBlockBetween(mth, pred, loopEnd);
-							change = true;
-						}
-					}
-					if (change) {
-						return true;
-					}
+				if (change) {
+					return true;
 				}
 			}
 		}
-		return splitReturn(mth);
+		return false;
 	}
 
 	/**
@@ -334,7 +337,7 @@ public class BlockProcessor extends AbstractVisitor {
 		}
 		boolean first = true;
 		for (BlockNode pred : preds) {
-			BlockNode newRetBlock = BlockSplitter.startNewBlock(mth, exitBlock.getStartOffset());
+			BlockNode newRetBlock = BlockSplitter.startNewBlock(mth, -1);
 			newRetBlock.add(AFlag.SYNTHETIC);
 			InsnNode newRetInsn;
 			if (first) {
@@ -345,8 +348,7 @@ public class BlockProcessor extends AbstractVisitor {
 				newRetInsn = duplicateReturnInsn(returnInsn);
 			}
 			newRetBlock.getInstructions().add(newRetInsn);
-			removeConnection(pred, exitBlock);
-			connect(pred, newRetBlock);
+			BlockSplitter.replaceConnection(pred, exitBlock, newRetBlock);
 		}
 		cleanExitNodes(mth);
 		return true;
@@ -389,35 +391,32 @@ public class BlockProcessor extends AbstractVisitor {
 		return insn;
 	}
 
-	private static void removeBlocks(MethodNode mth) {
-		Iterator<BlockNode> it = mth.getBasicBlocks().iterator();
-		while (it.hasNext()) {
-			BlockNode block = it.next();
+	private static void removeMarkedBlocks(MethodNode mth) {
+		mth.getBasicBlocks().removeIf(block -> {
 			if (block.contains(AFlag.REMOVE)) {
-				if (!block.getPredecessors().isEmpty()
-						|| !block.getSuccessors().isEmpty()) {
-					LOG.error("Block {} not deleted, method: {}", block, mth);
+				if (!block.getPredecessors().isEmpty() || !block.getSuccessors().isEmpty()) {
+					LOG.warn("Block {} not deleted, method: {}", block, mth);
 				} else {
 					CatchAttr catchAttr = block.get(AType.CATCH_BLOCK);
 					if (catchAttr != null) {
 						catchAttr.getTryBlock().removeBlock(mth, block);
 					}
-					it.remove();
+					return true;
 				}
 			}
-		}
+			return false;
+		});
 	}
 
 	private static void clearBlocksState(MethodNode mth) {
-		for (BlockNode block : mth.getBasicBlocks()) {
+		mth.getBasicBlocks().forEach(block -> {
 			block.remove(AType.LOOP);
 			block.remove(AFlag.LOOP_START);
 			block.remove(AFlag.LOOP_END);
-
 			block.setDoms(null);
 			block.setIDom(null);
 			block.setDomFrontier(null);
 			block.getDominatesOn().clear();
-		}
+		});
 	}
 }

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

@@ -11,6 +11,7 @@ import jadx.core.dex.attributes.AType;
 import jadx.core.dex.attributes.nodes.JumpInfo;
 import jadx.core.dex.instructions.IfNode;
 import jadx.core.dex.instructions.InsnType;
+import jadx.core.dex.instructions.TargetInsnNode;
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.InsnNode;
 import jadx.core.dex.nodes.MethodNode;
@@ -18,6 +19,7 @@ import jadx.core.dex.trycatch.CatchAttr;
 import jadx.core.dex.trycatch.ExceptionHandler;
 import jadx.core.dex.trycatch.SplitterBlockAttr;
 import jadx.core.dex.visitors.AbstractVisitor;
+import jadx.core.utils.BlockUtils;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
 public class BlockSplitter extends AbstractVisitor {
@@ -42,7 +44,20 @@ public class BlockSplitter extends AbstractVisitor {
 		mth.initBasicBlocks();
 		splitBasicBlocks(mth);
 		removeInsns(mth);
-		removeEmptyBlocks(mth);
+		removeEmptyDetachedBlocks(mth);
+		initBlocksInTargetNodes(mth);
+	}
+
+	/**
+	 * Init 'then' and 'else' blocks for 'if' instruction.
+	 */
+	private static void initBlocksInTargetNodes(MethodNode mth) {
+		mth.getBasicBlocks().forEach(block -> {
+			InsnNode lastInsn = BlockUtils.getLastInsn(block);
+			if (lastInsn instanceof TargetInsnNode) {
+				((TargetInsnNode) lastInsn).initBlocks(block);
+			}
+		});
 	}
 
 	private static void splitBasicBlocks(MethodNode mth) {
@@ -133,15 +148,31 @@ public class BlockSplitter extends AbstractVisitor {
 		to.getPredecessors().remove(from);
 	}
 
+	static void replaceConnection(BlockNode source, BlockNode oldDest, BlockNode newDest) {
+		removeConnection(source, oldDest);
+		connect(source, newDest);
+		replaceTarget(source, oldDest, newDest);
+	}
+
 	static BlockNode insertBlockBetween(MethodNode mth, BlockNode source, BlockNode target) {
 		BlockNode newBlock = startNewBlock(mth, target.getStartOffset());
 		newBlock.add(AFlag.SYNTHETIC);
 		removeConnection(source, target);
 		connect(source, newBlock);
 		connect(newBlock, target);
+		replaceTarget(source, target, newBlock);
+		source.updateCleanSuccessors();
+		newBlock.updateCleanSuccessors();
 		return newBlock;
 	}
 
+	static void replaceTarget(BlockNode source, BlockNode oldTarget, BlockNode newTarget) {
+		InsnNode lastInsn = BlockUtils.getLastInsn(source);
+		if (lastInsn instanceof TargetInsnNode) {
+			((TargetInsnNode) lastInsn).replaceTargetBlock(oldTarget, newTarget);
+		}
+	}
+
 	private static void setupConnections(MethodNode mth, Map<Integer, BlockNode> blocksMap) {
 		for (BlockNode block : mth.getBasicBlocks()) {
 			for (InsnNode insn : block.getInstructions()) {
@@ -222,7 +253,7 @@ public class BlockSplitter extends AbstractVisitor {
 		}
 	}
 
-	private void removeEmptyBlocks(MethodNode mth) {
+	static void removeEmptyDetachedBlocks(MethodNode mth) {
 		mth.getBasicBlocks().removeIf(block ->
 				block.getInstructions().isEmpty()
 						&& block.getPredecessors().isEmpty()

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

@@ -2,8 +2,6 @@ package jadx.core.dex.visitors.regions;
 
 import java.util.ArrayList;
 import java.util.BitSet;
-import java.util.Collections;
-import java.util.Comparator;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -46,13 +44,11 @@ import jadx.core.utils.ErrorsCounter;
 import jadx.core.utils.InstructionRemover;
 import jadx.core.utils.RegionUtils;
 import jadx.core.utils.exceptions.JadxOverflowException;
-import jadx.core.utils.exceptions.JadxRuntimeException;
 
 import static jadx.core.dex.visitors.regions.IfMakerHelper.confirmMerge;
 import static jadx.core.dex.visitors.regions.IfMakerHelper.makeIfInfo;
 import static jadx.core.dex.visitors.regions.IfMakerHelper.mergeNestedIfNodes;
 import static jadx.core.dex.visitors.regions.IfMakerHelper.searchNestedIf;
-import static jadx.core.utils.BlockUtils.getBlockByOffset;
 import static jadx.core.utils.BlockUtils.getNextBlock;
 import static jadx.core.utils.BlockUtils.isPathExists;
 import static jadx.core.utils.BlockUtils.skipSyntheticSuccessor;
@@ -691,27 +687,14 @@ public class RegionMaker {
 
 		int len = insn.getTargets().length;
 		// sort by target
-		Map<Integer, List<Object>> casesMap = new LinkedHashMap<>(len);
+		Map<BlockNode, List<Object>> blocksMap = new LinkedHashMap<>(len);
 		for (int i = 0; i < len; i++) {
 			Object key = insn.getKeys()[i];
-			int targ = insn.getTargets()[i];
-			List<Object> keys = casesMap.get(targ);
-			if (keys == null) {
-				keys = new ArrayList<>(2);
-				casesMap.put(targ, keys);
-			}
+			BlockNode targ = insn.getTargetBlocks()[i];
+			List<Object> keys = blocksMap.computeIfAbsent(targ, k -> new ArrayList<>(2));
 			keys.add(key);
 		}
-
-		Map<BlockNode, List<Object>> blocksMap = new LinkedHashMap<>(len);
-		for (Map.Entry<Integer, List<Object>> entry : casesMap.entrySet()) {
-			BlockNode c = getBlockByOffset(entry.getKey(), block.getSuccessors());
-			if (c == null) {
-				throw new JadxRuntimeException("Switch block not found by offset: " + entry.getKey());
-			}
-			blocksMap.put(c, entry.getValue());
-		}
-		BlockNode defCase = getBlockByOffset(insn.getDefaultCaseOffset(), block.getSuccessors());
+		BlockNode defCase = insn.getDefTargetBlock();
 		if (defCase != null) {
 			blocksMap.remove(defCase);
 		}
@@ -860,19 +843,16 @@ public class RegionMaker {
 			Map<BlockNode, BlockNode> fallThroughCases) {
 		List<BlockNode> list = new ArrayList<>(blocksMap.size());
 		list.addAll(blocksMap.keySet());
-		Collections.sort(list, new Comparator<BlockNode>() {
-			@Override
-			public int compare(BlockNode a, BlockNode b) {
-				BlockNode nextA = fallThroughCases.get(a);
-				if (nextA != null) {
-					if (b.equals(nextA)) {
-						return -1;
-					}
-				} else if (a.equals(fallThroughCases.get(b))) {
-					return 1;
+		list.sort((a, b) -> {
+			BlockNode nextA = fallThroughCases.get(a);
+			if (nextA != null) {
+				if (b.equals(nextA)) {
+					return -1;
 				}
-				return 0;
+			} else if (a.equals(fallThroughCases.get(b))) {
+				return 1;
 			}
+			return 0;
 		});
 
 		Map<BlockNode, List<Object>> newBlocksMap = new LinkedHashMap<>(blocksMap.size());

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

@@ -453,13 +453,12 @@ public class BlockUtils {
 	 */
 	public static List<BlockNode> collectBlocksDominatedBy(BlockNode dominator, BlockNode start) {
 		List<BlockNode> result = new ArrayList<>();
-		HashSet<BlockNode> visited = new HashSet<BlockNode>();
+		Set<BlockNode> visited = new HashSet<>();
 		collectWhileDominates(dominator, start, result, visited);
 		return result;
 	}
 
-	private static void collectWhileDominates(BlockNode dominator, BlockNode child, List<BlockNode> result,
-			HashSet<BlockNode> visited) {
+	private static void collectWhileDominates(BlockNode dominator, BlockNode child, List<BlockNode> result, Set<BlockNode> visited) {
 		if (visited.contains(child)) {
 			return;
 		}