fix: prefer early return for 'if-else-if' block (#2052)

skylot · Dec 5, 2023 · 2d5c0fd · 2d5c0fd
1 parent 5d56001
commit 2d5c0fd
Show file tree

Hide file tree

Showing 6 changed files with 672 additions and 36 deletions.
diff --git a/jadx-core/src/main/java/jadx/core/codegen/RegionGen.java b/jadx-core/src/main/java/jadx/core/codegen/RegionGen.java
@@ -32,7 +32,6 @@
 import jadx.core.dex.nodes.IBlock;
 import jadx.core.dex.nodes.IContainer;
 import jadx.core.dex.nodes.InsnNode;
-import jadx.core.dex.regions.Region;
 import jadx.core.dex.regions.SwitchRegion;
 import jadx.core.dex.regions.SwitchRegion.CaseInfo;
 import jadx.core.dex.regions.SynchronizedRegion;
@@ -147,15 +146,12 @@ public void makeIf(IfRegion region, ICodeWriter code, boolean newLine) throws Co
 	 * Connect if-else-if block
 	 */
 	private boolean connectElseIf(ICodeWriter code, IContainer els) throws CodegenException {
-		if (els.contains(AFlag.ELSE_IF_CHAIN) && els instanceof Region) {
-			List<IContainer> subBlocks = ((Region) els).getSubBlocks();
-			if (subBlocks.size() == 1) {
-				IContainer elseBlock = subBlocks.get(0);
-				if (elseBlock instanceof IfRegion) {
-					declareVars(code, elseBlock);
-					makeIf((IfRegion) elseBlock, code, false);
-					return true;
-				}
+		if (els.contains(AFlag.ELSE_IF_CHAIN)) {
+			IContainer elseBlock = RegionUtils.getSingleSubBlock(els);
+			if (elseBlock instanceof IfRegion) {
+				declareVars(code, elseBlock);
+				makeIf((IfRegion) elseBlock, code, false);
+				return true;
 			}
 		}
 		return false;

diff --git a/jadx-core/src/main/java/jadx/core/dex/visitors/regions/IfRegionVisitor.java b/jadx-core/src/main/java/jadx/core/dex/visitors/regions/IfRegionVisitor.java
@@ -13,6 +13,7 @@
 import jadx.core.dex.regions.conditions.IfCondition.Mode;
 import jadx.core.dex.regions.conditions.IfRegion;
 import jadx.core.dex.visitors.AbstractVisitor;
+import jadx.core.utils.InsnUtils;
 import jadx.core.utils.RegionUtils;
 
 import static jadx.core.utils.RegionUtils.insnsCount;
@@ -81,21 +82,23 @@ private static void orderBranches(MethodNode mth, IfRegion ifRegion) {
 					return;
 				}
 			}
-			boolean lastRegion = RegionUtils.hasExitEdge(ifRegion);
-			if (elseSize == 1 && lastRegion && mth.isVoidReturn()) {
-				InsnNode lastElseInsn = RegionUtils.getLastInsn(ifRegion.getElseRegion());
-				if (lastElseInsn != null && lastElseInsn.getType() == InsnType.THROW) {
-					// move `throw` into `then` block
+			if (elseSize == 1) {
+				boolean lastRegion = RegionUtils.hasExitEdge(ifRegion);
+				if (lastRegion && mth.isVoidReturn()) {
+					InsnNode lastElseInsn = RegionUtils.getLastInsn(ifRegion.getElseRegion());
+					if (InsnUtils.isInsnType(lastElseInsn, InsnType.THROW)) {
+						// move `throw` into `then` block
+						invertIfRegion(ifRegion);
+					} else {
+						// single return at method end will be removed later
+					}
+					return;
+				}
+				if (thenSize > 2 && !(lastRegion && thenSize < 4 /* keep small code block inside else */)) {
 					invertIfRegion(ifRegion);
-				} else {
-					// single return at method end will be removed later
+					return;
 				}
-				return;
-			}
-			if (!lastRegion) {
-				invertIfRegion(ifRegion);
 			}
-			return;
 		}
 		boolean thenExit = RegionUtils.hasExitBlock(ifRegion.getThenRegion());
 		boolean elseExit = RegionUtils.hasExitBlock(ifRegion.getElseRegion());
@@ -155,22 +158,32 @@ public boolean visitRegion(MethodNode mth, IRegion region) {
 		}
 	}
 
+	@SuppressWarnings("StatementWithEmptyBody")
 	private static boolean removeRedundantElseBlock(MethodNode mth, IfRegion ifRegion) {
-		if (ifRegion.getElseRegion() == null
-				|| ifRegion.contains(AFlag.ELSE_IF_CHAIN)
-				|| ifRegion.getElseRegion().contains(AFlag.ELSE_IF_CHAIN)) {
+		if (ifRegion.getElseRegion() == null) {
 			return false;
 		}
 		if (!RegionUtils.hasExitBlock(ifRegion.getThenRegion())) {
 			return false;
 		}
-		// code style check:
-		// will remove 'return;' from 'then' and 'else' with one instruction
-		// see #jadx.tests.integration.conditions.TestConditions9
-		if (mth.isVoidReturn()
-				&& insnsCount(ifRegion.getThenRegion()) == 2
-				&& insnsCount(ifRegion.getElseRegion()) == 2) {
-			return false;
+		InsnNode lastThanInsn = RegionUtils.getLastInsn(ifRegion.getThenRegion());
+		if (InsnUtils.isInsnType(lastThanInsn, InsnType.THROW)) {
+			// always omit else after 'throw'
+		} else {
+			// code style check:
+			// will remove 'return;' from 'then' and 'else' with one instruction
+			// see #jadx.tests.integration.conditions.TestConditions9
+			if (mth.isVoidReturn()) {
+				int thenSize = insnsCount(ifRegion.getThenRegion());
+				// keep small blocks with same or 'similar' size unchanged
+				if (thenSize < 5) {
+					int elseSize = insnsCount(ifRegion.getElseRegion());
+					int range = ifRegion.getElseRegion().contains(AFlag.ELSE_IF_CHAIN) ? 4 : 2;
+					if (thenSize == elseSize || (thenSize * range > elseSize && thenSize < elseSize * range)) {
+						return false;
+					}
+				}
+			}
 		}
 		IRegion parent = ifRegion.getParent();
 		Region newRegion = new Region(parent);

diff --git a/jadx-core/src/main/java/jadx/core/utils/RegionUtils.java b/jadx-core/src/main/java/jadx/core/utils/RegionUtils.java
@@ -324,6 +324,30 @@ public static boolean isRegionContainsBlock(IContainer container, BlockNode bloc
 		}
 	}
 
+	public static IContainer getSingleSubBlock(IContainer container) {
+		if (container instanceof Region) {
+			List<IContainer> subBlocks = ((Region) container).getSubBlocks();
+			if (subBlocks.size() == 1) {
+				return ignoreSimpleRegionWrapper(subBlocks.get(0));
+			}
+		}
+		return null;
+	}
+
+	private static IContainer ignoreSimpleRegionWrapper(IContainer container) {
+		while (true) {
+			if (container instanceof Region) {
+				List<IContainer> subBlocks = ((Region) container).getSubBlocks();
+				if (subBlocks.size() != 1) {
+					return container;
+				}
+				container = subBlocks.get(0);
+			} else {
+				return container;
+			}
+		}
+	}
+
 	public static List<IContainer> getExcHandlersForRegion(IContainer region) {
 		TryCatchBlockAttr tb = region.get(AType.TRY_BLOCK);
 		if (tb != null) {

diff --git a/jadx-core/src/test/java/jadx/tests/integration/conditions/TestIfCodeStyle.java b/jadx-core/src/test/java/jadx/tests/integration/conditions/TestIfCodeStyle.java
@@ -129,8 +129,9 @@ public void test() {
 		noDebugInfo();
 		assertThat(getClassNode(TestCls.class))
 				.code()
-				.doesNotContain("else")
-				.countString(8, "return;")
+				// allow one last 'else'
+				.oneOf(c -> c.doesNotContain("else").countString(8, "return;"),
+						c -> c.countString(1, "else").countString(7, "return;"))
 				.containsLines(2,
 						"if (readInt < 0) {",
 						indent() + "if (dataPosition > Integer.MAX_VALUE - readInt) {",
@@ -146,8 +147,9 @@ public void testSmali() {
 		disableCompilation();
 		assertThat(getClassNodeFromSmali())
 				.code()
-				.doesNotContain("else")
-				.countString(8, "return;")
+				// allow one last 'else'
+				.oneOf(c -> c.doesNotContain("else").countString(8, "return;"),
+						c -> c.countString(1, "else").countString(7, "return;"))
 				.containsLines(2,
 						"if (_aidl_parcelable_size < 0) {",
 						indent() + "if (_aidl_start_pos > Integer.MAX_VALUE - _aidl_parcelable_size) {",

diff --git a/jadx-core/src/test/java/jadx/tests/integration/conditions/TestIfCodeStyle2.java b/jadx-core/src/test/java/jadx/tests/integration/conditions/TestIfCodeStyle2.java
@@ -0,0 +1,23 @@
+package jadx.tests.integration.conditions;
+
+import org.junit.jupiter.api.Test;
+
+import jadx.tests.api.SmaliTest;
+
+import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;
+
+/**
+ * Issue #2052
+ */
+public class TestIfCodeStyle2 extends SmaliTest {
+
+	@Test
+	public void testSmali() {
+		disableCompilation();
+		assertThat(getClassNodeFromSmali())
+				.code()
+				.countString(1, "} else if (")
+				.countString(1, "} else {")
+				.countString(19, "return ");
+	}
+}