From 8e6594165015b15b889f7f6db04b73e2972a018d Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Wed, 20 Dec 2023 12:19:21 +0100
Subject: [PATCH] wifi: Extend WifiHelper::AssignStreams to support MLDs

---
 examples/wireless/examples-to-run.py |  2 +-
 src/wifi/helper/wifi-helper.cc       | 25 +++++++++++++++----------
 src/wifi/test/wifi-mlo-test.cc       |  2 +-
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/examples/wireless/examples-to-run.py b/examples/wireless/examples-to-run.py
index 2e2683c37..a8cb35952 100755
--- a/examples/wireless/examples-to-run.py
+++ b/examples/wireless/examples-to-run.py
@@ -180,7 +180,7 @@ cpp_examples = [
         "True",
     ),
     (
-        "wifi-eht-network --simulationTime=0.2 --udp=1 --downlink=0 --useRts=0 --nStations=5 --dlAckType=AGGR-MU-BAR --enableUlOfdma=1 --enableBsrp=1 --mpduBufferSize=1024 --mcs=5 --muSchedAccessReqInterval=45ms --frequency2=6 --minExpectedThroughput=70 --maxExpectedThroughput=715",
+        "wifi-eht-network --simulationTime=0.2 --udp=1 --downlink=0 --useRts=0 --nStations=4 --dlAckType=AGGR-MU-BAR --enableUlOfdma=1 --enableBsrp=1 --mpduBufferSize=1024 --mcs=4 --muSchedAccessReqInterval=45ms --frequency2=6 --minExpectedThroughput=50 --maxExpectedThroughput=415",
         "True",
         "True",
     ),
diff --git a/src/wifi/helper/wifi-helper.cc b/src/wifi/helper/wifi-helper.cc
index f7f409d56..800f6cdaa 100644
--- a/src/wifi/helper/wifi-helper.cc
+++ b/src/wifi/helper/wifi-helper.cc
@@ -992,40 +992,45 @@ WifiHelper::AssignStreams(NetDeviceContainer c, int64_t stream)
     for (auto i = c.Begin(); i != c.End(); ++i)
     {
         netDevice = (*i);
-        Ptr<WifiNetDevice> wifi = DynamicCast<WifiNetDevice>(netDevice);
-        if (wifi)
+        if (auto wifi = DynamicCast<WifiNetDevice>(netDevice))
         {
             // Handle any random numbers in the PHY objects.
-            currentStream += wifi->GetPhy()->AssignStreams(currentStream);
+            for (auto& phy : wifi->GetPhys())
+            {
+                currentStream += phy->AssignStreams(currentStream);
+            }
 
             // Handle any random numbers in the station managers.
-            currentStream += wifi->GetRemoteStationManager()->AssignStreams(currentStream);
+            for (auto& manager : wifi->GetRemoteStationManagers())
+            {
+                currentStream += manager->AssignStreams(currentStream);
+            }
 
             // Handle any random numbers in the MAC objects.
-            Ptr<WifiMac> mac = wifi->GetMac();
+            auto mac = wifi->GetMac();
             PointerValue ptr;
             if (!mac->GetQosSupported())
             {
                 mac->GetAttribute("Txop", ptr);
-                Ptr<Txop> txop = ptr.Get<Txop>();
+                auto txop = ptr.Get<Txop>();
                 currentStream += txop->AssignStreams(currentStream);
             }
             else
             {
                 mac->GetAttribute("VO_Txop", ptr);
-                Ptr<QosTxop> vo_txop = ptr.Get<QosTxop>();
+                auto vo_txop = ptr.Get<QosTxop>();
                 currentStream += vo_txop->AssignStreams(currentStream);
 
                 mac->GetAttribute("VI_Txop", ptr);
-                Ptr<QosTxop> vi_txop = ptr.Get<QosTxop>();
+                auto vi_txop = ptr.Get<QosTxop>();
                 currentStream += vi_txop->AssignStreams(currentStream);
 
                 mac->GetAttribute("BE_Txop", ptr);
-                Ptr<QosTxop> be_txop = ptr.Get<QosTxop>();
+                auto be_txop = ptr.Get<QosTxop>();
                 currentStream += be_txop->AssignStreams(currentStream);
 
                 mac->GetAttribute("BK_Txop", ptr);
-                Ptr<QosTxop> bk_txop = ptr.Get<QosTxop>();
+                auto bk_txop = ptr.Get<QosTxop>();
                 currentStream += bk_txop->AssignStreams(currentStream);
             }
 
diff --git a/src/wifi/test/wifi-mlo-test.cc b/src/wifi/test/wifi-mlo-test.cc
index 72c2d0e74..437121666 100644
--- a/src/wifi/test/wifi-mlo-test.cc
+++ b/src/wifi/test/wifi-mlo-test.cc
@@ -613,7 +613,7 @@ void
 MultiLinkOperationsTestBase::DoSetup()
 {
     RngSeedManager::SetSeed(1);
-    RngSeedManager::SetRun(2);
+    RngSeedManager::SetRun(3);
     int64_t streamNumber = 30;
 
     NodeContainer wifiApNode;