1// Copyright (c) 2018-2022, The rav1e contributors. All rights reserved
2//
3// This source code is subject to the terms of the BSD 2 Clause License and
4// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
5// was not distributed with this source code in the LICENSE file, you can
6// obtain it at www.aomedia.org/license/software. If the Alliance for Open
7// Media Patent License 1.0 was not distributed with this source code in the
8// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
9#![deny(missing_docs)]
10
11use crate::api::color::*;
12use crate::api::config::*;
13use crate::api::internal::*;
14use crate::api::util::*;
15
16use bitstream_io::*;
17
18use crate::encoder::*;
19use crate::frame::*;
20use crate::util::Pixel;
21
22use std::fmt;
23use std::io;
24use std::sync::Arc;
25
26/// The encoder context.
27///
28/// Contains the encoding state.
29pub struct Context<T: Pixel> {
30 pub(crate) inner: ContextInner<T>,
31 pub(crate) config: EncoderConfig,
32 pub(crate) pool: Option<Arc<rayon::ThreadPool>>,
33 pub(crate) is_flushing: bool,
34}
35
36impl<T: Pixel> Context<T> {
37 /// Allocates and returns a new frame.
38 ///
39 /// # Examples
40 ///
41 /// ```
42 /// use rav1e::prelude::*;
43 ///
44 /// # fn main() -> Result<(), InvalidConfig> {
45 /// let cfg = Config::default();
46 /// let ctx: Context<u8> = cfg.new_context()?;
47 /// let frame = ctx.new_frame();
48 /// # Ok(())
49 /// # }
50 /// ```
51 #[inline]
52 pub fn new_frame(&self) -> Frame<T> {
53 Frame::new(
54 self.config.width,
55 self.config.height,
56 self.config.chroma_sampling,
57 )
58 }
59
60 /// Sends the frame for encoding.
61 ///
62 /// This method adds the frame into the frame queue and runs the first passes
63 /// of the look-ahead computation.
64 ///
65 /// Passing `None` is equivalent to calling [`flush`].
66 ///
67 /// The caller is responsible for padding the invisible portion of the frame,
68 /// if multiple references to the frame are held.
69 /// Calling [`Plane::pad()`] after filling each plane or equivalent is required.
70 ///
71 /// # Errors
72 ///
73 /// If this method is called with a frame after the encoder has been flushed
74 /// or the encoder internal limit is hit (`std::i32::MAX` frames) the
75 /// [`EncoderStatus::EnoughData`] error is returned.
76 ///
77 /// # Examples
78 ///
79 /// ```
80 /// use rav1e::prelude::*;
81 ///
82 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
83 /// # if false {
84 /// let cfg = Config::default();
85 /// let mut ctx: Context<u8> = cfg.new_context().unwrap();
86 /// let f1 = ctx.new_frame();
87 /// let f2 = f1.clone();
88 /// let info = FrameParameters {
89 /// frame_type_override: FrameTypeOverride::Key,
90 /// opaque: None,
91 /// ..Default::default()
92 /// };
93 ///
94 /// // Send the plain frame data
95 /// ctx.send_frame(f1)?;
96 /// // Send the data and the per-frame parameters
97 /// // In this case the frame is forced to be a keyframe.
98 /// ctx.send_frame((f2, info))?;
99 /// // Flush the encoder, it is equivalent to a call to `flush()`
100 /// ctx.send_frame(None)?;
101 /// # }
102 /// # Ok(())
103 /// # }
104 /// ```
105 ///
106 /// [`flush`]: #method.flush
107 /// [`EncoderStatus::EnoughData`]: enum.EncoderStatus.html#variant.EnoughData
108 #[inline]
109 pub fn send_frame<F>(&mut self, frame: F) -> Result<(), EncoderStatus>
110 where
111 F: IntoFrame<T>,
112 {
113 let (frame, params) = frame.into();
114
115 if frame.is_none() {
116 if self.is_flushing {
117 return Ok(());
118 }
119 self.inner.limit = Some(self.inner.frame_count);
120 self.is_flushing = true;
121 } else if self.is_flushing
122 || (self.inner.config.still_picture && self.inner.frame_count > 0)
123 {
124 return Err(EncoderStatus::EnoughData);
125 // The rate control can process at most std::i32::MAX frames
126 } else if self.inner.frame_count == std::i32::MAX as u64 - 1 {
127 self.inner.limit = Some(self.inner.frame_count);
128 self.is_flushing = true;
129 }
130
131 let inner = &mut self.inner;
132 let run = move || inner.send_frame(frame, params);
133
134 match &self.pool {
135 Some(pool) => pool.install(run),
136 None => run(),
137 }
138 }
139
140 /// Returns the first-pass data of a two-pass encode for the frame that was
141 /// just encoded.
142 ///
143 /// This should be called BEFORE every call to [`receive_packet`] (including
144 /// the very first one), even if no packet was produced by the last call to
145 /// [`receive_packet`], if any (i.e., [`EncoderStatus::Encoded`] was
146 /// returned). It needs to be called once more after
147 /// [`EncoderStatus::LimitReached`] is returned, to retrieve the header that
148 /// should be written to the front of the stats file (overwriting the
149 /// placeholder header that was emitted at the start of encoding).
150 ///
151 /// It is still safe to call this function when [`receive_packet`] returns
152 /// any other error. It will return `None` instead of returning a duplicate
153 /// copy of the previous frame's data.
154 ///
155 /// [`receive_packet`]: #method.receive_packet
156 /// [`EncoderStatus::Encoded`]: enum.EncoderStatus.html#variant.Encoded
157 /// [`EncoderStatus::LimitReached`]:
158 /// enum.EncoderStatus.html#variant.LimitReached
159 #[inline]
160 pub fn twopass_out(&mut self) -> Option<&[u8]> {
161 self.inner.rc_state.twopass_out(self.inner.done_processing())
162 }
163
164 /// Returns the number of bytes of the stats file needed before the next
165 /// frame of the second pass in a two-pass encode can be encoded.
166 ///
167 /// This is a lower bound (more might be required), but if `0` is returned,
168 /// then encoding can proceed. This is just a hint to the application, and
169 /// does not need to be called for encoding the second pass to work, so long
170 /// as the application continues to provide more data to [`twopass_in`] in a
171 /// loop until [`twopass_in`] returns `0`.
172 ///
173 /// [`twopass_in`]: #method.twopass_in
174 #[inline]
175 pub fn twopass_bytes_needed(&mut self) -> usize {
176 self.inner.rc_state.twopass_in(None).unwrap_or(0)
177 }
178
179 /// Provides the stats data produced in the first pass of a two-pass encode
180 /// to the second pass.
181 ///
182 /// On success this returns the number of bytes of the data which were
183 /// consumed. When encoding the second pass of a two-pass encode, this should
184 /// be called repeatedly in a loop before every call to [`receive_packet`]
185 /// (including the very first one) until no bytes are consumed, or until
186 /// [`twopass_bytes_needed`] returns `0`.
187 ///
188 /// [`receive_packet`]: #method.receive_packet
189 /// [`twopass_bytes_needed`]: #method.twopass_bytes_needed
190 ///
191 /// # Errors
192 ///
193 /// Returns `Err(EncoderStatus::Failure)` if the two-pass data is invalid.
194 #[inline]
195 pub fn twopass_in(&mut self, buf: &[u8]) -> Result<usize, EncoderStatus> {
196 self.inner.rc_state.twopass_in(Some(buf)).or(Err(EncoderStatus::Failure))
197 }
198
199 /// Encodes the next frame and returns the encoded data.
200 ///
201 /// This method is where the main encoding work is done.
202 ///
203 /// # Errors
204 ///
205 /// May return `Err(EncoderStatus)`, which should be handled by the caller.
206 ///
207 /// # Examples
208 ///
209 /// Encoding a single frame:
210 ///
211 /// ```
212 /// use rav1e::prelude::*;
213 ///
214 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
215 /// # if false {
216 /// let cfg = Config::default();
217 /// let mut ctx: Context<u8> = cfg.new_context()?;
218 /// let frame = ctx.new_frame();
219 ///
220 /// ctx.send_frame(frame)?;
221 /// ctx.flush();
222 ///
223 /// loop {
224 /// match ctx.receive_packet() {
225 /// Ok(packet) => { /* Mux the packet. */ },
226 /// Err(EncoderStatus::Encoded) => (),
227 /// Err(EncoderStatus::LimitReached) => break,
228 /// Err(err) => Err(err)?,
229 /// }
230 /// }
231 /// # }
232 /// # Ok(())
233 /// # }
234 /// ```
235 ///
236 /// Encoding a sequence of frames:
237 ///
238 /// ```
239 /// use std::sync::Arc;
240 /// use rav1e::prelude::*;
241 ///
242 /// fn encode_frames(
243 /// ctx: &mut Context<u8>,
244 /// mut frames: impl Iterator<Item=Frame<u8>>
245 /// ) -> Result<(), EncoderStatus> {
246 /// // This is a slightly contrived example, intended to showcase the
247 /// // various statuses that can be returned from receive_packet().
248 /// // Assume that, for example, there are a lot of frames in the
249 /// // iterator, which are produced lazily, so you don't want to send
250 /// // them all in at once as to not exhaust the memory.
251 /// loop {
252 /// match ctx.receive_packet() {
253 /// Ok(packet) => { /* Mux the packet. */ },
254 /// Err(EncoderStatus::Encoded) => {
255 /// // A frame was encoded without emitting a packet. This is
256 /// // normal, just proceed as usual.
257 /// },
258 /// Err(EncoderStatus::LimitReached) => {
259 /// // All frames have been encoded. Time to break out of the
260 /// // loop.
261 /// break;
262 /// },
263 /// Err(EncoderStatus::NeedMoreData) => {
264 /// // The encoder has requested additional frames. Push the
265 /// // next frame in, or flush the encoder if there are no
266 /// // frames left (on None).
267 /// ctx.send_frame(frames.next().map(Arc::new))?;
268 /// },
269 /// Err(EncoderStatus::EnoughData) => {
270 /// // Since we aren't trying to push frames after flushing,
271 /// // this should never happen in this example.
272 /// unreachable!();
273 /// },
274 /// Err(EncoderStatus::NotReady) => {
275 /// // We're not doing two-pass encoding, so this can never
276 /// // occur.
277 /// unreachable!();
278 /// },
279 /// Err(EncoderStatus::Failure) => {
280 /// return Err(EncoderStatus::Failure);
281 /// },
282 /// }
283 /// }
284 ///
285 /// Ok(())
286 /// }
287 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
288 /// # if false {
289 /// # let mut enc = EncoderConfig::default();
290 /// # // So it runs faster.
291 /// # enc.width = 16;
292 /// # enc.height = 16;
293 /// # let cfg = Config::new().with_encoder_config(enc);
294 /// # let mut ctx: Context<u8> = cfg.new_context()?;
295 /// #
296 /// # let frames = vec![ctx.new_frame(); 4].into_iter();
297 /// # encode_frames(&mut ctx, frames);
298 /// # }
299 /// # Ok(())
300 /// # }
301 /// ```
302 #[inline]
303 pub fn receive_packet(&mut self) -> Result<Packet<T>, EncoderStatus> {
304 let inner = &mut self.inner;
305 let mut run = move || inner.receive_packet();
306
307 match &self.pool {
308 Some(pool) => pool.install(run),
309 None => run(),
310 }
311 }
312
313 /// Flushes the encoder.
314 ///
315 /// Flushing signals the end of the video. After the encoder has been
316 /// flushed, no additional frames are accepted.
317 ///
318 /// # Panics
319 ///
320 /// Panics if `send_frame` returns an `Err`.
321 /// This should never happen when calling it with `None`
322 /// and indicates a development error.
323 #[inline]
324 pub fn flush(&mut self) {
325 self.send_frame(None).unwrap();
326 }
327
328 /// Produces a sequence header matching the current encoding context.
329 ///
330 /// Its format is compatible with the AV1 Matroska and ISOBMFF specification.
331 /// Note that the returned header does not include any config OBUs which are
332 /// required for some uses. See [the specification].
333 ///
334 /// [the specification]:
335 /// https://aomediacodec.github.io/av1-isobmff/#av1codecconfigurationbox-section
336 ///
337 /// # Panics
338 ///
339 /// Panics if the header cannot be written in memory. This is unrecoverable,
340 /// and usually indicates the system is out of memory.
341 #[inline]
342 pub fn container_sequence_header(&self) -> Vec<u8> {
343 fn sequence_header_inner(seq: &Sequence) -> io::Result<Vec<u8>> {
344 let mut buf = Vec::new();
345
346 {
347 let mut bw = BitWriter::endian(&mut buf, BigEndian);
348 bw.write_bit(true)?; // marker
349 bw.write(7, 1)?; // version
350 bw.write(3, seq.profile)?;
351 bw.write(5, 31)?; // level
352 bw.write_bit(false)?; // tier
353 bw.write_bit(seq.bit_depth > 8)?; // high_bitdepth
354 bw.write_bit(seq.bit_depth == 12)?; // twelve_bit
355 bw.write_bit(seq.chroma_sampling == ChromaSampling::Cs400)?; // monochrome
356 bw.write_bit(seq.chroma_sampling != ChromaSampling::Cs444)?; // chroma_subsampling_x
357 bw.write_bit(seq.chroma_sampling == ChromaSampling::Cs420)?; // chroma_subsampling_y
358 bw.write(2, 0)?; // chroma_sample_position
359 bw.write(3, 0)?; // reserved
360 bw.write_bit(false)?; // initial_presentation_delay_present
361
362 bw.write(4, 0)?; // reserved
363 }
364
365 Ok(buf)
366 }
367
368 let seq = Sequence::new(&self.config);
369
370 sequence_header_inner(&seq).unwrap()
371 }
372}
373
374/// Rate Control Data
375pub enum RcData {
376 /// A Rate Control Summary Packet
377 ///
378 /// It is emitted once, after the encoder is flushed.
379 ///
380 /// It contains a summary of the rate control information for the
381 /// encoding process that just terminated.
382 Summary(Box<[u8]>),
383 /// A Rate Control Frame-specific Packet
384 ///
385 /// It is emitted every time a frame is processed.
386 ///
387 /// The information contained is required to encode its matching
388 /// frame in a second pass encoding.
389 Frame(Box<[u8]>),
390}
391
392impl<T: Pixel> Context<T> {
393 /// Return the Rate Control Summary Packet size
394 ///
395 /// It is useful mainly to preserve space when saving
396 /// both Rate Control Summary and Frame Packets in a single file.
397 pub fn rc_summary_size(&self) -> usize {
398 crate::rate::TWOPASS_HEADER_SZ
399 }
400
401 /// Return the first pass data
402 ///
403 /// Call it after `receive_packet`, it returns a packet or the encoder
404 /// lifecycle statuses [`EncoderStatus::Encoded`] and
405 /// [`EncoderStatus::LimitReached`].
406 ///
407 /// [`EncoderStatus::Encoded`]: enum.EncoderStatus.html#variant.Encoded
408 /// [`EncoderStatus::LimitReached`]:
409 /// enum.EncoderStatus.html#variant.LimitReached
410 ///
411 /// It will return a `RcData::Summary` once the encoder is flushed.
412 pub fn rc_receive_pass_data(&mut self) -> Option<RcData> {
413 if self.inner.done_processing() && self.inner.rc_state.pass1_data_retrieved
414 {
415 let data = self.inner.rc_state.emit_summary();
416 Some(RcData::Summary(data.to_vec().into_boxed_slice()))
417 } else if self.inner.rc_state.pass1_data_retrieved {
418 None
419 } else if let Some(data) = self.inner.rc_state.emit_frame_data() {
420 Some(RcData::Frame(data.to_vec().into_boxed_slice()))
421 } else {
422 unreachable!(
423 "The encoder received more frames than its internal limit allows"
424 )
425 }
426 }
427
428 /// Lower bound number of pass data packets required to progress the
429 /// encoding process.
430 ///
431 /// It should be called iteratively until it returns 0.
432 pub fn rc_second_pass_data_required(&self) -> usize {
433 if self.inner.done_processing() {
434 0
435 } else {
436 self.inner.rc_state.twopass_in_frames_needed() as usize
437 }
438 }
439
440 /// Feed the first pass Rate Control data to the encoder,
441 /// Frame-specific Packets only.
442 ///
443 /// Call it before `receive_packet()`
444 ///
445 /// # Errors
446 ///
447 /// Returns `EncoderStatus::Failure` if the data provided is incorrect
448 pub fn rc_send_pass_data(
449 &mut self, data: &[u8],
450 ) -> Result<(), EncoderStatus> {
451 self
452 .inner
453 .rc_state
454 .parse_frame_data_packet(data)
455 .map_err(|_| EncoderStatus::Failure)
456 }
457}
458
459impl<T: Pixel> fmt::Debug for Context<T> {
460 fn fmt(
461 &self, f: &mut fmt::Formatter<'_>,
462 ) -> std::result::Result<(), fmt::Error> {
463 write!(
464 f,
465 "{{ \
466 config: {:?}, \
467 is_flushing: {}, \
468 }}",
469 self.config, self.is_flushing,
470 )
471 }
472}
473