1
use std::borrow::Cow;
2
use std::collections::hash_map::RandomState;
3
use std::fmt::Debug;
4
use std::hash::{BuildHasher, Hash};
5
use std::path::{Path, PathBuf};
6
use std::sync::{Mutex, OnceLock};
7

            
8
use crate::pool::{Pool, PoolKindSealed, Poolable};
9
use crate::{PoolKind, Pooled};
10

            
11
/// A pooled string that is stored in a [`GlobalPool`].
12
///
13
/// This type implements `From<String>` and `From<&str>`.
14
pub type GlobalString<S = RandomState> = Pooled<&'static GlobalPool<String, S>, S>;
15
/// A pooled path that is stored in a [`GlobalPool`].
16
///
17
/// This type implements `From<PathBuf>` and `From<&Path>`.
18
pub type GlobalPath<S = RandomState> = Pooled<&'static GlobalPool<PathBuf, S>, S>;
19
/// A pooled buffer (`Vec<u8>`) that is stored in a [`GlobalPool`].
20
///
21
/// This type implements `From<Vec<u8>>` and `From<&[u8]>`.
22
pub type GlobalBuffer<S = RandomState> = Pooled<&'static GlobalPool<Vec<u8>, S>, S>;
23

            
24
/// A global string interning pool that manages [`GlobalString`]s.
25
pub type StringPool<S = RandomState> = GlobalPool<String, S>;
26
/// A global path interning pool that manages [`GlobalPath`]s.
27
///
28
/// Each [`PathPool`] has its own storage. When comparing [`GlobalPath`]s
29
/// from separate pools, the full string comparison function must be used.
30
pub type PathPool<S = RandomState> = GlobalPool<PathBuf, S>;
31
/// A global byte buffer interning pool that manages [`GlobalBuffer`]s.
32
///
33
/// Each [`BufferPool`] has its own storage. When comparing [`GlobalBuffer`]s
34
/// from separate pools, the full string comparison function must be used.
35
pub type BufferPool<S = RandomState> = GlobalPool<Vec<u8>, S>;
36

            
37
/// A global interned pool.
38
///
39
/// This type is used to create globally allocated pools.
40
///
41
/// ```rust
42
/// use interner::global::{GlobalPool, GlobalString};
43
///
44
/// static GLOBAL_STRINGS: GlobalPool<String> = GlobalPool::new();
45
///
46
/// let interned = GLOBAL_STRINGS.get(String::from("hello"));
47
/// let second = GLOBAL_STRINGS.get("hello");
48
///
49
/// assert!(GlobalString::ptr_eq(&interned, &second));
50
/// ```
51
#[derive(Debug)]
52
pub struct GlobalPool<T, S = RandomState>(Mutex<GlobalPoolState<T, S>>)
53
where
54
    T: Poolable + Debug + Clone + Eq + PartialEq + Hash + Ord + PartialOrd + 'static,
55
    S: BuildHasher + 'static;
56

            
57
#[derive(Debug)]
58
enum GlobalPoolState<T, S>
59
where
60
    &'static GlobalPool<T, S>: PoolKind<S>,
61
    T: Poolable + Debug + Clone + Eq + PartialEq + Hash + Ord + PartialOrd + 'static,
62
    S: BuildHasher + 'static,
