Skip to main content

irox_bits/buf/
alloc.rs

1// SPDX-License-Identifier: MIT
2// Copyright 2025 IROX Contributors
3//
4
5use crate::{Bits, BitsWrapper, BufBits, Error, MutBits};
6use alloc::collections::VecDeque;
7
8pub struct BitsBuffer<'a, T> {
9    inner: BitsWrapper<'a, T>,
10    buf: VecDeque<u8>,
11    flush_every_n: Option<usize>,
12}
13impl<'a, T> BitsBuffer<'a, T> {
14    pub fn new(inner: BitsWrapper<'a, T>) -> Self {
15        Self {
16            inner,
17            buf: VecDeque::with_capacity(16384),
18            flush_every_n: None,
19        }
20    }
21}
22
23impl<'a, T: 'a> Bits for BitsBuffer<'a, T>
24where
25    BitsWrapper<'a, T>: Bits,
26{
27    fn next_u8(&mut self) -> Result<Option<u8>, Error> {
28        if let Some(b) = self.buf.pop_front() {
29            return Ok(Some(b));
30        }
31
32        let _read = self.inner.read_some_into(&mut self.buf)?;
33
34        Ok(self.buf.pop_front())
35    }
36
37    fn read_some_into<R: MutBits>(&mut self, buf: &mut R) -> Result<usize, Error> {
38        self.inner.read_some_into(buf)
39    }
40}
41impl<'a, T: 'a> BufBits for BitsBuffer<'a, T>
42where
43    BitsWrapper<'a, T>: Bits,
44{
45    fn fill_buf(&mut self) -> Result<&[u8], Error> {
46        if self.buf.is_empty() {
47            let _read = self.inner.read_some_into(&mut self.buf)?;
48        }
49        let (a, b) = self.buf.as_slices();
50        if a.is_empty() {
51            return Ok(b);
52        }
53        Ok(a)
54    }
55
56    fn consume(&mut self, amt: usize) {
57        self.buf.drain(..amt);
58    }
59}
60impl<T> BitsBuffer<'_, T> {
61    pub fn flush_every_n(&mut self, n: usize) {
62        self.flush_every_n = Some(n);
63    }
64}
65impl<'a, T: 'a> MutBits for BitsBuffer<'a, T>
66where
67    BitsWrapper<'a, T>: MutBits,
68{
69    fn write_u8(&mut self, val: u8) -> Result<(), Error> {
70        self.buf.push_back(val);
71        if let Some(flush_every) = self.flush_every_n {
72            if self.buf.len() >= flush_every {
73                self.flush()?;
74            }
75        }
76        Ok(())
77    }
78
79    fn write_some_bytes(&mut self, mut val: &[u8]) -> usize {
80        let len = val.len();
81        let buflen = self.buf.len();
82        let mut written = 0;
83        if let Some(flush_every) = self.flush_every_n {
84            if buflen + len > flush_every {
85                let (a, b) = val.split_at(flush_every - buflen);
86                val = b;
87                let _ = self.buf.write_all_bytes(a);
88                if let Err(_e) = self.flush() {
89                    return written;
90                }
91                written += a.len();
92            }
93        }
94        let _ = self.buf.write_all_bytes(val);
95        written += val.len();
96        if let Some(flush_every) = self.flush_every_n {
97            if self.buf.len() == flush_every {
98                let _ = self.flush();
99            }
100        }
101        written
102    }
103
104    fn flush(&mut self) -> Result<(), Error> {
105        if self.buf.is_empty() {
106            return Ok(());
107        }
108        let slice = self.buf.make_contiguous();
109        self.inner.write_all_bytes(slice)?;
110        self.buf.clear();
111        Ok(())
112    }
113}