Update ReValue model to support tracking array contents

Added an example test case showing new tracking capabilities being leveraged by constant folding transformer.

Author: Col-E

recaf-core/src/main/java/software/coley/recaf/util/Types.java

@@ -142,6 +142,19 @@ public static Type array(@Nonnull Type type, int dimensions) {
 		return Type.getType("[".repeat(dimensions) + type.getDescriptor());
 	}
 
+	/**
+	 * @param arrayType
+	 * 		Some array type.
+	 *
+	 * @return The type when removing one dimension from the array.
+	 */
+	@Nonnull
+	public static Type undimension(@Nonnull Type arrayType) {
+		if (arrayType.getSort() != Type.ARRAY)
+			throw new IllegalStateException("Not an array: " + arrayType);
+		return Type.getType(arrayType.getDescriptor().substring(1));
+	}
+
 	/**
 	 * @param methodType
 	 * 		Parsed method descriptor type.

recaf-core/src/main/java/software/coley/recaf/util/analysis/ReAnalyzer.java

@@ -2,6 +2,7 @@
 
 import jakarta.annotation.Nonnull;
 import org.objectweb.asm.tree.analysis.Analyzer;
+import org.objectweb.asm.tree.analysis.Frame;
 import software.coley.recaf.util.analysis.value.ReValue;
 
 /**
@@ -28,4 +29,14 @@ public ReAnalyzer(@Nonnull ReInterpreter interpreter) {
 	public ReInterpreter getInterpreter() {
 		return interpreter;
 	}
+
+	@Override
+	protected Frame<ReValue> newFrame(int numLocals, int numStack) {
+		return new ReFrame(numLocals, numStack);
+	}
+
+	@Override
+	protected Frame<ReValue> newFrame(Frame<? extends ReValue> frame) {
+		return new ReFrame(frame);
+	}
 }

recaf-core/src/main/java/software/coley/recaf/util/analysis/ReFrame.java

@@ -0,0 +1,66 @@
+package software.coley.recaf.util.analysis;
+
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.AbstractInsnNode;
+import org.objectweb.asm.tree.analysis.AnalyzerException;
+import org.objectweb.asm.tree.analysis.Frame;
+import org.objectweb.asm.tree.analysis.Interpreter;
+import software.coley.recaf.util.analysis.value.ArrayValue;
+import software.coley.recaf.util.analysis.value.ReValue;
+
+/**
+ * Analysis frame for some enhanced value tracking.
+ *
+ * @author Matt Coley
+ */
+public class ReFrame extends Frame<ReValue> {
+	public ReFrame(int numLocals, int maxStack) {
+		super(numLocals, maxStack);
+	}
+
+	public ReFrame(Frame<? extends ReValue> frame) {
+		super(frame);
+	}
+
+	@Override
+	public void execute(AbstractInsnNode insn, Interpreter<ReValue> interpreter) throws AnalyzerException {
+		switch (insn.getOpcode()) {
+			case Opcodes.IASTORE:
+			case Opcodes.LASTORE:
+			case Opcodes.FASTORE:
+			case Opcodes.DASTORE:
+			case Opcodes.AASTORE:
+			case Opcodes.BASTORE:
+			case Opcodes.CASTORE:
+			case Opcodes.SASTORE:
+				// Support updating array contents in this frame when values are stored into it.
+				ReValue value = pop();
+				ReValue index = pop();
+				ReValue array = pop();
+				ReValue updatedArray = interpreter.ternaryOperation(insn, array, index, value);
+				if (updatedArray != array) {
+					for (int i = 0; i < getStackSize(); i++) {
+						ReValue stack = getStack(i);
+						if (stack == array) {
+							setStack(i, updatedArray);
+						} else if (stack instanceof ArrayValue stackArray) {
+							ArrayValue updatedStackArray = stackArray.updatedCopyIfContained(array, updatedArray);
+							setStack(i, updatedStackArray);
+						}
+					}
+					for (int i = 0; i < getLocals(); i++) {
+						ReValue local = getLocal(i);
+						if (local == array) {
+							setLocal(i, updatedArray);
+						} else if (local instanceof ArrayValue stackArray) {
+							ArrayValue updatedStackArray = stackArray.updatedCopyIfContained(array, updatedArray);
+							setLocal(i, updatedStackArray);
+						}
+					}
+				}
+				break;
+			default:
+				super.execute(insn, interpreter);
+		}
+	}
+}

recaf-core/src/main/java/software/coley/recaf/util/analysis/ReInterpreter.java

