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 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 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];