From 0d7bd988e0d5b60ac1aaef27a16dae84b53c3de2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Deronne?= <sebastien.deronne@gmail.com>
Date: Thu, 27 Oct 2022 20:29:09 +0200
Subject: [PATCH] wifi: add function to reconstruct per user information in
 TXVECTOR for DL OFDMA

---
 src/wifi/model/eht/eht-ppdu.cc   | 10 ++++-
 src/wifi/model/he/he-ppdu.cc     | 68 ++++++++++++++++++++++++++++++++
 src/wifi/model/he/he-ppdu.h      |  8 ++++
 src/wifi/model/wifi-tx-vector.cc |  6 +++
 src/wifi/model/wifi-tx-vector.h  |  6 +++
 5 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/src/wifi/model/eht/eht-ppdu.cc b/src/wifi/model/eht/eht-ppdu.cc
index b2953302b..6dc48bfa7 100644
--- a/src/wifi/model/eht/eht-ppdu.cc
+++ b/src/wifi/model/eht/eht-ppdu.cc
@@ -97,9 +97,15 @@ EhtPpdu::SetTxVectorFromPhyHeaders(WifiTxVector& txVector,
     txVector.SetLength(lSig.GetLength());
     txVector.SetAggregation(m_psdus.size() > 1 || m_psdus.begin()->second->IsAggregate());
     txVector.SetEhtPpduType(m_ehtPpduType); // FIXME: PPDU type should be read from U-SIG
-    for (const auto& muUserInfo : m_muUserInfos)
+    if (txVector.IsDlMu())
     {
-        txVector.SetHeMuUserInfo(muUserInfo.first, muUserInfo.second);
+        auto copyTxVector = txVector;
+        for (const auto& muUserInfo : m_muUserInfos)
+        {
+            txVector.SetHeMuUserInfo(muUserInfo.first, muUserInfo.second);
+        }
+        SetHeMuUserInfos(copyTxVector, heSig);
+        NS_ASSERT(txVector.GetHeMuUserInfoMap() == copyTxVector.GetHeMuUserInfoMap());
     }
     if (ns3::IsDlMu(m_preamble))
     {
diff --git a/src/wifi/model/he/he-ppdu.cc b/src/wifi/model/he/he-ppdu.cc
index 1e71801bf..de38d2cc3 100644
--- a/src/wifi/model/he/he-ppdu.cc
+++ b/src/wifi/model/he/he-ppdu.cc
@@ -223,10 +223,13 @@ HePpdu::SetTxVectorFromPhyHeaders(WifiTxVector& txVector,
     }
     if (IsDlMu())
     {
+        auto copyTxVector = txVector;
         for (const auto& muUserInfo : m_muUserInfos)
         {
             txVector.SetHeMuUserInfo(muUserInfo.first, muUserInfo.second);
         }
+        SetHeMuUserInfos(copyTxVector, heSig);
+        NS_ASSERT(txVector.GetHeMuUserInfoMap() == copyTxVector.GetHeMuUserInfoMap());
     }
     if (txVector.IsDlMu())
     {
@@ -241,6 +244,71 @@ HePpdu::SetTxVectorFromPhyHeaders(WifiTxVector& txVector,
     }
 }
 
+void
+HePpdu::SetHeMuUserInfos(WifiTxVector& txVector, const HeSigHeader& heSig) const
+{
+    WifiTxVector::HeMuUserInfoMap userInfos{};
+    const auto contentChannels = heSig.GetHeSigBContentChannels();
+    const auto ruAllocation = heSig.GetRuAllocation();
+    auto contentChannelIndex = 0;
+    for (const auto& contentChannel : contentChannels)
+    {
+        auto numRusLeft = 0;
+        auto ruAllocIndex = contentChannelIndex;
+        for (const auto& userInfo : contentChannel)
+        {
+            auto ruSpecs = HeRu::GetRuSpecs(ruAllocation.at(ruAllocIndex));
+            if (ruSpecs.empty())
+            {
+                continue;
+            }
+            if (numRusLeft == 0)
+            {
+                numRusLeft = ruSpecs.size();
+            }
+            auto ruIndex = (ruSpecs.size() - numRusLeft);
+            auto ruSpec = ruSpecs.at(ruIndex);
+            auto ruType = ruSpec.GetRuType();
+            if ((ruAllocation.size() == 8) && (ruType == HeRu::RU_996_TONE) &&
+                (std::all_of(
+                    contentChannel.cbegin(),
+                    contentChannel.cend(),
+                    [&userInfo](const auto& item) { return userInfo.staId == item.staId; })))
+            {
+                NS_ASSERT(txVector.GetChannelWidth() == 160);
+                ruType = HeRu::RU_2x996_TONE;
+            }
+            const auto ruBw = HeRu::GetBandwidth(ruType);
+            auto primary80 = ruAllocIndex < 4;
+            auto num20MhzSubchannelsInRu = (ruBw < 20) ? 1 : (ruBw / 20);
+            auto numRuAllocsInContentChannel = std::max(1, num20MhzSubchannelsInRu / 2);
+            auto ruIndexOffset = (ruBw < 20) ? (ruSpecs.size() * ruAllocIndex)
+                                             : (ruAllocIndex / num20MhzSubchannelsInRu);
+            if (!primary80)
+            {
+                ruIndexOffset -= HeRu::GetRusOfType(80, ruType).size();
+            }
+            if (!txVector.IsAllocated(userInfo.staId))
+            {
+                txVector.SetHeMuUserInfo(userInfo.staId,
+                                         {{ruType, ruSpec.GetIndex() + ruIndexOffset, primary80},
+                                          userInfo.mcs,
+                                          userInfo.nss});
+            }
+            if (ruType == HeRu::RU_2x996_TONE)
+            {
+                return;
+            }
+            numRusLeft--;
+            if (numRusLeft == 0)
+            {
+                ruAllocIndex += (2 * numRuAllocsInContentChannel);
+            }
+        }
+        contentChannelIndex++;
+    }
+}
+
 Time
 HePpdu::GetTxDuration() const
 {
diff --git a/src/wifi/model/he/he-ppdu.h b/src/wifi/model/he/he-ppdu.h
index cdd8c2890..361e4d09d 100644
--- a/src/wifi/model/he/he-ppdu.h
+++ b/src/wifi/model/he/he-ppdu.h
@@ -370,6 +370,14 @@ class HePpdu : public OfdmPpdu
                                            const LSigHeader& lSig,
                                            const HeSigHeader& heSig) const;
 
+    /**
+     * Reconstruct HeMuUserInfoMap from HE-SIG-B header.
+     *
+     * \param txVector the TXVECTOR to set its HeMuUserInfoMap
+     * \param heSig the HE-SIG-B to use to reconstruct HeMuUserInfoMap
+     */
+    void SetHeMuUserInfos(WifiTxVector& txVector, const HeSigHeader& heSig) const;
+
 #ifndef NS3_BUILD_PROFILE_DEBUG
     HeSigHeader m_heSig; //!< the HE-SIG PHY header
 #endif
diff --git a/src/wifi/model/wifi-tx-vector.cc b/src/wifi/model/wifi-tx-vector.cc
index 7342e0dfa..d5488ff2d 100644
--- a/src/wifi/model/wifi-tx-vector.cc
+++ b/src/wifi/model/wifi-tx-vector.cc
@@ -482,6 +482,12 @@ WifiTxVector::IsUlMu() const
     return ns3::IsUlMu(m_preamble);
 }
 
+bool
+WifiTxVector::IsAllocated(uint16_t staId) const
+{
+    return m_muUserInfos.count(staId) > 0;
+}
+
 HeRu::RuSpec
 WifiTxVector::GetRu(uint16_t staId) const
 {
diff --git a/src/wifi/model/wifi-tx-vector.h b/src/wifi/model/wifi-tx-vector.h
index 4ed8e0d5b..8a81c3821 100644
--- a/src/wifi/model/wifi-tx-vector.h
+++ b/src/wifi/model/wifi-tx-vector.h
@@ -378,6 +378,12 @@ class WifiTxVector
      * transmission
      */
     bool IsUlMu() const;
+    /**
+     * Check if STA ID is allocated
+     * \param staId STA ID
+     * \return true if allocated, false otherwise
+     */
+    bool IsAllocated(uint16_t staId) const;
     /**
      * Get the RU specification for the STA-ID.
      * This is applicable only for MU.