From 4f307c0085666462cd95125e88941653dbfebe40 Mon Sep 17 00:00:00 2001
From: Skylot <skylot@gmail.com>
Date: Thu, 14 Aug 2014 22:38:29 +0400
Subject: [PATCH] core: allow subblock replace for 'if' region

---
 .../jadx/core/dex/regions/AbstractRegion.java |  6 ++++-
 .../java/jadx/core/dex/regions/IfRegion.java  | 13 ++++++++++
 .../regions/ProcessTryCatchRegions.java       | 25 ++++++++++++-------
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/jadx-core/src/main/java/jadx/core/dex/regions/AbstractRegion.java b/jadx-core/src/main/java/jadx/core/dex/regions/AbstractRegion.java
index 3f4728546..c1af10fe2 100644
--- a/jadx-core/src/main/java/jadx/core/dex/regions/AbstractRegion.java
+++ b/jadx-core/src/main/java/jadx/core/dex/regions/AbstractRegion.java
@@ -4,7 +4,11 @@ import jadx.core.dex.attributes.AttrNode;
 import jadx.core.dex.nodes.IContainer;
 import jadx.core.dex.nodes.IRegion;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 public abstract class AbstractRegion extends AttrNode implements IRegion {
+	private static final Logger LOG = LoggerFactory.getLogger(AbstractRegion.class);
 
 	private IRegion parent;
 
@@ -23,7 +27,7 @@ public abstract class AbstractRegion extends AttrNode implements IRegion {
 
 	@Override
 	public boolean replaceSubBlock(IContainer oldBlock, IContainer newBlock) {
-		// TODO: implement for others regions
+		LOG.warn("Replace sub block not supported for class \"{}\"", this.getClass());
 		return false;
 	}
 }
diff --git a/jadx-core/src/main/java/jadx/core/dex/regions/IfRegion.java b/jadx-core/src/main/java/jadx/core/dex/regions/IfRegion.java
index ad82bfd2b..faee43f21 100644
--- a/jadx-core/src/main/java/jadx/core/dex/regions/IfRegion.java
+++ b/jadx-core/src/main/java/jadx/core/dex/regions/IfRegion.java
@@ -101,6 +101,19 @@ public final class IfRegion extends AbstractRegion {
 		return Collections.unmodifiableList(all);
 	}
 
+	@Override
+	public boolean replaceSubBlock(IContainer oldBlock, IContainer newBlock) {
+		if (oldBlock == thenRegion) {
+			thenRegion = newBlock;
+			return true;
+		}
+		if (oldBlock == elseRegion) {
+			elseRegion = newBlock;
+			return true;
+		}
+		return false;
+	}
+
 	@Override
 	public String baseString() {
 		if (ternRegion != null) {
diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/regions/ProcessTryCatchRegions.java b/jadx-core/src/main/java/jadx/core/dex/visitors/regions/ProcessTryCatchRegions.java
index bbbd2bbd9..058c50db8 100644
--- a/jadx-core/src/main/java/jadx/core/dex/visitors/regions/ProcessTryCatchRegions.java
+++ b/jadx-core/src/main/java/jadx/core/dex/visitors/regions/ProcessTryCatchRegions.java
@@ -1,5 +1,6 @@
 package jadx.core.dex.visitors.regions;
 
+import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.attributes.AType;
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.IContainer;
@@ -44,7 +45,7 @@ public class ProcessTryCatchRegions extends AbstractRegionVisitor {
 			DepthRegionTraversal.traverseAll(mth, new AbstractRegionVisitor() {
 				@Override
 				public void leaveRegion(MethodNode mth, IRegion region) {
-					checkAndWrap(tryBlocksMap, region);
+					checkAndWrap(mth, tryBlocksMap, region);
 				}
 			});
 			if (k++ > 100) {
@@ -98,13 +99,16 @@ public class ProcessTryCatchRegions extends AbstractRegionVisitor {
 		}
 	}
 
-	private static void checkAndWrap(Map<BlockNode, TryCatchBlock> tryBlocksMap, IRegion region) {
+	private static void checkAndWrap(MethodNode mth, Map<BlockNode, TryCatchBlock> tryBlocksMap, IRegion region) {
 		// search dominator blocks in this region (don't need to go deeper)
 		for (Map.Entry<BlockNode, TryCatchBlock> entry : tryBlocksMap.entrySet()) {
 			BlockNode dominator = entry.getKey();
 			if (region.getSubBlocks().contains(dominator)) {
 				TryCatchBlock tb = tryBlocksMap.get(dominator);
-				wrapBlocks(region, tb, dominator);
+				if (!wrapBlocks(region, tb, dominator)) {
+					LOG.warn("Can't wrap try/catch for {}, method: {}", dominator, mth);
+					mth.add(AFlag.INCONSISTENT_CODE);
+				}
 				tryBlocksMap.remove(dominator);
 				return;
 			}
@@ -114,9 +118,10 @@ public class ProcessTryCatchRegions extends AbstractRegionVisitor {
 	/**
 	 * Extract all block dominated by 'dominator' to separate region and mark as try/catch block
 	 */
-	private static void wrapBlocks(IRegion region, TryCatchBlock tb, BlockNode dominator) {
+	private static boolean wrapBlocks(IRegion region, TryCatchBlock tb, BlockNode dominator) {
 		Region newRegion = new Region(region);
-		for (IContainer cont : region.getSubBlocks()) {
+		List<IContainer> subBlocks = region.getSubBlocks();
+		for (IContainer cont : subBlocks) {
 			if (RegionUtils.isDominatedBy(dominator, cont)) {
 				if (isHandlerPath(tb, cont)) {
 					break;
@@ -125,13 +130,14 @@ public class ProcessTryCatchRegions extends AbstractRegionVisitor {
 			}
 		}
 		if (newRegion.getSubBlocks().isEmpty()) {
-			return;
+			return false;
 		}
 		// replace first node by region
 		IContainer firstNode = newRegion.getSubBlocks().get(0);
-		int i = region.getSubBlocks().indexOf(firstNode);
-		region.getSubBlocks().set(i, newRegion);
-		region.getSubBlocks().removeAll(newRegion.getSubBlocks());
+		if (!region.replaceSubBlock(firstNode, newRegion)) {
+			return false;
+		}
+		subBlocks.removeAll(newRegion.getSubBlocks());
 
 		newRegion.addAttr(tb.getCatchAttr());
 
@@ -142,6 +148,7 @@ public class ProcessTryCatchRegions extends AbstractRegionVisitor {
 				aReg.setParent(newRegion);
 			}
 		}
+		return true;
 	}
 
 	private static boolean isHandlerPath(TryCatchBlock tb, IContainer cont) {