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 {
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) {
final Map<LabelNode, Edge> labelToSource = new HashMap<>();
final MultiMap<LabelNode, Edge> labelToIncomingEdges = new MultiMap<>();
// get the list of all instructions in that method
final InsnList instructions = method.instructions;
@ -24,38 +41,39 @@ public class ControlFlowGraphBuilder {
final AbstractInsnNode instruction = instructions.get(i);
switch (instruction.getType()) {
case AbstractInsnNode.JUMP_INSN:
case AbstractInsnNode.JUMP_INSN -> {
final JumpInsnNode jumpInsn = (JumpInsnNode) instruction;
labelToSource.put(jumpInsn.label, new Edge(instruction, "T"));
case AbstractInsnNode.LOOKUPSWITCH_INSN:
labelToIncomingEdges.put(jumpInsn.label, new Edge(instruction, "T"));
}
case AbstractInsnNode.LOOKUPSWITCH_INSN -> {
final LookupSwitchInsnNode lookupSwitchInsnNode = (LookupSwitchInsnNode) instruction;
for (int j = 0; j < lookupSwitchInsnNode.labels.size(); j++) {
final LabelNode label = lookupSwitchInsnNode.labels.get(i);
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)
labelToSource.put(lookupSwitchInsnNode.dflt, new Edge(instruction, "default"));
case AbstractInsnNode.TABLESWITCH_INSN:
labelToIncomingEdges.put(lookupSwitchInsnNode.dflt, new Edge(instruction, "default"));
}
case AbstractInsnNode.TABLESWITCH_INSN -> {
final TableSwitchInsnNode tableSwitchInsnNode = (TableSwitchInsnNode) instruction;
for (int k = 0; k < tableSwitchInsnNode.labels.size(); k++) {
final LabelNode label = tableSwitchInsnNode.labels.get(i);
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)
labelToSource.put(tableSwitchInsnNode.dflt, new Edge(instruction, "default"));
labelToIncomingEdges.put(tableSwitchInsnNode.dflt, new Edge(instruction, "default"));
}
}
}
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);
graph.addNode(currentBasicBlock);
blocksById.put(0, currentBasicBlock);
graph.addEntryEdge(currentBasicBlock);
@ -66,18 +84,20 @@ public class ControlFlowGraphBuilder {
currentBasicBlock.appendInstruction(Disassembler.disassembleInstruction(instruction, i, instructions));
if (opcode == Opcodes.RETURN
|| opcode == Opcodes.ARETURN
|| opcode == Opcodes.LRETURN
|| opcode == Opcodes.IRETURN
|| opcode == Opcodes.FRETURN) {
if (isReturnInstruction(opcode)) {
graph.addExitEdge(currentBasicBlock);
}
if (type == AbstractInsnNode.JUMP_INSN
|| type == AbstractInsnNode.LOOKUPSWITCH_INSN
|| type == AbstractInsnNode.TABLESWITCH_INSN
|| labelToSource.containsKey(instruction)) {
final boolean isPointedLabel = instruction instanceof LabelNode &&
labelToIncomingEdges.containsKey((LabelNode) instruction);
insnToBlock.put(instruction, currentBasicBlock);
if (isPointedLabel) {
labels.add((LabelNode) instruction);
}
if (isBranchingInstruction(type, isPointedLabel)) {
final int nextI = i + 1;
// if we're not at the end
@ -86,7 +106,6 @@ public class ControlFlowGraphBuilder {
currentBasicBlock = new BasicBlock(nextI);
graph.addNode(currentBasicBlock);
blocksById.put(nextI, currentBasicBlock);
// Handle "implicit" edges from this basic block to the next one, or the exit node
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;
}
private static class Edge {
final AbstractInsnNode instruction;
final String condition;
private static boolean isReturnInstruction(int opcode) {
return opcode == Opcodes.RETURN
|| opcode == Opcodes.ARETURN
|| opcode == Opcodes.LRETURN
|| opcode == Opcodes.IRETURN
|| opcode == Opcodes.FRETURN;
}
Edge(AbstractInsnNode instruction, String condition) {
this.instruction = instruction;
this.condition = condition;
}
private static boolean isBranchingInstruction(int type, boolean isPointedLabel) {
return type == AbstractInsnNode.JUMP_INSN
|| type == AbstractInsnNode.LOOKUPSWITCH_INSN
|| type == AbstractInsnNode.TABLESWITCH_INSN
|| isPointedLabel;
}
private record Edge(AbstractInsnNode instruction, String condition) {
}
}