1
use std::borrow::{Borrow, Cow};
2
use std::collections::hash_map::RandomState;
3
use std::collections::HashSet;
4
use std::fmt::Debug;
5
use std::hash::{BuildHasher, Hash};
6
use std::marker::PhantomData;
7
use std::path::{Path, PathBuf};
8
use std::sync::atomic::AtomicBool;
9
use std::sync::{atomic, Arc};
10

            
11
use crate::{PoolKind, Pooled};
12

            
13
pub trait PoolKindSealed<Hasher> {
14
    type Owned: Poolable<Boxed = Self::Pooled> + Debug + Clone + Eq + Hash + Ord;
15
    type Pooled: Debug + Clone + Eq + Hash + Ord;
16

            
17
    fn with_active_symbols<T>(&self, logic: impl FnOnce(&mut Pool<Self, Hasher>) -> T) -> T;
18
    fn address_of(&self) -> *const ();
19
}
20

            
21
pub trait Poolable {
22
    type Boxed: Debug + Clone + Eq + Hash + Ord;
23

            
24
    fn boxed(self) -> Self::Boxed;
25
}
26

            
27
impl Poolable for String {
28
    type Boxed = Box<str>;
29

            
30
1108
    fn boxed(self) -> Self::Boxed {
31
1108
        self.into_boxed_str()
32
1108
    }
33
}
34

            
35
impl Poolable for PathBuf {
36
    type Boxed = Box<Path>;
37

            
38
8
    fn boxed(self) -> Self::Boxed {
39
8
        self.into_boxed_path()
40
8
    }
41
}
42

            
43
impl Poolable for Vec<u8> {
44
    type Boxed = Box<[u8]>;
45

            
46
17
    fn boxed(self) -> Self::Boxed {
47
17
        self.into_boxed_slice()
48
17
    }
49
}
50

            
51
#[derive(Debug)]
52
pub struct SharedData<P, S>(pub Arc<Data<P, S>>)
53
where
54
    P: PoolKind<S>,
55
    S: BuildHasher;
56

            
57
impl<P, S> Clone for SharedData<P, S>
58
where
59
    P: PoolKind<S>,
60
    S: BuildHasher,
61
{
62
5174
    fn clone(&self) -> Self {
63
5174
        Self(self.0.clone())
64
5174
    }
65
}
66

            
67
impl<P, S> Hash for SharedData<P, S>
68
where
69
    P: PoolKind<S>,
70
    S: BuildHasher,
71
{
72
2230
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
73
2230
        self.0.value.hash(state);
74
2230
    }
75
}
76

            
77
impl<P, S> Eq for SharedData<P, S>
78
where
79
    P: PoolKind<S>,
80
    S: BuildHasher,
81
{
82
}
83

            
84
impl<P, S> PartialEq for SharedData<P, S>
85
where
86
    P: PoolKind<S>,
87
    S: BuildHasher,
88
{
89
1111
    fn eq(&self, other: &Self) -> bool {
90
1111
        self.0.index == other.0.index
91
1111
    }
92
}
93

            
94
impl<P, S> Borrow<str> for SharedData<P, S>
95
where
96
    P: PoolKind<S, Pooled = Box<str>>,
97
    S: BuildHasher,
98
{
99
2920
    fn borrow(&self) -> &str {
100
2920
        &self.0.value
101
2920
    }
102
}
103

            
104
impl<P, S> Borrow<Path> for SharedData<P, S>
105
where
106
    P: PoolKind<S, Pooled = Box<Path>>,
107
    S: BuildHasher,
108
{
109
3
    fn borrow(&self) -> &Path {
110
3
        &self.0.value
111
3
    }
112
}
113

            
114
impl<P, S> Borrow<[u8]> for SharedData<P, S>
115
where
116
    P: PoolKind<S, Pooled = Box<[u8]>>,
117
    S: BuildHasher,
118
{
119
3
    fn borrow(&self) -> &[u8] {
120
3
        &self.0.value
121
3
    }
122
}
123

            
124
impl<P, S> Drop for SharedData<P, S>
125
where
126
    P: PoolKind<S>,
127
    S: BuildHasher,
