1 files changed, 23 insertions, 17 deletions

diff --git a/viz/schedule.py b/viz/schedule.py
index e525b17..ae7284b 100644
--- a/viz/schedule.py
+++ b/viz/schedule.py
@@ -35,25 +35,29 @@ class TimeSlotArray(object):
         for type in self.list_sizes:
             num = self.list_sizes[type]
             self.array[type] = []
-            for i in range(0, int((end - start) // self.time_per_maj + 1)):
+            for j in range(0, num):
                 # for each slot in the array, we need a list of all events under this type
                 # (for example, a list of all events that occur in this time slot, indexed
                 # by task).
-                self.array[type].append([])
-                for j in range(0, num):
-                    self.array[type][i].append(dict(zip(EVENT_LIST, \
-                                    [[] for j in range(0, len(EVENT_LIST))])))
+                for i in range(0, num):
+                    self.array[type].append(dict(zip(EVENT_LIST, \
+                                    [{} for j in range(0, len(EVENT_LIST))])))
     
     def get_time_slot(self, time):
         return int((time - self.start) // self.time_per_maj)
-            
+    
+    def _put_event_in_slot(self, list_type, no, klass, slot, event):
+        if slot not in self.array[list_type][no][klass]:
+            self.array[list_type][no][klass][slot] = []
+        self.array[list_type][no][klass][slot].append(event)
+                
     def add_event_to_time_slot(self, event):
         task_no = event.get_job().get_task().get_task_no()
         cpu = event.get_cpu()
         time_slot = self.get_time_slot(event.get_time())
         
-        self.array[TimeSlotArray.TASK_LIST][time_slot][task_no][event.__class__].append(event)
-        self.array[TimeSlotArray.CPU_LIST][time_slot][cpu][event.__class__].append(event)
+        self._put_event_in_slot(TimeSlotArray.TASK_LIST, task_no, event.__class__, time_slot, event)
+        self._put_event_in_slot(TimeSlotArray.CPU_LIST, cpu, event.__class__, time_slot, event)
             
         span_events = { SwitchAwayEvent : IsRunningDummy, InversionEndEvent : InversionDummy}
         
@@ -64,8 +68,9 @@ class TimeSlotArray(object):
                 for slot in range(start_slot + 1, end_slot):
                     dummy = span_events[span_event](task_no, cpu)
                     dummy.corresp_start_event = event.corresp_start_event
-                    self.array[TimeSlotArray.TASK_LIST][slot][task_no][dummy.__class__].append(dummy)
-                    self.array[TimeSlotArray.CPU_LIST][slot][cpu][dummy.__class__].append(dummy)
+                    
+                    self._put_event_in_slot(TimeSlotArray.TASK_LIST, task_no, dummy.__class__, slot, dummy)
+                    self._put_event_in_slot(TimeSlotArray.CPU_LIST, cpu, dummy.__class__, slot, dummy)
             
     def iter_over_period(self, start, end, start_no, end_no, list_type, event_types):
         if self.array is None:
@@ -76,17 +81,18 @@ class TimeSlotArray(object):
         if start_no > end_no:
             raise ValueError('start no should be less than end no')
         
-        start_slot = max(0, self.get_time_slot(start))
-        end_slot = min(len(self.array[list_type]), self.get_time_slot(end) + 2)
+        start_slot = self.get_time_slot(start)
+        end_slot = self.get_time_slot(end) + 2
         
         start_no = max(0, start_no)
         end_no = min(self.list_sizes[list_type] - 1, end_no)
         
-        for slot in range(start_slot, end_slot):
-            for no in range(start_no, end_no + 1):
-                for type in event_types:
-                    for event in self.array[list_type][slot][no][type]:
-                        yield event
+        for no in range(start_no, end_no + 1):
+            for type in event_types:
+                for slot in range(start_slot, end_slot):
+                    if slot in self.array[list_type][no][type]:
+                        for event in self.array[list_type][no][type][slot]:
+                            yield event
     
 class Schedule(object):
     """The total schedule (task system), consisting of a certain number of