ax_codec_core/
pool.rs

1use core::cell::RefCell;
2
3#[cfg(feature = "alloc")]
4use alloc::vec::Vec;
5
6#[cfg(feature = "std")]
7thread_local! {
8    static POOL: RefCell<Vec<Vec<u8>>> = RefCell::new(Vec::with_capacity(8));
9}
10
11#[cfg(feature = "std")]
12const POOL_MAX_BUFS: usize = 8;
13
14#[cfg(feature = "std")]
15const POOL_MAX_SIZE: usize = 64 * 1024;
16
17#[cfg(feature = "std")]
18pub fn take(min_capacity: usize) -> Vec<u8> {
19    POOL.with(|pool| {
20        if let Ok(mut pool) = pool.try_borrow_mut() {
21            if let Some(pos) = pool.iter().position(|buf| buf.capacity() >= min_capacity) {
22                let mut buf = pool.swap_remove(pos);
23                buf.clear();
24                buf
25            } else {
26                Vec::with_capacity(min_capacity)
27            }
28        } else {
29            Vec::with_capacity(min_capacity)
30        }
31    })
32}
33
34#[cfg(feature = "std")]
35pub fn recycle(mut buf: Vec<u8>) {
36    if buf.capacity() > POOL_MAX_SIZE {
37        return;
38    }
39    buf.clear();
40    POOL.with(|pool| {
41        if let Ok(mut pool) = pool.try_borrow_mut() {
42            if pool.len() < POOL_MAX_BUFS {
43                pool.push(buf);
44            }
45        }
46    });
47}
48
49#[cfg(feature = "std")]
50pub fn clear() {
51    POOL.with(|pool| {
52        pool.borrow_mut().clear();
53    });
54}
55
56#[cfg(feature = "std")]
57pub struct PooledVec {
58    buf: Option<Vec<u8>>,
59}
60
61#[cfg(feature = "std")]
62impl PooledVec {
63    pub fn with_capacity(min_capacity: usize) -> Self {
64        Self {
65            buf: Some(take(min_capacity)),
66        }
67    }
68
69    pub fn into_vec(mut self) -> Vec<u8> {
70        self.buf.take().unwrap()
71    }
72
73    pub fn take(&mut self) -> Vec<u8> {
74        self.buf.take().unwrap()
75    }
76
77    pub fn put(&mut self, buf: Vec<u8>) {
78        self.buf = Some(buf);
79    }
80}
81
82#[cfg(feature = "std")]
83impl core::ops::Deref for PooledVec {
84    type Target = Vec<u8>;
85
86    fn deref(&self) -> &Self::Target {
87        self.buf.as_ref().unwrap()
88    }
89}
90
91#[cfg(feature = "std")]
92impl core::ops::DerefMut for PooledVec {
93    fn deref_mut(&mut self) -> &mut Self::Target {
94        self.buf.as_mut().unwrap()
95    }
96}
97
98#[cfg(feature = "std")]
99impl Drop for PooledVec {
100    fn drop(&mut self) {
101        if let Some(buf) = self.buf.take() {
102            recycle(buf);
103        }
104    }
105}
106
107#[cfg(not(feature = "std"))]
108pub fn take(min_capacity: usize) -> Vec<u8> {
109    Vec::with_capacity(min_capacity)
110}
111
112#[cfg(not(feature = "std"))]
113pub fn recycle(_buf: Vec<u8>) {}
114
115#[cfg(not(feature = "std"))]
116pub fn clear() {}
117
118#[cfg(not(feature = "std"))]
119pub struct PooledVec {
120    buf: Vec<u8>,
121}
122
123#[cfg(not(feature = "std"))]
124impl PooledVec {
125    pub fn with_capacity(min_capacity: usize) -> Self {
126        Self {
127            buf: Vec::with_capacity(min_capacity),
128        }
129    }
130
131    pub fn into_vec(self) -> Vec<u8> {
132        self.buf
133    }
134
135    pub fn take(&mut self) -> Vec<u8> {
136        core::mem::replace(&mut self.buf, Vec::new())
137    }
138
139    pub fn put(&mut self, buf: Vec<u8>) {
140        self.buf = buf;
141    }
142}
143
144#[cfg(not(feature = "std"))]
145impl core::ops::Deref for PooledVec {
146    type Target = Vec<u8>;
147
148    fn deref(&self) -> &Self::Target {
149        &self.buf
150    }
151}
152
153#[cfg(not(feature = "std"))]
154impl core::ops::DerefMut for PooledVec {
155    fn deref_mut(&mut self) -> &mut Self::Target {
156        &mut self.buf
157    }
158}
159
160#[cfg(all(test, feature = "std"))]
161mod tests {
162    use super::*;
163
164    #[test]
165    fn pooled_vec_recycles() {
166        clear();
167
168        let addr = {
169            let v = PooledVec::with_capacity(1024);
170            v.buf.as_ref().unwrap().as_ptr() as usize
171        };
172
173        let v2 = PooledVec::with_capacity(1024);
174        let addr2 = v2.buf.as_ref().unwrap().as_ptr() as usize;
175
176        assert_eq!(addr, addr2, "pooled buffer was not recycled");
177    }
178
179    #[test]
180    fn oversized_not_pooled() {
181        clear();
182
183        let _ = PooledVec::with_capacity(POOL_MAX_SIZE + 1);
184
185        let v = PooledVec::with_capacity(1024);
186        assert!(v.capacity() >= 1024);
187    }
188}
ax_codec_core/pool.rs

ax_codec_core/
pool.rs
