luau-syntax 0.732.0

use super::*;
use crate::allocator::AstArena;
use std::fmt;
use std::marker::PhantomData;
use std::ptr::NonNull;

#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TypeTag {
    Reference,
    Table,
    Function,
    Typeof,
    SingletonBool,
    SingletonString,
    Group,
    Optional,
    Union,
    Intersection,
    Error,
}

#[repr(C)]
#[derive(Debug, PartialEq)]
pub struct TypeHeader<'ast> {
    pub tag: TypeTag,
    pub location: Location,
    _marker: PhantomData<&'ast ()>,
}

#[derive(Clone, Copy)]
pub struct Type<'ast> {
    ptr: NonNull<TypeHeader<'ast>>,
    _marker: PhantomData<&'ast TypeHeader<'ast>>,
}

impl<'ast> std::ops::Deref for Type<'ast> {
    type Target = TypeHeader<'ast>;

    fn deref(&self) -> &Self::Target {
        self.header()
    }
}

impl fmt::Debug for Type<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("Type")
            .field("location", &self.location())
            .field("kind", &self.kind())
            .finish()
    }
}

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

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TypeKind<'ast> {
    Reference {
        prefix: Option<AstName<'ast>>,
        prefix_location: Option<Location>,
        prefix_local: Option<&'ast Local<'ast>>,
        name: AstName<'ast>,
        location: Location,
        name_location: Location,
        has_parameter_list: bool,
        parameters: &'ast [TypeOrPack<'ast>],
    },
    Table {
        props: &'ast [TableTypeProp<'ast>],
        indexer: Option<&'ast TableTypeIndexer<'ast>>,
    },
    Function {
        attributes: &'ast [&'ast Attribute<'ast>],
        generics: &'ast [&'ast GenericType<'ast>],
        generic_packs: &'ast [&'ast GenericTypePack<'ast>],
        arg_types: TypeList<'ast>,
        arg_names: &'ast [Option<ArgumentName<'ast>>],
        return_types: TypePack<'ast>,
    },
    Typeof {
        expr: Expression<'ast>,
    },
    SingletonBool {
        value: bool,
    },
    SingletonString {
        value: AstString<'ast>,
    },
    Group {
        ty: Type<'ast>,
    },
    Optional,
    Union {
        types: &'ast [Type<'ast>],
    },
    Intersection {
        types: &'ast [Type<'ast>],
    },
    Error {
        types: &'ast [Type<'ast>],
        missing: bool,
        message_index: usize,
    },
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TypeOrPack<'ast> {
    Type(Type<'ast>),
    Pack(TypePack<'ast>),
}

macro_rules! type_node {
    ($name:ident { $($field:ident : $ty:ty),* $(,)? }, $tag:ident) => {
        #[repr(C)]
        #[derive(Debug, PartialEq)]
        pub struct $name<'ast> {
            pub base: TypeHeader<'ast>,
            $(pub $field: $ty),*
        }
    };
}

#[repr(C)]
#[derive(Debug, PartialEq)]
pub struct TypeUnit<'ast> {
    pub base: TypeHeader<'ast>,
}

