wifi: Handle transmission of HE TB PPDUs

src/wifi/model/wifi-phy.cc

@@ -3154,11 +3154,11 @@ WifiPhy::StartReceivePayload (Ptr<Event> event)
       Ptr<const WifiPsdu> psdu = GetAddressedPsduInPpdu (ppdu);
       if (psdu)
         {
-          WifiMode txMode = txVector.GetMode (GetStaId ());
+          uint16_t staId = GetStaId (ppdu);
+          WifiMode txMode = txVector.GetMode (staId);
           uint8_t nss = txVector.GetNssMax();
           if (txVector.GetPreambleType () == WIFI_PREAMBLE_HE_MU)
             {
-              uint16_t staId = GetStaId ();
               for (const auto & info : txVector.GetHeMuUserInfoMap ())
                 {
                   if (info.first == staId)
@@ -3279,7 +3279,7 @@ WifiPhy::EndOfMpdu (Ptr<Event> event, Ptr<const WifiPsdu> psdu, size_t mpduIndex
 {
   NS_LOG_FUNCTION (this << *event << mpduIndex << relativeStart << mpduDuration);
   Ptr<const WifiPpdu> ppdu = event->GetPpdu ();
-  uint16_t staId = GetStaId ();
+  uint16_t staId = GetStaId (ppdu);
   WifiTxVector txVector = event->GetTxVector ();
   uint16_t channelWidth = std::min (GetChannelWidth (), txVector.GetChannelWidth ());
   double snr = m_interference.CalculateSnr (event, channelWidth, txVector.GetNss (staId), GetBand (channelWidth));
@@ -3302,11 +3302,15 @@ WifiPhy::EndReceive (Ptr<Event> event)
 {
   Time psduDuration = event->GetEndTime () - event->GetStartTime ();
   NS_LOG_FUNCTION (this << *event << psduDuration);
-  NS_ASSERT (GetLastRxEndTime () == Simulator::Now ());
+  Ptr<const WifiPpdu> ppdu = event->GetPpdu ();
+  if (!ppdu->IsUlMu ())
+    {
+      NS_ASSERT (GetLastRxEndTime () == Simulator::Now ());
+    }
   NS_ASSERT (event->GetEndTime () == Simulator::Now ());
 
-  uint16_t staId = GetStaId ();
-  Ptr<const WifiPsdu> psdu = GetAddressedPsduInPpdu (event->GetPpdu ());
+  uint16_t staId = GetStaId (ppdu);
+  Ptr<const WifiPsdu> psdu = GetAddressedPsduInPpdu (ppdu);
   if (psdu->GetNMpdus () == 1)
     {
       //We do not enter here for A-MPDU since this is done in WifiPhy::EndOfMpdu
@@ -3426,7 +3430,11 @@ void
 WifiPhy::ResetReceive (Ptr<Event> event)
 {
   NS_LOG_FUNCTION (this << *event);
-  NS_ASSERT (event->GetEndTime () == Simulator::Now ());
+  Ptr<const WifiPpdu> ppdu = event->GetPpdu ();
+  if (!ppdu->IsUlMu ())
+    {
+      NS_ASSERT (event->GetEndTime () == Simulator::Now ());
+    }
   NS_ASSERT (!IsStateRx ());
   m_interference.NotifyRxEnd ();
   m_currentEvent = 0;
@@ -4782,7 +4790,7 @@ Ptr<const WifiPsdu>
 WifiPhy::GetAddressedPsduInPpdu (Ptr<const WifiPpdu> ppdu) const
 {
   Ptr<const WifiPsdu> psdu;
-  if (!ppdu->IsMu ())
+  if (!ppdu->IsDlMu ())
     {
       psdu = ppdu->GetPsdu ();
     }
@@ -4800,25 +4808,32 @@ WifiPhy::GetAddressedPsduInPpdu (Ptr<const WifiPpdu> ppdu) const
               bssColor = bssColorAttribute.Get ();
             }
         }
-      uint16_t staId = GetStaId ();
-      psdu = ppdu->GetPsdu (bssColor, staId);
+      psdu = ppdu->GetPsdu (bssColor, GetStaId (ppdu));
     }
     return psdu;
 }
 
 uint16_t
-WifiPhy::GetStaId (void) const
+WifiPhy::GetStaId (const Ptr<const WifiPpdu> ppdu) const
 {
-  Ptr<WifiNetDevice> device = DynamicCast<WifiNetDevice> (GetDevice ());
-  if (device)
+  uint16_t staId = SU_STA_ID;
+  if (ppdu->IsUlMu ())
     {
-      Ptr<StaWifiMac> mac = DynamicCast<StaWifiMac> (device->GetMac ());
-      if (mac && mac->IsAssociated ())
+      staId = ppdu->GetStaId ();
+    }
+  else if (ppdu->IsDlMu ())
+    {
+      Ptr<WifiNetDevice> device = DynamicCast<WifiNetDevice> (GetDevice ());
+      if (device)
         {
-          return mac->GetAssociationId ();
+          Ptr<StaWifiMac> mac = DynamicCast<StaWifiMac> (device->GetMac ());
+          if (mac && mac->IsAssociated ())
+            {
+              return mac->GetAssociationId ();
+            }
         }
     }
-  return SU_STA_ID;
+  return staId;
 }
 
 WifiSpectrumBand

src/wifi/model/wifi-phy.h

@@ -1789,9 +1789,10 @@ protected:
    * Return the STA ID that has been assigned to the station this PHY belongs to.
    * This is typically called for MU PPDUs, in order to pick the correct PSDU.
    *
+   * \param ppdu the PPDU for which the STA ID is requested
    * \return the STA ID
    */
-  virtual uint16_t GetStaId (void) const;
+  virtual uint16_t GetStaId (const Ptr<const WifiPpdu> ppdu) const;
 
   /**
    * Get the start band index and the stop band index for a given band

src/wifi/model/wifi-ppdu.cc

@@ -377,6 +377,11 @@ WifiPpdu::GetPsdu (uint8_t bssColor, uint16_t staId) const
       NS_ASSERT (m_psdus.size () == 1);
       return m_psdus.at (SU_STA_ID);
     }
+  else if (IsUlMu ())
+    {
+      NS_ASSERT (m_psdus.size () == 1);
+      return m_psdus.begin ()->second;
+    }
   else
     {
       if (bssColor == m_heSig.GetBssColor ())
@@ -391,6 +396,13 @@ WifiPpdu::GetPsdu (uint8_t bssColor, uint16_t staId) const
   return nullptr;
 }
 
+uint16_t
+WifiPpdu::GetStaId (void) const
+{
+  NS_ASSERT (IsUlMu ());
+  return m_psdus.begin ()->first;
+}
+
 bool
 WifiPpdu::IsTruncatedTx (void) const
 {
@@ -469,7 +481,19 @@ WifiPpdu::GetTxDuration (void) const
 bool
 WifiPpdu::IsMu (void) const
 {
-  return ((m_preamble == WIFI_PREAMBLE_VHT_MU) || (m_preamble == WIFI_PREAMBLE_HE_MU) || (m_preamble == WIFI_PREAMBLE_HE_TB));
+  return (IsDlMu () || IsUlMu ());
+}
+
+bool
+WifiPpdu::IsDlMu (void) const
+{
+  return ((m_preamble == WIFI_PREAMBLE_VHT_MU) || (m_preamble == WIFI_PREAMBLE_HE_MU));
+}
+
+bool
+WifiPpdu::IsUlMu (void) const
+{
+  return (m_preamble == WIFI_PREAMBLE_HE_TB);
 }
 
 WifiModulationClass

src/wifi/model/wifi-ppdu.h

@@ -101,6 +101,16 @@ public:
    * \return true if the PPDU is a MU PPDU
    */
   bool IsMu (void) const;
+  /**
+   * Return true if the PPDU is a DL MU PPDU
+   * \return true if the PPDU is a DL MU PPDU
+   */
+  bool IsDlMu (void) const;
+  /**
+   * Return true if the PPDU is an UL MU PPDU
+   * \return true if the PPDU is an UL MU PPDU
+   */
+  bool IsUlMu (void) const;
 
   /**
    * Get the modulation used for the PPDU.
@@ -108,6 +118,12 @@ public:
    */
   WifiModulationClass GetModulation (void) const;
 
+  /**
+   * Get the ID of the STA that transmitted the PPDU (for HE TB PPDU).
+   * \return the ID of the STA that transmitted the PPDU
+   */
+  uint16_t GetStaId (void) const;
+
   /**
    * \brief Print the PPDU contents.
    * \param os output stream in which the data should be printed.

src/wifi/test/wifi-phy-ofdma-test.cc

@@ -56,13 +56,8 @@ public:
   virtual ~OfdmaSpectrumWifiPhy ();
 
 private:
-  /**
-   * Return the STA ID that has been assigned to the station this PHY belongs to.
-   * This is typically called for MU PPDUs, in order to pick the correct PSDU.
-   *
-   * \return the STA ID
-   */
-  uint16_t GetStaId (void) const override;
+  // Inherited
+  uint16_t GetStaId (const Ptr<const WifiPpdu> ppdu) const override;
 
   uint16_t m_staId; ///< ID of the STA to which this PHY belongs to
 };
@@ -78,9 +73,13 @@ OfdmaSpectrumWifiPhy::~OfdmaSpectrumWifiPhy()
 }
 
 uint16_t
-OfdmaSpectrumWifiPhy::GetStaId (void) const
+OfdmaSpectrumWifiPhy::GetStaId (const Ptr<const WifiPpdu> ppdu) const
 {
-  return m_staId;
+  if (ppdu->IsDlMu ())
+    {
+      return m_staId;
+    }
+  return SpectrumWifiPhy::GetStaId (ppdu);
 }
 
 /**