From e6709c8beec58a8d53f5c9e516bf4e686cf15125 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Fri, 4 Mar 2022 13:37:55 +0100
Subject: [PATCH] wifi: Align behavior of WifiMac::GetTxopQueue to that of
 WifiMac::GetQosTxop

---
 src/wifi/model/wifi-mac.cc | 9 ++-------
 src/wifi/model/wifi-mac.h  | 6 ++++--
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/src/wifi/model/wifi-mac.cc b/src/wifi/model/wifi-mac.cc
index aca373cb3..8a1bb517e 100644
--- a/src/wifi/model/wifi-mac.cc
+++ b/src/wifi/model/wifi-mac.cc
@@ -492,13 +492,8 @@ WifiMac::GetBKQueue () const
 Ptr<WifiMacQueue>
 WifiMac::GetTxopQueue (AcIndex ac) const
 {
-  if (ac == AC_BE_NQOS)
-    {
-      NS_ASSERT (m_txop != nullptr);
-      return m_txop->GetWifiMacQueue ();
-    }
-  NS_ASSERT (ac == AC_BE || ac == AC_BK || ac == AC_VI || ac == AC_VO);
-  return m_edca.find (ac)->second->GetWifiMacQueue ();
+  Ptr<Txop> txop = (ac == AC_BE_NQOS ? m_txop : StaticCast<Txop> (GetQosTxop (ac)));
+  return (txop != nullptr ? txop->GetWifiMacQueue () : nullptr);
 }
 
 void
diff --git a/src/wifi/model/wifi-mac.h b/src/wifi/model/wifi-mac.h
index d62ba8fa8..b95f5b492 100644
--- a/src/wifi/model/wifi-mac.h
+++ b/src/wifi/model/wifi-mac.h
@@ -134,10 +134,12 @@ public:
    */
   Ptr<QosTxop> GetQosTxop (uint8_t tid) const;
   /**
-   * Get the wifi MAC queue of the (Qos)Txop associated with the given AC.
+   * Get the wifi MAC queue of the (Qos)Txop associated with the given AC,
+   * if such (Qos)Txop is installed, or a null pointer, otherwise.
    *
    * \param ac the given Access Category
-   * \return the wifi MAC queue of the (Qos)Txop associated with the given AC
+   * \return the wifi MAC queue of the (Qos)Txop associated with the given AC,
+   *         if such (Qos)Txop is installed, or a null pointer, otherwise
    */
   virtual Ptr<WifiMacQueue> GetTxopQueue (AcIndex ac) const;