string_cache/

atom.rs

// Copyright 2014 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use crate::dynamic_set::{Entry, dynamic_set};
use crate::static_sets::StaticAtomSet;
use debug_unreachable::debug_unreachable;

use std::borrow::Cow;
use std::cmp::Ordering::{self, Equal};
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::mem;
use std::num::NonZeroU64;
use std::ops;
use std::slice;
use std::str;
use std::sync::atomic::Ordering::SeqCst;

const DYNAMIC_TAG: u8 = 0b_00;
const INLINE_TAG: u8 = 0b_01; // len in upper nybble
const STATIC_TAG: u8 = 0b_10;
const TAG_MASK: u64 = 0b_11;

/// With alignment, a `*const Entry` pointer always has zeroes in its lowest `TAG_BITS` bits
const _: () = assert!(mem::align_of::<Entry>() >= TAG_MASK.next_power_of_two() as usize);

const LEN_OFFSET: u64 = 4;
const LEN_MASK: u64 = 0xF0;

const MAX_INLINE_LEN: usize = 7;
const STATIC_SHIFT_BITS: usize = 32;

/// Represents a string that has been interned.
///
/// While the type definition for `Atom` indicates that it generic on a particular
/// implementation of an atom set, you don't need to worry about this.  Atoms can be static
/// and come from a `StaticAtomSet` generated by the `string_cache_codegen` crate, or they
/// can be dynamic and created by you on an `EmptyStaticAtomSet`.
///
/// `Atom` implements `Clone` but not `Copy`, since internally atoms are reference-counted;
/// this means that you may need to `.clone()` an atom to keep copies to it in different
/// places, or when passing it to a function that takes an `Atom` rather than an `&Atom`.
///
/// ## Creating an atom at runtime
///
/// If you use `string_cache_codegen` to generate a precomputed list of atoms, your code
/// may then do something like read data from somewhere and extract tokens that need to be
/// compared to the atoms.  In this case, you can use `Atom::from(&str)` or
/// `Atom::from(String)`.  These create a reference-counted atom which will be
/// automatically freed when all references to it are dropped.
///
/// This means that your application can safely have a loop which tokenizes data, creates
/// atoms from the tokens, and compares the atoms to a predefined set of keywords, without
/// running the risk of arbitrary memory consumption from creating large numbers of atoms —
/// as long as your application does not store clones of the atoms it creates along the
/// way.
///
/// For example, the following is safe and will not consume arbitrary amounts of memory:
///
/// ```ignore
/// let untrusted_data = "large amounts of text ...";
///
/// for token in untrusted_data.split_whitespace() {
///     let atom = Atom::from(token); // interns the string
///
///     if atom == Atom::from("keyword") {
///         // handle that keyword
///     } else if atom == Atom::from("another_keyword") {
///         // handle that keyword
///     } else {
///         println!("unknown keyword");
///     }
/// } // atom is dropped here, so it is not kept around in memory
/// ```
///
/// ## Internal representation
///
/// An `Atom` is always 64 bits / 8 bytes.
/// The least-significant two bits form a tag to distinguish three different representations:
///
/// * `0b01`: A short string up to 7 bytes, stored inline in most-significant 56 bits.
///   Bits #4 to #7 (the upper nibble of the lower byte) are the length of the string.
/// * `0b10`: A string part of a statically-known indexed set with [perfect hashing].
///   The most-significant 32 bits are the index in the set.
/// * `0b00`: For other cases, the entire 64 bits are a heap-allocated pointer
///   to an entry in a global hash map.
///   Alignment of the allocation ensures the tag bits are indeed zero.
///   The entry is atomically reference-counted.
///   It is removed from the map and deallocated when its last `Atom` is dropped.
///   The map exists so that interning the same string again gives another pointer to the same entry.
///
/// In all cases, shallow 64-bit equality is equivalent to string equality.
///
/// [perfect hashing]: https://docs.rs/phf/latest/phf/
#[derive(PartialEq, Eq)]
// NOTE: Deriving PartialEq requires that a given string must always be interned the same way.
pub struct Atom<Static> {
    unsafe_data: NonZeroU64,
    phantom: PhantomData<Static>,
}