type_node!(
    TypeReference {
        has_parameter_list: bool,
        prefix: Option<AstName<'ast>>,
        prefix_location: Option<Location>,
        prefix_local: Option<&'ast Local<'ast>>,
        name: AstName<'ast>,
        name_location: Location,
        parameters: &'ast [TypeOrPack<'ast>]
    },
    Reference
);
type_node!(
    TypeTable {
        props: &'ast [TableTypeProp<'ast>],
        indexer: Option<&'ast TableTypeIndexer<'ast>>
    },
    Table
);
type_node!(
    TypeFunction {
        attributes: &'ast [&'ast Attribute<'ast>],
        generics: &'ast [&'ast GenericType<'ast>],
        generic_packs: &'ast [&'ast GenericTypePack<'ast>],
        arg_types: TypeList<'ast>,
        arg_names: &'ast [Option<ArgumentName<'ast>>],
        return_types: TypePack<'ast>
    },
    Function
);
type_node!(TypeTypeof { expr: Expression<'ast> }, Typeof);
type_node!(TypeSingletonBool { value: bool }, SingletonBool);
type_node!(TypeSingletonString { value: AstString<'ast> }, SingletonString);
type_node!(TypeGroup { ty: Type<'ast> }, Group);
type_node!(TypeUnion { types: &'ast [Type<'ast>] }, Union);
type_node!(TypeIntersection { types: &'ast [Type<'ast>] }, Intersection);
type_node!(
    TypeError {
        types: &'ast [Type<'ast>],
        missing: bool,
        message_index: usize
    },
    Error
);

impl<'ast> Type<'ast> {
    const fn new_header(tag: TypeTag, location: Location) -> TypeHeader<'ast> {
        TypeHeader {
            tag,
            location,
            _marker: PhantomData,
        }
    }

    pub(crate) fn from_node<T>(node: &'ast mut T) -> Self {
        Self {
            ptr: NonNull::from(node).cast(),
            _marker: PhantomData,
        }
    }

    pub fn as_ptr(self) -> *const () {
        self.ptr.as_ptr().cast()
    }

    #[inline(always)]
    fn header(&self) -> &TypeHeader<'ast> {
        unsafe { self.ptr.as_ref() }
    }

    #[inline(always)]
    pub fn location(self) -> Location {
        self.header().location
    }

    #[inline(always)]
    pub fn kind(&self) -> TypeKind<'ast> {
        match self.tag {
            TypeTag::Reference => {
                let node = self.cast_ref::<TypeReference>();
                TypeKind::Reference {
                    prefix: node.prefix,
                    prefix_location: node.prefix_location,
                    prefix_local: node.prefix_local,
                    name: node.name,
                    location: self.location,
                    name_location: node.name_location,
                    has_parameter_list: node.has_parameter_list,
                    parameters: node.parameters,
                }
            }
            TypeTag::Table => {
                let node = self.cast_ref::<TypeTable>();
                TypeKind::Table {
                    props: node.props,
                    indexer: node.indexer,
                }
            }
            TypeTag::Function => {
                let node = self.cast_ref::<TypeFunction>();
                TypeKind::Function {
                    attributes: node.attributes,
                    generics: node.generics,
                    generic_packs: node.generic_packs,
                    arg_types: node.arg_types,
                    arg_names: node.arg_names,
                    return_types: node.return_types,
                }
            }
            TypeTag::Typeof => TypeKind::Typeof {
                expr: self.cast_ref::<TypeTypeof>().expr,
            },
            TypeTag::SingletonBool => TypeKind::SingletonBool {
                value: self.cast_ref::<TypeSingletonBool>().value,
            },
            TypeTag::SingletonString => TypeKind::SingletonString {
                value: self.cast_ref::<TypeSingletonString>().value,
            },
            TypeTag::Group => TypeKind::Group {
                ty: self.cast_ref::<TypeGroup>().ty,
            },
            TypeTag::Optional => TypeKind::Optional,
            TypeTag::Union => TypeKind::Union {
                types: self.cast_ref::<TypeUnion>().types,
            },
            TypeTag::Intersection => TypeKind::Intersection {
                types: self.cast_ref::<TypeIntersection>().types,
            },
            TypeTag::Error => {
                let node = self.cast_ref::<TypeError>();
                TypeKind::Error {
                    types: node.types,
                    missing: node.missing,
                    message_index: node.message_index,
                }
            }
        }
    }

    pub fn visit<V: AstVisitor>(self, visitor: &mut V) {
        let should_visit = match self.tag {
            TypeTag::Reference => visitor.visit_reference_type(self),
            TypeTag::Table => visitor.visit_table_type(self),
            TypeTag::Function => visitor.visit_function_type(self),
            TypeTag::Typeof => visitor.visit_typeof_type(self),
            TypeTag::SingletonBool => visitor.visit_singleton_bool_type(self),
            TypeTag::SingletonString => visitor.visit_singleton_string_type(self),
            TypeTag::Group => visitor.visit_group_type(self),
            TypeTag::Optional => visitor.visit_optional_type(self),
            TypeTag::Union => visitor.visit_union_type(self),
            TypeTag::Intersection => visitor.visit_intersection_type(self),
            TypeTag::Error => visitor.visit_error_type(self),
        };

        if !should_visit {
            return;
        }

        match self.kind() {
            TypeKind::Reference { parameters, .. } => {
                for parameter in parameters {
                    match parameter {
                        TypeOrPack::Type(annotation) => annotation.visit(visitor),
                        TypeOrPack::Pack(pack) => pack.visit(visitor),
                    }
                }
            }
            TypeKind::Union { types, .. } | TypeKind::Intersection { types, .. } => {
                for annotation in types {
                    annotation.visit(visitor);
                }
            }
            TypeKind::Table { props, indexer, .. } => {
                for prop in props {
                    prop.ty.visit(visitor);
                }
                if let Some(indexer) = indexer {
                    indexer.index_type.visit(visitor);
                    indexer.result_type.visit(visitor);
                }
            }
            TypeKind::Function {
                arg_types,
                return_types,
                ..
            } => {
                arg_types.visit(visitor);
                return_types.visit(visitor);
            }
            TypeKind::Typeof { expr, .. } => expr.visit(visitor),
            TypeKind::Group { ty, .. } => ty.visit(visitor),
            TypeKind::Error { types, .. } => {
                for annotation in types {
                    annotation.visit(visitor);
                }
            }
            TypeKind::Optional
            | TypeKind::SingletonBool { .. }
            | TypeKind::SingletonString { .. } => {}
        }
    }

    #[inline(always)]
    fn cast_ref<T>(self) -> &'ast T {
        unsafe { self.ptr.cast::<T>().as_ref() }
    }

    #[inline(always)]
    fn cast_if_tag<T>(self, tag: TypeTag) -> Option<&'ast T> {
        (self.tag == tag).then(|| self.cast_ref())
    }

    #[inline(always)]
    pub fn as_reference(self) -> Option<&'ast TypeReference<'ast>> {
        self.cast_if_tag(TypeTag::Reference)
    }

    #[inline(always)]
    pub fn as_table(self) -> Option<&'ast TypeTable<'ast>> {
        self.cast_if_tag(TypeTag::Table)
    }

    #[inline(always)]
    pub fn as_function(self) -> Option<&'ast TypeFunction<'ast>> {
        self.cast_if_tag(TypeTag::Function)
    }

    #[inline(always)]
    pub fn as_typeof(self) -> Option<&'ast TypeTypeof<'ast>> {
        self.cast_if_tag(TypeTag::Typeof)
    }

    #[inline(always)]
    pub fn as_singleton_bool(self) -> Option<&'ast TypeSingletonBool<'ast>> {
        self.cast_if_tag(TypeTag::SingletonBool)
    }

    #[inline(always)]
    pub fn as_singleton_string(self) -> Option<&'ast TypeSingletonString<'ast>> {
        self.cast_if_tag(TypeTag::SingletonString)
    }

    #[inline(always)]
    pub fn as_group(self) -> Option<&'ast TypeGroup<'ast>> {
        self.cast_if_tag(TypeTag::Group)
    }

    #[inline(always)]
    pub fn as_union(self) -> Option<&'ast TypeUnion<'ast>> {
        self.cast_if_tag(TypeTag::Union)
    }

    #[inline(always)]
    pub fn as_intersection(self) -> Option<&'ast TypeIntersection<'ast>> {
        self.cast_if_tag(TypeTag::Intersection)
    }

    #[inline(always)]
    pub fn as_error(self) -> Option<&'ast TypeError<'ast>> {
        self.cast_if_tag(TypeTag::Error)
    }
}

