almost works

This commit is contained in:
Claudio Maggioni 2023-10-10 17:01:51 +02:00
parent e64edeecbe
commit d19481edca

View file

@ -15,8 +15,25 @@ import org.objectweb.asm.tree.TableSwitchInsnNode;
public class ControlFlowGraphBuilder { public class ControlFlowGraphBuilder {
private static class MultiMap<K, V> {
private final Map<K, List<V>> innerMap = new HashMap<>();
public void put(K key, V value) {
innerMap.computeIfAbsent(key, (k) -> new ArrayList<>());
innerMap.get(key).add(value);
}
public List<V> getAll(K key) {
return new ArrayList<>(innerMap.getOrDefault(key, List.of()));
}
public boolean containsKey(K key) {
return innerMap.containsKey(key);
}
}
public static ControlFlowGraph createControlFlowGraph(final MethodNode method) { public static ControlFlowGraph createControlFlowGraph(final MethodNode method) {
final Map<LabelNode, Edge> labelToSource = new HashMap<>(); final MultiMap<LabelNode, Edge> labelToIncomingEdges = new MultiMap<>();
// get the list of all instructions in that method // get the list of all instructions in that method
final InsnList instructions = method.instructions; final InsnList instructions = method.instructions;
@ -24,38 +41,39 @@ public class ControlFlowGraphBuilder {
final AbstractInsnNode instruction = instructions.get(i); final AbstractInsnNode instruction = instructions.get(i);
switch (instruction.getType()) { switch (instruction.getType()) {
case AbstractInsnNode.JUMP_INSN: case AbstractInsnNode.JUMP_INSN -> {
final JumpInsnNode jumpInsn = (JumpInsnNode) instruction; final JumpInsnNode jumpInsn = (JumpInsnNode) instruction;
labelToSource.put(jumpInsn.label, new Edge(instruction, "T")); labelToIncomingEdges.put(jumpInsn.label, new Edge(instruction, "T"));
}
case AbstractInsnNode.LOOKUPSWITCH_INSN: case AbstractInsnNode.LOOKUPSWITCH_INSN -> {
final LookupSwitchInsnNode lookupSwitchInsnNode = (LookupSwitchInsnNode) instruction; final LookupSwitchInsnNode lookupSwitchInsnNode = (LookupSwitchInsnNode) instruction;
for (int j = 0; j < lookupSwitchInsnNode.labels.size(); j++) { for (int j = 0; j < lookupSwitchInsnNode.labels.size(); j++) {
final LabelNode label = lookupSwitchInsnNode.labels.get(i); final LabelNode label = lookupSwitchInsnNode.labels.get(i);
final String value = lookupSwitchInsnNode.keys.get(i).toString(); final String value = lookupSwitchInsnNode.keys.get(i).toString();
labelToSource.put(label, new Edge(instruction, value)); labelToIncomingEdges.put(label, new Edge(instruction, value));
} }
if (lookupSwitchInsnNode.dflt != null) if (lookupSwitchInsnNode.dflt != null)
labelToSource.put(lookupSwitchInsnNode.dflt, new Edge(instruction, "default")); labelToIncomingEdges.put(lookupSwitchInsnNode.dflt, new Edge(instruction, "default"));
}
case AbstractInsnNode.TABLESWITCH_INSN: case AbstractInsnNode.TABLESWITCH_INSN -> {
final TableSwitchInsnNode tableSwitchInsnNode = (TableSwitchInsnNode) instruction; final TableSwitchInsnNode tableSwitchInsnNode = (TableSwitchInsnNode) instruction;
for (int k = 0; k < tableSwitchInsnNode.labels.size(); k++) { for (int k = 0; k < tableSwitchInsnNode.labels.size(); k++) {
final LabelNode label = tableSwitchInsnNode.labels.get(i); final LabelNode label = tableSwitchInsnNode.labels.get(i);
final String value = Integer.toString(tableSwitchInsnNode.min + k); final String value = Integer.toString(tableSwitchInsnNode.min + k);
labelToSource.put(label, new Edge(instruction, value)); labelToIncomingEdges.put(label, new Edge(instruction, value));
} }
if (tableSwitchInsnNode.dflt != null) if (tableSwitchInsnNode.dflt != null)
labelToSource.put(tableSwitchInsnNode.dflt, new Edge(instruction, "default")); labelToIncomingEdges.put(tableSwitchInsnNode.dflt, new Edge(instruction, "default"));
}
} }
} }
final ControlFlowGraph graph = new ControlFlowGraph(); final ControlFlowGraph graph = new ControlFlowGraph();
final Map<Integer, BasicBlock> blocksById = new HashMap<>(); final List<LabelNode> labels = new ArrayList<>();
final Map<AbstractInsnNode, BasicBlock> insnToBlock = new HashMap<>();
BasicBlock currentBasicBlock = new BasicBlock(0); BasicBlock currentBasicBlock = new BasicBlock(0);
graph.addNode(currentBasicBlock); graph.addNode(currentBasicBlock);
blocksById.put(0, currentBasicBlock);
graph.addEntryEdge(currentBasicBlock); graph.addEntryEdge(currentBasicBlock);
@ -66,18 +84,20 @@ public class ControlFlowGraphBuilder {
currentBasicBlock.appendInstruction(Disassembler.disassembleInstruction(instruction, i, instructions)); currentBasicBlock.appendInstruction(Disassembler.disassembleInstruction(instruction, i, instructions));
if (opcode == Opcodes.RETURN if (isReturnInstruction(opcode)) {
|| opcode == Opcodes.ARETURN
|| opcode == Opcodes.LRETURN
|| opcode == Opcodes.IRETURN
|| opcode == Opcodes.FRETURN) {
graph.addExitEdge(currentBasicBlock); graph.addExitEdge(currentBasicBlock);
} }
if (type == AbstractInsnNode.JUMP_INSN final boolean isPointedLabel = instruction instanceof LabelNode &&
|| type == AbstractInsnNode.LOOKUPSWITCH_INSN labelToIncomingEdges.containsKey((LabelNode) instruction);
|| type == AbstractInsnNode.TABLESWITCH_INSN
|| labelToSource.containsKey(instruction)) { insnToBlock.put(instruction, currentBasicBlock);
if (isPointedLabel) {
labels.add((LabelNode) instruction);
}
if (isBranchingInstruction(type, isPointedLabel)) {
final int nextI = i + 1; final int nextI = i + 1;
// if we're not at the end // if we're not at the end
@ -86,7 +106,6 @@ public class ControlFlowGraphBuilder {
currentBasicBlock = new BasicBlock(nextI); currentBasicBlock = new BasicBlock(nextI);
graph.addNode(currentBasicBlock); graph.addNode(currentBasicBlock);
blocksById.put(nextI, currentBasicBlock);
// Handle "implicit" edges from this basic block to the next one, or the exit node // Handle "implicit" edges from this basic block to the next one, or the exit node
if (instructions.get(nextI).getOpcode() != Opcodes.GOTO) { if (instructions.get(nextI).getOpcode() != Opcodes.GOTO) {
@ -96,19 +115,36 @@ public class ControlFlowGraphBuilder {
} }
} }
// TODO: add remaining branching edges and refactor code for (final LabelNode label : labels) {
final BasicBlock toBB = Objects.requireNonNull(insnToBlock.get(label));
for (final Edge incomingEdge : labelToIncomingEdges.getAll(label)) {
final BasicBlock fromBB = Objects.requireNonNull(insnToBlock.get(incomingEdge.instruction));
final ControlFlowEdge edge = new ControlFlowEdge(incomingEdge.condition);
graph.addEdge(edge);
graph.connect(fromBB, edge, toBB);
}
}
return graph; return graph;
} }
private static class Edge { private static boolean isReturnInstruction(int opcode) {
final AbstractInsnNode instruction; return opcode == Opcodes.RETURN
final String condition; || opcode == Opcodes.ARETURN
|| opcode == Opcodes.LRETURN
|| opcode == Opcodes.IRETURN
|| opcode == Opcodes.FRETURN;
}
Edge(AbstractInsnNode instruction, String condition) { private static boolean isBranchingInstruction(int type, boolean isPointedLabel) {
this.instruction = instruction; return type == AbstractInsnNode.JUMP_INSN
this.condition = condition; || type == AbstractInsnNode.LOOKUPSWITCH_INSN
} || type == AbstractInsnNode.TABLESWITCH_INSN
|| isPointedLabel;
}
private record Edge(AbstractInsnNode instruction, String condition) {
} }
} }