/// Static and inline atoms don’t have any allocated space. Dynamic atoms do but
/// are accounted for by [`malloc_size_of_dynamic_set`][crate::malloc_size_of_dynamic_set].
#[cfg(feature = "malloc_size_of")]
impl<Static: StaticAtomSet> malloc_size_of::MallocSizeOf for Atom<Static> {
    fn size_of(&self, _ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
        0
    }
}

// FIXME: bound removed from the struct definition before of this error for pack_static:
// "error[E0723]: trait bounds other than `Sized` on const fn parameters are unstable"
// https://github.com/rust-lang/rust/issues/57563
impl<Static> Atom<Static> {
    /// For the atom!() macros
    #[inline(always)]
    #[doc(hidden)]
    pub const fn pack_static(n: u32) -> Self {
        Self {
            unsafe_data: unsafe {
                // STATIC_TAG ensures this is non-zero
                NonZeroU64::new_unchecked((STATIC_TAG as u64) | ((n as u64) << STATIC_SHIFT_BITS))
            },
            phantom: PhantomData,
        }
    }

    /// For the atom!() macros
    #[inline(always)]
    #[doc(hidden)]
    pub const fn pack_inline(mut n: u64, len: u8) -> Self {
        if cfg!(target_endian = "big") {
            // Reverse order of top 7 bytes.
            // Bottom 8 bits of `n` are zero, and we need that to remain so.
            // String data is stored in top 7 bytes, tag and length in bottom byte.
            n = n.to_le() << 8;
        }

        let data: u64 = (INLINE_TAG as u64) | ((len as u64) << LEN_OFFSET) | n;
        Self {
            // INLINE_TAG ensures this is never zero
            unsafe_data: unsafe { NonZeroU64::new_unchecked(data) },
            phantom: PhantomData,
        }
    }

    fn tag(&self) -> u8 {
        (self.unsafe_data.get() & TAG_MASK) as u8
    }

    /// Assuming DYNAMIC_TAG, return the pointer to `Entry`
    fn dynamic_ptr(&self) -> *const Entry {
        std::ptr::with_exposed_provenance(self.unsafe_data.get() as usize)
    }
}

impl<Static: StaticAtomSet> Atom<Static> {
    /// Return the internal representation. For testing.
    #[doc(hidden)]
    pub fn unsafe_data(&self) -> u64 {
        self.unsafe_data.get()
    }

    /// Return true if this is a static Atom. For testing.
    #[doc(hidden)]
    pub fn is_static(&self) -> bool {
        self.tag() == STATIC_TAG
    }

    /// Return true if this is a dynamic Atom. For testing.
    #[doc(hidden)]
    pub fn is_dynamic(&self) -> bool {
        self.tag() == DYNAMIC_TAG
    }

    /// Return true if this is an inline Atom. For testing.
    #[doc(hidden)]
    pub fn is_inline(&self) -> bool {
        self.tag() == INLINE_TAG
    }

    fn static_index(&self) -> u64 {
        self.unsafe_data.get() >> STATIC_SHIFT_BITS
    }

    /// Returns a hash of the string
    ///
    /// For static or dynamic atoms, it is a pre-computed high-quality hash.
    ///
    /// For inline atoms however (short strings 7 bytes or less),
    /// the returned value is the literal inline representation
    /// with string bytes packed directly in the `u64` value,
    /// which makes it a relatively poor-quality hash if used directly.
    pub fn get_hash(&self) -> u64 {
        match self.tag() {
            DYNAMIC_TAG => {
                let entry = self.dynamic_ptr();
                unsafe { (*entry).hash }
            }
            STATIC_TAG => Static::get().hashes[self.static_index() as usize],
            INLINE_TAG => self.unsafe_data.get(),
            _ => unsafe { debug_unreachable!() },
        }
    }

    pub fn try_static(string_to_add: &str) -> Option<Self> {
        Self::try_static_internal(string_to_add).ok()
    }

    fn try_static_internal(string_to_add: &str) -> Result<Self, phf_shared::Hashes> {
        let static_set = Static::get();
        let hash = phf_shared::hash(string_to_add, &static_set.key);
        let index = phf_shared::get_index(&hash, static_set.disps, static_set.atoms.len());

        if static_set.atoms[index as usize] == string_to_add {
            Ok(Self::pack_static(index))
        } else {
            Err(hash)
        }
    }

