1use crate::TsacBackendKind;
4use crate::audio::{self, PcmCompare, prepare_tsac_wav};
5use crate::codec::{EncodeStats, TsacCodec, TsacOptions};
6use crate::download::{ensure_tsac, resolve_install_dir};
7use crate::platform::tsac_binary_supported;
8use anyhow::{Result, bail};
9use rlx_runtime::Device;
10use std::path::{Path, PathBuf};
11use std::time::Instant;
12
13#[derive(Debug, Clone)]
14pub struct ParityOptions {
15 pub quality: u8,
16 pub fast: bool,
18 pub native_device: Device,
19 pub bellard_cuda: bool,
20 pub min_correlation: f32,
21 pub max_mse: f32,
22}
23
24impl Default for ParityOptions {
25 fn default() -> Self {
26 Self {
27 quality: 9,
28 fast: true,
29 native_device: Device::Cpu,
30 bellard_cuda: false,
31 min_correlation: env_f32("RLX_TSAC_PARITY_MIN_CORR", 0.92),
32 max_mse: env_f32("RLX_TSAC_PARITY_MAX_MSE", 0.0025),
33 }
34 }
35}
36
37#[derive(Debug, Clone)]
38pub struct EngineRoundtrip {
39 pub encode: EncodeStats,
40 pub decode_ms: f64,
41 pub pcm: Vec<f32>,
42 pub tsac_path: PathBuf,
43 pub wav_path: PathBuf,
44}
45
46#[derive(Debug, Clone)]
47pub struct ParityBenchReport {
48 pub input_wav: PathBuf,
49 pub prepared_wav: PathBuf,
50 pub channels: u16,
51 pub native: EngineRoundtrip,
52 pub bellard: EngineRoundtrip,
53 pub roundtrip_pcm: PcmCompare,
54 pub native_enc_bellard_dec: PcmCompare,
55 pub bellard_enc_native_dec: PcmCompare,
56}
57
58impl ParityBenchReport {
59 pub fn passes(&self, opts: &ParityOptions) -> bool {
60 self.roundtrip_pcm
61 .passes(opts.min_correlation, opts.max_mse)
62 && self
63 .native_enc_bellard_dec
64 .passes(opts.min_correlation, opts.max_mse)
65 && self
66 .bellard_enc_native_dec
67 .passes(opts.min_correlation, opts.max_mse)
68 }
69
70 pub fn print_summary(&self, opts: &ParityOptions) {
71 let pass = self.passes(opts);
72 eprintln!(
73 "\n=== TSAC parity (native vs Bellard binary) ===\n\
74 input: {}\n\
75 prepared @ 44.1 kHz: {} ({} ch)\n\
76 quality={} fast={}\n",
77 self.input_wav.display(),
78 self.prepared_wav.display(),
79 self.channels,
80 opts.quality,
81 opts.fast
82 );
83 print_engine("native", &self.native);
84 print_engine("bellard", &self.bellard);
85 print_compare(
86 "roundtrip PCM (native vs bellard)",
87 &self.roundtrip_pcm,
88 opts,
89 );
90 print_compare(
91 "native encode → bellard decode",
92 &self.native_enc_bellard_dec,
93 opts,
94 );
95 print_compare(
96 "bellard encode → native decode",
97 &self.bellard_enc_native_dec,
98 opts,
99 );
100 eprintln!("overall: {}", if pass { "PASS" } else { "FAIL" });
101 }
102}
103
104pub fn bench_bellard_parity(
105 install_dir: impl AsRef<Path>,
106 in_wav: impl AsRef<Path>,
107 opts: &ParityOptions,
108) -> Result<ParityBenchReport> {
109 if !tsac_binary_supported() {
110 bail!(
111 "Bellard parity bench requires Linux x86_64 (public tsac binary); \
112 run on a Linux host or set RLX_TSAC_PARITY=1 in CI"
113 );
114 }
115 let install_dir = install_dir.as_ref();
116 ensure_tsac(install_dir)?;
117
118 let in_wav = in_wav.as_ref().to_path_buf();
119 let tag = std::process::id();
120 let tmp = std::env::temp_dir().join(format!("rlx-tsac-parity-{tag}"));
121 std::fs::create_dir_all(&tmp)?;
122
123 let prepared = tmp.join("input_44100.wav");
124 let channels = prepare_tsac_wav(&in_wav, &prepared)?;
125
126 let base_opts = TsacOptions {
127 quality: Some(opts.quality),
128 fast: opts.fast,
129 verbose: false,
130 separate_stereo: false,
131 channels: None,
132 device: Device::Cpu,
133 backend: TsacBackendKind::Auto,
134 };
135
136 let native_opts = TsacOptions {
137 device: opts.native_device,
138 backend: TsacBackendKind::Native,
139 ..base_opts.clone()
140 };
141 let bellard_opts = TsacOptions {
142 device: if opts.bellard_cuda {
143 Device::Cuda
144 } else {
145 Device::Cpu
146 },
147 backend: TsacBackendKind::Bellard,
148 ..base_opts
149 };
150
151 let native_codec = TsacCodec::open_with_options(install_dir, native_opts)?;
152 let bellard_codec = TsacCodec::open_with_options(install_dir, bellard_opts)?;
153
154 let native = run_engine_roundtrip(&native_codec, "native", &prepared, &tmp, channels)?;
155 let bellard = run_engine_roundtrip(&bellard_codec, "bellard", &prepared, &tmp, channels)?;
156
157 let native_enc_bellard_dec = cross_decode(
158 &bellard_codec,
159 &native.tsac_path,
160 &tmp.join("native_enc_bellard_dec.wav"),
161 )?;
162 let bellard_enc_native_dec = cross_decode(
163 &native_codec,
164 &bellard.tsac_path,
165 &tmp.join("bellard_enc_native_dec.wav"),
166 )?;
167
168 let roundtrip_pcm = PcmCompare::compare(&native.pcm, &bellard.pcm);
169 let native_enc_bellard_dec_cmp = PcmCompare::compare(&native.pcm, &native_enc_bellard_dec);
170 let bellard_enc_native_dec_cmp = PcmCompare::compare(&bellard.pcm, &bellard_enc_native_dec);
171
172 Ok(ParityBenchReport {
173 input_wav: in_wav,
174 prepared_wav: prepared,
175 channels,
176 native,
177 bellard,
178 roundtrip_pcm,
179 native_enc_bellard_dec: native_enc_bellard_dec_cmp,
180 bellard_enc_native_dec: bellard_enc_native_dec_cmp,
181 })
182}
183
184pub fn bench_bellard_parity_default(in_wav: impl AsRef<Path>) -> Result<ParityBenchReport> {
185 let dir = resolve_install_dir(None);
186 bench_bellard_parity(dir, in_wav, &ParityOptions::default())
187}
188
189fn run_engine_roundtrip(
190 codec: &TsacCodec,
191 label: &str,
192 prepared_wav: &Path,
193 tmp: &Path,
194 _channels: u16,
195) -> Result<EngineRoundtrip> {
196 let tsac_path = tmp.join(format!("{label}.tsac"));
197 let wav_path = tmp.join(format!("{label}_roundtrip.wav"));
198 let encode = codec.encode(prepared_wav, &tsac_path)?;
199 let t0 = Instant::now();
200 codec.decode(&tsac_path, &wav_path)?;
201 let decode_ms = t0.elapsed().as_secs_f64() * 1000.0;
202 let pcm = audio::load_pcm_from_wav(&wav_path)?;
203 Ok(EngineRoundtrip {
204 encode,
205 decode_ms,
206 pcm,
207 tsac_path,
208 wav_path,
209 })
210}
211
212fn cross_decode(codec: &TsacCodec, in_tsac: &Path, out_wav: &Path) -> Result<Vec<f32>> {
213 codec.decode(in_tsac, out_wav)?;
214 audio::load_pcm_from_wav(out_wav)
215}
216
217fn print_engine(label: &str, rt: &EngineRoundtrip) {
218 eprintln!(
219 "{label}: encode {:.1} ms, decode {:.1} ms, {} bytes compressed, {} samples",
220 rt.encode.encode_ms,
221 rt.decode_ms,
222 rt.encode.output_bytes,
223 rt.pcm.len()
224 );
225}
226
227fn print_compare(label: &str, cmp: &PcmCompare, opts: &ParityOptions) {
228 let ok = cmp.passes(opts.min_correlation, opts.max_mse);
229 eprintln!(
230 "{label}: corr={:.4} mse={:.6} max_abs={:.5} n={} [{}]",
231 cmp.correlation,
232 cmp.mse,
233 cmp.max_abs,
234 cmp.samples,
235 if ok { "ok" } else { "over tolerance" }
236 );
237}
238
239fn env_f32(name: &str, default: f32) -> f32 {
240 std::env::var(name)
241 .ok()
242 .and_then(|s| s.parse().ok())
243 .unwrap_or(default)
244}