peerspeak/src/codec/opus_impl.rs

use crate::codec::{AudioEncoder, AudioDecoder, CodecError};
use opus::{Encoder, Decoder, Application, Channels};

pub struct OpusEncoder {
    encoder: Encoder,
}

impl OpusEncoder {
    /// Creates a new Opus encoder.
    /// Standard voice parameters: sample_rate = 48000, channels = Channels::Mono, application = Application::Voip
    pub fn new(sample_rate: u32, channels: Channels, application: Application) -> Result<Self, CodecError> {
        let encoder = Encoder::new(sample_rate, channels, application)
            .map_err(|e| CodecError::Init(format!("Failed to create Opus encoder: {}", e)))?;
        Ok(Self { encoder })
    }
}

impl AudioEncoder for OpusEncoder {
    fn encode(&mut self, pcm: &[i16]) -> Result<Vec<u8>, CodecError> {
        // We allocate a buffer for the compressed output.
        // A maximum packet size of 4000 bytes is more than enough for a single voice frame.
        let mut compressed = vec![0u8; 4000];
        let len = self.encoder.encode(pcm, &mut compressed)
            .map_err(|e| CodecError::Encode(format!("Opus encoding failed: {}", e)))?;
        
        compressed.truncate(len);
        Ok(compressed)
    }
}

pub struct OpusDecoder {
    decoder: Decoder,
    channels: Channels,
    /// Samples-per-channel of the frames we transmit (20ms @ 48kHz = 960).
    /// Used to size the Packet Loss Concealment output, since libopus conceals
    /// `frame_size` samples when given no input — passing the full max buffer
    /// would synthesize a 120ms burst instead of a single 20ms frame.
    frame_samples: usize,
}

impl OpusDecoder {
    /// Creates a new Opus decoder.
    /// Standard voice parameters: sample_rate = 48000, channels = Channels::Mono.
    /// `frame_samples` is the per-channel length of one transmitted frame (e.g. 960).
    pub fn new(sample_rate: u32, channels: Channels, frame_samples: usize) -> Result<Self, CodecError> {
        let decoder = Decoder::new(sample_rate, channels)
            .map_err(|e| CodecError::Init(format!("Failed to create Opus decoder: {}", e)))?;
        Ok(Self { decoder, channels, frame_samples })
    }

    fn channels_count(&self) -> usize {
        match self.channels {
            Channels::Mono => 1,
            Channels::Stereo => 2,
        }
    }
}

impl AudioDecoder for OpusDecoder {
    fn decode(&mut self, compressed: Option<&[u8]>) -> Result<Vec<i16>, CodecError> {
        let channels_count = self.channels_count();

        let (mut pcm, input): (Vec<i16>, &[u8]) = match compressed {
            Some(data) if !data.is_empty() => {
                // Normal decode. Size the buffer to the maximum Opus frame (120ms =
                // 5760 samples/channel); libopus decodes the packet's true duration.
                (vec![0i16; 5760 * channels_count], data)
            }
            _ => {
                // Packet Loss Concealment: an empty input makes libopus synthesize
                // exactly `frame_samples` of concealment, so size the buffer to match.
                (vec![0i16; self.frame_samples * channels_count], &[])
            }
        };

        let decoded_per_channel = self.decoder.decode(input, &mut pcm, false)
            .map_err(|e| CodecError::Decode(format!("Opus decoding failed: {}", e)))?;

        pcm.truncate(decoded_per_channel * channels_count);
        Ok(pcm)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_round_trip() {
        let mut encoder = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap();
        let mut decoder = OpusDecoder::new(48000, Channels::Mono, 960).unwrap();

        // 1. Round-trip shape: generate a 960-sample mono PCM frame (440Hz sine wave)
        let mut pcm = vec![0i16; 960];
        for (i, sample) in pcm.iter_mut().enumerate() {
            let t = (i as f32) / 48000.0;
            let val = (t * 440.0 * 2.0 * std::f32::consts::PI).sin();
            *sample = (val * 10000.0) as i16;
        }

        // encode it
        let compressed = encoder.encode(&pcm).unwrap();
        assert!(!compressed.is_empty(), "Compressed buffer should not be empty");
        assert!(
            compressed.len() < pcm.len() * std::mem::size_of::<i16>(),
            "Compressed size ({}) should be smaller than raw PCM size ({})",
            compressed.len(),
            pcm.len() * std::mem::size_of::<i16>()
        );

        // decode it
        let decoded = decoder.decode(Some(&compressed)).unwrap();
        assert_eq!(decoded.len(), 960, "Decoded sample count should be exactly 960");

        // 2. Round-trip carries signal energy (not silence)
        let sum_sq: f64 = decoded.iter().map(|&x| (x as f64).powi(2)).sum();
        let rms = (sum_sq / decoded.len() as f64).sqrt();
        // Since input had amplitude ~10000, let's verify RMS is significantly above 0 (e.g. > 100.0)
        assert!(rms > 100.0, "Decoded signal should carry energy (RMS was {})", rms);
    }

    #[test]
    fn test_plc_sizing() {
        let mut decoder = OpusDecoder::new(48000, Channels::Mono, 960).unwrap();

        // decode(None) returns exactly frame_samples (960) samples
        let plc_none = decoder.decode(None).unwrap();
        assert_eq!(plc_none.len(), 960, "decode(None) should yield exactly 960 samples");

        // decode(Some(&[])) (empty slice) does the same
        let plc_empty = decoder.decode(Some(&[])).unwrap();
        assert_eq!(plc_empty.len(), 960, "decode(Some(&[])) should yield exactly 960 samples");
    }

    #[test]
    fn test_decoded_duration_follows_packet() {
        let mut encoder = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap();
        let mut decoder = OpusDecoder::new(48000, Channels::Mono, 960).unwrap();

        let pcm = vec![0i16; 960];
        let compressed = encoder.encode(&pcm).unwrap();
        let decoded = decoder.decode(Some(&compressed)).unwrap();
        assert_eq!(decoded.len(), 960, "Decoded sample count should match packet duration");
    }

    #[test]
    fn test_stereo_plc_sizing() {
        let mut decoder = OpusDecoder::new(48000, Channels::Stereo, 960).unwrap();

        // decode(None) returns exactly frame_samples * 2 (1920) samples
        let plc_none = decoder.decode(None).unwrap();
        assert_eq!(plc_none.len(), 960 * 2, "Stereo decode(None) should yield exactly 1920 samples");

        // decode(Some(&[])) (empty slice) does the same
        let plc_empty = decoder.decode(Some(&[])).unwrap();
        assert_eq!(plc_empty.len(), 960 * 2, "Stereo decode(Some(&[])) should yield exactly 1920 samples");
    }
}