From 2de357aac273229ea45f7a7c5a548778688fa4ae Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Wed, 18 Nov 2020 13:16:30 +0100
Subject: [PATCH] wifi: Add WifiTxVector::GetModulationClass for convenience

In case the TxVector is MU, we don't need to specify any STAID
because the modulation class is unique. Getting the modulation
class through the WifiMode requires to specify the STAID instead.
---
 src/wifi/model/wifi-tx-vector.cc | 17 +++++++++++++++++
 src/wifi/model/wifi-tx-vector.h  |  7 +++++++
 2 files changed, 24 insertions(+)

diff --git a/src/wifi/model/wifi-tx-vector.cc b/src/wifi/model/wifi-tx-vector.cc
index a259a7ad4..1c487867b 100644
--- a/src/wifi/model/wifi-tx-vector.cc
+++ b/src/wifi/model/wifi-tx-vector.cc
@@ -119,6 +119,23 @@ WifiTxVector::GetMode (uint16_t staId) const
   return m_mode;
 }
 
+WifiModulationClass
+WifiTxVector::GetModulationClass (void) const
+{
+  NS_ABORT_MSG_IF (!m_modeInitialized, "WifiTxVector mode must be set before using");
+
+  // TODO Replace the following with IsMu when available
+  if (m_preamble == WIFI_PREAMBLE_VHT_MU
+      || m_preamble == WIFI_PREAMBLE_HE_MU
+      || m_preamble == WIFI_PREAMBLE_HE_TB)
+    {
+      NS_ASSERT (!m_muUserInfos.empty ());
+      // all the modes belong to the same modulation class
+      return m_muUserInfos.begin ()->second.mcs.GetModulationClass ();
+    }
+  return m_mode.GetModulationClass ();
+}
+
 uint8_t
 WifiTxVector::GetTxPowerLevel (void) const
 {
diff --git a/src/wifi/model/wifi-tx-vector.h b/src/wifi/model/wifi-tx-vector.h
index 5365f8be4..65dc49c24 100644
--- a/src/wifi/model/wifi-tx-vector.h
+++ b/src/wifi/model/wifi-tx-vector.h
@@ -137,6 +137,13 @@ public:
   */
   void SetMode (WifiMode mode, uint16_t staId);
 
+  /**
+   * Get the modulation class specified by this TXVECTOR.
+   *
+   * \return the Modulation Class specified by this TXVECTOR
+   */
+  WifiModulationClass GetModulationClass (void) const;
+
   /**
    * \returns the transmission power level
    */