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

dsi_bitstream/impls/
buf_bit_writer.rs