    /// Get a reference to the underlying str.
    #[inline]
    pub fn as_str(&self) -> &str {
        self // auto-deref
    }

    /// Get a reference to the bytes of the underlying str.
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        self.as_str().as_bytes()
    }
}

impl<Static: StaticAtomSet> Default for Atom<Static> {
    #[inline]
    fn default() -> Self {
        Atom::pack_inline(0, 0)
    }
}

impl<Static: StaticAtomSet> Hash for Atom<Static> {
    #[inline]
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        state.write_u64(self.get_hash())
    }
}

impl<'a, Static: StaticAtomSet> From<Cow<'a, str>> for Atom<Static> {
    fn from(string_to_add: Cow<'a, str>) -> Self {
        let len = string_to_add.len();
        if len <= MAX_INLINE_LEN {
            let mut data: u64 = (INLINE_TAG as u64) | ((len as u64) << LEN_OFFSET);
            {
                let dest = inline_atom_slice_mut(&mut data);
                dest[..len].copy_from_slice(string_to_add.as_bytes());
            }
            Atom {
                // INLINE_TAG ensures this is never zero
                unsafe_data: unsafe { NonZeroU64::new_unchecked(data) },
                phantom: PhantomData,
            }
        } else {
            Self::try_static_internal(&string_to_add).unwrap_or_else(|hash| {
                // Reconstitute 64-bit `Hash128::h1`
                // https://docs.rs/phf_shared/0.14.0/src/phf_shared/lib.rs.html#45-54
                let hash = (hash.g as u64) << 32 | (hash.f1 as u64);
                let ptr: std::ptr::NonNull<Entry> = dynamic_set().insert(string_to_add, hash);
                let data = ptr.as_ptr().expose_provenance() as u64;
                debug_assert!(0 == data & TAG_MASK);
                Atom {
                    // The address of a ptr::NonNull is non-zero
                    unsafe_data: unsafe { NonZeroU64::new_unchecked(data) },
                    phantom: PhantomData,
                }
            })
        }
    }
}

impl<Static: StaticAtomSet> Clone for Atom<Static> {
    #[inline(always)]
    fn clone(&self) -> Self {
        if self.tag() == DYNAMIC_TAG {
            let entry = self.dynamic_ptr();
            // SAFETY: `self` is a valid Atom, meaning its `unsafe_data` points to a live `Entry`
            // kept alive by `self`'s reference count. We can safely dereference it.
            if unsafe { &*entry }.ref_count.fetch_add(1, SeqCst) == isize::MAX {
                std::process::abort();
            }
        }
        Atom { ..*self }
    }
}

impl<Static> Drop for Atom<Static> {
    #[inline]
    fn drop(&mut self) {
        if self.tag() == DYNAMIC_TAG {
            let entry = self.dynamic_ptr();
            if unsafe { &*entry }.ref_count.fetch_sub(1, SeqCst) == 1 {
                drop_slow(self)
            }
        }

        // Out of line to guide inlining.
        #[cold]
        fn drop_slow<Static>(this: &mut Atom<Static>) {
            dynamic_set().remove(this.dynamic_ptr().cast_mut());
        }
    }
}

impl<Static: StaticAtomSet> ops::Deref for Atom<Static> {
    type Target = str;

    #[inline]
    fn deref(&self) -> &str {
        unsafe {
            match self.tag() {
                DYNAMIC_TAG => {
                    let entry = self.dynamic_ptr();
                    &(*entry).string
                }
                INLINE_TAG => {
                    let len = (self.unsafe_data() & LEN_MASK) >> LEN_OFFSET;
                    debug_assert!(len as usize <= MAX_INLINE_LEN);
                    let src = inline_atom_slice(&self.unsafe_data);
                    str::from_utf8_unchecked(src.get_unchecked(..(len as usize)))
                }
                STATIC_TAG => Static::get().atoms[self.static_index() as usize],
                _ => debug_unreachable!(),
            }
        }
    }
}

impl<Static: StaticAtomSet> fmt::Debug for Atom<Static> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let ty_str = unsafe {
            match self.tag() {
                DYNAMIC_TAG => "dynamic",
                INLINE_TAG => "inline",
                STATIC_TAG => "static",
                _ => debug_unreachable!(),
            }
        };

        write!(f, "Atom('{}' type={})", self, ty_str)
    }
}

