3 files changed, 12 insertions, 23 deletions

diff --git a/src/arch/arm/table_walker.cc b/src/arch/arm/table_walker.cc
index 297054131..3e61a4bd6 100644
--- a/src/arch/arm/table_walker.cc
+++ b/src/arch/arm/table_walker.cc
@@ -57,7 +57,7 @@
 using namespace ArmISA;
 
 TableWalker::TableWalker(const Params *p)
-    : MemObject(p), drainManager(NULL),
+    : MemObject(p),
       stage2Mmu(NULL), port(NULL), masterId(Request::invldMasterId),
       isStage2(p->is_stage2), tlb(NULL),
       currState(NULL), pending(false),
@@ -137,17 +137,17 @@ TableWalker::WalkerState::WalkerState() :
 void
 TableWalker::completeDrain()
 {
-    if (drainManager && stateQueues[L1].empty() && stateQueues[L2].empty() &&
+    if (drainState() == DrainState::Draining &&
+        stateQueues[L1].empty() && stateQueues[L2].empty() &&
         pendingQueue.empty()) {
-        setDrainState(DrainState::Drained);
+
         DPRINTF(Drain, "TableWalker done draining, processing drain event\n");
-        drainManager->signalDrainDone();
-        drainManager = NULL;
+        signalDrainDone();
     }
 }
 
-unsigned int
-TableWalker::drain(DrainManager *dm)
+DrainState
+TableWalker::drain()
 {
     bool state_queues_not_empty = false;
 
@@ -159,25 +159,17 @@ TableWalker::drain(DrainManager *dm)
     }
 
     if (state_queues_not_empty || pendingQueue.size()) {
-        drainManager = dm;
-        setDrainState(DrainState::Draining);
         DPRINTF(Drain, "TableWalker not drained\n");
-
-        // return port drain count plus the table walker itself needs to drain
-        return 1;
+        return DrainState::Draining;
     } else {
-        setDrainState(DrainState::Drained);
         DPRINTF(Drain, "TableWalker free, no need to drain\n");
-
-        // table walker is drained, but its ports may still need to be drained
-        return 0;
+        return DrainState::Drained;
     }
 }
 
 void
 TableWalker::drainResume()
 {
-    Drainable::drainResume();
     if (params()->sys->isTimingMode() && currState) {
         delete currState;
         currState = NULL;
diff --git a/src/arch/arm/table_walker.hh b/src/arch/arm/table_walker.hh
index a5327cd95..e973e9a74 100644
--- a/src/arch/arm/table_walker.hh
+++ b/src/arch/arm/table_walker.hh
@@ -819,9 +819,6 @@ class TableWalker : public MemObject
      * currently busy. */
     std::list<WalkerState *> pendingQueue;
 
-    /** If we're draining keep the drain event around until we're drained */
-    DrainManager *drainManager;
-
     /** The MMU to forward second stage look upts to */
     Stage2MMU *stage2Mmu;
 
@@ -894,8 +891,8 @@ class TableWalker : public MemObject
     bool haveLargeAsid64() const { return _haveLargeAsid64; }
     /** Checks if all state is cleared and if so, completes drain */
     void completeDrain();
-    unsigned int drain(DrainManager *dm);
-    virtual void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    virtual void drainResume() M5_ATTR_OVERRIDE;
 
     virtual BaseMasterPort& getMasterPort(const std::string &if_name,
                                           PortID idx = InvalidPortID);
diff --git a/src/arch/arm/tlb.hh b/src/arch/arm/tlb.hh
index 28b99a8e0..63707dba2 100644
--- a/src/arch/arm/tlb.hh
+++ b/src/arch/arm/tlb.hh
@@ -284,7 +284,7 @@ class TLB : public BaseTLB
             bool callFromS2);
     Fault finalizePhysical(RequestPtr req, ThreadContext *tc, Mode mode) const;
 
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
     // Checkpointing
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;