mode: compute per-DIF enrollments

1 files changed, 86 insertions, 10 deletions

diff --git a/rumba/model.py b/rumba/model.py
index f48b0dd..5c22185 100644
--- a/rumba/model.py
+++ b/rumba/model.py
@@ -236,6 +236,7 @@ class Experiment:
             nodes = list()
         self.nodes = nodes
         self.testbed = testbed
+        self.enrollment_strategy = 'minimal'  # 'full-mesh', 'manual'
 
     def __repr__(self):
         s = ""
@@ -269,12 +270,8 @@ class Experiment:
     def del_node(self, node):
         self.nodes.remove(node)
 
-    # Examine the nodes, DIFs, registrations and compute the registration
-    # and enrollment order, etc.
-    def generate(self):
-
-        ###### Compute registration/enrollment order for DIFs #######
-
+    # Compute registration/enrollment order for DIFs
+    def compute_dif_ordering(self):
         # Compute DIFs dependency graph, as both adjacency and incidence list.
         difsdeps_adj = dict()
         difsdeps_inc = dict()
@@ -326,11 +323,90 @@ class Experiment:
                   "circular dependencies, involving the following"\
                   " DIFs: %s" % circular_set)
 
-        del difsdeps_inc_cnt
-        del difsdeps_adj
-        del circular_set
+        print("DIF topological ordering: %s" % self.dif_ordering)
+
+    # Compute per-DIF graphs, to be called after compute_dif_ordering()
+    def compute_enrollments(self):
+        dif_graphs = dict()
+        self.enrollments = dict()
+
+        for dif in self.dif_ordering:
+            neighsets = dict()
+            dif_graphs[dif] = dict()
+            first = None
+
+            # For each N-1-DIF supporting this DIF, compute the set of nodes that
+            # share such N-1-DIF. This set will be called the 'neighset' of
+            # the N-1-DIF for the current DIF.
+
+            for node in self.nodes:
+                if dif in node.dif_registrations:
+                    dif_graphs[dif][node] = [] # init for later use
+                    if first == None: # pick any node for later use
+                        first = node
+                    for lower_dif in node.dif_registrations[dif]:
+                        if lower_dif not in neighsets:
+                            neighsets[lower_dif] = []
+                        neighsets[lower_dif].append(node)
+
+            # Build the graph, represented as adjacency list
+            for lower_dif in neighsets:
+                # Each neighset corresponds to a complete (sub)graph.
+                for node1 in neighsets[lower_dif]:
+                    for node2 in neighsets[lower_dif]:
+                        if node1 != node2:
+                            dif_graphs[dif][node1].append((node2, lower_dif))
+
+            if first == None:
+                # This is a shim DIF, nothing to do
+                continue
+
+            print("DIF graphs for %s" % dif)
+            for node in dif_graphs[dif]:
+                for edge in dif_graphs[dif][node]:
+                    print("%s --> %s [%s]" % (node.name, edge[0].name, edge[1]))
+
+            self.enrollments[dif] = []
+
+            if self.enrollment_strategy == 'minimal':
+                # To generate the list of enrollments, we simulate one,
+                # using breadth-first trasversal.
+                enrolled = set([first])
+                frontier = set([first])
+                while len(frontier):
+                    cur = frontier.pop()
+                    for edge in dif_graphs[dif][cur]:
+                        if edge[0] not in enrolled:
+                            enrolled.add(edge[0])
+                            self.enrollments[dif].append({'enrollee': edge[0],
+                                                          'enroller': cur,
+                                                          'lower_dif': edge[1]})
+                            frontier.add(edge[0])
+
+            elif self.enrollment_strategy == 'full-mesh':
+                for cur in dif_graphs[dif]:
+                    for edge in dif_graphs[dif][cur]:
+                        if cur < edge[0]:
+                            self.enrollments[dif].append({'enrollee': cur,
+                                                          'enroller': edge[0],
+                                                          'lower_dif': edge[1]})
+
+            else:
+                # This is a bug
+                assert(False)
+
+            print("Enrollments for %s" % dif)
+            for e in self.enrollments[dif]:
+                print("    %s --> %s through N-1-DIF %s" % \
+                                    (e['enrollee'].name,
+                                     e['enroller'].name,
+                                     e['lower_dif']))
 
-        print(self.dif_ordering)
+    # Examine the nodes, DIFs, registrations and compute the registration
+    # and enrollment order, etc.
+    def generate(self):
+        self.compute_dif_ordering()
+        self.compute_enrollments()
 
     # Realize the experiment, using a testbed-specific setup
     def swap_in(self):