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}