From 0901e95ce479eefd7e2600f699f614410f735864 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Thu, 18 May 2023 19:57:08 +0200
Subject: [PATCH] wifi: Add Txop support for link swapping

---
 src/wifi/model/qos-txop.cc |  8 ++--
 src/wifi/model/txop.cc     | 87 ++++++++++++++++++++++++--------------
 src/wifi/model/txop.h      | 22 +++++++---
 src/wifi/model/wifi-mac.cc | 13 ++++++
 4 files changed, 88 insertions(+), 42 deletions(-)

diff --git a/src/wifi/model/qos-txop.cc b/src/wifi/model/qos-txop.cc
index fb5233e93..8d02ea937 100644
--- a/src/wifi/model/qos-txop.cc
+++ b/src/wifi/model/qos-txop.cc
@@ -640,9 +640,9 @@ QosTxop::GotAddBaResponse(const MgtAddBaResponseHeader& respHdr, Mac48Address re
         m_baManager->NotifyOriginatorAgreementRejected(recipient, tid);
     }
 
-    for (uint8_t linkId = 0; linkId < GetNLinks(); linkId++)
+    for (const auto& [id, link] : GetLinks())
     {
-        StartAccessIfNeeded(linkId);
+        StartAccessIfNeeded(id);
     }
 }
 
@@ -662,9 +662,9 @@ QosTxop::NotifyOriginatorAgreementNoReply(const Mac48Address& recipient, uint8_t
     m_baManager->NotifyOriginatorAgreementNoReply(recipient, tid);
     // the recipient has been "unblocked" and transmissions can resume using normal
     // acknowledgment, hence start access (if needed) on all the links
-    for (uint8_t linkId = 0; linkId < GetNLinks(); linkId++)
+    for (const auto& [id, link] : GetLinks())
     {
-        StartAccessIfNeeded(linkId);
+        StartAccessIfNeeded(id);
     }
 }
 
diff --git a/src/wifi/model/txop.cc b/src/wifi/model/txop.cc
index 011b97deb..1a645db79 100644
--- a/src/wifi/model/txop.cc
+++ b/src/wifi/model/txop.cc
@@ -169,15 +169,33 @@ Txop::CreateLinkEntity() const
 Txop::LinkEntity&
 Txop::GetLink(uint8_t linkId) const
 {
-    NS_ASSERT(linkId < m_links.size());
-    NS_ASSERT(m_links.at(linkId)); // check that the pointer owns an object
-    return *m_links.at(linkId);
+    auto it = m_links.find(linkId);
+    NS_ASSERT(it != m_links.cend());
+    NS_ASSERT(it->second); // check that the pointer owns an object
+    return *it->second;
 }
 
-uint8_t
-Txop::GetNLinks() const
+const std::map<uint8_t, std::unique_ptr<Txop::LinkEntity>>&
+Txop::GetLinks() const
 {
-    return m_links.size();
+    return m_links;
+}
+
+void
+Txop::SwapLinks(std::map<uint8_t, uint8_t> links)
+{
+    NS_LOG_FUNCTION(this);
+
+    decltype(m_links) tmp;
+    tmp.swap(m_links); // move all links to temporary map
+    for (const auto& [from, to] : links)
+    {
+        auto nh = tmp.extract(from);
+        nh.key() = to;
+        m_links.insert(std::move(nh));
+    }
+    // move links remaining in tmp to m_links
+    m_links.merge(tmp);
 }
 
 void
@@ -192,12 +210,9 @@ Txop::SetWifiMac(const Ptr<WifiMac> mac)
 {
     NS_LOG_FUNCTION(this << mac);
     m_mac = mac;
-    m_links.resize(m_mac->GetNLinks());
-    uint8_t linkId = 0;
-    for (auto& link : m_links)
+    for (const auto linkId : m_mac->GetLinkIds())
     {
-        link = CreateLinkEntity();
-        link->id = linkId++;
+        m_links.emplace(linkId, CreateLinkEntity());
     }
 }
 
@@ -229,9 +244,10 @@ void
 Txop::SetMinCws(std::vector<uint32_t> minCws)
 {
     NS_ABORT_IF(minCws.size() != m_links.size());
-    for (std::size_t linkId = 0; linkId < minCws.size(); linkId++)
+    std::size_t i = 0;
+    for (const auto& [id, link] : m_links)
     {
-        SetMinCw(minCws[linkId], linkId);
+        SetMinCw(minCws[i++], id);
     }
 }
 
