Additional instruction context actions (goto dec for type/field/methods, goto labels for switches, number radix conversion)

Author: Col-E

recaf-ui/src/main/java/software/coley/recaf/services/cell/context/BasicAssemblerContextMenuProviderFactory.java

@@ -7,17 +7,18 @@
 import javafx.scene.control.ContextMenu;
 import javafx.scene.control.Menu;
 import me.darknet.assembler.ast.ASTElement;
+import me.darknet.assembler.ast.primitive.ASTArray;
 import me.darknet.assembler.ast.primitive.ASTIdentifier;
 import me.darknet.assembler.ast.primitive.ASTInstruction;
 import me.darknet.assembler.ast.primitive.ASTLabel;
+import me.darknet.assembler.ast.primitive.ASTNumber;
+import me.darknet.assembler.ast.primitive.ASTObject;
 import me.darknet.assembler.ast.specific.ASTMethod;
+import me.darknet.assembler.parser.BytecodeFormat;
 import me.darknet.assembler.query.AssemblyQueries;
 import me.darknet.assembler.query.AssemblyUtils;
 import me.darknet.assembler.query.LabelQueryResult;
-import me.darknet.assembler.query.LabelReferenceKind;
-import me.darknet.assembler.query.LabelUsage;
 import me.darknet.assembler.query.VariableQueryResult;
-import me.darknet.assembler.query.VariableUsage;
 import me.darknet.assembler.query.resolution.ClassAnnotationResolution;
 import me.darknet.assembler.query.resolution.ClassExtends;
 import me.darknet.assembler.query.resolution.ClassImplements;
@@ -38,12 +39,14 @@
 import me.darknet.assembler.util.Location;
 import me.darknet.assembler.util.Range;
 import org.fxmisc.richtext.CodeArea;
+import org.objectweb.asm.Type;
 import org.slf4j.Logger;
 import software.coley.collections.Unchecked;
 import software.coley.recaf.analytics.logging.Logging;
 import software.coley.recaf.info.ClassInfo;
 import software.coley.recaf.info.JvmClassInfo;
 import software.coley.recaf.path.AssemblerPathData;
+import software.coley.recaf.path.ClassMemberPathNode;
 import software.coley.recaf.path.ClassPathNode;
 import software.coley.recaf.path.IncompletePathException;
 import software.coley.recaf.services.cell.icon.IconProviderService;
@@ -52,29 +55,28 @@
 import software.coley.recaf.ui.control.ActionMenuItem;
 import software.coley.recaf.ui.control.FontIconView;
 import software.coley.recaf.ui.control.richtext.Editor;
+import software.coley.recaf.util.Lang;
 import software.coley.recaf.util.SVG;
 import software.coley.recaf.workspace.model.Workspace;
 import software.coley.recaf.workspace.model.bundle.ClassBundle;
 import software.coley.recaf.workspace.model.bundle.JvmClassBundle;
 import software.coley.recaf.workspace.model.resource.WorkspaceResource;
 
+import java.math.BigInteger;
+import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 
 import static org.kordamp.ikonli.carbonicons.CarbonIcons.ARROW_RIGHT;
 import static org.kordamp.ikonli.carbonicons.CarbonIcons.HEALTH_CROSS;
 import static org.kordamp.ikonli.carbonicons.CarbonIcons.LIST_BOXES;
-import static org.objectweb.asm.Opcodes.CHECKCAST;
 import static org.objectweb.asm.Opcodes.IFEQ;
-import static org.objectweb.asm.Opcodes.INSTANCEOF;
 import static org.objectweb.asm.Opcodes.JSR;
