fix: store condition blocks in 'if' region for correct blocks list (#669)

jadx-core/src/main/java/jadx/core/dex/nodes/BlockNode.java

@@ -4,6 +4,8 @@ import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.List;
 
+import org.jetbrains.annotations.NotNull;
+
 import jadx.core.dex.attributes.AFlag;
 import jadx.core.dex.attributes.AType;
 import jadx.core.dex.attributes.AttrNode;
@@ -16,7 +18,7 @@ import jadx.core.utils.exceptions.JadxRuntimeException;
 
 import static jadx.core.utils.Utils.lockList;
 
-public class BlockNode extends AttrNode implements IBlock {
+public final class BlockNode extends AttrNode implements IBlock, Comparable<BlockNode> {
 
 	private int id;
 	private final int startOffset;
@@ -186,6 +188,11 @@ public class BlockNode extends AttrNode implements IBlock {
 	}
 
 	@Override
+	public int compareTo(@NotNull BlockNode o) {
+		return Integer.compare(id, o.id);
+	}
+
+	@Override
 	public String baseString() {
 		return Integer.toString(id);
 	}

jadx-core/src/main/java/jadx/core/dex/nodes/InsnContainer.java

 package jadx.core.dex.nodes;
 
+import java.util.Collections;
 import java.util.List;
 
 import jadx.core.dex.attributes.AttrNode;
 
-public class InsnContainer extends AttrNode implements IBlock {
+/**
+ * Lightweight replacement for BlockNode in regions.
+ * Use with caution! Some passes still expect BlockNode in method blocks list (mth.getBlockNodes())
+ */
+public final class InsnContainer extends AttrNode implements IBlock {
 
 	private final List<InsnNode> insns;
 
+	public InsnContainer(InsnNode insn) {
+		this.insns = Collections.singletonList(insn);
+	}
+
 	public InsnContainer(List<InsnNode> insns) {
 		this.insns = insns;
 	}

jadx-core/src/main/java/jadx/core/dex/regions/conditions/IfCondition.java

@@ -8,9 +8,7 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.Objects;
 
-import jadx.core.dex.attributes.AFlag;
-import jadx.core.dex.attributes.AttributeStorage;
-import jadx.core.dex.attributes.EmptyAttrStorage;
+import jadx.core.dex.attributes.AttrNode;
 import jadx.core.dex.instructions.ArithNode;
 import jadx.core.dex.instructions.ArithOp;
 import jadx.core.dex.instructions.IfNode;
@@ -24,7 +22,7 @@ import jadx.core.dex.nodes.InsnNode;
 import jadx.core.utils.BlockUtils;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
-public final class IfCondition {
+public final class IfCondition extends AttrNode {
 
 	public enum Mode {
 		COMPARE,
@@ -34,10 +32,6 @@ public final class IfCondition {
 		OR
 	}
 
-	private static final AttributeStorage EMPTY_ATTR_STORAGE = new EmptyAttrStorage();
-
-	private AttributeStorage storage = EMPTY_ATTR_STORAGE;
-
 	private final Mode mode;
 	private final List<IfCondition> args;
 	private final Compare compare;
@@ -268,23 +262,6 @@ public final class IfCondition {
 		return list;
 	}
 
-	public void add(AFlag flag) {
-		initStorage().add(flag);
-	}
-
-	public boolean contains(AFlag flag) {
-		return storage.contains(flag);
-	}
-
-	private AttributeStorage initStorage() {
-		AttributeStorage store = storage;
-		if (store == EMPTY_ATTR_STORAGE) {
-			store = new AttributeStorage();
-			storage = store;
-		}
-		return store;
-	}
-
 	@Override
 	public String toString() {
 		switch (mode) {

jadx-core/src/main/java/jadx/core/dex/regions/conditions/IfRegion.java

@@ -3,6 +3,7 @@ package jadx.core.dex.regions.conditions;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.IBranchRegion;
@@ -14,16 +15,14 @@ import jadx.core.utils.BlockUtils;
 
 public final class IfRegion extends AbstractRegion implements IBranchRegion {
 
-	private final BlockNode header;
+	private List<BlockNode> conditionBlocks;
 
 	private IfCondition condition;
 	private IContainer thenRegion;
 	private IContainer elseRegion;
 
-	public IfRegion(IRegion parent, BlockNode header) {
+	public IfRegion(IRegion parent) {
 		super(parent);
-		this.header = header;
-		this.condition = IfCondition.fromIfBlock(header);
 	}
 
 	public IfCondition getCondition() {
@@ -50,8 +49,18 @@ public final class IfRegion extends AbstractRegion implements IBranchRegion {
 		this.elseRegion = elseRegion;
 	}
 
-	public BlockNode getHeader() {
-		return header;
+	public List<BlockNode> getConditionBlocks() {
+		return conditionBlocks;
+	}
+
+	public void setConditionBlocks(List<BlockNode> conditionBlocks) {
+		this.conditionBlocks = conditionBlocks;
+	}
+
+	public void setConditionBlocks(Set<BlockNode> conditionBlocks) {
+		List<BlockNode> list = new ArrayList<>(conditionBlocks);
+		Collections.sort(list);
+		this.conditionBlocks = list;
 	}
 
 	public boolean simplifyCondition() {
@@ -72,14 +81,22 @@ public final class IfRegion extends AbstractRegion implements IBranchRegion {
 	}
 
 	public int getSourceLine() {
-		InsnNode lastInsn = BlockUtils.getLastInsn(header);
-		return lastInsn == null ? 0 : lastInsn.getSourceLine();
+		for (BlockNode block : conditionBlocks) {
+			InsnNode lastInsn = BlockUtils.getLastInsn(block);
+			if (lastInsn != null) {
+				int sourceLine = lastInsn.getSourceLine();
+				if (sourceLine != 0) {
+					return sourceLine;
+				}
+			}
+		}
+		return 0;
 	}
 
 	@Override
 	public List<IContainer> getSubBlocks() {
-		List<IContainer> all = new ArrayList<>(3);
-		all.add(header);
+		List<IContainer> all = new ArrayList<>(conditionBlocks.size() + 2);
+		all.addAll(conditionBlocks);
 		if (thenRegion != null) {
 			all.add(thenRegion);
 		}
@@ -126,6 +143,6 @@ public final class IfRegion extends AbstractRegion implements IBranchRegion {
 
 	@Override
 	public String toString() {
-		return "IF " + header + " THEN:" + thenRegion + " ELSE:" + elseRegion;
+		return "IF " + conditionBlocks + " THEN: " + thenRegion + " ELSE: " + elseRegion;
 	}
 }

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

@@ -18,7 +18,7 @@ import static jadx.core.utils.RegionUtils.insnsCount;
 
 public class IfRegionVisitor extends AbstractVisitor {
 
-	private static final TernaryVisitor TERNARY_VISITOR = new TernaryVisitor();
+	private static final TernaryMod TERNARY_VISITOR = new TernaryMod();
 	private static final ProcessIfRegionVisitor PROCESS_IF_REGION_VISITOR = new ProcessIfRegionVisitor();
 	private static final RemoveRedundantElseVisitor REMOVE_REDUNDANT_ELSE_VISITOR = new RemoveRedundantElseVisitor();
 
@@ -29,17 +29,6 @@ public class IfRegionVisitor extends AbstractVisitor {
 		DepthRegionTraversal.traverseIterative(mth, REMOVE_REDUNDANT_ELSE_VISITOR);
 	}
 
-	/**
-	 * Collapse ternary operators
-	 */
-	private static class TernaryVisitor implements IRegionIterativeVisitor {
-		@Override
-		public boolean visitRegion(MethodNode mth, IRegion region) {
-			return region instanceof IfRegion
-					&& TernaryMod.makeTernaryInsn(mth, (IfRegion) region);
-		}
-	}
-
 	private static class ProcessIfRegionVisitor extends AbstractRegionVisitor {
 		@Override
 		public boolean enterRegion(MethodNode mth, IRegion region) {

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

@@ -678,8 +678,9 @@ public class RegionMaker {
 		}
 		confirmMerge(currentIf);
 
-		IfRegion ifRegion = new IfRegion(currentRegion, block);
+		IfRegion ifRegion = new IfRegion(currentRegion);
 		ifRegion.setCondition(currentIf.getCondition());
+		ifRegion.setConditionBlocks(currentIf.getMergedBlocks());
 		currentRegion.getSubBlocks().add(ifRegion);
 
 		BlockNode outBlock = currentIf.getOutBlock();

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

 package jadx.core.dex.visitors.regions;
 
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import jadx.core.dex.attributes.AFlag;
@@ -13,6 +14,7 @@ import jadx.core.dex.instructions.args.RegisterArg;
 import jadx.core.dex.instructions.mods.TernaryInsn;
 import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.IContainer;
+import jadx.core.dex.nodes.IRegion;
 import jadx.core.dex.nodes.InsnNode;
 import jadx.core.dex.nodes.MethodNode;
 import jadx.core.dex.regions.Region;
@@ -20,12 +22,20 @@ import jadx.core.dex.regions.conditions.IfRegion;
 import jadx.core.dex.visitors.shrink.CodeShrinkVisitor;
 import jadx.core.utils.InsnList;
 
-public class TernaryMod {
+/**
+ * Convert 'if' to ternary operation
+ */
+public class TernaryMod implements IRegionIterativeVisitor {
 
-	private TernaryMod() {
+	@Override
+	public boolean visitRegion(MethodNode mth, IRegion region) {
+		if (region instanceof IfRegion) {
+			return makeTernaryInsn(mth, (IfRegion) region);
+		}
+		return false;
 	}
 
-	static boolean makeTernaryInsn(MethodNode mth, IfRegion ifRegion) {
+	private static boolean makeTernaryInsn(MethodNode mth, IfRegion ifRegion) {
 		if (ifRegion.contains(AFlag.ELSE_IF_CHAIN)) {
 			return false;
 		}
@@ -39,7 +49,12 @@ public class TernaryMod {
 		if (tb == null || eb == null) {
 			return false;
 		}
-		BlockNode header = ifRegion.getHeader();
+		List<BlockNode> conditionBlocks = ifRegion.getConditionBlocks();
+		if (conditionBlocks.isEmpty()) {
+			return false;
+		}
+
+		BlockNode header = conditionBlocks.get(0);
 		InsnNode thenInsn = tb.getInstructions().get(0);
 		InsnNode elseInsn = eb.getInstructions().get(0);
 
@@ -88,6 +103,8 @@ public class TernaryMod {
 			header.getInstructions().clear();
 			header.getInstructions().add(ternInsn);
 
+			clearConditionBlocks(conditionBlocks, header);
+
 			// shrink method again
 			CodeShrinkVisitor.shrinkMethod(mth);
 			return true;
@@ -120,12 +137,23 @@ public class TernaryMod {
 			header.getInstructions().add(retInsn);
 			header.add(AFlag.RETURN);
 
+			clearConditionBlocks(conditionBlocks, header);
+
 			CodeShrinkVisitor.shrinkMethod(mth);
 			return true;
 		}
 		return false;
 	}
 
+	private static void clearConditionBlocks(List<BlockNode> conditionBlocks, BlockNode header) {
+		for (BlockNode block : conditionBlocks) {
+			if (block != header) {
+				block.getInstructions().clear();
+				block.add(AFlag.REMOVE);
+			}
+		}
+	}
+
 	private static BlockNode getTernaryInsnBlock(IContainer thenRegion) {
 		if (thenRegion instanceof Region) {
 			Region r = (Region) thenRegion;
@@ -181,12 +209,7 @@ public class TernaryMod {
 			}
 			int sourceLine = assignInsn.getSourceLine();
 			if (sourceLine != 0) {
-				Integer count = map.get(sourceLine);
-				if (count != null) {
-					map.put(sourceLine, count + 1);
-				} else {
-					map.put(sourceLine, 1);
-				}
+				map.merge(sourceLine, 1, Integer::sum);
 			}
 		}
 		for (Map.Entry<Integer, Integer> entry : map.entrySet()) {

jadx-core/src/test/java/jadx/tests/integration/conditions/TestConditions4.java

@@ -24,7 +24,7 @@ public class TestConditions4 extends IntegrationTest {
 		String code = cls.getCode().toString();
 
 		assertThat(code, containsString("num >= 59 && num <= 66"));
-		assertThat(code, containsString("return inRange ? num + 1 : num;"));
+		assertThat(code, containsString("? num + 1 : num;"));
 		assertThat(code, not(containsString("else")));
 	}
 }

jadx-core/src/test/java/jadx/tests/integration/inline/TestInlineInLoop.java

@@ -42,6 +42,6 @@ public class TestInlineInLoop extends IntegrationTest {
 		assertThat(code, containsOne("int c3 = b;"));
 		assertThat(code, containsOne("int b2 = b + 1;"));
 		assertThat(code, containsOne("b = c3"));
-		assertThat(code, containsOne("a++;"));
+		assertThat(code, containsOne("a++"));
 	}
 }