luau-syntax 0.732.0

use luau_common::{DenseHashHasher, DenseHashSet};
use std::hash::{Hash, Hasher};

use crate::allocator::AstArena;

#[derive(Debug, Clone, Copy)]
pub struct AstName<'ast> {
    bytes: &'ast [u8],
}

impl<'ast> AstName<'ast> {
    pub const fn empty_key() -> Self {
        Self { bytes: b"" }
    }

    pub fn from_static(value: &'static str) -> Self {
        Self {
            bytes: value.as_bytes(),
        }
    }

    pub fn bytes(self) -> &'ast [u8] {
        self.bytes
    }

    pub(crate) fn from_bytes(bytes: &'ast [u8]) -> Self {
        Self { bytes }
    }

    pub(crate) fn narrow<'short>(self) -> AstName<'short>
    where
        'ast: 'short,
    {
        AstName { bytes: self.bytes }
    }
}

impl PartialEq for AstName<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.bytes.len() == other.bytes.len()
            && std::ptr::eq(self.bytes.as_ptr(), other.bytes.as_ptr())
    }
}

impl Eq for AstName<'_> {}

impl Hash for AstName<'_> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let key = self.bytes.as_ptr() as usize;
        ((key >> 4) ^ (key >> 9)).hash(state);
    }
}

impl PartialEq<&str> for AstName<'_> {
    fn eq(&self, other: &&str) -> bool {
        self.bytes == other.as_bytes()
    }
}

impl PartialEq<str> for AstName<'_> {
    fn eq(&self, other: &str) -> bool {
        self.bytes == other.as_bytes()
    }
}

impl PartialOrd for AstName<'_> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for AstName<'_> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.bytes.cmp(other.bytes)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LexemeType {
    Eof,
    Name,
    Attribute,
    ReservedAnd,
    ReservedBreak,
    ReservedDo,
    ReservedElse,
    ReservedElseif,
    ReservedEnd,
    ReservedFalse,
    ReservedFor,
    ReservedFunction,
    ReservedIf,
    ReservedIn,
    ReservedLocal,
    ReservedNil,
    ReservedNot,
    ReservedOr,
    ReservedRepeat,
    ReservedReturn,
    ReservedThen,
    ReservedTrue,
    ReservedUntil,
    ReservedWhile,
}

#[derive(Debug)]
pub struct AstNameTable<'ast> {
    entries: DenseHashSet<NameEntry<'ast>, NameEntryHasher>,
    arena: &'ast AstArena,
}

impl<'ast> AstNameTable<'ast> {
    pub fn new(arena: &'ast AstArena) -> Self {
        let mut table = Self {
            entries: DenseHashSet::with_buckets(NameEntry::empty(), 128),
            arena,
        };

        for &(name, kind) in RESERVED {
            table.add_static(name, kind);
        }

        table
    }

    pub fn add_static(&mut self, name: &'static str, kind: LexemeType) -> AstName<'ast> {
        let bytes = name.as_bytes();
        let entry = self.entries.insert_mut(NameEntry::interned(bytes, kind));
        debug_assert!(entry.kind == kind);
        entry.name
    }

    pub fn get_or_add_with_type(&mut self, name: &str) -> (AstName<'ast>, LexemeType) {
        self.get_or_add_bytes_with_type(name.as_bytes())
    }

    pub fn get_or_add_bytes_with_type(&mut self, name: &[u8]) -> (AstName<'ast>, LexemeType) {
        let entry = self.entries.insert_mut(NameEntry::lookup(name));
        if entry.kind == LexemeType::Eof {
            let interned = self.arena.alloc_bytes(name);
            *entry = NameEntry::owned(interned);
        }
        (entry.name, entry.kind)
    }

    pub fn get_with_type(&self, name: &str) -> Option<(AstName<'ast>, LexemeType)> {
        self.get_bytes_with_type(name.as_bytes())
    }

    pub fn get_bytes_with_type(&self, name: &[u8]) -> Option<(AstName<'ast>, LexemeType)> {
        self.entries
            .get(&NameEntry::lookup(name))
            .map(|entry| (entry.name, entry.kind))
    }

    pub fn get_or_add(&mut self, name: &str) -> AstName<'ast> {
        self.get_or_add_with_type(name).0
    }

    pub fn get_or_add_bytes(&mut self, name: &[u8]) -> AstName<'ast> {
        self.get_or_add_bytes_with_type(name).0
    }

    pub fn get(&self, name: &str) -> Option<AstName<'ast>> {
        self.get_with_type(name).map(|(name, _)| name)
    }

    pub fn get_bytes(&self, name: &[u8]) -> Option<AstName<'ast>> {
        self.get_bytes_with_type(name).map(|(name, _)| name)
    }

    pub fn arena(&self) -> &'ast AstArena {
        self.arena
    }
}

