tsl-devkit/lsp-server/src/language/symbol/builder.cpp

#include <utility>
#include "./builder.hpp"

namespace lsp::language::symbol
{
    Builder::Builder(SymbolTable& table) : table_(table), current_scope_id_(kInvalidScopeId), in_interface_section_(false)
    {
    }

    void Builder::Build(ast::ASTNode& root)
    {
        root.Accept(*this);
    }

    // ===== Helper Methods =====

    SymbolId Builder::CreateSymbol(const std::string& name, SymbolKind kind, const ast::Location& location, const std::optional<std::string>& type_hint)
    {
        std::optional<UnitVisibility> visibility;
        if (in_interface_section_)
            visibility = UnitVisibility::kInterface;

        Symbol symbol = [&]() -> Symbol {
            switch (kind)
            {
            case SymbolKind::Class:
            {
                Class cls;
                cls.name = name;
                cls.selection_range = location;
                cls.range = location;
                cls.unit_visibility = visibility;
                return Symbol(std::move(cls));
            }
            case SymbolKind::Function:
            {
                Function fn;
                fn.name = name;
                fn.selection_range = location;
                fn.range = location;
                fn.declaration_range = location;
                fn.return_type = type_hint;
                fn.unit_visibility = visibility;
                return Symbol(std::move(fn));
            }
            case SymbolKind::Method:
            {
                Method method;
                method.name = name;
                method.selection_range = location;
                method.range = location;
                method.declaration_range = location;
                method.return_type = type_hint;
                return Symbol(std::move(method));
            }
            case SymbolKind::Property:
            {
                Property property;
                property.name = name;
                property.selection_range = location;
                property.range = location;
                property.type = type_hint;
                return Symbol(std::move(property));
            }
            case SymbolKind::Field:
            {
                Field field;
                field.name = name;
                field.selection_range = location;
                field.range = location;
                field.type = type_hint;
                return Symbol(std::move(field));
            }
            case SymbolKind::Variable:
            {
                Variable var;
                var.name = name;
                var.selection_range = location;
                var.range = location;
                var.type = type_hint;
                var.unit_visibility = visibility;
                return Symbol(std::move(var));
            }
            case SymbolKind::Constant:
            {
                Constant constant;
                constant.name = name;
                constant.selection_range = location;
                constant.range = location;
                constant.type = type_hint;
                return Symbol(std::move(constant));
            }
            default:
            {
                Variable fallback;
                fallback.name = name;
                fallback.selection_range = location;
                fallback.range = location;
                fallback.type = type_hint;
                fallback.unit_visibility = visibility;
                return Symbol(std::move(fallback));
            }
            }
        }();

        SymbolId id = table_.CreateSymbol(std::move(symbol));

        if (current_scope_id_ != kInvalidScopeId)
        {
            table_.AddSymbolToScope(current_scope_id_, name, id);
        }

        return id;
    }

    SymbolId Builder::CreateSymbol(const std::string& name, SymbolKind kind, const ast::ASTNode& node, const std::optional<std::string>& type_hint)
    {
        return CreateSymbol(name, kind, node.span, type_hint);
    }

    SymbolId Builder::CreateFunctionSymbol(
        const std::string& name,
        const ast::Location& location,
        const std::vector<std::unique_ptr<ast::Parameter>>& parameters,
        const std::optional<ast::TypeAnnotation>& return_type)
    {
        std::optional<UnitVisibility> visibility;
        if (in_interface_section_)
            visibility = UnitVisibility::kInterface;

        Function fn;
        fn.name = name;
        fn.selection_range = location;
        fn.range = location;
        fn.declaration_range = location;
        fn.return_type = ExtractTypeName(return_type);
        fn.unit_visibility = visibility;
        fn.parameters = BuildParameters(parameters);

        Symbol symbol(std::move(fn));
        SymbolId id = table_.CreateSymbol(std::move(symbol));

        if (current_scope_id_ != kInvalidScopeId)
        {
            table_.AddSymbolToScope(current_scope_id_, name, id);
        }

        return id;
    }