@@ -258,9 +274,10 @@ void
 Txop::SetMaxCws(std::vector<uint32_t> maxCws)
 {
     NS_ABORT_IF(maxCws.size() != m_links.size());
-    for (std::size_t linkId = 0; linkId < maxCws.size(); linkId++)
+    std::size_t i = 0;
+    for (const auto& [id, link] : m_links)
     {
-        SetMaxCw(maxCws[linkId], linkId);
+        SetMaxCw(maxCws[i++], id);
     }
 }
 
@@ -355,9 +372,10 @@ void
 Txop::SetAifsns(std::vector<uint8_t> aifsns)
 {
     NS_ABORT_IF(aifsns.size() != m_links.size());
-    for (std::size_t linkId = 0; linkId < aifsns.size(); linkId++)
+    std::size_t i = 0;
+    for (const auto& [id, link] : m_links)
     {
-        SetAifsn(aifsns[linkId], linkId);
+        SetAifsn(aifsns[i++], id);
     }
 }
 
@@ -381,9 +399,10 @@ Txop::SetTxopLimits(const std::vector<Time>& txopLimits)
                     "The size of the given vector (" << txopLimits.size()
                                                      << ") does not match the number of links ("
                                                      << m_links.size() << ")");
-    for (std::size_t linkId = 0; linkId < txopLimits.size(); linkId++)
+    std::size_t i = 0;
+    for (const auto& [id, link] : m_links)
     {
-        SetTxopLimit(txopLimits[linkId], linkId);
+        SetTxopLimit(txopLimits[i++], id);
     }
 }
 
@@ -406,9 +425,10 @@ std::vector<uint32_t>
 Txop::GetMinCws() const
 {
     std::vector<uint32_t> ret;
-    for (std::size_t linkId = 0; linkId < m_links.size(); linkId++)
+    ret.reserve(m_links.size());
+    for (const auto& [id, link] : m_links)
     {
-        ret.push_back(GetMinCw(linkId));
+        ret.push_back(link->cwMin);
     }
     return ret;
 }
@@ -429,9 +449,10 @@ std::vector<uint32_t>
 Txop::GetMaxCws() const
 {
     std::vector<uint32_t> ret;
-    for (std::size_t linkId = 0; linkId < m_links.size(); linkId++)
+    ret.reserve(m_links.size());
+    for (const auto& [id, link] : m_links)
     {
-        ret.push_back(GetMaxCw(linkId));
+        ret.push_back(link->cwMax);
     }
     return ret;
 }
@@ -452,9 +473,10 @@ std::vector<uint8_t>
 Txop::GetAifsns() const
 {
     std::vector<uint8_t> ret;
-    for (std::size_t linkId = 0; linkId < m_links.size(); linkId++)
+    ret.reserve(m_links.size());
+    for (const auto& [id, link] : m_links)
     {
-        ret.push_back(GetAifsn(linkId));
+        ret.push_back(link->aifsn);
     }
     return ret;
 }
@@ -475,9 +497,10 @@ std::vector<Time>
 Txop::GetTxopLimits() const
 {
     std::vector<Time> ret;
-    for (std::size_t linkId = 0; linkId < m_links.size(); linkId++)
+    ret.reserve(m_links.size());
+    for (const auto& [id, link] : m_links)
     {
-        ret.push_back(GetTxopLimit(linkId));
+        ret.push_back(link->txopLimit);
     }
     return ret;
 }
@@ -548,10 +571,10 @@ void
 Txop::DoInitialize()
 {
     NS_LOG_FUNCTION(this);
-    for (std::size_t linkId = 0; linkId < m_links.size(); linkId++)
+    for (const auto& [id, link] : m_links)
     {
-        ResetCw(linkId);
-        GenerateBackoff(linkId);
+        ResetCw(id);
+        GenerateBackoff(id);
     }
 }
 
@@ -630,9 +653,9 @@ void
 Txop::NotifyOn()
 {
     NS_LOG_FUNCTION(this);
-    for (std::size_t linkId = 0; linkId < m_links.size(); linkId++)
+    for (const auto& [id, link] : m_links)
     {
-        StartAccessIfNeeded(linkId);
+        StartAccessIfNeeded(id);
     }
 }
 
