nub-lang/compiler/NubLang.LSP/DefinitionHandler.cs

using NubLang.Ast;
using OmniSharp.Extensions.LanguageServer.Protocol.Client.Capabilities;
using OmniSharp.Extensions.LanguageServer.Protocol.Document;
using OmniSharp.Extensions.LanguageServer.Protocol.Models;

namespace NubLang.LSP;

internal class DefinitionHandler(WorkspaceManager workspaceManager) : DefinitionHandlerBase
{
    protected override DefinitionRegistrationOptions CreateRegistrationOptions(DefinitionCapability capability, ClientCapabilities clientCapabilities)
    {
        return new DefinitionRegistrationOptions();
    }

    public override Task<LocationOrLocationLinks?> Handle(DefinitionParams request, CancellationToken cancellationToken)
    {
        return Task.FromResult(HandleSync(request, cancellationToken));
    }

    private LocationOrLocationLinks? HandleSync(DefinitionParams request, CancellationToken cancellationToken)
    {
        var uri = request.TextDocument.Uri;
        var topLevelNodes = workspaceManager.GetTopLevelNodes(uri.GetFileSystemPath());

        var line = request.Position.Line;
        var character = request.Position.Character;

        var node = topLevelNodes.DeepestNodeAtPosition(line, character);

        switch (node)
        {
            case VariableIdentifierNode variableIdentifierNode:
            {
                var funcNode = topLevelNodes.FunctionAtPosition(line, character);

                var parameter = funcNode?.Prototype.Parameters.FirstOrDefault(x => x.NameToken.Value == variableIdentifierNode.NameToken.Value);
                if (parameter != null)
                {
                    return new LocationOrLocationLinks(parameter.NameToken.ToLocation());
                }

                var variable = funcNode?.Body?
                    .Descendants()
                    .OfType<VariableDeclarationNode>()
                    .FirstOrDefault(x => x.NameToken.Value == variableIdentifierNode.NameToken.Value);

                if (variable != null)
                {
                    return new LocationOrLocationLinks(variable.NameToken.ToLocation());
                }

                return null;
            }
            case LocalFuncIdentifierNode localFuncIdentifierNode:
            {
                var funcNode = topLevelNodes.OfType<FuncNode>().FirstOrDefault(x => x.NameToken.Value == localFuncIdentifierNode.NameToken.Value);
                if (funcNode != null)
                {
                    return new LocationOrLocationLinks(funcNode.NameToken.ToLocation());
                }

                return null;
            }
            case ModuleFuncIdentifierNode localFuncIdentifierNode:
            {
                var repository = workspaceManager.GetModuleRepository();
                if (!repository.TryGet(localFuncIdentifierNode.ModuleToken, out var module))
                {
                    return null;
                }

                if (module.TryResolveFunc(localFuncIdentifierNode.NameToken, out var func, out _))
                {
                    return new LocationOrLocationLinks(func.NameToken.ToLocation());
                }

                return null;
            }
            default:
            {
                return null;
            }
        }
    }
}