-import static org.objectweb.asm.Opcodes.NEW;
-import static software.coley.collections.Unchecked.runnable;
-import static software.coley.recaf.util.Menus.action;
-import static software.coley.recaf.util.Menus.menu;
+import static software.coley.recaf.util.Menus.*;
 
 /**
  * Basic implementation for {@link AssemblerContextMenuProviderFactory}.
@@ -195,35 +197,19 @@ private static <T extends Resolution> void register(@Nonnull Class<T> type, @Non
 			// No items
 		});
 		register(InstructionResolution.class, (provider, menu, editor, workspace, resolution) -> {
-			// TODO:
-			//  - Different actions for different instructions (only differentiable by instruction name string atm)
-			//  - Convert integer representations (Hex, binary, decimal)
-			//  - Goto declaration for type/field/method references
-			//  - Goto declaration for jump instructions (with case hint)
 			ASTInstruction instruction = resolution.instruction();
 			ASTIdentifier identifier = instruction.identifier();
-			if (identifier != null && identifier.content() != null) {
-				int opcode = JvmOpcodes.opcode(identifier.content());
-				if (opcode >= IFEQ && opcode <= JSR && !instruction.arguments().isEmpty()) {
-					String labelName = instruction.arguments().getLast().content();
-					if (labelName != null)
-						menu.getItems().add(action("menu.goto.label", ARROW_RIGHT, () ->
-								gotoLabelDeclaration(resolution.method(), editor.getCodeArea(), labelName)));
-				} else if (opcode == NEW || opcode == CHECKCAST || opcode == INSTANCEOF) {
-					if (instruction.arguments().getLast() instanceof ASTIdentifier typeIdentifier) {
-						String typeName = typeIdentifier.content();
-						if (typeName != null) {
-							ClassPathNode typePath = workspace.findClass(typeName);
-							if (typePath != null) {
-								// TODO: We want to show the type that is targeted with this, similar to how we have in
-								//  other context menus. This code is also messy and there's probably some abstractions
-								//  we can make to reduce the indentation hell.
-								menu.getItems().add(action("menu.goto.class", ARROW_RIGHT, runnable(() -> provider.actions.gotoDeclaration(typePath))));
-							}
-						}
-					}
-				}
-			}
+			String instructionName = identifier == null ? null : identifier.content();
+			if (instructionName == null)
+				return;
+
+			CodeArea area = editor.getCodeArea();
+			ASTElement selectedElement = instruction.pick(area.getCaretPosition());
+			addIntegerConversionActions(menu, area, selectedElement);
+			addMemberReferenceAction(provider, menu, workspace, instructionName, instruction);
+			addTypeReferenceAction(provider, menu, workspace, instructionName, instruction);
+			addFlowLabelAction(menu, area, resolution.method(), instructionName, instruction);
+			addSwitchTargetActions(menu, area, resolution.method(), instructionName, instruction, selectedElement);
 		});
 		register(VariableDeclarationResolution.class, (provider, menu, editor, workspace, resolution) ->
 				addUsageMenu(menu, editor, variableUsageTargets(resolution.method(), resolution.variable().identity().name())));
@@ -251,6 +237,120 @@ private static <T extends Resolution> void register(@Nonnull Class<T> type, @Non
 		});
 	}
 
+	private static void addFlowLabelAction(@Nonnull ContextMenu menu, @Nonnull CodeArea area, @Nonnull ASTMethod method,
+	                                       @Nonnull String instructionName, @Nonnull ASTInstruction instruction) {
+		int opcode = JvmOpcodes.opcode(instructionName);
+		if (opcode < IFEQ || opcode > JSR || instruction.arguments().isEmpty())
+			return;
+
+		String labelName = instruction.arguments().getLast().content();
+		if (labelName != null) {
+			menu.getItems().add(action("menu.goto.label", ARROW_RIGHT, () ->
+					gotoLabelDeclaration(method, area, labelName)));
+		}
+	}
+
+	private static void addTypeReferenceAction(@Nonnull AbstractContextMenuProviderFactory provider, @Nonnull ContextMenu menu,
+	                                           @Nonnull Workspace workspace, @Nonnull String instructionName,
+	                                           @Nonnull ASTInstruction instruction) {
+		if (!AssemblyUtils.isTypeReferenceInstruction(BytecodeFormat.JVM, instructionName))
+			return;
+
+		int argumentIndex = AssemblyUtils.typeReferenceArgumentIndex(BytecodeFormat.JVM, instructionName);
+		if (argumentIndex < 0 || argumentIndex >= instruction.arguments().size())
+			return;
+
+		String typeName = normalizeTypeReference(instruction.arguments().get(argumentIndex).content());
+		if (typeName == null)
+			return;
+
+		ClassPathNode typePath = workspace.findClass(typeName);
+		ActionMenuItem gotoAction = action("menu.goto.class", ARROW_RIGHT, () -> {
+			try {
+				provider.actions.gotoDeclaration(Objects.requireNonNull(typePath));
+			} catch (IncompletePathException ex) {
+				logger.error("Cannot go to class due to incomplete path", ex);
+			}
+		});
+		if (typePath == null) gotoAction.setDisable(true);
+		menu.getItems().add(gotoAction);
+	}
+
+	private static void addMemberReferenceAction(@Nonnull AbstractContextMenuProviderFactory provider, @Nonnull ContextMenu menu,
+	                                             @Nonnull Workspace workspace, @Nonnull String instructionName,
+	                                             @Nonnull ASTInstruction instruction) {
+		MemberReference memberReference = parseMemberReference(instructionName, instruction);
+		if (memberReference == null)
+			return;
+
+		ClassPathNode ownerPath = workspace.findClass(memberReference.owner());
+		ClassMemberPathNode memberPath = ownerPath == null ? null : ownerPath.child(memberReference.name(), memberReference.descriptor());
+		String key = memberReference.isMethod() ? "menu.goto.method" : "menu.goto.field";
+		ActionMenuItem gotoAction = action(key, ARROW_RIGHT, () -> {
+			try {
+				provider.actions.gotoDeclaration(Objects.requireNonNull(memberPath));
+			} catch (IncompletePathException ex) {
+				logger.error("Cannot go to member due to incomplete path", ex);
+			}
+		});
+		if (memberPath == null) gotoAction.setDisable(true);
+		menu.getItems().add(gotoAction);
+	}
+
+	private static void addIntegerConversionActions(@Nonnull ContextMenu menu, @Nonnull CodeArea area,
+	                                                @Nullable ASTElement selectedElement) {
+		if (!(selectedElement instanceof ASTNumber number) || number.isFloatingPoint())
+			return;
+
+		String content = number.content();
+		Number value = number.number();
+		if (content == null || value == null)
+			return;
+
+		long longValue = value.longValue();
+		boolean wide = content.endsWith("L") || content.endsWith("l") || value instanceof Long;
+		int currentRadix = integerLiteralRadix(content);
+		if (currentRadix != 10)
+			menu.getItems().add(actionLiteral("Convert to decimal (%s)".formatted(formatIntegerLiteral(longValue, wide, 10)),
+					ARROW_RIGHT, () -> replaceAstText(area, number, formatIntegerLiteral(longValue, wide, 10))));
+		if (currentRadix != 16)
+			menu.getItems().add(actionLiteral("Convert to hex (%s)".formatted(formatIntegerLiteral(longValue, wide, 16)),
+					ARROW_RIGHT, () -> replaceAstText(area, number, formatIntegerLiteral(longValue, wide, 16))));
+		if (currentRadix != 2)
+			menu.getItems().add(actionLiteral("Convert to binary (%s)".formatted(formatIntegerLiteral(longValue, wide, 2)),
+					ARROW_RIGHT, () -> replaceAstText(area, number, formatIntegerLiteral(longValue, wide, 2))));
+	}
+
+	private static void addSwitchTargetActions(@Nonnull ContextMenu menu, @Nonnull CodeArea area, @Nonnull ASTMethod method,
+	                                           @Nonnull String instructionName, @Nonnull ASTInstruction instruction,
+	                                           @Nullable ASTElement selectedElement) {
+		if (!"lookupswitch".equals(instructionName) && !"tableswitch".equals(instructionName))
+			return;
+
+		List<SwitchTarget> targets = parseSwitchTargets(instructionName, instruction);
+		if (targets.isEmpty())
+			return;
+
+		if (selectedElement != null) {
+			for (SwitchTarget target : targets) {
+				if (target.matches(selectedElement)) {
+					menu.getItems().add(action("menu.goto.label", ARROW_RIGHT, () ->
+							gotoLabelDeclaration(method, area, target.labelName())));
+					break;
+				}
+			}
+		}
+
+		Menu switchMenu = new Menu();
+		switchMenu.textProperty().bind(Lang.getBinding("menu.goto.switch.title"));
+		switchMenu.setGraphic(new FontIconView(ARROW_RIGHT));
+		for (SwitchTarget target : targets) {
+			switchMenu.getItems().add(actionLiteral(target.description(), ARROW_RIGHT, () ->
+					gotoLabelDeclaration(method, area, target.labelName())));
+		}
+		menu.getItems().add(switchMenu);
+	}
+
 	@Nonnull
 	private static List<UsageTarget> variableUsageTargets(@Nonnull ASTMethod method, @Nonnull String variableName) {
 		VariableQueryResult queryResult = AssemblyQueries.variables(method);
@@ -330,8 +430,157 @@ private static void gotoLabelDeclaration(@Nonnull ASTMethod method, @Nonnull Cod
 			gotoAstElement(area, label);
 	}
 
+	// TODO: Surely some of this can be pulled up into JASM upstream...
+
+	private static void replaceAstText(@Nonnull CodeArea area, @Nonnull ASTElement element, @Nonnull String replacement) {
+		Range range = element.range();
+		area.replaceText(range.start(), range.end(), replacement);
+	}
+
+	@Nullable
+	private static String normalizeTypeReference(@Nullable String typeReference) {
+		if (typeReference == null || typeReference.isBlank())
+			return null;
+		if (typeReference.charAt(0) != '[' && typeReference.charAt(0) != 'L')
+			return typeReference;
+
+		Type type = Type.getType(typeReference);
+		while (type.getSort() == Type.ARRAY) type = type.getElementType();
+		return type.getSort() == Type.OBJECT ? type.getInternalName() : null;
+	}
+
+	@Nullable
+	private static MemberReference parseMemberReference(@Nonnull String instructionName, @Nonnull ASTInstruction instruction) {
+		if ((!isFieldReferenceInstruction(instructionName) && !isMethodReferenceInstruction(instructionName)) ||
+				instruction.arguments().size() < 2)
+			return null;
+
+		String ownerAndName = instruction.arguments().get(0).content();
+		String descriptor = instruction.arguments().get(1).content();
+		if (ownerAndName == null || descriptor == null)
+			return null;
+
+		int split = ownerAndName.lastIndexOf('.');
+		if (split <= 0 || split >= ownerAndName.length() - 1)
+			return null;
+
+		return new MemberReference(
+				ownerAndName.substring(0, split),
+				ownerAndName.substring(split + 1),
+				descriptor,
+				isMethodReferenceInstruction(instructionName)
+		);
+	}
+
+	private static boolean isFieldReferenceInstruction(@Nonnull String instructionName) {
+		return "getfield".equals(instructionName) || "getstatic".equals(instructionName) ||
+				"putfield".equals(instructionName) || "putstatic".equals(instructionName);
+	}
+
+	private static boolean isMethodReferenceInstruction(@Nonnull String instructionName) {
+		return instructionName.startsWith("invoke") && !"invokedynamic".equals(instructionName);
+	}
+
+	@Nonnull
+	private static List<SwitchTarget> parseSwitchTargets(@Nonnull String instructionName, @Nonnull ASTInstruction instruction) {
+		if (instruction.arguments().isEmpty())
+			return List.of();
+
+		ASTObject switchObject = instruction.argumentObject(instruction.arguments().size() - 1);
+		if (switchObject == null)
+			return List.of();
+
+		return "lookupswitch".equals(instructionName) ? parseLookupSwitchTargets(switchObject) : parseTableSwitchTargets(switchObject);
+	}
+
+	@Nonnull
+	private static List<SwitchTarget> parseLookupSwitchTargets(@Nonnull ASTObject switchObject) {
+		List<SwitchTarget> targets = new ArrayList<>();
+		for (int i = 0; i < switchObject.values().size(); i++) {
+			ASTIdentifier key = switchObject.values().key(i);
+			ASTElement value = switchObject.values().get(i);
+			String context = key == null ? null : key.literal();
+			String labelName = value == null ? null : value.content();
+			if (context == null || labelName == null)
+				continue;
+
+			targets.add(new SwitchTarget(
+					context,
+					labelName,
+					"default".equals(context) ? "Default -> " + labelName : "Case " + context + " -> " + labelName,
+					key,
+					value
+			));
+		}
+		return targets;
+	}
+
+	@Nonnull
+	private static List<SwitchTarget> parseTableSwitchTargets(@Nonnull ASTObject switchObject) {
+		List<SwitchTarget> targets = new ArrayList<>();
+		ASTElement defaultValue = switchObject.value("default");
+		String defaultLabel = defaultValue == null ? null : defaultValue.content();
+		if (defaultLabel != null)
+			targets.add(new SwitchTarget("default", defaultLabel, "Default -> " + defaultLabel,
+					switchObject.values().key("default"), defaultValue));
+
+		long min = 0;
+		ASTElement minValue = switchObject.value("min");
+		if (minValue instanceof ASTNumber number)
+			min = number.asLong();
+
+
+		ASTArray cases = switchObject.value("cases");
+		if (cases == null)
+			return targets;
+
+		long currentValue = min;
+		for (ASTElement caseTarget : cases.values()) {
+			String labelName = caseTarget.content();
+			if (labelName == null || labelName.isBlank())
+				continue;
+
+			String caseContext = Long.toString(currentValue++);
+			targets.add(new SwitchTarget(caseContext, labelName, "Case " + caseContext + " -> " + labelName,
+					caseTarget, caseTarget));
+		}
+		return targets;
+	}
+
+	private static int integerLiteralRadix(@Nonnull String content) {
+		int offset = content.startsWith("-") || content.startsWith("+") ? 1 : 0;
+		if (content.regionMatches(true, offset, "0x", 0, 2))
+			return 16;
+		if (content.regionMatches(true, offset, "0b", 0, 2))
+			return 2;
+		return 10;
+	}
+
+	@Nonnull
+	private static String formatIntegerLiteral(long value, boolean wide, int radix) {
+		BigInteger magnitude = BigInteger.valueOf(value).abs();
+		String sign = value < 0 ? "-" : "";
+		String text = switch (radix) {
+			case 16 -> sign + "0x" + magnitude.toString(16).toUpperCase(Locale.ROOT);
+			case 2 -> sign + "0b" + magnitude.toString(2);
+			default -> Long.toString(value);
+		};
+		return wide ? text + 'L' : text;
+	}
+
 	private record UsageTarget(@Nonnull String description, @Nonnull ASTElement element, @Nullable String iconPath) {}
 
+	private record MemberReference(@Nonnull String owner, @Nonnull String name, @Nonnull String descriptor,
+	                               boolean isMethod) {}
+
+	private record SwitchTarget(@Nonnull String context, @Nonnull String labelName, @Nonnull String description,
+	                            @Nullable ASTElement contextElement, @Nullable ASTElement labelElement) {
+		private boolean matches(@Nonnull ASTElement element) {
+			return element == contextElement || element == labelElement ||
+					context.equals(element.content()) || labelName.equals(element.content());
+		}
+	}
+
 	/**
 	 * @param <R>
 	 * 		Resolution impl type.