intern/
symbol.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
//! Attempt at flexible symbol interning, allowing to intern and free strings at runtime while also
//! supporting compile time declaration of symbols that will never be freed.

use std::{
    borrow::Borrow,
    fmt,
    hash::{BuildHasherDefault, Hash, Hasher},
    mem::{self, ManuallyDrop},
    ptr::NonNull,
    sync::OnceLock,
};

use dashmap::{DashMap, SharedValue};
use hashbrown::{hash_map::RawEntryMut, HashMap};
use rustc_hash::FxHasher;
use sptr::Strict;
use triomphe::Arc;

pub mod symbols;

// some asserts for layout compatibility
const _: () = assert!(std::mem::size_of::<Box<str>>() == std::mem::size_of::<&str>());
const _: () = assert!(std::mem::align_of::<Box<str>>() == std::mem::align_of::<&str>());

const _: () = assert!(std::mem::size_of::<Arc<Box<str>>>() == std::mem::size_of::<&&str>());
const _: () = assert!(std::mem::align_of::<Arc<Box<str>>>() == std::mem::align_of::<&&str>());

const _: () =
    assert!(std::mem::size_of::<*const *const str>() == std::mem::size_of::<TaggedArcPtr>());
const _: () =
    assert!(std::mem::align_of::<*const *const str>() == std::mem::align_of::<TaggedArcPtr>());

const _: () = assert!(std::mem::size_of::<Arc<Box<str>>>() == std::mem::size_of::<TaggedArcPtr>());
const _: () =
    assert!(std::mem::align_of::<Arc<Box<str>>>() == std::mem::align_of::<TaggedArcPtr>());

/// A pointer that points to a pointer to a `str`, it may be backed as a `&'static &'static str` or
/// `Arc<Box<str>>` but its size is that of a thin pointer. The active variant is encoded as a tag
/// in the LSB of the alignment niche.
// Note, Ideally this would encode a `ThinArc<str>` and `ThinRef<str>`/`ThinConstPtr<str>` instead of the double indirection.
#[derive(PartialEq, Eq, Hash, Copy, Clone, Debug)]
struct TaggedArcPtr {
    packed: NonNull<*const str>,
}

unsafe impl Send for TaggedArcPtr {}
unsafe impl Sync for TaggedArcPtr {}

impl TaggedArcPtr {
    const BOOL_BITS: usize = true as usize;

    const fn non_arc(r: &'static &'static str) -> Self {
        assert!(
            mem::align_of::<&'static &'static str>().trailing_zeros() as usize > Self::BOOL_BITS
        );
        // SAFETY: The pointer is non-null as it is derived from a reference
        // Ideally we would call out to `pack_arc` but for a `false` tag, unfortunately the
        // packing stuff requires reading out the pointer to an integer which is not supported
        // in const contexts, so here we make use of the fact that for the non-arc version the
        // tag is false (0) and thus does not need touching the actual pointer value.ext)

        let packed =
            unsafe { NonNull::new_unchecked((r as *const &str).cast::<*const str>().cast_mut()) };
        Self { packed }
    }

    fn arc(arc: Arc<Box<str>>) -> Self {
        assert!(
            mem::align_of::<&'static &'static str>().trailing_zeros() as usize > Self::BOOL_BITS
        );
        Self {
            packed: Self::pack_arc(
                // Safety: `Arc::into_raw` always returns a non null pointer
                unsafe { NonNull::new_unchecked(Arc::into_raw(arc).cast_mut().cast()) },
            ),
        }
    }

    /// Retrieves the tag.
    ///
    /// # Safety
    ///
    /// You can only drop the `Arc` if the instance is dropped.
    #[inline]
    pub(crate) unsafe fn try_as_arc_owned(self) -> Option<ManuallyDrop<Arc<Box<str>>>> {
        // Unpack the tag from the alignment niche
        let tag = Strict::addr(self.packed.as_ptr()) & Self::BOOL_BITS;
        if tag != 0 {
            // Safety: We checked that the tag is non-zero -> true, so we are pointing to the data offset of an `Arc`
            Some(ManuallyDrop::new(unsafe {
                Arc::from_raw(self.pointer().as_ptr().cast::<Box<str>>())
            }))
        } else {
            None
        }
    }

