Skip to main content

stringtheory/interning/
fixed_size.rs

1#[cfg(not(feature = "loom"))]
2use std::sync::{atomic::AtomicUsize, Arc, Mutex};
3use std::{
4    num::NonZeroUsize,
5    ptr::NonNull,
6    sync::atomic::Ordering::{AcqRel, Acquire},
7};
8
9#[cfg(feature = "loom")]
10use loom::sync::{atomic::AtomicUsize, Arc, Mutex};
11
12use super::{
13    helpers::{aligned, aligned_string, hash_string, layout_for_data, PackedLengthCapacity},
14    InternedString, Interner,
15};
16use crate::interning::helpers::{ReclaimedEntries, ReclaimedEntry};
17
18const HEADER_LEN: usize = std::mem::size_of::<EntryHeader>();
19const HEADER_ALIGN: usize = std::mem::align_of::<EntryHeader>();
20
21/// The minimum possible length of an entry.
22///
23/// For any entry in the interner, there is already an `EntryHeader` followed by the string data itself. In order to
24/// ensure that entries can be written contiguously, we additionally ensure that the number of bytes we utilize for the
25/// string data is aligned at least as much as `EntryHeader` itself.
26///
27/// This means that the minimum possible length of an entry, or the minimum number of bytes a valid entry could consume,
28/// is the length of the header plus the alignment of the header.
29const MINIMUM_ENTRY_LEN: usize = HEADER_LEN + HEADER_ALIGN;
30
31#[derive(Debug)]
32pub(crate) struct StringState {
33    interner: Arc<Mutex<InternerShardState>>,
34    header: NonNull<EntryHeader>,
35}
36
37impl StringState {
38    #[inline]
39    pub const fn as_str(&self) -> &str {
40        // SAFETY: We ensure `self.header` is well-aligned and points to an initialized `EntryHeader` value when creating `StringState`.
41        unsafe { get_entry_string(self.header) }
42    }
43}
44
45impl PartialEq for StringState {
46    fn eq(&self, other: &Self) -> bool {
47        self.header == other.header
48    }
49}
50
51impl Clone for StringState {
52    fn clone(&self) -> Self {
53        // SAFETY: The caller that creates `StringState` is responsible for ensuring that `self.header` is well-aligned
54        // and points to an initialized `EntryHeader` value.
55        let header = unsafe { self.header.as_ref() };
56        header.increment_active_refs();
57
58        Self {
59            interner: self.interner.clone(),
60            header: self.header,
61        }
62    }
63}
64
65impl Drop for StringState {
66    fn drop(&mut self) {
67        // SAFETY: The caller that creates `StringState` is responsible for ensuring that `self.header` is well-aligned
68        // and points to an initialized `EntryHeader` value.
69        let header = unsafe { self.header.as_ref() };
70        if header.decrement_active_refs() {
71            // We decremented the reference count to zero, so try to mark this entry for reclamation.
72            let mut interner = self.interner.lock().unwrap();
73            interner.mark_for_reclamation(self.header);
74        }
75    }
76}
77
78// SAFETY: We don't take references to the entry header pointer that outlast `StringState`, and the only modification we
79// do to the entry header is through atomic operations, so it's safe to both send and share `StringState` between
80// threads.
81unsafe impl Send for StringState {}
82unsafe impl Sync for StringState {}
83
84/// Metadata about an interner entry.
85///
86/// `EntryHeader` represents the smallest amount of information about an interned entry that's needed to support both
87/// lookup of existing interned strings, as well as the ability to reclaim space in the interner when an entry is no
88/// longer in use.
89struct EntryHeader {
90    /// The hash of the string that this entry represents.
91    hash: u64,
92
93    /// The number of active references to this entry.
94    ///
95    /// Only incremented by the interner itself, and decremented by `InternedString` when it's dropped.
96    refs: AtomicUsize,
97
98    /// Combined length/capacity of the entry, in terms of the string itself.
99    ///
100    /// Notably, this does _not_ include the length of the header itself. For example, an entry holding the string
101    /// "hello, world!" has a string length of 13 bytes, but since we've to pad out to meet our alignment requirements
102    /// for `EntryHeader`, we would end up with a capacity of 16 bytes. As such, `EntryHeader::len` would report `13`,
103    /// while `EntryHeader::capacity` would report `16`. Likewise, `EntryHeader::entry_len` would report `40`,
104    /// accounting for the string capacity (16) as well as the header length itself (24).
105    ///
106    /// As explained in the description of `PackedLengthCapacity`, this does mean strings can't be larger than ~4 GB on
107    /// 64-bit platforms, which isn't a problem we've.
108    len_cap: PackedLengthCapacity,
109}
110
111impl EntryHeader {
112    /// Creates a tombstone entry with the given capacity.
113    ///
114    /// This is to allow for updating a region in the data buffer, which has been reclaimed, such that it's
115    /// identifiable as being unused.
116    fn tombstone(entry: ReclaimedEntry) -> Self {
117        // The usable capacity for a reclaimed entry is the full capacity minus the size of `EntryHeader` itself, as
118        // reclaimed entries represent the _entire_ region in the data buffer, but `EntryHeader` only cares about the
119        // string portion itself.
120        let cap = Self::usable_from_reclaimed(entry);
121
122        Self {
123            hash: 0,
124            refs: AtomicUsize::new(0),
125            len_cap: PackedLengthCapacity::new(cap, 0),
126        }
127    }
128
129    /// Creates a new entry for the given string.
130    fn from_string(hash: u64, s: &str) -> Self {
131        // We're dictating the necessary capacity here, which is the length of the string rounded to the nearest
132        // multiple of the alignment of `EntryHeader`, which ensures that any subsequent entry will be properly aligned.
133        let cap = aligned_string::<Self>(s);
134
135        Self {
136            hash,
137            refs: AtomicUsize::new(1),
138            len_cap: PackedLengthCapacity::new(cap, s.len()),
139        }
140    }
141
142    /// Creates a new entry for the given string, based on the given reclaimed entry.
143    ///
144    /// This maps the entry header to the underlying capacity of the given reclaimed entry, which is done in cases where
145    /// a reclaimed entry is being used when interning a new string, and the reclaimed entry is larger than the string
146    /// being interned, but not large enough that we could split the excess capacity into a new reclaimed entry.
147    fn from_reclaimed_entry(mut entry: ReclaimedEntry, hash: u64, s: &str) -> (Self, Option<ReclaimedEntry>) {
148        // The usable capacity for a reclaimed entry is the full capacity minus the size of `EntryHeader` itself, as
149        // reclaimed entries represent the _entire_ region in the data buffer, but `EntryHeader` only cares about the
150        // string portion itself.
151        let entry_cap = EntryHeader::usable_from_reclaimed(entry);
152        let required_cap = aligned_string::<Self>(s);
153
154        // If the reclaimed entry has enough additional space beyond what we need for the string, we'll split it off and
155        // return it for the caller to keep around in the reclaimed entries list.
156        let remainder = entry_cap - required_cap;
157        let (adjusted_cap, maybe_split_entry) = if remainder >= MINIMUM_ENTRY_LEN {
158            let entry_len = EntryHeader::len_for(s);
159            let split_entry = entry.split_off(entry_len);
160
161            (entry_len - HEADER_LEN, Some(split_entry))
162        } else {
163            (entry_cap, None)
164        };
165
166        let header = Self {
167            hash,
168            refs: AtomicUsize::new(1),
169            len_cap: PackedLengthCapacity::new(adjusted_cap, s.len()),
170        };
171
172        (header, maybe_split_entry)
173    }
174
175    /// Returns the computed length of a complete entry, in bytes, for the given string.
176    ///
177    /// This includes the size of the entry header itself and the string data, when padded for alignment, and represents
178    /// the number of bytes that would be consumed in the data buffer.
179    const fn len_for(s: &str) -> usize {
180        HEADER_LEN + aligned_string::<Self>(s)
181    }
182
183    /// Returns the usable capacity of a reclaimed entry, in bytes.
184    ///
185    /// Usable refers to the number of bytes in a reclaimed entry that could be used for string data, after accounting
186    /// for the size of `EntryHeader` itself.
187    const fn usable_from_reclaimed(entry: ReclaimedEntry) -> usize {
188        entry.capacity() - HEADER_LEN
189    }
190
191    /// Returns the total number of bytes that this entry takes up in the data buffer.
192    const fn entry_len(&self) -> usize {
193        HEADER_LEN + self.capacity()
194    }
195
196    /// Returns the size of the string, in bytes, that this entry can hold.
197    const fn capacity(&self) -> usize {
198        self.len_cap.capacity()
199    }
200
201    /// Returns the size of the string, in bytes, that this entry _actually_ holds.
202    const fn len(&self) -> usize {
203        self.len_cap.len()
204    }
205
206    /// Returns `true` if this entry is currently referenced.
207    fn is_active(&self) -> bool {
208        self.refs.load(Acquire) != 0
209    }
210
211    /// Increments the active reference count by one.
212    fn increment_active_refs(&self) {
213        self.refs.fetch_add(1, AcqRel);
214    }
215
216    /// Decrements the active reference count by one.
217    ///
218    /// Returns `true` if the active reference count is zero _after_ calling this method.
219    fn decrement_active_refs(&self) -> bool {
220        self.refs.fetch_sub(1, AcqRel) == 1
221    }
222
223    /// Attempts to increment the active reference count, but only if it is currently non-zero.
224    ///
225    /// Returns `true` if the reference count was incremented. A return value of `false` means the entry has already
226    /// been released -- its reference count reached zero -- and is pending reclamation, so it must not be reused.
227    fn try_increment_active_refs(&self) -> bool {
228        self.refs
229            .fetch_update(AcqRel, Acquire, |refs| (refs != 0).then_some(refs + 1))
230            .is_ok()
231    }
232}
233
234#[derive(Debug)]
235struct InternerShardState {
236    // Direct pieces of our buffer allocation.
237    ptr: NonNull<u8>,
238    offset: usize,
239    capacity: NonZeroUsize,
240
241    // Number of active entries (strings) in the interner.
242    entries: usize,
243
244    // Length of all active entries, in bytes.
245    //
246    // This is equivalent to `self.offset` minus the total size of all reclaimed entries.
247    len: usize,
248
249    // Markers for entries that can be reused.
250    reclaimed: ReclaimedEntries,
251}
252
253impl InternerShardState {
254    /// Creates a new `InternerShardState` with a pre-allocated buffer that has the given capacity.
255    pub fn with_capacity(capacity: NonZeroUsize) -> Self {
256        assert!(
257            aligned::<EntryHeader>(capacity.get()) <= isize::MAX as usize,
258            "capacity would overflow isize::MAX, which violates layout constraints"
259        );
260
261        // Allocate our data buffer. This is the main backing allocation for all interned strings, and is well-aligned
262        // for `EntryHeader`.
263        //
264        // SAFETY: `layout_for_data` ensures the layout is non-zero.
265        let data_layout = layout_for_data::<EntryHeader>(capacity);
266        let data_ptr = unsafe { std::alloc::alloc(data_layout) };
267        let ptr = match NonNull::new(data_ptr) {
268            Some(ptr) => ptr,
269            None => std::alloc::handle_alloc_error(data_layout),
270        };
271
272        Self {
273            ptr,
274            offset: 0,
275            capacity,
276            entries: 0,
277            len: 0,
278            reclaimed: ReclaimedEntries::new(),
279        }
280    }
281
282    /// Returns the total number of unused bytes that are available for interning.
283    fn available(&self) -> usize {
284        self.capacity.get() - self.offset
285    }
286
287    fn get_entry_ptr(&self, offset: usize) -> NonNull<EntryHeader> {
288        debug_assert!(
289            offset + MINIMUM_ENTRY_LEN <= self.capacity.get(),
290            "offset would point to entry that cannot possibly avoid extending past end of data buffer"
291        );
292
293        // SAFETY: The caller is responsible for ensuring that `offset` is within the bounds of the data buffer, and
294        // that `offset` is well-aligned for `EntryHeader`.
295        let entry_ptr = unsafe { self.ptr.as_ptr().add(offset).cast::<EntryHeader>() };
296        debug_assert!(entry_ptr.is_aligned(), "entry header pointer must be well-aligned");
297
298        // SAFETY: `entry_ptr` is derived from `self.ptr`, which itself is `NonNull<u8>`, and the caller is responsible
299        // for ensuring that `offset` is within the bounds of the data buffer, so we know `entry_ptr` is non-null.
300        unsafe { NonNull::new_unchecked(entry_ptr) }
301    }
302
303    fn find_entry(&self, hash: u64, s: &str) -> Option<NonNull<EntryHeader>> {
304        let mut offset = 0;
305
306        while offset < self.offset {
307            // Construct a pointer to the entry at `offset`, and get a reference to the header value.
308            let header_ptr = self.get_entry_ptr(offset);
309            let header = unsafe { header_ptr.as_ref() };
310
311            // See if this entry is active or not. If it's active, then we'll quickly check the hash/length of the
312            // string to see if this is likely to be a match for `s`.
313            if header.is_active() && header.hash == hash && header.len() == s.len() {
314                // As a final check, we make sure that the entry string and `s` are equal. If they are, then we
315                // have an exact match and will return the entry.
316                //
317                // SAFETY: We know that our header is valid and initialized.
318                let s_entry = unsafe { get_entry_string(header_ptr) };
319                if s_entry == s {
320                    // Reuse this entry, but only if it's still live. The `is_active()` check above is not enough on its
321                    // own: a concurrent `StringState::drop` decrements the reference count _before_ taking this lock, so
322                    // it can drop the count to zero in the window between that check and here. Resurrecting a
323                    // zero-refcount entry would let two droppers each observe a 1 -> 0 transition and each reclaim the
324                    // same slot -- a double reclaim that corrupts the data buffer. If the conditional increment fails,
325                    // the entry is already pending reclamation by its dropper, so we skip it and let the caller intern a
326                    // fresh entry instead.
327                    if header.try_increment_active_refs() {
328                        return Some(header_ptr);
329                    }
330                }
331            }
332
333            // Either this was a reclaimed entry or we didn't have a match, so we move on to the next entry.
334            offset += header.entry_len();
335        }
336
337        None
338    }
339
340    fn write_entry(&mut self, offset: usize, entry_header: EntryHeader, s: &str) -> NonNull<EntryHeader> {
341        debug_assert_eq!(
342            entry_header.len(),
343            s.len(),
344            "entry header length must match string length"
345        );
346
347        let entry_ptr = self.get_entry_ptr(offset);
348        let entry_len = entry_header.entry_len();
349
350        // Write the entry header.
351        unsafe { entry_ptr.as_ptr().write(entry_header) };
352
353        let s_buf = s.as_bytes();
354
355        // Write the string.
356        let entry_s_buf = unsafe {
357            // Take the entry pointer and add 1, which sets our pointer to right _after_ the header.
358            let entry_s_ptr = entry_ptr.as_ptr().add(1).cast::<u8>();
359            std::slice::from_raw_parts_mut(entry_s_ptr, s_buf.len())
360        };
361        entry_s_buf.copy_from_slice(s_buf);
362
363        // Update our internal statistics.
364        self.entries += 1;
365        self.len += entry_len;
366
367        entry_ptr
368    }
369
370    fn write_to_unoccupied(&mut self, s_hash: u64, s: &str) -> NonNull<EntryHeader> {
371        let entry_header = EntryHeader::from_string(s_hash, s);
372
373        // Write the entry to the end of the data buffer.
374        let entry_offset = self.offset;
375        self.offset += entry_header.entry_len();
376
377        self.write_entry(entry_offset, entry_header, s)
378    }
379
380    fn write_to_reclaimed_entry(&mut self, entry: ReclaimedEntry, s_hash: u64, s: &str) -> NonNull<EntryHeader> {
381        let entry_offset = entry.offset();
382        let (entry_header, maybe_split_entry) = EntryHeader::from_reclaimed_entry(entry, s_hash, s);
383
384        // If we had enough capacity in the reclaimed entry to hold this string _and_ potentially hold another entry, we
385        // split it off and store that remainder entry.
386        if let Some(split_entry) = maybe_split_entry {
387            self.add_reclaimed(split_entry);
388        }
389
390        // Write the entry in place of the reclaimed entry.
391        self.write_entry(entry_offset, entry_header, s)
392    }
393
394    fn add_reclaimed_from_header(&mut self, header_ptr: NonNull<EntryHeader>) {
395        // Get the offset of the header within the data buffer.
396        //
397        // SAFETY: The caller is responsible for ensuring the entry header reference belongs to this interner. If that
398        // is upheld, then we know that entry header belongs to our data buffer, and that the pointer to the entry
399        // header is not less than the base pointer of the data buffer, ensuring the offset is non-negative.
400        let entry_offset = unsafe {
401            header_ptr
402                .cast::<u8>()
403                .as_ptr()
404                .offset_from(self.ptr.as_ptr().cast_const())
405        };
406        debug_assert!(entry_offset >= 0, "entry offset must be non-negative");
407
408        let header = unsafe { header_ptr.as_ref() };
409
410        let entry = ReclaimedEntry::new(entry_offset as usize, header.entry_len());
411        self.len -= entry.capacity();
412        self.add_reclaimed(entry);
413    }
414
415    fn add_reclaimed(&mut self, entry: ReclaimedEntry) {
416        // Reclamation is a two-step process: first, we have to actually keep track of the reclaimed entry, which
417        // potentially involves merging adjacent reclaimed entries, and then once all of that has happened, we tombstone
418        // the entry (whether merged or not).
419        //
420        // However, if the merged reclaimed entry immediately precedes any available capacity, we can skip tombstoning
421        // it, since we can just wind back `offset` to reclaim the space.
422        let merged_entry = self.reclaimed.insert(entry);
423        if merged_entry.offset() + merged_entry.capacity() == self.offset {
424            self.offset -= merged_entry.capacity();
425            self.reclaimed.remove(&merged_entry);
426            return;
427        }
428
429        self.clear_reclaimed_entry(merged_entry);
430    }
431
432    fn mark_for_reclamation(&mut self, header_ptr: NonNull<EntryHeader>) {
433        // Reclaim the entry only if its reference count is zero, and rely on `find_entry` to never resurrect it.
434        //
435        // `StringState::drop` decrements the reference count _before_ taking this lock, so a concurrent `try_intern`
436        // can observe this entry during that window. But the only reuse path, `find_entry`, increments via
437        // `try_increment_active_refs`, which refuses to increment a zero reference count -- so once the count reaches
438        // zero it stays zero, and exactly one dropper (the one that transitioned it 1 -> 0) reaches this method. That
439        // makes reclamation exactly-once: no second thread can resurrect the entry and then drop it again, which would
440        // otherwise double-reclaim the same slot and corrupt the data buffer.
441        //
442        // The reference count also cannot rise while we hold this lock (cloning needs a live handle, and `try_intern`
443        // needs this lock), so observing zero here means it stays zero through the reclamation below -- no
444        // use-after-free.
445        //
446        // SAFETY: The caller is responsible for ensuring that `header_ptr` is well-aligned and points to an initialized
447        // `EntryHeader` value.
448        let header = unsafe { header_ptr.as_ref() };
449        if !header.is_active() {
450            self.entries -= 1;
451            self.add_reclaimed_from_header(header_ptr);
452        }
453    }
454
455    fn clear_reclaimed_entry(&mut self, entry: ReclaimedEntry) {
456        let entry_ptr = self.get_entry_ptr(entry.offset);
457
458        // Write the entry tombstone itself, which clears out the hash and sets the reference count to zero.
459        //
460        // SAFETY: We know that `entry_ptr` is valid for writes (reclaimed entries are, by definition, inactive regions
461        // in the data buffer) and is well-aligned for `EntryHeader`.
462        let tombstone = EntryHeader::tombstone(entry);
463        let str_cap = tombstone.capacity();
464
465        unsafe {
466            entry_ptr.as_ptr().write(EntryHeader::tombstone(entry));
467        }
468
469        // Write a magic value to the entire string capacity for the entry. This ensures that there's a known repeating
470        // value which, in the case of debugging issues, can be a signal that offsets/reclaimed entries are incorrect
471        // and overlapping with active entries.
472        //
473        // SAFETY: Like above, the caller is responsible for ensuring that `offset` is within the bounds of the data
474        // buffer, and that `offset + capacity` does not extend past the bounds of the data buffer.
475        unsafe {
476            // Take the entry pointer and add 1, which sets our pointer to right _after_ the header.
477            let str_ptr = entry_ptr.as_ptr().add(1).cast::<u8>();
478            let str_buf = std::slice::from_raw_parts_mut(str_ptr, str_cap);
479            str_buf.fill(0xAA);
480        }
481    }
482
483    fn try_intern(&mut self, s_hash: u64, s: &str) -> Option<NonNull<EntryHeader>> {
484        // We can only intern strings with a size that fits within a packed length/capacity value, so if `s` is larger
485        // than that, we can't intern it, and there's existing entry we could have for it either.
486        if s.len() > PackedLengthCapacity::maximum_value() {
487            return None;
488        }
489
490        // Try and find an existing entry for this string.
491        if self.entries != 0 {
492            if let Some(existing_entry) = self.find_entry(s_hash, s) {
493                return Some(existing_entry);
494            }
495        }
496
497        let required_cap = EntryHeader::len_for(s);
498
499        // We didn't find an existing entry, so we're going to intern it.
500        //
501        // First, try and see if we have a reclaimed entry that can fit this string. If nothing suitable is found, or we
502        // have no reclaimed entries, then we'll just try to fit it in the remaining capacity of our data buffer.
503        if !self.reclaimed.is_empty() {
504            let maybe_reclaimed_entry = self.reclaimed.take_if(|entry| entry.capacity() >= required_cap);
505            if let Some(reclaimed_entry) = maybe_reclaimed_entry {
506                return Some(self.write_to_reclaimed_entry(reclaimed_entry, s_hash, s));
507            }
508        }
509
510        if required_cap <= self.available() {
511            Some(self.write_to_unoccupied(s_hash, s))
512        } else {
513            None
514        }
515    }
516}
517
518impl Drop for InternerShardState {
519    fn drop(&mut self) {
520        // SAFETY: We allocated this buffer with the global allocator, and we're generating the same layout that was
521        // used to allocate it in the first place.
522        unsafe {
523            std::alloc::dealloc(self.ptr.as_ptr(), layout_for_data::<EntryHeader>(self.capacity));
524        }
525    }
526}
527
528// SAFETY: We don't take references to the data buffer pointer that outlast `InternerShardState`, and all access to
529// `InternerShardState` itself is mediated through a mutex, so we're safe to send it around and share it between
530// threads.
531unsafe impl Send for InternerShardState {}
532unsafe impl Sync for InternerShardState {}
533
534#[derive(Debug)]
535struct InternerState<const SHARD_FACTOR: usize> {
536    shards: [Arc<Mutex<InternerShardState>>; SHARD_FACTOR],
537    capacity: usize,
538}
539
540impl<const SHARD_FACTOR: usize> InternerState<SHARD_FACTOR> {
541    // Ensure our shard factor is a power of two at compile, so that we can just use a mask to get the shard index.
542    const _POWER_OF_TWO_SHARD_FACTOR: () = {
543        if !SHARD_FACTOR.is_power_of_two() {
544            panic!("shard factor must be a power of two")
545        }
546    };
547
548    pub fn with_capacity(capacity: NonZeroUsize) -> Self {
549        let shard_capacity = match NonZeroUsize::new(capacity.get() / SHARD_FACTOR) {
550            Some(shard_capacity) => shard_capacity,
551            // If we can't divide the capacity evenly, we just specify a capacity of one which will force every shard to
552            // upsize the capacity so that a single entry can fit, and satisfies the need for `NonZeroUsize`.
553            //
554            // SAFETY: One is obviously not zero.
555            None => unsafe { NonZeroUsize::new_unchecked(1) },
556        };
557
558        let shards = std::iter::repeat_with(|| Arc::new(Mutex::new(InternerShardState::with_capacity(shard_capacity))))
559            .take(SHARD_FACTOR)
560            .collect::<Vec<_>>();
561        let capacity = shards
562            .iter()
563            .map(|shard| {
564                let shard = shard.lock().unwrap();
565                shard.capacity.get()
566            })
567            .sum();
568        let shards: [Arc<Mutex<InternerShardState>>; SHARD_FACTOR] =
569            shards.try_into().expect("should not fail to convert to array");
570
571        Self { shards, capacity }
572    }
573
574    fn is_empty(&self) -> bool {
575        self.shards.iter().any(|shard| shard.lock().unwrap().entries != 0)
576    }
577
578    fn len(&self) -> usize {
579        self.shards.iter().map(|shard| shard.lock().unwrap().entries).sum()
580    }
581
582    fn len_bytes(&self) -> usize {
583        self.shards.iter().map(|shard| shard.lock().unwrap().offset).sum()
584    }
585
586    fn capacity_bytes(&self) -> usize {
587        self.capacity
588    }
589
590    fn try_intern(&self, s: &str) -> Option<InternedString> {
591        let hash = hash_string(s);
592        let shard_idx = (hash as usize) & (SHARD_FACTOR - 1);
593
594        let shard = &self.shards[shard_idx];
595        intern_with_shard_and_hash(shard, hash, s)
596    }
597}
598
599/// A string interner based on a single, fixed-size backing buffer with support for reclamation.
600///
601/// ## Overview
602///
603/// This interner uses a single, fixed-size backing buffer where interned strings are stored contiguously. This provides
604/// bounded memory usage, and the interner won't allocate additional memory for new strings once the buffer is full.
605/// Since interned strings aren't likely to need to live for the life of the program, the interner supports
606/// reclamation. Once all references to an interned string have been dropped, the storage for that string is reclaimed
607/// and can be used to hold new strings.
608///
609/// ## Storage layout
610///
611/// The backing buffer stores strings contiguously, with an entry "header" prepended to each string. The header contains
612/// relevant data -- hash of the string, reference count, and length of the string -- needed to work with the entry
613/// either when searching for existing entries or when using the entry itself.
614///
615/// The layout of an entry is as follows:
616///
617/// ```text
618/// ┌───────────────────────── entry #1 ──────────────────────────┐ ┌─ entry #2 ─┐ ┌─ entry .. ─┐
619/// ▼                                                             ▼ ▼            ▼ ▼            ▼
620/// ┏━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━┓
621/// ┃ str hash  │  ref cnt  │  str len  │ str data  │   padding   ┃ ┃   header   ┃ ┃   header   ┃
622/// ┃ (8 bytes) │ (8 bytes) │ (8 bytes) │ (N bytes) │ (1-7 bytes) ┃ ┃  & string  ┃ ┃  & string  ┃
623/// ┗━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━┛
624/// ▲                                   ▲                           ▲
625/// └────────── `EntryHeader` ──────────┘                           └── aligned for `EntryHeader`
626///          (8 byte alignment)                                         via trailing padding
627/// ```
628///
629/// The backing buffer is always aligned properly for `EntryHeader`, so that the first entry can be referenced
630/// correctly. However, when appending additional entries to the buffer, we need to ensure that those entries also have
631/// an aligned start for accessing the header. This is complicated due to the variable number of bytes for the string
632/// data.
633///
634/// Alignment padding is added to the end of the entry to ensure that when appending the next entry, the start of the
635/// entry is properly aligned for `EntryHeader`. In the worst case, up to 7 bytes could be added (and thus wasted) on
636/// this alignment padding.
637///
638/// ## `InternedString`
639///
640/// The `InternedString` type is a handle to the entry header, and thus the string data, for an interned string. It's
641/// designed to be small -- 8 bytes. -- and cheap to clone, as it contains an atomic reference to the entry header and a
642/// reference to the interner that owns the string. It dereferences to the underlying string with relatively low
643/// overhead: two pointer indirections.
644///
645/// When an `InternedString` is dropped, it decrements the reference count for the entry it points to. If the reference
646/// count drops to zero, it will attempt to mark the entry for reclamation.
647///
648/// ## Reclamation
649///
650/// As we want to bound the memory used by the interner, but also not allow it to be filled up with strings that
651/// eventually end up going entirely unused, we need a way to remove those unused strings so their underlying storage
652/// can be used for new strings. This is where reclamation comes in.
653///
654/// When a string is interned, the entry header tracks how many active references there are to it. When that reference
655/// count drops to zero, the last reference to the string attempts to mark the entry for reclamation. Assuming no other
656/// reference has been taken out on the entry in the meantime, the entry gets added to a list of "reclaimed" entries.
657///
658/// Reclaimed entries are simple markers -- start and end position in the data buffer -- which are stored in a freelist.
659/// When attempting to intern a new string, this freelist is searched to see if there's an entry large enough to fit the
660/// new string, and if so, it's used.
661///
662/// Additionally, when entries are reclaimed, adjacent entries are merged together where possible. This helps to avoid
663/// unnecessary fragmentation over time, although not as effectively as reconstructing the data buffer to re-pack
664/// entries.
665#[derive(Clone, Debug)]
666pub struct FixedSizeInterner<const SHARD_FACTOR: usize> {
667    state: Arc<InternerState<SHARD_FACTOR>>,
668}
669
670impl<const SHARD_FACTOR: usize> FixedSizeInterner<SHARD_FACTOR> {
671    /// Creates a new `FixedSizeInterner` with the given capacity.
672    ///
673    /// The given capacity will potentially be rounded up by a small number of bytes (up to 7) in order to ensure the
674    /// backing buffer is properly aligned.
675    pub fn new(capacity: NonZeroUsize) -> Self {
676        Self {
677            state: Arc::new(InternerState::with_capacity(capacity)),
678        }
679    }
680}
681
682impl<const SHARD_FACTOR: usize> Interner for FixedSizeInterner<SHARD_FACTOR> {
683    fn is_empty(&self) -> bool {
684        self.state.is_empty()
685    }
686
687    fn len(&self) -> usize {
688        self.state.len()
689    }
690
691    fn len_bytes(&self) -> usize {
692        self.state.len_bytes()
693    }
694
695    fn capacity_bytes(&self) -> usize {
696        self.state.capacity_bytes()
697    }
698
699    fn try_intern(&self, s: &str) -> Option<InternedString> {
700        self.state.try_intern(s)
701    }
702}
703
704#[inline]
705const unsafe fn get_entry_string_parts(header_ptr: NonNull<EntryHeader>) -> (NonNull<u8>, usize) {
706    // SAFETY: The caller is responsible for ensuring that `header_ptr` is well-aligned and points to an initialized
707    // `EntryHeader` value.
708    let header = header_ptr.as_ref();
709
710    // Advance past the header and get the pointer to the string.
711    //
712    // SAFETY: We know that we're simply skipping over the header by advancing the pointer by one when it's still typed
713    // as `*mut EntryHeader`.
714    let s_ptr = header_ptr.add(1).cast::<u8>();
715    (s_ptr, header.len())
716}
717
718#[inline]
719const unsafe fn get_entry_string<'a>(header_ptr: NonNull<EntryHeader>) -> &'a str {
720    let (s_ptr, s_len) = get_entry_string_parts(header_ptr);
721
722    // SAFETY: We depend on `get_entry_string_parts` to give us a valid pointer and length for the string.
723    std::str::from_utf8_unchecked(std::slice::from_raw_parts(s_ptr.as_ptr() as *const _, s_len))
724}
725
726fn intern_with_shard_and_hash(shard: &Arc<Mutex<InternerShardState>>, hash: u64, s: &str) -> Option<InternedString> {
727    let header = {
728        let mut shard = shard.lock().unwrap();
729        shard.try_intern(hash, s)?
730    };
731
732    Some(InternedString::from(StringState {
733        interner: Arc::clone(shard),
734        header,
735    }))
736}
737
738#[cfg(test)]
739mod tests {
740    use std::{
741        collections::HashSet,
742        ops::{Deref as _, RangeInclusive},
743    };
744
745    use prop::sample::Index;
746    use proptest::{
747        collection::{hash_set, vec as arb_vec},
748        prelude::*,
749    };
750
751    use super::*;
752    use crate::interning::InternedStringState;
753
754    pub(super) fn create_shard(capacity: NonZeroUsize) -> Arc<Mutex<InternerShardState>> {
755        Arc::new(Mutex::new(InternerShardState::with_capacity(capacity)))
756    }
757
758    pub(super) fn intern_for_shard(shard: &Arc<Mutex<InternerShardState>>, s: &str) -> Option<InternedString> {
759        let hash = hash_string(s);
760        intern_with_shard_and_hash(shard, hash, s)
761    }
762
763    fn shard_capacity(shard: &Arc<Mutex<InternerShardState>>) -> usize {
764        shard.lock().unwrap().capacity.get()
765    }
766
767    fn shard_available(shard: &Arc<Mutex<InternerShardState>>) -> usize {
768        shard.lock().unwrap().available()
769    }
770
771    fn shard_entries(shard: &Arc<Mutex<InternerShardState>>) -> usize {
772        shard.lock().unwrap().entries
773    }
774
775    fn shard_reclaimed_len(shard: &Arc<Mutex<InternerShardState>>) -> usize {
776        shard.lock().unwrap().reclaimed.len()
777    }
778
779    fn shard_first_reclaimed_entry(shard: &Arc<Mutex<InternerShardState>>) -> ReclaimedEntry {
780        shard.lock().unwrap().reclaimed.first().unwrap()
781    }
782
783    pub(super) fn get_reclaimed_entry_for_string(s: &InternedString) -> ReclaimedEntry {
784        let state = match &s.state {
785            InternedStringState::FixedSize(state) => state,
786            _ => panic!("unexpected string state"),
787        };
788
789        let ptr = state.interner.lock().unwrap().ptr.as_ptr();
790        let header = unsafe { state.header.as_ref() };
791        let offset = unsafe { state.header.as_ptr().cast::<u8>().offset_from(ptr) as usize };
792        ReclaimedEntry::new(offset, header.entry_len())
793    }
794
795    fn entry_len(s: &str) -> usize {
796        EntryHeader::len_for(s)
797    }
798
799    fn arb_alphanum_strings(
800        str_len: RangeInclusive<usize>, unique_strs: RangeInclusive<usize>,
801    ) -> impl Strategy<Value = Vec<String>> {
802        // Create characters between 0x20 (32) and 0x7E (126), which are all printable ASCII characters.
803        let char_gen = any::<u8>().prop_map(|c| std::cmp::max(c % 127, 32));
804
805        let str_gen = any::<usize>()
806            .prop_map(move |n| std::cmp::max(n % *str_len.end(), *str_len.start()))
807            .prop_flat_map(move |len| arb_vec(char_gen.clone(), len))
808            // SAFETY: We know our characters are all valid UTF-8 because they're in the ASCII range.
809            .prop_map(|xs| unsafe { String::from_utf8_unchecked(xs) });
810
811        // Create a hash set, which handles the deduplication aspect for us, ensuring we have N unique strings where N
812        // is within the `unique_strs` range... and then convert it to `Vec<String>` for easier consumption.
813        hash_set(str_gen, unique_strs).prop_map(|unique_strs| unique_strs.into_iter().collect::<Vec<_>>())
814    }
815
816    #[test]
817    fn basic() {
818        let interner = FixedSizeInterner::<1>::new(NonZeroUsize::new(1024).unwrap());
819
820        let s1 = interner.try_intern("hello").unwrap();
821        let s2 = interner.try_intern("world").unwrap();
822        let s3 = interner.try_intern("hello").unwrap();
823
824        assert_eq!(s1.deref(), "hello");
825        assert_eq!(s2.deref(), "world");
826        assert_eq!(s3.deref(), "hello");
827
828        // The pointers from the interned strings should be the same, but not between the interned string and a pointer
829        // to an equivalent (but not interned) string:
830        assert!(std::ptr::eq(s1.deref() as *const _, s3.deref() as *const _));
831
832        let local_hello = "hello";
833        assert!(!std::ptr::eq(s1.deref() as *const _, local_hello as *const _));
834    }
835
836    #[test]
837    fn try_intern_without_capacity() {
838        // Big enough to fit a single "hello world!" string, but not big enough to fit two.
839        let interner = FixedSizeInterner::<1>::new(NonZeroUsize::new(64).unwrap());
840
841        let s1 = interner.try_intern("hello world!");
842        assert!(s1.is_some());
843
844        let s2 = interner.try_intern("hello, world");
845        assert!(s2.is_none());
846    }
847
848    #[test]
849    fn reclaim_after_dropped() {
850        let shard = create_shard(NonZeroUsize::new(1024).unwrap());
851
852        let s1 = intern_for_shard(&shard, "hello").expect("should not fail to intern");
853        let s1_entry_len = entry_len(&s1);
854
855        assert_eq!(shard_entries(&shard), 1);
856        assert_eq!(shard_available(&shard), 1024 - s1_entry_len);
857        assert_eq!(shard_reclaimed_len(&shard), 0);
858
859        // Drop the interned string, which should decrement the reference count to zero and then reclaim the entry.
860        drop(s1);
861
862        assert_eq!(shard_entries(&shard), 0);
863        assert_eq!(shard_available(&shard), 1024);
864        assert_eq!(shard_reclaimed_len(&shard), 0);
865    }
866
867    #[test]
868    fn interns_to_reclaimed_entry_with_leftover() {
869        // We want to intern a string initially, which takes up almost all of the capacity, and then drop it so it gets
870        // reclaimed. After that, we'll intern a much smaller string which should lead to utilizing that reclaimed
871        // entry, but only a part of it. Finally, we'll intern another new string.
872        //
873        // The point is to demonstrate that our reclamation logic is sound in terms of allowing reclaimed entries to be
874        // split while the search/insertion logic is operating.
875        let shard_capacity = NonZeroUsize::new(256).unwrap();
876        let shard = create_shard(shard_capacity);
877
878        // We craft four strings such that the first two (`s_large` and `s_medium1`) will take up enough capacity that
879        // `s_small` can't possibly be interned in the available capacity. We'll also craft `s_medium2` so it can fit
880        // within the reclaimed entry for `s_large` but takes enough capacity that `s_small` cannot fit in the leftover
881        // reclaimed entry that we split off.
882        let s_large = "99 bottles of beer on the wall, 99 bottles of beer! take one down, pass it around, 98 bottles of beer on the wall!";
883        let s_medium1 = "no act of kindness, no matter how small, is ever wasted";
884        let s_medium2 = "if you want to go fast, go alone; if you want to go far, go together";
885        let s_small = "are you there god? it's me, margaret";
886
887        let phase1_available_capacity = shard_capacity.get() - entry_len(s_large) - entry_len(s_medium1);
888        assert!(phase1_available_capacity < entry_len(s_small));
889        assert!((entry_len(s_large) - entry_len(s_medium2)) < entry_len(s_small));
890        assert!(entry_len(s_medium2) < entry_len(s_large));
891
892        // Phase 1: intern our two larger strings.
893        let s1 = intern_for_shard(&shard, s_large).expect("should not fail to intern");
894        let s2 = intern_for_shard(&shard, s_medium1).expect("should not fail to intern");
895
896        assert_eq!(shard_entries(&shard), 2);
897        assert_eq!(shard_available(&shard), phase1_available_capacity);
898        assert_eq!(shard_reclaimed_len(&shard), 0);
899
900        // Phase 2: drop `s_large` so it gets reclaimed.
901        drop(s1);
902        assert_eq!(shard_entries(&shard), 1);
903        assert_eq!(shard_reclaimed_len(&shard), 1);
904
905        // Phase 3: intern `s_medium2`, which should fit in the reclaimed entry for `s_large`. This should leave a
906        // small, split off reclaimed entry.
907        let s3 = intern_for_shard(&shard, s_medium2).expect("should not fail to intern");
908
909        assert_eq!(shard_entries(&shard), 2);
910        assert_eq!(shard_reclaimed_len(&shard), 1);
911
912        // Phase 4: intern `s_small`, which should not fit in the leftover reclaimed entry from `s_large` _or_ the
913        // available capacity.
914        let s4 = intern_for_shard(&shard, s_small);
915        assert_eq!(s4, None);
916
917        assert_eq!(shard_entries(&shard), 2);
918        assert_eq!(shard_reclaimed_len(&shard), 1);
919
920        // And make sure we can still dereference the interned strings we _do_ have left:
921        assert_eq!(s2.deref(), s_medium1);
922        assert_eq!(s3.deref(), s_medium2);
923    }
924
925    #[test]
926    fn has_reclaimed_entries_string_fits_exactly() {
927        // The shard is large enough to fit two of the identically-sized strings, but not all three. We show that when we drop
928        // an entry and it is reclaimed, a string of identical size should always be able to reuse that reclaimed entry.
929        const S1_VALUE: &str = "hello world!";
930        const S2_VALUE: &str = "hello, world";
931        const S3_VALUE: &str = "hello--world";
932
933        let shard = create_shard(NonZeroUsize::new(80).unwrap());
934
935        // Intern the first two strings, which should fit without issue.
936        let s1 = intern_for_shard(&shard, S1_VALUE).expect("should not fail to intern");
937        let s1_reclaimed_expected = get_reclaimed_entry_for_string(&s1);
938        let _s2 = intern_for_shard(&shard, S2_VALUE).expect("should not fail to intern");
939
940        assert_eq!(shard_entries(&shard), 2);
941        assert_eq!(shard_reclaimed_len(&shard), 0);
942
943        // Try to intern the third string, which should fail as we don't have the space.
944        let s3 = intern_for_shard(&shard, S3_VALUE);
945        assert_eq!(s3, None);
946
947        // Drop the first string, which should decrement the reference count to zero and then reclaim the entry.
948        drop(s1);
949
950        assert_eq!(shard_entries(&shard), 1);
951        assert_eq!(shard_reclaimed_len(&shard), 1);
952
953        let s1_reclaimed = shard_first_reclaimed_entry(&shard);
954        assert_eq!(s1_reclaimed_expected, s1_reclaimed);
955
956        // Try again to intern the third string, which should now succeed and take over the reclaimed entry entirely
957        // as the strings are identical in length.
958        let _s3 = intern_for_shard(&shard, S3_VALUE).expect("should not fail to intern");
959
960        assert_eq!(shard_entries(&shard), 2);
961        assert_eq!(shard_reclaimed_len(&shard), 0);
962    }
963
964    #[test]
965    fn reclaimed_entry_reuse_split_too_small() {
966        // This situation is slightly contrived, but: we want to test that when there's a reclaimed entry of a certain
967        // size, reusing that reclaimed entry won't lead to it being split if the resulting split entry would be too
968        // "small": unable to hold another minimum-sized entry.
969        //
970        // We have to intern three strings to do this because we only track reclaimed entries when they're followed by
971        // in-use entries, and the string we drop to create a reclaimed entry has to be big enough, but not _too_ big,
972        // to hold the string we want to intern after it.
973        let shard = create_shard(NonZeroUsize::new(128).unwrap());
974
975        // Declare our strings to intern and just check some preconditions by hand.
976        let s_one = "a horse, a horse, my kingdom for a horse!";
977        let s_one_entry_len = entry_len(s_one);
978        let s_two = "why hello there, beautiful";
979        let s_two_entry_len = entry_len(s_two);
980        let s_three = "real gs move in silence like lasagna";
981        let s_three_entry_len = entry_len(s_three);
982
983        assert!(s_one_entry_len <= shard_capacity(&shard));
984        assert!(s_two_entry_len <= shard_capacity(&shard));
985        assert!(s_one_entry_len + s_two_entry_len + s_three_entry_len > shard_capacity(&shard));
986        assert!(s_one_entry_len > s_two_entry_len);
987        assert!(s_three_entry_len > s_two_entry_len);
988        assert!((s_one_entry_len - s_three_entry_len) < MINIMUM_ENTRY_LEN);
989
990        // Intern the first two strings, which should fit without issue.
991        let s1 = intern_for_shard(&shard, s_one).expect("should not fail to intern");
992        let s1_reclaimed_expected = get_reclaimed_entry_for_string(&s1);
993        let _s2 = intern_for_shard(&shard, s_two).expect("should not fail to intern");
994
995        assert_eq!(shard_entries(&shard), 2);
996        assert_eq!(shard_reclaimed_len(&shard), 0);
997
998        // Try to intern the third string, which should fail as we don't have the space.
999        let s3 = intern_for_shard(&shard, s_three);
1000        assert_eq!(s3, None);
1001
1002        // Drop the first string, which should decrement the reference count to zero and then reclaim the entry.
1003        drop(s1);
1004
1005        assert_eq!(shard_entries(&shard), 1);
1006        assert_eq!(shard_reclaimed_len(&shard), 1);
1007
1008        let s1_reclaimed = shard_first_reclaimed_entry(&shard);
1009        assert_eq!(s1_reclaimed_expected, s1_reclaimed);
1010
1011        // Try again to intern the third string, which should now succeed and take over the reclaimed entry, but since
1012        // the remainder of the reclaimed entry after taking the necessary capacity for `s_three` is not large enough
1013        // (`MINIMUM_ENTRY_LEN`), we shouldn't end up splitting the reclaimed entry, and instead, `s3` should consume
1014        // the entire reclaimed entry.
1015        let s3 = intern_for_shard(&shard, s_three).expect("should not fail to intern");
1016        let s3_reclaimed_expected = get_reclaimed_entry_for_string(&s3);
1017
1018        assert_eq!(shard_entries(&shard), 2);
1019        assert_eq!(shard_reclaimed_len(&shard), 0);
1020        assert_eq!(s1_reclaimed_expected, s3_reclaimed_expected);
1021    }
1022
1023    #[test]
1024    fn reclaimed_entry_adjacent_to_spare_capacity() {
1025        let shard = create_shard(NonZeroUsize::new(128).unwrap());
1026
1027        // Intern two smallish strings that fit without issue.
1028        let s1 = intern_for_shard(&shard, "hello, world!").expect("should not fail to intern");
1029        let s2 = intern_for_shard(&shard, "cheeeeeehooooo!").expect("should not fail to intern");
1030        let s1_entry_len = entry_len(&s1);
1031        let s2_entry_len = entry_len(&s2);
1032
1033        assert_eq!(shard_reclaimed_len(&shard), 0);
1034        assert_eq!(shard_available(&shard), 128 - s1_entry_len - s2_entry_len);
1035
1036        drop(s2);
1037        assert_eq!(shard_reclaimed_len(&shard), 0);
1038        assert_eq!(shard_available(&shard), 128 - s1_entry_len);
1039
1040        drop(s1);
1041        assert_eq!(shard_reclaimed_len(&shard), 0);
1042        assert_eq!(shard_available(&shard), 128);
1043    }
1044
1045    proptest! {
1046        #[test]
1047        #[cfg_attr(miri, ignore)]
1048        fn property_test_entry_count_accurate(
1049            strs in arb_alphanum_strings(1..=128, 16..=512),
1050            indices in arb_vec(any::<Index>(), 1..=1000),
1051        ) {
1052            // We ask `proptest` to generate a bunch of unique strings of varying lengths (1-128 bytes, 16-512 unique
1053            // strings) which we then randomly select out of those strings which strings we want to intern. The goal
1054            // here is to potentially select the same string multiple times, to exercise the actual interning logic...
1055            // but practically, to ensure that when we intern a string that has already been interned, we're not
1056            // incrementing the entries count again.
1057
1058            // Create an interner with enough capacity to hold all of the strings we've generated. This is the maximum
1059            // string size multiplied by the number of strings we've generated... plus a little constant factor per
1060            // string to account for the entry header.
1061            const ENTRY_SIZE: usize = 128 + HEADER_LEN;
1062            let interner = FixedSizeInterner::<1>::new(NonZeroUsize::new(ENTRY_SIZE * indices.len()).unwrap());
1063
1064            // For each index, pull out the string and both track it in `unique_strs` and intern it. We hold on to the
1065            // interned string handle to make sure the interned string is actually kept alive, keeping the entry count
1066            // stable.
1067            let mut interned = Vec::new();
1068            let mut unique_strs = HashSet::new();
1069            for index in &indices {
1070                let s = index.get(&strs);
1071                unique_strs.insert(s);
1072
1073                let s_interned = interner.try_intern(s).expect("should never fail to intern");
1074                interned.push(s_interned);
1075            }
1076
1077            assert_eq!(unique_strs.len(), interner.len());
1078        }
1079    }
1080}
1081
1082#[cfg(all(test, feature = "loom"))]
1083mod loom_tests {
1084    use std::{num::NonZeroUsize, ops::Deref};
1085
1086    use loom::sync::{Arc, Mutex};
1087
1088    use super::{
1089        tests::{create_shard, get_reclaimed_entry_for_string, intern_for_shard},
1090        *,
1091    };
1092
1093    #[test]
1094    fn concurrent_drop_and_intern() {
1095        fn shard_reclaimed_entries(shard: &Arc<Mutex<InternerShardState>>) -> Vec<ReclaimedEntry> {
1096            shard.lock().unwrap().reclaimed.iter().copied().collect()
1097        }
1098
1099        fn do_reclaimed_entries_overlap(a: ReclaimedEntry, b: ReclaimedEntry) -> bool {
1100            let a_start = a.offset;
1101            let a_end = a.offset + a.capacity - 1;
1102
1103            let b_start = b.offset;
1104            let b_end = b.offset + b.capacity - 1;
1105
1106            (a_start <= b_start && b_start <= a_end) || (a_start <= b_end && b_end <= a_end)
1107        }
1108
1109        const STRING_TO_INTERN: &str = "hello, world!";
1110
1111        // This test is meant to explore the thread orderings when one thread is trying to drop (and thus reclaim) the
1112        // last active reference to an interned string, and another thread is trying to intern that very same string.
1113        //
1114        // We accept, as a caveat, that a possible outcome is that we intern the "new" string again, even though an
1115        // existing entry to that string may have existed in an alternative ordering.
1116        loom::model(|| {
1117            let shard = create_shard(NonZeroUsize::new(1024).unwrap());
1118            let t2_shard = Arc::clone(&shard);
1119
1120            // Intern the string from thread T1.
1121            let t1_interned_s = intern_for_shard(&shard, STRING_TO_INTERN).expect("should not fail to intern");
1122            assert_eq!(t1_interned_s.deref(), STRING_TO_INTERN);
1123            let t1_reclaimed_entry = get_reclaimed_entry_for_string(&t1_interned_s);
1124
1125            // Spawn thread T2, which tries to intern the same string and returns the handle to it.
1126            let t2_result = loom::thread::spawn(move || {
1127                let interned_s = intern_for_shard(&t2_shard, STRING_TO_INTERN).expect("should not fail to intern");
1128                let reclaimed_entry = get_reclaimed_entry_for_string(&interned_s);
1129
1130                (interned_s, reclaimed_entry)
1131            });
1132
1133            drop(t1_interned_s);
1134
1135            let (t2_interned_s, t2_reclaimed_entry) = t2_result.join().expect("should not fail to join T2");
1136            assert_eq!(t2_interned_s.deref(), STRING_TO_INTERN);
1137
1138            // What we're checking for here is that either:
1139            // - there's no reclaimed entries (T2 found the existing entry for the string before T1 dropped it)
1140            // - there's a reclaimed entry (T2 didn't find the existing entry for the string before T1 marked it as
1141            //   inactive) but the reclaimed entry does _not_ overlap with the interned string from T2, meaning we
1142            //   didn't get confused and allow T2 to use an existing entry that T1 then later marked as reclaimed
1143            let reclaimed_entries = shard_reclaimed_entries(&shard);
1144            assert!(reclaimed_entries.len() <= 1, "should have at most one reclaimed entry");
1145
1146            if !reclaimed_entries.is_empty() {
1147                // If we do have a reclaimed entry, it needs to match exactly with only one of the interned strings.
1148                let is_t1_entry = reclaimed_entries.first().unwrap() == &t1_reclaimed_entry;
1149                let is_t2_entry = reclaimed_entries.first().unwrap() == &t2_reclaimed_entry;
1150
1151                assert!(
1152                    (is_t1_entry || is_t2_entry) && !(is_t1_entry && is_t2_entry),
1153                    "should only match one interned string"
1154                );
1155
1156                // Additionally, we ensure that the reclaimed entry does not overlap with the other interned string.
1157                assert!(
1158                    !do_reclaimed_entries_overlap(t1_reclaimed_entry, t2_reclaimed_entry),
1159                    "reclaimed entry should not overlap with remaining interned string"
1160                );
1161            }
1162        });
1163    }
1164
1165    #[test]
1166    fn concurrent_resurrect_and_double_drop() {
1167        // Regression test for a double-reclamation race, mirroring the `GenericMapInterner` test of the same name.
1168        //
1169        // Unlike `concurrent_drop_and_intern` -- where the second thread keeps its handle alive -- here the second
1170        // thread interns the string _and_ drops it within the race window. This models the dangerous interleaving:
1171        //
1172        //   1. T1 drops the last reference: `refs` 1 -> 0 (but T1 hasn't taken the shard lock yet).
1173        //   2. T2 calls `try_intern`; `find_entry` observes the entry as active and resurrects it: `refs` 0 -> 1.
1174        //   3. T2 drops its handle: `refs` 1 -> 0.
1175        //   4. T1 and T2 each take the lock and observe `refs == 0`, so BOTH mark the same slot for reclamation.
1176        //
1177        // A double reclamation either underflows `entries`/`len` (a panic under overflow-checks, which the loom test
1178        // profile enables) or inserts a spurious/overlapping reclaimed entry. After every handle to the single interned
1179        // string has been dropped, the only correct end state is a fully empty shard.
1180        const STRING_TO_INTERN: &str = "hello, world!";
1181
1182        loom::model(|| {
1183            let shard = create_shard(NonZeroUsize::new(1024).unwrap());
1184            let t2_shard = Arc::clone(&shard);
1185
1186            // T1 interns the string and holds the only reference.
1187            let t1_interned_s = intern_for_shard(&shard, STRING_TO_INTERN).expect("should not fail to intern");
1188
1189            // T2 interns the same string and immediately drops it. The interesting orderings are those where this runs
1190            // while T1's drop has decremented the refcount to zero but not yet taken the shard lock.
1191            let t2 = loom::thread::spawn(move || {
1192                let t2_interned_s = intern_for_shard(&t2_shard, STRING_TO_INTERN).expect("should not fail to intern");
1193                assert_eq!(t2_interned_s.deref(), STRING_TO_INTERN);
1194                drop(t2_interned_s);
1195            });
1196
1197            // T1 drops its (originally last) reference.
1198            drop(t1_interned_s);
1199
1200            t2.join().expect("should not fail to join T2");
1201
1202            // Every handle has been dropped, so the slot must have been reclaimed exactly once: the shard is fully
1203            // empty, with consistent accounting and no leftover reclaimed entries. A double reclaim breaks at least one
1204            // of these (and typically panics first via the `entries`/`len` underflow).
1205            let shard = shard.lock().unwrap();
1206            assert_eq!(
1207                shard.entries, 0,
1208                "entries should be zero (it is double-subtracted on a double reclaim)"
1209            );
1210            assert_eq!(
1211                shard.len, 0,
1212                "len should be zero (it is double-subtracted on a double reclaim)"
1213            );
1214            assert_eq!(shard.offset, 0, "offset should wind back to zero");
1215            assert!(
1216                shard.reclaimed.is_empty(),
1217                "no reclaimed entries should remain (a double reclaim inserts a spurious/overlapping entry)"
1218            );
1219        });
1220    }
1221}