63
{
64
    Initializing,
65
    LazyInitialize { capacity: usize, hasher: fn() -> S },
66
    StaticInitialize { capacity: usize, hasher: S },
67
    Initialized(Pool<&'static GlobalPool<T, S>, S>),
68
}
69

            
70
impl<T> GlobalPool<T>
71
where
72
    T: Poolable + Debug + Clone + Eq + PartialEq + Hash + Ord + PartialOrd,
73
{
74
    /// Returns a new instance using [`RandomState`] for the internal hashing.
75
    #[must_use]
76
    pub const fn new() -> Self {
77
        Self::with_capacity_and_hasher_init(0, RandomState::new)
78
    }
79
}
80
impl<T, S> GlobalPool<T, S>
81
where
82
    T: Poolable + Debug + Clone + Eq + PartialEq + Hash + Ord + PartialOrd,
83
    S: BuildHasher,
84
{
85
    /// Returns a collection of the currently pooled items.
86
    #[must_use]
87
3
    pub fn pooled<C>(&'static self) -> C
88
3
    where
89
3
        C: FromIterator<Pooled<&'static Self, S>>,
90
3
    {
91
3
        self.with_active_symbols(|pool| {
92
3
            pool.active
93
3
                .iter()
94
3
                .map(|data| Pooled(data.clone()))
95
3
                .collect()
96
3
        })
97
3
    }
98
}
99
impl<T, S, S2> PartialEq<GlobalPool<T, S2>> for GlobalPool<T, S>
100
where
101
    T: Poolable + Poolable + Debug + Clone + Eq + PartialEq + Hash + Ord + PartialOrd,
102
    S: BuildHasher,
103
    S2: BuildHasher,
104
{
105
15
    fn eq(&self, other: &GlobalPool<T, S2>) -> bool {
106
15
        (self as *const Self).cast::<()>() == (other as *const GlobalPool<T, S2>).cast::<()>()
107
15
    }
108
}
109

            
110
impl<T, S> PoolKindSealed<S> for &'static GlobalPool<T, S>
111
where
112
    T: Poolable + Debug + Clone + Eq + PartialEq + Hash + Ord + PartialOrd,
113
    S: BuildHasher,
114
{
115
    type Owned = T;
116
    type Pooled = T::Boxed;
117

            
118
5125
    fn with_active_symbols<R>(&self, logic: impl FnOnce(&mut Pool<Self, S>) -> R) -> R {
119
5125
        let mut symbols = self.0.lock().expect("poisoned");
120
5125
        if !matches!(*symbols, GlobalPoolState::Initialized(_)) {
121
7
            let pool = match std::mem::replace(&mut *symbols, GlobalPoolState::Initializing) {
122
6
                GlobalPoolState::LazyInitialize { capacity, hasher } => {
123
6
                    Pool::with_capacity_and_hasher(capacity, hasher())
124
                }
125
1
                GlobalPoolState::StaticInitialize { capacity, hasher } => {
126
1
                    Pool::with_capacity_and_hasher(capacity, hasher)
127
                }
128

            
129
                _ => unreachable!("invalid state"),
130
            };
131
7
            *symbols = GlobalPoolState::Initialized(pool);
132
5119
        }
133

            
134
5126
        let GlobalPoolState::Initialized(pool) = &mut *symbols else { unreachable!("always initialized above") };
135
5126
        logic(pool)
136
5126
    }
137

            
138
    fn address_of(&self) -> *const () {
139
        std::ptr::addr_of!(*self).cast()
140
    }
141
}
142

            
143
impl<T, S> PoolKind<S> for &'static GlobalPool<T, S>
144
where
145
    T: Poolable + Debug + Clone + Eq + PartialEq + Hash + Ord + PartialOrd,
146
    S: BuildHasher,
147
{
148
}
149

            
150
impl<T, S> GlobalPool<T, S>
151
where
152
    T: Poolable + Debug + Clone + Eq + PartialEq + Hash + Ord + PartialOrd,
153
    S: BuildHasher,
154
{
155
    /// Returns a new instance using the provided hasher.
156
    pub const fn with_hasher(hasher: S) -> Self {
157
        Self::with_capacity_and_hasher(0, hasher)
158
    }
159

            
160
    /// Returns a new instance using the function to load the hasher when the
161
    /// pool is initialized on first use.
162
    pub const fn with_hasher_init(init: fn() -> S) -> Self {
163
        Self::with_capacity_and_hasher_init(0, init)
164
    }
165

            
166
    /// Returns a new instance using the provided hasher with enough capacity to
167
    /// hold the requested number of items without reallocating.
168
    pub const fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self {
169
        Self(Mutex::new(GlobalPoolState::StaticInitialize {
170
            capacity,
171
            hasher,
172
        }))
173
    }
174

            
175
    /// Returns a new instance using the function to load the hasher when the
176
    /// pool is initialized on first use. The returned instance has enough
177
    /// capacity to hold the requested number of items without reallocating.
178
    pub const fn with_capacity_and_hasher_init(capacity: usize, init: fn() -> S) -> Self {
179
        Self(Mutex::new(GlobalPoolState::LazyInitialize {
180
            capacity,
181
            hasher: init,
182
        }))
183
    }
184
}
185

            
186
impl<S> GlobalPool<String, S>
187
where
188
    S: BuildHasher,
189
{
190
    /// Returns a copy of an existing [`GlobalString`] if one is found.
191
    /// Otherwise, a new [`GlobalString`] is created and returned.
192
    ///
193
    /// While any copies of the returned [`GlobalString`] are still allocated,
194
    /// calling this function is guaranteed to return a copy of the same string.
195
4010
    pub fn get<'a, V>(&'static self, value: V) -> GlobalString<S>
196
4010
    where
197
4010
        V: Into<Cow<'a, str>>,
198
4010
    {
199
4010
        let value = value.into();
200
4010
        self.with_active_symbols(|symbols| symbols.get(value, &self))
201
4010
    }
202

            
203
    /// Returns a static pooled string, which keeps the pooled string allocated
204
    /// for the duration of the process.
205
    ///
206
    /// The string is not initialized until it is retrieved for the first time.
207
    pub const fn get_static(&'static self, value: &'static str) -> StaticPooledString<S> {
208
        StaticPooledString::new(self, value)
209
    }
210

            
211
    /// Returns a static pooled string, which keeps the pooled string allocated
212
    /// for the duration of the process. The string is initialized using the
213
    /// function provided when it is retrieved for the first time.
214
    pub const fn get_static_with(
215
        &'static self,
216
        function: fn() -> Cow<'static, str>,
217
    ) -> StaticPooledString<S> {
218
        StaticPooledString::new_fn(self, function)
219
    }
220
}
221

            
222
impl<S> GlobalPool<PathBuf, S>
223
where
224
    S: BuildHasher,
225
{
226
    /// Returns a copy of an existing [`GlobalPath`] if one is found.
227
    /// Otherwise, a new [`GlobalPath`] is created and returned.
228
    ///
229
    /// While any copies of the returned [`GlobalPath`] are still allocated,
230
    /// calling this function is guaranteed to return a copy of the same path.
231
5
    pub fn get<'a, V>(&'static self, value: V) -> GlobalPath<S>
232
5
    where
233
5
        V: Into<Cow<'a, Path>>,
234
5
    {
235
5
        let value = value.into();
236
5
        self.with_active_symbols(|symbols| symbols.get(value, &self))
237
5
    }
238

            
239
    // This function serves no purpose, currently, as there's no way to get a
240
    // static path in a const context -- Path::new() isn't const.
241
    // /// Returns a static pooled path, which keeps the pooled path allocated for
242
    // /// the duration of the process.
243
    // ///
244
    // /// The path is not initialized until it is retrieved for the first time.
245
    // pub const fn get_static(&'static self, value: &'static Path) -> StaticPooledPath<S> {
246
    //     StaticPooledPath::new(self, value)
247
    // }
248

            
249
    /// Returns a static pooled path, which keeps the pooled path allocated for
250
    /// the duration of the process. The path is initialized using the function
251
    /// provided when it is retrieved for the first time.
252
    pub const fn get_static_with(
253
        &'static self,
254
        function: fn() -> Cow<'static, Path>,
255
    ) -> StaticPooledPath<S> {
256
        StaticPooledPath::new_fn(self, function)
257
    }
258
}
259

            
260
impl<S> GlobalPool<Vec<u8>, S>
261
where
262
    S: BuildHasher,
