fix: add merged condition blocks for loop region (#1307)

jadx-core/src/main/java/jadx/core/codegen/ConditionGen.java

@@ -42,7 +42,7 @@ public class ConditionGen extends InsnGen {
 		super(insnGen.mgen, insnGen.fallback);
 	}
 
-	void add(ICodeWriter code, IfCondition condition) throws CodegenException {
+	public void add(ICodeWriter code, IfCondition condition) throws CodegenException {
 		add(code, new CondStack(), condition);
 	}
 

jadx-core/src/main/java/jadx/core/codegen/RegionGen.java

@@ -161,7 +161,7 @@ public class RegionGen extends InsnGen {
 	}
 
 	public void makeLoop(LoopRegion region, ICodeWriter code) throws CodegenException {
-		code.startLineWithNum(region.getConditionSourceLine());
+		code.startLineWithNum(region.getSourceLine());
 		LoopLabelAttr labelAttr = region.getInfo().getStart().get(AType.LOOP_LABEL);
 		if (labelAttr != null) {
 			code.add(mgen.getNameGen().getLoopLabel(labelAttr)).add(": ");
@@ -213,7 +213,7 @@ public class RegionGen extends InsnGen {
 			code.add("do {");
 			CodeGenUtils.addCodeComments(code, mth, condInsn);
 			makeRegionIndent(code, region.getBody());
-			code.startLineWithNum(region.getConditionSourceLine());
+			code.startLineWithNum(region.getSourceLine());
 			code.add("} while (");
 			conditionGen.add(code, condition);
 			code.add(");");

jadx-core/src/main/java/jadx/core/dex/attributes/nodes/JadxCommentsAttr.java

@@ -11,6 +11,7 @@ import jadx.api.CommentsLevel;
 import jadx.api.plugins.input.data.attributes.IJadxAttrType;
 import jadx.api.plugins.input.data.attributes.IJadxAttribute;
 import jadx.core.dex.attributes.AType;
+import jadx.core.utils.Utils;
 
 public class JadxCommentsAttr implements IJadxAttribute {
 
@@ -44,4 +45,12 @@ public class JadxCommentsAttr implements IJadxAttribute {
 	public IJadxAttrType<JadxCommentsAttr> getAttrType() {
 		return AType.JADX_COMMENTS;
 	}
+
+	@Override
+	public String toString() {
+		return "JadxCommentsAttr{\n "
+				+ Utils.listToString(comments.entrySet(), "\n ",
+						e -> e.getKey() + ": \n -> " + Utils.listToString(e.getValue(), "\n -> "))
+				+ '}';
+	}
 }

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

+package jadx.core.dex.nodes;
+
+import java.util.List;
+
+import org.jetbrains.annotations.Nullable;
+
+import jadx.core.dex.regions.conditions.IfCondition;
+
+public interface IConditionRegion extends IRegion {
+
+	@Nullable
+	IfCondition getCondition();
+
+	/**
+	 * Blocks merged into condition
+	 * Needed for backtracking
+	 * TODO: merge into condition object ???
+	 */
+	List<BlockNode> getConditionBlocks();
+
+	void invertCondition();
+
+	boolean simplifyCondition();
+
+	int getConditionSourceLine();
+}

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

@@ -120,12 +120,7 @@ public class InsnNode extends LineAttrNode {
 		if (getArgsCount() == 0) {
 			return false;
 		}
-		for (InsnArg insnArg : arguments) {
-			if (insnArg == arg || arg.sameRegAndSVar(insnArg)) {
-				return true;
-			}
-		}
-		return false;
+		return InsnUtils.containsVar(arguments, arg);
 	}
 
 	/**

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

+package jadx.core.dex.regions.conditions;
+
+import java.util.Collections;
+import java.util.List;
+
+import org.jetbrains.annotations.Nullable;
+
+import jadx.core.dex.nodes.BlockNode;
+import jadx.core.dex.nodes.IConditionRegion;
+import jadx.core.dex.nodes.IRegion;
+import jadx.core.dex.nodes.InsnNode;
+import jadx.core.dex.regions.AbstractRegion;
+import jadx.core.utils.BlockUtils;
+
+public abstract class ConditionRegion extends AbstractRegion implements IConditionRegion {
+
+	@Nullable
+	private IfCondition condition;
+	private List<BlockNode> conditionBlocks = Collections.emptyList();
+
+	public ConditionRegion(IRegion parent) {
+		super(parent);
+	}
+
+	@Override
+	@Nullable
+	public IfCondition getCondition() {
+		return condition;
+	}
+
+	@Override
+	public List<BlockNode> getConditionBlocks() {
+		return conditionBlocks;
+	}
+
+	@Override
+	public void invertCondition() {
+		if (condition != null) {
+			condition = IfCondition.invert(condition);
+		}
+	}
+
+	@Override
+	public boolean simplifyCondition() {
+		if (condition == null) {
+			return false;
+		}
+		IfCondition updated = IfCondition.simplify(condition);
+		if (updated != condition) {
+			condition = updated;
+			return true;
+		}
+		return false;
+	}
+
+	@Override
+	public int getConditionSourceLine() {
+		for (BlockNode block : conditionBlocks) {
+			InsnNode lastInsn = BlockUtils.getLastInsn(block);
+			if (lastInsn != null) {
+				int sourceLine = lastInsn.getSourceLine();
+				if (sourceLine != 0) {
+					return sourceLine;
+				}
+			}
+		}
+		return 0;
+	}
+
+	/**
+	 * Prefer way for update condition info
+	 */
+	public void updateCondition(IfInfo info) {
+		this.condition = info.getCondition();
+		this.conditionBlocks = info.getMergedBlocks();
+	}
+
+	public void updateCondition(IfCondition condition, List<BlockNode> conditionBlocks) {
+		this.condition = condition;
+		this.conditionBlocks = conditionBlocks;
+	}
+
+	public void updateCondition(BlockNode block) {
+		this.condition = IfCondition.fromIfBlock(block);
+		this.conditionBlocks = Collections.singletonList(block);
+	}
+}

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

@@ -272,6 +272,12 @@ public final class IfCondition extends AttrNode {
 		}
 	}
 
+	public List<InsnNode> collectInsns() {
+		List<InsnNode> list = new ArrayList<>();
+		visitInsns(list::add);
+		return list;
+	}
+
 	@Nullable
 	public InsnNode getFirstInsn() {
 		if (mode == Mode.COMPARE) {

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

@@ -3,7 +3,6 @@ package jadx.core.dex.regions.conditions;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.Set;
 
 import jadx.api.ICodeWriter;
 import jadx.core.codegen.RegionGen;
@@ -11,16 +10,9 @@ import jadx.core.dex.nodes.BlockNode;
 import jadx.core.dex.nodes.IBranchRegion;
 import jadx.core.dex.nodes.IContainer;
 import jadx.core.dex.nodes.IRegion;
-import jadx.core.dex.nodes.InsnNode;
-import jadx.core.dex.regions.AbstractRegion;
-import jadx.core.utils.BlockUtils;
 import jadx.core.utils.exceptions.CodegenException;
 
-public final class IfRegion extends AbstractRegion implements IBranchRegion {
-
-	private List<BlockNode> conditionBlocks;
-
-	private IfCondition condition;
+public final class IfRegion extends ConditionRegion implements IBranchRegion {
 	private IContainer thenRegion;
 	private IContainer elseRegion;
 
@@ -28,14 +20,6 @@ public final class IfRegion extends AbstractRegion implements IBranchRegion {
 		super(parent);
 	}
 
-	public IfCondition getCondition() {
-		return condition;
-	}
-
-	public void setCondition(IfCondition condition) {
-		this.condition = condition;
-	}
-
 	public IContainer getThenRegion() {
 		return thenRegion;
 	}
@@ -52,31 +36,8 @@ public final class IfRegion extends AbstractRegion implements IBranchRegion {
 		this.elseRegion = elseRegion;
 	}
 
-	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() {
-		IfCondition cond = IfCondition.simplify(condition);
-		if (cond != condition) {
-			condition = cond;
-			return true;
-		}
-		return false;
-	}
-
 	public void invert() {
-		condition = IfCondition.invert(condition);
+		invertCondition();
 		// swap regions
 		IContainer tmp = thenRegion;
 		thenRegion = elseRegion;
@@ -84,20 +45,12 @@ public final class IfRegion extends AbstractRegion implements IBranchRegion {
 	}
 
 	public int getSourceLine() {
-		for (BlockNode block : conditionBlocks) {
-			InsnNode lastInsn = BlockUtils.getLastInsn(block);
-			if (lastInsn != null) {
-				int sourceLine = lastInsn.getSourceLine();
-				if (sourceLine != 0) {
-					return sourceLine;
-				}
-			}
-		}
-		return 0;
+		return getConditionSourceLine();
 	}
 
 	@Override
 	public List<IContainer> getSubBlocks() {
+		List<BlockNode> conditionBlocks = getConditionBlocks();
 		List<IContainer> all = new ArrayList<>(conditionBlocks.size() + 2);
 		all.addAll(conditionBlocks);
 		if (thenRegion != null) {
@@ -151,6 +104,6 @@ public final class IfRegion extends AbstractRegion implements IBranchRegion {
 
 	@Override
 	public String toString() {
-		return "IF " + conditionBlocks + " THEN: " + thenRegion + " ELSE: " + elseRegion;
+		return "IF " + getConditionBlocks() + " THEN: " + thenRegion + " ELSE: " + elseRegion;
 	}
 }

jadx-core/src/main/java/jadx/core/dex/regions/loops/LoopRegion.java

 package jadx.core.dex.regions.loops;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 
 import org.jetbrains.annotations.Nullable;
@@ -9,55 +8,45 @@ import org.jetbrains.annotations.Nullable;
 import jadx.api.ICodeWriter;
 import jadx.core.codegen.RegionGen;
 import jadx.core.dex.attributes.nodes.LoopInfo;
-import jadx.core.dex.instructions.IfNode;
 import jadx.core.dex.instructions.args.RegisterArg;
 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.regions.AbstractRegion;
+import jadx.core.dex.regions.conditions.ConditionRegion;
 import jadx.core.dex.regions.conditions.IfCondition;
 import jadx.core.utils.BlockUtils;
+import jadx.core.utils.InsnUtils;
 import jadx.core.utils.exceptions.CodegenException;
 
-public final class LoopRegion extends AbstractRegion {
+public final class LoopRegion extends ConditionRegion {
 
 	private final LoopInfo info;
-	/**
-	 * loop header contains one 'if' insn, equals null for infinite loop
-	 */
-	@Nullable
-	private IfCondition condition;
-	private final BlockNode conditionBlock;
-	// instruction which must be executed before condition in every loop
-	private BlockNode preCondition;
-	private IRegion body;
 	private final boolean conditionAtEnd;
+	private final @Nullable BlockNode header;
+	// instruction which must be executed before condition in every loop
+	private @Nullable BlockNode preCondition;
 
+	private IRegion body;
 	private LoopType type;
 
 	public LoopRegion(IRegion parent, LoopInfo info, @Nullable BlockNode header, boolean reversed) {
 		super(parent);
 		this.info = info;
-		this.conditionBlock = header;
-		this.condition = IfCondition.fromIfBlock(header);
+		this.header = header;
 		this.conditionAtEnd = reversed;
+		if (header != null) {
+			updateCondition(header);
+		}
 	}
 
 	public LoopInfo getInfo() {
 		return info;
 	}
 
-	public IfCondition getCondition() {
-		return condition;
-	}
-
-	public void setCondition(IfCondition condition) {
-		this.condition = condition;
-	}
-
+	@Nullable
 	public BlockNode getHeader() {
-		return conditionBlock;
+		return header;
 	}
 
 	public IRegion getBody() {
@@ -79,10 +68,6 @@ public final class LoopRegion extends AbstractRegion {
 		this.preCondition = preCondition;
 	}
 
-	private IfNode getIfInsn() {
-		return (IfNode) BlockUtils.getLastInsn(conditionBlock);
-	}
-
 	/**
 	 * Check if pre-conditions can be inlined into loop condition
 	 */
@@ -91,7 +76,14 @@ public final class LoopRegion extends AbstractRegion {
 		if (insns.isEmpty()) {
 			return true;
 		}
-		IfNode ifInsn = getIfInsn();
+		IfCondition condition = getCondition();
+		if (condition == null) {
+			return false;
+		}
+		List<RegisterArg> conditionArgs = condition.getRegisterArgs();
+		if (conditionArgs.isEmpty()) {
+			return false;
+		}
 		int size = insns.size();
 		for (int i = 0; i < size; i++) {
 			InsnNode insn = insns.get(i);
@@ -110,7 +102,7 @@ public final class LoopRegion extends AbstractRegion {
 				}
 			}
 			// or in if insn
-			if (!found && ifInsn.containsVar(res)) {
+			if (!found && InsnUtils.containsVar(conditionArgs, res)) {
 				found = true;
 			}
 			if (!found) {
@@ -124,8 +116,8 @@ public final class LoopRegion extends AbstractRegion {
 	 * Move all preCondition block instructions before conditionBlock instructions
 	 */
 	public void mergePreCondition() {
-		if (preCondition != null && conditionBlock != null) {
-			List<InsnNode> condInsns = conditionBlock.getInstructions();
+		if (preCondition != null && header != null) {
+			List<InsnNode> condInsns = header.getInstructions();
 			List<InsnNode> preCondInsns = preCondition.getInstructions();
 			preCondInsns.addAll(condInsns);
 			condInsns.clear();
@@ -135,9 +127,13 @@ public final class LoopRegion extends AbstractRegion {
 		}
 	}
 
-	public int getConditionSourceLine() {
-		InsnNode lastInsn = BlockUtils.getLastInsn(conditionBlock);
-		return lastInsn == null ? 0 : lastInsn.getSourceLine();
+	public int getSourceLine() {
+		InsnNode lastInsn = BlockUtils.getLastInsn(header);
+		int headerLine = lastInsn == null ? 0 : lastInsn.getSourceLine();
+		if (headerLine != 0) {
+			return headerLine;
+		}
+		return getConditionSourceLine();
 	}
 
 	public LoopType getType() {
@@ -150,17 +146,15 @@ public final class LoopRegion extends AbstractRegion {
 
 	@Override
 	public List<IContainer> getSubBlocks() {
-		List<IContainer> all = new ArrayList<>(3);
+		List<IContainer> all = new ArrayList<>(2 + getConditionBlocks().size());
 		if (preCondition != null) {
 			all.add(preCondition);
 		}
-		if (conditionBlock != null) {
-			all.add(conditionBlock);
-		}
+		all.addAll(getConditionBlocks());
 		if (body != null) {
 			all.add(body);
 		}
-		return Collections.unmodifiableList(all);
+		return all;
 	}
 
 	@Override

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

@@ -216,7 +216,7 @@ public class RegionMaker {
 			// invert loop condition if 'then' points to exit
 			condInfo = IfInfo.invert(condInfo);
 		}
-		loopRegion.setCondition(condInfo.getCondition());
+		loopRegion.updateCondition(condInfo);
 		exitBlocks.removeAll(condInfo.getMergedBlocks());
 
 		if (!exitBlocks.isEmpty()) {
@@ -720,8 +720,7 @@ public class RegionMaker {
 		confirmMerge(currentIf);
 
 		IfRegion ifRegion = new IfRegion(currentRegion);
-		ifRegion.setCondition(currentIf.getCondition());
-		ifRegion.setConditionBlocks(currentIf.getMergedBlocks());
+		ifRegion.updateCondition(currentIf);
 		currentRegion.getSubBlocks().add(ifRegion);
 
 		BlockNode outBlock = currentIf.getOutBlock();

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

@@ -18,6 +18,7 @@ import org.slf4j.LoggerFactory;
 
 import jadx.api.ICodeWriter;
 import jadx.api.impl.SimpleCodeWriter;
+import jadx.core.codegen.ConditionGen;
 import jadx.core.codegen.InsnGen;
 import jadx.core.codegen.MethodGen;
 import jadx.core.dex.attributes.AType;
@@ -31,6 +32,8 @@ import jadx.core.dex.nodes.InsnNode;
 import jadx.core.dex.nodes.MethodNode;
 import jadx.core.dex.nodes.RootNode;
 import jadx.core.dex.regions.Region;
+import jadx.core.dex.regions.conditions.IfCondition;
+import jadx.core.dex.regions.loops.LoopRegion;
 import jadx.core.dex.visitors.AbstractVisitor;
 import jadx.core.dex.visitors.DotGraphVisitor;
 import jadx.core.dex.visitors.IDexTreeVisitor;
@@ -139,6 +142,7 @@ public class DebugUtils {
 	private static void printRegion(MethodNode mth, IRegion region, ICodeWriter cw, String indent, boolean printInsns) {
 		printWithAttributes(cw, indent, region.toString(), region);
 		indent += "|  ";
+		printRegionSpecificInfo(cw, indent, mth, region, printInsns);
 		for (IContainer container : region.getSubBlocks()) {
 			if (container instanceof IRegion) {
 				printRegion(mth, (IRegion) container, cw, indent, printInsns);
@@ -152,6 +156,23 @@ public class DebugUtils {
 		}
 	}
 
+	private static void printRegionSpecificInfo(ICodeWriter cw, String indent,
+			MethodNode mth, IRegion region, boolean printInsns) {
+		if (region instanceof LoopRegion) {
+			LoopRegion loop = (LoopRegion) region;
+			IfCondition condition = loop.getCondition();
+			if (printInsns && condition != null) {
+				ConditionGen conditionGen = new ConditionGen(new InsnGen(MethodGen.getFallbackMethodGen(mth), true));
+				cw.startLine(indent).add("|> ");
+				try {
+					conditionGen.add(cw, condition);
+				} catch (Exception e) {
+					cw.startLine(indent).add(">!! ").add(condition.toString());
+				}
+			}
+		}
+	}
+
 	private static void printInsns(MethodNode mth, ICodeWriter cw, String indent, IBlock block) {
 		for (InsnNode insn : block.getInstructions()) {
 			try {

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

@@ -206,4 +206,16 @@ public class InsnUtils {
 		insn.add(AFlag.DONT_GENERATE);
 		return true;
 	}
+
+	public static <T extends InsnArg> boolean containsVar(List<T> list, RegisterArg arg) {
+		if (list == null || list.isEmpty()) {
+			return false;
+		}
+		for (InsnArg insnArg : list) {
+			if (insnArg == arg || arg.sameRegAndSVar(insnArg)) {
+				return true;
+			}
+		}
+		return false;
+	}
 }

jadx-core/src/test/java/jadx/tests/integration/variables/TestVariablesInInlinedAssign.java

+package jadx.tests.integration.variables;
+
+import jadx.tests.api.IntegrationTest;
+import jadx.tests.api.extensions.inputs.InputPlugin;
+import jadx.tests.api.extensions.inputs.TestWithInputPlugins;
+
+import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;
+
+public class TestVariablesInInlinedAssign extends IntegrationTest {
+
+	public static class TestCls {
+		public final int test(final char[] s) {
+			int i;
+			for (i = 0; i < s.length; i++) {
+				final char c = s[i];
+				if (c != 'a' && c != 'b') {
+					break;
+				}
+			}
+			return i;
+		}
+	}
+
+	@TestWithInputPlugins({ InputPlugin.DEX, InputPlugin.JAVA })
+	public void test() {
+		assertThat(getClassNode(TestCls.class))
+				.code()
+				.containsOne("char c");
+	}
+}