less_avc/writer.rs
1// Copyright 2022-2023 Andrew D. Straw.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT
5// or http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8//! Associates an encoder with a writer to allow writing encoded frames.
9
10use std::io::Write;
11
12use super::{Error, LessEncoder, Result, YCbCrImage};
13
14/// An encoding session ready to start but which has not yet necessarily encoded
15/// its first frame.
16///
17/// This mainly exists to hold the writer but defer writing until we have the
18/// first frame (in the `Configured` variant). After the first frame is written,
19/// it will be in the `Recording` variant. (The `MovedOut` variant should never
20/// be observed and represents a temporary internal state.)
21enum WriteState<W> {
22 Configured(W),
23 Recording(RecordingState<W>),
24 MovedOut,
25}
26
27impl<W: Write> WriteState<W> {
28 fn write_frame(&mut self, frame: &YCbCrImage) -> Result<()> {
29 // Temporarily replace ourself with a dummy value.
30 let orig_state = std::mem::replace(self, WriteState::MovedOut);
31 let state = match orig_state {
32 WriteState::Configured(fd) => {
33 let (initial_nal_data, encoder) = LessEncoder::new(frame)?;
34 let mut state = RecordingState { wtr: fd, encoder };
35 state
36 .wtr
37 .write_all(&initial_nal_data.sps.to_annex_b_data())?;
38 state
39 .wtr
40 .write_all(&initial_nal_data.pps.to_annex_b_data())?;
41 state
42 .wtr
43 .write_all(&initial_nal_data.frame.to_annex_b_data())?;
44 state
45 }
46 WriteState::Recording(mut state) => {
47 let encoded = state.encoder.encode(frame)?;
48 state.wtr.write_all(&encoded.to_annex_b_data())?;
49 state
50 }
51 WriteState::MovedOut => {
52 return Err(Error::InconsistentState {
53 #[cfg(feature = "backtrace")]
54 backtrace: std::backtrace::Backtrace::capture(),
55 })
56 }
57 };
58
59 // Restore ourself to the correct state.
60 *self = WriteState::Recording(state);
61
62 Ok(())
63 }
64}
65
66/// Small helper struct holding writer and encoder for an ongoing encoding
67/// session.
68struct RecordingState<W> {
69 wtr: W,
70 encoder: LessEncoder,
71}
72
73/// Write images to an [std::io::Write] implementation in `.h264` file format.
74pub struct H264Writer<W> {
75 inner: WriteState<W>,
76}
77
78impl<W: Write> H264Writer<W> {
79 /// Create a new [H264Writer] from an [std::io::Write] implementation.
80 pub fn new(wtr: W) -> Result<Self> {
81 Ok(Self {
82 inner: WriteState::Configured(wtr),
83 })
84 }
85
86 /// Retrieve the underlying [std::io::Write] implementation.
87 pub fn into_inner(self) -> W {
88 match self.inner {
89 WriteState::Configured(w) => w,
90 WriteState::Recording(state) => state.wtr,
91 WriteState::MovedOut => {
92 unreachable!("inconsistent internal state");
93 }
94 }
95 }
96
97 /// Encode and write a frame
98 pub fn write(&mut self, frame: &YCbCrImage) -> Result<()> {
99 self.inner.write_frame(frame)
100 }
101}