From cec80188dd3d44db242a5ee2c0bdf646248c61e7 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Mon, 30 Oct 2023 19:40:40 +0100
Subject: [PATCH] wifi: Main PHY can switch from non-primary link to
 non-primary link

---
 src/wifi/model/eht/default-emlsr-manager.cc | 51 ++++++++++++---------
 src/wifi/model/eht/default-emlsr-manager.h  |  3 --
 2 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/src/wifi/model/eht/default-emlsr-manager.cc b/src/wifi/model/eht/default-emlsr-manager.cc
index 90a952791..7befd67be 100644
--- a/src/wifi/model/eht/default-emlsr-manager.cc
+++ b/src/wifi/model/eht/default-emlsr-manager.cc
@@ -98,16 +98,32 @@ DefaultEmlsrManager::NotifyMainPhySwitch(uint8_t currLinkId, uint8_t nextLinkId)
 {
     NS_LOG_FUNCTION(this << currLinkId << nextLinkId);
 
-    if (!m_switchAuxPhy)
+    if (m_switchAuxPhy)
     {
-        // record that the main PHY will have to switch back to its current link
-        m_linkIdForMainPhyAfterTxop = currLinkId;
-        m_auxPhyToReconnect = GetStaMac()->GetWifiPhy(nextLinkId);
+        // switch channel on Aux PHY so that it operates on the link on which the main PHY was
+        // operating
+        SwitchAuxPhy(nextLinkId, currLinkId);
         return;
     }
 
-    // switch channel on Aux PHY so that it operates on the link on which the main PHY was operating
-    SwitchAuxPhy(nextLinkId, currLinkId);
+    if (currLinkId != GetMainPhyId())
+    {
+        // the main PHY is leaving a non-primary link, hence an aux PHY needs to be reconnected
+        NS_ASSERT_MSG(
+            m_auxPhyToReconnect,
+            "There should be an aux PHY to reconnect when the main PHY leaves a non-primary link");
+
+        // the Aux PHY is not actually switching (hence no switching delay)
+        GetStaMac()->NotifySwitchingEmlsrLink(m_auxPhyToReconnect, currLinkId, Seconds(0));
+        SetCcaEdThresholdOnLinkSwitch(m_auxPhyToReconnect, currLinkId);
+        m_auxPhyToReconnect = nullptr;
+    }
+
+    if (nextLinkId != GetMainPhyId())
+    {
+        // the main PHY is moving to a non-primary link and the aux PHY does not switch link
+        m_auxPhyToReconnect = GetStaMac()->GetWifiPhy(nextLinkId);
+    }
 }
 
 void
@@ -128,20 +144,8 @@ DefaultEmlsrManager::DoNotifyTxopEnd(uint8_t linkId)
     NS_LOG_FUNCTION(this << linkId);
 
     // switch main PHY to the previous link, if needed
-    if (m_linkIdForMainPhyAfterTxop && linkId != m_linkIdForMainPhyAfterTxop)
+    if (!m_switchAuxPhy && m_auxPhyToReconnect)
     {
-        auto auxPhy = m_auxPhyToReconnect;
-
-        // lambda to switch the main PHY back to its previous link and reconnect the aux PHY to
-        // its original link
-        auto restorePhys = [=, this]() {
-            SwitchMainPhy(*m_linkIdForMainPhyAfterTxop, false);
-            // the Aux PHY is not actually switching (hence no switching delay)
-            GetStaMac()->NotifySwitchingEmlsrLink(auxPhy, linkId, Seconds(0));
-            SetCcaEdThresholdOnLinkSwitch(auxPhy, linkId);
-            m_linkIdForMainPhyAfterTxop.reset();
-        };
-
         auto mainPhy = GetStaMac()->GetDevice()->GetPhy(m_mainPhyId);
 
         // the main PHY may be switching at the end of a TXOP when, e.g., the main PHY starts
@@ -151,15 +155,18 @@ DefaultEmlsrManager::DoNotifyTxopEnd(uint8_t linkId)
         // a new channel switch.
         if (!mainPhy->IsStateSwitching())
         {
-            restorePhys();
+            SwitchMainPhy(GetMainPhyId(), false);
         }
         else
         {
-            Simulator::Schedule(mainPhy->GetDelayUntilIdle(), restorePhys);
+            Simulator::Schedule(mainPhy->GetDelayUntilIdle(),
+                                &DefaultEmlsrManager::SwitchMainPhy,
+                                this,
+                                GetMainPhyId(),
+                                false);
         }
         return;
     }
-    m_linkIdForMainPhyAfterTxop.reset();
 }
 
 } // namespace ns3
diff --git a/src/wifi/model/eht/default-emlsr-manager.h b/src/wifi/model/eht/default-emlsr-manager.h
index 0d989f1b4..d62ee06ae 100644
--- a/src/wifi/model/eht/default-emlsr-manager.h
+++ b/src/wifi/model/eht/default-emlsr-manager.h
@@ -58,9 +58,6 @@ class DefaultEmlsrManager : public EmlsrManager
 
     bool m_switchAuxPhy; /**< whether Aux PHY should switch channel to operate on the link on which
                               the Main PHY was operating before moving to the link of the Aux PHY */
-    std::optional<uint8_t>
-        m_linkIdForMainPhyAfterTxop;  //!< ID of the link the main PHY has to switch to once
-                                      //!< the current TXOP terminates
     Ptr<WifiPhy> m_auxPhyToReconnect; //!< Aux PHY the ChannelAccessManager of the link on which
                                       //!< the main PHY is operating has to connect a listener to
                                       //!< when the main PHY is back operating on its previous link