1 files changed, 9 insertions, 17 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;