xenia-rs/crates/xenia-apu/src/xma2_codec.rs

//! Thin unsafe wrapper around the mainline FFmpeg `AV_CODEC_ID_XMA2` decoder.
//!
//! Unlike canary's vendored `XMAFRAMES` (one frame per packet, custom padding
//! header), the distro xma2 decoder consumes whole 2 KB XMA2 packets
//! (`block_align == 2048`), needs `extradata` declaring the channel/stream
//! layout, and buffers samples internally across packets. We drive it with the
//! guest's raw 2 KB packets and pull whatever 512-sample float-planar frames it
//! emits, returning them as interleaved S16 big-endian PCM (canary `ConvertFrame`).

use std::os::raw::c_int;
use std::ptr;

use ffmpeg_sys_next as ff;

/// One xma2 decoder instance, configured for a fixed (sample_rate, channels).
pub struct Xma2Codec {
    codec: *const ff::AVCodec,
    ctx: *mut ff::AVCodecContext,
    frame: *mut ff::AVFrame,
    packet: *mut ff::AVPacket,
    extradata: Vec<u8>,
    channels: u32,
}

// FFmpeg objects are not Send/Sync by default; the decoder is only ever touched
// on the CPU scheduler thread (decode_pending), so this is sound for our use.
unsafe impl Send for Xma2Codec {}

impl Xma2Codec {
    /// Build XMA2WAVEFORMATEX extradata (34 bytes) for a single XMA2 stream.
    /// Layout (little-endian, per FFmpeg `xma_decode_init` / xma2defs.h):
    ///   [0..2]  NumStreams (u16)         = 1
    ///   [2..6]  ChannelMask (u32)        = mono/stereo mask
    ///   [6..34] remaining XMA2WAVEFORMATEX fields (unused by the decoder)
    fn build_extradata(channels: u32) -> Vec<u8> {
        let mut e = vec![0u8; 34];
        // NumStreams = 1
        e[0..2].copy_from_slice(&1u16.to_le_bytes());
        // ChannelMask: 0x3 (FL|FR) for stereo, 0x4 (FC) for mono.
        let mask: u32 = if channels >= 2 { 0x3 } else { 0x4 };
        e[2..6].copy_from_slice(&mask.to_le_bytes());
        e
    }

    pub fn new(sample_rate: u32, channels: u32) -> Result<Self, String> {
        unsafe {
            let codec = ff::avcodec_find_decoder(ff::AVCodecID::AV_CODEC_ID_XMA2);
            if codec.is_null() {
                return Err("xma2 decoder not found in libavcodec".into());
            }
            let ctx = ff::avcodec_alloc_context3(codec);
            if ctx.is_null() {
                return Err("avcodec_alloc_context3 failed".into());
            }

            let mut extradata = Self::build_extradata(channels);
            // FFmpeg requires extradata to be allocated with av_malloc and
            // padded; copy our bytes into an av_malloc'd buffer.
            let pad = ff::AV_INPUT_BUFFER_PADDING_SIZE as usize;
            let raw = ff::av_mallocz(extradata.len() + pad) as *mut u8;
            if raw.is_null() {
                ff::avcodec_free_context(&mut (ctx as *mut _));
                return Err("av_mallocz extradata failed".into());
            }
            ptr::copy_nonoverlapping(extradata.as_ptr(), raw, extradata.len());
            (*ctx).extradata = raw;
            (*ctx).extradata_size = extradata.len() as c_int;

            (*ctx).sample_rate = sample_rate as c_int;
            (*ctx).block_align = 2048;
            ff::av_channel_layout_default(&mut (*ctx).ch_layout, channels as c_int);

            let ret = ff::avcodec_open2(ctx, codec, ptr::null_mut());
            if ret < 0 {
                let mut ctxm = ctx;
                ff::avcodec_free_context(&mut ctxm);
                return Err(format!("avcodec_open2 failed: {}", av_err(ret)));
            }

            let frame = ff::av_frame_alloc();
            let packet = ff::av_packet_alloc();
            if frame.is_null() || packet.is_null() {
                let mut ctxm = ctx;
                ff::avcodec_free_context(&mut ctxm);
                return Err("av_frame_alloc/av_packet_alloc failed".into());
            }

            // keep our Vec alive as the source of truth for length
            extradata.shrink_to_fit();

            Ok(Self {
                codec,
                ctx,
                frame,
                packet,
                extradata,
                channels,
            })
        }
    }

    pub fn channels(&self) -> u32 {
        self.channels
    }

