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