#[derive(Debug, Clone, Copy)]
struct NameEntry<'ast> {
    key: NameKey,
    name: AstName<'ast>,
    kind: LexemeType,
}

#[derive(Debug, Clone, Copy, Default)]
struct NameKey {
    bytes: *const u8,
    len: u32,
}

impl NameKey {
    fn empty() -> Self {
        Self {
            bytes: std::ptr::null(),
            len: 0,
        }
    }

    fn new(bytes: &[u8]) -> Self {
        Self {
            bytes: bytes.as_ptr(),
            len: u32::try_from(bytes.len()).expect("name length must fit u32"),
        }
    }

    fn equals(self, other: Self) -> bool {
        if self.len != other.len {
            return false;
        }

        if self.bytes == other.bytes || self.len == 0 {
            return true;
        }

        let len = self.len as usize;
        for i in 0..len {
            // Safety: `NameKey` only points at either source input that remains valid for the
            // duration of the lookup or arena/static bytes stored by the table itself.
            let left = unsafe { *self.bytes.add(i) };
            let right = unsafe { *other.bytes.add(i) };
            if left != right {
                return false;
            }
        }

        true
    }
}

#[derive(Debug, Clone, Copy)]
struct NameEntryHasher;

impl<'ast> NameEntry<'ast> {
    fn empty() -> Self {
        Self {
            key: NameKey::empty(),
            name: AstName::from_static(""),
            kind: LexemeType::Eof,
        }
    }

    fn lookup(bytes: &[u8]) -> Self {
        Self {
            key: NameKey::new(bytes),
            name: AstName::empty_key(),
            kind: LexemeType::Eof,
        }
    }

    fn interned(bytes: &'ast [u8], kind: LexemeType) -> Self {
        Self {
            key: NameKey::new(bytes),
            name: AstName::from_bytes(bytes),
            kind,
        }
    }

    fn owned(bytes: &'ast [u8]) -> Self {
        Self::interned(
            bytes,
            if bytes.first() == Some(&b'@') {
                LexemeType::Attribute
            } else {
                LexemeType::Name
            },
        )
    }
}

impl PartialEq for NameKey {
    fn eq(&self, other: &Self) -> bool {
        self.equals(*other)
    }
}

impl Eq for NameKey {}

impl PartialEq for NameEntry<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.key == other.key
    }
}

impl Eq for NameEntry<'_> {}

#[derive(Debug, Clone, Copy)]
pub struct AstNameDenseHasher;

impl DenseHashHasher<AstName<'_>> for AstNameDenseHasher {
    fn hash(key: &AstName<'_>) -> u64 {
        let key = key.bytes.as_ptr() as usize;
        ((key >> 4) ^ (key >> 9)) as u64
    }
}

impl DenseHashHasher<NameEntry<'_>> for NameEntryHasher {
    fn hash(entry: &NameEntry<'_>) -> u64 {
        name_hash(entry.key.bytes, entry.key.len as usize)
    }
}

fn name_hash(bytes: *const u8, len: usize) -> u64 {
    const FNV_OFFSET_BASIS: u32 = 2_166_136_261;
    const FNV_PRIME: u32 = 16_777_619;

    let mut hash = FNV_OFFSET_BASIS;

    for i in 0..len {
        // Safety: `NameKey` only points at either source input that remains valid for the
        // duration of the lookup or arena/static bytes stored by the table itself.
        let byte = unsafe { *bytes.add(i) };
        hash ^= u32::from(byte);
        hash = hash.wrapping_mul(FNV_PRIME);
    }

    u64::from(hash)
}

const RESERVED: &[(&str, LexemeType)] = &[
    ("and", LexemeType::ReservedAnd),
    ("break", LexemeType::ReservedBreak),
    ("do", LexemeType::ReservedDo),
    ("else", LexemeType::ReservedElse),
    ("elseif", LexemeType::ReservedElseif),
    ("end", LexemeType::ReservedEnd),
    ("false", LexemeType::ReservedFalse),
    ("for", LexemeType::ReservedFor),
    ("function", LexemeType::ReservedFunction),
    ("if", LexemeType::ReservedIf),
    ("in", LexemeType::ReservedIn),
    ("local", LexemeType::ReservedLocal),
    ("nil", LexemeType::ReservedNil),
    ("not", LexemeType::ReservedNot),
    ("or", LexemeType::ReservedOr),
    ("repeat", LexemeType::ReservedRepeat),
    ("return", LexemeType::ReservedReturn),
    ("then", LexemeType::ReservedThen),
    ("true", LexemeType::ReservedTrue),
    ("until", LexemeType::ReservedUntil),
    ("while", LexemeType::ReservedWhile),
];