--- a/python/ppci/codegen/graph.py	Mon Jan 27 19:58:07 2014 +0100
+++ b/python/ppci/codegen/graph.py	Thu Jan 30 19:03:24 2014 +0100
@@ -7,9 +7,12 @@
     def __init__(self):
         self.nodes = set()
         self.edges = set()
+        self.adj_map = {}
 
     def addNode(self, n):
         self.nodes.add(n)
+        if n not in self.adj_map:
+            self.adj_map[n] = set()
 
     def delNode(self, n):
         self.nodes.remove(n)
@@ -18,6 +21,8 @@
         """ Add an edge from n to m """
         self.edges.add((n, m))
         self.edges.add((m, n))
+        self.adj_map[n].add(m)
+        self.adj_map[m].add(n)
 
     def hasEdge(self, n, m):
         return (n, m) in self.edges
@@ -28,8 +33,7 @@
 
     def adjecent(self, n):
         """ Return all nodes with edges to n """
-        # TODO: this can be optimized for speed:
-        return set(m for m in self.nodes if self.hasEdge(n, m))
+        return self.adj_map[n] & self.nodes
 
     def to_dot(self, f):
         """ Generate graphviz dot representation """
@@ -57,21 +61,37 @@
 
 
 class DiGraph(Graph):
-    """
-        Directed graph.
-    """
+    """ Directed graph. """
+    def __init__(self):
+        super().__init__()
+        self.suc_map = {}
+        self.pre_map = {}
+
     def addEdge(self, n, m):
-        """ Add an edge from n to m """
+        """ Add a directed edge from n to m """
+        assert n in self.nodes
+        assert m in self.nodes
         self.edges.add((n, m))
+        self.suc_map[n].add(m)
+        self.pre_map[m].add(n)
+        self.adj_map[n].add(m)
+        self.adj_map[m].add(n)
+
+    def addNode(self, n):
+        super().addNode(n)
+        if n not in self.suc_map:
+            self.suc_map[n] = set()
+        if n not in self.pre_map:
+            self.pre_map[n] = set()
 
     def hasEdge(self, n, m):
         return (n, m) in self.edges
 
     def successors(self, n):
-        return set(m for m in self.nodes if self.hasEdge(n, m))
+        return self.suc_map[n] & self.nodes
 
     def predecessors(self, n):
-        return set(m for m in self.nodes if self.hasEdge(m, n))
+        return self.pre_map[n] & self.nodes
 
 
 class DiNode(Node):