hermes-atom-table 0.1.0

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Ported from juno atom_table for the Hermes Rust lexer; carries encapsulated
//! unsafe; adds a byte/WTF-8 intern path.

use std::cell::Cell;
use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::fmt::Formatter;
use std::ptr::null;

/// Type used to hold a string index internally.
type NumIndex = u32;

/// A string uniquing table - only one copy of a string is stored and all attempts
/// to add the same string again return the same atom. This table is intended to
/// be easily shareable, so it utilizes interior mutability. UnsafeCell<> is safe
/// because we never allow reference to it to escape.
#[derive(Debug, Default)]
pub struct AtomTable(UnsafeCell<Inner>);

/// A string uniquing table - only one copy of a string is stored and all attempts
/// to add the same string again return the same atom.
#[derive(Default)]
struct Inner {
    /// Strings are added here and never removed or mutated.
    strings: Vec<String>,
    /// Maps from a reference inside [`Inner::strings`] to the index in [`Inner::strings`].
    /// Since strings are never removed or modified, the lifetime of the key
    /// is effectively static.
    map: HashMap<&'static str, NumIndex>,

    /// Strings are added here and never removed or mutated.
    strings_u16: Vec<Vec<u16>>,
    /// Maps from a reference inside [`Inner::strings_u16`] to the index in [`Inner::strings_u16`].
    /// Since strings are never removed or modified, the lifetime of the key
    /// is effectively static.
    map_u16: HashMap<&'static [u16], NumIndex>,

    /// Byte strings are added here and never removed or mutated.
    /// The bytes need not be valid UTF-8 (they may be WTF-8 or arbitrary byte
    /// sequences, e.g. JS string literals containing lone surrogates).
    strings_bytes: Vec<Vec<u8>>,
    /// Maps from a reference inside [`Inner::strings_bytes`] to the index in
    /// [`Inner::strings_bytes`]. Since strings are never removed or modified,
    /// the lifetime of the key is effectively static.
    map_bytes: HashMap<&'static [u8], NumIndex>,
}

/// This represents a unique string index in the table.
#[derive(Copy, Clone, Eq, PartialEq, Hash)]
pub struct Atom(NumIndex);

/// This represents a unique string index in the table.
#[derive(Copy, Clone, Eq, PartialEq, Hash)]
pub struct AtomU16(NumIndex);

/// This represents a unique byte-string index in the table.
/// The bytes need not be valid UTF-8.
#[derive(Copy, Clone, Eq, PartialEq, Hash)]
pub struct AtomBytes(NumIndex);

thread_local! {
    /// Stores the active table used for debug formatting.
    static DEBUG_TABLE: Cell<* const AtomTable> = Cell::new(null());
}

// An implementation of Debug which optionally obtains the Atom value from the
// active debug map.
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);

        // If the debug table is set and the atom is valid in it, add the value
        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()
    }
}

// An implementation of Debug which optionally obtains the Atom value from the
// active debug map.
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);

        // If the debug table is set and the atom is valid in it, add the value
        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()
    }
}

// An implementation of Debug which optionally obtains the AtomBytes value from
// the active debug map.
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);

        // If the debug table is set and the atom is valid in it, add the value
        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()
    }
}

/// A special value reserved for the invalid atom.
pub const INVALID_ATOM: Atom = Atom(NumIndex::MAX);

/// A special value reserved for the invalid atom bytes.
pub const INVALID_ATOM_BYTES: AtomBytes = AtomBytes(NumIndex::MAX);

