Merge ktlim@zizzer:/bk/newmem

into  zamp.eecs.umich.edu:/z/ktlim2/clean/newmem-merge

--HG--
extra : convert_revision : 942c43e2fdd68cde7aaaba5e88a667f80feab162

1 files changed, 15 insertions, 0 deletions

diff --git a/src/python/m5/__init__.py b/src/python/m5/__init__.py
index 579785a46..dc3af7000 100644
--- a/src/python/m5/__init__.py
+++ b/src/python/m5/__init__.py
@@ -213,14 +213,28 @@ atexit.register(cc_main.doExitCleanup)
 # matter since most scripts will probably 'from m5.objects import *'.
 import objects
 
+# This loops until all objects have been fully drained.
 def doDrain(root):
+    all_drained = drain(root)
+    while (not all_drained):
+        all_drained = drain(root)
+
+# Tries to drain all objects.  Draining might not be completed unless
+# all objects return that they are drained on the first call.  This is
+# because as objects drain they may cause other objects to no longer
+# be drained.
+def drain(root):
+    all_drained = False
     drain_event = cc_main.createCountedDrain()
     unready_objects = root.startDrain(drain_event, True)
     # If we've got some objects that can't drain immediately, then simulate
     if unready_objects > 0:
         drain_event.setCount(unready_objects)
         simulate()
+    else:
+        all_drained = True
     cc_main.cleanupCountedDrain(drain_event)
+    return all_drained
 
 def resume(root):
     root.resume()
@@ -236,6 +250,7 @@ def checkpoint(root):
 def restoreCheckpoint(root):
     print "Restoring from checkpoint"
     cc_main.unserializeAll()
+    resume(root)
 
 def changeToAtomic(system):
     if not isinstance(system, objects.Root) and not isinstance(system, System):