diff --git a/src/ch/usi/inf/sp/cfg/ControlFlowGraphBuilder.java b/src/ch/usi/inf/sp/cfg/ControlFlowGraphBuilder.java index 718f208..2f3c190 100644 --- a/src/ch/usi/inf/sp/cfg/ControlFlowGraphBuilder.java +++ b/src/ch/usi/inf/sp/cfg/ControlFlowGraphBuilder.java @@ -3,6 +3,9 @@ package ch.usi.inf.sp.cfg; import java.util.*; import ch.usi.inf.sp.bytecode.Disassembler; +import ch.usi.inf.sp.cfg.builder.CFGBuilder; +import ch.usi.inf.sp.cfg.builder.Edge; +import ch.usi.inf.sp.cfg.builder.MultiMap; import org.objectweb.asm.Opcodes; import org.objectweb.asm.tree.AbstractInsnNode; import org.objectweb.asm.tree.InsnList; @@ -15,146 +18,7 @@ import org.objectweb.asm.tree.TableSwitchInsnNode; public class ControlFlowGraphBuilder { - private static class MultiMap { - private final Map> innerMap = new HashMap<>(); - - public void put(K key, V value) { - innerMap.computeIfAbsent(key, (k) -> new ArrayList<>()); - innerMap.get(key).add(value); - } - - public List 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 MultiMap labelToIncomingEdges = new MultiMap<>(); - - // get the list of all instructions in that method - final InsnList instructions = method.instructions; - for (int i = 0; i < instructions.size(); i++) { - final AbstractInsnNode instruction = instructions.get(i); - - switch (instruction.getType()) { - case AbstractInsnNode.JUMP_INSN -> { - final JumpInsnNode jumpInsn = (JumpInsnNode) instruction; - 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(); - labelToIncomingEdges.put(label, new Edge(instruction, value)); - } - if (lookupSwitchInsnNode.dflt != null) - 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); - labelToIncomingEdges.put(label, new Edge(instruction, value)); - } - if (tableSwitchInsnNode.dflt != null) - labelToIncomingEdges.put(tableSwitchInsnNode.dflt, new Edge(instruction, "default")); - } - } - } - - final ControlFlowGraph graph = new ControlFlowGraph(); - final List labels = new ArrayList<>(); - final Map insnToBlock = new HashMap<>(); - - BasicBlock currentBasicBlock = new BasicBlock(0); - graph.addNode(currentBasicBlock); - - graph.addEntryEdge(currentBasicBlock); - - for (int i = 0; i < instructions.size(); i++) { - final AbstractInsnNode instruction = instructions.get(i); - final int opcode = instruction.getOpcode(); - final int type = instruction.getType(); - - currentBasicBlock.appendInstruction(Disassembler.disassembleInstruction(instruction, i, instructions)); - - if (isReturnInstruction(opcode)) { - graph.addExitEdge(currentBasicBlock); - } - - final boolean isPointedLabel = instruction instanceof LabelNode && - labelToIncomingEdges.containsKey((LabelNode) instruction); - - insnToBlock.put(instruction, currentBasicBlock); - - if (isPointedLabel) { - labels.add((LabelNode) instruction); - } - - if (isEndOfBlock(instruction, labelToIncomingEdges)) { - final int nextI = i + 1; - - // if we're not at the end - if (nextI < instructions.size()) { - final BasicBlock previousBasicBlock = currentBasicBlock; - - currentBasicBlock = new BasicBlock(nextI); - graph.addNode(currentBasicBlock); - - // Handle "implicit" edges from this basic block to the next one, or the exit node - if (instructions.get(nextI).getOpcode() != Opcodes.GOTO) { - graph.addFallthroughEdge(previousBasicBlock, currentBasicBlock); - } - } - } - } - - 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 new CFGBuilder(method).create(); } - - private static boolean isReturnInstruction(int opcode) { - return opcode == Opcodes.RETURN - || opcode == Opcodes.ARETURN - || opcode == Opcodes.LRETURN - || opcode == Opcodes.IRETURN - || opcode == Opcodes.FRETURN; - } - - private static boolean isEndOfBlock(AbstractInsnNode instruction, MultiMap labelToInstruction) { - final AbstractInsnNode nextInsn = instruction.getNext(); - if (nextInsn == null) { - return false; // cannot start another bb at the end of the method with 0 instructions - } - - final int type = instruction.getType(); - if (type == AbstractInsnNode.JUMP_INSN - || type == AbstractInsnNode.LOOKUPSWITCH_INSN - || type == AbstractInsnNode.TABLESWITCH_INSN) { - return true; // if we're branching or jumping after this instruction, we NEED to cut the current bb short - } - - // if the next instruction is a label some other bb may jump into then cut, otherwise continue - return nextInsn instanceof LabelNode && labelToInstruction.containsKey((LabelNode) nextInsn); - } - - private record Edge(AbstractInsnNode instruction, String condition) { - } - } diff --git a/src/ch/usi/inf/sp/cfg/builder/CFGBuilder.java b/src/ch/usi/inf/sp/cfg/builder/CFGBuilder.java new file mode 100644 index 0000000..7622b13 --- /dev/null +++ b/src/ch/usi/inf/sp/cfg/builder/CFGBuilder.java @@ -0,0 +1,134 @@ +package ch.usi.inf.sp.cfg.builder; + +import ch.usi.inf.sp.bytecode.Disassembler; +import ch.usi.inf.sp.cfg.BasicBlock; +import ch.usi.inf.sp.cfg.ControlFlowEdge; +import ch.usi.inf.sp.cfg.ControlFlowGraph; +import org.objectweb.asm.Opcodes; +import org.objectweb.asm.tree.*; + +import java.util.*; + +public final class CFGBuilder { + private final MultiMap labelToIncomingEdges = new MultiMap<>(); + private final List labels = new ArrayList<>(); + private final Map insnToBlock = new HashMap<>(); + + private final ControlFlowGraph graph = new ControlFlowGraph(); + private final MethodNode method; + private BasicBlock currentBasicBlock; + + public CFGBuilder(MethodNode method) { + this.method = method; + setNewBasicBlock(0); + graph.addEntryEdge(currentBasicBlock); + } + + private void setNewBasicBlock(int id) { + currentBasicBlock = new BasicBlock(id); + graph.addNode(currentBasicBlock); + } + + public ControlFlowGraph create() { + // get the list of all instructions in that method + final InsnList instructions = method.instructions; + for (int i = 0; i < instructions.size(); i++) { + final AbstractInsnNode instruction = instructions.get(i); + + switch (instruction.getType()) { + case AbstractInsnNode.JUMP_INSN -> labelToIncomingEdges.put( + ((JumpInsnNode) instruction).label, + new Edge(instruction, "T") + ); + case AbstractInsnNode.LOOKUPSWITCH_INSN -> + captureSwitchEdges(new LookupSwitchInstructionNodeInfo((LookupSwitchInsnNode) instruction)); + case AbstractInsnNode.TABLESWITCH_INSN -> + captureSwitchEdges(new TableSwitchInstructionNodeInfo((TableSwitchInsnNode) instruction)); + } + } + + for (int i = 0; i < instructions.size(); i++) { + final AbstractInsnNode instruction = instructions.get(i); + + currentBasicBlock.appendInstruction(Disassembler.disassembleInstruction(instruction, i, instructions)); + + if (isReturnInstruction(instruction)) { + graph.addExitEdge(currentBasicBlock); + } + + insnToBlock.put(instruction, currentBasicBlock); + + if (isInsnSignificantLabel(instruction)) { + labels.add((LabelNode) instruction); + } + + if (isEndOfBlock(instruction)) { + final int nextI = i + 1; + + final BasicBlock previousBasicBlock = currentBasicBlock; + setNewBasicBlock(nextI); + + // Handle "implicit" edges from this basic block to the next one, or the exit node + if (instructions.get(nextI).getOpcode() != Opcodes.GOTO) { + graph.addFallthroughEdge(previousBasicBlock, currentBasicBlock); + } + } + } + + for (final LabelNode label : labels) { + final BasicBlock toBB = Objects.requireNonNull(insnToBlock.get(label)); + for (final ch.usi.inf.sp.cfg.builder.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 void captureSwitchEdges(final SwitchInstructionNodeInfo info) { + for (int j = 0; j < info.getCaseCount(); j++) { + final LabelNode label = info.getLabelForCase(j); + final String value = info.getKeyForCase(j); + labelToIncomingEdges.put(label, new Edge(info.getNode(), value)); + } + if (info.getDefaultCase() != null) { + labelToIncomingEdges.put(info.getDefaultCase(), new Edge(info.getNode(), "default")); + } + } + + private boolean isInsnSignificantLabel(final AbstractInsnNode node) { + return node instanceof LabelNode && labelToIncomingEdges.containsKey((LabelNode) node); + } + + private boolean isReturnInstruction(final AbstractInsnNode instruction) { + final int opcode = instruction.getOpcode(); + return opcode == Opcodes.RETURN || + opcode == Opcodes.ARETURN || + opcode == Opcodes.LRETURN || + opcode == Opcodes.IRETURN || + opcode == Opcodes.FRETURN; + } + + private boolean isEndOfBlock(final AbstractInsnNode instruction) { + final AbstractInsnNode nextInsn = instruction.getNext(); + + if (nextInsn == null) { + return false; // cannot start another bb at the end of the method with 0 instructions + } + + final int type = instruction.getType(); + if (type == AbstractInsnNode.JUMP_INSN || + type == AbstractInsnNode.LOOKUPSWITCH_INSN || + type == AbstractInsnNode.TABLESWITCH_INSN) { + return true; // if we're branching or jumping after this instruction, we NEED to cut the current bb short + } + + // if the next instruction is a label some other bb may jump into then cut, otherwise continue + return isInsnSignificantLabel(nextInsn); + } +} diff --git a/src/ch/usi/inf/sp/cfg/builder/Edge.java b/src/ch/usi/inf/sp/cfg/builder/Edge.java new file mode 100644 index 0000000..5602d66 --- /dev/null +++ b/src/ch/usi/inf/sp/cfg/builder/Edge.java @@ -0,0 +1,6 @@ +package ch.usi.inf.sp.cfg.builder; + +import org.objectweb.asm.tree.AbstractInsnNode; + +public record Edge(AbstractInsnNode instruction, String condition) { +} diff --git a/src/ch/usi/inf/sp/cfg/builder/LookupSwitchInstructionNodeInfo.java b/src/ch/usi/inf/sp/cfg/builder/LookupSwitchInstructionNodeInfo.java new file mode 100644 index 0000000..bb88912 --- /dev/null +++ b/src/ch/usi/inf/sp/cfg/builder/LookupSwitchInstructionNodeInfo.java @@ -0,0 +1,38 @@ +package ch.usi.inf.sp.cfg.builder; + +import org.objectweb.asm.tree.AbstractInsnNode; +import org.objectweb.asm.tree.LabelNode; +import org.objectweb.asm.tree.LookupSwitchInsnNode; + +public class LookupSwitchInstructionNodeInfo implements SwitchInstructionNodeInfo { + private final LookupSwitchInsnNode node; + + public LookupSwitchInstructionNodeInfo(LookupSwitchInsnNode node) { + this.node = node; + } + + @Override + public AbstractInsnNode getNode() { + return node; + } + + @Override + public int getCaseCount() { + return node.labels.size(); + } + + @Override + public LabelNode getLabelForCase(int index) { + return node.labels.get(index); + } + + @Override + public String getKeyForCase(int index) { + return Integer.toString(node.keys.get(index)); + } + + @Override + public LabelNode getDefaultCase() { + return node.dflt; + } +} diff --git a/src/ch/usi/inf/sp/cfg/builder/MultiMap.java b/src/ch/usi/inf/sp/cfg/builder/MultiMap.java new file mode 100644 index 0000000..0d243b7 --- /dev/null +++ b/src/ch/usi/inf/sp/cfg/builder/MultiMap.java @@ -0,0 +1,23 @@ +package ch.usi.inf.sp.cfg.builder; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class MultiMap { + private final Map> innerMap = new HashMap<>(); + + public void put(K key, V value) { + innerMap.computeIfAbsent(key, (k) -> new ArrayList<>()); + innerMap.get(key).add(value); + } + + public List getAll(K key) { + return new ArrayList<>(innerMap.getOrDefault(key, List.of())); + } + + public boolean containsKey(K key) { + return innerMap.containsKey(key); + } +} diff --git a/src/ch/usi/inf/sp/cfg/builder/SwitchInstructionNodeInfo.java b/src/ch/usi/inf/sp/cfg/builder/SwitchInstructionNodeInfo.java new file mode 100644 index 0000000..c4993b0 --- /dev/null +++ b/src/ch/usi/inf/sp/cfg/builder/SwitchInstructionNodeInfo.java @@ -0,0 +1,13 @@ +package ch.usi.inf.sp.cfg.builder; + +import org.objectweb.asm.tree.AbstractInsnNode; +import org.objectweb.asm.tree.LabelNode; + +public interface SwitchInstructionNodeInfo { + AbstractInsnNode getNode(); + int getCaseCount(); + LabelNode getLabelForCase(int index); + String getKeyForCase(int index); + + LabelNode getDefaultCase(); +} diff --git a/src/ch/usi/inf/sp/cfg/builder/TableSwitchInstructionNodeInfo.java b/src/ch/usi/inf/sp/cfg/builder/TableSwitchInstructionNodeInfo.java new file mode 100644 index 0000000..6091c2a --- /dev/null +++ b/src/ch/usi/inf/sp/cfg/builder/TableSwitchInstructionNodeInfo.java @@ -0,0 +1,38 @@ +package ch.usi.inf.sp.cfg.builder; + +import org.objectweb.asm.tree.AbstractInsnNode; +import org.objectweb.asm.tree.LabelNode; +import org.objectweb.asm.tree.TableSwitchInsnNode; + +public class TableSwitchInstructionNodeInfo implements SwitchInstructionNodeInfo { + private final TableSwitchInsnNode node; + + public TableSwitchInstructionNodeInfo(TableSwitchInsnNode node) { + this.node = node; + } + + @Override + public AbstractInsnNode getNode() { + return node; + } + + @Override + public int getCaseCount() { + return node.labels.size(); + } + + @Override + public LabelNode getLabelForCase(int index) { + return node.labels.get(index); + } + + @Override + public String getKeyForCase(int index) { + return Integer.toString(node.min + index); + } + + @Override + public LabelNode getDefaultCase() { + return node.dflt; + } +}