    SymbolId Builder::CreateMethodSymbol(
        const std::string& name,
        const ast::Location& location,
        const std::vector<std::unique_ptr<ast::Parameter>>& parameters,
        const std::optional<ast::TypeAnnotation>& return_type)
    {
        Method method;
        method.name = name;
        method.selection_range = location;
        method.range = location;
        method.declaration_range = location;
        method.return_type = ExtractTypeName(return_type);
        method.parameters = BuildParameters(parameters);

        Symbol symbol(std::move(method));
        SymbolId id = table_.CreateSymbol(std::move(symbol));

        if (current_scope_id_ != kInvalidScopeId)
        {
            table_.AddSymbolToScope(current_scope_id_, name, id);
        }

        return id;
    }

    ScopeId Builder::EnterScopeWithSymbol(ScopeKind kind, SymbolId symbol_id, const ast::Location& range)
    {
        current_scope_id_ = table_.CreateScope(kind, range, current_scope_id_, symbol_id);
        return current_scope_id_;
    }

    ScopeId Builder::EnterScope(ScopeKind kind, const ast::Location& range)
    {
        current_scope_id_ = table_.CreateScope(kind, range, current_scope_id_, std::nullopt);
        return current_scope_id_;
    }

    void Builder::ExitScope()
    {
        const auto* scope_info = table_.scopes().scope(current_scope_id_);
        if (scope_info && scope_info->parent)
        {
            current_scope_id_ = *scope_info->parent;
        }
    }

    void Builder::VisitStatements(const std::vector<ast::StatementPtr>& statements)
    {
        for (const auto& stmt : statements)
        {
            if (stmt)
                stmt->Accept(*this);
        }
    }

    void Builder::VisitExpression(ast::Expression& expr)
    {
        expr.Accept(*this);
    }

    std::optional<std::string> Builder::ExtractTypeName(const std::optional<ast::TypeAnnotation>& type) const
    {
        if (type)
            return type->name;
        return std::nullopt;
    }

    std::vector<language::symbol::Parameter> Builder::BuildParameters(const std::vector<std::unique_ptr<ast::Parameter>>& parameters) const
    {
        std::vector<language::symbol::Parameter> result;
        result.reserve(parameters.size());

        for (const auto& param : parameters)
        {
            if (!param)
                continue;

            language::symbol::Parameter p;
            p.name = param->name;
            if (param->type)
                p.type = param->type->name;
            if (param->default_value)
                p.default_value = "<expr>";
            result.push_back(std::move(p));
        }

        return result;
    }

    // ===== Node Visitation =====

    void Builder::VisitProgram(ast::Program& node)
    {
        current_scope_id_ = table_.CreateScope(ScopeKind::kGlobal, node.span, std::nullopt, std::nullopt);
        VisitStatements(node.statements);
    }

    void Builder::VisitUnitDefinition(ast::UnitDefinition& node)
    {
        auto unit_scope = EnterScope(ScopeKind::kUnit, node.span);

        // Process interface section
        in_interface_section_ = true;
        VisitStatements(node.interface_statements);

        // Process implementation section
        in_interface_section_ = false;
        VisitStatements(node.implementation_statements);

        ExitScope();
    }

    void Builder::VisitClassDefinition(ast::ClassDefinition& node)
    {
        auto class_id = CreateSymbol(node.name, SymbolKind::Class, node.location);

        auto class_scope = EnterScopeWithSymbol(ScopeKind::kClass, class_id, node.span);

        auto prev_parent = current_parent_symbol_id_;
        current_parent_symbol_id_ = class_id;

        for (auto& member : node.members)
        {
            if (member)
                member->Accept(*this);
        }

        current_parent_symbol_id_ = prev_parent;
        ExitScope();
    }

