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 compilationUnit = workspaceManager.GetCompilationUnit(uri);
        if (compilationUnit == null)
        {
            return null;
        }

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

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

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

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

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

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

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