use rustc_hash::FxHashMap;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::cell::Cell;
use std::hash::{Hash, Hasher};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InternerCollision {
pub key: u64,
pub existing: String,
pub conflicting: String,
}
impl std::fmt::Display for InternerCollision {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"InternedKey hash collision at {}: '{}' conflicts with '{}'",
self.key, self.conflicting, self.existing
)
}
}
impl std::error::Error for InternerCollision {}
#[repr(transparent)]
#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
pub struct InternedKey(u64);
impl InternedKey {
#[inline]
pub fn from_str(s: &str) -> Self {
const FNV_OFFSET: u64 = 0xcbf29ce484222325;
const FNV_PRIME: u64 = 0x100000001b3;
let mut h = FNV_OFFSET;
for &byte in s.as_bytes() {
h ^= byte as u64;
h = h.wrapping_mul(FNV_PRIME);
}
InternedKey(h)
}
#[inline]
pub fn as_u64(&self) -> u64 {
self.0
}
#[inline]
pub fn from_u64(v: u64) -> Self {
InternedKey(v)
}
}
impl Hash for InternedKey {
#[inline]
fn hash<H: Hasher>(&self, state: &mut H) {
state.write_u64(self.0);
}
}
impl Serialize for InternedKey {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
SERIALIZE_INTERNER.with(|cell| {
let ptr = cell
.get()
.expect("BUG: SERIALIZE_INTERNER not set during InternedKey serialization");
let interner = unsafe { &*ptr };
interner.resolve(*self).serialize(serializer)
})
}
}
impl<'de> Deserialize<'de> for InternedKey {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
struct KeyVisitor;
impl<'de> serde::de::Visitor<'de> for KeyVisitor {
type Value = InternedKey;
fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
f.write_str("a string key")
}
fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
let key = InternedKey::from_str(v);
DESERIALIZE_INTERNER.with(|cell| -> Result<(), E> {
if let Some(ptr) = cell.get() {
let interner = unsafe { &mut *ptr };
interner.try_register(key, v).map_err(E::custom)?;
}
Ok(())
})?;
Ok(key)
}
fn visit_string<E: serde::de::Error>(self, v: String) -> Result<Self::Value, E> {
self.visit_str(&v)
}
}
deserializer.deserialize_str(KeyVisitor)
}
}
#[derive(Debug, Clone, Default)]
pub struct StringInterner {
strings: FxHashMap<InternedKey, String>,
}
impl StringInterner {
pub fn new() -> Self {
Self::default()
}
#[inline]
pub fn try_register(&mut self, key: InternedKey, s: &str) -> Result<(), InternerCollision> {
if let Some(existing) = self.strings.get(&key) {
if existing != s {
return Err(InternerCollision {
key: key.as_u64(),
existing: existing.clone(),
conflicting: s.to_string(),
});
}
return Ok(());
}
self.strings.insert(key, s.to_string());
Ok(())
}
#[inline]
pub fn try_get_or_intern(&mut self, s: &str) -> Result<InternedKey, InternerCollision> {
let key = InternedKey::from_str(s);
self.try_register(key, s)?;
Ok(key)
}
#[inline]
pub(crate) fn get_or_intern(&mut self, s: &str) -> InternedKey {
self.try_get_or_intern(s)
.unwrap_or_else(|collision| panic!("{collision}"))
}
pub fn validate_names<'a>(
&self,
names: impl IntoIterator<Item = &'a str>,
) -> Result<Vec<InternedKey>, InternerCollision> {
let names = names.into_iter();
let (lower, _) = names.size_hint();
let mut keys = Vec::with_capacity(lower);
let mut staged: FxHashMap<InternedKey, &'a str> = FxHashMap::default();
for name in names {
let key = InternedKey::from_str(name);
if let Some(existing) = self.strings.get(&key).map(String::as_str) {
if existing != name {
return Err(InternerCollision {
key: key.as_u64(),
existing: existing.to_string(),
conflicting: name.to_string(),
});
}
} else if let Some(existing) = staged.get(&key).copied() {
if existing != name {
return Err(InternerCollision {
key: key.as_u64(),
existing: existing.to_string(),
conflicting: name.to_string(),
});
}
} else {
staged.insert(key, name);
}
keys.push(key);
}
Ok(keys)
}
#[inline]
pub fn resolve(&self, key: InternedKey) -> &str {
self.strings
.get(&key)
.map(|s| s.as_str())
.unwrap_or_else(|| {
panic!(
"InternedKey {} not found in StringInterner ({} entries)",
key.as_u64(),
self.strings.len()
)
})
}
pub fn iter(&self) -> impl Iterator<Item = (InternedKey, &str)> {
self.strings.iter().map(|(&k, v)| (k, v.as_str()))
}
#[inline]
pub fn try_resolve(&self, key: InternedKey) -> Option<&str> {
self.strings.get(&key).map(|s| s.as_str())
}
#[inline]
pub fn try_resolve_to_key(&self, s: &str) -> Option<InternedKey> {
let key = InternedKey::from_str(s);
if self.strings.get(&key).is_some_and(|existing| existing == s) {
Some(key)
} else {
None
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn raw_collision_is_typed_and_keeps_first_mapping() {
let mut interner = StringInterner::new();
let key = InternedKey::from_str("second");
interner.try_register(key, "first").unwrap();
let err = interner.try_register(key, "second").unwrap_err();
assert_eq!(err.key, key.as_u64());
assert_eq!(err.existing, "first");
assert_eq!(err.conflicting, "second");
assert_eq!(interner.resolve(key), "first");
assert_eq!(interner.try_resolve_to_key("second"), None);
}
#[test]
fn batch_validation_detects_internal_collision_without_mutation() {
let mut interner = StringInterner::new();
let key = InternedKey::from_str("incoming");
interner
.try_register(key, "existing-with-forced-key")
.unwrap();
let before: Vec<_> = interner.iter().map(|(k, v)| (k, v.to_string())).collect();
assert!(interner.validate_names(["ordinary", "incoming"]).is_err());
let after: Vec<_> = interner.iter().map(|(k, v)| (k, v.to_string())).collect();
assert_eq!(after, before);
}
#[test]
fn serde_collision_is_a_deserializer_error() {
let incoming = "persisted-name";
let mut interner = StringInterner::new();
interner
.try_register(InternedKey::from_str(incoming), "conflicting-existing")
.unwrap();
let bytes = bincode::serialize(incoming).unwrap();
let guard = SerdeDeserializeGuard::new(&mut interner);
let decoded = bincode::deserialize::<InternedKey>(&bytes);
drop(guard);
assert!(decoded.unwrap_err().to_string().contains("hash collision"));
assert_eq!(
interner.resolve(InternedKey::from_str(incoming)),
"conflicting-existing"
);
}
}
thread_local! {
static SERIALIZE_INTERNER: Cell<Option<*const StringInterner>> = const { Cell::new(None) };
static DESERIALIZE_INTERNER: Cell<Option<*mut StringInterner>> = const { Cell::new(None) };
pub(crate) static STRIP_PROPERTIES: Cell<bool> = const { Cell::new(false) };
}
pub(crate) struct SerdeSerializeGuard<'a> {
_phantom: std::marker::PhantomData<&'a StringInterner>,
}
impl<'a> SerdeSerializeGuard<'a> {
pub fn new(interner: &'a StringInterner) -> Self {
SERIALIZE_INTERNER.with(|cell| cell.set(Some(interner as *const StringInterner)));
SerdeSerializeGuard {
_phantom: std::marker::PhantomData,
}
}
}
impl Drop for SerdeSerializeGuard<'_> {
fn drop(&mut self) {
SERIALIZE_INTERNER.with(|cell| cell.set(None));
}
}
pub(crate) struct SerdeDeserializeGuard<'a> {
_phantom: std::marker::PhantomData<&'a mut StringInterner>,
}
impl<'a> SerdeDeserializeGuard<'a> {
pub fn new(interner: &'a mut StringInterner) -> Self {
DESERIALIZE_INTERNER.with(|cell| cell.set(Some(interner as *mut StringInterner)));
SerdeDeserializeGuard {
_phantom: std::marker::PhantomData,
}
}
}
impl Drop for SerdeDeserializeGuard<'_> {
fn drop(&mut self) {
DESERIALIZE_INTERNER.with(|cell| cell.set(None));
}
}
pub(crate) struct StripPropertiesGuard;
impl StripPropertiesGuard {
pub fn new() -> Self {
STRIP_PROPERTIES.with(|cell| cell.set(true));
StripPropertiesGuard
}
}
impl Drop for StripPropertiesGuard {
fn drop(&mut self) {
STRIP_PROPERTIES.with(|cell| cell.set(false));
}
}