263
{
264
    /// Returns a copy of an existing [`GlobalBuffer`] if one is found.
265
    /// Otherwise, a new [`GlobalBuffer`] is created and returned.
266
    ///
267
    /// While any copies of the returned [`GlobalBuffer`] are still allocated,
268
    /// calling this function is guaranteed to return a copy of the same byte
269
    /// buffer.
270
6
    pub fn get<'a, V>(&'static self, value: V) -> GlobalBuffer<S>
271
6
    where
272
6
        V: Into<Cow<'a, [u8]>>,
273
6
    {
274
6
        let value = value.into();
275
6
        self.with_active_symbols(|symbols| symbols.get(value, &self))
276
6
    }
277

            
278
    /// Returns a static pooled buffer, which keeps the pooled buffer allocated for
279
    /// the duration of the process.
280
    ///
281
    /// The buffer is not initialized until it is retrieved for the first time.
282
    pub const fn get_static(&'static self, value: &'static [u8]) -> StaticPooledBuffer<S> {
283
        StaticPooledBuffer::new(self, value)
284
    }
285

            
286
    /// Returns a static pooled buffer, which keeps the pooled buffer allocated
287
    /// for the duration of the process. The buffer is initialized using the
288
    /// function provided when it is retrieved for the first time.
289
    pub const fn get_static_with(
290
        &'static self,
291
        function: fn() -> Cow<'static, [u8]>,
292
    ) -> StaticPooledBuffer<S> {
293
        StaticPooledBuffer::new_fn(self, function)
294
    }
295
}
296

            
297
/// A lazily-initialized [`GlobalString`] that stays allocated for the duration
298
/// of the process.
299
#[derive(Debug)]
300
pub struct StaticPooledString<S = RandomState>
301
where
302
    S: BuildHasher + 'static,