impl<'ast> TypeKind<'ast> {
    pub fn is_checked_function(&self) -> bool {
        self.has_function_attribute(AttributeKind::Checked)
    }

    pub fn has_function_attribute(&self, kind: AttributeKind) -> bool {
        self.get_function_attribute(kind).is_some()
    }

    pub fn get_function_attribute(&self, kind: AttributeKind) -> Option<&'ast Attribute<'ast>> {
        match self {
            Self::Function { attributes, .. } => find_attribute(attributes, kind),
            _ => None,
        }
    }
}

#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TypePackTag {
    Explicit,
    Variadic,
    Generic,
}

#[repr(C)]
#[derive(Debug, PartialEq)]
pub struct TypePackHeader<'ast> {
    pub tag: TypePackTag,
    pub location: Location,
    _marker: PhantomData<&'ast ()>,
}

#[derive(Clone, Copy)]
pub struct TypePack<'ast> {
    ptr: NonNull<TypePackHeader<'ast>>,
    _marker: PhantomData<&'ast TypePackHeader<'ast>>,
}

impl<'ast> std::ops::Deref for TypePack<'ast> {
    type Target = TypePackHeader<'ast>;

    fn deref(&self) -> &Self::Target {
        self.header()
    }
}

impl fmt::Debug for TypePack<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("TypePack")
            .field("location", &self.location())
            .field("kind", &self.kind())
            .finish()
    }
}

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

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct TypeList<'ast> {
    pub types: &'ast [Type<'ast>],
    pub tail_type: Option<TypePack<'ast>>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum TypePackKind<'ast> {
    Explicit { type_list: TypeList<'ast> },
    Variadic { variadic_type: Type<'ast> },
    Generic { generic_name: AstName<'ast> },
}

macro_rules! type_pack_node {
    ($name:ident { $($field:ident : $ty:ty),* $(,)? }, $tag:ident) => {
        #[repr(C)]
        #[derive(Debug, PartialEq)]
        pub struct $name<'ast> {
            pub base: TypePackHeader<'ast>,
            $(pub $field: $ty),*
        }
    };
}

type_pack_node!(TypePackExplicitNode { type_list: TypeList<'ast> }, Explicit);
type_pack_node!(TypePackVariadicNode { variadic_type: Type<'ast> }, Variadic);
type_pack_node!(TypePackGenericNode { generic_name: AstName<'ast> }, Generic);

impl TypeList<'_> {
    pub fn visit<V: AstVisitor>(&self, visitor: &mut V) {
        for annotation in self.types {
            annotation.visit(visitor);
        }
        if let Some(tail_type) = self.tail_type {
            tail_type.visit(visitor);
        }
    }
}

impl<'ast> TypePack<'ast> {
    const fn new_header(tag: TypePackTag, location: Location) -> TypePackHeader<'ast> {
        TypePackHeader {
            tag,
            location,
            _marker: PhantomData,
        }
    }

    pub(crate) fn from_node<T>(node: &'ast mut T) -> Self {
        Self {
            ptr: NonNull::from(node).cast(),
            _marker: PhantomData,
        }
    }

    pub fn as_ptr(self) -> *const () {
        self.ptr.as_ptr().cast()
    }

    #[inline(always)]
    fn header(&self) -> &TypePackHeader<'ast> {
        unsafe { self.ptr.as_ref() }
    }

    #[inline(always)]
    pub fn location(self) -> Location {
        self.header().location
    }

    #[inline(always)]
    pub fn kind(&self) -> TypePackKind<'ast> {
        match self.tag {
            TypePackTag::Explicit => TypePackKind::Explicit {
                type_list: self.cast_ref::<TypePackExplicitNode>().type_list,
            },
            TypePackTag::Variadic => TypePackKind::Variadic {
                variadic_type: self.cast_ref::<TypePackVariadicNode>().variadic_type,
            },
            TypePackTag::Generic => TypePackKind::Generic {
                generic_name: self.cast_ref::<TypePackGenericNode>().generic_name,
            },
        }
    }

    pub fn explicit_type_list(&self) -> Option<TypeList<'ast>> {
        match self.kind() {
            TypePackKind::Explicit { type_list } => Some(type_list),
            _ => None,
        }
    }

    pub fn visit<V: AstVisitor>(self, visitor: &mut V) {
        let should_visit = match self.tag {
            TypePackTag::Explicit => visitor.visit_explicit_type_pack(self),
            TypePackTag::Variadic => visitor.visit_variadic_type_pack(self),
            TypePackTag::Generic => visitor.visit_generic_type_pack_type(self),
        };

        if !should_visit {
            return;
        }

        match self.kind() {
            TypePackKind::Explicit { type_list } => type_list.visit(visitor),
            TypePackKind::Variadic { variadic_type } => variadic_type.visit(visitor),
            TypePackKind::Generic { .. } => {}
        }
    }

    #[inline(always)]
    fn cast_ref<T>(self) -> &'ast T {
        unsafe { self.ptr.cast::<T>().as_ref() }
    }

    #[inline(always)]
    fn cast_if_tag<T>(self, tag: TypePackTag) -> Option<&'ast T> {
        (self.tag == tag).then(|| self.cast_ref())
    }

    #[inline(always)]
    pub fn as_explicit(self) -> Option<&'ast TypePackExplicitNode<'ast>> {
        self.cast_if_tag(TypePackTag::Explicit)
    }

    #[inline(always)]
    pub fn as_variadic(self) -> Option<&'ast TypePackVariadicNode<'ast>> {
        self.cast_if_tag(TypePackTag::Variadic)
    }

    #[inline(always)]
    pub fn as_generic(self) -> Option<&'ast TypePackGenericNode<'ast>> {
        self.cast_if_tag(TypePackTag::Generic)
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Attribute<'ast> {
    pub location: Location,
    pub kind: AttributeKind,
    pub args: &'ast [Expression<'ast>],
    pub name: AstName<'ast>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AttributeKind {
    Checked,
    Native,
    Deprecated,
    DebugNoinline,
    Unknown,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeprecatedInfo {
    pub deprecated: bool,
    pub use_replacement: Option<Vec<u8>>,
    pub reason: Option<Vec<u8>>,
}

impl Attribute<'_> {
    pub fn kind(&self) -> AttributeKind {
        self.kind
    }

    pub fn visit<V: AstVisitor>(&self, visitor: &mut V) {
        let _ = visitor.visit_attribute(self);
    }

    pub fn deprecated_info(&self) -> DeprecatedInfo {
        let mut info = DeprecatedInfo {
            deprecated: self.kind == AttributeKind::Deprecated,
            use_replacement: None,
            reason: None,
        };

        if !info.deprecated {
            return info;
        }

        let Some(argument) = self.args.first() else {
            return info;
        };
        let ExpressionKind::Table { items } = argument.kind() else {
            return info;
        };

        for item in items {
            let TableItem::Record { key, value } = item else {
                continue;
            };
            let ExpressionKind::String {
                value: key_value, ..
            } = key.kind()
            else {
                continue;
            };
            let ExpressionKind::String { value, .. } = value.kind() else {
                continue;
            };

            match key_value.as_bytes() {
                b"use" => info.use_replacement = Some(value.as_bytes().to_vec()),
                b"reason" => info.reason = Some(value.as_bytes().to_vec()),
                _ => {}
            }
        }

        info
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct TableTypeProp<'ast> {
    pub name: AstName<'ast>,
    pub location: Location,
    pub ty: Type<'ast>,
    pub access: TableAccess,
    pub access_location: Option<Location>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct TableTypeIndexer<'ast> {
    pub index_type: Type<'ast>,
    pub result_type: Type<'ast>,
    pub location: Location,
    pub access: TableAccess,
    pub access_location: Option<Location>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct DeclaredExternTypeProperty<'ast> {
    pub name: AstName<'ast>,
    pub name_location: Location,
    pub ty: Type<'ast>,
    pub is_method: bool,
    pub location: Location,
    pub access: TableAccess,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TableAccess {
    Read,
    Write,
    ReadWrite,
}

impl AstArena {
    fn alloc_type_node<'ast, T: 'ast>(&'ast self, node: T) -> Type<'ast> {
        Type::from_node(self.alloc(node))
    }

    fn alloc_type_pack_node<'ast, T: 'ast>(&'ast self, node: T) -> TypePack<'ast> {
        TypePack::from_node(self.alloc(node))
    }

    pub fn alloc_type_kind<'ast>(
        &'ast self,
        location: Location,
        kind: TypeKind<'ast>,
    ) -> Type<'ast> {
        match kind {
            TypeKind::Reference {
                prefix,
                prefix_location,
                prefix_local,
                name,
                name_location,
                has_parameter_list,
                parameters,
                ..
            } => self.alloc_type_node(TypeReference {
                base: Type::new_header(TypeTag::Reference, location),
                has_parameter_list,
                prefix,
                prefix_location,
                prefix_local,
                name,
                name_location,
                parameters,
            }),
            TypeKind::Table { props, indexer } => self.alloc_type_node(TypeTable {
                base: Type::new_header(TypeTag::Table, location),
                props,
                indexer,
            }),
            TypeKind::Function {
                attributes,
                generics,
                generic_packs,
                arg_types,
                arg_names,
                return_types,
            } => self.alloc_type_node(TypeFunction {
                base: Type::new_header(TypeTag::Function, location),
                attributes,
                generics,
                generic_packs,
                arg_types,
                arg_names,
                return_types,
            }),
            TypeKind::Typeof { expr } => self.alloc_type_node(TypeTypeof {
                base: Type::new_header(TypeTag::Typeof, location),
                expr,
            }),
            TypeKind::SingletonBool { value } => self.alloc_type_node(TypeSingletonBool {
                base: Type::new_header(TypeTag::SingletonBool, location),
                value,
            }),
            TypeKind::SingletonString { value } => self.alloc_type_node(TypeSingletonString {
                base: Type::new_header(TypeTag::SingletonString, location),
                value,
            }),
            TypeKind::Group { ty } => self.alloc_type_node(TypeGroup {
                base: Type::new_header(TypeTag::Group, location),
                ty,
            }),
            TypeKind::Optional => self.alloc_type_node(TypeUnit {
                base: Type::new_header(TypeTag::Optional, location),
            }),
            TypeKind::Union { types } => self.alloc_type_node(TypeUnion {
                base: Type::new_header(TypeTag::Union, location),
                types,
            }),
            TypeKind::Intersection { types } => self.alloc_type_node(TypeIntersection {
                base: Type::new_header(TypeTag::Intersection, location),
                types,
            }),
            TypeKind::Error {
                types,
                missing,
                message_index,
            } => self.alloc_type_node(TypeError {
                base: Type::new_header(TypeTag::Error, location),
                types,
                missing,
                message_index,
            }),
        }
    }

    pub fn alloc_type_pack_kind<'ast>(
        &'ast self,
        location: Location,
        kind: TypePackKind<'ast>,
    ) -> TypePack<'ast> {
        match kind {
            TypePackKind::Explicit { type_list } => {
                self.alloc_type_pack_node(TypePackExplicitNode {
                    base: TypePack::new_header(TypePackTag::Explicit, location),
                    type_list,
                })
            }
            TypePackKind::Variadic { variadic_type } => {
                self.alloc_type_pack_node(TypePackVariadicNode {
                    base: TypePack::new_header(TypePackTag::Variadic, location),
                    variadic_type,
                })
            }
            TypePackKind::Generic { generic_name } => {
                self.alloc_type_pack_node(TypePackGenericNode {
                    base: TypePack::new_header(TypePackTag::Generic, location),
                    generic_name,
                })
            }
        }
    }
}