dsi_bitstream/utils/
count.rs

1/*
2 * SPDX-FileCopyrightText: 2023 Sebastiano Vigna
3 *
4 * SPDX-License-Identifier: Apache-2.0 OR MIT
5 */
6
7use crate::{
8    prelude::{
9        DeltaRead, DeltaWrite, GammaRead, GammaWrite, OmegaRead, OmegaWrite, PiRead, PiWrite,
10        ZetaRead, ZetaWrite, len_delta, len_gamma, len_omega, len_pi, len_zeta,
11    },
12    traits::*,
13};
14#[cfg(feature = "mem_dbg")]
15use mem_dbg::{MemDbg, MemSize};
16
17/// A wrapper around a [`BitWrite`] that keeps track of the number
18/// of bits written and optionally prints on standard error the
19/// operations performed on the stream.
20#[derive(Debug, Clone)]
21#[cfg_attr(feature = "mem_dbg", derive(MemDbg, MemSize))]
22pub struct CountBitWriter<E: Endianness, BW: BitWrite<E>, const PRINT: bool = false> {
23    bit_write: BW,
24    /// The number of bits written so far on the underlying [`BitWrite`]. Uses
25    /// u64 (bit positions are u64) so it does not wrap after 2^32 bits on
26    /// 32-bit targets. Increments widen usize losslessly (usize <= 64 bits).
27    pub bits_written: u64,
28    _marker: core::marker::PhantomData<E>,
29}
30
31impl<E: Endianness, BW: BitWrite<E>, const PRINT: bool> CountBitWriter<E, BW, PRINT> {
32    #[must_use]
33    pub const fn new(bit_write: BW) -> Self {
34        Self {
35            bit_write,
36            bits_written: 0,
37            _marker: core::marker::PhantomData,
38        }
39    }
40
41    /// Consumes this writer and returns the underlying [`BitWrite`].
42    #[must_use]
43    pub fn into_inner(self) -> BW {
44        self.bit_write
45    }
46}
47
48impl<E: Endianness, BW: BitWrite<E>, const PRINT: bool> BitWrite<E>
49    for CountBitWriter<E, BW, PRINT>
50{
51    type Error = <BW as BitWrite<E>>::Error;
52
53    fn write_bits(&mut self, value: u64, num_bits: usize) -> Result<usize, Self::Error> {
54        self.bit_write.write_bits(value, num_bits).inspect(|x| {
55            self.bits_written += *x as u64;
56            if PRINT {
57                #[cfg(feature = "std")]
58                eprintln!(
59                    "write_bits({:#016x}, {}) = {} (total = {})",
60                    value, num_bits, x, self.bits_written
61                );
62            }
63        })
64    }
65
66    fn write_unary(&mut self, n: u64) -> Result<usize, Self::Error> {
67        self.bit_write.write_unary(n).inspect(|x| {
68            self.bits_written += *x as u64;
69            if PRINT {
70                #[cfg(feature = "std")]
71                eprintln!("write_unary({}) = {} (total = {})", n, x, self.bits_written);
72            }
73        })
74    }
75
76    fn flush(&mut self) -> Result<usize, Self::Error> {
77        // Note that the flushed bits have already been counted by the
78        // write methods, so they are not added to `bits_written`
79        self.bit_write.flush().inspect(|x| {
80            let _ = x;
81            if PRINT {
82                #[cfg(feature = "std")]
83                eprintln!("flush() = {} (total = {})", x, self.bits_written);
84            }
85        })
86    }
87}
88
89impl<E: Endianness, BW: BitWrite<E> + GammaWrite<E>, const PRINT: bool> GammaWrite<E>
90    for CountBitWriter<E, BW, PRINT>
91{
92    fn write_gamma(&mut self, n: u64) -> Result<usize, BW::Error> {
93        self.bit_write.write_gamma(n).inspect(|x| {
94            self.bits_written += *x as u64;
95            if PRINT {
96                #[cfg(feature = "std")]
97                eprintln!("write_gamma({}) = {} (total = {})", n, x, self.bits_written);
98            }
99        })
100    }
101}
102
103impl<E: Endianness, BW: BitWrite<E> + DeltaWrite<E>, const PRINT: bool> DeltaWrite<E>
104    for CountBitWriter<E, BW, PRINT>
105{
106    fn write_delta(&mut self, n: u64) -> Result<usize, BW::Error> {
107        self.bit_write.write_delta(n).inspect(|x| {
108            self.bits_written += *x as u64;
109            if PRINT {
110                #[cfg(feature = "std")]
111                eprintln!("write_delta({}) = {} (total = {})", n, x, self.bits_written);
112            }
113        })
114    }
115}
116
117impl<E: Endianness, BW: BitWrite<E> + ZetaWrite<E>, const PRINT: bool> ZetaWrite<E>
118    for CountBitWriter<E, BW, PRINT>
119{
120    fn write_zeta(&mut self, n: u64, k: usize) -> Result<usize, BW::Error> {
121        self.bit_write.write_zeta(n, k).inspect(|x| {
122            self.bits_written += *x as u64;
123            if PRINT {
124                #[cfg(feature = "std")]
125                eprintln!(
126                    "write_zeta({}, {}) = {} (total = {})",
127                    n, k, x, self.bits_written
128                );
129            }
130        })
131    }
132
133    fn write_zeta3(&mut self, n: u64) -> Result<usize, BW::Error> {
134        self.bit_write.write_zeta3(n).inspect(|x| {
135            self.bits_written += *x as u64;
136            if PRINT {
137                #[cfg(feature = "std")]
138                eprintln!("write_zeta3({}) = {} (total = {})", n, x, self.bits_written);
139            }
140        })
141    }
142}
143
144impl<E: Endianness, BW: BitWrite<E> + OmegaWrite<E>, const PRINT: bool> OmegaWrite<E>
145    for CountBitWriter<E, BW, PRINT>
146{
147    fn write_omega(&mut self, n: u64) -> Result<usize, BW::Error> {
148        self.bit_write.write_omega(n).inspect(|x| {
149            self.bits_written += *x as u64;
150            if PRINT {
151                #[cfg(feature = "std")]
152                eprintln!("write_omega({}) = {} (total = {})", n, x, self.bits_written);
153            }
154        })
155    }
156}
157
158impl<E: Endianness, BW: BitWrite<E> + PiWrite<E>, const PRINT: bool> PiWrite<E>
159    for CountBitWriter<E, BW, PRINT>
160{
161    fn write_pi(&mut self, n: u64, k: usize) -> Result<usize, BW::Error> {
162        self.bit_write.write_pi(n, k).inspect(|x| {
163            self.bits_written += *x as u64;
164            if PRINT {
165                #[cfg(feature = "std")]
166                eprintln!(
167                    "write_pi({}, {}) = {} (total = {})",
168                    n, k, x, self.bits_written
169                );
170            }
171        })
172    }
173
174    fn write_pi2(&mut self, n: u64) -> Result<usize, BW::Error> {
175        self.bit_write.write_pi2(n).inspect(|x| {
176            self.bits_written += *x as u64;
177            if PRINT {
178                #[cfg(feature = "std")]
179                eprintln!("write_pi2({}) = {} (total = {})", n, x, self.bits_written);
180            }
181        })
182    }
183}
184
185impl<E: Endianness, BW: BitWrite<E> + BitSeek, const PRINT: bool> BitSeek
186    for CountBitWriter<E, BW, PRINT>
187{
188    type Error = <BW as BitSeek>::Error;
189
190    fn bit_pos(&mut self) -> Result<u64, Self::Error> {
191        self.bit_write.bit_pos()
192    }
193
194    fn set_bit_pos(&mut self, bit_pos: u64) -> Result<(), Self::Error> {
195        self.bit_write.set_bit_pos(bit_pos)
196    }
197}
198
199/// A wrapper around a [`BitRead`] that keeps track of the number
200/// of bits read and optionally prints on standard error the
201/// operations performed on the stream.
202#[derive(Debug, Clone)]
203#[cfg_attr(feature = "mem_dbg", derive(MemDbg, MemSize))]
204pub struct CountBitReader<E: Endianness, BR: BitRead<E>, const PRINT: bool = false> {
205    bit_read: BR,
206    /// The number of bits read (or skipped) so far from the underlying
207    /// [`BitRead`]. Uses u64 (see `bits_written`).
208    pub bits_read: u64,
209    _marker: core::marker::PhantomData<E>,
210}
211
212impl<E: Endianness, BR: BitRead<E>, const PRINT: bool> CountBitReader<E, BR, PRINT> {
213    #[must_use]
214    pub const fn new(bit_read: BR) -> Self {
215        Self {
216            bit_read,
217            bits_read: 0,
218            _marker: core::marker::PhantomData,
219        }
220    }
221
222    /// Consumes this reader and returns the underlying [`BitRead`].
223    #[must_use]
224    pub fn into_inner(self) -> BR {
225        self.bit_read
226    }
227}
228
229impl<E: Endianness, BR: BitRead<E>, const PRINT: bool> BitRead<E> for CountBitReader<E, BR, PRINT> {
230    type Error = <BR as BitRead<E>>::Error;
231    type PeekWord = BR::PeekWord;
232    const PEEK_BITS: usize = BR::PEEK_BITS;
233
234    fn read_bits(&mut self, num_bits: usize) -> Result<u64, Self::Error> {
235        self.bit_read.read_bits(num_bits).inspect(|x| {
236            let _ = x;
237            self.bits_read += num_bits as u64;
238            if PRINT {
239                #[cfg(feature = "std")]
240                eprintln!(
241                    "read_bits({}) = {:#016x} (total = {})",
242                    num_bits, x, self.bits_read
243                );
244            }
245        })
246    }
247
248    fn read_unary(&mut self) -> Result<u64, Self::Error> {
249        self.bit_read.read_unary().inspect(|x| {
250            self.bits_read += *x + 1;
251            if PRINT {
252                #[cfg(feature = "std")]
253                eprintln!("read_unary() = {} (total = {})", x, self.bits_read);
254            }
255        })
256    }
257
258    fn peek_bits(&mut self, n_bits: usize) -> Result<Self::PeekWord, Self::Error> {
259        self.bit_read.peek_bits(n_bits)
260    }
261
262    fn skip_bits(&mut self, n_bits: usize) -> Result<(), Self::Error> {
263        // Count only on success: a failed skip must not report the bits as
264        // skipped (previously the counter was incremented before the fallible
265        // call).
266        self.bit_read.skip_bits(n_bits).inspect(|_| {
267            self.bits_read += n_bits as u64;
268            if PRINT {
269                #[cfg(feature = "std")]
270                eprintln!("skip_bits({}) (total = {})", n_bits, self.bits_read);
271            }
272        })
273    }
274
275    fn skip_bits_after_peek(&mut self, n: usize) {
276        // This fast path advances the underlying stream, so count it too;
277        // it is infallible.
278        self.bits_read += n as u64;
279        self.bit_read.skip_bits_after_peek(n)
280    }
281}
282
283impl<E: Endianness, BR: BitRead<E> + GammaRead<E>, const PRINT: bool> GammaRead<E>
284    for CountBitReader<E, BR, PRINT>
285{
286    fn read_gamma(&mut self) -> Result<u64, BR::Error> {
287        self.bit_read.read_gamma().inspect(|x| {
288            self.bits_read += len_gamma(*x) as u64;
289            if PRINT {
290                #[cfg(feature = "std")]
291                eprintln!("read_gamma() = {} (total = {})", x, self.bits_read);
292            }
293        })
294    }
295}
296
297impl<E: Endianness, BR: BitRead<E> + DeltaRead<E>, const PRINT: bool> DeltaRead<E>
298    for CountBitReader<E, BR, PRINT>
299{
300    fn read_delta(&mut self) -> Result<u64, BR::Error> {
301        self.bit_read.read_delta().inspect(|x| {
302            self.bits_read += len_delta(*x) as u64;
303            if PRINT {
304                #[cfg(feature = "std")]
305                eprintln!("read_delta() = {} (total = {})", x, self.bits_read);
306            }
307        })
308    }
309}
310
311impl<E: Endianness, BR: BitRead<E> + ZetaRead<E>, const PRINT: bool> ZetaRead<E>
312    for CountBitReader<E, BR, PRINT>
313{
314    fn read_zeta(&mut self, k: usize) -> Result<u64, BR::Error> {
315        self.bit_read.read_zeta(k).inspect(|x| {
316            self.bits_read += len_zeta(*x, k) as u64;
317            if PRINT {
318                #[cfg(feature = "std")]
319                eprintln!("read_zeta({}) = {} (total = {})", k, x, self.bits_read);
320            }
321        })
322    }
323
324    fn read_zeta3(&mut self) -> Result<u64, BR::Error> {
325        self.bit_read.read_zeta3().inspect(|x| {
326            self.bits_read += len_zeta(*x, 3) as u64;
327            if PRINT {
328                #[cfg(feature = "std")]
329                eprintln!("read_zeta3() = {} (total = {})", x, self.bits_read);
330            }
331        })
332    }
333}
334
335impl<E: Endianness, BR: BitRead<E> + OmegaRead<E>, const PRINT: bool> OmegaRead<E>
336    for CountBitReader<E, BR, PRINT>
337{
338    fn read_omega(&mut self) -> Result<u64, BR::Error> {
339        self.bit_read.read_omega().inspect(|x| {
340            self.bits_read += len_omega(*x) as u64;
341            if PRINT {
342                #[cfg(feature = "std")]
343                eprintln!("read_omega() = {} (total = {})", x, self.bits_read);
344            }
345        })
346    }
347}
348
349impl<E: Endianness, BR: BitRead<E> + PiRead<E>, const PRINT: bool> PiRead<E>
350    for CountBitReader<E, BR, PRINT>
351{
352    fn read_pi(&mut self, k: usize) -> Result<u64, BR::Error> {
353        self.bit_read.read_pi(k).inspect(|x| {
354            self.bits_read += len_pi(*x, k) as u64;
355            if PRINT {
356                #[cfg(feature = "std")]
357                eprintln!("read_pi({}) = {} (total = {})", k, x, self.bits_read);
358            }
359        })
360    }
361
362    fn read_pi2(&mut self) -> Result<u64, BR::Error> {
363        self.bit_read.read_pi2().inspect(|x| {
364            self.bits_read += len_pi(*x, 2) as u64;
365            if PRINT {
366                #[cfg(feature = "std")]
367                eprintln!("read_pi2() = {} (total = {})", x, self.bits_read);
368            }
369        })
370    }
371}
372
373impl<E: Endianness, BR: BitRead<E> + BitSeek, const PRINT: bool> BitSeek
374    for CountBitReader<E, BR, PRINT>
375{
376    type Error = <BR as BitSeek>::Error;
377
378    fn bit_pos(&mut self) -> Result<u64, Self::Error> {
379        self.bit_read.bit_pos()
380    }
381
382    fn set_bit_pos(&mut self, bit_pos: u64) -> Result<(), Self::Error> {
383        self.bit_read.set_bit_pos(bit_pos)
384    }
385}
386
387#[cfg(test)]
388#[cfg(feature = "std")]
389mod tests {
390    use super::*;
391    use crate::prelude::*;
392
393    #[test]
394    fn test_count() -> Result<(), Box<dyn core::error::Error + Send + Sync + 'static>> {
395        let mut buffer = <Vec<u64>>::new();
396        let bit_write = <BufBitWriter<LE, _>>::new(MemWordWriterVec::new(&mut buffer));
397        let mut count_bit_write = CountBitWriter::<_, _, true>::new(bit_write);
398
399        count_bit_write.write_unary(5)?;
400        assert_eq!(count_bit_write.bits_written, 6);
401        count_bit_write.write_unary(100)?;
402        assert_eq!(count_bit_write.bits_written, 107);
403        count_bit_write.write_bits(1, 20)?;
404        assert_eq!(count_bit_write.bits_written, 127);
405        count_bit_write.write_bits(1, 33)?;
406        assert_eq!(count_bit_write.bits_written, 160);
407        count_bit_write.write_gamma(2)?;
408        assert_eq!(count_bit_write.bits_written, 163);
409        count_bit_write.write_delta(1)?;
410        assert_eq!(count_bit_write.bits_written, 167);
411        count_bit_write.write_zeta(0, 4)?;
412        assert_eq!(count_bit_write.bits_written, 171);
413        count_bit_write.write_zeta3(0)?;
414        assert_eq!(count_bit_write.bits_written, 174);
415        count_bit_write.write_omega(3)?;
416        assert_eq!(count_bit_write.bits_written, 174 + len_omega(3) as u64);
417        let after_omega = count_bit_write.bits_written;
418        count_bit_write.write_pi(5, 3)?;
419        assert_eq!(
420            count_bit_write.bits_written,
421            after_omega + len_pi(5, 3) as u64
422        );
423        let after_pi = count_bit_write.bits_written;
424        count_bit_write.write_pi2(7)?;
425        assert_eq!(count_bit_write.bits_written, after_pi + len_pi(7, 2) as u64);
426        let after_pi2 = count_bit_write.bits_written;
427        count_bit_write.flush()?;
428        // Flushing must not change the count
429        assert_eq!(count_bit_write.bits_written, after_pi2);
430        drop(count_bit_write);
431
432        let bit_read = <BufBitReader<LE, _>>::new(MemWordReader::<u64, _>::new(&buffer));
433        let mut count_bit_read = CountBitReader::<_, _, true>::new(bit_read);
434
435        assert_eq!(count_bit_read.peek_bits(5)?, 0);
436        assert_eq!(count_bit_read.read_unary()?, 5);
437        assert_eq!(count_bit_read.bits_read, 6);
438        assert_eq!(count_bit_read.read_unary()?, 100);
439        assert_eq!(count_bit_read.bits_read, 107);
440        assert_eq!(count_bit_read.read_bits(20)?, 1);
441        assert_eq!(count_bit_read.bits_read, 127);
442        count_bit_read.skip_bits(33)?;
443        assert_eq!(count_bit_read.bits_read, 160);
444        assert_eq!(count_bit_read.read_gamma()?, 2);
445        assert_eq!(count_bit_read.bits_read, 163);
446        assert_eq!(count_bit_read.read_delta()?, 1);
447        assert_eq!(count_bit_read.bits_read, 167);
448        assert_eq!(count_bit_read.read_zeta(4)?, 0);
449        assert_eq!(count_bit_read.bits_read, 171);
450        assert_eq!(count_bit_read.read_zeta3()?, 0);
451        assert_eq!(count_bit_read.bits_read, 174);
452        assert_eq!(count_bit_read.read_omega()?, 3);
453        assert_eq!(count_bit_read.bits_read, after_omega);
454        assert_eq!(count_bit_read.read_pi(3)?, 5);
455        assert_eq!(count_bit_read.bits_read, after_pi);
456        assert_eq!(count_bit_read.read_pi2()?, 7);
457        assert_eq!(count_bit_read.bits_read, after_pi2);
458
459        Ok(())
460    }
461
462    #[test]
463    fn test_skip_bits_after_peek_is_counted() {
464        let buffer = vec![0u64; 4];
465        let bit_read = <BufBitReader<LE, _>>::new(MemWordReader::<u64, _>::new(&buffer));
466        let mut r = CountBitReader::<_, _, false>::new(bit_read);
467        let _ = r.peek_bits(5).unwrap();
468        r.skip_bits_after_peek(5);
469        // peek does not consume; the post-peek skip advances the stream, so it
470        // must be counted (previously it was not).
471        assert_eq!(r.bits_read, 5);
472    }
473
474    #[test]
475    fn test_failed_skip_bits_is_not_counted() {
476        // Strict reader over an empty backend: skip_bits hits EOF.
477        let buffer = Vec::<u64>::new();
478        let bit_read = <BufBitReader<LE, _>>::new(MemWordReader::<u64, _>::new(&buffer));
479        let mut r = CountBitReader::<_, _, false>::new(bit_read);
480        assert!(r.skip_bits(1).is_err());
481        // A failed skip must not report the bits as skipped.
482        assert_eq!(r.bits_read, 0);
483    }
484}
dsi_bitstream/utils/count.rs

dsi_bitstream/utils/
count.rs
