string_cache/atom.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 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388
// 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;
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
/// ```
#[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>,
}
// 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,
}
}
fn tag(&self) -> u8 {
(self.unsafe_data.get() & TAG_MASK) as u8
}
}
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
}
/// Get the hash of the string as it is stored in the set.
pub fn get_hash(&self) -> u32 {
match self.tag() {
DYNAMIC_TAG => {
let entry = self.unsafe_data.get() as *const Entry;
unsafe { (*entry).hash }
}
STATIC_TAG => Static::get().hashes[self.static_index() as usize],
INLINE_TAG => {
let data = self.unsafe_data.get();
// This may or may not be great...
((data >> 32) ^ data) as u32
}
_ => 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)
}
}
}
impl<Static: StaticAtomSet> Default for Atom<Static> {
#[inline]
fn default() -> Self {
Atom::pack_static(Static::empty_string_index())
}
}
impl<Static: StaticAtomSet> Hash for Atom<Static> {
#[inline]
fn hash<H>(&self, state: &mut H)
where
H: Hasher,
{
state.write_u32(self.get_hash())
}
}
impl<'a, Static: StaticAtomSet> From<Cow<'a, str>> for Atom<Static> {
fn from(string_to_add: Cow<'a, str>) -> Self {
Self::try_static_internal(&*string_to_add).unwrap_or_else(|hash| {
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 {
let ptr: std::ptr::NonNull<Entry> = DYNAMIC_SET.insert(string_to_add, hash.g);
let data = ptr.as_ptr() 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.unsafe_data.get() as *const Entry;
unsafe { &*entry }.ref_count.fetch_add(1, SeqCst);
}
Atom { ..*self }
}
}
impl<Static> Drop for Atom<Static> {
#[inline]
fn drop(&mut self) {
if self.tag() == DYNAMIC_TAG {
let entry = self.unsafe_data.get() as *const Entry;
if unsafe { &*entry }.ref_count.fetch_sub(1, SeqCst) == 1 {
drop_slow(self)
}
}
// Out of line to guide inlining.
fn drop_slow<Static>(this: &mut Atom<Static>) {
DYNAMIC_SET.remove(this.unsafe_data.get() as *mut Entry);
}
}
}
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.unsafe_data.get() as *const Entry;
&(*entry).string
}
INLINE_TAG => {
let len = (self.unsafe_data() & LEN_MASK) >> LEN_OFFSET;
let src = inline_atom_slice(&self.unsafe_data);
str::from_utf8_unchecked(&src[..(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> {
if self.unsafe_data == other.unsafe_data {
return Some(Equal);
}
self.as_ref().partial_cmp(other.as_ref())
}
}
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_ref().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 = mem::MaybeUninit::<[u8; 64]>::uninit();
let buffer = unsafe { &mut *buffer.as_mut_ptr() };
if let Some(buffer_prefix) = buffer.get_mut(..s.len()) {
buffer_prefix.copy_from_slice(s.as_bytes());
let as_str = unsafe { ::std::str::from_utf8_unchecked_mut(buffer_prefix) };
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).eq_ignore_ascii_case(other)
}
}
#[inline(always)]
fn inline_atom_slice(x: &NonZeroU64) -> &[u8] {
unsafe {
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 = data.offset(1);
}
let len = 7;
slice::from_raw_parts(data, len)
}
}
#[inline(always)]
fn inline_atom_slice_mut(x: &mut u64) -> &mut [u8] {
unsafe {
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 = data.offset(1);
}
let len = 7;
slice::from_raw_parts_mut(data, len)
}
}