    #[inline]
    fn pack_arc(ptr: NonNull<*const str>) -> NonNull<*const str> {
        let packed_tag = true as usize;

        // can't use this strict provenance stuff here due to trait methods not being const
        // unsafe {
        //     // Safety: The pointer is derived from a non-null
        //     NonNull::new_unchecked(Strict::map_addr(ptr.as_ptr(), |addr| {
        //         // Safety:
        //         // - The pointer is `NonNull` => it's address is `NonZero<usize>`
        //         // - `P::BITS` least significant bits are always zero (`Pointer` contract)
        //         // - `T::BITS <= P::BITS` (from `Self::ASSERTION`)
        //         //
        //         // Thus `addr >> T::BITS` is guaranteed to be non-zero.
        //         //
        //         // `{non_zero} | packed_tag` can't make the value zero.

        //         (addr >> Self::BOOL_BITS) | packed_tag
        //     }))
        // }
        // so what follows is roughly what the above looks like but inlined

        let self_addr = ptr.as_ptr() as *const *const str as usize;
        let addr = self_addr | packed_tag;
        let dest_addr = addr as isize;
        let offset = dest_addr.wrapping_sub(self_addr as isize);

        // SAFETY: The resulting pointer is guaranteed to be NonNull as we only modify the niche bytes
        unsafe { NonNull::new_unchecked(ptr.as_ptr().cast::<u8>().wrapping_offset(offset).cast()) }
    }

    #[inline]
    pub(crate) fn pointer(self) -> NonNull<*const str> {
        // SAFETY: The resulting pointer is guaranteed to be NonNull as we only modify the niche bytes
        unsafe {
            NonNull::new_unchecked(Strict::map_addr(self.packed.as_ptr(), |addr| {
                addr & !Self::BOOL_BITS
            }))
        }
    }

    #[inline]
    pub(crate) fn as_str(&self) -> &str {
        // SAFETY: We always point to a pointer to a str no matter what variant is active
        unsafe { *self.pointer().as_ptr().cast::<&str>() }
    }
}

#[derive(PartialEq, Eq, Hash)]
pub struct Symbol {
    repr: TaggedArcPtr,
}

impl fmt::Debug for Symbol {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_str().fmt(f)
    }
}

const _: () = assert!(std::mem::size_of::<Symbol>() == std::mem::size_of::<NonNull<()>>());
const _: () = assert!(std::mem::align_of::<Symbol>() == std::mem::align_of::<NonNull<()>>());

static MAP: OnceLock<DashMap<SymbolProxy, (), BuildHasherDefault<FxHasher>>> = OnceLock::new();

impl Symbol {
    pub fn intern(s: &str) -> Self {
        let (mut shard, hash) = Self::select_shard(s);
        // Atomically,
        // - check if `obj` is already in the map
        //   - if so, copy out its entry, conditionally bumping the backing Arc and return it
        //   - if not, put it into a box and then into an Arc, insert it, bump the ref-count and return the copy
        // This needs to be atomic (locking the shard) to avoid races with other thread, which could
        // insert the same object between us looking it up and inserting it.
        match shard.raw_entry_mut().from_key_hashed_nocheck(hash, s) {
            RawEntryMut::Occupied(occ) => Self { repr: increase_arc_refcount(occ.key().0) },
            RawEntryMut::Vacant(vac) => Self {
                repr: increase_arc_refcount(
                    vac.insert_hashed_nocheck(
                        hash,
                        SymbolProxy(TaggedArcPtr::arc(Arc::new(Box::<str>::from(s)))),
                        SharedValue::new(()),
                    )
                    .0
                     .0,
                ),
            },
        }
    }

    pub fn integer(i: usize) -> Self {
        match i {
            0 => symbols::INTEGER_0.clone(),
            1 => symbols::INTEGER_1.clone(),
            2 => symbols::INTEGER_2.clone(),
            3 => symbols::INTEGER_3.clone(),
            4 => symbols::INTEGER_4.clone(),
            5 => symbols::INTEGER_5.clone(),
            6 => symbols::INTEGER_6.clone(),
            7 => symbols::INTEGER_7.clone(),
            8 => symbols::INTEGER_8.clone(),
            9 => symbols::INTEGER_9.clone(),
            10 => symbols::INTEGER_10.clone(),
            11 => symbols::INTEGER_11.clone(),
            12 => symbols::INTEGER_12.clone(),
            13 => symbols::INTEGER_13.clone(),
            14 => symbols::INTEGER_14.clone(),
            15 => symbols::INTEGER_15.clone(),
            i => Symbol::intern(&format!("{i}")),
        }
    }

    pub fn empty() -> Self {
        symbols::__empty.clone()
    }

    pub fn as_str(&self) -> &str {
        self.repr.as_str()
    }

