pito_codecs/opus/enc/
mod.rs

1pub(crate) mod celt;
2pub(crate) mod silk;
3
4use std::io::Write;
5use std::panic::{AssertUnwindSafe, catch_unwind};
6
7use crate::media::{AudioBuffer, Error, Packet, Result};
8use crate::opus::OpusHead;
9use crate::opus::ogg;
10use crate::opus::packet::Bandwidth;
11use crate::opus::range::RangeEncoder;
12
13use celt::{CeltEncoder, FrameParams};
14use silk::SilkEncoder;
15
16pub const FRAME: usize = 960;
17const DELAY: usize = 192;
18const CELT_DELAY: usize = 120;
19const MAX_PACKET: usize = 1276;
20const CHUNK: usize = 256;
21const HISTORY: usize = 1032;
22const SILK_DELAY: usize = 131;
23const SILK_PAD: usize = DELAY + CELT_DELAY - SILK_DELAY;
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum Application {
27    Audio,
28    Voip,
29    LowDelay,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct EncoderConfig {
34    pub rate: u32,
35    pub channels: u16,
36    pub bitrate: u32,
37    pub application: Application,
38    pub vbr: bool,
39}
40
41impl EncoderConfig {
42    pub fn new(channels: u16, bitrate: u32) -> EncoderConfig {
43        EncoderConfig {
44            rate: 48000,
45            channels,
46            bitrate,
47            application: Application::Audio,
48            vbr: true,
49        }
50    }
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54enum Mode {
55    Silk,
56    Celt,
57}
58
59fn mode_for(cfg: &EncoderConfig) -> Mode {
60    match cfg.application {
61        Application::LowDelay => Mode::Celt,
62        Application::Voip if cfg.bitrate <= 40_000 => Mode::Silk,
63        _ if cfg.bitrate < 16_000 => Mode::Silk,
64        _ => Mode::Celt,
65    }
66}
67
68fn silk_bandwidth(bitrate: u32) -> Bandwidth {
69    match bitrate {
70        0..8_000 => Bandwidth::Narrowband,
71        8_000..11_000 => Bandwidth::Mediumband,
72        _ => Bandwidth::Wideband,
73    }
74}
75
76impl Mode {
77    fn pad(self) -> usize {
78        match self {
79            Mode::Silk => SILK_PAD,
80            Mode::Celt => DELAY,
81        }
82    }
83}
84
85pub struct Encoder {
86    cfg: EncoderConfig,
87    mode: Mode,
88    silk: Option<Box<SilkEncoder>>,
89    scratch: Vec<f32>,
90    celt: CeltEncoder,
91    rc: RangeEncoder,
92    pending: Vec<f32>,
93    consumed: usize,
94    frames: u64,
95    input: u64,
96    final_range: u32,
97}
98
99impl Encoder {
100    pub fn new(cfg: EncoderConfig) -> Result<Encoder> {
101        if cfg.rate != 48000 {
102            return Err(Error::Unsupported(format!(
103                "Opus encoder input at {} Hz; resample to 48000 first",
104                cfg.rate
105            )));
106        }
107        if !(1..=2).contains(&cfg.channels) {
108            return Err(Error::Unsupported(format!(
109                "Opus encoder with {} channels",
110                cfg.channels
111            )));
112        }
113        if !(6_000..=510_000).contains(&cfg.bitrate) {
114            return Err(Error::Invalid(format!(
115                "Opus bitrate {} outside 6000..=510000",
116                cfg.bitrate
117            )));
118        }
119        let c = cfg.channels as usize;
120        let mode = mode_for(&cfg);
121        Ok(Encoder {
122            cfg,
123            mode,
124            silk: (mode == Mode::Silk).then(|| Box::new(SilkEncoder::new(c))),
125            scratch: Vec::new(),
126            celt: CeltEncoder::new(c),
127            rc: RangeEncoder::new(MAX_PACKET),
128            pending: vec![0.0; mode.pad() * c],
129            consumed: 0,
130            frames: 0,
131            input: 0,
132            final_range: 0,
133        })
134    }
135
136    pub fn config(&self) -> &EncoderConfig {
137        &self.cfg
138    }
139
140    pub fn pre_skip(&self) -> u16 {
141        (DELAY + CELT_DELAY) as u16
142    }
143
144    pub fn head(&self) -> OpusHead {
145        OpusHead {
146            channels: self.cfg.channels as usize,
147            pre_skip: self.pre_skip(),
148            input_sample_rate: self.cfg.rate,
149            output_gain: 0,
150            mapping_family: 0,
151            streams: 1,
152            coupled_streams: self.cfg.channels as usize - 1,
153            mapping: (0..self.cfg.channels as u8).collect(),
154        }
155    }
156
157    pub fn codec_private(&self) -> Vec<u8> {
158        self.head().to_bytes()
159    }
160
161    pub fn reset(&mut self) {
162        let c = self.cfg.channels as usize;
163        self.celt.reset();
164        if let Some(silk) = self.silk.as_mut() {
165            silk.reset();
166        }
167        self.pending.clear();
168        self.pending.resize(self.mode.pad() * c, 0.0);
169        self.consumed = 0;
170        self.frames = 0;
171        self.input = 0;
172    }
173
174    pub fn final_range(&self) -> u32 {
175        self.final_range
176    }
177
178    pub fn frame_samples(&self) -> usize {
179        FRAME
180    }
181
182    pub fn encode(&mut self, audio: &AudioBuffer) -> Result<Vec<Packet>> {
183        if audio.rate != self.cfg.rate || audio.channels != self.cfg.channels {
184            return Err(Error::Invalid(format!(
185                "audio is {} Hz x{}, encoder expects {} Hz x{}",
186                audio.rate, audio.channels, self.cfg.rate, self.cfg.channels
187            )));
188        }
189        if !audio
190            .samples
191            .len()
192            .is_multiple_of(self.cfg.channels as usize)
193        {
194            return Err(Error::Invalid(format!(
195                "{} samples do not split into {} channels",
196                audio.samples.len(),
197                self.cfg.channels
198            )));
199        }
200        self.input += audio.frames() as u64;
201        self.pending.extend_from_slice(&audio.samples);
202        let mut out = Vec::new();
203        self.drain(&mut out)?;
204        Ok(out)
205    }
206
207    pub fn flush(&mut self) -> Result<Vec<Packet>> {
208        let c = self.cfg.channels as usize;
209        let need = self.input + self.pre_skip() as u64;
210        let mut out = Vec::new();
211        while self.frames * (FRAME as u64) < need {
212            let have = self.pending.len() / c - self.consumed / c;
213            if have < FRAME {
214                self.pending
215                    .resize(self.pending.len() + (FRAME - have) * c, 0.0);
216            }
217            self.drain(&mut out)?;
218        }
219        Ok(out)
220    }
221
222    fn drain(&mut self, out: &mut Vec<Packet>) -> Result<()> {
223        let c = self.cfg.channels as usize;
224        let step = FRAME * c;
225        loop {
226            let avail = (self.pending.len() - self.consumed) / step;
227            if avail == 0 {
228                break;
229            }
230            let at_chunk = self.frames.is_multiple_of(CHUNK as u64);
231            let celt = self.mode == Mode::Celt;
232            if celt && at_chunk && avail >= 2 * CHUNK && rayon::current_num_threads() > 1 {
233                self.encode_chunks(avail / CHUNK, out)?;
234                continue;
235            }
236            if celt && at_chunk && self.frames > 0 {
237                let h = self.consumed.saturating_sub(HISTORY * c);
238                self.celt.prime(&self.pending[h..self.consumed]);
239            }
240            let start = self.consumed;
241            self.consumed += step;
242            let pcm = &self.pending[start..start + step];
243            let (cfg, rc, index) = (&self.cfg, &mut self.rc, self.frames);
244            let (pkt, rng) = match self.silk.as_deref_mut() {
245                Some(silk) => {
246                    let scratch = &mut self.scratch;
247                    catch_unwind(AssertUnwindSafe(|| {
248                        encode_silk(cfg, silk, rc, pcm, index, scratch)
249                    }))
250                }
251                None => {
252                    let celt = &mut self.celt;
253                    catch_unwind(AssertUnwindSafe(|| {
254                        encode_packet(cfg, celt, rc, pcm, index)
255                    }))
256                }
257            }
258            .map_err(|_| crashed())?;
259            self.final_range = rng;
260            self.frames += 1;
261            out.push(pkt);
262        }
263        let keep = HISTORY * c;
264        if self.consumed > keep && (self.consumed - keep) * 2 >= self.pending.len() {
265            self.pending.drain(..self.consumed - keep);
266            self.consumed = keep;
267        }
268        Ok(())
269    }
270
271    fn encode_chunks(&mut self, chunks: usize, out: &mut Vec<Packet>) -> Result<()> {
272        use rayon::prelude::*;
273        let c = self.cfg.channels as usize;
274        let step = FRAME * c;
275        let cfg = self.cfg;
276        let base = self.consumed;
277        let first = self.frames;
278        let pending = &self.pending;
279        let done: Vec<Option<(Vec<Packet>, CeltEncoder, u32)>> = (0..chunks)
280            .into_par_iter()
281            .map(|k| {
282                catch_unwind(|| {
283                    let mut celt = CeltEncoder::new(c);
284                    let mut rc = RangeEncoder::new(MAX_PACKET);
285                    let start = base + k * CHUNK * step;
286                    let f0 = first + (k * CHUNK) as u64;
287                    if f0 > 0 {
288                        celt.prime(&pending[start.saturating_sub(HISTORY * c)..start]);
289                    }
290                    let mut packets = Vec::with_capacity(CHUNK);
291                    let mut rng = 0;
292                    for i in 0..CHUNK {
293                        let s = start + i * step;
294                        let (p, r) = encode_packet(
295                            &cfg,
296                            &mut celt,
297                            &mut rc,
298                            &pending[s..s + step],
299                            f0 + i as u64,
300                        );
301                        packets.push(p);
302                        rng = r;
303                    }
304                    (packets, celt, rng)
305                })
306                .ok()
307            })
308            .collect();
309        if done.iter().any(Option::is_none) {
310            return Err(crashed());
311        }
312        let mut last = None;
313        for (packets, celt, rng) in done.into_iter().flatten() {
314            out.extend(packets);
315            last = Some((celt, rng));
316        }
317        if let Some((celt, rng)) = last {
318            self.celt = celt;
319            self.final_range = rng;
320        }
321        self.frames += (chunks * CHUNK) as u64;
322        self.consumed += chunks * CHUNK * step;
323        Ok(())
324    }
325}
326
327fn crashed() -> Error {
328    Error::Backend("Opus encoder failed on this input".into())
329}
330
331fn encode_silk(
332    cfg: &EncoderConfig,
333    silk: &mut SilkEncoder,
334    rc: &mut RangeEncoder,
335    pcm: &[f32],
336    index: u64,
337    scratch: &mut Vec<f32>,
338) -> (Packet, u32) {
339    let c = cfg.channels as usize;
340    let bw = silk_bandwidth(cfg.bitrate);
341    scratch.clear();
342    scratch.extend(pcm.iter().map(|&v| {
343        if v.is_finite() {
344            v.clamp(-1.0, 1.0)
345        } else {
346            0.0
347        }
348    }));
349    let bytes = ((cfg.bitrate as usize * FRAME) / (48000 * 8)).clamp(3, MAX_PACKET);
350    rc.reset(MAX_PACKET - 1);
351    silk.encode(rc, scratch, bw, ((bytes - 1) * 8) as u32);
352    let used = (((rc.tell() + 7) >> 3).max(1) as usize).min(MAX_PACKET - 1);
353    let config: u8 = match bw {
354        Bandwidth::Narrowband => 1,
355        Bandwidth::Mediumband => 5,
356        _ => 9,
357    };
358    let toc = (config << 3) | (((c == 2) as u8) << 2);
359    let frame = &rc.finish()[..used];
360    let mut data = Vec::with_capacity(bytes.max(used + 1));
361    if !cfg.vbr && used + 3 <= bytes {
362        let rem = bytes - 2 - used;
363        data.push(toc | 3);
364        data.push(0x41);
365        data.extend(std::iter::repeat_n(255u8, (rem - 1) / 255));
366        data.push(((rem - 1) % 255) as u8);
367        data.extend_from_slice(frame);
368        data.resize(bytes, 0);
369    } else {
370        data.push(toc);
371        data.extend_from_slice(frame);
372    }
373    let pts = index as i64 * FRAME as i64;
374    (
375        Packet {
376            data,
377            pts,
378            dts: pts,
379            keyframe: true,
380        },
381        rc.rng(),
382    )
383}
384
385fn encode_packet(
386    cfg: &EncoderConfig,
387    celt: &mut CeltEncoder,
388    rc: &mut RangeEncoder,
389    pcm: &[f32],
390    index: u64,
391) -> (Packet, u32) {
392    let c = cfg.channels as usize;
393    let end = Bandwidth::Fullband.celt_end_band();
394    let allow_pf = index % CHUNK as u64 != CHUNK as u64 - 1;
395    let a = celt.analyse(pcm, FRAME, end, allow_pf);
396    let base = (cfg.bitrate as usize * FRAME) as f32 / (48000.0 * 8.0);
397    let total = if cfg.vbr {
398        celt.vbr.bytes(base, MAX_PACKET, &a)
399    } else {
400        (base as usize).clamp(3, MAX_PACKET)
401    };
402    let nbytes = total - 1;
403    let params = FrameParams {
404        start: 0,
405        end,
406        nbytes,
407        intensity: intensity_band(cfg.bitrate, c, end),
408    };
409    rc.reset(nbytes);
410    celt.encode_frame(FRAME, &params, &a, rc);
411    let toc = (31u8 << 3) | (((c == 2) as u8) << 2);
412    let mut data = Vec::with_capacity(total);
413    data.push(toc);
414    data.extend_from_slice(rc.finish());
415    let pts = index as i64 * FRAME as i64;
416    (
417        Packet {
418            data,
419            pts,
420            dts: pts,
421            keyframe: true,
422        },
423        rc.rng(),
424    )
425}
426
427pub struct OggWriter<W: Write> {
428    enc: Encoder,
429    out: ogg::Writer<W>,
430}
431
432impl<W: Write> OggWriter<W> {
433    pub fn new(cfg: EncoderConfig, w: W) -> Result<OggWriter<W>> {
434        let enc = Encoder::new(cfg)?;
435        let out = ogg::Writer::new(w, 0x5049_544f, &enc.head(), "PITO MD codecs")?;
436        Ok(OggWriter { enc, out })
437    }
438
439    pub fn write(&mut self, audio: &AudioBuffer) -> Result<()> {
440        for p in self.enc.encode(audio)? {
441            self.out.push(&p.data, p.pts + FRAME as i64)?;
442        }
443        Ok(())
444    }
445
446    pub fn finish(mut self) -> Result<W> {
447        let end = (self.enc.input + self.enc.pre_skip() as u64) as i64;
448        for p in self.enc.flush()? {
449            self.out.push(&p.data, (p.pts + FRAME as i64).min(end))?;
450        }
451        self.out.finish()
452    }
453}
454
455pub fn encode_ogg(cfg: EncoderConfig, audio: &AudioBuffer) -> Result<Vec<u8>> {
456    let mut w = OggWriter::new(cfg, Vec::new())?;
457    w.write(audio)?;
458    w.finish()
459}
460
461fn intensity_band(bitrate: u32, channels: usize, end: usize) -> usize {
462    if channels == 1 {
463        return end;
464    }
465    let kbps = bitrate / 1000;
466    let band = match kbps {
467        0..=31 => 4,
468        32..=47 => 8,
469        48..=63 => 11,
470        64..=79 => 13,
471        80..=95 => 15,
472        96..=111 => 17,
473        112..=127 => 19,
474        128..=143 => 20,
475        _ => 21,
476    };
477    band.min(end)
478}
479
480#[cfg(test)]
481mod tests;
pito_codecs/opus/enc/mod.rs

pito_codecs/opus/enc/
mod.rs
