luau-syntax 0.732.0

use luau_syntax::allocator::AstArena;
use luau_syntax::ast::{AstNodeRef, AstVisitor, ExpressionKind, StatementTag, Type, TypeKind};
use luau_syntax::ast_names::AstNameTable;
use luau_syntax::parser::{self, ParseOptions, ParseResult};

#[derive(Default)]
struct TrackingVisitor {
    nodes: Vec<&'static str>,
    visit_types: bool,
}

impl TrackingVisitor {
    fn new() -> Self {
        Self::default()
    }

    fn with_types() -> Self {
        Self {
            visit_types: true,
            ..Self::default()
        }
    }
}

impl AstVisitor for TrackingVisitor {
    fn visit_node(&mut self, node: AstNodeRef<'_, '_>) -> bool {
        self.nodes.push(match node {
            AstNodeRef::Attribute(_) => "Attribute",
            AstNodeRef::GenericType(_) => "GenericType",
            AstNodeRef::GenericTypePack(_) => "GenericTypePack",
            AstNodeRef::Statement(statement) => match statement.tag {
                StatementTag::Assign => "StatementTag::Assign",
                StatementTag::CompoundAssign => "StatementTag::CompoundAssign",
                StatementTag::Break => "StatementTag::Break",
                StatementTag::Continue => "StatementTag::Continue",
                StatementTag::Class => "StatementTag::Class",
                StatementTag::Block => "StatementTag::Block",
                StatementTag::Expression => "StatementTag::Expression",
                StatementTag::NumericFor => "StatementTag::NumericFor",
                StatementTag::GenericFor => "StatementTag::GenericFor",
                StatementTag::FunctionDeclaration => "StatementTag::FunctionDeclaration",
                StatementTag::If => "StatementTag::If",
                StatementTag::LocalFunction => "StatementTag::LocalFunction",
                StatementTag::Local => "StatementTag::Local",
                StatementTag::TypeAlias => "StatementTag::TypeAlias",
                StatementTag::TypeFunction => "StatementTag::TypeFunction",
                StatementTag::DeclareGlobal => "StatementTag::DeclareGlobal",
                StatementTag::DeclareFunction => "StatementTag::DeclareFunction",
                StatementTag::DeclareExternType => "StatementTag::DeclareExternType",
                StatementTag::Repeat => "StatementTag::Repeat",
                StatementTag::Return => "StatementTag::Return",
                StatementTag::While => "StatementTag::While",
                StatementTag::Error => "StatementTag::Error",
            },
            AstNodeRef::Expression(expression) => match expression.kind() {
                ExpressionKind::Boolean(_) => "Expression::Boolean",
                ExpressionKind::Call { .. } => "Expression::Call",
                ExpressionKind::FunctionLiteral(_) => "Expression::FunctionLiteral",
                ExpressionKind::Grouped(_) => "Expression::Grouped",
                ExpressionKind::Integer { .. } => "Expression::Integer",
                ExpressionKind::Local { .. } => "Expression::Local",
                ExpressionKind::Nil => "Expression::Nil",
                ExpressionKind::Number { .. } => "Expression::Number",
                ExpressionKind::String { .. } => "Expression::String",
                ExpressionKind::InterpString { .. } => "Expression::InterpString",
                ExpressionKind::Table { .. } => "Expression::Table",
                ExpressionKind::If { .. } => "Expression::If",
                ExpressionKind::Varargs => "Expression::Varargs",
                ExpressionKind::IndexExpr { .. } => "Expression::IndexExpr",
                ExpressionKind::IndexName { .. } => "Expression::IndexName",
                ExpressionKind::TypeAssertion { .. } => "Expression::TypeAssertion",
                ExpressionKind::Instantiate { .. } => "Expression::Instantiate",
                ExpressionKind::Unary { .. } => "Expression::Unary",
                ExpressionKind::Global(_) => "Expression::Global",
                ExpressionKind::Binary { .. } => "Expression::Binary",
                ExpressionKind::Error { .. } => "Expression::Error",
            },
            AstNodeRef::Type(_) => "Type",
            AstNodeRef::TypePack(_) => "TypePack",
        });
        true
    }

    fn visit_type(&mut self, annotation: Type) -> bool {
        if !self.visit_types {
            return false;
        }

        self.nodes.push(match &annotation.kind() {
            TypeKind::Reference { .. } => "TypeKind::Reference",
            TypeKind::Table { .. } => "TypeKind::Table",
            TypeKind::Function { .. } => "TypeKind::Function",
            TypeKind::Typeof { .. } => "TypeKind::Typeof",
            TypeKind::SingletonBool { .. } => "TypeKind::SingletonBool",
            TypeKind::SingletonString { .. } => "TypeKind::SingletonString",
            TypeKind::Group { .. } => "TypeKind::Group",
            TypeKind::Optional => "TypeKind::Optional",
            TypeKind::Union { .. } => "TypeKind::Union",
            TypeKind::Intersection { .. } => "TypeKind::Intersection",
            TypeKind::Error { .. } => "TypeKind::Error",
        });
        true
    }
}

fn with_parse<R>(source: &str, f: impl for<'ast> FnOnce(ParseResult<'ast>) -> R) -> R {
    let arena = AstArena::new();
    let mut names = AstNameTable::new(&arena);
    let result = parser::parse(source, &arena, &mut names, ParseOptions::default()).unwrap();
    f(result)
}

// AstVisitor.test.cpp: TypeAnnotationsAreNotVisited
#[test]
fn type_annotations_are_not_visited_by_default() {
    let mut visitor = TrackingVisitor::new();

    with_parse("local a: A<number>", |result| {
        result.root.visit(&mut visitor)
    });

    assert_eq!(
        visitor.nodes,
        ["StatementTag::Block", "StatementTag::Local"]
    );
}

// AstVisitor.test.cpp: LocalTwoBindings
#[test]
fn local_two_bindings() {
    let mut visitor = TrackingVisitor::new();

    with_parse("local a, b", |result| result.root.visit(&mut visitor));

    assert_eq!(
        visitor.nodes,
        ["StatementTag::Block", "StatementTag::Local"]
    );
}

// AstVisitor.test.cpp: LocalTwoAnnotatedBindings
#[test]
fn local_two_annotated_bindings_visit_types_when_enabled() {
    let mut visitor = TrackingVisitor::with_types();

    with_parse("local a: A, b: B<number>", |result| {
        result.root.visit(&mut visitor)
    });

    assert_eq!(
        visitor.nodes,
        [
            "StatementTag::Block",
            "StatementTag::Local",
            "TypeKind::Reference",
            "TypeKind::Reference",
            "TypeKind::Reference",
        ]
    );
}

// AstVisitor.test.cpp: LocalTwoAnnotatedBindingsWithTwoValues
#[test]
fn local_two_annotated_bindings_with_two_values_visit_types_then_values() {
    let mut visitor = TrackingVisitor::with_types();

    with_parse("local a: A, b: B<number> = 1, 2", |result| {
        result.root.visit(&mut visitor)
    });

    assert_eq!(
        visitor.nodes,
        [
            "StatementTag::Block",
            "StatementTag::Local",
            "TypeKind::Reference",
            "TypeKind::Reference",
            "TypeKind::Reference",
            "Expression::Number",
            "Expression::Number",
        ]
    );
}