Skip to main content

zenith_float_num/common/
buf.rs

1//! Buffer for holding mantissa digits.
2//!
3//! Lengths of at most [`INLINE_WORDS`] words are stored inline (no heap allocation).
4//! Larger buffers promote to `Vec`.
5
6use crate::defs::Error;
7use crate::defs::Word;
8use crate::defs::WORD_BIT_SIZE;
9use core::hash::{Hash, Hasher};
10use core::ops::Deref;
11use core::ops::DerefMut;
12use core::ops::Index;
13use core::ops::IndexMut;
14use core::slice::SliceIndex;
15
16use crate::common::util::shift_slice_left;
17use crate::common::util::shift_slice_right;
18
19use alloc::vec::Vec;
20
21/// Number of words kept on the stack before allocating.
22pub const INLINE_WORDS: usize = 2;
23
24/// Buffer for holding mantissa digits.
25#[derive(Debug)]
26pub struct WordBuf {
27    inner: Storage,
28}
29
30#[derive(Debug)]
31enum Storage {
32    Inline { data: [Word; INLINE_WORDS], len: u8 },
33    Heap(Vec<Word>),
34}
35
36impl WordBuf {
37    #[inline]
38    pub fn new(sz: usize) -> Result<Self, Error> {
39        if sz <= INLINE_WORDS {
40            Ok(WordBuf {
41                inner: Storage::Inline {
42                    data: [0; INLINE_WORDS],
43                    len: sz as u8,
44                },
45            })
46        } else {
47            let mut inner = Vec::new();
48            inner.try_reserve_exact(sz)?;
49            inner.resize(sz, 0);
50            Ok(WordBuf {
51                inner: Storage::Heap(inner),
52            })
53        }
54    }
55
56    #[inline]
57    fn as_slice(&self) -> &[Word] {
58        match &self.inner {
59            Storage::Inline { data, len } => &data[..*len as usize],
60            Storage::Heap(v) => v.as_slice(),
61        }
62    }
63
64    #[inline]
65    fn as_mut_slice(&mut self) -> &mut [Word] {
66        match &mut self.inner {
67            Storage::Inline { data, len } => &mut data[..*len as usize],
68            Storage::Heap(v) => v.as_mut_slice(),
69        }
70    }
71
72    fn resize_len(&mut self, new_len: usize) -> Result<(), Error> {
73        if new_len <= self.len() {
74            self.truncate_words(new_len);
75            return Ok(());
76        }
77        match &mut self.inner {
78            Storage::Inline { len, .. } if new_len <= INLINE_WORDS => {
79                *len = new_len as u8;
80                Ok(())
81            }
82            Storage::Heap(v) => {
83                v.try_reserve(new_len - v.len())?;
84                v.resize(new_len, 0);
85                Ok(())
86            }
87            Storage::Inline { data, len } => {
88                let mut v = Vec::new();
89                v.try_reserve_exact(new_len)?;
90                v.extend_from_slice(&data[..*len as usize]);
91                v.resize(new_len, 0);
92                self.inner = Storage::Heap(v);
93                Ok(())
94            }
95        }
96    }
97
98    #[inline]
99    pub fn fill(&mut self, d: Word) {
100        self.as_mut_slice().fill(d);
101    }
102
103    #[inline]
104    pub fn len(&self) -> usize {
105        match &self.inner {
106            Storage::Inline { len, .. } => *len as usize,
107            Storage::Heap(v) => v.len(),
108        }
109    }
110
111    /// True when the buffer is stored without a heap allocation.
112    #[inline]
113    pub fn is_inline(&self) -> bool {
114        matches!(self.inner, Storage::Inline { .. })
115    }
116
117    /// Decrease length of the buffer to l bits. Data is shifted.
118    pub fn trunc_to(&mut self, l: usize) {
119        let n = (l + WORD_BIT_SIZE - 1) / WORD_BIT_SIZE;
120        let sz = self.len();
121        if n >= sz {
122            return;
123        }
124        shift_slice_right(self.as_mut_slice(), (sz - n) * WORD_BIT_SIZE);
125        self.truncate_words(n);
126    }
127
128    /// Decrease length of the buffer to l bits. Data is not moved.
129    pub fn trunc_to_2(&mut self, l: usize) {
130        let n = (l + WORD_BIT_SIZE - 1) / WORD_BIT_SIZE;
131        self.truncate_words(n);
132    }
133
134    fn truncate_words(&mut self, n: usize) {
135        match &mut self.inner {
136            Storage::Inline { len, .. } => {
137                *len = (*len).min(n as u8);
138            }
139            Storage::Heap(v) => v.truncate(n),
140        }
141    }
142
143    /// Try to extend the size to fit the precision p. Data is shifted to the left.
144    pub fn try_extend(&mut self, p: usize) -> Result<(), Error> {
145        let n = (p + WORD_BIT_SIZE - 1) / WORD_BIT_SIZE;
146        let l = self.len();
147        if n > l {
148            self.resize_len(n)?;
149            shift_slice_left(self.as_mut_slice(), (n - l) * WORD_BIT_SIZE);
150        }
151        Ok(())
152    }
153
154    /// Try to extend the size to fit the precision p. Fill new elements with 0. Data is not moved.
155    pub fn try_extend_2(&mut self, p: usize) -> Result<(), Error> {
156        let n = (p + WORD_BIT_SIZE - 1) / WORD_BIT_SIZE;
157        if n > self.len() {
158            self.resize_len(n)?;
159        }
160        Ok(())
161    }
162
163    /// Try to extend the size to fit the precision p. Data is shifted to the left by d bits.
164    pub fn try_extend_3(&mut self, p: usize, d: usize) -> Result<(), Error> {
165        let n = (p + WORD_BIT_SIZE - 1) / WORD_BIT_SIZE;
166        let l = self.len();
167        if n > l {
168            self.resize_len(n)?;
169        }
170        shift_slice_left(self.as_mut_slice(), d);
171        Ok(())
172    }
173
174    // Remove trailing words containing zeroes.
175    pub fn trunc_trailing_zeroes(&mut self) {
176        let mut n = 0;
177        for v in self.as_slice().iter() {
178            if *v == 0 {
179                n += 1;
180            } else {
181                break;
182            }
183        }
184        if n > 0 {
185            let sz = self.len();
186            shift_slice_right(self.as_mut_slice(), n * WORD_BIT_SIZE);
187            self.truncate_words(sz - n);
188        }
189    }
190
191    // Remove leading words containing zeroes.
192    pub fn trunc_leading_zeroes(&mut self) {
193        let mut n = 0;
194        for v in self.as_slice().iter().rev() {
195            if *v == 0 {
196                n += 1;
197            } else {
198                break;
199            }
200        }
201        if n > 0 {
202            let sz = self.len();
203            self.truncate_words(sz - n);
204        }
205    }
206}
207
208impl Hash for WordBuf {
209    fn hash<H: Hasher>(&self, state: &mut H) {
210        self.as_slice().hash(state);
211    }
212}
213
214impl<I: SliceIndex<[Word]>> IndexMut<I> for WordBuf {
215    #[inline]
216    fn index_mut(&mut self, index: I) -> &mut Self::Output {
217        self.as_mut_slice().index_mut(index)
218    }
219}
220
221impl<I: SliceIndex<[Word]>> Index<I> for WordBuf {
222    type Output = I::Output;
223
224    #[inline]
225    fn index(&self, index: I) -> &Self::Output {
226        self.as_slice().index(index)
227    }
228}
229
230impl Deref for WordBuf {
231    type Target = [Word];
232
233    #[inline]
234    fn deref(&self) -> &[Word] {
235        self.as_slice()
236    }
237}
238
239impl DerefMut for WordBuf {
240    #[inline]
241    fn deref_mut(&mut self) -> &mut [Word] {
242        self.as_mut_slice()
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    #[test]
251    fn test_inline_small_and_promote() {
252        let mut b = WordBuf::new(1).unwrap();
253        assert!(b.is_inline());
254        assert_eq!(b.len(), 1);
255        b[0] = 7;
256        b.try_extend_2(WORD_BIT_SIZE * 8).unwrap();
257        assert!(!b.is_inline());
258        assert_eq!(b.len(), 8);
259        assert_eq!(b[0], 7);
260    }
261
262    #[test]
263    fn huge_reserve_returns_memory_error() {
264        let words = (isize::MAX as usize) / core::mem::size_of::<Word>() + 1;
265        assert!(matches!(WordBuf::new(words), Err(Error::MemoryAllocation)));
266        let layout_err = core::alloc::Layout::from_size_align(usize::MAX, 3).unwrap_err();
267        assert_eq!(Error::from(layout_err), Error::MemoryAllocation);
268    }
269}