    void Builder::VisitFunctionDeclaration(ast::FunctionDeclaration& node)
    {
        auto location = ast::Location{};
        if (!node.name.empty() && !node.parameters.empty())
        {
            location = node.span;
        }
        CreateFunctionSymbol(node.name, location, node.parameters, node.return_type);
    }

    void Builder::VisitFunctionDefinition(ast::FunctionDefinition& node)
    {
        auto func_id = CreateFunctionSymbol(node.name, node.location, node.parameters, node.return_type);

        if (node.body)
        {
            auto func_scope = EnterScopeWithSymbol(ScopeKind::kFunction, func_id, node.body->span);

            auto prev_function = current_function_id_;
            current_function_id_ = func_id;

            for (auto& param : node.parameters)
            {
                if (param)
                {
                    CreateSymbol(param->name, SymbolKind::Variable, param->location);
                }
            }

            VisitStatements(node.body->statements);

            current_function_id_ = prev_function;
            ExitScope();
        }
    }

    void Builder::VisitMethodDeclaration(ast::MethodDeclaration& node)
    {
        auto method_id = CreateMethodSymbol(node.name, node.location, node.parameters, node.return_type);

        if (node.body)
        {
            auto method_scope = EnterScopeWithSymbol(ScopeKind::kFunction, method_id, node.body->span);

            auto prev_function = current_function_id_;
            current_function_id_ = method_id;

            for (auto& param : node.parameters)
            {
                if (param)
                {
                    CreateSymbol(param->name, SymbolKind::Variable, param->location);
                }
            }

            VisitStatements(node.body->statements);

            current_function_id_ = prev_function;
            ExitScope();
        }
    }

    void Builder::VisitPropertyDeclaration(ast::PropertyDeclaration& node)
    {
        auto type = ExtractTypeName(node.type);
        CreateSymbol(node.name, SymbolKind::Property, node.location, type);
    }

    void Builder::VisitExternalMethodDefinition(ast::ExternalMethodDefinition& node)
    {
        std::string method_name = node.name;
        size_t dot_pos = method_name.find_last_of('.');
        if (dot_pos != std::string::npos)
        {
            method_name = method_name.substr(dot_pos + 1);
        }

        std::optional<SymbolId> method_id;

        if (current_parent_symbol_id_)
        {
            const auto* scope_info = table_.scopes().scope(current_scope_id_);
            if (scope_info)
            {
                method_id = table_.scopes().FindSymbolInScope(current_scope_id_, method_name);
            }
        }

        if (!method_id)
        {
            method_id = CreateMethodSymbol(method_name, node.location, node.parameters, node.return_type);
        }

        if (node.body)
        {
            auto method_scope = EnterScopeWithSymbol(ScopeKind::kFunction, *method_id, node.body->span);

            auto prev_function = current_function_id_;
            current_function_id_ = *method_id;

            for (auto& param : node.parameters)
            {
                if (param)
                {
                    CreateSymbol(param->name, SymbolKind::Variable, param->location);
                }
            }

            VisitStatements(node.body->statements);

            current_function_id_ = prev_function;
            ExitScope();
        }
    }

    void Builder::VisitClassMember(ast::ClassMember& node)
    {
        std::visit([this](auto& member) {
            if (member)
                member->Accept(*this);
        },
                   node.member);
    }

    void Builder::VisitVarDeclaration(ast::VarDeclaration& node)
    {
        CreateSymbol(node.name, SymbolKind::Variable, node.location, ExtractTypeName(node.type));

        if (node.initializer)
        {
            VisitExpression(*node.initializer.value());
        }
    }

    void Builder::VisitStaticDeclaration(ast::StaticDeclaration& node)
    {
        CreateSymbol(node.name, SymbolKind::Variable, node.location, ExtractTypeName(node.type));

        if (node.initializer)
        {
            VisitExpression(*node.initializer.value());
        }
    }

    void Builder::VisitGlobalDeclaration(ast::GlobalDeclaration& node)
    {
        CreateSymbol(node.name, SymbolKind::Variable, node.location, ExtractTypeName(node.type));

        if (node.initializer)
        {
            VisitExpression(*node.initializer.value());
        }
    }