303
{
304
    init: StaticStringInit<S>,
305
    cell: OnceLock<GlobalString<S>>,
306
}
307

            
308
/// A lazily-initialized [`GlobalBuffer`] that stays allocated for the duration
309
/// of the process.
310
#[derive(Debug)]
311
pub struct StaticPooledBuffer<S = RandomState>
312
where
313
    S: BuildHasher + 'static,
314
{
315
    init: StaticBufferInit<S>,
316
    cell: OnceLock<GlobalBuffer<S>>,
317
}
318

            
319
/// A lazily-initialized [`GlobalPath`] that stays allocated for the duration
320
/// of the process.
321
#[derive(Debug)]
322
pub struct StaticPooledPath<S = RandomState>
323
where
324
    S: BuildHasher + 'static,
325
{
326
    init: StaticPathInit<S>,
327
    cell: OnceLock<GlobalPath<S>>,
328
}
329

            
330
macro_rules! impl_static_pooled {
331
    ($name:ident, $pooled:ident, $statename:ident, $owned:ty, $borrowed:ty) => {
332
        impl<S> $name<S>
333
        where
334
            S: BuildHasher + 'static,
335
        {
336
            #[allow(dead_code)] // This function isn't called for StaticPooledPath, because there's no way to get a static Path.
337
            const fn new(pool: &'static GlobalPool<$owned, S>, value: &'static $borrowed) -> Self {
338
                Self {
339
                    init: $statename::Static(pool, value),
340
                    cell: OnceLock::new(),
341
                }
342
            }
343

            
344
            const fn new_fn(
345
                pool: &'static GlobalPool<$owned, S>,
346
                init: fn() -> Cow<'static, $borrowed>,
347
            ) -> Self {
348
                Self {
349
                    init: $statename::Fn(pool, init),
350
                    cell: OnceLock::new(),
351
                }
352
            }
353

            
354
            /// Returns a reference-counted clone of the contained resource.
355
            ///
356
            /// If this is the first time the contents are accessed, the value
357
            /// will be initialized from the pool on the first access. This
358
            /// requires synchronization and can block the current thraed very
359
            /// briefly.
360
            ///
361
            /// All subsequent accesses will be non-blocking.
362
26
            pub fn get(&self) -> &$pooled<S> {
363
26
                self.cell.get_or_init(|| match self.init {
364
                    $statename::Static(pool, value) => pool.get(value).clone(),
365
                    $statename::Fn(pool, init) => pool.get(init()).clone(),
366
26
                })
367
26
            }
368
        }
369

            
370
        #[derive(Debug, Clone, Copy)]
371
        #[allow(dead_code)] // Path can't use the Static variant.
372
        enum $statename<S>
373
        where
374
            S: BuildHasher + 'static,
375
        {
376
            Static(&'static GlobalPool<$owned, S>, &'static $borrowed),
377
            Fn(
378
                &'static GlobalPool<$owned, S>,
379
                fn() -> Cow<'static, $borrowed>,
380
            ),
381
        }
382

            
383
        impl<S> std::ops::Deref for $name<S>
384
        where
385
            S: BuildHasher,
386
        {
387
            type Target = $pooled<S>;
388

            
389
8
            fn deref(&self) -> &Self::Target {
390
8
                self.get()
391
8
            }
392
        }
393

            
394
        impl<S, S2> PartialEq<Pooled<&'static GlobalPool<$owned, S2>, S2>> for $name<S>
395
        where
396
            S: BuildHasher,
397
            S2: BuildHasher,
398
        {
399
5
            fn eq(&self, other: &Pooled<&'static GlobalPool<$owned, S2>, S2>) -> bool {
400
5
                self.get() == other
401
5
            }
402
        }
403

            
404
        impl<S, S2> PartialEq<$name<S>> for Pooled<&'static GlobalPool<$owned, S2>, S2>
405
        where
406
            S: BuildHasher,
407
            S2: BuildHasher,
408
        {
409
            fn eq(&self, other: &$name<S>) -> bool {
410
                self == other.get()
411
            }
412
        }
413
    };
414
}
415

            
416
3
impl_static_pooled!(
417
3
    StaticPooledString,
418
3
    GlobalString,
419
3
    StaticStringInit,
420
3
    String,
421
3
    str
422
3
);
423
3
impl_static_pooled!(
424
3
    StaticPooledBuffer,
425
3
    GlobalBuffer,
426
3
    StaticBufferInit,
427
3
    Vec<u8>,
428
3
    [u8]
429
3
);
430
2
impl_static_pooled!(StaticPooledPath, GlobalPath, StaticPathInit, PathBuf, Path);