Change to PQ

1 files changed, 18 insertions, 10 deletions

diff --git a/unit_trace/trace_reader.py b/unit_trace/trace_reader.py
index c895cf2..dc7cb11 100644
--- a/unit_trace/trace_reader.py
+++ b/unit_trace/trace_reader.py
@@ -27,6 +27,7 @@
 
 import struct
 import sys
+import heapq
 
 
 ###############################################################################
@@ -59,7 +60,8 @@ def trace_reader(files, buffsize):
         file_iter = _get_file_iter(file)
         file_iters.append(file_iter)
         try:
-            file_iter_buff.append([file_iter.next()])
+            rec = file_iter.next()
+            file_iter_buff.append([(rec.when, rec)])
         # What if there isn't a single valid record in a trace file?
         # file_iter.next() will raise a StopIteration that we need to catch
         except:
@@ -75,11 +77,13 @@ def trace_reader(files, buffsize):
     for x in range(0,len(file_iter_buff)):
         try:
             for y in range(0,buffsize):
-                file_iter_buff[x].append(file_iters[x].next())
+                rec = file_iters[x].next()
+                heapq.heappush(file_iter_buff[x], (rec.when, rec))
         except StopIteration:
             pass
-    for x in range(0,len(file_iter_buff)):
-        file_iter_buff[x] = sorted(file_iter_buff[x],key=lambda rec: rec.when)
+    #Obsolete when using heap
+    #for x in range(0,len(file_iter_buff)):
+    #    file_iter_buff[x] = sorted(file_iter_buff[x],key=lambda rec: rec.when)
 
     # Remember the time of the last record. This way, we can make sure records
     # truly are produced in monotonically increasing order by time and terminate
@@ -95,19 +99,23 @@ def trace_reader(files, buffsize):
         earliest = -1
         buff_to_refill = -1
         for x in range(0,len(file_iter_buff)):
-            if earliest==-1 or file_iter_buff[x][0].when < earliest.when:
-                earliest = file_iter_buff[x][0]
+            if earliest==-1 or file_iter_buff[x][0][1].when < earliest.when:
+                earliest = file_iter_buff[x][0][1]
                 buff_to_refill = x
 
         # Take it out of the buffer
-        del file_iter_buff[buff_to_refill][0]
+        #del file_iter_buff[buff_to_refill][0]
+        heapq.heappop(file_iter_buff[buff_to_refill])
 
         # Try to append a new record to the buffer (if there is another) and
         #     then keep the buffer sorted
         try:
-            file_iter_buff[buff_to_refill].append(file_iters[buff_to_refill].next())
-            file_iter_buff[buff_to_refill] = sorted(file_iter_buff[buff_to_refill],
-                key=lambda rec: rec.when)
+            rec = file_iters[buff_to_refill].next()
+            rec_tup = (rec.when, rec)
+            heapq.heappush(file_iter_buff[buff_to_refill], rec_tup)
+            #file_iter_buff[buff_to_refill].append(file_iters[buff_to_refill].next())
+            #file_iter_buff[buff_to_refill] = sorted(file_iter_buff[buff_to_refill],
+            #    key=lambda rec: rec.when)
 
         # If there aren't any more records, fine. Unless the buffer is also empty.
         # If that is the case, delete the buffer.