1use std::fmt;
4use std::rc::Rc;
5
6use gwr_models::data_frame::DataFrame;
7use gwr_models::fabric::FabricConfig;
8use gwr_track::entity::Entity;
9use rand::SeedableRng;
10use rand::seq::{IteratorRandom, SliceRandom};
11use rand_xoshiro::Xoshiro256PlusPlus;
12use serde::Serialize;
13
14#[derive(clap::ValueEnum, Clone, Copy, Default, Debug, Serialize, PartialEq)]
15#[serde(rename_all = "kebab-case")]
16pub enum TrafficPattern {
17 #[default]
19 AllToOne,
20
21 Random,
23
24 AllToAllFixed,
26
27 AllToAllSeq,
29
30 AllToAllRandom,
33}
34
35impl fmt::Display for TrafficPattern {
36 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
37 write!(f, "{self:?}")
38 }
39}
40
41pub struct FrameGen {
47 pub entity: Rc<Entity>,
48 config: Rc<FabricConfig>,
49 source_index: usize,
50 dest_index: usize,
51 traffic_pattern: TrafficPattern,
52 overhead_size_bytes: usize,
53 payload_size_bytes: usize,
54 num_send_frames: usize,
55 num_sent_frames: usize,
56 rng: Xoshiro256PlusPlus,
57 next_dests: Vec<usize>,
58}
59
60impl FrameGen {
61 #[expect(clippy::too_many_arguments)]
62 #[must_use]
63 pub fn new(
64 parent: &Rc<Entity>,
65 config: Rc<FabricConfig>,
66 source_index: usize,
67 initial_dest_index: usize,
68 traffic_pattern: TrafficPattern,
69 overhead_size_bytes: usize,
70 payload_size_bytes: usize,
71 num_send_frames: usize,
72 seed: u64,
73 ) -> Self {
74 let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed ^ (source_index as u64));
76 let num_ports = config.num_ports();
77
78 let dest_index = match traffic_pattern {
79 TrafficPattern::Random => (0..num_ports).choose(&mut rng).unwrap(),
80 TrafficPattern::AllToOne
81 | TrafficPattern::AllToAllFixed
82 | TrafficPattern::AllToAllSeq
83 | TrafficPattern::AllToAllRandom => initial_dest_index,
84 };
85
86 let mut next_dests: Vec<usize> = (0..num_ports).collect();
87 next_dests.shuffle(&mut rng);
88
89 Self {
90 entity: Rc::new(Entity::new(parent, format!("gen{source_index}").as_str())),
91 config,
92 source_index,
93 dest_index,
94 traffic_pattern,
95 overhead_size_bytes,
96 payload_size_bytes,
97 num_send_frames,
98 num_sent_frames: 0,
99 rng,
100 next_dests,
101 }
102 }
103
104 #[must_use]
105 fn next_dest(&mut self) -> usize {
106 let num_ports = self.config.max_num_ports();
107 match self.traffic_pattern {
108 TrafficPattern::Random => (0..num_ports).choose(&mut self.rng).unwrap(),
109 TrafficPattern::AllToOne => self.dest_index,
110 TrafficPattern::AllToAllFixed => self.dest_index,
111 TrafficPattern::AllToAllSeq => (self.dest_index + 1) % num_ports,
112 TrafficPattern::AllToAllRandom => {
113 let dest = self.next_dests.pop().unwrap();
114 if self.next_dests.is_empty() {
115 self.next_dests = (0..num_ports).collect();
116 self.next_dests.shuffle(&mut self.rng);
117 }
118 dest
119 }
120 }
121 }
122}
123
124impl Iterator for FrameGen {
125 type Item = DataFrame;
126 fn next(&mut self) -> Option<Self::Item> {
127 let no_valid_packets = (self.dest_index == self.source_index)
129 && (self.traffic_pattern == TrafficPattern::AllToOne
130 || self.traffic_pattern == TrafficPattern::AllToAllFixed);
131
132 if no_valid_packets {
133 return None;
134 }
135
136 if self.num_sent_frames < self.num_send_frames {
137 while self.dest_index == self.source_index {
138 self.dest_index = self.next_dest();
139 }
140
141 let label = (self.num_sent_frames | (self.source_index << 32)) as u64;
142 self.num_sent_frames += 1;
143
144 let dest = self.config.port_indices()[self.dest_index];
147 let frame = Some(
148 DataFrame::new(
149 &self.entity,
150 self.overhead_size_bytes,
151 self.payload_size_bytes,
152 )
153 .set_dest(dest as u64)
154 .set_label(label),
155 );
156
157 self.dest_index = self.next_dest();
158 frame
159 } else {
160 None
161 }
162 }
163}