Skip to main content

luau_syntax/
ast_names.rs

1use luau_common::{DenseHashHasher, DenseHashSet};
2use std::hash::{Hash, Hasher};
3
4use crate::allocator::AstArena;
5
6#[derive(Debug, Clone, Copy)]
7pub struct AstName<'ast> {
8    bytes: &'ast [u8],
9}
10
11impl<'ast> AstName<'ast> {
12    pub const fn empty_key() -> Self {
13        Self { bytes: b"" }
14    }
15
16    pub fn from_static(value: &'static str) -> Self {
17        Self {
18            bytes: value.as_bytes(),
19        }
20    }
21
22    pub fn bytes(self) -> &'ast [u8] {
23        self.bytes
24    }
25
26    pub(crate) fn from_bytes(bytes: &'ast [u8]) -> Self {
27        Self { bytes }
28    }
29
30    pub(crate) fn narrow<'short>(self) -> AstName<'short>
31    where
32        'ast: 'short,
33    {
34        AstName { bytes: self.bytes }
35    }
36}
37
38impl PartialEq for AstName<'_> {
39    fn eq(&self, other: &Self) -> bool {
40        self.bytes.len() == other.bytes.len()
41            && std::ptr::eq(self.bytes.as_ptr(), other.bytes.as_ptr())
42    }
43}
44
45impl Eq for AstName<'_> {}
46
47impl Hash for AstName<'_> {
48    fn hash<H: Hasher>(&self, state: &mut H) {
49        let key = self.bytes.as_ptr() as usize;
50        ((key >> 4) ^ (key >> 9)).hash(state);
51    }
52}
53
54impl PartialEq<&str> for AstName<'_> {
55    fn eq(&self, other: &&str) -> bool {
56        self.bytes == other.as_bytes()
57    }
58}
59
60impl PartialEq<str> for AstName<'_> {
61    fn eq(&self, other: &str) -> bool {
62        self.bytes == other.as_bytes()
63    }
64}
65
66impl PartialOrd for AstName<'_> {
67    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
68        Some(self.cmp(other))
69    }
70}
71
72impl Ord for AstName<'_> {
73    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
74        self.bytes.cmp(other.bytes)
75    }
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum LexemeType {
80    Eof,
81    Name,
82    Attribute,
83    ReservedAnd,
84    ReservedBreak,
85    ReservedDo,
86    ReservedElse,
87    ReservedElseif,
88    ReservedEnd,
89    ReservedFalse,
90    ReservedFor,
91    ReservedFunction,
92    ReservedIf,
93    ReservedIn,
94    ReservedLocal,
95    ReservedNil,
96    ReservedNot,
97    ReservedOr,
98    ReservedRepeat,
99    ReservedReturn,
100    ReservedThen,
101    ReservedTrue,
102    ReservedUntil,
103    ReservedWhile,
104}
105
106#[derive(Debug)]
107pub struct AstNameTable<'ast> {
108    entries: DenseHashSet<NameEntry<'ast>, NameEntryHasher>,
109    arena: &'ast AstArena,
110}
111
112impl<'ast> AstNameTable<'ast> {
113    pub fn new(arena: &'ast AstArena) -> Self {
114        let mut table = Self {
115            entries: DenseHashSet::with_buckets(NameEntry::empty(), 128),
116            arena,
117        };
118
119        for &(name, kind) in RESERVED {
120            table.add_static(name, kind);
121        }
122
123        table
124    }
125
126    pub fn add_static(&mut self, name: &'static str, kind: LexemeType) -> AstName<'ast> {
127        let bytes = name.as_bytes();
128        let entry = self.entries.insert_mut(NameEntry::interned(bytes, kind));
129        debug_assert!(entry.kind == kind);
130        entry.name
131    }
132
133    pub fn get_or_add_with_type(&mut self, name: &str) -> (AstName<'ast>, LexemeType) {
134        self.get_or_add_bytes_with_type(name.as_bytes())
135    }
136
137    pub fn get_or_add_bytes_with_type(&mut self, name: &[u8]) -> (AstName<'ast>, LexemeType) {
138        let entry = self.entries.insert_mut(NameEntry::lookup(name));
139        if entry.kind == LexemeType::Eof {
140            let interned = self.arena.alloc_bytes(name);
141            *entry = NameEntry::owned(interned);
142        }
143        (entry.name, entry.kind)
144    }
145
146    pub fn get_with_type(&self, name: &str) -> Option<(AstName<'ast>, LexemeType)> {
147        self.get_bytes_with_type(name.as_bytes())
148    }
149
150    pub fn get_bytes_with_type(&self, name: &[u8]) -> Option<(AstName<'ast>, LexemeType)> {
151        self.entries
152            .get(&NameEntry::lookup(name))
153            .map(|entry| (entry.name, entry.kind))
154    }
155
156    pub fn get_or_add(&mut self, name: &str) -> AstName<'ast> {
157        self.get_or_add_with_type(name).0
158    }
159
160    pub fn get_or_add_bytes(&mut self, name: &[u8]) -> AstName<'ast> {
161        self.get_or_add_bytes_with_type(name).0
162    }
163
164    pub fn get(&self, name: &str) -> Option<AstName<'ast>> {
165        self.get_with_type(name).map(|(name, _)| name)
166    }
167
168    pub fn get_bytes(&self, name: &[u8]) -> Option<AstName<'ast>> {
169        self.get_bytes_with_type(name).map(|(name, _)| name)
170    }
171
172    pub fn arena(&self) -> &'ast AstArena {
173        self.arena
174    }
175}
176
177#[derive(Debug, Clone, Copy)]
178struct NameEntry<'ast> {
179    key: NameKey,
180    name: AstName<'ast>,
181    kind: LexemeType,
182}
183
184#[derive(Debug, Clone, Copy, Default)]
185struct NameKey {
186    bytes: *const u8,
187    len: u32,
188}
189
190impl NameKey {
191    fn empty() -> Self {
192        Self {
193            bytes: std::ptr::null(),
194            len: 0,
195        }
196    }
197
198    fn new(bytes: &[u8]) -> Self {
199        Self {
200            bytes: bytes.as_ptr(),
201            len: u32::try_from(bytes.len()).expect("name length must fit u32"),
202        }
203    }
204
205    fn equals(self, other: Self) -> bool {
206        if self.len != other.len {
207            return false;
208        }
209
210        if self.bytes == other.bytes || self.len == 0 {
211            return true;
212        }
213
214        let len = self.len as usize;
215        for i in 0..len {
216            // Safety: `NameKey` only points at either source input that remains valid for the
217            // duration of the lookup or arena/static bytes stored by the table itself.
218            let left = unsafe { *self.bytes.add(i) };
219            let right = unsafe { *other.bytes.add(i) };
220            if left != right {
221                return false;
222            }
223        }
224
225        true
226    }
227}
228
229#[derive(Debug, Clone, Copy)]
230struct NameEntryHasher;
231
232impl<'ast> NameEntry<'ast> {
233    fn empty() -> Self {
234        Self {
235            key: NameKey::empty(),
236            name: AstName::from_static(""),
237            kind: LexemeType::Eof,
238        }
239    }
240
241    fn lookup(bytes: &[u8]) -> Self {
242        Self {
243            key: NameKey::new(bytes),
244            name: AstName::empty_key(),
245            kind: LexemeType::Eof,
246        }
247    }
248
249    fn interned(bytes: &'ast [u8], kind: LexemeType) -> Self {
250        Self {
251            key: NameKey::new(bytes),
252            name: AstName::from_bytes(bytes),
253            kind,
254        }
255    }
256
257    fn owned(bytes: &'ast [u8]) -> Self {
258        Self::interned(
259            bytes,
260            if bytes.first() == Some(&b'@') {
261                LexemeType::Attribute
262            } else {
263                LexemeType::Name
264            },
265        )
266    }
267}
268
269impl PartialEq for NameKey {
270    fn eq(&self, other: &Self) -> bool {
271        self.equals(*other)
272    }
273}
274
275impl Eq for NameKey {}
276
277impl PartialEq for NameEntry<'_> {
278    fn eq(&self, other: &Self) -> bool {
279        self.key == other.key
280    }
281}
282
283impl Eq for NameEntry<'_> {}
284
285#[derive(Debug, Clone, Copy)]
286pub struct AstNameDenseHasher;
287
288impl DenseHashHasher<AstName<'_>> for AstNameDenseHasher {
289    fn hash(key: &AstName<'_>) -> u64 {
290        let key = key.bytes.as_ptr() as usize;
291        ((key >> 4) ^ (key >> 9)) as u64
292    }
293}
294
295impl DenseHashHasher<NameEntry<'_>> for NameEntryHasher {
296    fn hash(entry: &NameEntry<'_>) -> u64 {
297        name_hash(entry.key.bytes, entry.key.len as usize)
298    }
299}
300
301fn name_hash(bytes: *const u8, len: usize) -> u64 {
302    const FNV_OFFSET_BASIS: u32 = 2_166_136_261;
303    const FNV_PRIME: u32 = 16_777_619;
304
305    let mut hash = FNV_OFFSET_BASIS;
306
307    for i in 0..len {
308        // Safety: `NameKey` only points at either source input that remains valid for the
309        // duration of the lookup or arena/static bytes stored by the table itself.
310        let byte = unsafe { *bytes.add(i) };
311        hash ^= u32::from(byte);
312        hash = hash.wrapping_mul(FNV_PRIME);
313    }
314
315    u64::from(hash)
316}
317
318const RESERVED: &[(&str, LexemeType)] = &[
319    ("and", LexemeType::ReservedAnd),
320    ("break", LexemeType::ReservedBreak),
321    ("do", LexemeType::ReservedDo),
322    ("else", LexemeType::ReservedElse),
323    ("elseif", LexemeType::ReservedElseif),
324    ("end", LexemeType::ReservedEnd),
325    ("false", LexemeType::ReservedFalse),
326    ("for", LexemeType::ReservedFor),
327    ("function", LexemeType::ReservedFunction),
328    ("if", LexemeType::ReservedIf),
329    ("in", LexemeType::ReservedIn),
330    ("local", LexemeType::ReservedLocal),
331    ("nil", LexemeType::ReservedNil),
332    ("not", LexemeType::ReservedNot),
333    ("or", LexemeType::ReservedOr),
334    ("repeat", LexemeType::ReservedRepeat),
335    ("return", LexemeType::ReservedReturn),
336    ("then", LexemeType::ReservedThen),
337    ("true", LexemeType::ReservedTrue),
338    ("until", LexemeType::ReservedUntil),
339    ("while", LexemeType::ReservedWhile),
340];