From 99a66d19a58c5eef846905af59fc5c5e10854a01 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Fri, 24 Nov 2023 14:27:02 +0100
Subject: [PATCH] wifi: Allow links blocked for specified reasons to be
 returned by GetLinkIds

---
 src/wifi/model/wifi-mac-queue-scheduler-impl.h | 17 ++++++++++++++---
 src/wifi/model/wifi-mac-queue-scheduler.h      |  6 +++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/wifi/model/wifi-mac-queue-scheduler-impl.h b/src/wifi/model/wifi-mac-queue-scheduler-impl.h
index 0bf92b47a..a28200066 100644
--- a/src/wifi/model/wifi-mac-queue-scheduler-impl.h
+++ b/src/wifi/model/wifi-mac-queue-scheduler-impl.h
@@ -79,7 +79,9 @@ class WifiMacQueueSchedulerImpl : public WifiMacQueueScheduler
                                                 std::optional<uint8_t> linkId,
                                                 const WifiContainerQueueId& prevQueueId) final;
     /** \copydoc ns3::WifiMacQueueScheduler::GetLinkIds */
-    std::list<uint8_t> GetLinkIds(AcIndex ac, Ptr<const WifiMpdu> mpdu) final;
+    std::list<uint8_t> GetLinkIds(AcIndex ac,
+                                  Ptr<const WifiMpdu> mpdu,
+                                  const std::list<WifiQueueBlockedReason>& ignoredReasons) final;
     /** \copydoc ns3::WifiMacQueueScheduler::BlockQueues */
     void BlockQueues(WifiQueueBlockedReason reason,
                      AcIndex ac,
@@ -446,14 +448,23 @@ WifiMacQueueSchedulerImpl<Priority, Compare>::SetPriority(AcIndex ac,
 
 template <class Priority, class Compare>
 std::list<uint8_t>
-WifiMacQueueSchedulerImpl<Priority, Compare>::GetLinkIds(AcIndex ac, Ptr<const WifiMpdu> mpdu)
+WifiMacQueueSchedulerImpl<Priority, Compare>::GetLinkIds(
+    AcIndex ac,
+    Ptr<const WifiMpdu> mpdu,
+    const std::list<WifiQueueBlockedReason>& ignoredReasons)
 {
     auto queueInfoIt = InitQueueInfo(ac, mpdu);
     std::list<uint8_t> linkIds;
 
     // include only links that are not blocked in the returned list
-    for (const auto& [linkId, mask] : queueInfoIt->second.linkIds)
+    for (auto [linkId, mask] : queueInfoIt->second.linkIds)
     {
+        // reset the bits of the mask corresponding to the reasons to ignore
+        for (const auto reason : ignoredReasons)
+        {
+            mask.reset(static_cast<std::size_t>(reason));
+        }
+
         if (mask.none())
         {
             linkIds.emplace_back(linkId);
diff --git a/src/wifi/model/wifi-mac-queue-scheduler.h b/src/wifi/model/wifi-mac-queue-scheduler.h
index 51ccbcc62..3dcc17e6b 100644
--- a/src/wifi/model/wifi-mac-queue-scheduler.h
+++ b/src/wifi/model/wifi-mac-queue-scheduler.h
@@ -135,9 +135,13 @@ class WifiMacQueueScheduler : public Object
      *
      * \param ac the given Access Category
      * \param mpdu the given MPDU
+     * \param ignoredReasons list of reasons for blocking a link that are ignored
      * \return the list of the IDs of the links the given MPDU can be sent over
      */
-    virtual std::list<uint8_t> GetLinkIds(AcIndex ac, Ptr<const WifiMpdu> mpdu) = 0;
+    virtual std::list<uint8_t> GetLinkIds(
+        AcIndex ac,
+        Ptr<const WifiMpdu> mpdu,
+        const std::list<WifiQueueBlockedReason>& ignoredReasons = {}) = 0;
 
     /**
      * Block the given set of links for the container queues of the given types and