Skip to main content

oximedia_bitstream/write/
bit_recorder.rs

1// Copyright 2017 Brian Langenberger
2// Copyright 2024-2026 COOLJAPAN OU (Team Kitasan)
3//
4// Licensed under the Apache License, Version 2.0 or the MIT license,
5// at your option. See the LICENSE-APACHE / LICENSE-MIT files for details.
6
7//! `BitRecorder` — feature-gated playback-style writer (requires `alloc`).
8
9#![cfg(feature = "alloc")]
10
11use std::{io, vec::Vec};
12
13use super::{
14    BitCount, BitWrite, BitWriter, Counter, Endianness, Integer, Overflowed, PhantomData,
15    Primitive, SignedBitCount, SignedInteger, ToBitStream, ToBitStreamWith, UnsignedInteger,
16};
17
18/// For recording writes in order to play them back on another writer
19/// # Example
20/// ```
21/// use std::io::Write;
22/// use oximedia_bitstream::{BigEndian, BitWriter, BitWrite, BitRecorder};
23/// let mut recorder: BitRecorder<u32, BigEndian> = BitRecorder::new();
24/// recorder.write_var(1, 0b1u8).unwrap();
25/// recorder.write_var(2, 0b01u8).unwrap();
26/// recorder.write_var(5, 0b10111u8).unwrap();
27/// assert_eq!(recorder.written(), 8);
28/// let mut writer = BitWriter::endian(Vec::new(), BigEndian);
29/// recorder.playback(&mut writer);
30/// assert_eq!(writer.into_writer(), [0b10110111]);
31/// ```
32#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
33pub struct BitRecorder<N, E: Endianness> {
34    writer: BitWriter<Vec<u8>, E>,
35    phantom: PhantomData<N>,
36}
37
38impl<N: Counter, E: Endianness> BitRecorder<N, E> {
39    /// Creates new recorder
40    #[inline]
41    pub fn new() -> Self {
42        BitRecorder {
43            writer: BitWriter::new(Vec::new()),
44            phantom: PhantomData,
45        }
46    }
47
48    /// Creates new recorder sized for the given number of bytes
49    #[inline]
50    pub fn with_capacity(bytes: usize) -> Self {
51        BitRecorder {
52            writer: BitWriter::new(Vec::with_capacity(bytes)),
53            phantom: PhantomData,
54        }
55    }
56
57    /// Creates new recorder with the given endianness
58    #[inline]
59    pub fn endian(endian: E) -> Self {
60        BitRecorder {
61            writer: BitWriter::endian(Vec::new(), endian),
62            phantom: PhantomData,
63        }
64    }
65
66    /// Returns number of bits written
67    ///
68    /// # Panics
69    ///
70    /// Panics if the number of bits written is
71    /// larger than the maximum supported by the counter type.
72    /// Use [`BitRecorder::written_checked`] for a non-panicking
73    /// alternative.
74    #[inline]
75    pub fn written(&self) -> N {
76        self.written_checked()
77            .expect("writer maintains checked count when tracking is enabled")
78    }
79
80    /// Returns number of bits written
81    ///
82    /// # Errors
83    ///
84    /// Returns an error if the number of bits written overflows
85    /// our counter type.
86    #[inline]
87    pub fn written_checked(&self) -> Result<N, Overflowed> {
88        let mut written = N::try_from(self.writer.writer.len())
89            .map_err(|_| Overflowed)?
90            .checked_mul(8u8.into())?;
91
92        written.checked_add_assign(N::try_from(self.writer.bits).map_err(|_| Overflowed)?)?;
93
94        Ok(written)
95    }
96
97    /// Plays recorded writes to the given writer
98    #[inline]
99    pub fn playback<W: BitWrite>(&self, writer: &mut W) -> io::Result<()> {
100        writer.write_bytes(self.writer.writer.as_slice())?;
101        writer.write_var(self.writer.bits, self.writer.value)?;
102        Ok(())
103    }
104
105    /// Clears recorder, removing all values
106    #[inline]
107    pub fn clear(&mut self) {
108        self.writer = BitWriter::new({
109            let mut v = core::mem::take(&mut self.writer.writer);
110            v.clear();
111            v
112        });
113    }
114}
115
116impl<N: Counter, E: Endianness> Default for BitRecorder<N, E> {
117    #[inline]
118    fn default() -> Self {
119        Self::new()
120    }
121}
122
123impl<N, E> BitWrite for BitRecorder<N, E>
124where
125    E: Endianness,
126    N: Counter,
127{
128    #[inline]
129    fn write_bit(&mut self, bit: bool) -> io::Result<()> {
130        BitWrite::write_bit(&mut self.writer, bit)
131    }
132
133    #[inline]
134    fn write<const BITS: u32, I>(&mut self, value: I) -> io::Result<()>
135    where
136        I: Integer,
137    {
138        BitWrite::write::<BITS, I>(&mut self.writer, value)
139    }
140
141    #[inline]
142    fn write_const<const BITS: u32, const VALUE: u32>(&mut self) -> io::Result<()> {
143        self.writer.write_const::<BITS, VALUE>()
144    }
145
146    #[inline]
147    fn write_var<I>(&mut self, bits: u32, value: I) -> io::Result<()>
148    where
149        I: Integer,
150    {
151        self.writer.write_var(bits, value)
152    }
153
154    #[inline]
155    fn write_unsigned<const BITS: u32, U>(&mut self, value: U) -> io::Result<()>
156    where
157        U: UnsignedInteger,
158    {
159        BitWrite::write_unsigned::<BITS, U>(&mut self.writer, value)
160    }
161
162    #[inline]
163    fn write_unsigned_var<U>(&mut self, bits: u32, value: U) -> io::Result<()>
164    where
165        U: UnsignedInteger,
166    {
167        self.writer.write_unsigned_var(bits, value)
168    }
169
170    #[inline]
171    fn write_signed<const BITS: u32, S>(&mut self, value: S) -> io::Result<()>
172    where
173        S: SignedInteger,
174    {
175        BitWrite::write_signed::<BITS, S>(&mut self.writer, value)
176    }
177
178    #[inline(always)]
179    fn write_signed_var<S>(&mut self, bits: u32, value: S) -> io::Result<()>
180    where
181        S: SignedInteger,
182    {
183        self.writer.write_signed_var(bits, value)
184    }
185
186    #[inline]
187    fn write_count<const MAX: u32>(&mut self, count: BitCount<MAX>) -> io::Result<()> {
188        self.writer.write_count::<MAX>(count)
189    }
190
191    #[inline]
192    fn write_counted<const MAX: u32, I>(&mut self, bits: BitCount<MAX>, value: I) -> io::Result<()>
193    where
194        I: Integer + Sized,
195    {
196        self.writer.write_counted::<MAX, I>(bits, value)
197    }
198
199    #[inline]
200    fn write_unsigned_counted<const BITS: u32, U>(
201        &mut self,
202        bits: BitCount<BITS>,
203        value: U,
204    ) -> io::Result<()>
205    where
206        U: UnsignedInteger,
207    {
208        self.writer.write_unsigned_counted::<BITS, U>(bits, value)
209    }
210
211    #[inline]
212    fn write_signed_counted<const MAX: u32, S>(
213        &mut self,
214        bits: impl TryInto<SignedBitCount<MAX>>,
215        value: S,
216    ) -> io::Result<()>
217    where
218        S: SignedInteger,
219    {
220        self.writer.write_signed_counted::<MAX, S>(bits, value)
221    }
222
223    #[inline]
224    fn write_from<V>(&mut self, value: V) -> io::Result<()>
225    where
226        V: Primitive,
227    {
228        BitWrite::write_from::<V>(&mut self.writer, value)
229    }
230
231    #[inline]
232    fn write_as_from<F, V>(&mut self, value: V) -> io::Result<()>
233    where
234        F: Endianness,
235        V: Primitive,
236    {
237        BitWrite::write_as_from::<F, V>(&mut self.writer, value)
238    }
239
240    #[inline]
241    fn pad(&mut self, bits: u32) -> io::Result<()> {
242        BitWrite::pad(&mut self.writer, bits)
243    }
244
245    #[inline]
246    fn write_bytes(&mut self, buf: &[u8]) -> io::Result<()> {
247        BitWrite::write_bytes(&mut self.writer, buf)
248    }
249
250    #[inline]
251    fn write_unary<const STOP_BIT: u8>(&mut self, value: u32) -> io::Result<()> {
252        self.writer.write_unary::<STOP_BIT>(value)
253    }
254
255    #[inline]
256    fn build<T: ToBitStream>(&mut self, build: &T) -> Result<(), T::Error> {
257        BitWrite::build(&mut self.writer, build)
258    }
259
260    #[inline]
261    fn build_with<'a, T: ToBitStreamWith<'a>>(
262        &mut self,
263        build: &T,
264        context: &T::Context,
265    ) -> Result<(), T::Error> {
266        BitWrite::build_with(&mut self.writer, build, context)
267    }
268
269    #[inline]
270    fn byte_aligned(&self) -> bool {
271        BitWrite::byte_aligned(&self.writer)
272    }
273
274    #[inline]
275    fn byte_align(&mut self) -> io::Result<()> {
276        BitWrite::byte_align(&mut self.writer)
277    }
278
279    #[inline]
280    fn write_huffman<T>(&mut self, value: T::Symbol) -> io::Result<()>
281    where
282        T: crate::huffman::ToBits,
283    {
284        BitWrite::write_huffman::<T>(&mut self.writer, value)
285    }
286}
287
288impl<N: PartialOrd + Counter + Copy, E: Endianness> BitRecorder<N, E> {
289    /// Returns shortest option between ourself and candidate
290    ///
291    /// Executes fallible closure on emptied candidate recorder,
292    /// compares the lengths of ourself and the candidate,
293    /// and returns the shorter of the two.
294    ///
295    /// If the new candidate is shorter, we swap ourself and
296    /// the candidate so any recorder capacity can be reused.
297    ///
298    /// # Example
299    ///
300    /// ```
301    /// use oximedia_bitstream::{BitRecorder, BitWrite, BigEndian};
302    ///
303    /// let mut best = BitRecorder::<u8, BigEndian>::new();
304    /// let mut candidate = BitRecorder::new();
305    ///
306    /// // write an 8 bit value to our initial candidate
307    /// best.write::<8, u8>(0);
308    /// assert_eq!(best.written(), 8);
309    ///
310    /// // try another candidate which writes 4 bits
311    /// best = best.best(&mut candidate, |w| {
312    ///     w.write::<4, u8>(0)
313    /// }).unwrap();
314    ///
315    /// // which becomes our new best candidate
316    /// assert_eq!(best.written(), 4);
317    ///
318    /// // finally, try a not-so-best candidate
319    /// // which writes 10 bits
320    /// best = best.best(&mut candidate, |w| {
321    ///     w.write::<10, u16>(0)
322    /// }).unwrap();
323    ///
324    /// // so our best candidate remains 4 bits
325    /// assert_eq!(best.written(), 4);
326    /// ```
327    pub fn best<F>(
328        mut self,
329        candidate: &mut Self,
330        f: impl FnOnce(&mut Self) -> Result<(), F>,
331    ) -> Result<Self, F> {
332        candidate.clear();
333
334        f(candidate)?;
335
336        if candidate.written() < self.written() {
337            core::mem::swap(&mut self, candidate);
338        }
339
340        Ok(self)
341    }
342}