feat: use queue instead of recursion for type updates

2026-03-21 21:55:09 +00:00
parent 15ea9a56b9
commit b3db337abd
9 changed files with 560 additions and 195 deletions
@@ -0,0 +1,121 @@
+package jadx.core.dex.visitors.typeinference;
+
+import java.util.Iterator;
+import java.util.List;
+import java.util.function.Predicate;
+
+import org.jetbrains.annotations.Nullable;
+
+import jadx.core.dex.instructions.args.ArgType;
+import jadx.core.dex.instructions.args.InsnArg;
+
+import static jadx.core.dex.visitors.typeinference.TypeUpdateResult.CHANGED;
+import static jadx.core.dex.visitors.typeinference.TypeUpdateResult.REJECT;
+import static jadx.core.dex.visitors.typeinference.TypeUpdateResult.SAME;
+
+/**
+ * Type update callback to set same type for args from list.
+ */
+public class ArgsListUpdateCallback<T extends InsnArg> implements ITypeUpdateCallback {
+	private final TypeUpdate typeUpdate;
+	private final TypeUpdateInfo updateInfo;
+	private final Iterator<T> argsIterator;
+	private final ArgType candidateType;
+	private final boolean direct;
+
+	private @Nullable Predicate<T> argsFilter;
+	private @Nullable ITypeUpdateCallback finalResultCallback;
+	private boolean ignoreReject = false;
+
+	private boolean allSame = true;
+	private boolean firstQueue = false;
+
+	public ArgsListUpdateCallback(TypeUpdate typeUpdate, TypeUpdateInfo updateInfo,
+			List<T> args, ArgType candidateType, boolean direct) {
+		this.typeUpdate = typeUpdate;
+		this.updateInfo = updateInfo;
+		this.argsIterator = args.iterator();
+		this.candidateType = candidateType;
+		this.direct = direct;
+	}
+
+	@Override
+	public @Nullable TypeUpdateResult updateCallback(TypeUpdateResult result) {
+		while (true) {
+			if (!ignoreReject) {
+				if (result == REJECT) {
+					return finalResult(result);
+				}
+			}
+			if (result != SAME) {
+				allSame = false;
+			}
+			T next = getNextArg();
+			if (next == null) {
+				return finalResult(allSame ? SAME : CHANGED);
+			}
+			result = queueUpdate(next);
+			if (result == null) {
+				// keep this callback
+				return null;
+			}
+		}
+	}
+
+	private @Nullable TypeUpdateResult queueUpdate(T next) {
+		ITypeUpdateCallback cb;
+		if (firstQueue) {
+			cb = this;
+			firstQueue = false;
+		} else {
+			cb = null;
+		}
+		if (direct) {
+			return typeUpdate.queueDirectTypeUpdate(updateInfo, next, candidateType, cb);
+		}
+		return typeUpdate.queueTypeUpdate(updateInfo, next, candidateType, cb);
+	}
+
+	public @Nullable TypeUpdateResult runFirstQueue() {
+		firstQueue = true;
+		return updateCallback(SAME);
+	}
+
+	public void setFinalResultCallback(@Nullable ITypeUpdateCallback finalResultCallback) {
+		this.finalResultCallback = finalResultCallback;
+	}
+
+	public void setArgsFilter(@Nullable Predicate<T> argsFilter) {
+		this.argsFilter = argsFilter;
+	}
+
+	public void setIgnoreReject(boolean ignoreReject) {
+		this.ignoreReject = ignoreReject;
+	}
+
+	private @Nullable TypeUpdateResult finalResult(TypeUpdateResult result) {
+		if (finalResultCallback != null) {
+			return finalResultCallback.updateCallback(result);
+		}
+		return result;
+	}
+
+	private @Nullable T getNextArg() {
+		Iterator<T> it = argsIterator;
+		Predicate<T> filter = argsFilter;
+		while (true) {
+			if (!it.hasNext()) {
+				return null;
+			}
+			T next = it.next();
+			if (filter == null || filter.test(next)) {
+				return next;
+			}
+		}
+	}
+
+	@Override
+	public String toString() {
+		return "ArgsListUpdateCallback";
+	}
+}
@@ -1,6 +1,7 @@
 package jadx.core.dex.visitors.typeinference;

 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;

 import jadx.core.dex.instructions.args.ArgType;
 import jadx.core.dex.instructions.args.InsnArg;
@@ -16,5 +17,6 @@ public interface ITypeListener {
 	 * @param arg           apply suggested type for this arg
 	 * @param candidateType suggest new type
 	 */
+	@Nullable
 	TypeUpdateResult update(TypeUpdateInfo updateInfo, InsnNode insn, InsnArg arg, @NotNull ArgType candidateType);
 }
