4 files changed, 18 insertions, 31 deletions

diff --git a/src/cpu/minor/cpu.cc b/src/cpu/minor/cpu.cc
index ac1a18bf9..a93d0037d 100644
--- a/src/cpu/minor/cpu.cc
+++ b/src/cpu/minor/cpu.cc
@@ -47,8 +47,7 @@
 #include "debug/Quiesce.hh"
 
 MinorCPU::MinorCPU(MinorCPUParams *params) :
-    BaseCPU(params),
-    drainManager(NULL)
+    BaseCPU(params)
 {
     /* This is only written for one thread at the moment */
     Minor::MinorThread *thread;
@@ -194,39 +193,33 @@ MinorCPU::startup()
     activateContext(0);
 }
 
-unsigned int
-MinorCPU::drain(DrainManager *drain_manager)
+DrainState
+MinorCPU::drain()
 {
     DPRINTF(Drain, "MinorCPU drain\n");
 
-    drainManager = drain_manager;
-
     /* Need to suspend all threads and wait for Execute to idle.
      * Tell Fetch1 not to fetch */
-    unsigned int ret = pipeline->drain(drain_manager);
-
-    if (ret == 0)
+    if (pipeline->drain()) {
         DPRINTF(Drain, "MinorCPU drained\n");
-    else
+        return DrainState::Drained;
+    } else {
         DPRINTF(Drain, "MinorCPU not finished draining\n");
-
-    return ret;
+        return DrainState::Draining;
+    }
 }
 
 void
 MinorCPU::signalDrainDone()
 {
     DPRINTF(Drain, "MinorCPU drain done\n");
-    setDrainState(DrainState::Drained);
-    drainManager->signalDrainDone();
-    drainManager = NULL;
+    signalDrainDone();
 }
 
 void
 MinorCPU::drainResume()
 {
-    assert(getDrainState() == DrainState::Drained ||
-        getDrainState() == DrainState::Running);
+    assert(drainState() == DrainState::Drained);
 
     if (switchedOut()) {
         DPRINTF(Drain, "drainResume while switched out.  Ignoring\n");
@@ -242,8 +235,6 @@ MinorCPU::drainResume()
 
     wakeup();
     pipeline->drainResume();
-
-    setDrainState(DrainState::Running);
 }
 
 void
diff --git a/src/cpu/minor/cpu.hh b/src/cpu/minor/cpu.hh
index fba54b515..2e877d786 100644
--- a/src/cpu/minor/cpu.hh
+++ b/src/cpu/minor/cpu.hh
@@ -112,10 +112,6 @@ class MinorCPU : public BaseCPU
         virtual void recvTimingSnoopReq(PacketPtr pkt) { }
     };
 
-    /** The DrainManager passed into drain that needs be signalled when
-     *  draining is complete */
-    DrainManager *drainManager;
-
   protected:
      /** Return a reference to the data port. */
     MasterPort &getDataPort();
@@ -155,10 +151,10 @@ class MinorCPU : public BaseCPU
     void unserialize(CheckpointIn &cp);
 
     /** Drain interface */
-    unsigned int drain(DrainManager *drain_manager);
-    void drainResume();
-    /** Signal from Pipeline that MinorCPU should signal the DrainManager
-     *  that a drain is complete and set its drainState */
+    DrainState drain() M5_ATTR_OVERRIDE;
+    void drainResume() M5_ATTR_OVERRIDE;
+    /** Signal from Pipeline that MinorCPU should signal that a drain
+     *  is complete and set its drainState */
     void signalDrainDone();
     void memWriteback();
 
diff --git a/src/cpu/minor/pipeline.cc b/src/cpu/minor/pipeline.cc
index 9d802234b..39b7f31f9 100644
--- a/src/cpu/minor/pipeline.cc
+++ b/src/cpu/minor/pipeline.cc
@@ -192,8 +192,8 @@ Pipeline::wakeupFetch()
     execute.wakeupFetch();
 }
 
-unsigned int
-Pipeline::drain(DrainManager *manager)
+bool
+Pipeline::drain()
 {
     DPRINTF(MinorCPU, "Draining pipeline by halting inst fetches. "
         " Execution should drain naturally\n");
@@ -205,7 +205,7 @@ Pipeline::drain(DrainManager *manager)
     bool drained = isDrained();
     needToSignalDrained = !drained;
 
-    return (drained ? 0 : 1);
+    return drained;
 }
 
 void
diff --git a/src/cpu/minor/pipeline.hh b/src/cpu/minor/pipeline.hh
index 355a3c6c2..bf2071b02 100644
--- a/src/cpu/minor/pipeline.hh
+++ b/src/cpu/minor/pipeline.hh
@@ -115,7 +115,7 @@ class Pipeline : public Ticked
     void wakeupFetch();
 
     /** Try to drain the CPU */
-    unsigned int drain(DrainManager *manager);
+    bool drain();
 
     void drainResume();