impl<Static: StaticAtomSet> PartialOrd for Atom<Static> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<Static: StaticAtomSet> Ord for Atom<Static> {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        if self.unsafe_data == other.unsafe_data {
            return Equal;
        }
        self.as_str().cmp(other.as_ref())
    }
}

// AsciiExt requires mutating methods, so we just implement the non-mutating ones.
// We don't need to implement is_ascii because there's no performance improvement
// over the one from &str.
impl<Static: StaticAtomSet> Atom<Static> {
    fn from_mutated_str<F: FnOnce(&mut str)>(s: &str, f: F) -> Self {
        let mut buffer = [const { mem::MaybeUninit::<u8>::uninit() }; 64];

        if let Some(buffer_prefix) = buffer.get_mut(..s.len()) {
            let buffer_ptr = buffer_prefix.as_mut_ptr().cast::<u8>();
            // SAFETY: `buffer_ptr` points to the `MaybeUninit` array.
            // We use `copy_nonoverlapping` to write valid data into it,
            // and then create a slice covering ONLY the initialized portion.
            let as_str = unsafe {
                buffer_ptr.copy_from_nonoverlapping(s.as_ptr(), s.len());
                let buffer_slice = slice::from_raw_parts_mut(buffer_ptr, s.len());
                std::str::from_utf8_unchecked_mut(buffer_slice)
            };
            f(as_str);
            Atom::from(&*as_str)
        } else {
            let mut string = s.to_owned();
            f(&mut string);
            Atom::from(string)
        }
    }

    /// Like [`to_ascii_uppercase`].
    ///
    /// [`to_ascii_uppercase`]: https://doc.rust-lang.org/std/ascii/trait.AsciiExt.html#tymethod.to_ascii_uppercase
    pub fn to_ascii_uppercase(&self) -> Self {
        for (i, b) in self.bytes().enumerate() {
            if let b'a'..=b'z' = b {
                return Atom::from_mutated_str(self, |s| s[i..].make_ascii_uppercase());
            }
        }
        self.clone()
    }

    /// Like [`to_ascii_lowercase`].
    ///
    /// [`to_ascii_lowercase`]: https://doc.rust-lang.org/std/ascii/trait.AsciiExt.html#tymethod.to_ascii_lowercase
    pub fn to_ascii_lowercase(&self) -> Self {
        for (i, b) in self.bytes().enumerate() {
            if let b'A'..=b'Z' = b {
                return Atom::from_mutated_str(self, |s| s[i..].make_ascii_lowercase());
            }
        }
        self.clone()
    }

    /// Like [`eq_ignore_ascii_case`].
    ///
    /// [`eq_ignore_ascii_case`]: https://doc.rust-lang.org/std/ascii/trait.AsciiExt.html#tymethod.eq_ignore_ascii_case
    pub fn eq_ignore_ascii_case(&self, other: &Self) -> bool {
        (self == other) || self.eq_str_ignore_ascii_case(other)
    }

    /// Like [`eq_ignore_ascii_case`], but takes an unhashed string as `other`.
    ///
    /// [`eq_ignore_ascii_case`]: https://doc.rust-lang.org/std/ascii/trait.AsciiExt.html#tymethod.eq_ignore_ascii_case
    pub fn eq_str_ignore_ascii_case(&self, other: &str) -> bool {
        self.as_str().eq_ignore_ascii_case(other)
    }
}

#[inline(always)]
fn inline_atom_slice(x: &NonZeroU64) -> &[u8] {
    let x: *const NonZeroU64 = x;
    let mut data = x as *const u8;
    // All except the lowest byte, which is first in little-endian, last in big-endian.
    if cfg!(target_endian = "little") {
        data = unsafe { data.offset(1) };
    }
    let len = 7;
    unsafe { slice::from_raw_parts(data, len) }
}

#[inline(always)]
fn inline_atom_slice_mut(x: &mut u64) -> &mut [u8] {
    let x: *mut u64 = x;
    let mut data = x as *mut u8;
    // All except the lowest byte, which is first in little-endian, last in big-endian.
    if cfg!(target_endian = "little") {
        data = unsafe { data.offset(1) };
    }
    let len = 7;
    unsafe { slice::from_raw_parts_mut(data, len) }
}