1use rustfft::{num_complex::Complex32, FftPlanner};
4
5use crate::config::{Modulation, OfdmConfig, PassbandMode};
6use crate::equalizer::{
7 equalize_symbol_with_pilots, equalizer_initial_channel, equalizer_refresh_channel,
8 equalizer_reset_tracking, EqualizerTrackingState,
9};
10use crate::packet::{
11 bits_to_bytes, build_packet_bytes, bytes_to_bits, fec_decode_bits, fec_encode_bits,
12 fec_encoded_bits_len, parse_packet_bytes, PacketInfo,
13};
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub(crate) enum PacketSymbolKind {
17 Training,
18 Data,
19}
20
21pub fn encode_single_packet_baseband(payload: &[u8], cfg: &OfdmConfig) -> Vec<Complex32> {
22 let pkt_bytes = build_packet_bytes(payload, 0, 1, cfg);
23 tx_one_packet_baseband(&pkt_bytes, cfg)
24}
25
26pub fn decode_packet_baseband(rbb: &[Complex32], cfg: &OfdmConfig) -> Option<Vec<u8>> {
27 decode_packet_info_baseband(rbb, cfg).map(|p| p.payload)
28}
29
30pub fn expected_single_packet_data_symbols(payload: &[u8], cfg: &OfdmConfig) -> Vec<Complex32> {
31 let pkt_bytes = build_packet_bytes(payload, 0, 1, cfg);
32 expected_packet_data_symbols(&pkt_bytes, cfg)
33}
34
35pub fn recover_single_packet_data_symbols(
36 rbb: &[Complex32],
37 cfg: &OfdmConfig,
38) -> Option<Vec<Complex32>> {
39 equalized_data_symbols_baseband(rbb, cfg)
40}
41
42pub(crate) fn recover_decided_packet_bytes_baseband(
43 rbb: &[Complex32],
44 cfg: &OfdmConfig,
45) -> Option<Vec<u8>> {
46 let rx_syms = equalized_data_symbols_baseband(rbb, cfg)?;
47 let bits = demap_bits(&rx_syms, cfg.modulation);
48 let raw_bits = fec_decode_bits(&bits, cfg.fec_mode);
49 Some(bits_to_bytes(&raw_bits))
50}
51
52pub(crate) fn tx_one_packet_baseband(pkt_bytes: &[u8], cfg: &OfdmConfig) -> Vec<Complex32> {
53 let (used_bins, pilot_bins, data_bins) = ofdm_bin_plan(cfg);
54 let raw_bits = bytes_to_bits(pkt_bytes);
55 let bits = fec_encode_bits(&raw_bits, cfg.fec_mode);
56 let mut payload_syms = map_bits(&bits, cfg.modulation);
57 let syms_per_ofdm = data_bins.len();
58 if syms_per_ofdm == 0 {
59 return Vec::new();
60 }
61 let n_data = payload_syms.len().div_ceil(syms_per_ofdm);
62 payload_syms.resize(n_data * syms_per_ofdm, Complex32::new(0.0, 0.0));
63
64 let mut xbb = Vec::<Complex32>::new();
65 let sync_half = known_sync_half(cfg);
66 xbb.extend_from_slice(&sync_half);
67 xbb.extend_from_slice(&sync_half);
68
69 let train_time = training_symbol_time_domain(&used_bins, cfg);
70 append_cp_symbol(&mut xbb, &train_time, cfg.ncp);
71
72 let symbol_plan = packet_symbol_plan(n_data, cfg);
73 let mut data_idx = 0usize;
74 let mut data_symbol_idx = 0usize;
75 for kind in symbol_plan {
76 if kind == PacketSymbolKind::Training {
77 append_cp_symbol(&mut xbb, &train_time, cfg.ncp);
78 continue;
79 }
80 let mut x = vec![Complex32::new(0.0, 0.0); cfg.nfft];
81 for (k, &bin) in data_bins.iter().enumerate() {
82 x[bin] = payload_syms[data_idx * syms_per_ofdm + k];
83 }
84 if !pilot_bins.is_empty() {
85 let pref = known_pilot_symbols(pilot_bins.len(), data_symbol_idx + 1);
86 for (k, &bin) in pilot_bins.iter().enumerate() {
87 x[bin] = pref[k];
88 }
89 }
90 let xt = ifft(&x);
91 append_cp_symbol(&mut xbb, &xt, cfg.ncp);
92 data_idx += 1;
93 data_symbol_idx += 1;
94 }
95
96 xbb
97}
98
99pub(crate) fn decode_packet_info_baseband(
100 rbb: &[Complex32],
101 cfg: &OfdmConfig,
102) -> Option<PacketInfo> {
103 let rx_syms = equalized_data_symbols_baseband(rbb, cfg)?;
104 recover_packet_from_symbols(&rx_syms, cfg)
105}
106
107fn equalized_data_symbols_baseband(rbb: &[Complex32], cfg: &OfdmConfig) -> Option<Vec<Complex32>> {
108 let (used_bins, pilot_bins, data_bins) = ofdm_bin_plan(cfg);
109 let n_data_carriers = data_bins.len();
110 if n_data_carriers == 0 {
111 return None;
112 }
113 let xsync_len = 2 * cfg.sync_half_len;
114 let train_len = cfg.nfft + cfg.ncp;
115 if rbb.len() < xsync_len + train_len {
116 return None;
117 }
118
119 let train_start = xsync_len;
120 let train_no_cp = &rbb[train_start + cfg.ncp..train_start + cfg.ncp + cfg.nfft];
121 let ytrain = fft(train_no_cp);
122
123 let train_known = known_training_symbols(used_bins.len(), cfg.modulation);
124 let mut hest = equalizer_initial_channel(cfg, &ytrain, &used_bins, &train_known);
125 let mut eq_state = EqualizerTrackingState::default();
126 equalizer_reset_tracking(&mut eq_state);
127
128 let data_start = xsync_len + train_len;
129 let max_payload_bytes = cfg.packet_payload_bytes + 16;
130 let max_bits = fec_encoded_bits_len(max_payload_bytes * 8, cfg.fec_mode);
131 let max_data_ofdm = max_bits.div_ceil(n_data_carriers * cfg.modulation.bits_per_symbol()) + 2;
132 let symbol_plan = packet_symbol_plan(max_data_ofdm, cfg);
133 let mut rx_syms = Vec::<Complex32>::new();
134 let mut data_symbol_idx = 0usize;
135
136 for (sym_idx, kind) in symbol_plan.into_iter().enumerate() {
137 let s0 = data_start + sym_idx * (cfg.nfft + cfg.ncp);
138 let s1 = s0 + cfg.nfft + cfg.ncp;
139 if s1 > rbb.len() {
140 break;
141 }
142 let y = fft(&rbb[s0 + cfg.ncp..s0 + cfg.ncp + cfg.nfft]);
143 if kind == PacketSymbolKind::Training {
144 equalizer_refresh_channel(cfg, &mut hest, &y, &used_bins, &train_known);
145 equalizer_reset_tracking(&mut eq_state);
146 continue;
147 }
148 let pref = known_pilot_symbols(pilot_bins.len(), data_symbol_idx + 1);
149 let xeq_used = equalize_symbol_with_pilots(
150 cfg,
151 &mut eq_state,
152 &y,
153 &used_bins,
154 &pilot_bins,
155 &pref,
156 &hest,
157 );
158 for dbin in &data_bins {
159 if let Some(pos) = used_bins.iter().position(|b| b == dbin) {
160 rx_syms.push(xeq_used[pos]);
161 }
162 }
163 data_symbol_idx += 1;
164 }
165
166 Some(rx_syms)
167}
168
169fn expected_packet_data_symbols(pkt_bytes: &[u8], cfg: &OfdmConfig) -> Vec<Complex32> {
170 let (_, _, data_bins) = ofdm_bin_plan(cfg);
171 let syms_per_ofdm = data_bins.len();
172 if syms_per_ofdm == 0 {
173 return Vec::new();
174 }
175 let bits = bytes_to_bits(pkt_bytes);
176 let coded_bits = fec_encode_bits(&bits, cfg.fec_mode);
177 let mut payload_syms = map_bits(&coded_bits, cfg.modulation);
178 let n_data = payload_syms.len().div_ceil(syms_per_ofdm);
179 payload_syms.resize(n_data * syms_per_ofdm, Complex32::new(0.0, 0.0));
180 payload_syms
181}
182
183pub(crate) fn packet_symbol_plan(n_data_symbols: usize, cfg: &OfdmConfig) -> Vec<PacketSymbolKind> {
184 let mut plan = Vec::with_capacity(
185 n_data_symbols
186 + cfg
187 .retrain_interval_data_symbols
188 .map(|intv| if intv > 0 { n_data_symbols / intv } else { 0 })
189 .unwrap_or(0)
190 + usize::from(cfg.terminal_training_symbol),
191 );
192 for data_idx in 0..n_data_symbols {
193 if let Some(interval) = cfg.retrain_interval_data_symbols {
194 if interval > 0 && data_idx > 0 && data_idx % interval == 0 {
195 plan.push(PacketSymbolKind::Training);
196 }
197 }
198 plan.push(PacketSymbolKind::Data);
199 }
200 if cfg.terminal_training_symbol {
201 plan.push(PacketSymbolKind::Training);
202 }
203 plan
204}
205
206pub(crate) fn training_symbol_time_domain(used_bins: &[usize], cfg: &OfdmConfig) -> Vec<Complex32> {
207 let train_known = known_training_symbols(used_bins.len(), cfg.modulation);
208 let mut xtrain = vec![Complex32::new(0.0, 0.0); cfg.nfft];
209 for (k, &bin) in used_bins.iter().enumerate() {
210 xtrain[bin] = train_known[k];
211 }
212 ifft(&xtrain)
213}
214
215pub(crate) fn map_bits(bits: &[u8], modulation: Modulation) -> Vec<Complex32> {
216 match modulation {
217 Modulation::Bpsk => bits
218 .iter()
219 .map(|&b| {
220 if b == 0 {
221 Complex32::new(1.0, 0.0)
222 } else {
223 Complex32::new(-1.0, 0.0)
224 }
225 })
226 .collect(),
227 Modulation::Qpsk => {
228 let mut out = Vec::with_capacity(bits.len().div_ceil(2));
229 let mut i = 0;
230 while i < bits.len() {
231 let b0 = bits[i];
232 let b1 = if i + 1 < bits.len() { bits[i + 1] } else { 0 };
233 let re = if b0 == 0 { 1.0 } else { -1.0 };
234 let im = if b1 == 0 { 1.0 } else { -1.0 };
235 out.push(Complex32::new(re, im) * (1.0 / 2.0f32.sqrt()));
236 i += 2;
237 }
238 out
239 }
240 }
241}
242
243pub(crate) fn demap_bits(syms: &[Complex32], modulation: Modulation) -> Vec<u8> {
244 match modulation {
245 Modulation::Bpsk => syms.iter().map(|s| (s.re < 0.0) as u8).collect(),
246 Modulation::Qpsk => {
247 let mut bits = Vec::with_capacity(syms.len() * 2);
248 for s in syms {
249 bits.push((s.re < 0.0) as u8);
250 bits.push((s.im < 0.0) as u8);
251 }
252 bits
253 }
254 }
255}
256
257fn recover_packet_from_symbols(syms: &[Complex32], cfg: &OfdmConfig) -> Option<PacketInfo> {
258 let bits = demap_bits(syms, cfg.modulation);
259 let raw_bits = fec_decode_bits(&bits, cfg.fec_mode);
260 let bytes = bits_to_bytes(&raw_bits);
261 parse_packet_bytes(&bytes).map(|(pkt, _)| pkt)
262}
263
264pub(crate) fn known_training_symbols(n: usize, modulation: Modulation) -> Vec<Complex32> {
265 match modulation {
266 Modulation::Bpsk => (0..n)
267 .map(|i| {
268 if i % 2 == 0 {
269 Complex32::new(1.0, 0.0)
270 } else {
271 Complex32::new(-1.0, 0.0)
272 }
273 })
274 .collect(),
275 Modulation::Qpsk => {
276 let base = [
277 Complex32::new(1.0, 1.0),
278 Complex32::new(1.0, -1.0),
279 Complex32::new(-1.0, 1.0),
280 Complex32::new(-1.0, -1.0),
281 ];
282 (0..n)
283 .map(|i| base[i % 4] * (1.0 / 2.0f32.sqrt()))
284 .collect()
285 }
286 }
287}
288
289pub(crate) fn known_pilot_symbols(n: usize, sym_idx: usize) -> Vec<Complex32> {
290 (0..n)
291 .map(|k| {
292 let m = ((sym_idx - 1) + k) % 4;
293 Complex32::from_polar(1.0, std::f32::consts::FRAC_PI_2 * (m as f32))
294 })
295 .collect()
296}
297
298pub(crate) fn ofdm_bin_plan(cfg: &OfdmConfig) -> (Vec<usize>, Vec<usize>, Vec<usize>) {
299 let (used_bins, pilot_candidates) = resolve_bins_with_base_freq(cfg);
300 let pilots_on = cfg
301 .use_pilots
302 .unwrap_or(matches!(cfg.modulation, Modulation::Qpsk));
303 let pilot_bins = if pilots_on {
304 let mut pilots = pilot_candidates
305 .iter()
306 .copied()
307 .filter(|b| used_bins.contains(b))
308 .collect::<Vec<_>>();
309 if let Some(n) = cfg.num_pilots {
310 pilots.truncate(n.min(pilots.len()));
311 }
312 pilots
313 } else {
314 Vec::new()
315 };
316 let data_bins = used_bins
317 .iter()
318 .copied()
319 .filter(|b| !pilot_bins.contains(b))
320 .collect::<Vec<_>>();
321 (used_bins, pilot_bins, data_bins)
322}
323
324fn dedup_stable(v: Vec<usize>) -> Vec<usize> {
325 let mut out = Vec::with_capacity(v.len());
326 for x in v {
327 if !out.contains(&x) {
328 out.push(x);
329 }
330 }
331 out
332}
333
334fn resolve_bins_with_base_freq(cfg: &OfdmConfig) -> (Vec<usize>, Vec<usize>) {
335 let mut used_bins = cfg.used_bins.clone();
336 let mut pilot_candidates = if cfg.pilot_bins.is_empty() {
337 used_bins.clone()
338 } else {
339 cfg.pilot_bins.clone()
340 };
341
342 if let Some(base_hz) = cfg.base_freq_hz {
343 let numerology_fs = match cfg.passband_mode {
344 PassbandMode::Legacy => cfg.fs,
345 PassbandMode::Iq => cfg.fs_baseband,
346 };
347 let df = numerology_fs / (cfg.nfft as f32);
348 if df > 0.0 && !used_bins.is_empty() {
349 let target_bin = (base_hz / df).round() as isize;
350 let current_min = *used_bins.iter().min().unwrap_or(&1) as isize;
351 let shift = target_bin.max(1) - current_min;
352 let kmax = (cfg.nfft / 2).saturating_sub(1) as isize;
353 used_bins = used_bins
354 .iter()
355 .map(|&b| ((b as isize + shift).clamp(1, kmax)) as usize)
356 .collect();
357 pilot_candidates = pilot_candidates
358 .iter()
359 .map(|&b| ((b as isize + shift).clamp(1, kmax)) as usize)
360 .collect();
361 }
362 }
363
364 (dedup_stable(used_bins), dedup_stable(pilot_candidates))
365}
366
367pub(crate) fn known_sync_half(cfg: &OfdmConfig) -> Vec<Complex32> {
368 let l = cfg.sync_half_len;
369 if l == 0 {
370 return Vec::new();
371 }
372
373 let (used_bins, _pilot_bins, _data_bins) = ofdm_bin_plan(cfg);
374 if used_bins.is_empty() {
375 return vec![Complex32::new(0.0, 0.0); l];
376 }
377
378 let mut out = Vec::with_capacity(l);
379 for n in 0..l {
380 let mut acc = Complex32::new(0.0, 0.0);
381 for (i, &bin) in used_bins.iter().enumerate() {
382 let phase0 = std::f32::consts::FRAC_PI_2 * (((3 * i + 1) % 4) as f32);
383 let phase =
384 phase0 + 2.0 * std::f32::consts::PI * (bin as f32) * (n as f32) / (cfg.nfft as f32);
385 acc += Complex32::from_polar(1.0, phase);
386 }
387 out.push(acc);
388 }
389
390 let scale = 1.0 / (used_bins.len() as f32).sqrt();
391 for v in &mut out {
392 *v *= scale;
393 }
394 out
395}
396
397pub(crate) fn append_cp_symbol(out: &mut Vec<Complex32>, x: &[Complex32], ncp: usize) {
398 out.extend_from_slice(&x[x.len() - ncp..]);
399 out.extend_from_slice(x);
400}
401
402pub(crate) fn ifft(x: &[Complex32]) -> Vec<Complex32> {
403 let mut planner = FftPlanner::<f32>::new();
404 let fft = planner.plan_fft_inverse(x.len());
405 let mut buf = x.to_vec();
406 fft.process(&mut buf);
407 let scale = 1.0 / (x.len() as f32).sqrt();
408 for v in &mut buf {
409 *v *= scale;
410 }
411 buf
412}
413
414pub(crate) fn fft(x: &[Complex32]) -> Vec<Complex32> {
415 let mut planner = FftPlanner::<f32>::new();
416 let fft = planner.plan_fft_forward(x.len());
417 let mut buf = x.to_vec();
418 fft.process(&mut buf);
419 let scale = 1.0 / (x.len() as f32).sqrt();
420 for v in &mut buf {
421 *v *= scale;
422 }
423 buf
424}
425
426#[cfg(test)]
427mod tests {
428 use super::*;
429
430 #[test]
431 fn decode_single_packet_baseband() {
432 let cfg = OfdmConfig::default();
433 let payload: Vec<u8> = (100..120).collect();
434 let xbb = encode_single_packet_baseband(&payload, &cfg);
435 let out = decode_packet_baseband(&xbb, &cfg).expect("baseband decode failed");
436 assert_eq!(out, payload);
437 }
438
439 #[test]
440 fn recover_single_packet_bpsk_symbols() {
441 let cfg = OfdmConfig::default();
442 let payload: Vec<u8> = (0..24).collect();
443 let xbb = encode_single_packet_baseband(&payload, &cfg);
444 let got = recover_single_packet_data_symbols(&xbb, &cfg).expect("symbol recovery failed");
445 let expect = expected_single_packet_data_symbols(&payload, &cfg);
446 assert_eq!(got.len(), expect.len());
447 for (g, e) in got.iter().zip(expect.iter()) {
448 assert!((*g - *e).norm() < 0.05, "got={g:?} expect={e:?}");
449 }
450 }
451
452 #[test]
453 fn recover_single_packet_qpsk_symbols() {
454 let mut cfg = OfdmConfig::default();
455 cfg.modulation = Modulation::Qpsk;
456 cfg.use_pilots = Some(true);
457 let payload: Vec<u8> = (0..24).map(|x| x ^ 0x5a).collect();
458 let xbb = encode_single_packet_baseband(&payload, &cfg);
459 let got = recover_single_packet_data_symbols(&xbb, &cfg).expect("symbol recovery failed");
460 let expect = expected_single_packet_data_symbols(&payload, &cfg);
461 assert_eq!(got.len(), expect.len());
462 for (g, e) in got.iter().zip(expect.iter()) {
463 assert!((*g - *e).norm() < 0.08, "got={g:?} expect={e:?}");
464 }
465 }
466
467 #[test]
468 fn pilot_count_capped_by_num_pilots() {
469 let mut cfg = OfdmConfig::default();
470 cfg.modulation = Modulation::Qpsk;
471 cfg.use_pilots = Some(true);
472 cfg.used_bins = vec![2, 3, 4, 5];
473 cfg.pilot_bins = vec![2, 4, 5];
474 cfg.num_pilots = Some(2);
475 let (_used, pilots, data) = ofdm_bin_plan(&cfg);
476 assert_eq!(pilots, vec![2, 4]);
477 assert_eq!(data, vec![3, 5]);
478 }
479
480 #[test]
481 fn base_frequency_shifts_bins() {
482 let mut cfg = OfdmConfig::default();
483 cfg.base_freq_hz = Some(2_000.0);
484 cfg.use_pilots = Some(false);
485 let (used, pilots, data) = ofdm_bin_plan(&cfg);
486 assert_eq!(
487 used,
488 vec![93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108]
489 );
490 assert!(pilots.is_empty());
491 assert_eq!(
492 data,
493 vec![93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108]
494 );
495 }
496
497 #[test]
498 fn decode_single_packet_baseband_with_hamming_fec() {
499 let mut cfg = OfdmConfig::default();
500 cfg.fec_mode = crate::config::FecMode::Hamming74;
501 let payload: Vec<u8> = (0..24).map(|x| x ^ 0x33).collect();
502 let xbb = encode_single_packet_baseband(&payload, &cfg);
503 let out = decode_packet_baseband(&xbb, &cfg).expect("baseband decode failed");
504 assert_eq!(out, payload);
505 }
506}