Skip to main content

quantwave_core/
utils.rs

1use std::ops::{Index, IndexMut};
2
3#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
4pub struct RingBuffer<T> {
5    buffer: Vec<T>,
6    mask: usize,
7    head: usize,
8    tail: usize,
9    pub count: usize,
10}
11
12impl<T: Default + Clone> RingBuffer<T> {
13    pub fn new() -> Self {
14        Self::with_capacity(16)
15    }
16
17    pub fn with_capacity(capacity: usize) -> Self {
18        let pow2 = if capacity == 0 { 16 } else { (capacity + 1).next_power_of_two() };
19        Self {
20            buffer: vec![T::default(); pow2],
21            mask: pow2 - 1,
22            head: 0,
23            tail: 0,
24            count: 0,
25        }
26    }
27
28    #[inline]
29    pub fn push_back(&mut self, item: T) {
30        self.buffer[self.head & self.mask] = item;
31        self.head = self.head.wrapping_add(1);
32        self.count += 1;
33    }
34
35    #[inline]
36    pub fn pop_front(&mut self) -> Option<T> {
37        if self.count == 0 {
38            None
39        } else {
40            let item = self.buffer[self.tail & self.mask].clone();
41            self.tail = self.tail.wrapping_add(1);
42            self.count -= 1;
43            Some(item)
44        }
45    }
46
47    #[inline]
48    pub fn push_front(&mut self, item: T) {
49        self.tail = self.tail.wrapping_sub(1);
50        self.buffer[self.tail & self.mask] = item;
51        self.count += 1;
52    }
53
54    #[inline]
55    pub fn pop_back(&mut self) -> Option<T> {
56        if self.count == 0 {
57            None
58        } else {
59            self.head = self.head.wrapping_sub(1);
60            self.count -= 1;
61            Some(self.buffer[self.head & self.mask].clone())
62        }
63    }
64
65    #[inline]
66    pub fn len(&self) -> usize {
67        self.count
68    }
69
70    #[inline]
71    pub fn is_empty(&self) -> bool {
72        self.count == 0
73    }
74
75    #[inline]
76    pub fn front(&self) -> Option<&T> {
77        if self.count == 0 {
78            None
79        } else {
80            Some(&self.buffer[self.tail & self.mask])
81        }
82    }
83
84    #[inline]
85    pub fn back(&self) -> Option<&T> {
86        if self.count == 0 {
87            None
88        } else {
89            Some(&self.buffer[self.head.wrapping_sub(1) & self.mask])
90        }
91    }
92
93    #[inline]
94    pub fn clear(&mut self) {
95        self.head = 0;
96        self.tail = 0;
97        self.count = 0;
98    }
99
100    pub fn iter(&self) -> RingBufferIter<'_, T> {
101        RingBufferIter {
102            rb: self,
103            index: 0,
104        }
105    }
106
107    pub fn retain<F>(&mut self, mut f: F)
108    where
109        F: FnMut(&T) -> bool,
110    {
111        let mut new_count = 0;
112        let mut i = 0;
113        while i < self.count {
114            let idx = self.tail.wrapping_add(i) & self.mask;
115            if f(&self.buffer[idx]) {
116                if new_count != i {
117                    let dest_idx = self.tail.wrapping_add(new_count) & self.mask;
118                    self.buffer.swap(idx, dest_idx);
119                }
120                new_count += 1;
121            }
122            i += 1;
123        }
124        self.count = new_count;
125        self.head = self.tail.wrapping_add(new_count);
126    }
127    
128    pub fn get(&self, index: usize) -> Option<&T> {
129        if index < self.count {
130            Some(&self.buffer[self.tail.wrapping_add(index) & self.mask])
131        } else {
132            None
133        }
134    }
135}
136
137pub struct RingBufferIter<'a, T> {
138    rb: &'a RingBuffer<T>,
139    index: usize,
140}
141
142impl<'a, T> Iterator for RingBufferIter<'a, T> {
143    type Item = &'a T;
144
145    fn next(&mut self) -> Option<Self::Item> {
146        if self.index < self.rb.count {
147            let actual_idx = self.rb.tail.wrapping_add(self.index) & self.rb.mask;
148            self.index += 1;
149            Some(&self.rb.buffer[actual_idx])
150        } else {
151            None
152        }
153    }
154}
155
156impl<'a, T> ExactSizeIterator for RingBufferIter<'a, T> {
157    fn len(&self) -> usize {
158        self.rb.count - self.index
159    }
160}
161
162impl<T> Index<usize> for RingBuffer<T> {
163    type Output = T;
164    
165    #[inline]
166    fn index(&self, index: usize) -> &Self::Output {
167        assert!(index < self.count, "index out of bounds");
168        &self.buffer[self.tail.wrapping_add(index) & self.mask]
169    }
170}
171
172impl<T> IndexMut<usize> for RingBuffer<T> {
173    #[inline]
174    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
175        assert!(index < self.count, "index out of bounds");
176        let mask = self.mask;
177        let actual = self.tail.wrapping_add(index) & mask;
178        &mut self.buffer[actual]
179    }
180}
181
182impl<'a, T: Default + Clone> IntoIterator for &'a RingBuffer<T> {
183    type Item = &'a T;
184    type IntoIter = RingBufferIter<'a, T>;
185
186    fn into_iter(self) -> Self::IntoIter {
187        self.iter()
188    }
189}
190
191impl<T: Default + Clone> From<Vec<T>> for RingBuffer<T> {
192    fn from(vec: Vec<T>) -> Self {
193        let mut rb = RingBuffer::with_capacity(vec.len());
194        for item in vec {
195            rb.push_back(item);
196        }
197        rb
198    }
199}