diff --git a/src/ch/usi/inf/sp/cfg/ControlFlowGraphBuilder.java b/src/ch/usi/inf/sp/cfg/ControlFlowGraphBuilder.java
index 4806b89..f78a2c8 100644
--- a/src/ch/usi/inf/sp/cfg/ControlFlowGraphBuilder.java
+++ b/src/ch/usi/inf/sp/cfg/ControlFlowGraphBuilder.java
@@ -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) {
     }
 
 }