From 9651b2a9c4ce6dc7352690636bcf5653ce0d87be Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Deronne?= <sebastien.deronne@gmail.com>
Date: Sun, 26 Jan 2025 11:46:12 +0100
Subject: [PATCH] wifi: Simplify RU calculations using PHY index

---
 src/wifi/model/wifi-phy-operating-channel.cc | 34 +++++---------------
 src/wifi/model/wifi-tx-vector.cc             | 25 ++++----------
 2 files changed, 14 insertions(+), 45 deletions(-)

diff --git a/src/wifi/model/wifi-phy-operating-channel.cc b/src/wifi/model/wifi-phy-operating-channel.cc
index 45cd6eb04..e4eb3f191 100644
--- a/src/wifi/model/wifi-phy-operating-channel.cc
+++ b/src/wifi/model/wifi-phy-operating-channel.cc
@@ -812,8 +812,7 @@ WifiPhyOperatingChannel::GetAll20MHzChannelIndicesInSecondary(
 std::set<uint8_t>
 WifiPhyOperatingChannel::Get20MHzIndicesCoveringRu(WifiRu::RuSpec ru, MHz_u width) const
 {
-    auto ruType = WifiRu::GetRuType(ru);
-    auto ruIndex = WifiRu::GetIndex(ru);
+    const auto ruType = WifiRu::GetRuType(ru);
 
     NS_ASSERT_MSG(WifiRu::GetBandwidth(ruType) <= width,
                   "No RU of type " << ruType << " is contained in a " << width << " MHz channel");
@@ -822,7 +821,7 @@ WifiPhyOperatingChannel::Get20MHzIndicesCoveringRu(WifiRu::RuSpec ru, MHz_u widt
                                       << GetTotalWidth() << ")");
 
     // handle first the special case of center 26-tone RUs
-    if (ruType == RuType::RU_26_TONE && ruIndex == 19)
+    if (const auto ruIndex = WifiRu::GetIndex(ru); ruType == RuType::RU_26_TONE && ruIndex == 19)
     {
         NS_ASSERT_MSG(WifiRu::IsHe(ru), "Center 26-tone RUs can only be used with HE");
         NS_ASSERT_MSG(width >= MHz_u{80},
@@ -837,31 +836,14 @@ WifiPhyOperatingChannel::Get20MHzIndicesCoveringRu(WifiRu::RuSpec ru, MHz_u widt
         return indices;
     }
 
-    if (ruType == RuType::RU_26_TONE && ruIndex > 19)
+    auto ruPhyIndex = WifiRu::GetPhyIndex(ru, width, m_primary20Index);
+    if (ruType == RuType::RU_26_TONE && ruPhyIndex > 19)
     {
         // "ignore" the center 26-tone RUs in 80 MHz channels
-        ruIndex--;
-        if (ruIndex > 37)
+        ruPhyIndex--;
+        if (ruPhyIndex > 37)
         {
-            NS_ASSERT(WifiRu::IsEht(ru));
-            ruIndex -= (ruIndex - 19) / 37;
-        }
-    }
-
-    // if the HE RU refers to a 160 MHz channel, we have to update the RU index (which refers to an
-    // 80 MHz channel) if the RU is not in the lower 80 MHz channel
-    if (width == MHz_u{160} && WifiRu::IsHe(ru))
-    {
-        bool primary80IsLower80 = (m_primary20Index < 4);
-        if (primary80IsLower80 != std::get<HeRu::RuSpec>(ru).GetPrimary80MHz())
-        {
-            auto nRusIn80MHz = HeRu::GetNRus(MHz_u{80}, ruType);
-            // "ignore" the center 26-tone RU in an 80 MHz channel
-            if (ruType == RuType::RU_26_TONE)
-            {
-                nRusIn80MHz--;
-            }
-            ruIndex += nRusIn80MHz;
+            ruPhyIndex -= (ruPhyIndex - 19) / 37;
         }
     }
 
@@ -896,7 +878,7 @@ WifiPhyOperatingChannel::Get20MHzIndicesCoveringRu(WifiRu::RuSpec ru, MHz_u widt
                         ruType,
                         WifiRu::IsHe(ru) ? WIFI_MOD_CLASS_HE : WIFI_MOD_CLASS_EHT);
     // compute the index (starting at 0) of the covering channel within the given width
-    std::size_t indexOfCoveringChannelInGivenWidth = (ruIndex - 1) / nRusInCoveringChannel;
+    std::size_t indexOfCoveringChannelInGivenWidth = (ruPhyIndex - 1) / nRusInCoveringChannel;
 
     // expand the index of the covering channel in the indices of its constituent
     // 20 MHz channels (within the given width)
diff --git a/src/wifi/model/wifi-tx-vector.cc b/src/wifi/model/wifi-tx-vector.cc
index 84f877615..2253b72f1 100644
--- a/src/wifi/model/wifi-tx-vector.cc
+++ b/src/wifi/model/wifi-tx-vector.cc
@@ -799,7 +799,7 @@ WifiTxVector::DeriveRuAllocation(uint8_t p20Index) const
     for (const auto& [ru, staIds] : orderedMap)
     {
         const auto ruType = WifiRu::GetRuType(ru);
-        auto ruIndex = WifiRu::GetIndex(ru);
+        auto ruIndex = WifiRu::GetPhyIndex(ru, m_channelWidth, p20Index);
         if ((ruType == RuType::RU_26_TONE) && (ruIndex == 19))
         {
             continue;
@@ -811,25 +811,12 @@ WifiTxVector::DeriveRuAllocation(uint8_t p20Index) const
             WifiRu::GetRusOfType(ruBw > MHz_u{20} ? ruBw : MHz_u{20}, ruType, mc);
         if ((m_channelWidth >= MHz_u{80}) && (ruIndex > 19))
         {
-            // take into account the center 26-tone RU in the low 80 MHz
+            // "ignore" the center 26-tone RUs in 80 MHz channels
             ruIndex--;
-        }
-        const auto isPrimary80MHz = std::get<HeRu::RuSpec>(ru).GetPrimary80MHz();
-        const auto primary80IsLower80 = (p20Index < m_channelWidth / MHz_u{40});
-        const auto isLow80 =
-            (primary80IsLower80 && isPrimary80MHz) || (!primary80IsLower80 && !isPrimary80MHz);
-        if ((m_channelWidth > MHz_u{80}) && !isLow80 && (ruIndex > 19))
-        {
-            // take into account the center 26-tone RU in the high 80 MHz
-            ruIndex--;
-        }
-        if ((m_channelWidth > MHz_u{80}) && !isLow80 &&
-            (ruType != WifiRu::GetRuType(m_channelWidth)))
-        {
-            // adjust RU index for the high 80 MHz: in that case index is restarting at 1,
-            // hence we need to add an offset corresponding to the number of RUs of the same type in
-            // the low 80 MHz
-            ruIndex += HeRu::GetRusOfType(MHz_u{80}, ruType).size();
+            if (ruIndex > 37)
+            {
+                ruIndex -= (ruIndex - 19) / 37;
+            }
         }
         const auto numSubchannelsForRu = (ruBw < MHz_u{20}) ? 1 : Count20MHzSubchannels(ruBw);
         const auto index = (ruBw < MHz_u{20}) ? ((ruIndex - 1) / rusPerSubchannel.size())