dsi_bitstream/impls/
buf_bit_writer.rs

1/*
2 * SPDX-FileCopyrightText: 2023 Tommaso Fontana
3 * SPDX-FileCopyrightText: 2023 Inria
4 * SPDX-FileCopyrightText: 2023 Sebastiano Vigna
5 *
6 * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
7 */
8
9use core::any::TypeId;
10use core::{mem, ptr};
11use std::io::BufWriter;
12use std::path::Path;
13
14use crate::codes::params::{DefaultWriteParams, WriteParams};
15use crate::impls::WordAdapter;
16use crate::traits::*;
17use common_traits::{AsBytes, CastableInto, FiniteRangeNumber, Integer, Number};
18#[cfg(feature = "mem_dbg")]
19use mem_dbg::{MemDbg, MemSize};
20
21/// An implementation of [`BitWrite`] for a [`WordWrite`].
22///
23/// This implementation uses a bit buffer to store bits that are not yet
24/// written. The size of the bit buffer is the size of the word used by the
25/// [`WordWrite`], which on most platform should be `usize`.
26///
27/// The additional type parameter `WP` is used to select the parameters for the
28/// instantaneous codes, but the casual user should be happy with the default
29/// value. See [`WriteParams`] for more details.
30///
31/// The convenience methods [`from_path`] and [`from_file`] create a
32/// [`BufBitWriter`] around a buffered file writer.
33///
34/// For additional flexibility, this structures implements [`std::io::Write`].
35/// Note that because of coherence rules it is not possible to implement
36/// [`std::io::Write`] for a generic [`BitWrite`].
37
38#[derive(Debug)]
39#[cfg_attr(feature = "mem_dbg", derive(MemDbg, MemSize))]
40pub struct BufBitWriter<E: Endianness, WW: WordWrite, WP: WriteParams = DefaultWriteParams> {
41    /// The [`WordWrite`] to which we will write words.
42    backend: WW,
43    /// The buffer where we store bits until we have a word worth of them.
44    /// It might be empty, but it is never full.
45    /// Only the lower (BE) or upper (LE) `WW::Word::BITS - space_left_in_buffer`
46    /// bits are valid: the rest have undefined value.
47    buffer: WW::Word,
48    /// Number of upper (BE) or lower (LE) free bits in the buffer. It is always greater than zero.
49    space_left_in_buffer: usize,
50    _marker_endianness: core::marker::PhantomData<(E, WP)>,
51}
52
53/// Creates a new [`BufBitWriter`] with [default read
54/// parameters](`DefaultWriteParams`) from a file path using the provided
55/// endianness and read word.
56///
57/// # Examples
58///
59/// ```no_run
60/// use dsi_bitstream::prelude::*;
61/// let mut writer = buf_bit_writer::from_path::<LE, u32>("data.bin")?;
62/// # Ok::<(), Box<dyn std::error::Error>>(())
63/// ```
64pub fn from_path<E: Endianness, W: Word>(
65    path: impl AsRef<Path>,
66) -> anyhow::Result<BufBitWriter<E, WordAdapter<W, BufWriter<std::fs::File>>, DefaultWriteParams>> {
67    Ok(from_file::<E, W>(std::fs::File::create(path)?))
68}
69
70/// Creates a new [`BufBitWriter`] with [default read
71/// parameters](`DefaultWriteParams`) from a file path using the provided
72/// endianness and read word.
73///
74/// See also [`from_path`] for a version that takes a path.
75pub fn from_file<E: Endianness, W: Word>(
76    file: std::fs::File,
77) -> BufBitWriter<E, WordAdapter<W, BufWriter<std::fs::File>>, DefaultWriteParams> {
78    BufBitWriter::<E, WordAdapter<W, BufWriter<std::fs::File>>>::new(WordAdapter::new(
79        BufWriter::new(file),
80    ))
81}
82
83impl<E: Endianness, WW: WordWrite, WP: WriteParams> BufBitWriter<E, WW, WP> {
84    /// Creates a new [`BufBitWriter`] around a [`WordWrite`].
85    ///
86    /// ### Example
87    /// ```
88    /// #[cfg(not(feature = "std"))]
89    /// # fn main() {}
90    /// # #[cfg(feature = "std")]
91    /// # fn main() {
92    /// use dsi_bitstream::prelude::*;
93    /// let buffer = Vec::<usize>::new();
94    /// let word_writer = MemWordWriterVec::new(buffer);
95    /// let mut buf_bit_writer = <BufBitWriter<BE, _>>::new(word_writer);
96    /// # }
97    /// ```
98    pub fn new(backend: WW) -> Self {
99        Self {
100            backend,
101            buffer: WW::Word::ZERO,
102            space_left_in_buffer: WW::Word::BITS,
103            _marker_endianness: core::marker::PhantomData,
104        }
105    }
106}
107
108impl<E: Endianness, WW: WordWrite, WP: WriteParams> BufBitWriter<E, WW, WP>
109where
110    BufBitWriter<E, WW, WP>: BitWrite<E>,
111{
112    ///  Return the backend, consuming this writer after
113    /// [flushing it](BufBitWriter::flush).
114    pub fn into_inner(mut self) -> Result<WW, <Self as BitWrite<E>>::Error> {
115        self.flush()?;
116        // SAFETY: forget(self) prevents double dropping backend
117        let backend = unsafe { ptr::read(&self.backend) };
118        mem::forget(self);
119        Ok(backend)
120    }
121}
122
123impl<E: Endianness, WW: WordWrite, WP: WriteParams> core::ops::Drop for BufBitWriter<E, WW, WP> {
124    fn drop(&mut self) {
125        if TypeId::of::<E>() == TypeId::of::<LE>() {
126            flush_le(self).unwrap();
127        } else {
128            // TypeId::of::<E>() = TypeId::of::<BE>()
129            flush_be(self).unwrap();
130        }
131    }
132}
133
134/// Helper function flushing a [`BufBitWriter`] in big-endian fashion.
135///
136/// The endianness is hardwired because the function is called
137/// from [`BufBitWriter::drop`] using a check on the
138/// [`TypeId`] of the endianness.
139fn flush_be<E: Endianness, WW: WordWrite, WP: WriteParams>(
140    buf_bit_writer: &mut BufBitWriter<E, WW, WP>,
141) -> Result<usize, WW::Error> {
142    let to_flush = WW::Word::BITS - buf_bit_writer.space_left_in_buffer;
143    if to_flush != 0 {
144        buf_bit_writer.buffer <<= buf_bit_writer.space_left_in_buffer;
145        buf_bit_writer
146            .backend
147            .write_word(buf_bit_writer.buffer.to_be())?;
148        buf_bit_writer.space_left_in_buffer = WW::Word::BITS;
149    }
150    buf_bit_writer.backend.flush()?;
151    Ok(to_flush)
152}
153
154impl<WW: WordWrite, WP: WriteParams> BitWrite<BE> for BufBitWriter<BE, WW, WP>
155where
156    u64: CastableInto<WW::Word>,
157{
158    type Error = <WW as WordWrite>::Error;
159
160    fn flush(&mut self) -> Result<usize, Self::Error> {
161        flush_be(self)
162    }
163
164    #[allow(unused_mut)]
165    #[inline]
166    fn write_bits(&mut self, mut value: u64, n_bits: usize) -> Result<usize, Self::Error> {
167        debug_assert!(n_bits <= 64);
168        #[cfg(feature = "checks")]
169        assert!(
170            value & (1_u128 << n_bits).wrapping_sub(1) as u64 == value,
171            "Error: value {} does not fit in {} bits",
172            value,
173            n_bits
174        );
175        debug_assert!(self.space_left_in_buffer > 0);
176
177        #[cfg(test)]
178        if n_bits < 64 {
179            // We put garbage in the higher bits for testing
180            value |= u64::MAX << n_bits;
181        }
182
183        // Easy way out: we fit the buffer
184        if n_bits < self.space_left_in_buffer {
185            self.buffer <<= n_bits;
186            // Clean up bits higher than n_bits
187            self.buffer |= value.cast() & !(WW::Word::MAX << n_bits as u32);
188            self.space_left_in_buffer -= n_bits;
189            return Ok(n_bits);
190        }
191
192        // Load the bottom of the buffer, if necessary, and dump the whole buffer
193        self.buffer = self.buffer << (self.space_left_in_buffer - 1) << 1;
194        // The first shift discards bits higher than n_bits
195        self.buffer |= (value << (64 - n_bits) >> (64 - self.space_left_in_buffer)).cast();
196        self.backend.write_word(self.buffer.to_be())?;
197
198        let mut to_write = n_bits - self.space_left_in_buffer;
199
200        for _ in 0..to_write / WW::Word::BITS {
201            to_write -= WW::Word::BITS;
202            self.backend
203                .write_word((value >> to_write).cast().to_be())?;
204        }
205
206        self.space_left_in_buffer = WW::Word::BITS - to_write;
207        self.buffer = value.cast();
208        Ok(n_bits)
209    }
210
211    #[inline(always)]
212    #[allow(clippy::collapsible_if)]
213    fn write_unary(&mut self, mut value: u64) -> Result<usize, Self::Error> {
214        debug_assert!(value < u64::MAX);
215        debug_assert!(self.space_left_in_buffer > 0);
216
217        let code_length = value + 1;
218
219        // Easy way out: we fit the buffer
220        if code_length <= self.space_left_in_buffer as u64 {
221            self.space_left_in_buffer -= code_length as usize;
222            self.buffer = self.buffer << value << 1;
223            self.buffer |= WW::Word::ONE;
224            if self.space_left_in_buffer == 0 {
225                self.backend.write_word(self.buffer.to_be())?;
226                self.space_left_in_buffer = WW::Word::BITS;
227            }
228            return Ok(code_length as usize);
229        }
230
231        self.buffer = self.buffer << (self.space_left_in_buffer - 1) << 1;
232        self.backend.write_word(self.buffer.to_be())?;
233
234        value -= self.space_left_in_buffer as u64;
235
236        for _ in 0..value / WW::Word::BITS as u64 {
237            self.backend.write_word(WW::Word::ZERO)?;
238        }
239
240        value %= WW::Word::BITS as u64;
241
242        if value == WW::Word::BITS as u64 - 1 {
243            self.backend.write_word(WW::Word::ONE.to_be())?;
244            self.space_left_in_buffer = WW::Word::BITS;
245        } else {
246            self.buffer = WW::Word::ONE;
247            self.space_left_in_buffer = WW::Word::BITS - (value as usize + 1);
248        }
249
250        Ok(code_length as usize)
251    }
252
253    #[cfg(not(feature = "no_copy_impls"))]
254    fn copy_from<F: Endianness, R: BitRead<F>>(
255        &mut self,
256        bit_read: &mut R,
257        mut n: u64,
258    ) -> Result<(), CopyError<R::Error, Self::Error>> {
259        if n < self.space_left_in_buffer as u64 {
260            self.buffer = (self.buffer << n)
261                | bit_read
262                    .read_bits(n as usize)
263                    .map_err(CopyError::ReadError)?
264                    .cast();
265            self.space_left_in_buffer -= n as usize;
266            return Ok(());
267        }
268
269        self.buffer = (self.buffer << (self.space_left_in_buffer - 1) << 1)
270            | bit_read
271                .read_bits(self.space_left_in_buffer)
272                .map_err(CopyError::ReadError)?
273                .cast();
274        n -= self.space_left_in_buffer as u64;
275
276        self.backend
277            .write_word(self.buffer.to_be())
278            .map_err(CopyError::WriteError)?;
279
280        for _ in 0..n / WW::Word::BITS as u64 {
281            self.backend
282                .write_word(
283                    bit_read
284                        .read_bits(WW::Word::BITS)
285                        .map_err(CopyError::ReadError)?
286                        .cast()
287                        .to_be(),
288                )
289                .map_err(CopyError::WriteError)?;
290        }
291
292        n %= WW::Word::BITS as u64;
293        self.buffer = bit_read
294            .read_bits(n as usize)
295            .map_err(CopyError::ReadError)?
296            .cast();
297        self.space_left_in_buffer = WW::Word::BITS - n as usize;
298
299        Ok(())
300    }
301}
302
303/// Helper function flushing a [`BufBitWriter`] in big-endian fashion.
304///
305/// The endianness is hardwired because the function is called
306/// from [`BufBitWriter::drop`] using a check on the
307/// [`TypeId`] of the endianness.
308fn flush_le<E: Endianness, WW: WordWrite, WP: WriteParams>(
309    buf_bit_writer: &mut BufBitWriter<E, WW, WP>,
310) -> Result<usize, WW::Error> {
311    let to_flush = WW::Word::BITS - buf_bit_writer.space_left_in_buffer;
312    if to_flush != 0 {
313        buf_bit_writer.buffer >>= buf_bit_writer.space_left_in_buffer;
314        buf_bit_writer
315            .backend
316            .write_word(buf_bit_writer.buffer.to_le())?;
317        buf_bit_writer.space_left_in_buffer = WW::Word::BITS;
318    }
319    buf_bit_writer.backend.flush()?;
320    Ok(to_flush)
321}
322
323impl<WW: WordWrite, WP: WriteParams> BitWrite<LE> for BufBitWriter<LE, WW, WP>
324where
325    u64: CastableInto<WW::Word>,
326{
327    type Error = <WW as WordWrite>::Error;
328
329    fn flush(&mut self) -> Result<usize, Self::Error> {
330        flush_le(self)
331    }
332
333    #[inline]
334    fn write_bits(&mut self, mut value: u64, n_bits: usize) -> Result<usize, Self::Error> {
335        debug_assert!(n_bits <= 64);
336        #[cfg(feature = "checks")]
337        assert!(
338            value & (1_u128 << n_bits).wrapping_sub(1) as u64 == value,
339            "Error: value {} does not fit in {} bits",
340            value,
341            n_bits
342        );
343        debug_assert!(self.space_left_in_buffer > 0);
344
345        #[cfg(test)]
346        if n_bits < 64 {
347            // We put garbage in the higher bits for testing
348            value |= u64::MAX << n_bits;
349        }
350
351        // Easy way out: we fit the buffer
352        if n_bits < self.space_left_in_buffer {
353            self.buffer >>= n_bits;
354            // Clean up bits higher than n_bits
355            self.buffer |=
356                (value.cast() & !(WW::Word::MAX << n_bits as u32)).rotate_right(n_bits as u32);
357            self.space_left_in_buffer -= n_bits;
358            return Ok(n_bits);
359        }
360
361        // Load the top of the buffer, if necessary, and dump the whole buffer
362        self.buffer = self.buffer >> (self.space_left_in_buffer - 1) >> 1;
363        self.buffer |= value.cast() << (WW::Word::BITS - self.space_left_in_buffer);
364        self.backend.write_word(self.buffer.to_le())?;
365
366        let to_write = n_bits - self.space_left_in_buffer;
367        value = value >> (self.space_left_in_buffer - 1) >> 1;
368
369        for _ in 0..to_write / WW::Word::BITS {
370            self.backend.write_word(value.cast().to_le())?;
371            // This cannot be executed with WW::Word::BITS >= 64
372            value >>= WW::Word::BITS;
373        }
374
375        self.space_left_in_buffer = WW::Word::BITS - to_write % WW::Word::BITS;
376        self.buffer = value.cast().rotate_right(to_write as u32);
377        Ok(n_bits)
378    }
379
380    #[inline(always)]
381    #[allow(clippy::collapsible_if)]
382    fn write_unary(&mut self, mut value: u64) -> Result<usize, Self::Error> {
383        debug_assert!(value < u64::MAX);
384        debug_assert!(self.space_left_in_buffer > 0);
385
386        let code_length = value + 1;
387
388        // Easy way out: we fit the buffer
389        if code_length <= self.space_left_in_buffer as u64 {
390            self.space_left_in_buffer -= code_length as usize;
391            self.buffer = self.buffer >> value >> 1;
392            self.buffer |= WW::Word::ONE << (WW::Word::BITS - 1);
393            if self.space_left_in_buffer == 0 {
394                self.backend.write_word(self.buffer.to_le())?;
395                self.space_left_in_buffer = WW::Word::BITS;
396            }
397            return Ok(code_length as usize);
398        }
399
400        self.buffer = self.buffer >> (self.space_left_in_buffer - 1) >> 1;
401        self.backend.write_word(self.buffer.to_le())?;
402
403        value -= self.space_left_in_buffer as u64;
404
405        for _ in 0..value / WW::Word::BITS as u64 {
406            self.backend.write_word(WW::Word::ZERO)?;
407        }
408
409        value %= WW::Word::BITS as u64;
410
411        if value == WW::Word::BITS as u64 - 1 {
412            self.backend
413                .write_word((WW::Word::ONE << (WW::Word::BITS - 1)).to_le())?;
414            self.space_left_in_buffer = WW::Word::BITS;
415        } else {
416            self.buffer = WW::Word::ONE << (WW::Word::BITS - 1);
417            self.space_left_in_buffer = WW::Word::BITS - (value as usize + 1);
418        }
419
420        Ok(code_length as usize)
421    }
422
423    #[cfg(not(feature = "no_copy_impls"))]
424    fn copy_from<F: Endianness, R: BitRead<F>>(
425        &mut self,
426        bit_read: &mut R,
427        mut n: u64,
428    ) -> Result<(), CopyError<R::Error, Self::Error>> {
429        if n < self.space_left_in_buffer as u64 {
430            self.buffer = (self.buffer >> n)
431                | (bit_read
432                    .read_bits(n as usize)
433                    .map_err(CopyError::ReadError)?)
434                .cast()
435                .rotate_right(n as u32);
436            self.space_left_in_buffer -= n as usize;
437            return Ok(());
438        }
439
440        self.buffer = (self.buffer >> (self.space_left_in_buffer - 1) >> 1)
441            | (bit_read
442                .read_bits(self.space_left_in_buffer)
443                .map_err(CopyError::ReadError)?
444                .cast())
445            .rotate_right(self.space_left_in_buffer as u32);
446        n -= self.space_left_in_buffer as u64;
447
448        self.backend
449            .write_word(self.buffer.to_le())
450            .map_err(CopyError::WriteError)?;
451
452        for _ in 0..n / WW::Word::BITS as u64 {
453            self.backend
454                .write_word(
455                    bit_read
456                        .read_bits(WW::Word::BITS)
457                        .map_err(CopyError::ReadError)?
458                        .cast()
459                        .to_le(),
460                )
461                .map_err(CopyError::WriteError)?;
462        }
463
464        n %= WW::Word::BITS as u64;
465        self.buffer = bit_read
466            .read_bits(n as usize)
467            .map_err(CopyError::ReadError)?
468            .cast()
469            .rotate_right(n as u32);
470        self.space_left_in_buffer = WW::Word::BITS - n as usize;
471
472        Ok(())
473    }
474}
475
476#[cfg(feature = "std")]
477impl<WW: WordWrite, WP: WriteParams> std::io::Write for BufBitWriter<BE, WW, WP>
478where
479    u64: CastableInto<WW::Word>,
480{
481    #[inline(always)]
482    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
483        let mut iter = buf.chunks_exact(WW::Word::BYTES);
484
485        for word in &mut iter {
486            self.write_bits(u64::from_be_bytes(word.try_into().unwrap()), 64)
487                .map_err(|_| std::io::Error::other("Could not write bits to stream"))?;
488        }
489
490        let rem = iter.remainder();
491        if !rem.is_empty() {
492            let mut word = 0;
493            let bits = rem.len() * 8;
494            for byte in rem.iter() {
495                word <<= 8;
496                word |= *byte as u64;
497            }
498            self.write_bits(word, bits)
499                .map_err(|_| std::io::Error::other("Could not write bits to stream"))?;
500        }
501
502        Ok(buf.len())
503    }
504
505    fn flush(&mut self) -> std::io::Result<()> {
506        flush_be(self).map_err(|_| std::io::Error::other("Could not flush bits to stream"))?;
507        Ok(())
508    }
509}
510
511#[cfg(feature = "std")]
512impl<WW: WordWrite, WP: WriteParams> std::io::Write for BufBitWriter<LE, WW, WP>
513where
514    u64: CastableInto<WW::Word>,
515{
516    #[inline(always)]
517    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
518        let mut iter = buf.chunks_exact(WW::Word::BYTES);
519
520        for word in &mut iter {
521            self.write_bits(u64::from_le_bytes(word.try_into().unwrap()), 64)
522                .map_err(|_| std::io::Error::other("Could not write bits to stream"))?;
523        }
524
525        let rem = iter.remainder();
526        if !rem.is_empty() {
527            let mut word = 0;
528            let bits = rem.len() * 8;
529            for byte in rem.iter().rev() {
530                word <<= 8;
531                word |= *byte as u64;
532            }
533            self.write_bits(word, bits)
534                .map_err(|_| std::io::Error::other("Could not write bits to stream"))?;
535        }
536
537        Ok(buf.len())
538    }
539
540    fn flush(&mut self) -> std::io::Result<()> {
541        flush_le(self).map_err(|_| std::io::Error::other("Could not flush bits to stream"))?;
542        Ok(())
543    }
544}
545
546#[cfg(test)]
547#[cfg(feature = "std")]
548mod test {
549    use super::*;
550    use crate::prelude::MemWordWriterVec;
551    use std::io::Write;
552
553    #[test]
554    fn test_write() {
555        let data = [
556            0x90, 0x2d, 0xd0, 0x26, 0xdf, 0x89, 0xbb, 0x7e, 0x3a, 0xd6, 0xc6, 0x96, 0x73, 0xe9,
557            0x9d, 0xc9, 0x2a, 0x77, 0x82, 0xa9, 0xe6, 0x4b, 0x53, 0xcc, 0x83, 0x80, 0x4a, 0xf3,
558            0xcd, 0xe3, 0x50, 0x4e, 0x45, 0x4a, 0x3a, 0x42, 0x00, 0x4b, 0x4d, 0xbe, 0x4c, 0x88,
559            0x24, 0xf2, 0x4b, 0x6b, 0xbd, 0x79, 0xeb, 0x74, 0xbc, 0xe8, 0x7d, 0xff, 0x4b, 0x3d,
560            0xa7, 0xd6, 0x0d, 0xef, 0x9c, 0x5b, 0xb3, 0xec, 0x94, 0x97, 0xcc, 0x8b, 0x41, 0xe1,
561            0x9c, 0xcc, 0x1a, 0x03, 0x58, 0xc4, 0xfb, 0xd0, 0xc0, 0x10, 0xe2, 0xa0, 0xc9, 0xac,
562            0xa7, 0xbb, 0x50, 0xf6, 0x5c, 0x87, 0x68, 0x0f, 0x42, 0x93, 0x3f, 0x2e, 0x28, 0x28,
563            0x76, 0x83, 0x9b, 0xeb, 0x12, 0xe0, 0x4f, 0xc5, 0xb0, 0x8d, 0x14, 0xda, 0x3b, 0xdf,
564            0xd3, 0x4b, 0x80, 0xd1, 0xfc, 0x87, 0x85, 0xae, 0x54, 0xc7, 0x45, 0xc9, 0x38, 0x43,
565            0xa7, 0x9f, 0xdd, 0xa9, 0x71, 0xa7, 0x52, 0x36, 0x82, 0xff, 0x49, 0x55, 0xdb, 0x84,
566            0xc2, 0x95, 0xad, 0x45, 0x80, 0xc6, 0x02, 0x80, 0xf8, 0xfc, 0x86, 0x79, 0xae, 0xb9,
567            0x57, 0xe7, 0x3b, 0x33, 0x64, 0xa8,
568        ];
569
570        for i in 0..data.len() {
571            let mut buffer = Vec::<u64>::new();
572            let mut writer = BufBitWriter::<BE, _>::new(MemWordWriterVec::new(&mut buffer));
573
574            writer.write_all(&data[..i]).unwrap();
575            std::io::Write::flush(&mut writer).unwrap();
576
577            let buffer = writer.into_inner().unwrap().into_inner();
578            assert_eq!(unsafe { &buffer.align_to::<u8>().1[..i] }, &data[..i]);
579
580            let mut buffer = Vec::<u64>::new();
581            let mut writer = BufBitWriter::<LE, _>::new(MemWordWriterVec::new(&mut buffer));
582
583            writer.write_all(&data[..i]).unwrap();
584            std::io::Write::flush(&mut writer).unwrap();
585
586            let buffer = writer.into_inner().unwrap().into_inner();
587            assert_eq!(unsafe { &buffer.align_to::<u8>().1[..i] }, &data[..i]);
588        }
589    }
590
591    macro_rules! test_buf_bit_writer {
592        ($f: ident, $word:ty) => {
593            #[test]
594            fn $f() -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
595                #[allow(unused_imports)]
596                use crate::{
597                    codes::{GammaRead, GammaWrite},
598                    prelude::{
599                        BufBitReader, DeltaRead, DeltaWrite, MemWordReader, len_delta, len_gamma,
600                    },
601                };
602                use rand::Rng;
603                use rand::{SeedableRng, rngs::SmallRng};
604
605                let mut buffer_be: Vec<$word> = vec![];
606                let mut buffer_le: Vec<$word> = vec![];
607                let mut big = BufBitWriter::<BE, _>::new(MemWordWriterVec::new(&mut buffer_be));
608                let mut little = BufBitWriter::<LE, _>::new(MemWordWriterVec::new(&mut buffer_le));
609
610                let mut r = SmallRng::seed_from_u64(0);
611                const ITER: usize = 1_000_000;
612
613                for _ in 0..ITER {
614                    let value = r.random_range(0..128);
615                    assert_eq!(big.write_gamma(value)?, len_gamma(value));
616                    let value = r.random_range(0..128);
617                    assert_eq!(little.write_gamma(value)?, len_gamma(value));
618                    let value = r.random_range(0..128);
619                    assert_eq!(big.write_gamma(value)?, len_gamma(value));
620                    let value = r.random_range(0..128);
621                    assert_eq!(little.write_gamma(value)?, len_gamma(value));
622                    let value = r.random_range(0..128);
623                    assert_eq!(big.write_delta(value)?, len_delta(value));
624                    let value = r.random_range(0..128);
625                    assert_eq!(little.write_delta(value)?, len_delta(value));
626                    let value = r.random_range(0..128);
627                    assert_eq!(big.write_delta(value)?, len_delta(value));
628                    let value = r.random_range(0..128);
629                    assert_eq!(little.write_delta(value)?, len_delta(value));
630                    let n_bits = r.random_range(0..=64);
631                    if n_bits == 0 {
632                        big.write_bits(0, 0)?;
633                    } else {
634                        big.write_bits(r.random::<u64>() & u64::MAX >> 64 - n_bits, n_bits)?;
635                    }
636                    let n_bits = r.random_range(0..=64);
637                    if n_bits == 0 {
638                        little.write_bits(0, 0)?;
639                    } else {
640                        little.write_bits(r.random::<u64>() & u64::MAX >> 64 - n_bits, n_bits)?;
641                    }
642                    let value = r.random_range(0..128);
643                    assert_eq!(big.write_unary(value)?, value as usize + 1);
644                    let value = r.random_range(0..128);
645                    assert_eq!(little.write_unary(value)?, value as usize + 1);
646                }
647
648                drop(big);
649                drop(little);
650
651                type ReadWord = u16;
652                #[allow(clippy::size_of_in_element_count)] // false positive
653                let be_trans: &[ReadWord] = unsafe {
654                    core::slice::from_raw_parts(
655                        buffer_be.as_ptr() as *const ReadWord,
656                        buffer_be.len()
657                            * (core::mem::size_of::<$word>() / core::mem::size_of::<ReadWord>()),
658                    )
659                };
660                #[allow(clippy::size_of_in_element_count)] // false positive
661                let le_trans: &[ReadWord] = unsafe {
662                    core::slice::from_raw_parts(
663                        buffer_le.as_ptr() as *const ReadWord,
664                        buffer_le.len()
665                            * (core::mem::size_of::<$word>() / core::mem::size_of::<ReadWord>()),
666                    )
667                };
668
669                let mut big_buff = BufBitReader::<BE, _>::new(MemWordReader::new(be_trans));
670                let mut little_buff = BufBitReader::<LE, _>::new(MemWordReader::new(le_trans));
671
672                let mut r = SmallRng::seed_from_u64(0);
673
674                for _ in 0..ITER {
675                    assert_eq!(big_buff.read_gamma()?, r.random_range(0..128));
676                    assert_eq!(little_buff.read_gamma()?, r.random_range(0..128));
677                    assert_eq!(big_buff.read_gamma()?, r.random_range(0..128));
678                    assert_eq!(little_buff.read_gamma()?, r.random_range(0..128));
679                    assert_eq!(big_buff.read_delta()?, r.random_range(0..128));
680                    assert_eq!(little_buff.read_delta()?, r.random_range(0..128));
681                    assert_eq!(big_buff.read_delta()?, r.random_range(0..128));
682                    assert_eq!(little_buff.read_delta()?, r.random_range(0..128));
683                    let n_bits = r.random_range(0..=64);
684                    if n_bits == 0 {
685                        assert_eq!(big_buff.read_bits(0)?, 0);
686                    } else {
687                        assert_eq!(
688                            big_buff.read_bits(n_bits)?,
689                            r.random::<u64>() & u64::MAX >> 64 - n_bits
690                        );
691                    }
692                    let n_bits = r.random_range(0..=64);
693                    if n_bits == 0 {
694                        assert_eq!(little_buff.read_bits(0)?, 0);
695                    } else {
696                        assert_eq!(
697                            little_buff.read_bits(n_bits)?,
698                            r.random::<u64>() & u64::MAX >> 64 - n_bits
699                        );
700                    }
701
702                    assert_eq!(big_buff.read_unary()?, r.random_range(0..128));
703                    assert_eq!(little_buff.read_unary()?, r.random_range(0..128));
704                }
705
706                Ok(())
707            }
708        };
709    }
710
711    test_buf_bit_writer!(test_u128, u128);
712    test_buf_bit_writer!(test_u64, u64);
713    test_buf_bit_writer!(test_u32, u32);
714
715    test_buf_bit_writer!(test_u16, u16);
716    test_buf_bit_writer!(test_usize, usize);
717}
dsi_bitstream/impls/buf_bit_writer.rs

dsi_bitstream/impls/
buf_bit_writer.rs