    /// Feed one raw 2 KB XMA2 packet (header + data) to the decoder. Returns the
    /// number of bytes the decoder accepted (0 = buffered, needs no new packet
    /// yet / EAGAIN). Decoded frames are pulled via [`receive_frame`].
    pub fn send_packet(&mut self, packet: &[u8]) -> Result<(), String> {
        unsafe {
            // av_packet_from_data takes ownership of an av_malloc buffer; simpler
            // to point at our own bytes via a stack packet with a padded copy.
            let pad = ff::AV_INPUT_BUFFER_PADDING_SIZE as usize;
            let buf = ff::av_malloc(packet.len() + pad) as *mut u8;
            if buf.is_null() {
                return Err("av_malloc packet failed".into());
            }
            ptr::copy_nonoverlapping(packet.as_ptr(), buf, packet.len());
            ptr::write_bytes(buf.add(packet.len()), 0, pad);
            ff::av_packet_unref(self.packet);
            // Wrap buf so FFmpeg frees it.
            let ret = ff::av_packet_from_data(self.packet, buf, packet.len() as c_int);
            if ret < 0 {
                ff::av_free(buf as *mut _);
                return Err(format!("av_packet_from_data failed: {}", av_err(ret)));
            }
            let ret = ff::avcodec_send_packet(self.ctx, self.packet);
            if ret == ff::AVERROR(ff::EAGAIN) {
                // Decoder full — caller should drain frames first then retry.
                return Err("EAGAIN".into());
            }
            if ret < 0 {
                return Err(format!("avcodec_send_packet failed: {}", av_err(ret)));
            }
            Ok(())
        }
    }

    /// Signal end-of-stream so the decoder flushes its internal FIFO.
    pub fn send_eof(&mut self) {
        unsafe {
            let _ = ff::avcodec_send_packet(self.ctx, ptr::null());
        }
    }

    /// Pull one decoded frame as interleaved S16 big-endian PCM, or None if the
    /// decoder needs more input (EAGAIN) or is drained (EOF). Returns
    /// (samples_per_channel, interleaved_s16be_bytes).
    pub fn receive_frame(&mut self) -> Option<(u32, Vec<u8>)> {
        unsafe {
            let ret = ff::avcodec_receive_frame(self.ctx, self.frame);
            if ret < 0 {
                return None;
            }
            let nb = (*self.frame).nb_samples as u32;
            if nb == 0 {
                return None;
            }
            let ch = (*self.frame).ch_layout.nb_channels.max(1) as u32;
            let out = convert_frame_planar_to_s16be(self.frame, ch, nb);
            Some((nb, out))
        }
    }
}

impl Drop for Xma2Codec {
    fn drop(&mut self) {
        unsafe {
            if !self.frame.is_null() {
                ff::av_frame_free(&mut self.frame);
            }
            if !self.packet.is_null() {
                ff::av_packet_free(&mut self.packet);
            }
            if !self.ctx.is_null() {
                ff::avcodec_free_context(&mut self.ctx);
            }
            let _ = &self.codec;
            let _ = &self.extradata;
        }
    }
}

/// Convert FFmpeg planar-float output to interleaved S16 big-endian PCM
/// (faithful to canary `XmaContext::ConvertFrame`: saturate to [-1,1], scale by
/// 2^15-1, byte-swap each sample). `channels` planes of `nb_samples` floats.
unsafe fn convert_frame_planar_to_s16be(
    frame: *mut ff::AVFrame,
    channels: u32,
    nb_samples: u32,
) -> Vec<u8> {
    const SCALE: f32 = ((1i32 << 15) - 1) as f32;
    let mut out = Vec::with_capacity((nb_samples * channels * 2) as usize);
    unsafe {
        // extended_data[ch] points to a plane of f32 (AV_SAMPLE_FMT_FLTP).
        let ext = (*frame).extended_data;
        for i in 0..nb_samples as isize {
            for ch in 0..channels as isize {
                let plane = *ext.offset(ch) as *const f32;
                let s = if plane.is_null() { 0.0 } else { *plane.offset(i) };
                let clamped = s.clamp(-1.0, 1.0) * SCALE;
                let v = clamped as i16;
                out.extend_from_slice(&v.to_be_bytes());
            }
        }
    }
    out
}

fn av_err(code: c_int) -> String {
    unsafe {
        let mut buf = [0i8; ff::AV_ERROR_MAX_STRING_SIZE as usize];
        ff::av_strerror(code, buf.as_mut_ptr(), buf.len());
        let cstr = std::ffi::CStr::from_ptr(buf.as_ptr());
        cstr.to_string_lossy().into_owned()
    }
}