@@ -0,0 +1,19 @@
+package jadx.core.dex.visitors.typeinference;
+
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Callback to process and modify type update result
+ */
+@FunctionalInterface
+public interface ITypeUpdateCallback {
+
+	/**
+	 * Called on type update result being calculated
+	 *
+	 * @param result - type update result
+	 * @return modified result, can be null - will keep callback and wait for another result
+	 */
+	@Nullable
+	TypeUpdateResult updateCallback(TypeUpdateResult result);
+}
@@ -0,0 +1,144 @@
+package jadx.core.dex.visitors.typeinference;
+
+import java.util.Set;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+import org.jetbrains.annotations.Nullable;
+
+import jadx.core.dex.instructions.BaseInvokeNode;
+import jadx.core.dex.instructions.args.ArgType;
+import jadx.core.dex.instructions.args.InsnArg;
+import jadx.core.dex.instructions.args.RegisterArg;
+
+import static jadx.core.dex.visitors.typeinference.TypeUpdateResult.CHANGED;
+import static jadx.core.dex.visitors.typeinference.TypeUpdateResult.REJECT;
+import static jadx.core.dex.visitors.typeinference.TypeUpdateResult.SAME;
+
+public class InvokeUpdateCallback implements ITypeUpdateCallback {
+	private final TypeUpdate typeUpdate;
+	private final TypeUpdateInfo updateInfo;
+	private final BaseInvokeNode invoke;
+	private final int argsCount;
+	private final Set<ArgType> knownTypeVars;
+	private final Supplier<ArgType> getReturnType;
+	private final Function<Integer, ArgType> getArgType;
+
+	private boolean isAssign;
+	private boolean allSame = true;
+	private int currentArg = -1;
+
+	private InsnArg updateArg;
+	private ArgType updateType;
+	private boolean firstQueue = false;
+
+	public InvokeUpdateCallback(TypeUpdate typeUpdate, TypeUpdateInfo updateInfo, BaseInvokeNode invoke, int argsCount,
+			Set<ArgType> knownTypeVars, Supplier<ArgType> getReturnType, Function<Integer, ArgType> getArgType) {
+		this.typeUpdate = typeUpdate;
+		this.updateInfo = updateInfo;
+		this.invoke = invoke;
+		this.argsCount = argsCount;
+		this.knownTypeVars = knownTypeVars;
+		this.getReturnType = getReturnType;
+		this.getArgType = getArgType;
+	}
+
+	@Override
+	public @Nullable TypeUpdateResult updateCallback(TypeUpdateResult result) {
+		while (true) {
+			switch (result) {
+				case CHANGED:
+					allSame = false;
+					break;
+
+				case REJECT:
+					TypeCompareEnum compare = typeUpdate.getTypeCompare().compareTypes(updateType, updateArg.getType());
+					if (isAssign ? compare.isWider() : compare.isNarrow()) {
+						return REJECT;
+					}
+					break;
+			}
+			if (!getNextArg()) {
+				return allSame ? SAME : CHANGED;
+			}
+			ITypeUpdateCallback cb;
+			if (firstQueue) {
+				cb = this;
+				firstQueue = false;
+			} else {
+				cb = null;
+			}
+			result = typeUpdate.queueTypeUpdate(updateInfo, updateArg, updateType, cb);
+			if (result == null) {
+				return null;
+			}
+		}
+	}
+
+	public TypeUpdateResult runQueue() {
+		firstQueue = true;
+		TypeUpdateResult result = SAME;
+		RegisterArg resultArg = invoke.getResult();
+		if (resultArg != null && !resultArg.isTypeImmutable()) {
+			ArgType returnType = checkType(knownTypeVars, getReturnType.get());
+			if (returnType != null) {
+				updateArg = resultArg;
+				updateType = returnType;
+				isAssign = true;
+				firstQueue = false;
+				result = typeUpdate.queueTypeUpdate(updateInfo, updateArg, updateType, this);
+				if (result == null) {
+					return null;
+				}
+			}
+		}
+		return updateCallback(result);
+	}
+
+	private boolean getNextArg() {
+		while (true) {
+			currentArg++;
+			int i = currentArg;
+			if (i >= argsCount) {
+				return false;
+			}
+			int argOffset = invoke.getFirstArgOffset();
+			InsnArg invokeArg = invoke.getArg(argOffset + i);
+			if (!invokeArg.isTypeImmutable()) {
+				ArgType argType = checkType(knownTypeVars, getArgType.apply(i));
+				if (argType != null) {
+					updateArg = invokeArg;
+					updateType = argType;
+					isAssign = false;
+					return true;
+				}
+			}
+		}
+	}
+
+	private @Nullable ArgType checkType(Set<ArgType> knownTypeVars, @Nullable ArgType type) {
+		if (type == null) {
+			return null;
+		}
+		if (type.isWildcard()) {
+			return null;
+		}
+		if (type.containsTypeVariable()) {
+			if (knownTypeVars.isEmpty()) {
+				return null;
+			}
+			Boolean hasUnknown = type.visitTypes(this::isUnknown);
+			if (hasUnknown != null) {
+				return null;
+			}
+		}
+		return type;
+	}
+
+	private @Nullable Boolean isUnknown(ArgType t) {
+		if (t.isGenericType() && !knownTypeVars.contains(t)) {
+			return Boolean.TRUE;
+		}
+		return null;
+	}
+}
@@ -52,7 +52,7 @@ public final class TypeUpdate {
 	}

 	/**
-	 * Perform recursive type checking and type propagation for all related variables
+	 * Perform type checking and type propagation for all related variables
 	 */
 	public TypeUpdateResult apply(MethodNode mth, SSAVar ssaVar, ArgType candidateType) {
 		return apply(mth, ssaVar, candidateType, TypeUpdateFlags.FLAGS_EMPTY);
@@ -81,9 +81,14 @@ public final class TypeUpdate {
 			if (candidateType == null || !candidateType.isTypeKnown()) {
 				return REJECT;
 			}
-
+			if (Consts.DEBUG_TYPE_INFERENCE) {
+				LOG.debug("Start type update for {} to {}", ssaVar.toShortString(), candidateType);
+			}
 			TypeUpdateInfo updateInfo = new TypeUpdateInfo(mth, flags, args);
-			TypeUpdateResult result = updateTypeChecked(updateInfo, ssaVar.getAssign(), candidateType);
+			TypeUpdateResult result = queueTypeUpdate(updateInfo, ssaVar.getAssign(), candidateType, null);
+			if (result == null) {
+				result = runUpdate(updateInfo);
+			}
 			if (result == REJECT) {
 				return result;
 			}
@@ -103,16 +108,88 @@ public final class TypeUpdate {
 		}
 	}

-	private TypeUpdateResult updateTypeChecked(TypeUpdateInfo updateInfo, InsnArg arg, ArgType candidateType) {
-		if (candidateType == null) {
-			throw new JadxRuntimeException("Null type update for arg: " + arg);
+	private TypeUpdateResult runUpdate(TypeUpdateInfo updateInfo) {
+		TypeUpdateResult result = REJECT;
+		while (true) {
+			TypeUpdateRequest request = updateInfo.pollNextRequest();
+			if (request == null) {
+				return result;
+			}
+			InsnArg updateArg = request.getArg();
+			ArgType updateType = request.getCandidateType();
+			TypeUpdateResult newResult;
+			if (request.isDirect()) {
+				newResult = requestUpdate(updateInfo, updateArg, updateType);
+			} else {
+				newResult = updateTypeForArg(updateInfo, updateArg, updateType);
+			}
+			updateInfo.saveCallback(request);
+			if (newResult == null) {
+				// no result: continue
+			} else {
+				// propagate result back through callbacks
+				result = processCallbacks(updateInfo, newResult);
+			}
 		}
-		if (updateInfo.isProcessed(arg)) {
-			return CHANGED;
+	}
+
+	private static TypeUpdateResult processCallbacks(TypeUpdateInfo updateInfo, TypeUpdateResult result) {
+		TypeUpdateResult current = result;
+		while (true) {
+			TypeUpdateRequest cbReq = updateInfo.pollNextCallback();
+			if (cbReq == null) {
+				return current;
+			}
+			ITypeUpdateCallback callback = Objects.requireNonNull(cbReq.getCallback());
+			current = callback.updateCallback(current);
+			if (current == null) {
+				// no result, put callback back into queue
+				// so it can be executed once result is calculated
+				updateInfo.saveCallback(cbReq);
+				return null;
+			}
+			if (current == REJECT) {
+				updateInfo.rollbackUpdate(cbReq.getArg());
+			}
+			// proceed to next callback
 		}
+	}
+
+	/**
+	 * Queue type update for InsnArg.
+	 *
+	 * @param callback - will be executed when result for this update is calculated,
+	 *                 can be null - callback will pass result without change
+	 * @return null if update added into queue, non-null result if not queued (verify failed)
+	 */
+	public @Nullable TypeUpdateResult queueTypeUpdate(TypeUpdateInfo updateInfo,
+			InsnArg arg, ArgType candidateType, @Nullable ITypeUpdateCallback callback) {
+		// verify can be done in queue processing before request run, but kept here for faster processing
+		// this might increase code complexity because result should be checked for null every time
 		TypeUpdateResult res = verifyType(updateInfo, arg, candidateType);
 		if (res != null) {
-			return res;
+			if (callback == null) {
+				return res;
+			}
+			TypeUpdateResult result = callback.updateCallback(res);
+			if (result == null) {
+				updateInfo.saveCallback(new TypeUpdateRequest(arg, candidateType, false, callback));
+			}
+			return result;
+		}
+		updateInfo.queueRequest(new TypeUpdateRequest(arg, candidateType, false, callback));
+		return null;
+	}
+
+	public @Nullable TypeUpdateResult queueDirectTypeUpdate(TypeUpdateInfo updateInfo, InsnArg arg, ArgType candidateType,
+			@Nullable ITypeUpdateCallback callback) {
+		updateInfo.queueRequest(new TypeUpdateRequest(arg, candidateType, true, callback));
+		return null;
+	}
+
+	private TypeUpdateResult updateTypeForArg(TypeUpdateInfo updateInfo, InsnArg arg, ArgType candidateType) {
+		if (Consts.DEBUG_TYPE_INFERENCE) {
+			LOG.debug("-> update type for: {} to {}", arg, candidateType);
 		}
 		if (arg instanceof RegisterArg) {
 			RegisterArg reg = (RegisterArg) arg;
@@ -122,6 +199,12 @@ public final class TypeUpdate {
 	}

 	private @Nullable TypeUpdateResult verifyType(TypeUpdateInfo updateInfo, InsnArg arg, ArgType candidateType) {
+		if (candidateType == null) {
+			throw new JadxRuntimeException("Null type update for arg: " + arg);
+		}
+		if (updateInfo.isProcessed(arg)) {
+			return CHANGED;
+		}
 		ArgType currentType = arg.getType();
 		TypeUpdateFlags typeUpdateFlags = updateInfo.getFlags();
 		if (Objects.equals(currentType, candidateType)) {
@@ -194,27 +277,16 @@ public final class TypeUpdate {
 		if (!inBounds(updateInfo, ssaVar, typeInfo.getBounds(), candidateType)) {
 			return REJECT;
 		}
-		TypeUpdateResult result = requestUpdate(updateInfo, ssaVar.getAssign(), candidateType);
-		boolean allSame = result == SAME;
-		if (result != REJECT) {
-			List<RegisterArg> useList = ssaVar.getUseList();
-			for (RegisterArg arg : useList) {
-				result = requestUpdate(updateInfo, arg, candidateType);
-				if (result == REJECT) {
-					break;
-				}
-				if (result != SAME) {
-					allSame = false;
-				}
+		var updateCallback = new ArgsListUpdateCallback<>(this, updateInfo, ssaVar.getUseList(), candidateType, true);
+		updateCallback.setFinalResultCallback(result -> {
+			if (result == REJECT) {
+				// rollback update for all registers in current SSA var
+				updateInfo.rollbackUpdate(ssaVar.getAssign());
+				ssaVar.getUseList().forEach(updateInfo::rollbackUpdate);
 			}
-		}
-		if (result == REJECT) {
-			// rollback update for all registers in current SSA var
-			updateInfo.rollbackUpdate(ssaVar.getAssign());
-			ssaVar.getUseList().forEach(updateInfo::rollbackUpdate);
-			return REJECT;
-		}
-		return allSame ? SAME : CHANGED;
+			return result;
+		});
+		return queueDirectTypeUpdate(updateInfo, ssaVar.getAssign(), candidateType, updateCallback);
 	}

 	private TypeUpdateResult requestUpdate(TypeUpdateInfo updateInfo, InsnArg arg, ArgType candidateType) {
@@ -222,14 +294,6 @@ public final class TypeUpdate {
 			return CHANGED;
 		}
 		updateInfo.requestUpdate(arg, candidateType);
-		TypeUpdateResult result = runListeners(updateInfo, arg, candidateType);
-		if (result == REJECT) {
-			updateInfo.rollbackUpdate(arg);
-		}
-		return result;
-	}
-
-	private TypeUpdateResult runListeners(TypeUpdateInfo updateInfo, InsnArg arg, ArgType candidateType) {
 		InsnNode insn = arg.getParentInsn();
 		if (insn == null) {
 			return SAME;
@@ -238,6 +302,9 @@ public final class TypeUpdate {
 		if (listener == null) {
 			return CHANGED;
 		}
+		if (Consts.DEBUG_TYPE_INFERENCE) {
+			LOG.debug("Run listener for insn: {}, arg: {}, type: {}", insn.getType(), arg, candidateType);
+		}
 		return listener.update(updateInfo, insn, arg, candidateType);
 	}

@@ -359,85 +426,30 @@ public final class TypeUpdate {
 			ArgType returnType = methodDetails.getReturnType();
 			List<ArgType> argTypes = methodDetails.getArgTypes();
 			int argsCount = argTypes.size();
+
+			Supplier<ArgType> getReturnType;
+			Function<Integer, ArgType> getArgType;
 			if (typeVarsMap.isEmpty()) {
 				// generics can't be resolved => use as is
-				return applyInvokeTypes(updateInfo, invoke, argsCount, knownTypeVars, () -> returnType, argTypes::get);
+				getReturnType = () -> returnType;
+				getArgType = argTypes::get;
+			} else {
+				// resolve types before apply
+				getReturnType = () -> typeUtils.replaceTypeVariablesUsingMap(returnType, typeVarsMap);
+				getArgType = argNum -> typeUtils.replaceClassGenerics(candidateType, argTypes.get(argNum));
 			}
-			// resolve types before apply
-			return applyInvokeTypes(updateInfo, invoke, argsCount, knownTypeVars,
-					() -> typeUtils.replaceTypeVariablesUsingMap(returnType, typeVarsMap),
-					argNum -> typeUtils.replaceClassGenerics(candidateType, argTypes.get(argNum)));
+			return new InvokeUpdateCallback(this, updateInfo, invoke, argsCount, knownTypeVars, getReturnType, getArgType)
+					.runQueue();
 		}
 		return SAME;
 	}

-	private TypeUpdateResult applyInvokeTypes(TypeUpdateInfo updateInfo, BaseInvokeNode invoke, int argsCount,
-			Set<ArgType> knownTypeVars, Supplier<ArgType> getReturnType, Function<Integer, ArgType> getArgType) {
-		boolean allSame = true;
-		RegisterArg resultArg = invoke.getResult();
-		if (resultArg != null && !resultArg.isTypeImmutable()) {
-			ArgType returnType = checkType(knownTypeVars, getReturnType.get());
-			if (returnType != null) {
-				TypeUpdateResult result = updateTypeChecked(updateInfo, resultArg, returnType);
-				if (result == REJECT) {
-					TypeCompareEnum compare = comparator.compareTypes(returnType, resultArg.getType());
-					if (compare.isWider()) {
-						return REJECT;
-					}
-				}
-				if (result == CHANGED) {
-					allSame = false;
-				}
-			}
-		}
-		int argOffset = invoke.getFirstArgOffset();
-		for (int i = 0; i < argsCount; i++) {
-			InsnArg invokeArg = invoke.getArg(argOffset + i);
-			if (!invokeArg.isTypeImmutable()) {
-				ArgType argType = checkType(knownTypeVars, getArgType.apply(i));
-				if (argType != null) {
-					TypeUpdateResult result = updateTypeChecked(updateInfo, invokeArg, argType);
-					if (result == REJECT) {
-						TypeCompareEnum compare = comparator.compareTypes(argType, invokeArg.getType());
-						if (compare.isNarrow()) {
-							return REJECT;
-						}
-					}
-					if (result == CHANGED) {
-						allSame = false;
-					}
-				}
-			}
-		}
-		return allSame ? SAME : CHANGED;
-	}
-
-	@Nullable
-	private ArgType checkType(Set<ArgType> knownTypeVars, @Nullable ArgType type) {
-		if (type == null) {
-			return null;
-		}
-		if (type.isWildcard()) {
-			return null;
-		}
-		if (type.containsTypeVariable()) {
-			if (knownTypeVars.isEmpty()) {
-				return null;
-			}
-			Boolean hasUnknown = type.visitTypes(t -> t.isGenericType() && !knownTypeVars.contains(t) ? Boolean.TRUE : null);
-			if (hasUnknown != null) {
-				return null;
-			}
-		}
-		return type;
-	}
-
 	private TypeUpdateResult sameFirstArgListener(TypeUpdateInfo updateInfo, InsnNode insn, InsnArg arg, ArgType candidateType) {
 		InsnArg changeArg = isAssign(insn, arg) ? insn.getArg(0) : insn.getResult();
 		if (updateInfo.hasUpdateWithType(changeArg, candidateType)) {
 			return CHANGED;
 		}
-		return updateTypeChecked(updateInfo, changeArg, candidateType);
+		return queueTypeUpdate(updateInfo, changeArg, candidateType, null);
 	}

 	private TypeUpdateResult moveListener(TypeUpdateInfo updateInfo, InsnNode insn, InsnArg arg, ArgType candidateType) {
@@ -451,18 +463,19 @@ public final class TypeUpdate {
 		TypeCompareEnum cmp = comparator.compareTypes(candidateType, changeArg.getType());
 		boolean correctType = cmp.isEqual() || (assignChanged ? cmp.isWider() : cmp.isNarrow());

-		TypeUpdateResult result = updateTypeChecked(updateInfo, changeArg, candidateType);
-		if (result == SAME && !correctType) {
-			if (Consts.DEBUG_TYPE_INFERENCE) {
-				LOG.debug("Move insn types mismatch: {} -> {}, change arg: {}, insn: {}",
-						candidateType, changeArg.getType(), changeArg, insn);
+		return queueTypeUpdate(updateInfo, changeArg, candidateType, result -> {
+			if (result == SAME && !correctType) {
+				if (Consts.DEBUG_TYPE_INFERENCE) {
+					LOG.debug("Move insn types mismatch: {} -> {}, change arg: {}, insn: {}",
+							candidateType, changeArg.getType(), changeArg, insn);
+				}
+				return REJECT;
 			}
-			return REJECT;
-		}
-		if (result == REJECT && correctType) {
-			return CHANGED;
-		}
-		return result;
+			if (result == REJECT && correctType) {
+				return CHANGED;
+			}
+			return result;
+		});
 	}

 	/**
@@ -470,22 +483,12 @@ public final class TypeUpdate {
 	 */
 	private TypeUpdateResult allSameListener(TypeUpdateInfo updateInfo, InsnNode insn, InsnArg arg, ArgType candidateType) {
 		if (!isAssign(insn, arg)) {
-			return updateTypeChecked(updateInfo, insn.getResult(), candidateType);
+			return queueTypeUpdate(updateInfo, insn.getResult(), candidateType, null);
 		}
-		boolean allSame = true;
-		for (InsnArg insnArg : insn.getArguments()) {
-			if (insnArg == arg) {
-				continue;
-			}
-			TypeUpdateResult result = updateTypeChecked(updateInfo, insnArg, candidateType);
-			if (result == REJECT) {
-				return result;
-			}
-			if (result != SAME) {
-				allSame = false;
-			}
-		}
-		return allSame ? SAME : CHANGED;
+		// update args with same type
+		var updateCallback = new ArgsListUpdateCallback<>(this, updateInfo, insn.getArgList(), candidateType, false);
+		updateCallback.setArgsFilter(a -> a != arg);
+		return updateCallback.runFirstQueue();
 	}

 	private TypeUpdateResult arithListener(TypeUpdateInfo updateInfo, InsnNode insn, InsnArg arg, ArgType candidateType) {
@@ -501,32 +504,26 @@ public final class TypeUpdate {
 	 * Try to set candidate type to all args, don't fail on reject
 	 */
 	private TypeUpdateResult suggestAllSameListener(TypeUpdateInfo updateInfo, InsnNode insn, InsnArg arg, ArgType candidateType) {
+		var updateCallback = new ArgsListUpdateCallback<>(this, updateInfo, insn.getArgList(), candidateType, false);
+		updateCallback.setArgsFilter(a -> a != arg);
+		updateCallback.setIgnoreReject(true);
 		if (!isAssign(insn, arg)) {
 			RegisterArg resultArg = insn.getResult();
 			if (resultArg != null) {
-				updateTypeChecked(updateInfo, resultArg, candidateType);
+				// start with result
+				return queueTypeUpdate(updateInfo, resultArg, candidateType, updateCallback);
 			}
 		}
-		boolean allSame = true;
-		for (InsnArg insnArg : insn.getArguments()) {
-			if (insnArg != arg) {
-				TypeUpdateResult result = updateTypeChecked(updateInfo, insnArg, candidateType);
-				if (result == REJECT) {
-					// ignore
-				} else if (result != SAME) {
-					allSame = false;
-				}
-			}
-		}
-		return allSame ? SAME : CHANGED;
+		// start with first arg
+		return updateCallback.runFirstQueue();
 	}

 	private TypeUpdateResult checkCastListener(TypeUpdateInfo updateInfo, InsnNode insn, InsnArg arg, ArgType candidateType) {
 		IndexInsnNode checkCast = (IndexInsnNode) insn;
 		if (isAssign(insn, arg)) {
 			InsnArg insnArg = insn.getArg(0);
-			TypeUpdateResult result = updateTypeChecked(updateInfo, insnArg, candidateType);
-			return result == REJECT ? SAME : result;
+			return queueTypeUpdate(updateInfo, insnArg, candidateType,
+					r -> r == REJECT ? SAME : r);
 		}
 		ArgType castType = (ArgType) checkCast.getIndex();
 		TypeCompareEnum res = comparator.compareTypes(candidateType, castType);
@@ -543,7 +540,7 @@ public final class TypeUpdate {
 		}
 		if (res == TypeCompareEnum.NARROW_BY_GENERIC && candidateType.containsGeneric()) {
 			// propagate generic type to result
-			return updateTypeChecked(updateInfo, checkCast.getResult(), candidateType);
+			return queueTypeUpdate(updateInfo, checkCast.getResult(), candidateType, null);
 		}
 		ArgType currentType = checkCast.getArg(0).getType();
 		return candidateType.equals(currentType) ? SAME : CHANGED;
@@ -569,18 +566,19 @@ public final class TypeUpdate {

 	private TypeUpdateResult arrayGetListener(TypeUpdateInfo updateInfo, InsnNode insn, InsnArg arg, ArgType candidateType) {
 		if (isAssign(insn, arg)) {
-			TypeUpdateResult result = updateTypeChecked(updateInfo, insn.getArg(0), ArgType.array(candidateType));
-			if (result == REJECT) {
-				ArgType arrType = insn.getArg(0).getType();
-				if (arrType.isTypeKnown() && arrType.isArray() && arrType.getArrayElement().isPrimitive()) {
-					TypeCompareEnum compResult = comparator.compareTypes(candidateType, arrType.getArrayElement());
-					if (compResult == TypeCompareEnum.WIDER) {
-						// allow implicit upcast for primitive types (int a = byteArr[n])
-						return CHANGED;
+			return queueTypeUpdate(updateInfo, insn.getArg(0), ArgType.array(candidateType), result -> {
+				if (result == REJECT) {
+					ArgType arrType = insn.getArg(0).getType();
+					if (arrType.isTypeKnown() && arrType.isArray() && arrType.getArrayElement().isPrimitive()) {
+						TypeCompareEnum compResult = comparator.compareTypes(candidateType, arrType.getArrayElement());
+						if (compResult == TypeCompareEnum.WIDER) {
+							// allow implicit upcast for primitive types (int a = byteArr[n])
+							return CHANGED;
+						}
 					}
 				}
-			}
-			return result;
+				return result;
+			});
 		}
 		InsnArg arrArg = insn.getArg(0);
 		if (arrArg == arg) {
@@ -588,18 +586,19 @@ public final class TypeUpdate {
 			if (arrayElement == null) {
 				return REJECT;
 			}
-			TypeUpdateResult result = updateTypeChecked(updateInfo, insn.getResult(), arrayElement);
-			if (result == REJECT) {
-				ArgType resType = insn.getResult().getType();
-				if (resType.isTypeKnown() && resType.isPrimitive()) {
-					TypeCompareEnum compResult = comparator.compareTypes(resType, arrayElement);
-					if (compResult == TypeCompareEnum.WIDER) {
-						// allow implicit upcast for primitive types (int a = byteArr[n])
-						return CHANGED;
+			return queueTypeUpdate(updateInfo, insn.getResult(), arrayElement, result -> {
+				if (result == REJECT) {
+					ArgType resType = insn.getResult().getType();
+					if (resType.isTypeKnown() && resType.isPrimitive()) {
+						TypeCompareEnum compResult = comparator.compareTypes(resType, arrayElement);
+						if (compResult == TypeCompareEnum.WIDER) {
+							// allow implicit upcast for primitive types (int a = byteArr[n])
+							return CHANGED;
+						}
 					}
 				}
-			}
-			return result;
+				return result;
+			});
 		}
 		// index argument
 		return SAME;
@@ -613,21 +612,22 @@ public final class TypeUpdate {
 			if (arrayElement == null) {
 				return REJECT;
 			}
-			TypeUpdateResult result = updateTypeChecked(updateInfo, putArg, arrayElement);
-			if (result == REJECT) {
-				ArgType putType = putArg.getType();
-				if (putType.isTypeKnown()) {
-					TypeCompareEnum compResult = comparator.compareTypes(arrayElement, putType);
-					if (compResult == TypeCompareEnum.WIDER || compResult == TypeCompareEnum.WIDER_BY_GENERIC) {
-						// allow wider result (i.e. allow put any objects in Object[] or byte in int[])
-						return CHANGED;
+			return queueTypeUpdate(updateInfo, putArg, arrayElement, result -> {
+				if (result == REJECT) {
+					ArgType putType = putArg.getType();
+					if (putType.isTypeKnown()) {
+						TypeCompareEnum compResult = comparator.compareTypes(arrayElement, putType);
+						if (compResult == TypeCompareEnum.WIDER || compResult == TypeCompareEnum.WIDER_BY_GENERIC) {
+							// allow wider result (i.e. allow put any objects in Object[] or byte in int[])
+							return CHANGED;
+						}
 					}
 				}
-			}
-			return result;
+				return result;
+			});
 		}
 		if (arrArg == putArg) {
-			return updateTypeChecked(updateInfo, arrArg, ArgType.array(candidateType));
+			return queueTypeUpdate(updateInfo, arrArg, ArgType.array(candidateType), null);
 		}
 		// index
 		return SAME;
@@ -637,26 +637,27 @@ public final class TypeUpdate {
 		InsnArg firstArg = insn.getArg(0);
 		InsnArg secondArg = insn.getArg(1);
 		InsnArg updateArg = firstArg == arg ? secondArg : firstArg;
-		TypeUpdateResult result = updateTypeChecked(updateInfo, updateArg, candidateType);
-		if (result == REJECT) {
-			// soft checks for objects and array - exact type not compared
-			ArgType updateArgType = updateArg.getType();
-			if (candidateType.isObject() && updateArgType.canBeObject()) {
-				return SAME;
-			}
-			if (candidateType.isArray() && updateArgType.canBeArray()) {
-				return SAME;
-			}
-			if (candidateType.isPrimitive()) {
-				if (updateArgType.canBePrimitive(candidateType.getPrimitiveType())) {
+		return queueTypeUpdate(updateInfo, updateArg, candidateType, result -> {
+			if (result == REJECT) {
+				// soft checks for objects and array - exact type not compared
+				ArgType updateArgType = updateArg.getType();
+				if (candidateType.isObject() && updateArgType.canBeObject()) {
 					return SAME;
 				}
-				if (updateArgType.isTypeKnown() && candidateType.getRegCount() == updateArgType.getRegCount()) {
+				if (candidateType.isArray() && updateArgType.canBeArray()) {
 					return SAME;
 				}
+				if (candidateType.isPrimitive()) {
+					if (updateArgType.canBePrimitive(candidateType.getPrimitiveType())) {
+						return SAME;
+					}
+					if (updateArgType.isTypeKnown() && candidateType.getRegCount() == updateArgType.getRegCount()) {
+						return SAME;
+					}
+				}
 			}
-		}
-		return result;
+			return result;
+		});
 	}

 	private static boolean isAssign(InsnNode insn, InsnArg arg) {
@@ -1,14 +1,18 @@
 package jadx.core.dex.visitors.typeinference;

+import java.util.ArrayList;
 import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;

+import org.jetbrains.annotations.Nullable;
+
 import jadx.api.JadxArgs;
 import jadx.core.dex.instructions.args.ArgType;
 import jadx.core.dex.instructions.args.InsnArg;
 import jadx.core.dex.nodes.MethodNode;
+import jadx.core.utils.ListUtils;
 import jadx.core.utils.Utils;
 import jadx.core.utils.exceptions.JadxOverflowException;
 import jadx.core.utils.exceptions.JadxRuntimeException;
@@ -17,6 +21,8 @@ public class TypeUpdateInfo {
 	private final MethodNode mth;
 	private final TypeUpdateFlags flags;
 	private final Map<InsnArg, TypeUpdateEntry> updateMap = new IdentityHashMap<>();
+	private final List<TypeUpdateRequest> queue = new ArrayList<>();
+	private final List<TypeUpdateRequest> callbackQueue = new ArrayList<>();
 	private final int updatesLimitCount;
 	private int updateSeq = 0;

@@ -26,6 +32,24 @@ public class TypeUpdateInfo {
 		this.updatesLimitCount = mth.getInsnsCount() * args.getTypeUpdatesLimitCount();
 	}

+	public void queueRequest(TypeUpdateRequest request) {
+		queue.add(request);
+	}
+
+	public void saveCallback(TypeUpdateRequest request) {
+		if (request.getCallback() != null) {
+			callbackQueue.add(request);
+		}
+	}
+
+	public @Nullable TypeUpdateRequest pollNextRequest() {
+		return ListUtils.pollLast(queue);
+	}
+
+	public @Nullable TypeUpdateRequest pollNextCallback() {
+		return ListUtils.pollLast(callbackQueue);
+	}
+
 	public void requestUpdate(InsnArg arg, ArgType changeType) {
 		TypeUpdateEntry prev = updateMap.put(arg, new TypeUpdateEntry(updateSeq++, arg, changeType));
 		if (prev != null) {
@@ -0,0 +1,41 @@
+package jadx.core.dex.visitors.typeinference;
+
+import org.jetbrains.annotations.Nullable;
+
+import jadx.core.dex.instructions.args.ArgType;
+import jadx.core.dex.instructions.args.InsnArg;
+
+public class TypeUpdateRequest {
+	private final InsnArg arg;
+	private final ArgType candidateType;
+	private final boolean direct;
+	private final @Nullable ITypeUpdateCallback callback;
+
+	public TypeUpdateRequest(InsnArg arg, ArgType candidateType, boolean direct, @Nullable ITypeUpdateCallback callback) {
+		this.arg = arg;
+		this.candidateType = candidateType;
+		this.direct = direct;
+		this.callback = callback;
+	}
+
+	public InsnArg getArg() {
+		return arg;
+	}
+
+	public ArgType getCandidateType() {
+		return candidateType;
+	}
+
+	public boolean isDirect() {
+		return direct;
+	}
+
+	public @Nullable ITypeUpdateCallback getCallback() {
+		return callback;
+	}
+
+	@Override
+	public String toString() {
+		return "TypeUpdateRequest{arg=" + arg + ", candidateType=" + candidateType + '}';
+	}
+}
@@ -227,4 +227,15 @@ public class ListUtils {
 		}
 		return list;
 	}
+
+	public static <T> @Nullable T pollLast(List<T> list) {
+		if (list == null) {
+			return null;
+		}
+		int size = list.size();
+		if (size == 0) {
+			return null;
+		}
+		return list.remove(size - 1);
+	}
 }
@@ -3,10 +3,12 @@ package jadx.tests.integration.types;
 import org.junit.jupiter.api.Test;

 import jadx.tests.api.IntegrationTest;
-import jadx.tests.api.utils.assertj.JadxAssertions;
+
+import static jadx.tests.api.utils.assertj.JadxAssertions.assertThat;

 public class TestArrayTypes extends IntegrationTest {

+	@SuppressWarnings({ "ThrowablePrintedToSystemOut", "unused" })
 	public static class TestCls {

 		public void test() {
@@ -25,7 +27,7 @@ public class TestArrayTypes extends IntegrationTest {

 	@Test
 	public void test() {
-		JadxAssertions.assertThat(getClassNode(TestCls.class))
+		assertThat(getClassNode(TestCls.class))
 				.code()
 				.containsOne("use(new Object[]{e});");
 	}
@@ -33,7 +35,7 @@ public class TestArrayTypes extends IntegrationTest {
 	@Test
 	public void testNoDebug() {
 		noDebugInfo();
-		JadxAssertions.assertThat(getClassNode(TestCls.class))
+		assertThat(getClassNode(TestCls.class))
 				.code()
 				.containsOne("use(new Object[]{exc});");
 	}