Skip to main content

rave_core/
codec_traits.rs

1//! Shared codec traits used across crate boundaries.
2//!
3//! These traits break the circular dependency between `rave-nvcodec`,
4//! `rave-ffmpeg`, and `rave-pipeline` by providing a neutral home.
5
6use crate::error::Result;
7use crate::ffi_types::CUevent;
8use crate::types::FrameEnvelope;
9
10// ─── Bitstream source (demuxer → decoder) ────────────────────────────────
11
12/// Demuxed compressed bitstream packets (host-side, NOT raw pixels).
13///
14/// Implementations: file reader, network receiver, FFmpeg demuxer, etc.
15pub trait BitstreamSource: Send + 'static {
16    /// Read the next compressed bitstream packet, or `None` at end-of-stream.
17    fn read_packet(&mut self) -> Result<Option<BitstreamPacket>>;
18}
19
20/// A single demuxed NAL unit or access unit.
21pub struct BitstreamPacket {
22    /// Compressed bitstream data (Annex B or length-prefixed).
23    /// Host memory is acceptable — this is codec-compressed (~10 KB/frame).
24    pub data: Vec<u8>,
25    /// Presentation timestamp in microseconds.
26    ///
27    /// Container-specific time bases are converted at the demux boundary
28    /// (for example in `rave-ffmpeg`) so downstream decode/encode stages
29    /// operate on one stable unit.
30    pub pts: i64,
31    /// Whether this packet encodes an IDR/keyframe.
32    pub is_keyframe: bool,
33}
34
35// ─── Bitstream sink (encoder → muxer) ────────────────────────────────────
36
37/// Receives encoded bitstream output.
38///
39/// Implementations: file writer, muxer, network sender, etc.
40pub trait BitstreamSink: Send + 'static {
41    /// Write one encoded packet.
42    ///
43    /// `pts` and `dts` are in microseconds. Container muxers are responsible
44    /// for rescaling to stream time_base at the output boundary.
45    fn write_packet(&mut self, data: &[u8], pts: i64, dts: i64, is_keyframe: bool) -> Result<()>;
46    /// Flush any internal buffers and finalise the output stream.
47    fn flush(&mut self) -> Result<()>;
48}
49
50// ─── Pipeline stage traits ───────────────────────────────────────────────
51
52/// A decoded NV12 frame with an optional CUDA event for cross-stream sync.
53///
54/// When produced by a real hardware decoder (NVDEC), `decode_event` carries
55/// the event recorded on `decode_stream` after the D2D copy.  The preprocess
56/// stage calls `cuStreamWaitEvent(preprocess_stream, event)` before reading
57/// the texture — this ensures the decode data is ready without CPU-blocking.
58///
59/// Mock decoders set both fields to `None`.
60pub struct DecodedFrame {
61    pub envelope: FrameEnvelope,
62    /// CUDA event recorded after decode completes.  `None` for mock decoders.
63    pub decode_event: Option<CUevent>,
64    /// Channel to return used events to the decoder's `EventPool` for reuse.
65    pub event_return: Option<std::sync::mpsc::Sender<CUevent>>,
66}
67
68// SAFETY: CUevent is a CUDA handle (*mut c_void) that is safe to send across
69// threads — CUDA events have no thread affinity.
70unsafe impl Send for DecodedFrame {}
71
72/// Video frame decoder producing GPU-resident NV12 frames.
73pub trait FrameDecoder: Send + 'static {
74    /// Decode the next frame from the bitstream source.
75    ///
76    /// Returns `None` at end-of-stream.
77    fn decode_next(&mut self) -> Result<Option<DecodedFrame>>;
78}
79
80/// Video frame encoder consuming GPU-resident NV12 frames.
81pub trait FrameEncoder: Send + 'static {
82    /// Submit one GPU-resident NV12 frame for encoding.
83    fn encode(&mut self, frame: FrameEnvelope) -> Result<()>;
84    /// Flush any pending frames and finalise the bitstream.
85    fn flush(&mut self) -> Result<()>;
86}
87
88// ─── Model precision ─────────────────────────────────────────────────────
89
90/// Which floating-point precision the inference model expects.
91///
92/// Mirrors [`rave_cuda::kernels::ModelPrecision`] — the two types are
93/// identical and convert via pattern matching.
94#[derive(Clone, Copy, Debug, PartialEq, Eq)]
95pub enum ModelPrecision {
96    /// 32-bit single-precision float.
97    F32,
98    /// 16-bit half-precision float.
99    F16,
100}