From 338d625f558f6124eb850d8891c036a39d0cb132 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Fri, 12 May 2023 09:37:45 +0200
Subject: [PATCH] wifi: Update TID-to-Link Mapping IE to 802.11be D3.1

---
 .../model/eht/tid-to-link-mapping-element.cc  | 170 +++++++++++++++---
 .../model/eht/tid-to-link-mapping-element.h   |  54 ++++--
 src/wifi/test/wifi-eht-info-elems-test.cc     |  27 +--
 3 files changed, 201 insertions(+), 50 deletions(-)

diff --git a/src/wifi/model/eht/tid-to-link-mapping-element.cc b/src/wifi/model/eht/tid-to-link-mapping-element.cc
index 9a8feb6ab..63b291ad7 100644
--- a/src/wifi/model/eht/tid-to-link-mapping-element.cc
+++ b/src/wifi/model/eht/tid-to-link-mapping-element.cc
@@ -17,20 +17,25 @@
  * Author: Sharan Naribole <sharan.naribole@gmail.com>
  */
 
-#include <ns3/assert.h>
-#include <ns3/tid-to-link-mapping-element.h>
+#include "tid-to-link-mapping-element.h"
+
+#include "ns3/assert.h"
+#include "ns3/simulator.h"
 
 namespace ns3
 {
 
+/// Bitmask with all bits from 63 to 26 set to 1, all the others set to 0
+static constexpr uint64_t BIT_63_TO_26_MASK = 0xfffffffffc000000;
+
 uint16_t
 TidToLinkMapping::Control::GetSubfieldSize() const
 {
-    // IEEE 802.11be D2.0 Figure 9-1002an
-    NS_ASSERT_MSG(defaultMapping != presenceBitmap.has_value(),
+    // IEEE 802.11be D3.1 Figure 9-1002ap
+    NS_ASSERT_MSG(!defaultMapping || !presenceBitmap.has_value(),
                   "Presence bitmap not expected if default mapping is set");
     auto size = WIFI_TID_TO_LINK_MAPPING_CONTROL_BASIC_SIZE_B;
-    if (defaultMapping)
+    if (!presenceBitmap.has_value())
     {
         return size;
     }
@@ -40,10 +45,12 @@ TidToLinkMapping::Control::GetSubfieldSize() const
 void
 TidToLinkMapping::Control::Serialize(Buffer::Iterator& start) const
 {
-    uint8_t val = static_cast<uint8_t>(direction);
-    val |= ((defaultMapping ? 1 : 0) << 2);
+    uint8_t val = static_cast<uint8_t>(direction) | ((defaultMapping ? 1 : 0) << 2) |
+                  ((mappingSwitchTimePresent ? 1 : 0) << 3) |
+                  ((expectedDurationPresent ? 1 : 0) << 4) | ((linkMappingSize == 1 ? 1 : 0) << 5);
+
     start.WriteU8(val);
-    NS_ASSERT_MSG(defaultMapping != presenceBitmap.has_value(),
+    NS_ASSERT_MSG(!defaultMapping || !presenceBitmap.has_value(),
                   "Presence bitmap not expected if default mapping is set");
     if (presenceBitmap.has_value())
     {
@@ -61,8 +68,12 @@ TidToLinkMapping::Control::Deserialize(Buffer::Iterator start)
 
     direction = static_cast<TidLinkMapDir>(val & 0x03);
     defaultMapping = (((val >> 2) & 0x01) == 1);
+    mappingSwitchTimePresent = (((val >> 3) & 0x01) == 1);
+    expectedDurationPresent = (((val >> 4) & 0x01) == 1);
+    linkMappingSize = (((val >> 5) & 0x01) == 1 ? 1 : 2);
     if (defaultMapping)
     {
+        presenceBitmap.reset();
         return count;
     }
     presenceBitmap = i.ReadU8();
@@ -82,7 +93,60 @@ TidToLinkMapping::ElementIdExt() const
 }
 
 void
-TidToLinkMapping::SetLinkMappingOfTid(uint8_t tid, std::list<uint8_t> linkIds)
+TidToLinkMapping::SetMappingSwitchTime(Time mappingSwitchTime)
+{
+    // The 2 octet Mapping Switch Time field has units of TUs and is set to the time at which
+    // the new mapping is established using as a time-base the value of the TSF corresponding
+    // to the BSS identified by the BSSID of the frame containing the TID-To-Link Mapping
+    // element: i.e., bits 10 to 25 of the TSF. (Sec. 9.4.2.314 of 802.11be D3.1)
+    NS_ABORT_IF(mappingSwitchTime < Simulator::Now());
+    auto switchTimeUsec = static_cast<uint64_t>(mappingSwitchTime.GetMicroSeconds());
+    // set the Mapping Switch Time field to bits 10 to 25 of the given time
+    m_mappingSwitchTime = (switchTimeUsec & ~BIT_63_TO_26_MASK) >> 10;
+    m_control.mappingSwitchTimePresent = true;
+}
+
+std::optional<Time>
+TidToLinkMapping::GetMappingSwitchTime() const
+{
+    if (!m_control.mappingSwitchTimePresent)
+    {
+        return std::nullopt;
+    }
+
+    auto nowUsec = static_cast<uint64_t>(Simulator::Now().GetMicroSeconds());
+    uint64_t switchTimeUsec = (*m_mappingSwitchTime << 10) + (nowUsec & BIT_63_TO_26_MASK);
+    if (switchTimeUsec < nowUsec)
+    {
+        // The switch time derived from the value in the corresponding field may be less than the
+        // current time in case the bits 10 to 25 of TSF have been reset since the transmission
+        // of the frame carrying this field. In such a case we have to increase bits 63 to 26 by 1
+        switchTimeUsec += (1 << 26);
+    }
+    return MicroSeconds(switchTimeUsec);
+}
+
+void
+TidToLinkMapping::SetExpectedDuration(Time expectedDuration)
+{
+    auto durationTu = static_cast<uint64_t>(expectedDuration.GetMicroSeconds()) >> 10;
+    m_expectedDuration = (durationTu & 0x0000000000ffffff); // Expected Duration size is 3 bytes
+    m_control.expectedDurationPresent = true;
+}
+
+std::optional<Time>
+TidToLinkMapping::GetExpectedDuration() const
+{
+    if (!m_control.expectedDurationPresent)
+    {
+        return std::nullopt;
+    }
+
+    return MicroSeconds(*m_expectedDuration << 10);
+}
+
+void
+TidToLinkMapping::SetLinkMappingOfTid(uint8_t tid, std::set<uint8_t> linkIds)
 {
     NS_ABORT_MSG_IF(tid > 7, "Invalid tid: " << +tid);
     NS_ABORT_MSG_IF(m_control.defaultMapping,
@@ -94,35 +158,35 @@ TidToLinkMapping::SetLinkMappingOfTid(uint8_t tid, std::list<uint8_t> linkIds)
     for (const auto& linkId : linkIds)
     {
         linkMapping |= (1 << linkId);
+        if (linkId > 7)
+        {
+            m_control.linkMappingSize = 2;
+        }
     }
 
     m_linkMapping[tid] = linkMapping;
-
-    if (!m_control.presenceBitmap.has_value())
-    {
-        m_control.presenceBitmap = 0;
-    }
-    *m_control.presenceBitmap |= (1 << tid);
+    m_control.presenceBitmap = m_control.presenceBitmap.value_or(0) | (1 << tid);
 }
 
-std::list<uint8_t>
+std::set<uint8_t>
 TidToLinkMapping::GetLinkMappingOfTid(uint8_t tid) const
 {
     auto it = m_linkMapping.find(tid);
 
-    if (it == m_linkMapping.end())
+    if (it == m_linkMapping.cend())
     {
         return {};
     }
 
-    std::list<uint8_t> linkIds;
+    std::set<uint8_t> linkIds;
     for (uint8_t linkId = 0; linkId < 15; linkId++)
     {
         if (((it->second >> linkId) & 0x0001) == 1)
         {
-            linkIds.push_back(linkId);
+            linkIds.insert(linkId);
         }
     }
+    NS_ABORT_MSG_IF(linkIds.empty(), "TID " << +tid << " cannot be mapped to an empty link set");
 
     return linkIds;
 }
@@ -130,26 +194,53 @@ TidToLinkMapping::GetLinkMappingOfTid(uint8_t tid) const
 uint16_t
 TidToLinkMapping::GetInformationFieldSize() const
 {
-    // IEEE 802.11be D2.0 9.4.2.314 TID-To-Link Mapping element
+    // IEEE 802.11be D3.1 9.4.2.314 TID-To-Link Mapping element
     uint16_t ret = WIFI_IE_ELEMENT_ID_EXT_SIZE; // Element ID Extension
     ret += m_control.GetSubfieldSize();
-    NS_ASSERT_MSG(m_linkMapping.empty() == m_control.defaultMapping,
+    if (m_control.mappingSwitchTimePresent)
+    {
+        ret += 2; // Mapping Switch Time
+    }
+    if (m_control.expectedDurationPresent)
+    {
+        ret += 3; // Expected Duration
+    }
+
+    NS_ASSERT_MSG(!m_control.defaultMapping || m_linkMapping.empty(),
                   "Per-TID link mapping not expected if default mapping is set");
-    ret += WIFI_LINK_MAPPING_PER_TID_SIZE_B * (m_linkMapping.size());
+    ret += m_control.linkMappingSize * (m_linkMapping.size());
     return ret;
 }
 
 void
 TidToLinkMapping::SerializeInformationField(Buffer::Iterator start) const
 {
-    // IEEE 802.11be D2.0 9.4.2.314 TID-To-Link Mapping element
+    // IEEE 802.11be D3.1 9.4.2.314 TID-To-Link Mapping element
     m_control.Serialize(start);
-    NS_ASSERT_MSG(m_linkMapping.empty() == m_control.defaultMapping,
+    if (m_control.mappingSwitchTimePresent)
+    {
+        start.WriteHtolsbU16(*m_mappingSwitchTime);
+    }
+    if (m_control.expectedDurationPresent)
+    {
+        start.WriteU8((*m_expectedDuration >> 0) & 0xff);
+        start.WriteU8((*m_expectedDuration >> 8) & 0xff);
+        start.WriteU8((*m_expectedDuration >> 16) & 0xff);
+    }
+
+    NS_ASSERT_MSG(!m_control.defaultMapping || m_linkMapping.empty(),
                   "Per-TID link mapping not expected if default mapping is set");
 
     for (const auto& [tid, linkMapping] : m_linkMapping)
     {
-        start.WriteHtolsbU16(linkMapping);
+        if (m_control.linkMappingSize == 1)
+        {
+            start.WriteU8(linkMapping);
+        }
+        else
+        {
+            start.WriteHtolsbU16(linkMapping);
+        }
     }
 }
 
@@ -162,6 +253,23 @@ TidToLinkMapping::DeserializeInformationField(Buffer::Iterator start, uint16_t l
     NS_ASSERT_MSG(nCtrlOctets <= length, "Tid-to-Link Mapping deserialize error");
     i.Next(nCtrlOctets);
     count += nCtrlOctets;
+    if (m_control.mappingSwitchTimePresent)
+    {
+        m_mappingSwitchTime = i.ReadLsbtohU16();
+        count += 2;
+    }
+    if (m_control.expectedDurationPresent)
+    {
+        uint8_t byte0 = i.ReadU8();
+        uint8_t byte1 = i.ReadU8();
+        uint8_t byte2 = i.ReadU8();
+        m_expectedDuration = byte2;
+        m_expectedDuration.value() <<= 8;
+        m_expectedDuration.value() |= byte1;
+        m_expectedDuration.value() <<= 8;
+        m_expectedDuration.value() |= byte0;
+        count += 3;
+    }
     m_linkMapping.clear();
     if (m_control.presenceBitmap.has_value())
     {
@@ -172,8 +280,16 @@ TidToLinkMapping::DeserializeInformationField(Buffer::Iterator start, uint16_t l
         {
             if (((presenceBitmap >> tid) & 0x01) == 1)
             {
-                m_linkMapping[tid] = i.ReadLsbtohU16();
-                count += 2;
+                if (m_control.linkMappingSize == 1)
+                {
+                    m_linkMapping[tid] = i.ReadU8();
+                    count++;
+                }
+                else
+                {
+                    m_linkMapping[tid] = i.ReadLsbtohU16();
+                    count += 2;
+                }
             }
         }
     }
diff --git a/src/wifi/model/eht/tid-to-link-mapping-element.h b/src/wifi/model/eht/tid-to-link-mapping-element.h
index 3a34cb893..d9a48e192 100644
--- a/src/wifi/model/eht/tid-to-link-mapping-element.h
+++ b/src/wifi/model/eht/tid-to-link-mapping-element.h
@@ -20,12 +20,12 @@
 #ifndef TID_TO_LINK_MAPPING_H
 #define TID_TO_LINK_MAPPING_H
 
-#include <ns3/wifi-information-element.h>
+#include "ns3/nstime.h"
+#include "ns3/wifi-information-element.h"
 
-#include <list>
 #include <map>
 #include <optional>
-#include <vector>
+#include <set>
 
 namespace ns3
 {
@@ -45,8 +45,6 @@ enum class TidLinkMapDir : uint8_t
 constexpr auto DEFAULT_WIFI_TID_LINK_MAPPING{true};
 /// default value for the Direction subfield of the TID-To-Link Control field
 constexpr auto DEFAULT_WIFI_TID_LINK_MAP_DIR{TidLinkMapDir::BOTH_DIRECTIONS};
-/// size in bytes of the Link Mapping Of TID n field (IEEE 802.11be D2.0 9.4.2.314)
-constexpr uint16_t WIFI_LINK_MAPPING_PER_TID_SIZE_B = 2;
 /// size in bytes of the TID-To-Link Control field with default link mapping
 constexpr uint16_t WIFI_TID_TO_LINK_MAPPING_CONTROL_BASIC_SIZE_B =
     1; // IEEE 802.11be D2.0 9.4.2.314
@@ -59,7 +57,7 @@ constexpr uint16_t WIFI_LINK_MAPPING_PRESENCE_IND_SIZE_B = 1;
  *
  * This class serializes and deserializes
  * the TID-to-Link Mapping element
- * IEEE 802.11be D2.0 9.4.2.314
+ * IEEE 802.11be D3.1 9.4.2.314
  *
  */
 class TidToLinkMapping : public WifiInformationElement
@@ -67,13 +65,14 @@ class TidToLinkMapping : public WifiInformationElement
   public:
     /**
      * TID-to-Link Mapping Control subfield
-     * IEEE 802.11be D2.0 Figure 9-1002an
+     * IEEE 802.11be D3.1 Figure 9-1002ap
      */
     struct Control
     {
+        friend class TidToLinkMapping;
+
         TidLinkMapDir direction{DEFAULT_WIFI_TID_LINK_MAP_DIR}; ///< Direction
         bool defaultMapping{DEFAULT_WIFI_TID_LINK_MAPPING};     ///< Default link mapping
-        std::optional<uint8_t> presenceBitmap;                  ///< Link Mapping Presence Indicator
 
         /// \return Serialized size of TID-to-Link Mapping Control in octets
         uint16_t GetSubfieldSize() const;
@@ -91,11 +90,43 @@ class TidToLinkMapping : public WifiInformationElement
          * \return the number of octets read
          */
         uint16_t Deserialize(Buffer::Iterator start);
+
+      private:
+        /**
+         * These members are private to prevent users to manipulate them directly (their values
+         * depend on the values of other fields)
+         */
+        bool mappingSwitchTimePresent{false}; ///< Mapping Switch Time Present
+        bool expectedDurationPresent{false};  ///< Expected Duration Present
+        uint8_t linkMappingSize{1}; ///< length of the Link Mapping Of TID n field in octets
+        std::optional<uint8_t> presenceBitmap; ///< Link Mapping Presence Indicator
     };
 
     WifiInformationElementId ElementId() const override;
     WifiInformationElementId ElementIdExt() const override;
 
+    /**
+     * Set the Mapping Switch Time field.
+     *
+     * \param mappingSwitchTime the time when the new mapping is established (rounded to a
+     *                          multiple of a TU)
+     */
+    void SetMappingSwitchTime(Time mappingSwitchTime);
+
+    /// \return the value of the Mapping Switch Time field, if present.
+    std::optional<Time> GetMappingSwitchTime() const;
+
+    /**
+     * Set the Expected Duration field.
+     *
+     * \param expectedDuration the value for the Expected Duration field (rounded to a
+     *                         multiple of a TU)
+     */
+    void SetExpectedDuration(Time expectedDuration);
+
+    /// \return the value of the Expected Duration field, if present.
+    std::optional<Time> GetExpectedDuration() const;
+
     /**
      * Set the Link Mapping field of the given TID such that the given TID is mapped
      * to the links associated with the given link IDs.
@@ -103,19 +134,22 @@ class TidToLinkMapping : public WifiInformationElement
      * \param tid the given TID
      * \param linkIds the IDs of the links which the given TID is mapped to
      */
-    void SetLinkMappingOfTid(uint8_t tid, std::list<uint8_t> linkIds);
+    void SetLinkMappingOfTid(uint8_t tid, std::set<uint8_t> linkIds);
     /**
      * Get the Link Mapping field of the given TID.
      *
      * \param tid the given TID
      * \return the IDs of the links which the given TID is mapped to
      */
-    std::list<uint8_t> GetLinkMappingOfTid(uint8_t tid) const;
+    std::set<uint8_t> GetLinkMappingOfTid(uint8_t tid) const;
 
     TidToLinkMapping::Control m_control;       ///< TID-to-link Mapping Control
     std::map<uint8_t, uint16_t> m_linkMapping; ///< TID-indexed Link Mapping
 
   private:
+    std::optional<uint16_t> m_mappingSwitchTime; ///< Mapping Switch Time
+    std::optional<uint32_t> m_expectedDuration;  ///< Expected Duration
+
     uint16_t GetInformationFieldSize() const override;
     void SerializeInformationField(Buffer::Iterator start) const override;
     uint16_t DeserializeInformationField(Buffer::Iterator start, uint16_t length) override;
diff --git a/src/wifi/test/wifi-eht-info-elems-test.cc b/src/wifi/test/wifi-eht-info-elems-test.cc
index e26ced5d9..f60c85d88 100644
--- a/src/wifi/test/wifi-eht-info-elems-test.cc
+++ b/src/wifi/test/wifi-eht-info-elems-test.cc
@@ -27,6 +27,7 @@
 #include "ns3/wifi-phy-operating-channel.h"
 
 #include <optional>
+#include <set>
 #include <sstream>
 #include <vector>
 
@@ -1224,7 +1225,7 @@ class TidToLinkMappingElementTest : public HeaderSerializationTestCase
      * \param args A sequence of TID-Link mapping pairs
      */
     template <typename... Args>
-    void SetLinkMapping(uint8_t tid, const std::list<uint8_t>& linkIds, Args&&... args);
+    void SetLinkMapping(uint8_t tid, const std::set<uint8_t>& linkIds, Args&&... args);
 
     /**
      * Base case to stop the recursion performed by the templated version of this method.
@@ -1251,7 +1252,7 @@ TidToLinkMappingElementTest::TidToLinkMappingElementTest(TidLinkMapDir direction
 template <typename... Args>
 void
 TidToLinkMappingElementTest::SetLinkMapping(uint8_t tid,
-                                            const std::list<uint8_t>& linkIds,
+                                            const std::set<uint8_t>& linkIds,
                                             Args&&... args)
 {
     m_tidToLinkMapping.m_control.defaultMapping = false;
@@ -1364,31 +1365,31 @@ WifiEhtInfoElemsTestSuite::WifiEhtInfoElemsTestSuite()
     AddTestCase(new WifiEhtCapabilitiesIeTest(false, 320), TestCase::QUICK);
     AddTestCase(new TidToLinkMappingElementTest(TidLinkMapDir::DOWNLINK), TestCase::QUICK);
     AddTestCase(
-        new TidToLinkMappingElementTest(TidLinkMapDir::UPLINK, 3, std::list<uint8_t>{0, 4, 6}),
+        new TidToLinkMappingElementTest(TidLinkMapDir::UPLINK, 3, std::set<uint8_t>{0, 4, 6}),
         TestCase::QUICK);
     AddTestCase(new TidToLinkMappingElementTest(TidLinkMapDir::BOTH_DIRECTIONS,
                                                 3,
-                                                std::list<uint8_t>{0, 4, 6},
+                                                std::set<uint8_t>{0, 4, 6},
                                                 6,
-                                                std::list<uint8_t>{3, 7, 11, 14}),
+                                                std::set<uint8_t>{3, 7, 11, 14}),
                 TestCase::QUICK);
     AddTestCase(new TidToLinkMappingElementTest(TidLinkMapDir::DOWNLINK,
                                                 0,
-                                                std::list<uint8_t>{0, 1, 2},
+                                                std::set<uint8_t>{0, 1, 2},
                                                 1,
-                                                std::list<uint8_t>{3, 4, 5},
+                                                std::set<uint8_t>{3, 4, 5},
                                                 2,
-                                                std::list<uint8_t>{6, 7},
+                                                std::set<uint8_t>{6, 7},
                                                 3,
-                                                std::list<uint8_t>{8, 9, 10},
+                                                std::set<uint8_t>{8, 9, 10},
                                                 4,
-                                                std::list<uint8_t>{11, 12, 13},
+                                                std::set<uint8_t>{11, 12, 13},
                                                 5,
-                                                std::list<uint8_t>{14},
+                                                std::set<uint8_t>{14},
                                                 6,
-                                                std::list<uint8_t>{1, 3, 6},
+                                                std::set<uint8_t>{1, 3, 6},
                                                 7,
-                                                std::list<uint8_t>{11, 14}),
+                                                std::set<uint8_t>{11, 14}),
                 TestCase::QUICK);
     AddTestCase(new EhtOperationElementTest({0, 0, 0, 0, 0}, 1, 2, 3, 4, 5, 6, 7, 8, std::nullopt),
                 TestCase::QUICK);