4 files changed, 27 insertions, 16 deletions

diff --git a/src/python/m5/SimObject.py b/src/python/m5/SimObject.py
index c01db2a80..3aea55f5f 100644
--- a/src/python/m5/SimObject.py
+++ b/src/python/m5/SimObject.py
@@ -123,7 +123,8 @@ class MetaSimObject(type):
                       'cxx_class' : str,
                       'cxx_type' : str,
                       'cxx_header' : str,
-                      'type' : str }
+                      'type' : str,
+                      'cxx_bases' : list }
     # Attributes that can be set any time
     keywords = { 'check' : FunctionType }
 
@@ -148,6 +149,8 @@ class MetaSimObject(type):
                 value_dict[key] = val
         if 'abstract' not in value_dict:
             value_dict['abstract'] = False
+        if 'cxx_bases' not in value_dict:
+            value_dict['cxx_bases'] = []
         cls_dict['_value_dict'] = value_dict
         cls = super(MetaSimObject, mcls).__new__(mcls, name, bases, cls_dict)
         if 'type' in value_dict:
@@ -414,6 +417,7 @@ class MetaSimObject(type):
         code('%module(package="m5.internal") param_$cls')
         code()
         code('%{')
+        code('#include "sim/sim_object.hh"')
         code('#include "params/$cls.hh"')
         for param in params:
             param.cxx_predecls(code)
@@ -458,7 +462,17 @@ using std::ptrdiff_t;
         code('%nodefault $classname;')
         code('class $classname')
         if cls._base:
-            code('    : public ${{cls._base.cxx_class}}')
+            bases = [ cls._base.cxx_class ] + cls.cxx_bases
+        else:
+            bases = cls.cxx_bases
+        base_first = True
+        for base in bases:
+            if base_first:
+                code('    : public ${{base}}')
+                base_first = False
+            else:
+                code('    , public ${{base}}')
+
         code('{')
         code('  public:')
         cls.export_methods(code)
@@ -581,30 +595,25 @@ class SimObject(object):
     abstract = True
     cxx_header = "sim/sim_object.hh"
 
+    cxx_bases = [ "Drainable" ]
+
     @classmethod
     def export_method_swig_predecls(cls, code):
         code('''
 %include <std_string.i>
+
+%import "python/swig/drain.i"
 ''')
 
     @classmethod
     def export_methods(cls, code):
         code('''
-    enum State {
-      Running,
-      Draining,
-      Drained
-    };
-
     void init();
     void loadState(Checkpoint *cp);
     void initState();
     void regStats();
     void resetStats();
     void startup();
-
-    unsigned int drain(Event *drain_event);
-    void resume();
 ''')
 
     # Initialize new instance.  For objects with SimObject-valued
diff --git a/src/python/m5/__init__.py b/src/python/m5/__init__.py
index 930609b6b..dc6c5a923 100644
--- a/src/python/m5/__init__.py
+++ b/src/python/m5/__init__.py
@@ -51,3 +51,4 @@ if internal:
     from event import *
     from main import main
     from simulate import *
+
diff --git a/src/python/m5/internal/__init__.py b/src/python/m5/internal/__init__.py
index ca09ab468..30090549a 100644
--- a/src/python/m5/internal/__init__.py
+++ b/src/python/m5/internal/__init__.py
@@ -31,3 +31,4 @@ import debug
 import event
 import stats
 import trace
+from drain import DrainManager, Drainable
diff --git a/src/python/m5/simulate.py b/src/python/m5/simulate.py
index 89b6b1e9d..df30f654a 100644
--- a/src/python/m5/simulate.py
+++ b/src/python/m5/simulate.py
@@ -169,19 +169,19 @@ def doDrain(root):
 # be drained.
 def drain(root):
     all_drained = False
-    drain_event = internal.event.createCountedDrain()
-    unready_objs = sum(obj.drain(drain_event) for obj in root.descendants())
+    dm = internal.drain.createDrainManager()
+    unready_objs = sum(obj.drain(dm) for obj in root.descendants())
     # If we've got some objects that can't drain immediately, then simulate
     if unready_objs > 0:
-        drain_event.setCount(unready_objs)
+        dm.setCount(unready_objs)
         simulate()
     else:
         all_drained = True
-    internal.event.cleanupCountedDrain(drain_event)
+    internal.drain.cleanupDrainManager(dm)
     return all_drained
 
 def resume(root):
-    for obj in root.descendants(): obj.resume()
+    for obj in root.descendants(): obj.drainResume()
 
 def checkpoint(dir):
     root = objects.Root.getInstance()