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