diff --git a/src/wifi/model/txop.h b/src/wifi/model/txop.h
index 39f5be835..d52c2c601 100644
--- a/src/wifi/model/txop.h
+++ b/src/wifi/model/txop.h
@@ -26,6 +26,7 @@
 #include "ns3/object.h"
 #include "ns3/traced-value.h"
 
+#include <map>
 #include <memory>
 #include <vector>
 
@@ -408,6 +409,16 @@ class Txop : public Object
      */
     virtual bool HasFramesToTransmit(uint8_t linkId);
 
+    /**
+     * Swap the links based on the information included in the given map. This method
+     * is normally called by the WifiMac of a non-AP MLD upon completing ML setup to have
+     * its link IDs match AP MLD's link IDs.
+     *
+     * \param links a set of pairs (from, to) each mapping a current link ID to the
+     *              link ID it has to become (i.e., link 'from' becomes link 'to')
+     */
+    void SwapLinks(std::map<uint8_t, uint8_t> links);
+
   protected:
     ///< ChannelAccessManager associated class
     friend class ChannelAccessManager;
@@ -483,7 +494,6 @@ class Txop : public Object
         /// Destructor (a virtual method is needed to make this struct polymorphic)
         virtual ~LinkEntity() = default;
 
-        uint8_t id{0};                             //!< Link ID (starting at 0)
         uint32_t backoffSlots{0};                  //!< the number of backoff slots
         Time backoffStart{0};                      /**< the backoffStart variable is used to keep
                                                         track of the time at which a backoff was
@@ -504,12 +514,11 @@ class Txop : public Object
      * \return a reference to the link associated with the given ID
      */
     LinkEntity& GetLink(uint8_t linkId) const;
+
     /**
-     * Get the number of links.
-     *
-     * \return the number of links
+     * \return a const reference to the map of link entities
      */
-    uint8_t GetNLinks() const;
+    const std::map<uint8_t, std::unique_ptr<LinkEntity>>& GetLinks() const;
 
     DroppedMpdu m_droppedMpduCallback; //!< the dropped MPDU callback
     Ptr<WifiMacQueue> m_queue;         //!< the wifi MAC queue
@@ -533,7 +542,8 @@ class Txop : public Object
      */
     virtual std::unique_ptr<LinkEntity> CreateLinkEntity() const;
 
-    std::vector<std::unique_ptr<LinkEntity>> m_links; //!< vector of LinkEntity objects
+    std::map<uint8_t, std::unique_ptr<LinkEntity>>
+        m_links; //!< ID-indexed map of LinkEntity objects
 };
 
 } // namespace ns3
diff --git a/src/wifi/model/wifi-mac.cc b/src/wifi/model/wifi-mac.cc
index 5b76b159f..2ec91a9d3 100644
--- a/src/wifi/model/wifi-mac.cc
+++ b/src/wifi/model/wifi-mac.cc
@@ -996,6 +996,8 @@ WifiMac::SwapLinks(std::map<uint8_t, uint8_t> links)
 {
     NS_LOG_FUNCTION(this);
 
+    std::map<uint8_t, uint8_t> actualPairs;
+
     while (!links.empty())
     {
         auto from = links.cbegin()->first;
@@ -1018,6 +1020,7 @@ WifiMac::SwapLinks(std::map<uint8_t, uint8_t> links)
             auto [it, inserted] =
                 m_links.emplace(to, nullptr); // insert an element with key to if not present
             m_links[to].swap(linkToMove);     // to is the link to move now
+            actualPairs.emplace(from, to);
             UpdateLinkId(to);
             links.erase(from);
             if (!linkToMove)
@@ -1034,6 +1037,7 @@ WifiMac::SwapLinks(std::map<uint8_t, uint8_t> links)
             {
                 // no new position specified for 'to', use the current empty cell
                 m_links[empty].swap(linkToMove);
+                actualPairs.emplace(to, empty);
                 break;
             }
 
@@ -1047,6 +1051,15 @@ WifiMac::SwapLinks(std::map<uint8_t, uint8_t> links)
     {
         m_linkIds.insert(id);
     }
+
+    if (m_txop)
+    {
+        m_txop->SwapLinks(actualPairs);
+    }
+    for (auto& [ac, edca] : m_edca)
+    {
+        edca->SwapLinks(actualPairs);
+    }
 }
 
 void