feat(core): 添加 API 用于跳过重采样器直接获取原始 PCM 数据

Author: apoint123

crates/ffmpeg_audio/src/core/frame.rs

@@ -5,10 +5,27 @@ use std::{
 };
 
 use crate::{
+    AudioError,
+    AudioSample,
+    Result,
     TimeBase,
     sys,
 };
 
+/// A safe enum representing the memory layout of underlying FFmpeg PCM data.
+#[derive(Debug, Clone)]
+pub enum RawAudioData<'a, T> {
+    /// Packed (Interleaved) layout.
+    /// All channel data is interleaved in a single contiguous memory block, e.g., `[L, R, L, R, L,
+    /// R]`.
+    Packed(&'a [T]),
+
+    /// Planar layout.
+    /// Data for each channel is stored in independent, contiguous memory blocks, e.g., `[L, L, L]`
+    /// and `[R, R, R]`. The length of the outer Vec represents the number of channels.
+    Planar(Vec<&'a [T]>),
+}
+
 /// A safe, zero-copy wrapper around FFmpeg's raw `AVFrame`.
 ///
 /// This wrapper is useful for 1-to-N zero-copy dispatching to multiple downstream
@@ -21,7 +38,7 @@ pub struct AudioFrame<'a> {
     _marker: PhantomData<&'a mut ()>,
 }
 
-impl AudioFrame<'_> {
+impl<'a> AudioFrame<'a> {
     /// Creates a new `AudioFrame` wrapper.
     ///
     /// # Safety
@@ -136,6 +153,57 @@ impl AudioFrame<'_> {
                 duration
             })
     }
+
+    /// Zero-copy extraction of raw PCM audio data directly from the underlying FFmpeg AVFrame.
+    ///
+    /// # Returns
+    /// * `Ok(RawAudioData)` - The corresponding memory slice enum.
+    /// * `Err(AudioError::FormatMismatch)` - The requested `T` does not match the type actually
+    ///   output by the underlying decoder.
+    pub fn raw_data<T: AudioSample>(&self) -> Result<RawAudioData<'a, T>> {
+        let fmt = self.sample_fmt();
+        let is_packed = fmt == T::PACKED_FORMAT;
+        let is_planar = fmt == T::PLANAR_FORMAT;
+
+        if !is_packed && !is_planar {
+            return Err(AudioError::FormatMismatch);
+        }
+
+        let logical_samples = self.samples();
+        let channels = unsafe { (*self.ptr.as_ptr()).ch_layout.nb_channels } as usize;
+        let offset = self.offset();
+        let extended_data = unsafe { (*self.ptr.as_ptr()).extended_data };
+
+        if logical_samples == 0 || extended_data.is_null() {
+            return if is_packed {
+                Ok(RawAudioData::Packed(&[]))
+            } else {
+                Ok(RawAudioData::Planar(vec![&[]; channels]))
+            };
+        }
+
+        if is_packed {
+            unsafe {
+                let base_ptr = (*extended_data).cast::<T>();
+                let adjusted_ptr = base_ptr.add(offset * channels);
+                let slice = std::slice::from_raw_parts(adjusted_ptr, logical_samples * channels);
+
+                Ok(RawAudioData::Packed(slice))
+            }
+        } else {
+            let mut planes = Vec::with_capacity(channels);
+            unsafe {
+                for ch in 0..channels {
+                    let base_ptr = (*extended_data.add(ch)).cast::<T>();
+                    let adjusted_ptr = base_ptr.add(offset);
+                    let slice = std::slice::from_raw_parts(adjusted_ptr, logical_samples);
+
+                    planes.push(slice);
+                }
+            }
+            Ok(RawAudioData::Planar(planes))
+        }
+    }
 }
 
 #[cfg(test)]

crates/ffmpeg_audio/src/lib.rs

@@ -8,7 +8,10 @@ pub mod resample;
 pub use core::http::HttpAudioSource;
 pub use core::{
     format::AudioSample,
-    frame::AudioFrame,
+    frame::{
+        AudioFrame,
+        RawAudioData,
+    },
     info::SourceAudioInfo,
     time::TimeBase,
 };

crates/ffmpeg_audio/tests/integration_test.rs

@@ -5,6 +5,7 @@ use std::{
 
 use ffmpeg_audio::{
     AudioReader,
+    RawAudioData,
     ResampleOptions,
     SeekMode,
 };
@@ -419,6 +420,79 @@ fn test_planar_format_mismatch_rejection_planar_to_packed() {
     );
 }
 
