dsi_bitstream/impls/
buf_bit_reader.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 common_traits::*;
10
11use crate::codes::params::{DefaultReadParams, ReadParams};
12use crate::traits::*;
13use core::convert::Infallible;
14use core::{mem, ptr};
15#[cfg(feature = "mem_dbg")]
16use mem_dbg::{MemDbg, MemSize};
17use std::error::Error;
18
19/// An internal shortcut to the double type of the word of a
20/// [`WordRead`].
21type BB<WR> = <<WR as WordRead>::Word as DoubleType>::DoubleType;
22
23/// An implementation of [`BitRead`] and [`BitSeek`] for a [`WordRead`] and a
24/// [`WordSeek`].
25///
26/// This implementation uses a bit buffer to store bits that are not yet read.
27/// The buffer is sized as twice the word size of the underlying [`WordRead`].
28/// Typically, the best choice is to have a buffer that is sized as `usize`,
29/// which means that the word of the underlying [`WordRead`] should be half of
30/// that (i.e., `u32` for a 64-bit architecture). However, results will vary
31/// depending on the CPU.
32///
33/// The peek word is equal to the bit buffer. The value returned
34/// by [`peek_bits`](crate::traits::BitRead::peek_bits) contains at least as
35/// many bits as the word size plus one (extended with zeros beyond end of
36/// stream).
37///
38/// This implementation is usually faster than
39/// [`BitReader`](crate::impls::BitReader).
40///
41/// The additional type parameter `RP` is used to select the parameters for the
42/// instantanous codes, but the casual user should be happy with the default
43/// value. See [`ReadParams`] for more details.
44///
45/// For additional flexibility, this structures implements [`std::io::Read`].
46/// Note that because of coherence rules it is not possible to implement
47/// [`std::io::Read`] for a generic [`BitRead`].
48
49#[derive(Debug)]
50#[cfg_attr(feature = "mem_dbg", derive(MemDbg, MemSize))]
51pub struct BufBitReader<E: Endianness, WR: WordRead, RP: ReadParams = DefaultReadParams>
52where
53    WR::Word: DoubleType,
54{
55    /// The [`WordRead`] used to fill the buffer.
56    backend: WR,
57    /// The 2-word bit buffer that is used to read the codes. It is never full,
58    /// but it may be empty. Only the upper (BE) or lower (LE)
59    /// `bits_in_buffer` bits are valid; the other bits are always zeroes.
60    buffer: BB<WR>,
61    /// Number of valid upper (BE) or lower (LE) bits in the buffer.
62    /// It is always smaller than `BB::<WR>::BITS`.
63    bits_in_buffer: usize,
64    _marker: core::marker::PhantomData<(E, RP)>,
65}
66
67impl<E: Endianness, WR: WordRead + Clone, RP: ReadParams> core::clone::Clone
68    for BufBitReader<E, WR, RP>
69where
70    WR::Word: DoubleType,
71{
72    fn clone(&self) -> Self {
73        Self {
74            backend: self.backend.clone(),
75            buffer: self.buffer,
76            bits_in_buffer: self.bits_in_buffer,
77            _marker: core::marker::PhantomData,
78        }
79    }
80}
81
82impl<E: Endianness, WR: WordRead, RP: ReadParams> BufBitReader<E, WR, RP>
83where
84    WR::Word: DoubleType,
85{
86    /// Create a new [`BufBitReader`] around a [`WordRead`].
87    ///
88    /// # Example
89    /// ```
90    /// use dsi_bitstream::prelude::*;
91    /// let words: [u32; 2] = [0x0043b59f, 0xccf16077];
92    /// let word_reader = MemWordReader::new(&words);
93    /// let mut buf_bit_reader = <BufBitReader<BE, _>>::new(word_reader);
94    /// ```
95    #[must_use]
96    pub fn new(backend: WR) -> Self {
97        check_tables(WR::Word::BITS + 1);
98        Self {
99            backend,
100            buffer: BB::<WR>::ZERO,
101            bits_in_buffer: 0,
102            _marker: core::marker::PhantomData,
103        }
104    }
105
106    ///  Return the backend, consuming this reader.
107    pub fn into_inner(self) -> Result<WR, Infallible> {
108        // SAFETY: forget(self) prevents double dropping backend
109        let backend = unsafe { ptr::read(&self.backend) };
110        mem::forget(self);
111        Ok(backend)
112    }
113}
114
115//
116// Big-endian implementation
117//
118
119impl<WR: WordRead, RP: ReadParams> BufBitReader<BE, WR, RP>
120where
121    WR::Word: DoubleType,
122{
123    /// Ensure that in the buffer there are at least `WR::Word::BITS` bits to read.
124    /// This method can be called only if there are at least
125    /// `WR::Word::BITS` free bits in the buffer.
126    #[inline(always)]
127    fn refill(&mut self) -> Result<(), <WR as WordRead>::Error> {
128        debug_assert!(BB::<WR>::BITS - self.bits_in_buffer >= WR::Word::BITS);
129
130        let new_word: BB<WR> = self.backend.read_word()?.to_be().upcast();
131        self.bits_in_buffer += WR::Word::BITS;
132        self.buffer |= new_word << (BB::<WR>::BITS - self.bits_in_buffer);
133        Ok(())
134    }
135}
136
137impl<WR: WordRead, RP: ReadParams> BitRead<BE> for BufBitReader<BE, WR, RP>
138where
139    WR::Word: DoubleType + UpcastableInto<u64>,
140    BB<WR>: CastableInto<u64>,
141{
142    type Error = <WR as WordRead>::Error;
143    type PeekWord = BB<WR>;
144
145    #[inline(always)]
146    fn peek_bits(&mut self, n_bits: usize) -> Result<Self::PeekWord, Self::Error> {
147        debug_assert!(n_bits > 0);
148        debug_assert!(n_bits <= Self::PeekWord::BITS);
149
150        // A peek can do at most one refill, otherwise we might lose data
151        if n_bits > self.bits_in_buffer {
152            self.refill()?;
153        }
154
155        debug_assert!(n_bits <= self.bits_in_buffer);
156
157        // Move the n_bits highest bits of the buffer to the lowest
158        Ok(self.buffer >> (BB::<WR>::BITS - n_bits))
159    }
160
161    #[inline(always)]
162    fn skip_bits_after_peek(&mut self, n_bits: usize) {
163        self.bits_in_buffer -= n_bits;
164        self.buffer <<= n_bits;
165    }
166
167    #[inline]
168    fn read_bits(&mut self, mut n_bits: usize) -> Result<u64, Self::Error> {
169        debug_assert!(n_bits <= 64);
170        debug_assert!(self.bits_in_buffer < BB::<WR>::BITS);
171
172        // most common path, we just read the buffer
173        if n_bits <= self.bits_in_buffer {
174            // Valid right shift of BB::<WR>::BITS - n_bits, even when n_bits is zero
175            let result: u64 = (self.buffer >> (BB::<WR>::BITS - n_bits - 1) >> 1_u32).cast();
176            self.bits_in_buffer -= n_bits;
177            self.buffer <<= n_bits;
178            return Ok(result);
179        }
180
181        let mut result: u64 =
182            (self.buffer >> (BB::<WR>::BITS - 1 - self.bits_in_buffer) >> 1_u8).cast();
183        n_bits -= self.bits_in_buffer;
184
185        // Directly read to the result without updating the buffer
186        while n_bits > WR::Word::BITS {
187            let new_word: u64 = self.backend.read_word()?.to_be().upcast();
188            result = (result << WR::Word::BITS) | new_word;
189            n_bits -= WR::Word::BITS;
190        }
191
192        debug_assert!(n_bits > 0);
193        debug_assert!(n_bits <= WR::Word::BITS);
194
195        // get the final word
196        let new_word = self.backend.read_word()?.to_be();
197        self.bits_in_buffer = WR::Word::BITS - n_bits;
198        // compose the remaining bits
199        let upcasted: u64 = new_word.upcast();
200        let final_bits: u64 = (upcasted >> self.bits_in_buffer).downcast();
201        result = (result << (n_bits - 1) << 1) | final_bits;
202        // and put the rest in the buffer
203        self.buffer = (UpcastableInto::<BB<WR>>::upcast(new_word)
204            << (BB::<WR>::BITS - self.bits_in_buffer - 1))
205            << 1;
206
207        Ok(result)
208    }
209
210    #[inline]
211    fn read_unary(&mut self) -> Result<u64, Self::Error> {
212        debug_assert!(self.bits_in_buffer < BB::<WR>::BITS);
213
214        // count the zeros from the left
215        let zeros: usize = self.buffer.leading_zeros() as _;
216
217        // if we encountered an 1 in the bits_in_buffer we can return
218        if zeros < self.bits_in_buffer {
219            self.buffer = self.buffer << zeros << 1;
220            self.bits_in_buffer -= zeros + 1;
221            return Ok(zeros as u64);
222        }
223
224        let mut result: u64 = self.bits_in_buffer as _;
225
226        loop {
227            let new_word = self.backend.read_word()?.to_be();
228
229            if new_word != WR::Word::ZERO {
230                let zeros: usize = new_word.leading_zeros() as _;
231                self.buffer =
232                    UpcastableInto::<BB<WR>>::upcast(new_word) << (WR::Word::BITS + zeros) << 1;
233                self.bits_in_buffer = WR::Word::BITS - zeros - 1;
234                return Ok(result + zeros as u64);
235            }
236            result += WR::Word::BITS as u64;
237        }
238    }
239
240    #[inline]
241    fn skip_bits(&mut self, mut n_bits: usize) -> Result<(), Self::Error> {
242        debug_assert!(self.bits_in_buffer < BB::<WR>::BITS);
243        // happy case, just shift the buffer
244        if n_bits <= self.bits_in_buffer {
245            self.bits_in_buffer -= n_bits;
246            self.buffer <<= n_bits;
247            return Ok(());
248        }
249
250        n_bits -= self.bits_in_buffer;
251
252        // skip words as needed
253        while n_bits > WR::Word::BITS {
254            let _ = self.backend.read_word()?;
255            n_bits -= WR::Word::BITS;
256        }
257
258        // get the final word
259        let new_word = self.backend.read_word()?.to_be();
260        self.bits_in_buffer = WR::Word::BITS - n_bits;
261
262        self.buffer = UpcastableInto::<BB<WR>>::upcast(new_word)
263            << (BB::<WR>::BITS - 1 - self.bits_in_buffer)
264            << 1;
265
266        Ok(())
267    }
268
269    #[cfg(not(feature = "no_copy_impls"))]
270    fn copy_to<F: Endianness, W: BitWrite<F>>(
271        &mut self,
272        bit_write: &mut W,
273        mut n: u64,
274    ) -> Result<(), CopyError<Self::Error, W::Error>> {
275        let from_buffer = Ord::min(n, self.bits_in_buffer as _);
276        self.buffer = self.buffer.rotate_left(from_buffer as _);
277
278        #[allow(unused_mut)]
279        let mut self_buffer_u64: u64 = self.buffer.cast();
280
281        #[cfg(feature = "checks")]
282        {
283            // Clean up in case checks are enabled
284            if n < 64 {
285                self_buffer_u64 &= (1_u64 << n) - 1;
286            }
287        }
288
289        bit_write
290            .write_bits(self_buffer_u64, from_buffer as usize)
291            .map_err(CopyError::WriteError)?;
292        n -= from_buffer;
293
294        if n == 0 {
295            self.bits_in_buffer -= from_buffer as usize;
296            return Ok(());
297        }
298
299        while n > WR::Word::BITS as u64 {
300            bit_write
301                .write_bits(
302                    self.backend
303                        .read_word()
304                        .map_err(CopyError::ReadError)?
305                        .to_be()
306                        .upcast(),
307                    WR::Word::BITS,
308                )
309                .map_err(CopyError::WriteError)?;
310            n -= WR::Word::BITS as u64;
311        }
312
313        assert!(n > 0);
314        let new_word = self
315            .backend
316            .read_word()
317            .map_err(CopyError::ReadError)?
318            .to_be();
319        self.bits_in_buffer = WR::Word::BITS - n as usize;
320        bit_write
321            .write_bits((new_word >> self.bits_in_buffer).upcast(), n as usize)
322            .map_err(CopyError::WriteError)?;
323        self.buffer = UpcastableInto::<BB<WR>>::upcast(new_word)
324            .rotate_right(WR::Word::BITS as u32 - n as u32);
325
326        Ok(())
327    }
328}
329
330impl<
331        E: Error + Send + Sync + 'static,
332        WR: WordRead<Error = E> + WordSeek<Error = E>,
333        RP: ReadParams,
334    > BitSeek for BufBitReader<BE, WR, RP>
335where
336    WR::Word: DoubleType,
337{
338    type Error = <WR as WordSeek>::Error;
339
340    #[inline]
341    fn bit_pos(&mut self) -> Result<u64, Self::Error> {
342        Ok(self.backend.word_pos()? * WR::Word::BITS as u64 - self.bits_in_buffer as u64)
343    }
344
345    #[inline]
346    fn set_bit_pos(&mut self, bit_index: u64) -> Result<(), Self::Error> {
347        self.backend
348            .set_word_pos(bit_index / WR::Word::BITS as u64)?;
349        let bit_offset = (bit_index % WR::Word::BITS as u64) as usize;
350        self.buffer = BB::<WR>::ZERO;
351        self.bits_in_buffer = 0;
352        if bit_offset != 0 {
353            let new_word: BB<WR> = self.backend.read_word()?.to_be().upcast();
354            self.bits_in_buffer = WR::Word::BITS - bit_offset;
355            self.buffer = new_word << (BB::<WR>::BITS - self.bits_in_buffer);
356        }
357        Ok(())
358    }
359}
360
361//
362// Little-endian implementation
363//
364
365impl<WR: WordRead, RP: ReadParams> BufBitReader<LE, WR, RP>
366where
367    WR::Word: DoubleType,
368{
369    /// Ensure that in the buffer there are at least `WR::Word::BITS` bits to read.
370    /// This method can be called only if there are at least
371    /// `WR::Word::BITS` free bits in the buffer.
372    #[inline(always)]
373    fn refill(&mut self) -> Result<(), <WR as WordRead>::Error> {
374        debug_assert!(BB::<WR>::BITS - self.bits_in_buffer >= WR::Word::BITS);
375
376        let new_word: BB<WR> = self.backend.read_word()?.to_le().upcast();
377        self.buffer |= new_word << self.bits_in_buffer;
378        self.bits_in_buffer += WR::Word::BITS;
379        Ok(())
380    }
381}
382
383impl<WR: WordRead, RP: ReadParams> BitRead<LE> for BufBitReader<LE, WR, RP>
384where
385    WR::Word: DoubleType + UpcastableInto<u64>,
386    BB<WR>: CastableInto<u64>,
387{
388    type Error = <WR as WordRead>::Error;
389    type PeekWord = BB<WR>;
390
391    #[inline(always)]
392    fn peek_bits(&mut self, n_bits: usize) -> Result<Self::PeekWord, Self::Error> {
393        debug_assert!(n_bits > 0);
394        debug_assert!(n_bits <= Self::PeekWord::BITS);
395
396        // A peek can do at most one refill, otherwise we might lose data
397        if n_bits > self.bits_in_buffer {
398            self.refill()?;
399        }
400
401        debug_assert!(n_bits <= self.bits_in_buffer);
402
403        // Keep the n_bits lowest bits of the buffer
404        let shamt = BB::<WR>::BITS - n_bits;
405        Ok((self.buffer << shamt) >> shamt)
406    }
407
408    #[inline(always)]
409    fn skip_bits_after_peek(&mut self, n_bits: usize) {
410        self.bits_in_buffer -= n_bits;
411        self.buffer >>= n_bits;
412    }
413
414    #[inline]
415    fn read_bits(&mut self, mut n_bits: usize) -> Result<u64, Self::Error> {
416        debug_assert!(n_bits <= 64);
417        debug_assert!(self.bits_in_buffer < BB::<WR>::BITS);
418
419        // most common path, we just read the buffer
420        if n_bits <= self.bits_in_buffer {
421            let result: u64 = (self.buffer & ((BB::<WR>::ONE << n_bits) - BB::<WR>::ONE)).cast();
422            self.bits_in_buffer -= n_bits;
423            self.buffer >>= n_bits;
424            return Ok(result);
425        }
426
427        let mut result: u64 = self.buffer.cast();
428        let mut bits_in_res = self.bits_in_buffer;
429
430        // Directly read to the result without updating the buffer
431        while n_bits > WR::Word::BITS + bits_in_res {
432            let new_word: u64 = self.backend.read_word()?.to_le().upcast();
433            result |= new_word << bits_in_res;
434            bits_in_res += WR::Word::BITS;
435        }
436
437        n_bits -= bits_in_res;
438
439        debug_assert!(n_bits > 0);
440        debug_assert!(n_bits <= WR::Word::BITS);
441
442        // get the final word
443        let new_word = self.backend.read_word()?.to_le();
444        self.bits_in_buffer = WR::Word::BITS - n_bits;
445        // compose the remaining bits
446        let shamt = 64 - n_bits;
447        let upcasted: u64 = new_word.upcast();
448        let final_bits: u64 = ((upcasted << shamt) >> shamt).downcast();
449        result |= final_bits << bits_in_res;
450        // and put the rest in the buffer
451        self.buffer = UpcastableInto::<BB<WR>>::upcast(new_word) >> n_bits;
452
453        Ok(result)
454    }
455
456    #[inline]
457    fn read_unary(&mut self) -> Result<u64, Self::Error> {
458        debug_assert!(self.bits_in_buffer < BB::<WR>::BITS);
459
460        // count the zeros from the right
461        let zeros: usize = self.buffer.trailing_zeros() as usize;
462
463        // if we encountered an 1 in the bits_in_buffer we can return
464        if zeros < self.bits_in_buffer {
465            self.buffer = self.buffer >> zeros >> 1;
466            self.bits_in_buffer -= zeros + 1;
467            return Ok(zeros as u64);
468        }
469
470        let mut result: u64 = self.bits_in_buffer as _;
471
472        loop {
473            let new_word = self.backend.read_word()?.to_le();
474
475            if new_word != WR::Word::ZERO {
476                let zeros: usize = new_word.trailing_zeros() as _;
477                self.buffer = UpcastableInto::<BB<WR>>::upcast(new_word) >> zeros >> 1;
478                self.bits_in_buffer = WR::Word::BITS - zeros - 1;
479                return Ok(result + zeros as u64);
480            }
481            result += WR::Word::BITS as u64;
482        }
483    }
484
485    #[inline]
486    fn skip_bits(&mut self, mut n_bits: usize) -> Result<(), Self::Error> {
487        debug_assert!(self.bits_in_buffer < BB::<WR>::BITS);
488        // happy case, just shift the buffer
489        if n_bits <= self.bits_in_buffer {
490            self.bits_in_buffer -= n_bits;
491            self.buffer >>= n_bits;
492            return Ok(());
493        }
494
495        n_bits -= self.bits_in_buffer;
496
497        // skip words as needed
498        while n_bits > WR::Word::BITS {
499            let _ = self.backend.read_word()?;
500            n_bits -= WR::Word::BITS;
501        }
502
503        // get the final word
504        let new_word = self.backend.read_word()?.to_le();
505        self.bits_in_buffer = WR::Word::BITS - n_bits;
506        self.buffer = UpcastableInto::<BB<WR>>::upcast(new_word) >> n_bits;
507
508        Ok(())
509    }
510
511    #[cfg(not(feature = "no_copy_impls"))]
512    fn copy_to<F: Endianness, W: BitWrite<F>>(
513        &mut self,
514        bit_write: &mut W,
515        mut n: u64,
516    ) -> Result<(), CopyError<Self::Error, W::Error>> {
517        let from_buffer = Ord::min(n, self.bits_in_buffer as _);
518
519        #[allow(unused_mut)]
520        let mut self_buffer_u64: u64 = self.buffer.cast();
521
522        #[cfg(feature = "checks")]
523        {
524            // Clean up in case checks are enabled
525            if n < 64 {
526                self_buffer_u64 &= (1_u64 << n) - 1;
527            }
528        }
529
530        bit_write
531            .write_bits(self_buffer_u64, from_buffer as usize)
532            .map_err(CopyError::WriteError)?;
533
534        self.buffer >>= from_buffer;
535        n -= from_buffer;
536
537        if n == 0 {
538            self.bits_in_buffer -= from_buffer as usize;
539            return Ok(());
540        }
541
542        while n > WR::Word::BITS as u64 {
543            bit_write
544                .write_bits(
545                    self.backend
546                        .read_word()
547                        .map_err(CopyError::ReadError)?
548                        .to_le()
549                        .upcast(),
550                    WR::Word::BITS,
551                )
552                .map_err(CopyError::WriteError)?;
553            n -= WR::Word::BITS as u64;
554        }
555
556        assert!(n > 0);
557        let new_word = self
558            .backend
559            .read_word()
560            .map_err(CopyError::ReadError)?
561            .to_le();
562        self.bits_in_buffer = WR::Word::BITS - n as usize;
563
564        #[allow(unused_mut)]
565        let mut new_word_u64: u64 = new_word.upcast();
566
567        #[cfg(feature = "checks")]
568        {
569            // Clean up in case checks are enabled
570            if n < 64 {
571                new_word_u64 &= (1_u64 << n) - 1;
572            }
573        }
574
575        bit_write
576            .write_bits(new_word_u64, n as usize)
577            .map_err(CopyError::WriteError)?;
578        self.buffer = UpcastableInto::<BB<WR>>::upcast(new_word) >> n;
579        Ok(())
580    }
581}
582
583impl<
584        E: Error + Send + Sync + 'static,
585        WR: WordRead<Error = E> + WordSeek<Error = E>,
586        RP: ReadParams,
587    > BitSeek for BufBitReader<LE, WR, RP>
588where
589    WR::Word: DoubleType,
590{
591    type Error = <WR as WordSeek>::Error;
592
593    #[inline]
594    fn bit_pos(&mut self) -> Result<u64, Self::Error> {
595        Ok(self.backend.word_pos()? * WR::Word::BITS as u64 - self.bits_in_buffer as u64)
596    }
597
598    #[inline]
599    fn set_bit_pos(&mut self, bit_index: u64) -> Result<(), Self::Error> {
600        self.backend
601            .set_word_pos(bit_index / WR::Word::BITS as u64)?;
602
603        let bit_offset = (bit_index % WR::Word::BITS as u64) as usize;
604        self.buffer = BB::<WR>::ZERO;
605        self.bits_in_buffer = 0;
606        if bit_offset != 0 {
607            let new_word: BB<WR> = self.backend.read_word()?.to_le().upcast();
608            self.bits_in_buffer = WR::Word::BITS - bit_offset;
609            self.buffer = new_word >> bit_offset;
610        }
611        Ok(())
612    }
613}
614
615impl<WR: WordRead, RP: ReadParams> std::io::Read for BufBitReader<LE, WR, RP>
616where
617    WR::Word: DoubleType + UpcastableInto<u64>,
618    BB<WR>: CastableInto<u64>,
619{
620    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
621        let mut iter = buf.chunks_exact_mut(8);
622
623        for chunk in &mut iter {
624            let word = self
625                .read_bits(64)
626                .map_err(|_| std::io::ErrorKind::UnexpectedEof)?;
627            chunk.copy_from_slice(&word.to_le_bytes());
628        }
629
630        let rem = iter.into_remainder();
631        if !rem.is_empty() {
632            let word = self
633                .read_bits(rem.len() * 8)
634                .map_err(|_| std::io::ErrorKind::UnexpectedEof)?;
635            rem.copy_from_slice(&word.to_le_bytes()[..rem.len()]);
636        }
637
638        Ok(buf.len())
639    }
640}
641
642impl<WR: WordRead, RP: ReadParams> std::io::Read for BufBitReader<BE, WR, RP>
643where
644    WR::Word: DoubleType + UpcastableInto<u64>,
645    BB<WR>: CastableInto<u64>,
646{
647    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
648        let mut iter = buf.chunks_exact_mut(8);
649
650        for chunk in &mut iter {
651            let word = self
652                .read_bits(64)
653                .map_err(|_| std::io::ErrorKind::UnexpectedEof)?;
654            chunk.copy_from_slice(&word.to_be_bytes());
655        }
656
657        let rem = iter.into_remainder();
658        if !rem.is_empty() {
659            let word = self
660                .read_bits(rem.len() * 8)
661                .map_err(|_| std::io::ErrorKind::UnexpectedEof)?;
662            rem.copy_from_slice(&word.to_be_bytes()[8 - rem.len()..]);
663        }
664
665        Ok(buf.len())
666    }
667}
668
669#[cfg(test)]
670mod test {
671    use super::*;
672    use crate::prelude::{MemWordReader, MemWordWriterVec};
673    use std::io::Read;
674
675    #[test]
676    fn test_read() {
677        let data = [
678            0x90, 0x2d, 0xd0, 0x26, 0xdf, 0x89, 0xbb, 0x7e, 0x3a, 0xd6, 0xc6, 0x96, 0x73, 0xe9,
679            0x9d, 0xc9, 0x2a, 0x77, 0x82, 0xa9, 0xe6, 0x4b, 0x53, 0xcc, 0x83, 0x80, 0x4a, 0xf3,
680            0xcd, 0xe3, 0x50, 0x4e, 0x45, 0x4a, 0x3a, 0x42, 0x00, 0x4b, 0x4d, 0xbe, 0x4c, 0x88,
681            0x24, 0xf2, 0x4b, 0x6b, 0xbd, 0x79, 0xeb, 0x74, 0xbc, 0xe8, 0x7d, 0xff, 0x4b, 0x3d,
682            0xa7, 0xd6, 0x0d, 0xef, 0x9c, 0x5b, 0xb3, 0xec, 0x94, 0x97, 0xcc, 0x8b, 0x41, 0xe1,
683            0x9c, 0xcc, 0x1a, 0x03, 0x58, 0xc4, 0xfb, 0xd0, 0xc0, 0x10, 0xe2, 0xa0, 0xc9, 0xac,
684            0xa7, 0xbb, 0x50, 0xf6, 0x5c, 0x87, 0x68, 0x0f, 0x42, 0x93, 0x3f, 0x2e, 0x28, 0x28,
685            0x76, 0x83, 0x9b, 0xeb, 0x12, 0xe0, 0x4f, 0xc5, 0xb0, 0x8d, 0x14, 0xda, 0x3b, 0xdf,
686            0xd3, 0x4b, 0x80, 0xd1, 0xfc, 0x87, 0x85, 0xae, 0x54, 0xc7, 0x45, 0xc9, 0x38, 0x43,
687            0xa7, 0x9f, 0xdd, 0xa9, 0x71, 0xa7, 0x52, 0x36, 0x82, 0xff, 0x49, 0x55, 0xdb, 0x84,
688            0xc2, 0x95, 0xad, 0x45, 0x80, 0xc6, 0x02, 0x80, 0xf8, 0xfc, 0x86, 0x79, 0xae, 0xb9,
689            0x57, 0xe7, 0x3b, 0x33, 0x64, 0xa8,
690        ];
691        let data_u32 = unsafe { data.align_to::<u32>().1 };
692
693        for i in 0..data.len() {
694            let mut reader = BufBitReader::<LE, _>::new(MemWordReader::new(&data_u32));
695            let mut buffer = vec![0; i];
696            assert_eq!(reader.read(&mut buffer).unwrap(), i);
697            assert_eq!(&buffer, &data[..i]);
698
699            let mut reader = BufBitReader::<BE, _>::new(MemWordReader::new(&data_u32));
700            let mut buffer = vec![0; i];
701            assert_eq!(reader.read(&mut buffer).unwrap(), i);
702            assert_eq!(&buffer, &data[..i]);
703        }
704    }
705
706    macro_rules! test_buf_bit_reader {
707        ($f: ident, $word:ty) => {
708            #[test]
709            fn $f() -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
710                #[allow(unused_imports)]
711                use crate::{
712                    codes::{GammaRead, GammaWrite},
713                    prelude::{
714                        len_delta, len_gamma, BufBitWriter, DeltaRead, DeltaWrite, MemWordReader,
715                    },
716                };
717                use rand::Rng;
718                use rand::{rngs::SmallRng, SeedableRng};
719
720                let mut buffer_be: Vec<$word> = vec![];
721                let mut buffer_le: Vec<$word> = vec![];
722                let mut big = BufBitWriter::<BE, _>::new(MemWordWriterVec::new(&mut buffer_be));
723                let mut little = BufBitWriter::<LE, _>::new(MemWordWriterVec::new(&mut buffer_le));
724
725                let mut r = SmallRng::seed_from_u64(0);
726                const ITER: usize = 1_000_000;
727
728                for _ in 0..ITER {
729                    let value = r.random_range(0..128);
730                    assert_eq!(big.write_gamma(value)?, len_gamma(value));
731                    let value = r.random_range(0..128);
732                    assert_eq!(little.write_gamma(value)?, len_gamma(value));
733                    let value = r.random_range(0..128);
734                    assert_eq!(big.write_gamma(value)?, len_gamma(value));
735                    let value = r.random_range(0..128);
736                    assert_eq!(little.write_gamma(value)?, len_gamma(value));
737                    let value = r.random_range(0..128);
738                    assert_eq!(big.write_delta(value)?, len_delta(value));
739                    let value = r.random_range(0..128);
740                    assert_eq!(little.write_delta(value)?, len_delta(value));
741                    let value = r.random_range(0..128);
742                    assert_eq!(big.write_delta(value)?, len_delta(value));
743                    let value = r.random_range(0..128);
744                    assert_eq!(little.write_delta(value)?, len_delta(value));
745                    let n_bits = r.random_range(0..=64);
746                    if n_bits == 0 {
747                        big.write_bits(0, 0)?;
748                    } else {
749                        big.write_bits(1, n_bits)?;
750                    }
751                    let n_bits = r.random_range(0..=64);
752                    if n_bits == 0 {
753                        little.write_bits(0, 0)?;
754                    } else {
755                        little.write_bits(1, n_bits)?;
756                    }
757                    let value = r.random_range(0..128);
758                    assert_eq!(big.write_unary(value)?, value as usize + 1);
759                    let value = r.random_range(0..128);
760                    assert_eq!(little.write_unary(value)?, value as usize + 1);
761                }
762
763                drop(big);
764                drop(little);
765
766                type ReadWord = $word;
767
768                #[allow(clippy::size_of_in_element_count)] // false positive
769                let be_trans: &[ReadWord] = unsafe {
770                    core::slice::from_raw_parts(
771                        buffer_be.as_ptr() as *const ReadWord,
772                        buffer_be.len()
773                            * (core::mem::size_of::<$word>() / core::mem::size_of::<ReadWord>()),
774                    )
775                };
776
777                #[allow(clippy::size_of_in_element_count)] // false positive
778                let le_trans: &[ReadWord] = unsafe {
779                    core::slice::from_raw_parts(
780                        buffer_le.as_ptr() as *const ReadWord,
781                        buffer_le.len()
782                            * (core::mem::size_of::<$word>() / core::mem::size_of::<ReadWord>()),
783                    )
784                };
785
786                let mut big_buff = BufBitReader::<BE, _>::new(MemWordReader::new(be_trans));
787                let mut little_buff = BufBitReader::<LE, _>::new(MemWordReader::new(le_trans));
788
789                let mut r = SmallRng::seed_from_u64(0);
790
791                for _ in 0..ITER {
792                    assert_eq!(big_buff.read_gamma()?, r.random_range(0..128));
793                    assert_eq!(little_buff.read_gamma()?, r.random_range(0..128));
794                    assert_eq!(big_buff.read_gamma()?, r.random_range(0..128));
795                    assert_eq!(little_buff.read_gamma()?, r.random_range(0..128));
796                    assert_eq!(big_buff.read_delta()?, r.random_range(0..128));
797                    assert_eq!(little_buff.read_delta()?, r.random_range(0..128));
798                    assert_eq!(big_buff.read_delta()?, r.random_range(0..128));
799                    assert_eq!(little_buff.read_delta()?, r.random_range(0..128));
800                    let n_bits = r.random_range(0..=64);
801                    if n_bits == 0 {
802                        assert_eq!(big_buff.read_bits(0)?, 0);
803                    } else {
804                        assert_eq!(big_buff.read_bits(n_bits)?, 1);
805                    }
806                    let n_bits = r.random_range(0..=64);
807                    if n_bits == 0 {
808                        assert_eq!(little_buff.read_bits(0)?, 0);
809                    } else {
810                        assert_eq!(little_buff.read_bits(n_bits)?, 1);
811                    }
812
813                    assert_eq!(big_buff.read_unary()?, r.random_range(0..128));
814                    assert_eq!(little_buff.read_unary()?, r.random_range(0..128));
815                }
816
817                Ok(())
818            }
819        };
820    }
821
822    test_buf_bit_reader!(test_u64, u64);
823    test_buf_bit_reader!(test_u32, u32);
824
825    test_buf_bit_reader!(test_u16, u16);
826}