2 files changed, 22 insertions, 15 deletions

diff --git a/src/mem/ruby/network/garnet/fixed-pipeline/NetworkInterface_d.cc b/src/mem/ruby/network/garnet/fixed-pipeline/NetworkInterface_d.cc
index 7384cc6a7..47ded231c 100644
--- a/src/mem/ruby/network/garnet/fixed-pipeline/NetworkInterface_d.cc
+++ b/src/mem/ruby/network/garnet/fixed-pipeline/NetworkInterface_d.cc
@@ -106,20 +106,23 @@ NetworkInterface_d::addOutPort(NetworkLink_d *out_link,
 }
 
 void
-NetworkInterface_d::addNode(map<int, MessageBuffer *>& in,
-                            map<int, MessageBuffer *>& out)
+NetworkInterface_d::addNode(vector<MessageBuffer *>& in,
+                            vector<MessageBuffer *>& out)
 {
     inNode_ptr = in;
     outNode_ptr = out;
 
     for (auto& it : in) {
-        // the protocol injects messages into the NI
-        it.second->setConsumer(this);
-        it.second->setReceiver(this);
+        if (it != nullptr) {
+            it->setConsumer(this);
+            it->setReceiver(this);
+        }
     }
 
     for (auto& it : out) {
-        it.second->setSender(this);
+        if (it != nullptr) {
+            it->setSender(this);
+        }
     }
 }
 
@@ -223,9 +226,11 @@ NetworkInterface_d::wakeup()
 
     // Checking for messages coming from the protocol
     // can pick up a message/cycle for each virtual net
-    for (auto it = inNode_ptr.begin(); it != inNode_ptr.end(); ++it) {
-        int vnet = (*it).first;
-        MessageBuffer *b = (*it).second;
+    for (int vnet = 0; vnet < inNode_ptr.size(); ++vnet) {
+        MessageBuffer *b = inNode_ptr[vnet];
+        if (b == nullptr) {
+            continue;
+        }
 
         while (b->isReady()) { // Is there a message waiting
             msg_ptr = b->peekMsgPtr();
@@ -355,9 +360,11 @@ void
 NetworkInterface_d::checkReschedule()
 {
     for (const auto& it : inNode_ptr) {
-        MessageBuffer *b = it.second;
+        if (it == nullptr) {
+            continue;
+        }
 
-        while (b->isReady()) { // Is there a message waiting
+        while (it->isReady()) { // Is there a message waiting
             scheduleEvent(Cycles(1));
             return;
         }
diff --git a/src/mem/ruby/network/garnet/fixed-pipeline/NetworkInterface_d.hh b/src/mem/ruby/network/garnet/fixed-pipeline/NetworkInterface_d.hh
index 2494d05d1..0cb928d82 100644
--- a/src/mem/ruby/network/garnet/fixed-pipeline/NetworkInterface_d.hh
+++ b/src/mem/ruby/network/garnet/fixed-pipeline/NetworkInterface_d.hh
@@ -60,8 +60,8 @@ class NetworkInterface_d : public ClockedObject, public Consumer
     void addOutPort(NetworkLink_d *out_link, CreditLink_d *credit_link);
 
     void wakeup();
-    void addNode(std::map<int, MessageBuffer *> &inNode,
-                 std::map<int, MessageBuffer *> &outNode);
+    void addNode(std::vector<MessageBuffer *> &inNode,
+                 std::vector<MessageBuffer *> &outNode);
 
     void print(std::ostream& out) const;
     int get_vnet(int vc);
@@ -90,9 +90,9 @@ class NetworkInterface_d : public ClockedObject, public Consumer
     std::vector<Cycles> m_ni_enqueue_time;
 
     // The Message buffers that takes messages from the protocol
-    std::map<int, MessageBuffer *> inNode_ptr;
+    std::vector<MessageBuffer *> inNode_ptr;
     // The Message buffers that provides messages to the protocol
-    std::map<int, MessageBuffer *> outNode_ptr;
+    std::vector<MessageBuffer *> outNode_ptr;
 
     bool flitisizeMessage(MsgPtr msg_ptr, int vnet);
     int calculateVC(int vnet);