    void Builder::VisitConstDeclaration(ast::ConstDeclaration& node)
    {
        CreateSymbol(node.name, SymbolKind::Constant, node.location, ExtractTypeName(node.type));

        if (node.value)
        {
            VisitExpression(*node.value);
        }
    }

    void Builder::VisitFieldDeclaration(ast::FieldDeclaration& node)
    {
        CreateSymbol(node.name, SymbolKind::Field, node.location, ExtractTypeName(node.type));

        if (node.initializer)
        {
            VisitExpression(*node.initializer.value());
        }
    }

    void Builder::VisitUsesStatement(ast::UsesStatement& node)
    {
        // Unit imports would be stored in the Unit symbol
        // For now just skip since we don't have a way to store them
    }

    // ===== Expression Processing (references) =====

    void Builder::VisitIdentifier(ast::Identifier& node)
    {
        auto symbol_id = table_.scopes().FindSymbolInScopeChain(current_scope_id_, node.name);
        if (symbol_id)
        {
            table_.AddReference(*symbol_id, node.location, false, false);
        }
    }

    void Builder::VisitCallExpression(ast::CallExpression& node)
    {
        if (node.callee)
        {
            if (auto* id = dynamic_cast<ast::Identifier*>(node.callee.get()))
            {
                auto symbol_id = table_.scopes().FindSymbolInScopeChain(current_scope_id_, id->name);
                if (symbol_id)
                {
                    table_.AddReference(*symbol_id, id->location, false, false);

                    if (current_function_id_)
                    {
                        table_.AddCall(*current_function_id_, *symbol_id, node.span);
                    }
                }
            }
            else
            {
                node.callee->Accept(*this);
            }
        }

        for (auto& arg : node.arguments)
        {
            if (arg.value)
                VisitExpression(*arg.value);
        }
    }

    void Builder::VisitAttributeExpression(ast::AttributeExpression& node)
    {
        if (node.object)
        {
            node.object->Accept(*this);
        }

        if (node.attribute)
        {
            node.attribute->Accept(*this);
        }
    }

    void Builder::VisitAssignmentExpression(ast::AssignmentExpression& node)
    {
        ProcessLValue(node.left, true);

        if (node.right)
        {
            VisitExpression(*node.right);
        }
    }

    void Builder::ProcessLValue(const ast::LValue& lvalue, bool is_write)
    {
        std::visit([this, is_write](auto& val) {
            using T = std::decay_t<decltype(val)>;
            if constexpr (std::is_same_v<T, std::unique_ptr<ast::Identifier>>)
            {
                if (val)
                {
                    auto symbol_id = table_.scopes().FindSymbolInScopeChain(current_scope_id_, val->name);
                    if (symbol_id)
                    {
                        table_.AddReference(*symbol_id, val->location, false, is_write);
                    }
                }
            }
            else if constexpr (std::is_same_v<T, std::unique_ptr<ast::AttributeExpression>>)
            {
                if (val)
                {
                    if (val->object)
                    {
                        val->object->Accept(*this);
                    }
                    if (val->attribute)
                    {
                        val->attribute->Accept(*this);
                    }
                }
            }
            else if constexpr (std::is_same_v<T, std::unique_ptr<ast::SubscriptExpression>>)
            {
                if (val)
                {
                    if (val->base)
                    {
                        val->base->Accept(*this);
                    }
                    for (auto& idx : val->indices)
                    {
                        if (idx.start)
                            idx.start->Accept(*this);
                        if (idx.end)
                            idx.end->Accept(*this);
                        if (idx.step)
                            idx.step->Accept(*this);
                    }
                }
            }
            else if constexpr (std::is_same_v<T, std::unique_ptr<ast::ColumnReference>>)
            {
                if (val && val->value)
                {
                    val->value->Accept(*this);
                }
            }
        },
                   lvalue);
    }

    // ===== Statement Processing =====

