hermes_atom_table/
lib.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8//! Ported from juno atom_table for the Hermes Rust lexer; carries encapsulated
9//! unsafe; adds a byte/WTF-8 intern path.
10
11use std::cell::Cell;
12use std::cell::UnsafeCell;
13use std::collections::HashMap;
14use std::fmt::Formatter;
15use std::ptr::null;
16
17/// Type used to hold a string index internally.
18type NumIndex = u32;
19
20/// A string uniquing table - only one copy of a string is stored and all attempts
21/// to add the same string again return the same atom. This table is intended to
22/// be easily shareable, so it utilizes interior mutability. UnsafeCell<> is safe
23/// because we never allow reference to it to escape.
24#[derive(Debug, Default)]
25pub struct AtomTable(UnsafeCell<Inner>);
26
27/// A string uniquing table - only one copy of a string is stored and all attempts
28/// to add the same string again return the same atom.
29#[derive(Default)]
30struct Inner {
31    /// Strings are added here and never removed or mutated.
32    strings: Vec<String>,
33    /// Maps from a reference inside [`Inner::strings`] to the index in [`Inner::strings`].
34    /// Since strings are never removed or modified, the lifetime of the key
35    /// is effectively static.
36    map: HashMap<&'static str, NumIndex>,
37
38    /// Strings are added here and never removed or mutated.
39    strings_u16: Vec<Vec<u16>>,
40    /// Maps from a reference inside [`Inner::strings_u16`] to the index in [`Inner::strings_u16`].
41    /// Since strings are never removed or modified, the lifetime of the key
42    /// is effectively static.
43    map_u16: HashMap<&'static [u16], NumIndex>,
44
45    /// Byte strings are added here and never removed or mutated.
46    /// The bytes need not be valid UTF-8 (they may be WTF-8 or arbitrary byte
47    /// sequences, e.g. JS string literals containing lone surrogates).
48    strings_bytes: Vec<Vec<u8>>,
49    /// Maps from a reference inside [`Inner::strings_bytes`] to the index in
50    /// [`Inner::strings_bytes`]. Since strings are never removed or modified,
51    /// the lifetime of the key is effectively static.
52    map_bytes: HashMap<&'static [u8], NumIndex>,
53}
54
55/// This represents a unique string index in the table.
56#[derive(Copy, Clone, Eq, PartialEq, Hash)]
57pub struct Atom(NumIndex);
58
59/// This represents a unique string index in the table.
60#[derive(Copy, Clone, Eq, PartialEq, Hash)]
61pub struct AtomU16(NumIndex);
62
63/// This represents a unique byte-string index in the table.
64/// The bytes need not be valid UTF-8.
65#[derive(Copy, Clone, Eq, PartialEq, Hash)]
66pub struct AtomBytes(NumIndex);
67
68thread_local! {
69    /// Stores the active table used for debug formatting.
70    static DEBUG_TABLE: Cell<* const AtomTable> = Cell::new(null());
71}
72
73// An implementation of Debug which optionally obtains the Atom value from the
74// active debug map.
75impl std::fmt::Debug for Atom {
76    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
77        let mut t = f.debug_tuple("Atom");
78        t.field(&self.0);
79
80        // If the debug table is set and the atom is valid in it, add the value
81        DEBUG_TABLE.with(|debug_table| {
82            let p = debug_table.get();
83            if let Some(r) = unsafe { p.as_ref() } {
84                if let Some(value) = r.try_str(*self) {
85                    t.field(&value);
86                }
87            }
88        });
89        t.finish()
90    }
91}
92
93// An implementation of Debug which optionally obtains the Atom value from the
94// active debug map.
95impl std::fmt::Debug for AtomU16 {
96    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
97        let mut t = f.debug_tuple("Atom");
98        t.field(&self.0);
99
100        // If the debug table is set and the atom is valid in it, add the value
101        DEBUG_TABLE.with(|debug_table| {
102            let p = debug_table.get();
103            if let Some(r) = unsafe { p.as_ref() } {
104                if let Some(value) = r.try_str_u16(*self) {
105                    t.field(&value);
106                }
107            }
108        });
109        t.finish()
110    }
111}
112
113// An implementation of Debug which optionally obtains the AtomBytes value from
114// the active debug map.
115impl std::fmt::Debug for AtomBytes {
116    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
117        let mut t = f.debug_tuple("AtomBytes");
118        t.field(&self.0);
119
120        // If the debug table is set and the atom is valid in it, add the value
121        DEBUG_TABLE.with(|debug_table| {
122            let p = debug_table.get();
123            if let Some(r) = unsafe { p.as_ref() } {
124                if let Some(value) = r.try_bytes(*self) {
125                    t.field(&value);
126                }
127            }
128        });
129        t.finish()
130    }
131}
132
133/// A special value reserved for the invalid atom.
134pub const INVALID_ATOM: Atom = Atom(NumIndex::MAX);
135
136/// A special value reserved for the invalid atom bytes.
137pub const INVALID_ATOM_BYTES: AtomBytes = AtomBytes(NumIndex::MAX);
138
139impl Inner {
140    /// Add a string to the table and return its atom index. The same
141    /// string always returns the same index.
142    fn add_atom<V: Into<String> + AsRef<str>>(&mut self, value: V) -> Atom {
143        if let Some(index) = self.map.get(value.as_ref()) {
144            return Atom(*index);
145        }
146        self.add(value.into())
147    }
148
149    /// Perform the actual addition of the owned string.
150    fn add(&mut self, owned: String) -> Atom {
151        // Remember the index of the new element.
152        let index = self.strings.len();
153        assert!(index < INVALID_ATOM.0 as usize, "More than 4GB atoms?");
154
155        // Obtain a reference to the existing string on the heap. That reference
156        // is valid while `self` is valid.
157        let key: *const str = owned.as_str();
158
159        // Push the new string.
160        self.strings.push(owned);
161
162        self.map.insert(unsafe { &*key }, index as NumIndex);
163        Atom(index as NumIndex)
164    }
165
166    /// Return the contents of the specified atom.
167    #[inline]
168    fn str(&self, ident: Atom) -> &str {
169        self.strings[ident.0 as usize].as_str()
170    }
171
172    fn try_str(&self, ident: Atom) -> Option<&str> {
173        if (ident.0 as usize) < self.strings.len() {
174            Some(self.str(ident))
175        } else {
176            None
177        }
178    }
179
180    /// Add a string to the table and return its atom index. The same
181    /// string always returns the same index.
182    fn add_atom_u16<V: Into<Vec<u16>> + AsRef<[u16]>>(&mut self, value: V) -> AtomU16 {
183        if let Some(index) = self.map_u16.get(value.as_ref()) {
184            return AtomU16(*index);
185        }
186        self.add_u16(value.into())
187    }
188
189    /// Perform the actual addition of the owned string.
190    fn add_u16(&mut self, owned: Vec<u16>) -> AtomU16 {
191        // Remember the index of the new element.
192        let index = self.strings_u16.len();
193        assert!(index < INVALID_ATOM.0 as usize, "More than 4GB atoms?");
194
195        // Obtain a reference to the existing string on the heap. That reference
196        // is valid while `self` is valid.
197        let key: *const [u16] = owned.as_slice();
198
199        // Push the new string.
200        self.strings_u16.push(owned);
201
202        self.map_u16.insert(unsafe { &*key }, index as NumIndex);
203        AtomU16(index as NumIndex)
204    }
205
206    /// Return the contents of the specified atom.
207    #[inline]
208    fn str_u16(&self, ident: AtomU16) -> &[u16] {
209        self.strings_u16[ident.0 as usize].as_slice()
210    }
211
212    fn try_str_u16(&self, ident: AtomU16) -> Option<&[u16]> {
213        if (ident.0 as usize) < self.strings_u16.len() {
214            Some(self.str_u16(ident))
215        } else {
216            None
217        }
218    }
219
220    /// Add a byte string to the table and return its atom index. The same
221    /// byte sequence always returns the same index. The bytes need not be
222    /// valid UTF-8.
223    fn add_atom_bytes<V: Into<Vec<u8>> + AsRef<[u8]>>(&mut self, value: V) -> AtomBytes {
224        if let Some(index) = self.map_bytes.get(value.as_ref()) {
225            return AtomBytes(*index);
226        }
227        self.add_bytes(value.into())
228    }
229
230    /// Perform the actual addition of the owned byte string.
231    fn add_bytes(&mut self, owned: Vec<u8>) -> AtomBytes {
232        // Remember the index of the new element.
233        let index = self.strings_bytes.len();
234        assert!(index < INVALID_ATOM_BYTES.0 as usize, "More than 4GB atoms?");
235
236        // Obtain a reference to the existing bytes on the heap. That reference
237        // is valid while `self` is valid. Pushing an owned Vec into the outer
238        // Vec moves only the Vec struct, never its heap buffer — so a
239        // *const [u8] captured from owned.as_slice() before the push stays
240        // valid.
241        let key: *const [u8] = owned.as_slice();
242
243        // Push the new byte string.
244        self.strings_bytes.push(owned);
245
246        self.map_bytes.insert(unsafe { &*key }, index as NumIndex);
247        AtomBytes(index as NumIndex)
248    }
249
250    /// Return the contents of the specified atom bytes.
251    #[inline]
252    fn bytes(&self, ident: AtomBytes) -> &[u8] {
253        self.strings_bytes[ident.0 as usize].as_slice()
254    }
255
256    fn try_bytes(&self, ident: AtomBytes) -> Option<&[u8]> {
257        if (ident.0 as usize) < self.strings_bytes.len() {
258            Some(self.bytes(ident))
259        } else {
260            None
261        }
262    }
263}
264
265impl AtomTable {
266    /// Create a new empty atom table.
267    pub fn new() -> AtomTable {
268        Default::default()
269    }
270
271    /// Add a string to the table and return its atom index. The same
272    /// string always returns the same index.
273    pub fn atom<V: Into<String> + AsRef<str>>(&self, value: V) -> Atom {
274        unsafe { &mut *self.0.get() }.add_atom(value)
275    }
276
277    /// Return the contents of the specified atom.
278    #[inline]
279    pub fn str(&self, ident: Atom) -> &str {
280        unsafe { &*self.0.get() }.str(ident)
281    }
282
283    #[inline]
284    pub fn try_str(&self, ident: Atom) -> Option<&str> {
285        unsafe { &*self.0.get() }.try_str(ident)
286    }
287
288    /// Add a string to the table and return its atom index. The same
289    /// string always returns the same index.
290    pub fn atom_u16<V: Into<Vec<u16>> + AsRef<[u16]>>(&self, value: V) -> AtomU16 {
291        unsafe { &mut *self.0.get() }.add_atom_u16(value)
292    }
293
294    /// Return the contents of the specified atom.
295    #[inline]
296    pub fn str_u16(&self, ident: AtomU16) -> &[u16] {
297        unsafe { &*self.0.get() }.str_u16(ident)
298    }
299
300    #[inline]
301    pub fn try_str_u16(&self, ident: AtomU16) -> Option<&[u16]> {
302        unsafe { &*self.0.get() }.try_str_u16(ident)
303    }
304
305    /// Add a byte string to the table and return its atom index. The same
306    /// byte sequence always returns the same index. The bytes need not be
307    /// valid UTF-8 (e.g., WTF-8 sequences encoding lone surrogates are
308    /// accepted).
309    pub fn atom_bytes<V: Into<Vec<u8>> + AsRef<[u8]>>(&self, value: V) -> AtomBytes {
310        unsafe { &mut *self.0.get() }.add_atom_bytes(value)
311    }
312
313    /// Return the contents of the specified atom bytes.
314    #[inline]
315    pub fn bytes(&self, ident: AtomBytes) -> &[u8] {
316        unsafe { &*self.0.get() }.bytes(ident)
317    }
318
319    #[inline]
320    pub fn try_bytes(&self, ident: AtomBytes) -> Option<&[u8]> {
321        unsafe { &*self.0.get() }.try_bytes(ident)
322    }
323
324    /// Execute the callback in a context where this table is used for debug
325    /// printing of atoms.
326    pub fn in_debug_context<R, F: FnOnce() -> R>(&self, f: F) -> R {
327        DEBUG_TABLE.with(|debug_table| {
328            let prev_table = debug_table.replace(self);
329            let res = f();
330            debug_assert!(
331                debug_table.get() == self,
332                "debug context unexpectedly changed"
333            );
334            debug_table.set(prev_table);
335            res
336        })
337    }
338
339    /// Set a table or nullptr as the Atom debug context. If non-null, debug
340    /// printing of atoms will use it. Return the previous debug context.
341    ///
342    /// # Safety
343    /// The table must not be destroyed or moved while it is set.
344    pub unsafe fn unsafe_set_debug_context(ptr: *const Self) -> *const Self {
345        DEBUG_TABLE.with(|debug_table| debug_table.replace(ptr))
346    }
347}
348
349impl std::ops::Index<Atom> for AtomTable {
350    type Output = str;
351
352    fn index(&self, index: Atom) -> &Self::Output {
353        self.str(index)
354    }
355}
356
357impl std::ops::Index<AtomU16> for AtomTable {
358    type Output = [u16];
359
360    fn index(&self, index: AtomU16) -> &Self::Output {
361        self.str_u16(index)
362    }
363}
364
365impl std::ops::Index<AtomBytes> for AtomTable {
366    type Output = [u8];
367
368    fn index(&self, index: AtomBytes) -> &Self::Output {
369        self.bytes(index)
370    }
371}
372
373#[cfg(test)]
374mod tests {
375    use super::*;
376
377    #[test]
378    fn test_tab() {
379        let idtab = AtomTable::new();
380
381        let id_foo = idtab.atom("foo");
382        let p_foo: *const str = idtab.str(id_foo);
383        let id_bar = idtab.atom("bar");
384        assert_ne!(id_foo, id_bar);
385
386        assert_eq!(idtab.atom("foo"), id_foo);
387        assert_eq!(idtab.atom("bar"), id_bar);
388
389        assert_eq!(idtab.atom(String::from("foo")), id_foo);
390        assert_eq!(idtab.atom(String::from("bar")), id_bar);
391
392        assert_eq!(idtab.str(id_foo), "foo");
393        assert_eq!(idtab.str(id_bar), "bar");
394
395        assert_eq!(idtab.str(id_foo) as *const str, p_foo);
396    }
397
398    #[test]
399    fn test_bytes() {
400        let tab = AtomTable::new();
401        let foo = tab.atom_bytes(b"foo".as_slice());
402        let bar = tab.atom_bytes(b"bar".as_slice());
403        assert_ne!(foo, bar);
404        assert_eq!(tab.atom_bytes(b"foo".as_slice()), foo);
405        assert_eq!(tab.atom_bytes(Vec::from(*b"bar")), bar);
406        assert_eq!(tab.bytes(foo), b"foo");
407        assert_eq!(&tab[bar], b"bar");
408        let p_foo: *const [u8] = tab.bytes(foo);
409        let _ = tab.atom_bytes(b"baz".as_slice());
410        assert_eq!(tab.bytes(foo) as *const [u8], p_foo);
411    }
412
413    #[test]
414    fn test_bytes_ill_formed_utf8() {
415        let tab = AtomTable::new();
416        let lone_surrogate: &[u8] = &[0xed, 0xa0, 0x80];
417        let a = tab.atom_bytes(lone_surrogate);
418        assert_eq!(tab.bytes(a), lone_surrogate);
419        assert_eq!(tab.atom_bytes(lone_surrogate), a);
420        let s = tab.atom("foo");
421        let b = tab.atom_bytes(b"foo".as_slice());
422        assert_eq!(tab.str(s), "foo");
423        assert_eq!(tab.bytes(b), b"foo");
424    }
425
426    #[test]
427    fn test_bytes_try_and_invalid() {
428        let tab = AtomTable::new();
429        let a = tab.atom_bytes(b"x".as_slice());
430        assert_eq!(tab.try_bytes(a), Some(b"x".as_slice()));
431        assert_eq!(tab.try_bytes(INVALID_ATOM_BYTES), None);
432    }
433}
hermes_atom_table/lib.rs

hermes_atom_table/
lib.rs
