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

dsi_bitstream/impls/
buf_bit_writer.rs