    #[inline]
    fn select_shard(
        s: &str,
    ) -> (
        dashmap::RwLockWriteGuard<
            'static,
            HashMap<SymbolProxy, SharedValue<()>, BuildHasherDefault<FxHasher>>,
        >,
        u64,
    ) {
        let storage = MAP.get_or_init(symbols::prefill);
        let hash = {
            let mut hasher = std::hash::BuildHasher::build_hasher(storage.hasher());
            s.hash(&mut hasher);
            hasher.finish()
        };
        let shard_idx = storage.determine_shard(hash as usize);
        let shard = &storage.shards()[shard_idx];
        (shard.write(), hash)
    }

    #[cold]
    fn drop_slow(arc: &Arc<Box<str>>) {
        let (mut shard, hash) = Self::select_shard(arc);

        match Arc::count(arc) {
            0 => unreachable!(),
            1 => unreachable!(),
            2 => (),
            _ => {
                // Another thread has interned another copy
                return;
            }
        }

        let ptr = match shard.raw_entry_mut().from_key_hashed_nocheck::<str>(hash, arc.as_ref()) {
            RawEntryMut::Occupied(occ) => occ.remove_entry(),
            RawEntryMut::Vacant(_) => unreachable!(),
        }
        .0
         .0;
        // SAFETY: We're dropping, we have ownership.
        ManuallyDrop::into_inner(unsafe { ptr.try_as_arc_owned().unwrap() });
        debug_assert_eq!(Arc::count(arc), 1);

        // Shrink the backing storage if the shard is less than 50% occupied.
        if shard.len() * 2 < shard.capacity() {
            shard.shrink_to_fit();
        }
    }
}

impl Drop for Symbol {
    #[inline]
    fn drop(&mut self) {
        // SAFETY: We're dropping, we have ownership.
        let Some(arc) = (unsafe { self.repr.try_as_arc_owned() }) else {
            return;
        };
        // When the last `Ref` is dropped, remove the object from the global map.
        if Arc::count(&arc) == 2 {
            // Only `self` and the global map point to the object.

            Self::drop_slow(&arc);
        }
        // decrement the ref count
        ManuallyDrop::into_inner(arc);
    }
}

impl Clone for Symbol {
    fn clone(&self) -> Self {
        Self { repr: increase_arc_refcount(self.repr) }
    }
}

fn increase_arc_refcount(repr: TaggedArcPtr) -> TaggedArcPtr {
    // SAFETY: We're not dropping the `Arc`.
    let Some(arc) = (unsafe { repr.try_as_arc_owned() }) else {
        return repr;
    };
    // increase the ref count
    mem::forget(Arc::clone(&arc));
    repr
}

impl fmt::Display for Symbol {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_str().fmt(f)
    }
}

// only exists so we can use `from_key_hashed_nocheck` with a &str
#[derive(Debug, PartialEq, Eq)]
struct SymbolProxy(TaggedArcPtr);

impl Hash for SymbolProxy {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.as_str().hash(state);
    }
}

impl Borrow<str> for SymbolProxy {
    fn borrow(&self) -> &str {
        self.0.as_str()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn smoke_test() {
        Symbol::intern("isize");
        let base_len = MAP.get().unwrap().len();
        let hello = Symbol::intern("hello");
        let world = Symbol::intern("world");
        let more_worlds = world.clone();
        let bang = Symbol::intern("!");
        let q = Symbol::intern("?");
        assert_eq!(MAP.get().unwrap().len(), base_len + 4);
        let bang2 = Symbol::intern("!");
        assert_eq!(MAP.get().unwrap().len(), base_len + 4);
        drop(bang2);
        assert_eq!(MAP.get().unwrap().len(), base_len + 4);
        drop(q);
        assert_eq!(MAP.get().unwrap().len(), base_len + 3);
        let default = Symbol::intern("default");
        let many_worlds = world.clone();
        assert_eq!(MAP.get().unwrap().len(), base_len + 3);
        assert_eq!(
            "hello default world!",
            format!("{} {} {}{}", hello.as_str(), default.as_str(), world.as_str(), bang.as_str())
        );
        drop(default);
        assert_eq!(
            "hello world!",
            format!("{} {}{}", hello.as_str(), world.as_str(), bang.as_str())
        );
        drop(many_worlds);
        drop(more_worlds);
        drop(hello);
        drop(world);
        drop(bang);
        assert_eq!(MAP.get().unwrap().len(), base_len);
    }
}