    void Builder::VisitBlockStatement(ast::BlockStatement& node)
    {
        auto block_scope = EnterScope(ScopeKind::kBlock, node.span);
        VisitStatements(node.statements);
        ExitScope();
    }

    void Builder::VisitIfStatement(ast::IfStatement& node)
    {
        for (auto& branch : node.branches)
        {
            if (branch.condition)
                VisitExpression(*branch.condition);

            if (branch.body)
                branch.body->Accept(*this);
        }

        if (node.else_body && *node.else_body)
        {
            (*node.else_body)->Accept(*this);
        }
    }

    void Builder::VisitForInStatement(ast::ForInStatement& node)
    {
        auto for_scope = EnterScope(ScopeKind::kBlock, node.span);

        if (!node.key.empty())
        {
            CreateSymbol(node.key, SymbolKind::Variable, node.key_location);
        }

        if (!node.value.empty())
        {
            CreateSymbol(node.value, SymbolKind::Variable, node.value_location);
        }

        if (node.collection)
            VisitExpression(*node.collection);

        if (node.body)
            node.body->Accept(*this);

        ExitScope();
    }

    void Builder::VisitForToStatement(ast::ForToStatement& node)
    {
        auto for_scope = EnterScope(ScopeKind::kBlock, node.span);

        CreateSymbol(node.counter, SymbolKind::Variable, node.counter_location);

        if (node.start)
            VisitExpression(*node.start);
        if (node.end)
            VisitExpression(*node.end);
        if (node.step)
            VisitExpression(*node.step);
        if (node.body)
            node.body->Accept(*this);

        ExitScope();
    }

    void Builder::VisitWhileStatement(ast::WhileStatement& node)
    {
        if (node.condition)
            VisitExpression(*node.condition);
        if (node.body)
            node.body->Accept(*this);
    }

    void Builder::VisitRepeatStatement(ast::RepeatStatement& node)
    {
        VisitStatements(node.body);
        if (node.condition)
            VisitExpression(*node.condition);
    }

    void Builder::VisitCaseStatement(ast::CaseStatement& node)
    {
        if (node.discriminant)
            VisitExpression(*node.discriminant);

        for (auto& branch : node.branches)
        {
            for (auto& value : branch.values)
            {
                if (value)
                    VisitExpression(*value);
            }
            if (branch.body)
                branch.body->Accept(*this);
        }

        if (node.else_body && *node.else_body)
        {
            (*node.else_body)->Accept(*this);
        }
    }

    void Builder::VisitTryStatement(ast::TryStatement& node)
    {
        if (node.try_body)
        {
            VisitStatements(node.try_body->statements);
        }

        if (node.except_body)
        {
            VisitStatements(node.except_body->statements);
        }
    }

    void Builder::VisitMatrixIterationStatement(ast::MatrixIterationStatement& node)
    {
        auto iter_scope = EnterScope(ScopeKind::kBlock, node.span);

        if (node.target)
            VisitExpression(*node.target);

        if (node.body)
            VisitStatements(node.body->statements);

        ExitScope();
    }

    void Builder::VisitExpressionStatement(ast::ExpressionStatement& node)
    {
        if (node.expression)
            VisitExpression(*node.expression);
    }

    void Builder::VisitReturnStatement(ast::ReturnStatement& node)
    {
        if (node.value && *node.value)
            VisitExpression(**node.value);
    }

    // ===== Other Expressions =====

    void Builder::VisitBinaryExpression(ast::BinaryExpression& node)
    {
        if (node.left)
            VisitExpression(*node.left);
        if (node.right)
            VisitExpression(*node.right);
    }

    void Builder::VisitTernaryExpression(ast::TernaryExpression& node)
    {
        if (node.condition)
            VisitExpression(*node.condition);
        if (node.consequence)
            VisitExpression(*node.consequence);
        if (node.alternative)
            VisitExpression(*node.alternative);
    }

