feat: add method to update all blocks related info in method (#2335)

skylot · Nov 2, 2024 · 417bb7a · 417bb7a
1 parent 4c74e8e
commit 417bb7a
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 0 deletions.
diff --git a/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java b/jadx-core/src/main/java/jadx/core/dex/nodes/MethodNode.java
@@ -408,6 +408,10 @@ public boolean isPreExitBlock(BlockNode block) {
 		return exitBlock.getPredecessors().contains(block);
 	}
 
+	public void resetLoops() {
+		this.loops = new ArrayList<>();
+	}
+
 	public void registerLoop(LoopInfo loop) {
 		if (loops.isEmpty()) {
 			loops = new ArrayList<>(5);

diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/blocks/BlockProcessor.java b/jadx-core/src/main/java/jadx/core/dex/visitors/blocks/BlockProcessor.java
@@ -94,6 +94,33 @@ private static void processBlocksTree(MethodNode mth) {
 		updateCleanSuccessors(mth);
 	}
 
+	/**
+	 * Recalculate all additional info attached to blocks:
+	 *
+	 * <pre>
+	 * - dominators
+	 * - dominance frontier
+	 * - post dominators (only if {@link AFlag#COMPUTE_POST_DOM} added to method)
+	 * - loops and nested loop info
+	 * </pre>
+	 *
+	 * This method should be called after changing a block tree in custom passes added before
+	 * {@link BlockFinisher}.
+	 */
+	public static void updateBlocksData(MethodNode mth) {
+		clearBlocksState(mth);
+		DominatorTree.compute(mth);
+		markLoops(mth);
+
+		DominatorTree.computeDominanceFrontier(mth);
+		registerLoops(mth);
+		processNestedLoops(mth);
+
+		PostDominatorTree.compute(mth);
+
+		updateCleanSuccessors(mth);
+	}
+
 	static void updateCleanSuccessors(MethodNode mth) {
 		mth.getBasicBlocks().forEach(BlockNode::updateCleanSuccessors);
 	}
@@ -245,6 +272,7 @@ private static void markLoops(MethodNode mth) {
 	}
 
 	private static void registerLoops(MethodNode mth) {
+		mth.resetLoops();
 		mth.getBasicBlocks().forEach(block -> {
 			if (block.contains(AFlag.LOOP_START)) {
 				block.getAll(AType.LOOP).forEach(mth::registerLoop);