Add weights to digraph

Because in some cases (e.g. partitioning) the order is needed,
add weights to the digraph to get an (somewhat) stable
topological sort.

Change-Id: I5ef1acc6338ac93c593faa0eafe26cbed42ed887
Signed-off-by: Andreas Florath <andreas@florath.net>

diskimage_builder/graph/digraph.py

@@ -12,6 +12,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 #
+import bisect
 
 
 class Digraph(object):
@@ -20,6 +21,32 @@ class Digraph(object):
     Each node of the digraph must have a unique name.
     """
 
+    class Edge(object):
+        """Directed graph edge.
+
+        The digraph has weighted edges.  This class holds the weight and
+        a reference to the node.
+        """
+
+        def __init__(self, node, weight):
+            self.__node = node
+            self.__weight = weight
+
+        def __eq__(self, other):
+            return self.__weight == other.get_weight() \
+                and self.__node == other.get_node()
+
+        def __lt__(self, other):
+            return self.__weight < other.get_weight()
+
+        def get_node(self):
+            """Return the (pointed to) node"""
+            return self.__node
+
+        def get_weight(self):
+            """Return the edge's weight"""
+            return self.__weight
+
     class Node(object):
         """Directed graph node.
 
@@ -35,33 +62,38 @@ class Digraph(object):
             computed.
             """
             self.__name = name
-            self.__incoming = set()
+            self.__incoming = []
-            self.__outgoing = set()
+            self.__outgoing = []
+
+        def __repr__(self):
+            return "<Node [%s]>" % self.__name
 
         def get_name(self):
             """Returns the name of the node."""
             return self.__name
 
-        def add_incoming(self, node):
+        def add_incoming(self, node, weight):
             """Add node to the incoming list."""
+            bisect.insort(self.__incoming, Digraph.Edge(node, weight))
 
-            self.__incoming.add(node)
+        def add_outgoing(self, node, weight):
-
+            """Add node to the outgoing list."""
-        def add_outgoing(self, node):
+            bisect.insort(self.__outgoing, Digraph.Edge(node, weight))
-            """Add node to the incoming list."""
-
-            self.__outgoing.add(node)
 
         def get_iter_outgoing(self):
             """Return an iterator over the outgoing nodes."""
 
-            return iter(self.__outgoing)
+            return iter([x.get_node() for x in self.__outgoing])
+
+        def has_incoming(self):
+            """Returns True if the node has incoming edges"""
+            return self.__incoming
 
         @staticmethod
         def __as_named_list(inlist):
             """Return given list as list of names."""
 
-            return map(lambda x: x.get_name(), inlist)
+            return [x.get_node().get_name() for x in inlist]
 
         def get_outgoing_as_named_list(self):
             """Return the names of all outgoing nodes as a list."""
@@ -105,7 +137,7 @@ class Digraph(object):
                                    "exists" % node.get_name())
         self._named_nodes[anode.get_name()] = anode
 
-    def create_edge(self, anode, bnode):
+    def create_edge(self, anode, bnode, weight=0):
         """Creates an edge from a to b - both must be nodes."""
 
         assert issubclass(anode.__class__, Digraph.Node)
@@ -114,8 +146,8 @@ class Digraph(object):
         assert anode == self._named_nodes[anode.get_name()]
         assert bnode.get_name() in self._named_nodes.keys()
         assert bnode == self._named_nodes[bnode.get_name()]
-        anode.add_outgoing(bnode)
+        anode.add_outgoing(bnode, weight)
-        bnode.add_incoming(anode)
+        bnode.add_incoming(anode, weight)
 
     def get_iter_nodes_values(self):
         """Returns the nodes dict to the values.
@@ -144,7 +176,7 @@ class Digraph(object):
             rval[node.get_name()] = node.get_outgoing_as_named_list()
         return rval
 
-    def topological_sort(dg):
+    def topological_sort(self):
         """Digraph topological search.
 
         This algorithm is based upon a depth first search with
@@ -169,7 +201,9 @@ class Digraph(object):
                 tsort.insert(0, node)
 
         # The 'main' function of the topological sort