    void Builder::VisitSubscriptExpression(ast::SubscriptExpression& node)
    {
        if (node.base)
            VisitExpression(*node.base);

        for (auto& idx : node.indices)
        {
            if (idx.start)
                VisitExpression(*idx.start);
            if (idx.end)
                VisitExpression(*idx.end);
            if (idx.step)
                VisitExpression(*idx.step);
        }
    }

    void Builder::VisitArrayExpression(ast::ArrayExpression& node)
    {
        for (auto& elem : node.elements)
        {
            if (elem.key)
                VisitExpression(*elem.key.value());
            if (elem.value)
                VisitExpression(*elem.value);
        }
    }

    void Builder::VisitAnonymousFunctionExpression(ast::AnonymousFunctionExpression& node)
    {
        auto func_id = CreateFunctionSymbol("<anonymous>", node.span, node.parameters, node.return_type);

        if (node.body)
        {
            auto func_scope = EnterScopeWithSymbol(ScopeKind::kAnonymousFunction, func_id, node.body->span);

            auto prev_function = current_function_id_;
            current_function_id_ = func_id;

            for (auto& param : node.parameters)
            {
                if (param)
                {
                    CreateSymbol(param->name, SymbolKind::Variable, param->location);
                }
            }

            VisitStatements(node.body->statements);

            current_function_id_ = prev_function;
            ExitScope();
        }
    }

    void Builder::VisitUnaryPlusExpression(ast::UnaryPlusExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitUnaryMinusExpression(ast::UnaryMinusExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitPrefixIncrementExpression(ast::PrefixIncrementExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitPrefixDecrementExpression(ast::PrefixDecrementExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitPostfixIncrementExpression(ast::PostfixIncrementExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitPostfixDecrementExpression(ast::PostfixDecrementExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitLogicalNotExpression(ast::LogicalNotExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitBitwiseNotExpression(ast::BitwiseNotExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitDerivativeExpression(ast::DerivativeExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitMatrixTransposeExpression(ast::MatrixTransposeExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitExprOperatorExpression(ast::ExprOperatorExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitFunctionPointerExpression([[maybe_unused]] ast::FunctionPointerExpression& node)
    {
        if (node.argument)
            VisitExpression(*node.argument);
    }

    void Builder::VisitNewExpression(ast::NewExpression& node)
    {
        if (node.target)
            VisitExpression(*node.target);
    }

    void Builder::VisitEchoExpression(ast::EchoExpression& node)
    {
        for (auto& expr : node.expressions)
        {
            if (expr)
                VisitExpression(*expr);
        }
    }

    void Builder::VisitRaiseExpression(ast::RaiseExpression& node)
    {
        if (node.exception)
            VisitExpression(*node.exception);
    }

    void Builder::VisitInheritedExpression(ast::InheritedExpression& node)
    {
        if (node.call && *node.call)
        {
            (*node.call)->Accept(*this);
        }
    }

    void Builder::VisitParenthesizedExpression(ast::ParenthesizedExpression& node)
    {
        for (auto& elem : node.elements)
        {
            if (elem.key)
                VisitExpression(*elem.key.value());
            if (elem.value)
                VisitExpression(*elem.value);
        }
    }

    void Builder::VisitColumnReference(ast::ColumnReference& node)
    {
        if (node.value)
        {
            VisitExpression(*node.value);
        }
    }

    void Builder::VisitUnpackPattern([[maybe_unused]] ast::UnpackPattern& node)
    {
        // Unpack pattern processing - would need to handle LValue variants
    }

    void Builder::VisitCompilerDirective([[maybe_unused]] ast::CompilerDirective& node)
    {
        // Compiler directive processing
    }

    void Builder::VisitConditionalDirective([[maybe_unused]] ast::ConditionalDirective& node)
    {
        // Conditional compilation directive - no statements to visit
    }

    void Builder::VisitConditionalBlock(ast::ConditionalBlock& node)
    {
        VisitStatements(node.consequence);
        VisitStatements(node.alternative);
    }

} // namespace lsp::language::symbol