use std::cell::Cell;
use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::fmt::Formatter;
use std::ptr::null;
type NumIndex = u32;
#[derive(Debug, Default)]
pub struct AtomTable(UnsafeCell<Inner>);
#[derive(Default)]
struct Inner {
strings: Vec<String>,
map: HashMap<&'static str, NumIndex>,
strings_u16: Vec<Vec<u16>>,
map_u16: HashMap<&'static [u16], NumIndex>,
strings_bytes: Vec<Vec<u8>>,
map_bytes: HashMap<&'static [u8], NumIndex>,
converted_bytes: HashMap<AtomBytes, Converted>,
}
#[derive(Copy, Clone, Eq, PartialEq, Hash)]
pub struct Atom(NumIndex);
#[derive(Copy, Clone, Eq, PartialEq, Hash)]
pub struct AtomU16(NumIndex);
#[derive(Copy, Clone, Eq, PartialEq, Hash)]
pub struct AtomBytes(NumIndex);
thread_local! {
static DEBUG_TABLE: Cell<* const AtomTable> = Cell::new(null());
}
impl std::fmt::Debug for Atom {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let mut t = f.debug_tuple("Atom");
t.field(&self.0);
DEBUG_TABLE.with(|debug_table| {
let p = debug_table.get();
if let Some(r) = unsafe { p.as_ref() } {
if let Some(value) = r.try_str(*self) {
t.field(&value);
}
}
});
t.finish()
}
}
impl std::fmt::Debug for AtomU16 {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let mut t = f.debug_tuple("Atom");
t.field(&self.0);
DEBUG_TABLE.with(|debug_table| {
let p = debug_table.get();
if let Some(r) = unsafe { p.as_ref() } {
if let Some(value) = r.try_str_u16(*self) {
t.field(&value);
}
}
});
t.finish()
}
}
impl std::fmt::Debug for AtomBytes {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let mut t = f.debug_tuple("AtomBytes");
t.field(&self.0);
DEBUG_TABLE.with(|debug_table| {
let p = debug_table.get();
if let Some(r) = unsafe { p.as_ref() } {
if let Some(value) = r.try_bytes(*self) {
t.field(&value);
}
}
});
t.finish()
}
}
pub const INVALID_ATOM: Atom = Atom(NumIndex::MAX);
pub const INVALID_ATOM_BYTES: AtomBytes = AtomBytes(NumIndex::MAX);
impl Inner {
fn add_atom<V: Into<String> + AsRef<str>>(&mut self, value: V) -> Atom {
if let Some(index) = self.map.get(value.as_ref()) {
return Atom(*index);
}
self.add(value.into())
}
fn add(&mut self, owned: String) -> Atom {
let index = self.strings.len();
assert!(index < INVALID_ATOM.0 as usize, "More than 4GB atoms?");
let key: *const str = owned.as_str();
self.strings.push(owned);
self.map.insert(unsafe { &*key }, index as NumIndex);
Atom(index as NumIndex)
}
#[inline]
fn str(&self, ident: Atom) -> &str {
self.strings[ident.0 as usize].as_str()
}
fn try_str(&self, ident: Atom) -> Option<&str> {
if (ident.0 as usize) < self.strings.len() {
Some(self.str(ident))
} else {
None
}
}
fn add_atom_u16<V: Into<Vec<u16>> + AsRef<[u16]>>(&mut self, value: V) -> AtomU16 {
if let Some(index) = self.map_u16.get(value.as_ref()) {
return AtomU16(*index);
}
self.add_u16(value.into())
}
fn add_u16(&mut self, owned: Vec<u16>) -> AtomU16 {
let index = self.strings_u16.len();
assert!(index < INVALID_ATOM.0 as usize, "More than 4GB atoms?");
let key: *const [u16] = owned.as_slice();
self.strings_u16.push(owned);
self.map_u16.insert(unsafe { &*key }, index as NumIndex);
AtomU16(index as NumIndex)
}
#[inline]
fn str_u16(&self, ident: AtomU16) -> &[u16] {
self.strings_u16[ident.0 as usize].as_slice()
}
fn try_str_u16(&self, ident: AtomU16) -> Option<&[u16]> {
if (ident.0 as usize) < self.strings_u16.len() {
Some(self.str_u16(ident))
} else {
None
}
}
fn add_atom_bytes<V: Into<Vec<u8>> + AsRef<[u8]>>(&mut self, value: V) -> AtomBytes {
if let Some(index) = self.map_bytes.get(value.as_ref()) {
return AtomBytes(*index);
}
self.add_bytes(value.into())
}
fn add_bytes(&mut self, owned: Vec<u8>) -> AtomBytes {
let index = self.strings_bytes.len();
assert!(index < INVALID_ATOM_BYTES.0 as usize, "More than 4GB atoms?");
let key: *const [u8] = owned.as_slice();
self.strings_bytes.push(owned);
self.map_bytes.insert(unsafe { &*key }, index as NumIndex);
AtomBytes(index as NumIndex)
}
#[inline]
fn bytes(&self, ident: AtomBytes) -> &[u8] {
self.strings_bytes[ident.0 as usize].as_slice()
}
fn try_bytes(&self, ident: AtomBytes) -> Option<&[u8]> {
if (ident.0 as usize) < self.strings_bytes.len() {
Some(self.bytes(ident))
} else {
None
}
}
fn ensure_converted(&mut self, ident: AtomBytes) {
if !self.converted_bytes.contains_key(&ident) {
let converted = convert_wtf8(self.bytes(ident));
self.converted_bytes.insert(ident, converted);
}
}
#[inline]
fn converted(&self, ident: AtomBytes) -> &Converted {
&self.converted_bytes[&ident]
}
}
struct Converted {
text: String,
replaced: bool,
}
#[inline]
fn surrogate_at(bytes: &[u8]) -> Option<u32> {
match bytes {
[0xED, b1 @ 0xA0..=0xBF, b2 @ 0x80..=0xBF, ..] => {
Some(0xD000 | ((*b1 as u32 & 0x3F) << 6) | (*b2 as u32 & 0x3F))
}
_ => None,
}
}
fn convert_wtf8(bytes: &[u8]) -> Converted {
let mut out = String::with_capacity(bytes.len());
let mut replaced = false;
let mut rest = bytes;
loop {
let err = match std::str::from_utf8(rest) {
Ok(valid) => {
out.push_str(valid);
return Converted {
text: out,
replaced,
};
}
Err(err) => err,
};
let (valid, invalid) = rest.split_at(err.valid_up_to());
out.push_str(std::str::from_utf8(valid).unwrap());
rest = match surrogate_at(invalid) {
Some(high) if high < 0xDC00 => match surrogate_at(&invalid[3..]) {
Some(low) if low >= 0xDC00 => {
let cp = 0x10000 + ((high - 0xD800) << 10) + (low - 0xDC00);
out.push(char::from_u32(cp).unwrap_or(char::REPLACEMENT_CHARACTER));
&invalid[6..]
}
_ => {
out.push(char::REPLACEMENT_CHARACTER);
replaced = true;
&invalid[3..]
}
},
Some(_) => {
out.push(char::REPLACEMENT_CHARACTER);
replaced = true;
&invalid[3..]
}
None => {
out.push(char::REPLACEMENT_CHARACTER);
replaced = true;
&invalid[err.error_len().unwrap_or(invalid.len())..]
}
};
}
}
impl AtomTable {
pub fn new() -> AtomTable {
Default::default()
}
pub fn atom<V: Into<String> + AsRef<str>>(&self, value: V) -> Atom {
unsafe { &mut *self.0.get() }.add_atom(value)
}
#[inline]
pub fn str(&self, ident: Atom) -> &str {
unsafe { &*self.0.get() }.str(ident)
}
#[inline]
pub fn try_str(&self, ident: Atom) -> Option<&str> {
unsafe { &*self.0.get() }.try_str(ident)
}
pub fn atom_u16<V: Into<Vec<u16>> + AsRef<[u16]>>(&self, value: V) -> AtomU16 {
unsafe { &mut *self.0.get() }.add_atom_u16(value)
}
#[inline]
pub fn str_u16(&self, ident: AtomU16) -> &[u16] {
unsafe { &*self.0.get() }.str_u16(ident)
}
#[inline]
pub fn try_str_u16(&self, ident: AtomU16) -> Option<&[u16]> {
unsafe { &*self.0.get() }.try_str_u16(ident)
}
pub fn atom_bytes<V: Into<Vec<u8>> + AsRef<[u8]>>(&self, value: V) -> AtomBytes {
unsafe { &mut *self.0.get() }.add_atom_bytes(value)
}
#[inline]
pub fn bytes(&self, ident: AtomBytes) -> &[u8] {
unsafe { &*self.0.get() }.bytes(ident)
}
#[inline]
pub fn try_bytes(&self, ident: AtomBytes) -> Option<&[u8]> {
unsafe { &*self.0.get() }.try_bytes(ident)
}
#[inline]
pub fn bytes_str_lossy(&self, ident: AtomBytes) -> &str {
if let Ok(s) = std::str::from_utf8(unsafe { &*self.0.get() }.bytes(ident)) {
return s;
}
unsafe { &mut *self.0.get() }.ensure_converted(ident);
unsafe { &*self.0.get() }.converted(ident).text.as_str()
}
#[inline]
pub fn try_bytes_str(&self, ident: AtomBytes) -> Option<&str> {
let bytes = unsafe { &*self.0.get() }.try_bytes(ident)?;
if let Ok(s) = std::str::from_utf8(bytes) {
return Some(s);
}
unsafe { &mut *self.0.get() }.ensure_converted(ident);
let converted = unsafe { &*self.0.get() }.converted(ident);
if converted.replaced {
None
} else {
Some(converted.text.as_str())
}
}
pub fn in_debug_context<R, F: FnOnce() -> R>(&self, f: F) -> R {
DEBUG_TABLE.with(|debug_table| {
let prev_table = debug_table.replace(self);
let res = f();
debug_assert!(
debug_table.get() == self,
"debug context unexpectedly changed"
);
debug_table.set(prev_table);
res
})
}
pub unsafe fn unsafe_set_debug_context(ptr: *const Self) -> *const Self {
DEBUG_TABLE.with(|debug_table| debug_table.replace(ptr))
}
}
impl std::ops::Index<Atom> for AtomTable {
type Output = str;
fn index(&self, index: Atom) -> &Self::Output {
self.str(index)
}
}
impl std::ops::Index<AtomU16> for AtomTable {
type Output = [u16];
fn index(&self, index: AtomU16) -> &Self::Output {
self.str_u16(index)
}
}
impl std::ops::Index<AtomBytes> for AtomTable {
type Output = [u8];
fn index(&self, index: AtomBytes) -> &Self::Output {
self.bytes(index)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tab() {
let idtab = AtomTable::new();
let id_foo = idtab.atom("foo");
let p_foo: *const str = idtab.str(id_foo);
let id_bar = idtab.atom("bar");
assert_ne!(id_foo, id_bar);
assert_eq!(idtab.atom("foo"), id_foo);
assert_eq!(idtab.atom("bar"), id_bar);
assert_eq!(idtab.atom(String::from("foo")), id_foo);
assert_eq!(idtab.atom(String::from("bar")), id_bar);
assert_eq!(idtab.str(id_foo), "foo");
assert_eq!(idtab.str(id_bar), "bar");
assert_eq!(idtab.str(id_foo) as *const str, p_foo);
}
#[test]
fn test_bytes() {
let tab = AtomTable::new();
let foo = tab.atom_bytes(b"foo".as_slice());
let bar = tab.atom_bytes(b"bar".as_slice());
assert_ne!(foo, bar);
assert_eq!(tab.atom_bytes(b"foo".as_slice()), foo);
assert_eq!(tab.atom_bytes(Vec::from(*b"bar")), bar);
assert_eq!(tab.bytes(foo), b"foo");
assert_eq!(&tab[bar], b"bar");
let p_foo: *const [u8] = tab.bytes(foo);
let _ = tab.atom_bytes(b"baz".as_slice());
assert_eq!(tab.bytes(foo) as *const [u8], p_foo);
}
#[test]
fn test_bytes_ill_formed_utf8() {
let tab = AtomTable::new();
let lone_surrogate: &[u8] = &[0xed, 0xa0, 0x80];
let a = tab.atom_bytes(lone_surrogate);
assert_eq!(tab.bytes(a), lone_surrogate);
assert_eq!(tab.atom_bytes(lone_surrogate), a);
let s = tab.atom("foo");
let b = tab.atom_bytes(b"foo".as_slice());
assert_eq!(tab.str(s), "foo");
assert_eq!(tab.bytes(b), b"foo");
}
#[test]
fn test_bytes_try_and_invalid() {
let tab = AtomTable::new();
let a = tab.atom_bytes(b"x".as_slice());
assert_eq!(tab.try_bytes(a), Some(b"x".as_slice()));
assert_eq!(tab.try_bytes(INVALID_ATOM_BYTES), None);
}
#[test]
fn lone_surrogate_becomes_exactly_one_replacement_char() {
let t = AtomTable::new();
let a = t.atom_bytes(vec![0xED, 0xA0, 0x80]);
assert_eq!(t.try_bytes_str(a), None);
let s = t.bytes_str_lossy(a);
assert_eq!(
s.chars().filter(|c| *c == '\u{FFFD}').count(),
1,
"std::from_utf8_lossy would give 3 here; we must be WTF-8 aware"
);
assert_eq!(s, "\u{FFFD}");
}
#[test]
fn valid_utf8_is_borrowed_unchanged() {
let t = AtomTable::new();
let a = t.atom_bytes("greet".as_bytes().to_vec());
assert_eq!(t.try_bytes_str(a), Some("greet"));
assert_eq!(t.bytes_str_lossy(a), "greet");
assert_eq!(t.bytes_str_lossy(a).as_ptr(), t.bytes(a).as_ptr());
}
#[test]
fn surrogates_mixed_with_text_replace_only_the_surrogate() {
let t = AtomTable::new();
let mut v = b"a".to_vec();
v.extend_from_slice(&[0xED, 0xA0, 0x80]);
v.extend_from_slice("b".as_bytes());
let a = t.atom_bytes(v);
assert_eq!(t.bytes_str_lossy(a), "a\u{FFFD}b");
}
fn anchored(t: &AtomTable) -> usize {
unsafe { &*t.0.get() }.converted_bytes.len()
}
#[test]
fn surrogate_pair_folds_into_the_astral_char() {
let t = AtomTable::new();
let a = t.atom_bytes(vec![0xED, 0xA0, 0xBD, 0xED, 0xB8, 0x80]);
assert_eq!(t.try_bytes_str(a), Some("\u{1F600}"));
assert_eq!(t.bytes_str_lossy(a), "\u{1F600}");
assert_eq!(
t.try_bytes_str(a).unwrap().as_ptr(),
t.bytes_str_lossy(a).as_ptr()
);
assert_eq!(anchored(&t), 1);
}
#[test]
fn a_pair_beside_an_unpaired_surrogate_is_not_representable() {
let t = AtomTable::new();
let mut v = vec![0xED, 0xA0, 0xBD, 0xED, 0xB8, 0x80]; v.extend_from_slice(&[0xED, 0xA0, 0x80]); v.extend_from_slice(b"!");
let a = t.atom_bytes(v);
assert_eq!(t.try_bytes_str(a), None);
let s = t.bytes_str_lossy(a);
assert_eq!(s, "\u{1F600}\u{FFFD}!");
assert_eq!(s.chars().filter(|c| *c == '\u{FFFD}').count(), 1);
}
#[test]
fn the_valid_path_never_anchors() {
let t = AtomTable::new();
let a = t.atom_bytes(b"plain".as_slice());
assert_eq!(t.try_bytes_str(a), Some("plain"));
assert_eq!(t.try_bytes_str(a).unwrap().as_ptr(), t.bytes(a).as_ptr());
assert_eq!(t.bytes_str_lossy(a).as_ptr(), t.bytes(a).as_ptr());
assert_eq!(anchored(&t), 0);
}
#[test]
fn the_folded_result_is_stable_across_calls() {
let t = AtomTable::new();
let a = t.atom_bytes(vec![0xED, 0xA0, 0xBD, 0xED, 0xB8, 0x80]);
let first: &str = t.try_bytes_str(a).unwrap();
for i in 0..100u8 {
let b = t.atom_bytes(vec![0xED, 0xA0, 0xBD, 0xED, 0xB8, 0x80, i]);
assert!(t.try_bytes_str(b).unwrap().starts_with('\u{1F600}'));
}
assert_eq!(first, "\u{1F600}");
assert_eq!(t.try_bytes_str(a).unwrap().as_ptr(), first.as_ptr());
}
#[test]
fn unpaired_surrogates_are_one_replacement_each() {
let t = AtomTable::new();
let cases: &[(&[u8], &str)] = &[
(&[0xED, 0xA0, 0x80], "\u{FFFD}"),
(&[0xED, 0xA0, 0x80, 0xED, 0xA0, 0x80], "\u{FFFD}\u{FFFD}"),
(&[0xED, 0xB0, 0x80, 0xED, 0xA0, 0x80], "\u{FFFD}\u{FFFD}"),
(&[0xED, 0xB0, 0x80, 0xED, 0xB0, 0x80], "\u{FFFD}\u{FFFD}"),
(&[0xED, 0xB0, 0x80, b'z'], "\u{FFFD}z"),
(&[0xED, 0xA0, 0x80, b'z'], "\u{FFFD}z"),
];
for (bytes, expected) in cases {
let a = t.atom_bytes(*bytes);
assert_eq!(t.bytes_str_lossy(a), *expected, "bytes {bytes:02X?}");
assert_eq!(t.try_bytes_str(a), None, "bytes {bytes:02X?}");
}
}
#[test]
fn non_surrogate_garbage_is_replaced_per_ill_formed_sequence() {
let t = AtomTable::new();
let a = t.atom_bytes(vec![0xFF, 0xFE]);
assert_eq!(t.bytes_str_lossy(a), "\u{FFFD}\u{FFFD}");
let b = t.atom_bytes(vec![b'x', 0xE2, 0x82]);
assert_eq!(t.bytes_str_lossy(b), "x\u{FFFD}");
let c = t.atom_bytes(vec![b'o', 0x80, b'k']);
assert_eq!(t.bytes_str_lossy(c), "o\u{FFFD}k");
}
#[test]
fn an_earlier_lossy_str_survives_later_conversions() {
let t = AtomTable::new();
let a = t.atom_bytes(vec![0xED, 0xA0, 0x80]);
let first: &str = t.bytes_str_lossy(a);
for i in 0..100u8 {
let b = t.atom_bytes(vec![0xED, 0xA0, 0x80, b'a' + i % 26, i]);
assert!(t.bytes_str_lossy(b).starts_with('\u{FFFD}'));
}
assert_eq!(first, "\u{FFFD}");
}
#[test]
fn try_bytes_str_rejects_invalid_atoms() {
let t = AtomTable::new();
assert_eq!(t.try_bytes_str(INVALID_ATOM_BYTES), None);
}
#[test]
fn the_lossy_result_is_stable_across_calls() {
let t = AtomTable::new();
let a = t.atom_bytes(vec![0xED, 0xA0, 0x80]);
let p1 = t.bytes_str_lossy(a).as_ptr();
for i in 0..1000 {
t.atom_bytes(format!("filler{i}").into_bytes());
}
let p2 = t.bytes_str_lossy(a).as_ptr();
assert_eq!(p1, p2, "the anchored String must not be rebuilt or moved");
}
}