fix: improve checks for restore new filled array (#2289)

Author: skylot

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

@@ -434,7 +434,7 @@ private void makeInsnBody(ICodeWriter code, InsnNode insn, Set<Flags> state) thr
 					code.add(']');
 				}
 				int dim = arrayType.getArrayDimension();
-				for (; k < dim - 1; k++) {
+				for (; k < dim; k++) {
 					code.add("[]");
 				}
 				break;

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

@@ -14,6 +14,11 @@ public enum CodeFeature {
 		 * Code contains switch instruction
 		 */
 		SWITCH,
+
+		/**
+		 * Code contains new-array instruction
+		 */
+		NEW_ARRAY,
 	}
 
 	public static boolean contains(MethodNode mth, CodeFeature feature) {

jadx-core/src/main/java/jadx/core/dex/instructions/InsnDecoder.java

@@ -540,6 +540,7 @@ private InsnNode makeNewArray(InsnData insn) {
 		for (int i = 1; i < regsCount; i++) {
 			newArr.addArg(InsnArg.typeImmutableReg(insn, i, ArgType.INT));
 		}
+		CodeFeaturesAttr.add(method, CodeFeature.NEW_ARRAY);
 		return newArr;
 	}
 

jadx-core/src/main/java/jadx/core/dex/instructions/args/InsnArg.java

@@ -13,6 +13,7 @@
 import jadx.core.dex.nodes.InsnNode;
 import jadx.core.dex.nodes.MethodNode;
 import jadx.core.utils.InsnRemover;
+import jadx.core.utils.InsnUtils;
 import jadx.core.utils.exceptions.JadxRuntimeException;
 
 /**
@@ -301,6 +302,10 @@ public boolean isSameCodeVar(RegisterArg arg) {
 		return false;
 	}
 
+	public boolean isUseVar(RegisterArg arg) {
+		return InsnUtils.containsVar(this, arg);
+	}
+
 	protected final <T extends InsnArg> T copyCommonParams(T copy) {
 		copy.copyAttributesFrom(this);
 		copy.setParentInsn(parentInsn);

jadx-core/src/main/java/jadx/core/dex/visitors/ReplaceNewArray.java

@@ -1,12 +1,17 @@
 package jadx.core.dex.visitors;
 
-import java.util.HashSet;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.SortedMap;
 import java.util.TreeMap;
 
 import jadx.core.dex.attributes.AFlag;
+import jadx.core.dex.attributes.nodes.CodeFeaturesAttr;
+import jadx.core.dex.attributes.nodes.CodeFeaturesAttr.CodeFeature;
 import jadx.core.dex.instructions.FilledNewArrayNode;
 import jadx.core.dex.instructions.IndexInsnNode;
 import jadx.core.dex.instructions.InsnType;
@@ -35,21 +40,20 @@ public class ReplaceNewArray extends AbstractVisitor {
 
 	@Override
 	public void visit(MethodNode mth) throws JadxException {
-		if (mth.isNoCode()) {
+		if (!CodeFeaturesAttr.contains(mth, CodeFeature.NEW_ARRAY)) {
 			return;
 		}
+		InsnRemover remover = new InsnRemover(mth);
 		int k = 0;
 		while (true) {
 			boolean changed = false;
-			InsnRemover remover = new InsnRemover(mth);
 			for (BlockNode block : mth.getBasicBlocks()) {
-				remover.setBlock(block);
-				List<InsnNode> instructions = block.getInstructions();
-				int size = instructions.size();
+				List<InsnNode> insnList = block.getInstructions();
+				int size = insnList.size();
 				for (int i = 0; i < size; i++) {
-					changed |= processInsn(mth, instructions, i, remover);
+					changed |= processInsn(mth, insnList, i, remover);
 				}
-				remover.perform();
+				remover.performForBlock(block);
 			}
 			if (changed) {
 				CodeShrinkVisitor.shrinkMethod(mth);
@@ -107,12 +111,11 @@ private static boolean processNewArray(MethodNode mth,
 		SortedMap<Long, InsnNode> arrPuts = new TreeMap<>();
 		for (RegisterArg registerArg : useList) {
 			InsnNode parentInsn = registerArg.getParentInsn();
-			if (parentInsn == null || parentInsn.getType() != InsnType.APUT) {
+			if (parentInsn == null
+					|| parentInsn.getType() != InsnType.APUT
+					|| !arrArg.sameRegAndSVar(parentInsn.getArg(0))) {
 				continue;
 			}
-			if (!arrArg.sameRegAndSVar(parentInsn.getArg(0))) {
-				return false;
-			}
 			Object constVal = InsnUtils.getConstValueByArg(mth.root(), parentInsn.getArg(1));
 			if (!(constVal instanceof LiteralArg)) {
 				return false;
@@ -130,8 +133,7 @@ private static boolean processNewArray(MethodNode mth,
 		if (arrPuts.size() < minLen) {
 			return false;
 		}
-		// expect all puts to be in same block
-		if (!new HashSet<>(instructions).containsAll(arrPuts.values())) {
+		if (!verifyPutInsns(arrArg, instructions, arrPuts)) {
 			return false;
 		}
 
@@ -165,6 +167,28 @@ private static boolean processNewArray(MethodNode mth,
 		return true;
 	}
 
+	private static boolean verifyPutInsns(RegisterArg arrReg, List<InsnNode> insnList, SortedMap<Long, InsnNode> arrPuts) {
+		List<InsnNode> puts = new ArrayList<>(arrPuts.values());
+		int putsCount = puts.size();
+		// expect all puts to be in the same block
+		if (insnList.size() < putsCount) {
+			return false;
+		}
+		Set<InsnNode> insnSet = Collections.newSetFromMap(new IdentityHashMap<>());
+		insnSet.addAll(insnList);
+		if (!insnSet.containsAll(puts)) {
+			return false;
+		}
+		// array arg shouldn't be used in puts insns
+		for (InsnNode put : puts) {
+			InsnArg putArg = put.getArg(2);
+			if (putArg.isUseVar(arrReg)) {
+				return false;
+			}
+		}
+		return true;
+	}
+
 	private static InsnArg replaceConstInArg(MethodNode mth, InsnArg valueArg) {
 		if (valueArg.isLiteral()) {
 			IFieldInfoRef f = mth.getParentClass().getConstFieldByLiteralArg((LiteralArg) valueArg);

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

@@ -3,6 +3,7 @@
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 import java.util.stream.Collectors;
 
 import org.jetbrains.annotations.Nullable;
@@ -75,6 +76,16 @@ public void perform() {
 		toRemove.clear();
 	}
 
+	public void performForBlock(BlockNode block) {
+		if (toRemove.isEmpty()) {
+			return;
+		}
+		instrList = Objects.requireNonNull(block.getInstructions());
+		unbindInsns(mth, toRemove);
+		removeAll(instrList, toRemove);
+		toRemove.clear();
+	}
+
 	public static void unbindInsn(@Nullable MethodNode mth, InsnNode insn) {
 		unbindAllArgs(mth, insn);
 		unbindResult(mth, insn);

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

@@ -263,6 +263,36 @@ public static <T extends InsnArg> boolean containsVar(List<T> list, RegisterArg
 		return false;
 	}
 
+	public static boolean containsVar(InsnNode insn, RegisterArg arg) {
+		if (insn == null) {
+			return false;
+		}
+		RegisterArg result = insn.getResult();
+		if (result != null && result.sameRegAndSVar(arg)) {
+			return true;
+		}
+		if (insn.getArgsCount() == 0) {
+			return false;
+		}
+		for (InsnArg insnArg : insn.getArguments()) {
+			if (containsVar(insnArg, arg)) {
+				return true;
+			}
+		}
+		return false;
+	}
+
+	public static boolean containsVar(InsnArg insnArg, RegisterArg arg) {
+		if (insnArg.isRegister()) {
+			return ((RegisterArg) insnArg).sameRegAndSVar(arg);
+		}
+		if (insnArg.isInsnWrap()) {
+			InsnNode wrapInsn = ((InsnWrapArg) insnArg).getWrapInsn();
+			return containsVar(wrapInsn, arg);
+		}
+		return false;
+	}
+
 	public static boolean contains(InsnNode insn, AFlag flag) {
 		return insn != null && insn.contains(flag);
 	}

jadx-core/src/test/java/jadx/tests/integration/arrays/TestArrayFillNegative.java

@@ -0,0 +1,33 @@
+package jadx.tests.integration.arrays;
+
+import org.junit.jupiter.api.Test;
+
+import jadx.tests.api.IntegrationTest;
+
+import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;
+
+public class TestArrayFillNegative extends IntegrationTest {
+
+	public static class TestCls {
+		public int[] test() {
+			int[] arr = new int[3];
+			arr[0] = 1;
+			arr[1] = arr[0] + 1;
+			arr[2] = arr[1] + 1;
+			return arr;
+		}
+
+		public void check() {
+			assertThat(test()).isEqualTo(new int[] { 1, 2, 3 });
+		}
+	}
+
+	@Test
+	public void test() {
+		assertThat(getClassNode(TestCls.class))
+				.code()
+				.doesNotContain("int[] arr = {1, ")
+				.containsOne("int[] arr = new int[3];")
+				.containsOne("arr[1] = arr[0] + 1;");
+	}
+}

jadx-core/src/test/java/jadx/tests/integration/arrays/TestMultiDimArrayFill.java

@@ -11,13 +11,7 @@ public static class TestCls {
 
 		public static Obj test(int a, int b) {
 			return new Obj(
-					new int[][] {
-							new int[] { 1 },
-							new int[] { 2 },
-							{ 3 },
-							new int[] { 4, 5 },
-							new int[0]
-					},
+					new int[][] { { 1 }, { 2 }, { 3 }, { 4, 5 }, new int[0] },
 					new int[] { a, a, a, a, b });
 		}