Skip to main content

lavende_core/audio/
buffer.rs

1pub mod pool {
2    use crate::audio::constants::{MAX_BUCKET_ENTRIES, MAX_POOL_BYTES, POOL_IDLE_CLEAR_SECS};
3    use parking_lot::Mutex;
4    use std::{
5        collections::HashMap,
6        sync::OnceLock,
7        time::{Duration, Instant},
8    };
9    const CLEANUP_INTERVAL: Duration = Duration::from_secs(30);
10    struct PoolInner {
11        buckets: HashMap<usize, Vec<Vec<u8>>>,
12        total_bytes: usize,
13        last_activity: Instant,
14        last_cleanup: Instant,
15    }
16    impl PoolInner {
17        fn new() -> Self {
18            let now = Instant::now();
19            Self {
20                buckets: HashMap::new(),
21                total_bytes: 0,
22                last_activity: now,
23                last_cleanup: now,
24            }
25        }
26        fn aligned_size(size: usize) -> usize {
27            let aligned = size.max(1024).next_power_of_two();
28            aligned.min(1024 * 1024)
29        }
30        fn needs_cleanup(&self) -> bool {
31            self.total_bytes > 0 && self.last_cleanup.elapsed() >= CLEANUP_INTERVAL
32        }
33        fn acquire(&mut self, size: usize) -> Vec<u8> {
34            self.last_activity = Instant::now();
35            let aligned = Self::aligned_size(size);
36            if let Some(buf) = self
37                .buckets
38                .get_mut(&aligned)
39                .and_then(|bucket| bucket.pop())
40            {
41                self.total_bytes -= aligned;
42                return buf;
43            }
44            Vec::with_capacity(aligned)
45        }
46        fn release(&mut self, mut buf: Vec<u8>) {
47            self.last_activity = Instant::now();
48            let size = buf.capacity();
49            if !(1024..=10 * 1024 * 1024).contains(&size) {
50                return;
51            }
52            if self.total_bytes + size > MAX_POOL_BYTES {
53                return;
54            }
55            let bucket = self.buckets.entry(size).or_default();
56            if bucket.len() >= MAX_BUCKET_ENTRIES {
57                return;
58            }
59            buf.clear();
60            self.total_bytes += size;
61            bucket.push(buf);
62        }
63        fn cleanup(&mut self) {
64            self.last_cleanup = Instant::now();
65            let is_idle = self.last_activity.elapsed() >= Duration::from_secs(POOL_IDLE_CLEAR_SECS);
66            let is_over_limit = self.total_bytes > MAX_POOL_BYTES;
67            if is_idle || is_over_limit {
68                self.buckets.clear();
69                self.total_bytes = 0;
70            }
71        }
72    }
73    pub struct BufferPool {
74        inner: Mutex<PoolInner>,
75    }
76    impl BufferPool {
77        fn new() -> Self {
78            Self {
79                inner: Mutex::new(PoolInner::new()),
80            }
81        }
82        pub fn acquire(&self, size: usize) -> Vec<u8> {
83            let mut g = self.inner.lock();
84            if g.needs_cleanup() {
85                g.cleanup();
86            }
87            g.acquire(size)
88        }
89        pub fn release(&self, buf: Vec<u8>) {
90            self.inner.lock().release(buf);
91        }
92        pub fn stats(&self) -> PoolStats {
93            let g = self.inner.lock();
94            PoolStats {
95                total_bytes: g.total_bytes,
96                buckets: g.buckets.len(),
97                entries: g.buckets.values().map(|b| b.len()).sum(),
98            }
99        }
100    }
101    #[derive(Debug, Clone)]
102    pub struct PoolStats {
103        pub total_bytes: usize,
104        pub buckets: usize,
105        pub entries: usize,
106    }
107    static GLOBAL_BYTE_POOL: OnceLock<BufferPool> = OnceLock::new();
108    pub fn get_byte_pool() -> &'static BufferPool {
109        GLOBAL_BYTE_POOL.get_or_init(BufferPool::new)
110    }
111}
112pub mod ring {
113    use crate::audio::buffer::pool::get_byte_pool;
114    pub struct RingBuffer {
115        buf: Vec<u8>,
116        size: usize,
117        write_offset: usize,
118        read_offset: usize,
119        length: usize,
120    }
121    impl RingBuffer {
122        pub fn new(size: usize) -> Self {
123            let mut buf = get_byte_pool().acquire(size);
124            buf.resize(size, 0);
125            Self {
126                buf,
127                size,
128                write_offset: 0,
129                read_offset: 0,
130                length: 0,
131            }
132        }
133        pub fn len(&self) -> usize {
134            self.length
135        }
136        pub fn is_empty(&self) -> bool {
137            self.length == 0
138        }
139        pub fn remaining(&self) -> usize {
140            self.size - self.length
141        }
142        pub fn write(&mut self, chunk: &[u8]) {
143            let chunk = if chunk.len() > self.size {
144                &chunk[chunk.len() - self.size..]
145            } else {
146                chunk
147            };
148            let to_write = chunk.len();
149            let available_at_end = self.size - self.write_offset;
150            if to_write <= available_at_end {
151                self.buf[self.write_offset..self.write_offset + to_write].copy_from_slice(chunk);
152            } else {
153                self.buf[self.write_offset..].copy_from_slice(&chunk[..available_at_end]);
154                self.buf[..to_write - available_at_end].copy_from_slice(&chunk[available_at_end..]);
155            }
156            let new_len = self.length + to_write;
157            if new_len > self.size {
158                let overwritten = new_len - self.size;
159                self.read_offset = (self.read_offset + overwritten) % self.size;
160                self.length = self.size;
161            } else {
162                self.length = new_len;
163            }
164            self.write_offset = (self.write_offset + to_write) % self.size;
165        }
166        pub fn read(&mut self, n: usize) -> Option<Vec<u8>> {
167            let to_read = self.peek(n)?;
168            self.read_offset = (self.read_offset + to_read.len()) % self.size;
169            self.length -= to_read.len();
170            Some(to_read)
171        }
172        pub fn peek(&self, n: usize) -> Option<Vec<u8>> {
173            let to_read = n.min(self.length);
174            if to_read == 0 {
175                return None;
176            }
177            let mut out = get_byte_pool().acquire(to_read);
178            out.resize(to_read, 0);
179            self.copy_to(&mut out);
180            Some(out)
181        }
182        pub fn peek_slice<F, R>(&self, n: usize, f: F) -> Option<R>
183        where
184            F: FnOnce(&[u8], &[u8]) -> R,
185        {
186            let to_read = n.min(self.length);
187            if to_read == 0 {
188                return None;
189            }
190            let available_at_end = self.size - self.read_offset;
191            let result = if to_read <= available_at_end {
192                f(&self.buf[self.read_offset..self.read_offset + to_read], &[])
193            } else {
194                f(
195                    &self.buf[self.read_offset..],
196                    &self.buf[..to_read - available_at_end],
197                )
198            };
199            Some(result)
200        }
201        fn copy_to(&self, out: &mut [u8]) {
202            let to_copy = out.len();
203            let available_at_end = self.size - self.read_offset;
204            if to_copy <= available_at_end {
205                out.copy_from_slice(&self.buf[self.read_offset..self.read_offset + to_copy]);
206            } else {
207                out[..available_at_end].copy_from_slice(&self.buf[self.read_offset..]);
208                out[available_at_end..].copy_from_slice(&self.buf[..to_copy - available_at_end]);
209            }
210        }
211        pub fn skip(&mut self, n: usize) -> usize {
212            let to_skip = n.min(self.length);
213            self.read_offset = (self.read_offset + to_skip) % self.size;
214            self.length -= to_skip;
215            to_skip
216        }
217        pub fn clear(&mut self) {
218            self.write_offset = 0;
219            self.read_offset = 0;
220            self.length = 0;
221        }
222    }
223    impl Drop for RingBuffer {
224        fn drop(&mut self) {
225            if !self.buf.is_empty() {
226                let buf = std::mem::take(&mut self.buf);
227                get_byte_pool().release(buf);
228            }
229        }
230    }
231    #[cfg(test)]
232    mod tests {
233        use super::*;
234        #[test]
235        fn test_ring_buffer_basic() {
236            let mut rb = RingBuffer::new(10);
237            assert_eq!(rb.remaining(), 10);
238            rb.write(b"hello");
239            assert_eq!(rb.len(), 5);
240            assert_eq!(rb.remaining(), 5);
241            let data = rb.read(3).unwrap();
242            assert_eq!(data, b"hel");
243            assert_eq!(rb.len(), 2);
244            let data = rb.peek(2).unwrap();
245            assert_eq!(data, b"lo");
246            assert_eq!(rb.len(), 2);
247            let data = rb.read(5).unwrap();
248            assert_eq!(data, b"lo");
249            assert_eq!(rb.len(), 0);
250        }
251        #[test]
252        fn test_ring_buffer_wrap_around() {
253            let mut rb = RingBuffer::new(10);
254            rb.write(b"0123456789");
255            rb.skip(5);
256            rb.write(b"abcde");
257            let data = rb.read(10).unwrap();
258            assert_eq!(data, b"56789abcde");
259        }
260        #[test]
261        fn test_ring_buffer_overwrite() {
262            let mut rb = RingBuffer::new(5);
263            rb.write(b"12345");
264            rb.write(b"67");
265            let data = rb.read(5).unwrap();
266            assert_eq!(data, b"34567");
267        }
268        #[test]
269        fn test_ring_buffer_large_write() {
270            let mut rb = RingBuffer::new(5);
271            rb.write(b"12345678");
272            let data = rb.read(5).unwrap();
273            assert_eq!(data, b"45678");
274        }
275        #[test]
276        fn test_peek_slice_zero_copy() {
277            let mut rb = RingBuffer::new(10);
278            rb.write(b"hello");
279            let result = rb.peek_slice(5, |a, b| {
280                let mut v = a.to_vec();
281                v.extend_from_slice(b);
282                v
283            });
284            assert_eq!(result.unwrap(), b"hello");
285        }
286    }
287}
288pub use pool::{BufferPool, get_byte_pool};
289pub use ring::RingBuffer;
290pub type PooledBuffer = Vec<i16>;
291pub fn cast_to_bytes(v: PooledBuffer) -> Vec<u8> {
292    let mut v = std::mem::ManuallyDrop::new(v);
293    unsafe { Vec::from_raw_parts(v.as_mut_ptr() as *mut u8, v.len() * 2, v.capacity() * 2) }
294}
295pub fn cast_from_bytes(v: Vec<u8>) -> PooledBuffer {
296    debug_assert_eq!(v.len() % 2, 0, "byte buffer length must be even");
297    debug_assert_eq!(v.capacity() % 2, 0, "byte buffer capacity must be even");
298    let mut v = std::mem::ManuallyDrop::new(v);
299    unsafe { Vec::from_raw_parts(v.as_mut_ptr() as *mut i16, v.len() / 2, v.capacity() / 2) }
300}
301#[inline]
302pub fn as_byte_slice(v: &[i16]) -> &[u8] {
303    unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2) }
304}
305#[inline]
306pub fn as_i16_slice(v: &[u8]) -> &[i16] {
307    debug_assert_eq!(
308        v.as_ptr() as usize % std::mem::align_of::<i16>(),
309        0,
310        "byte slice must be 2-byte aligned for i16 reinterpretation"
311    );
312    debug_assert_eq!(v.len() % 2, 0, "byte slice length must be even");
313    unsafe { std::slice::from_raw_parts(v.as_ptr() as *const i16, v.len() / 2) }
314}
315#[inline]
316pub fn release_buffer(v: PooledBuffer) {
317    get_byte_pool().release(cast_to_bytes(v));
318}
319#[inline]
320pub fn acquire_buffer(capacity: usize) -> PooledBuffer {
321    cast_from_bytes(get_byte_pool().acquire(capacity * 2))
322}