128
{
129
6269
    fn drop(&mut self) {
130
6269
        // The main Symbols structure holds two strong references to the same
131
6269
        // Arc we hold. Thus, if we reach 3 strong count (our ref included), we
132
6269
        // need to remove the symbol so it can be freeed.
133
6269
        //
134
6269
        // We can use any form of atomics here because if the strong count is 3,
135
6269
        // we can be guaranteed the only thread able to free our data is this
136
6269
        // thread.
137
6269
        if Arc::strong_count(&self.0) == 3
138
2223
            && self
139
2223
                .0
140
2223
                .freeing
141
2223
                .compare_exchange(
142
2223
                    false,
143
2223
                    true,
144
2223
                    atomic::Ordering::Relaxed,
145
2223
                    atomic::Ordering::Relaxed,
146
2223
                )
147
2223
                .is_ok()
148
1112
        {
149
1112
            self.0.pool.with_active_symbols(|symbols| {
150
1112
                // Check that the strong count hasn't changed. If it has, we
151
1112
                // need to allow the symbol to stay alive.
152
1112
                if Arc::strong_count(&self.0) > 3 {
153
1
                    self.0.freeing.store(false, atomic::Ordering::Relaxed);
154
1111
                } else {
155
1111
                    symbols.active.remove(self);
156
1111
                    symbols.slots[self.0.index] = None;
157
1111
                    symbols.free_slots.push(self.0.index);
158
1111
                }
159
1112
            });
160
5157
        }
161
6269
    }
162
}
163

            
164
#[derive(Debug)]
165
pub struct Data<P, S>
166
where
167
    P: PoolKind<S>,
168
{
169
    pub index: usize,
170
    pub value: P::Pooled,
171
    pub freeing: AtomicBool,
172
    pub pool: P,
173
    _hasher: PhantomData<S>,
174
}
175

            
176
#[derive(Debug)]
177
pub struct Pool<P, S>
178
where
179
    P: PoolKind<S>,
180
    S: BuildHasher,
181
{
182
    pub active: HashSet<SharedData<P, S>, S>,
183
    pub slots: Vec<Option<Pooled<P, S>>>,
184
    pub free_slots: Vec<usize>,
185
}
186

            
187
impl<P, S> Pool<P, S>
188
where
189
    P: PoolKind<S>,
190
    S: BuildHasher,
191
{
192
8
    pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self {
193
8
        Self {
194
8
            active: HashSet::with_capacity_and_hasher(capacity, hasher),
195
8
            slots: Vec::with_capacity(capacity),
196
8
            free_slots: Vec::new(),
197
8
        }
198
8
    }
199

            
200
    pub fn get<K>(&mut self, pooled: Cow<'_, K>, pool: &P) -> Pooled<P, S>
201
    where
202
        K: ToOwned<Owned = P::Owned> + Hash + Eq + ?Sized,
203
        P::Owned: Borrow<K> + PartialEq<K>,
204
        SharedData<P, S>: Borrow<K>,
205
    {
206
4044
        if let Some(symbol) = self.active.get(pooled.as_ref()).cloned() {
207
2925
            Pooled(symbol)
208
        } else {
209
1119
            let value = pooled.into_owned();
210

            
211
1119
            let index = if let Some(free_slot) = self.free_slots.pop() {
212
1095
                free_slot
213
            } else {
214
24
                let slot_id = self.slots.len();
215
24
                self.slots.push(None);
216
24
                slot_id
217
            };
218

            
219
1119
            let symbol = Pooled(SharedData(Arc::new(Data {
220
1119
                index,
221
1119
                value: value.boxed(),
222
1119
                freeing: AtomicBool::new(false),
223
1119
                pool: pool.clone(),
224
1119
                _hasher: PhantomData,
225
1119
            })));
226
1119
            self.active.insert(symbol.0.clone());
227
1119
            self.slots[index] = Some(symbol.clone());
228
1119
            symbol
229
        }
230
4044
    }
231
}
232

            
233
impl<P> Default for Pool<P, RandomState>
234
where
235
    P: PoolKind<RandomState>,
236
{
237
10
    fn default() -> Self {
238
10
        Self {
239
10
            active: HashSet::with_hasher(RandomState::default()),
240
10
            slots: Vec::new(),
241
10
            free_slots: Vec::new(),
242
10
        }
243
10
    }
244
}