+#[test]
+fn test_raw_data_extraction_packed_and_type_safety() {
+    let wav_data = generate_sine_wav(0.1);
+    let mut reader = AudioReader::new(Cursor::new(wav_data)).unwrap();
+
+    let channels = reader.source_info().channels as usize;
+
+    let frame = reader.receive_frame().unwrap().expect("应能读取到 WAV 帧");
+
+    let err_result = frame.raw_data::<f32>();
+    assert!(
+        matches!(err_result, Err(ffmpeg_audio::AudioError::FormatMismatch)),
+        "期望类型不匹配错误，但得到了: {err_result:?}"
+    );
+
+    let raw = frame.raw_data::<i16>().expect("提取 i16 原始数据失败");
+
+    match raw {
+        RawAudioData::Packed(data) => {
+            let expected_len = frame.samples() * channels;
+
+            assert_eq!(
+                data.len(),
+                expected_len,
+                "Packed 切片的元素总数 ({}) 应等于逻辑样本数 ({}) 乘以声道数 ({})",
+                data.len(),
+                frame.samples(),
+                channels
+            );
+
+            let has_signal = data.iter().any(|&s| s != 0);
+            assert!(has_signal, "捞出的原始 PCM 不应全为静音");
+        }
+        RawAudioData::Planar(_) => panic!("WAV 应该被解码为 Packed 布局，却得到了 Planar"),
+    }
+}
+
+#[test]
+fn test_raw_data_signal_integrity_across_frames() {
+    let sample_rate = 44100;
+    let freq: f32 = 440.0;
+    let duration = 0.1;
+
+    let wav_data = generate_sine_wav(duration);
+    let mut reader = AudioReader::new(Cursor::new(wav_data)).unwrap();
+
+    let mut global_sample_index = 0;
+
+    while let Some(frame) = reader.receive_frame().unwrap() {
+        let raw = frame.raw_data::<i16>().expect("提取 i16 失败");
+
+        if let RawAudioData::Packed(data) = raw {
+            for &actual_sample in data {
+                let t = global_sample_index as f32 / sample_rate as f32;
+                let expected_sample =
+                    (f32::sin(2.0 * std::f32::consts::PI * freq * t) * 16000.0) as i16;
+
+                let diff = (i32::from(actual_sample) - i32::from(expected_sample)).abs();
+                assert!(
+                    diff <= 1,
+                    "信号失真！在全局样本索引 {global_sample_index} 期望 {expected_sample}, 实际得到 {actual_sample}"
+                );
+
+                global_sample_index += 1;
+            }
+        } else {
+            panic!("WAV 应为 Packed 格式");
+        }
+    }
+
+    assert_eq!(global_sample_index, 4410, "解码出的样本总数不对");
+}
+
 #[cfg(not(target_arch = "wasm32"))]
 mod file_tests {
     use std::fs::File;
@@ -495,4 +569,45 @@ mod file_tests {
             "Should have successfully decoded mutated stream"
         );
     }
+
+    #[test]
+    fn test_raw_data_extraction_planar_with_seek_offset() {
+        let file = File::open(AAC_SEEK_PATH).expect("Failed to open AAC test asset");
+        let mut reader = AudioReader::new(file).unwrap();
+
+        let expected_channels = reader.source_info().channels as usize;
+
+        let target = Duration::from_millis(501);
+        reader.seek(target, SeekMode::Accurate).unwrap();
+
+        let frame = reader.receive_frame().unwrap().expect("Seek 后读取帧失败");
+
+        assert!(frame.samples() > 0, "修剪后应仍有剩余数据");
+
+        let raw = frame.raw_data::<f32>().expect("提取 f32 原始数据失败");
+
+        match raw {
+            RawAudioData::Planar(planes) => {
+                assert_eq!(
+                    planes.len(),
+                    expected_channels,
+                    "返回的平面切片数量 ({}) 应等于文件的物理声道数 ({})",
+                    planes.len(),
+                    expected_channels
+                );
+
+                for (ch_idx, plane) in planes.iter().enumerate() {
+                    assert_eq!(
+                        plane.len(),
+                        frame.samples(),
+                        "声道 {ch_idx} 的切片长度未与修剪后的样本数对齐"
+                    );
+
+                    let has_signal = plane.iter().any(|&s| s.abs() > 0.0);
+                    assert!(has_signal, "提取出来的声道 {ch_idx} 数据异常");
+                }
+            }
+            RawAudioData::Packed(_) => panic!("AAC 应该被解码为 Planar 布局，却得到了 Packed"),
+        }
+    }
 }