@@ -347,19 +347,24 @@ public ReValue binaryOperation(@Nonnull AbstractInsnNode insn, @Nonnull ReValue
 			case BALOAD:
 			case CALOAD:
 			case SALOAD:
-				// We aren't tracking array contents, so nothing to do here.
+				if (value1 instanceof ArrayValue array && value2 instanceof IntValue index && index.value().isPresent())
+					return array.getValue(index.value().getAsInt());
 				return IntValue.UNKNOWN;
 			case FALOAD:
-				// We aren't tracking array contents, so nothing to do here.
+				if (value1 instanceof ArrayValue array && value2 instanceof IntValue index && index.value().isPresent())
+					return array.getValue(index.value().getAsInt());
 				return FloatValue.UNKNOWN;
 			case LALOAD:
-				// We aren't tracking array contents, so nothing to do here.
+				if (value1 instanceof ArrayValue array && value2 instanceof IntValue index && index.value().isPresent())
+					return array.getValue(index.value().getAsInt());
 				return LongValue.UNKNOWN;
 			case DALOAD:
-				// We aren't tracking array contents, so nothing to do here.
+				if (value1 instanceof ArrayValue array && value2 instanceof IntValue index && index.value().isPresent())
+					return array.getValue(index.value().getAsInt());
 				return DoubleValue.UNKNOWN;
 			case AALOAD:
-				// We aren't tracking array contents, so nothing to do here.
+				if (value1 instanceof ArrayValue array && value2 instanceof IntValue index && index.value().isPresent())
+					return array.getValue(index.value().getAsInt());
 				return ObjectValue.VAL_OBJECT_MAYBE_NULL;
 			case IADD:
 				if (value1 instanceof IntValue i1 && value2 instanceof IntValue i2) return i1.add(i2);
@@ -492,19 +497,36 @@ public ReValue binaryOperation(@Nonnull AbstractInsnNode insn, @Nonnull ReValue
 	}
 
 	@Override
-	public ReValue ternaryOperation(@Nonnull AbstractInsnNode insn, @Nonnull ReValue value1, @Nonnull ReValue value2, ReValue value3) {
-		// We don't track array operations, but this would cover:
+	public ReValue ternaryOperation(@Nonnull AbstractInsnNode insn, @Nonnull ReValue value1, @Nonnull ReValue value2, @Nonnull ReValue value3) {
+		// This method covers the following instructions:
 		//  IASTORE, LASTORE, FASTORE, DASTORE, AASTORE, BASTORE, CASTORE, SASTORE
-		// Load operations are handled in 'binaryOperation'
-		return null;
+		// It will always be an array-store operation.
+		// NOTE: Load operations are handled in 'binaryOperation'
+		if (value1 instanceof ArrayValue array && value2 instanceof IntValue index && index.value().isPresent())
+			return array.setValue(index.value().getAsInt(), value3);
+		return value1;
 	}
 
 	@Override
 	public ReValue naryOperation(@Nonnull AbstractInsnNode insn, @Nonnull List<? extends ReValue> values) {
 		int opcode = insn.getOpcode();
 		if (opcode == MULTIANEWARRAY) {
 			Type type = Type.getType(((MultiANewArrayInsnNode) insn).desc);
-			return newValue(type, Nullness.NOT_NULL);
+
+			// Extract dimensions from passed values.
+			// Unknown values will be negative, which we will sanity check for.
+			int[] dimensions = values.stream()
+					.mapToInt(v -> v instanceof IntValue intValue && intValue.value().isPresent() ?
+							intValue.value().getAsInt() : -1)
+					.toArray();
+			if (dimensions.length == 0) // Sanity check, should never occur.
+				return newValue(type, Nullness.NOT_NULL);
+			for (int dimension : dimensions)
+				if (dimension < 0) // Validate dimensions are valid
+					return newValue(type, Nullness.NOT_NULL);
+
+			// Make array of known dimensions.
+			return ArrayValue.multiANewArray(type, dimensions);
 		} else if (opcode == INVOKEDYNAMIC) {
 			Type returnType = Type.getReturnType(((InvokeDynamicInsnNode) insn).desc);
 			return newValue(returnType);

recaf-core/src/main/java/software/coley/recaf/util/analysis/value/ArrayValue.java

@@ -1,10 +1,14 @@
 package software.coley.recaf.util.analysis.value;
 
 import jakarta.annotation.Nonnull;
+import jakarta.annotation.Nullable;
+import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
+import software.coley.recaf.util.Types;
 import software.coley.recaf.util.analysis.Nullness;
 import software.coley.recaf.util.analysis.value.impl.ArrayValueImpl;
 
+import java.util.Arrays;
 import java.util.OptionalInt;
 
 /**
@@ -63,6 +67,46 @@ static ArrayValue of(@Nonnull Type type, @Nonnull Nullness nullness, int length)
 		return new ArrayValueImpl(type, nullness, length);
 	}
 
+	/**
+	 * @param type
+	 * 		Array type.
+	 * @param dimensions
+	 * 		Dimensions of the array to create.
+	 *
+	 * @return Array value created from a {@link Opcodes#MULTIANEWARRAY} instruction.
+	 */
+	@Nonnull
+	static ReValue multiANewArray(@Nonnull Type type, @Nonnull int[] dimensions) {
+		int length = dimensions[0];
+		if (dimensions.length == 1)
+			return of(type, Nullness.NOT_NULL, length);
+		return new ArrayValueImpl(type, Nullness.NOT_NULL, length,
+				i -> multiANewArray(Types.undimension(type), Arrays.copyOfRange(dimensions, 1, dimensions.length))
+		);
+	}
+
+	/**
+	 * @param index
+	 * 		Index to assign value at.
+	 * @param value
+	 * 		Value to assign.
+	 *
+	 * @return New array with the given value assigned at the given index.
+	 */
+	@Nonnull
+	ArrayValue setValue(int index, @Nonnull ReValue value);
+
+	/**
+	 * @param originalValue
+	 * 		Some value.
+	 * @param updatedValue
+	 * 		Some updated version of the value.
+	 *
+	 * @return New array with the given original value replaced with the updated value.
+	 */
+	@Nonnull
+	ArrayValue updatedCopyIfContained(@Nonnull ReValue originalValue, @Nonnull ReValue updatedValue);
+
 	@Override
 	default boolean hasKnownValue() {
 		return false;
@@ -116,4 +160,13 @@ default int dimensions() {
 	 */
 	@Nonnull
 	OptionalInt getFirstDimensionLength();
+
+	/**
+	 * @param index
+	 * 		Index within {@link #getFirstDimensionLength()}.
+	 *
+	 * @return Value, if known, at the given index. Otherwise, a {@link ReValue} of the array's element type..
+	 */
+	@Nullable
+	ReValue getValue(int index);
 }