Skip to main content

rlx_tsac/
parity.rs

1//! Parity + timing bench: native [`TsacCodec`] vs the public Bellard Linux `tsac` binary.
2
3use 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    /// Bellard `-f` / native fast path (no transformer).
17    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}