bitstream_io/huffman.rs
1// Copyright 2017 Brian Langenberger
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
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! Traits and implementations for reading or writing Huffman codes
10//! from or to a stream.
11
12#![warn(missing_docs)]
13
14/// A trait for building a final value from individual bits
15///
16/// Though similar to the [`crate::read::FromBitStream`] trait,
17/// this is intended to parse short symbols from a stream of bits
18/// while `FromBitStream` is meant for parsing larger structs from
19/// a whole reader.
20/// For example, one might have several [`FromBits`] implementations
21/// in a single program that all generate `i32` symbols from bits,
22/// but implementing `FromBitStream` multiple times for `i32`
23/// isn't possible (or practical).
24pub trait FromBits {
25 /// Our returned symbol type
26 type Symbol;
27
28 /// Given a fallable bit generator, return our output type
29 ///
30 /// # Errors
31 ///
32 /// Passes along any error from the bit generator
33 fn from_bits<F, E>(next: F) -> Result<Self::Symbol, E>
34 where
35 F: FnMut() -> Result<bool, E>;
36}
37
38/// For building individual bits from a final value
39///
40/// Though similar to the [`crate::write::ToBitStream`] trait,
41/// this is intended to generate a stream of bits from short symbols
42/// while `ToBitStream` is meant for writing larger structs to
43/// a whole writer.
44/// For example, one might have several [`ToBits`] implementations
45/// in a single program that all write `i32` symbols to bits,
46/// but implementing `ToBitStream` multiple times for `i32`
47/// isn't possible (or practical).
48pub trait ToBits {
49 /// The type we accept to output
50 type Symbol;
51
52 /// Given a value to generate, write out bits as needed.
53 ///
54 /// Outputs nothing if the symbol isn't defined.
55 ///
56 /// # Errors
57 ///
58 /// Passes along any error from the bit generator
59 fn to_bits<F, E>(value: Self::Symbol, write: F) -> Result<(), E>
60 where
61 F: FnMut(bool) -> Result<(), E>;
62}
63
64/// Defines a new Huffman tree for reading and writing
65///
66/// Its syntax is: `define_huffman_tree!(name : type = nodes)`
67/// where `name` is some identifier to identify the tree in the
68/// macro's current scope, `type` is the tree's output
69/// type (which should implement `Copy` and `Eq`), and `nodes` is either a
70/// final leaf value or a `[bit_0, bit_1]` pair where `bit_0` is
71/// the tree visited on a `0` bit, and `bit_1` is the tree visited
72/// on a `1` bit.
73///
74/// # Example
75///
76/// ```
77/// use bitstream_io::{define_huffman_tree, huffman::FromBits};
78/// define_huffman_tree!(TreeName : &'static str = ["bit 0", ["bit 1->0", "bit 1->1"]]);
79/// let mut bits = [true, false].iter().copied();
80/// assert_eq!(TreeName::from_bits(|| bits.next().ok_or(())).unwrap(), "bit 1->0");
81/// ```
82#[macro_export]
83macro_rules! define_huffman_tree {
84 ($name:ident : $type:ty = $nodes:tt) => {
85 #[derive(Copy, Clone, Debug)]
86 struct $name;
87
88 impl $crate::huffman::FromBits for $name {
89 type Symbol = $type;
90
91 fn from_bits<F, E>(mut next: F) -> Result<Self::Symbol, E>
92 where
93 F: FnMut() -> Result<bool, E>,
94 {
95 $crate::compile_read_tree_nodes!(next, $nodes)
96 }
97 }
98
99 impl $crate::huffman::ToBits for $name {
100 type Symbol = $type;
101
102 fn to_bits<F, E>(value: Self::Symbol, mut write: F) -> Result<(), E>
103 where
104 F: FnMut(bool) -> Result<(), E>
105 {
106 $crate::compile_write_tree_nodes!(value ; write ; $nodes ; );
107 Ok(())
108 }
109 }
110 };
111}
112
113/// A helper macro for compiling individual Huffman tree nodes
114#[macro_export]
115macro_rules! compile_read_tree_nodes {
116 ($next:ident , [$bit_0:tt, $bit_1:tt]) => {
117 if $next()? {
118 $crate::compile_read_tree_nodes!($next, $bit_1)
119 } else {
120 $crate::compile_read_tree_nodes!($next, $bit_0)
121 }
122 };
123 ($next:ident , $final:tt) => {
124 Ok($final)
125 };
126}
127
128/// A helper macro for compiling individual Huffman tree nodes
129#[macro_export]
130macro_rules! compile_write_tree_nodes {
131 ($value:ident ; $write:ident ; [$bit_0:tt, $bit_1:tt] ; ) => {
132 $crate::compile_write_tree_nodes!($value ; $write ; $bit_0 ; false);
133 $crate::compile_write_tree_nodes!($value ; $write ; $bit_1 ; true);
134 };
135 ($value:ident ; $write:ident ; [$bit_0:tt, $bit_1:tt] ; $($bits:tt),*) => {
136 $crate::compile_write_tree_nodes!($value ; $write ; $bit_0 ; $($bits),* , false);
137 $crate::compile_write_tree_nodes!($value ; $write ; $bit_1 ; $($bits),* , true);
138 };
139 ($value:ident ; $write:ident ; $final:tt ; $( $bits:tt),* ) => {
140 if $value == $final {
141 $( $write($bits)?; )*
142 return Ok(());
143 }
144 };
145}
146
147/// A limited unary reader which stops at the given maximum.
148///
149/// Counts non-`STOP_BIT` values (which must be 0 or 1)
150/// until `STOP_BIT`, or until `MAXIMUM` is reached.
151/// Returns the number of non-`STOP_BIT` bits, or `None` if
152/// maximum is reached beforehand.
153///
154/// # Examples
155/// ```
156/// use bitstream_io::{BitReader, BitRead, BigEndian, huffman::LimitedUnary};
157///
158/// let data: &[u8] = &[0b001_00000, 0b1111_1111];
159/// let mut r = BitReader::endian(data, BigEndian);
160/// // get 2 bits until the next 1 bit
161/// assert_eq!(r.read_huffman::<LimitedUnary<1, 5>>().unwrap(), Some(2));
162/// // but 5 bits in a row is our maximum
163/// assert_eq!(r.read_huffman::<LimitedUnary<1, 5>>().unwrap(), None);
164/// // the remaining 8 bits are ok to be read
165/// assert_eq!(r.read::<8, u8>().unwrap(), 0b1111_1111);
166/// ```
167///
168/// ```
169/// use bitstream_io::{BitWriter, BitWrite, BigEndian, huffman::LimitedUnary};
170///
171/// let mut w = BitWriter::endian(vec![], BigEndian);
172/// // writes 2 as a regular unary value which stops at the 1 bit
173/// w.write_huffman::<LimitedUnary<1, 5>>(Some(2)).unwrap();
174/// // writing values beyond the maximum does nothing
175/// w.write_huffman::<LimitedUnary<1, 5>>(Some(10)).unwrap();
176/// // writes 5, 0 bits (which is our maximum)
177/// w.write_huffman::<LimitedUnary<1, 5>>(None).unwrap();
178/// // write some 1 bits to pad out the stream
179/// w.write::<8, u8>(0b1111_1111);
180///
181/// assert_eq!(w.into_writer(), &[0b001_00000, 0b1111_1111]);
182/// ```
183#[derive(Copy, Clone, Debug)]
184pub struct LimitedUnary<const STOP_BIT: u8, const MAXIMUM: u32>;
185
186impl<const STOP_BIT: u8, const MAXIMUM: u32> FromBits for LimitedUnary<STOP_BIT, MAXIMUM> {
187 type Symbol = Option<u32>;
188
189 fn from_bits<F, E>(mut next: F) -> Result<Self::Symbol, E>
190 where
191 F: FnMut() -> Result<bool, E>,
192 {
193 const {
194 assert!(matches!(STOP_BIT, 0 | 1), "stop bit must be 0 or 1");
195 }
196
197 let mut bits = 0;
198 while bits < MAXIMUM {
199 if next()?
200 != match STOP_BIT {
201 0 => false,
202 1 => true,
203 _ => unreachable!(),
204 }
205 {
206 bits += 1;
207 } else {
208 return Ok(Some(bits));
209 }
210 }
211 Ok(None)
212 }
213}
214
215impl<const STOP_BIT: u8, const MAXIMUM: u32> ToBits for LimitedUnary<STOP_BIT, MAXIMUM> {
216 type Symbol = Option<u32>;
217
218 fn to_bits<F, E>(value: Option<u32>, mut write: F) -> Result<(), E>
219 where
220 F: FnMut(bool) -> Result<(), E>,
221 {
222 const {
223 assert!(matches!(STOP_BIT, 0 | 1), "stop bit must be 0 or 1");
224 }
225
226 match value {
227 Some(bits) if bits < MAXIMUM => {
228 (0..bits).try_for_each(|_| {
229 write(match STOP_BIT {
230 0 => true,
231 1 => false,
232 _ => unreachable!(),
233 })
234 })?;
235 write(match STOP_BIT {
236 0 => false,
237 1 => true,
238 _ => unreachable!(),
239 })
240 }
241 Some(_) => {
242 /*more bits than MAXIMUM, so output nothing*/
243 Ok(())
244 }
245 None => (0..MAXIMUM).try_for_each(|_| {
246 write(match STOP_BIT {
247 0 => true,
248 1 => false,
249 _ => unreachable!(),
250 })
251 }),
252 }
253 }
254}