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