impl Inner {
    /// Add a string to the table and return its atom index. The same
    /// string always returns the same index.
    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())
    }

    /// Perform the actual addition of the owned string.
    fn add(&mut self, owned: String) -> Atom {
        // Remember the index of the new element.
        let index = self.strings.len();
        assert!(index < INVALID_ATOM.0 as usize, "More than 4GB atoms?");

        // Obtain a reference to the existing string on the heap. That reference
        // is valid while `self` is valid.
        let key: *const str = owned.as_str();

        // Push the new string.
        self.strings.push(owned);

        self.map.insert(unsafe { &*key }, index as NumIndex);
        Atom(index as NumIndex)
    }

    /// Return the contents of the specified atom.
    #[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
        }
    }

    /// Add a string to the table and return its atom index. The same
    /// string always returns the same index.
    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())
    }

    /// Perform the actual addition of the owned string.
    fn add_u16(&mut self, owned: Vec<u16>) -> AtomU16 {
        // Remember the index of the new element.
        let index = self.strings_u16.len();
        assert!(index < INVALID_ATOM.0 as usize, "More than 4GB atoms?");

        // Obtain a reference to the existing string on the heap. That reference
        // is valid while `self` is valid.
        let key: *const [u16] = owned.as_slice();

        // Push the new string.
        self.strings_u16.push(owned);

        self.map_u16.insert(unsafe { &*key }, index as NumIndex);
        AtomU16(index as NumIndex)
    }

    /// Return the contents of the specified atom.
    #[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
        }
    }

    /// Add a byte string to the table and return its atom index. The same
    /// byte sequence always returns the same index. The bytes need not be
    /// valid UTF-8.
    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())
    }

    /// Perform the actual addition of the owned byte string.
    fn add_bytes(&mut self, owned: Vec<u8>) -> AtomBytes {
        // Remember the index of the new element.
        let index = self.strings_bytes.len();
        assert!(index < INVALID_ATOM_BYTES.0 as usize, "More than 4GB atoms?");

        // Obtain a reference to the existing bytes on the heap. That reference
        // is valid while `self` is valid. Pushing an owned Vec into the outer
        // Vec moves only the Vec struct, never its heap buffer — so a
        // *const [u8] captured from owned.as_slice() before the push stays
        // valid.
        let key: *const [u8] = owned.as_slice();

        // Push the new byte string.
        self.strings_bytes.push(owned);

        self.map_bytes.insert(unsafe { &*key }, index as NumIndex);
        AtomBytes(index as NumIndex)
    }

    /// Return the contents of the specified atom bytes.
    #[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
        }
    }
}

impl AtomTable {
    /// Create a new empty atom table.
    pub fn new() -> AtomTable {
        Default::default()
    }

    /// Add a string to the table and return its atom index. The same
    /// string always returns the same index.
    pub fn atom<V: Into<String> + AsRef<str>>(&self, value: V) -> Atom {
        unsafe { &mut *self.0.get() }.add_atom(value)
    }

    /// Return the contents of the specified atom.
    #[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)
    }

    /// Add a string to the table and return its atom index. The same
    /// string always returns the same index.
    pub fn atom_u16<V: Into<Vec<u16>> + AsRef<[u16]>>(&self, value: V) -> AtomU16 {
        unsafe { &mut *self.0.get() }.add_atom_u16(value)
    }

    /// Return the contents of the specified atom.
    #[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)
    }

    /// Add a byte string to the table and return its atom index. The same
    /// byte sequence always returns the same index. The bytes need not be
    /// valid UTF-8 (e.g., WTF-8 sequences encoding lone surrogates are
    /// accepted).
    pub fn atom_bytes<V: Into<Vec<u8>> + AsRef<[u8]>>(&self, value: V) -> AtomBytes {
        unsafe { &mut *self.0.get() }.add_atom_bytes(value)
    }

    /// Return the contents of the specified atom bytes.
    #[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)
    }

    /// Execute the callback in a context where this table is used for debug
    /// printing of atoms.
    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
        })
    }

    /// Set a table or nullptr as the Atom debug context. If non-null, debug
    /// printing of atoms will use it. Return the previous debug context.
    ///
    /// # Safety
    /// The table must not be destroyed or moved while it is set.
    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);
    }
}