1use crossbeam_channel::{bounded, Receiver, Sender};
2use std::ops::{Deref, DerefMut};
3
4pub struct OwnedPool<T> {
5 free: Vec<T>,
6 num_borrowed: usize,
7 receiver: Receiver<T>,
8 sender: Sender<T>,
9 reset_fn: fn(&mut T),
10 init_fn: fn() -> T,
11}
12
13impl<T> OwnedPool<T> {
14 pub fn with_capacity(
15 pool_size: usize,
16 init_fn: fn() -> T,
17 reset_fn: fn(&mut T),
18 ) -> Self {
19 let (sender, receiver) = bounded(pool_size);
20 Self {
21 free: Vec::with_capacity(pool_size),
22 num_borrowed: 0,
23 init_fn,
24 reset_fn,
25 sender,
26 receiver,
27 }
28 }
29
30 pub fn borrow(&mut self) -> Pooled<T> {
31 if let Some(mut pooled_val) = self.free.pop() {
32 self.num_borrowed += 1;
33 (self.reset_fn)(&mut pooled_val);
34 Pooled {
35 inner: Some(pooled_val),
36 sender: self.sender.clone(),
37 }
38 } else if self.num_borrowed < self.free.capacity() {
39 self.num_borrowed += 1;
40 Pooled {
41 inner: Some((self.init_fn)()),
42 sender: self.sender.clone(),
43 }
44 } else {
45 panic!("Cannot borrow more than `pool_size` entries.");
46 }
47 }
48
49 pub fn try_recv(&mut self) {
50 for pooled_val in self.receiver.try_iter() {
51 self.free.push(pooled_val);
52 self.num_borrowed -= 1;
53 }
54 }
55}
56
57pub struct Pooled<T> {
58 inner: Option<T>,
59 sender: Sender<T>,
60}
61
62impl<T> Pooled<T> {
63 pub fn as_ref(&self) -> &T {
64 self.inner.as_ref().unwrap()
65 }
66
67 pub fn as_mut(&mut self) -> &mut T {
68 self.inner.as_mut().unwrap()
69 }
70}
71
72impl<T> Deref for Pooled<T> {
73 type Target = T;
74
75 fn deref(&self) -> &Self::Target {
76 self.inner.as_ref().unwrap()
77 }
78}
79
80impl<T> DerefMut for Pooled<T> {
81 fn deref_mut(&mut self) -> &mut Self::Target {
82 self.inner.as_mut().unwrap()
83 }
84}
85
86impl<T> Drop for Pooled<T> {
87 fn drop(&mut self) {
88 let inner = std::mem::take(&mut self.inner);
89 let _ = self.sender.send(inner.unwrap());
90 }
91}