-        for node in dg.get_iter_nodes_values():
+        for node in self.get_iter_nodes_values():
+            if node.has_incoming():
+                continue
             visit(node)
 
         return tsort
@@ -189,6 +223,6 @@ def node_list_to_node_name_list(node_list):
     """Converts a node list into a list of the corresponding node names."""
 
     node_name_list = []
-    for n in node_list:
+    for node in node_list:
-        node_name_list.append(n.get_name())
+        node_name_list.append(node.get_name())
     return node_name_list

diskimage_builder/tests/functional/test_graph.py

@@ -121,3 +121,23 @@ class TestDigraph(testtools.TestCase):
             self.assertTrue(False)
         except RuntimeError:
             pass
+
+    def test_iter_outgoing_weight_01(self):
+        """Tests iter_outgoing in a graph with weights"""
+
+        digraph = Digraph()
+        node0 = Digraph.Node("R")
+        digraph.add_node(node0)
+        node1 = Digraph.Node("A")
+        digraph.add_node(node1)
+        node2 = Digraph.Node("B")
+        digraph.add_node(node2)
+        node3 = Digraph.Node("C")
+        digraph.add_node(node3)
+
+        digraph.create_edge(node0, node1, 1)
+        digraph.create_edge(node0, node2, 2)
+        digraph.create_edge(node0, node3, 3)
+
+        self.assertEqual([node1, node2, node3],
+                         list(node0.get_iter_outgoing()))

diskimage_builder/tests/functional/test_graph_toposort.py

@@ -12,9 +12,11 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import testtools
+
+from diskimage_builder.graph.digraph import Digraph
 from diskimage_builder.graph.digraph import digraph_create_from_dict
 from diskimage_builder.graph.digraph import node_list_to_node_name_list
-import testtools
 
 
 class TestTopologicalSearch(testtools.TestCase):
@@ -67,3 +69,48 @@ class TestTopologicalSearch(testtools.TestCase):
         self.assertTrue(tnames.index('A') < tnames.index('B'))
         self.assertTrue(tnames.index('B') < tnames.index('C'))
         self.assertTrue(tnames.index('D') < tnames.index('E'))
+
+    def test_tsort_006(self):
+        """Complex digraph with weights"""
+
+        digraph = Digraph()
+        node0 = Digraph.Node("R")
+        digraph.add_node(node0)
+        node1 = Digraph.Node("A")
+        digraph.add_node(node1)
+        node2 = Digraph.Node("B")
+        digraph.add_node(node2)
+        node3 = Digraph.Node("C")
+        digraph.add_node(node3)
+        node4 = Digraph.Node("B1")
+        digraph.add_node(node4)
+        node5 = Digraph.Node("B2")
+        digraph.add_node(node5)
+        node6 = Digraph.Node("B3")
+        digraph.add_node(node6)
+
+        digraph.create_edge(node0, node1, 1)
+        digraph.create_edge(node0, node2, 2)
+        digraph.create_edge(node0, node3, 3)
+
+        digraph.create_edge(node2, node4, 7)
+        digraph.create_edge(node2, node5, 14)
+        digraph.create_edge(node2, node6, 21)
+
+        tsort = digraph.topological_sort()
+        tnames = node_list_to_node_name_list(tsort)
+
+        # Also here: many possible solutions
+        self.assertTrue(tnames.index('R') < tnames.index('A'))
+        self.assertTrue(tnames.index('R') < tnames.index('B'))
+        self.assertTrue(tnames.index('R') < tnames.index('C'))
+        self.assertTrue(tnames.index('B') < tnames.index('B1'))
+        self.assertTrue(tnames.index('B') < tnames.index('B2'))
+        self.assertTrue(tnames.index('B') < tnames.index('B3'))
+
+        # In addition in the weighted graph the following
+        # must also hold:
+        self.assertTrue(tnames.index('B') < tnames.index('A'))
+        self.assertTrue(tnames.index('C') < tnames.index('B'))
+        self.assertTrue(tnames.index('B2') < tnames.index('B1'))
+        self.assertTrue(tnames.index('B3') < tnames.index('B2'))