rusty_h264_encoder/
lib.rs

1//! Pure-Rust H.264 (Constrained Baseline) encoder.
2//!
3//! Status: all-intra, `I_16x16` DC-predicted macroblocks with the full
4//! transform → quantization → CAVLC pipeline. The Annex-B output is bit-exactly
5//! decodable by reference decoders (verified against ffmpeg). Richer intra modes
6//! (I_4x4), inter prediction, and the in-loop deblocking filter (currently
7//! signalled disabled) are layered in by later generations behind this API.
8//!
9//! ```
10//! use rusty_h264_encoder::{Encoder, EncoderConfig};
11//! use rusty_h264_common::YuvFrame;
12//!
13//! let cfg = EncoderConfig::new(16, 16);
14//! let mut enc = Encoder::new(cfg).unwrap();
15//! let frame = YuvFrame::black(16, 16);
16//! let bitstream = enc.encode(&frame); // Annex-B bytes for one access unit
17//! assert!(!bitstream.is_empty());
18//! ```
19
20mod config;
21mod lookahead;
22mod mb16;
23mod params;
24mod rc;
25mod slice;
26
27pub use config::{EncoderConfig, Preset};
28pub use params::{Pps, Sps};
29pub use rc::RateControl;
30
31use rusty_h264_common::{BitWriter, ChromaFormat, NalUnit, NalUnitType, Profile, YuvFrame};
32
33/// Errors that can arise constructing or driving the encoder.
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub enum EncodeError {
36    /// A feature outside the implemented Constrained Baseline subset was asked for.
37    Unsupported(&'static str),
38    /// The supplied frame's dimensions or plane sizes don't match the config.
39    FrameMismatch,
40}
41
42impl core::fmt::Display for EncodeError {
43    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
44        match self {
45            EncodeError::Unsupported(s) => write!(f, "unsupported: {s}"),
46            EncodeError::FrameMismatch => write!(f, "frame dimensions do not match encoder config"),
47        }
48    }
49}
50
51impl std::error::Error for EncodeError {}
52
53/// A Constrained Baseline H.264 encoder.
54#[derive(Debug)]
55pub struct Encoder {
56    cfg: EncoderConfig,
57    sps: Sps,
58    pps: Pps,
59    /// Count of frames fed so far; drives IDR placement via `gop_size`.
60    frame_index: u32,
61    /// `frame_num` of the next picture (resets to 0 at each IDR).
62    next_frame_num: u32,
63    /// Index of the current picture within its GOP (0 at IDR), for POC.
64    gop_index: u32,
65    /// Decoded-picture buffer: recent **deblocked** reconstructions (coded size),
66    /// most-recent first, used as inter references (`ref_idx` 0 = front).
67    refs: Vec<RefFrame>,
68    /// Average-bitrate controller; `None` for constant-QP encoding.
69    rc: Option<RateControl>,
70}
71
72/// A reference picture: deblocked reconstruction at coded (MB-grid) resolution.
73/// Stored now (4a); read by motion compensation in 4b.
74#[derive(Clone, Debug)]
75#[allow(dead_code)]
76pub(crate) struct RefFrame {
77    // 16-byte aligned (moved from the encoder's aligned rec planes) so the openh264
78    // MC asm can load aligned reference row chunks.
79    pub y: rusty_h264_common::aligned::AlignedBytes,
80    pub u: rusty_h264_common::aligned::AlignedBytes,
81    pub v: rusty_h264_common::aligned::AlignedBytes,
82}
83
84impl Encoder {
85    /// Creates an encoder, validating that the configuration is within the
86    /// implemented subset.
87    pub fn new(cfg: EncoderConfig) -> Result<Self, EncodeError> {
88        if !matches!(cfg.profile, Profile::ConstrainedBaseline | Profile::Baseline) {
89            return Err(EncodeError::Unsupported("only Constrained Baseline profile"));
90        }
91        if cfg.chroma != ChromaFormat::Yuv420 {
92            return Err(EncodeError::Unsupported("only 4:2:0 chroma"));
93        }
94        if cfg.width == 0 || cfg.height == 0 || cfg.width % 2 != 0 || cfg.height % 2 != 0 {
95            return Err(EncodeError::Unsupported("dimensions must be positive and even"));
96        }
97        let sps = Sps::from_config(&cfg);
98        let pps = Pps::from_config(&cfg);
99        let rc = (cfg.bitrate > 0).then(|| RateControl::new(cfg.bitrate, cfg.framerate, cfg.qp));
100        Ok(Self {
101            cfg,
102            sps,
103            pps,
104            frame_index: 0,
105            next_frame_num: 0,
106            gop_index: 0,
107            refs: Vec::new(),
108            rc,
109        })
110    }
111
112    /// The active configuration.
113    pub fn config(&self) -> &EncoderConfig {
114        &self.cfg
115    }
116
117    /// Encodes one frame, returning the Annex-B access unit. Every `gop_size`
118    /// frames (and always the first) is coded as an IDR, prefixed with SPS/PPS.
119    ///
120    /// Generation 1 codes *every* picture as an IDR (all-intra); inter frames
121    /// arrive with motion compensation later.
122    pub fn encode(&mut self, frame: &YuvFrame) -> Vec<u8> {
123        self.try_encode(frame).expect("frame matched config")
124    }
125
126    /// Fallible [`encode`](Self::encode): validates the frame against the config.
127    pub fn try_encode(&mut self, frame: &YuvFrame) -> Result<Vec<u8>, EncodeError> {
128        if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
129            return Err(EncodeError::FrameMismatch);
130        }
131
132        // GOP placement: an IDR at each `gop_size` boundary, P-frames between.
133        let is_idr = self.cfg.gop_size <= 1 || self.frame_index % self.cfg.gop_size == 0;
134        if is_idr {
135            self.gop_index = 0;
136            self.next_frame_num = 0;
137            self.refs.clear();
138        }
139        let frame_num = self.next_frame_num;
140        let poc_lsb = (2 * self.gop_index) % 16;
141
142        // Rate control (if enabled) chooses this frame's QP from a cheap
143        // look-ahead complexity estimate; otherwise the QP is fixed.
144        let complexity = if self.rc.is_some() {
145            lookahead::complexity(&self.cfg, frame, if is_idr { None } else { self.refs.first() })
146        } else {
147            0.0
148        };
149        let qp = match &self.rc {
150            Some(rc) => rc.pick_qp(is_idr, complexity),
151            None => self.cfg.qp,
152        };
153
154        let mut out = Vec::new();
155        let mut w = BitWriter::new();
156        let (nal_type, reference) = if is_idr {
157            // SPS/PPS precede every IDR so the stream is independently decodable.
158            self.sps.to_nal().write_annex_b(&mut out);
159            self.pps.to_nal().write_annex_b(&mut out);
160            slice::write_idr_slice_header(&mut w, &self.cfg, qp);
161            let r = mb16::encode_slice_data(&mut w, &self.cfg, frame, qp, false, &[]);
162            (NalUnitType::IdrSlice, r)
163        } else {
164            slice::write_p_slice_header(&mut w, &self.cfg, qp, frame_num, poc_lsb, self.refs.len());
165            let r = mb16::encode_slice_data(&mut w, &self.cfg, frame, qp, true, &self.refs);
166            (NalUnitType::NonIdrSlice, r)
167        };
168        let slice_bytes = w.into_bytes();
169        // Feed the coded slice size (the picture's own bits) back to the controller.
170        if let Some(rc) = &mut self.rc {
171            rc.update(is_idr, slice_bytes.len() * 8, qp, complexity);
172        }
173        NalUnit::new(3, nal_type, slice_bytes).write_annex_b(&mut out);
174
175        // The deblocked reconstruction enters the DPB (most-recent first), which
176        // is kept to `max_num_ref_frames` by a sliding window.
177        self.refs.insert(0, reference);
178        self.refs.truncate(self.cfg.num_ref_frames.max(1) as usize);
179        self.frame_index += 1;
180        self.gop_index += 1;
181        self.next_frame_num = (self.next_frame_num + 1) % 16;
182        Ok(out)
183    }
184
185    /// Batch-encodes every frame, returning one Annex-B access unit per frame.
186    ///
187    /// At constant QP the GOPs are independent — each begins with an IDR that
188    /// resets the DPB, `frame_num` and POC, and SPS/PPS precede every IDR — so they
189    /// are encoded **in parallel across CPU cores** and the result is
190    /// **byte-identical** to calling [`encode`](Self::encode) frame-by-frame. With
191    /// rate control enabled the per-frame QP depends on history, so this falls back
192    /// to sequential encoding. Within a GOP, P-frames are inherently sequential
193    /// (each predicts from the previous reconstruction); the parallelism is across
194    /// GOPs, so it scales with the number of GOPs in the clip.
195    pub fn encode_all(&self, frames: &[YuvFrame]) -> Result<Vec<Vec<u8>>, EncodeError> {
196        for f in frames {
197            if f.width != self.cfg.width || f.height != self.cfg.height || !f.is_valid() {
198                return Err(EncodeError::FrameMismatch);
199            }
200        }
201        // Rate control threads state across frames → it must stay sequential.
202        if self.cfg.bitrate > 0 {
203            let mut enc = Encoder::new(self.cfg.clone())?;
204            return frames.iter().map(|f| enc.try_encode(f)).collect();
205        }
206        let gop = self.cfg.gop_size.max(1) as usize;
207        let gops: Vec<&[YuvFrame]> = frames.chunks(gop).collect();
208        if gops.is_empty() {
209            return Ok(Vec::new());
210        }
211        let n = std::env::var("RUSTY_THREADS")
212            .ok()
213            .and_then(|v| v.parse().ok())
214            .or_else(|| std::thread::available_parallelism().map(|n| n.get()).ok())
215            .unwrap_or(1)
216            .min(gops.len());
217        // Each GOP is encoded with a fresh encoder (an IDR resets all state), so
218        // GOPs distribute across `n` worker threads with no shared mutable state.
219        let mut out: Vec<Option<Vec<Vec<u8>>>> = (0..gops.len()).map(|_| None).collect();
220        let cfg = &self.cfg;
221        let gops_ref = &gops;
222        std::thread::scope(|s| {
223            let handles: Vec<_> = (0..n)
224                .map(|t| {
225                    s.spawn(move || {
226                        let mut local = Vec::new();
227                        let mut i = t;
228                        while i < gops_ref.len() {
229                            let mut enc = Encoder::new(cfg.clone()).expect("config");
230                            let aus: Vec<Vec<u8>> = gops_ref[i].iter().map(|f| enc.encode(f)).collect();
231                            local.push((i, aus));
232                            i += n;
233                        }
234                        local
235                    })
236                })
237                .collect();
238            for h in handles {
239                for (i, aus) in h.join().expect("encode worker panicked") {
240                    out[i] = Some(aus);
241                }
242            }
243        });
244        Ok(out.into_iter().flatten().flatten().collect())
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn rejects_unsupported_profile() {
254        let mut cfg = EncoderConfig::new(16, 16);
255        cfg.profile = Profile::High;
256        assert!(matches!(Encoder::new(cfg), Err(EncodeError::Unsupported(_))));
257    }
258
259    #[test]
260    fn encodes_access_unit_with_sps_pps_idr() {
261        use rusty_h264_common::nal::split_annex_b;
262        let cfg = EncoderConfig::new(32, 32);
263        let mut enc = Encoder::new(cfg).unwrap();
264        let frame = YuvFrame::black(32, 32);
265        let au = enc.encode(&frame);
266
267        let nals = split_annex_b(&au);
268        assert_eq!(nals.len(), 3);
269        assert_eq!(NalUnitType::from_id(nals[0][0]), NalUnitType::Sps);
270        assert_eq!(NalUnitType::from_id(nals[1][0]), NalUnitType::Pps);
271        assert_eq!(NalUnitType::from_id(nals[2][0]), NalUnitType::IdrSlice);
272    }
273
274    #[test]
275    fn encode_all_matches_sequential_cqp() {
276        // GOP-parallel batch encoding must be byte-identical to frame-by-frame
277        // sequential encoding at constant QP (GOPs are independent).
278        let (w, h) = (48usize, 32usize);
279        let mut cfg = EncoderConfig::new(w, h);
280        cfg.gop_size = 4; // 10 frames → 3 GOPs (4,4,2)
281        let frames: Vec<YuvFrame> = (0..10u8)
282            .map(|t| YuvFrame {
283                width: w,
284                height: h,
285                y: (0..w * h).map(|i| (i as u8).wrapping_add(t.wrapping_mul(7))).collect(),
286                u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
287                v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
288            })
289            .collect();
290        let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
291        let seq: Vec<Vec<u8>> = frames.iter().map(|f| seq_enc.encode(f)).collect();
292        let par = Encoder::new(cfg).unwrap().encode_all(&frames).unwrap();
293        assert_eq!(seq, par, "GOP-parallel must equal sequential at CQP");
294    }
295
296    #[test]
297    fn rejects_mismatched_frame() {
298        let cfg = EncoderConfig::new(16, 16);
299        let mut enc = Encoder::new(cfg).unwrap();
300        let frame = YuvFrame::black(32, 16);
301        assert_eq!(enc.try_encode(&frame), Err(EncodeError::FrameMismatch));
302    }
303}
rusty_h264_encoder/lib.rs

rusty_h264_encoder/
lib.rs
