nub-lang-archive-2/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;
using Range = OmniSharp.Extensions.LanguageServer.Protocol.Models.Range;

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.Name == variableIdentifierNode.Name);
                if (parameter != null)
                {
                    if (parameter.Tokens.Count == 0)
                    {
                        return null;
                    }

                    return new LocationOrLocationLinks(new Location
                    {
                        Uri = parameter.Tokens.First().Span.FilePath,
                        Range = new Range(parameter.Tokens.First().Span.Start.Line - 1, parameter.Tokens.First().Span.Start.Column - 1, parameter.Tokens.Last().Span.End.Line - 1, parameter.Tokens.Last().Span.End.Column - 1)
                    });
                }

                var variable = function?.Body?
                    .EnumerateDescendantsAndSelf()
                    .OfType<VariableDeclarationNode>()
                    .FirstOrDefault(x => x.Name == variableIdentifierNode.Name);

                if (variable != null)
                {
                    if (variable.Tokens.Count == 0)
                    {
                        return null;
                    }

                    return new LocationOrLocationLinks(new Location
                    {
                        Uri = variable.Tokens.First().Span.FilePath,
                        Range = new Range(variable.Tokens.First().Span.Start.Line - 1, variable.Tokens.First().Span.Start.Column - 1, variable.Tokens.Last().Span.End.Line - 1, variable.Tokens.Last().Span.End.Column - 1)
                    });
                }

                return null;
            }
            case FuncIdentifierNode funcIdentifierNode:
            {
                var functionBody = compilationUnit.ImportedFunctions.FirstOrDefault(x => x.Module == funcIdentifierNode.Module && x.Name == funcIdentifierNode.Name);
                if (functionBody == null || functionBody.Tokens.Count == 0)
                {
                    return null;
                }

                return new LocationOrLocationLinks(new Location
                {
                    Uri = functionBody.Tokens.First().Span.FilePath,
                    Range = new Range(functionBody.Tokens.First().Span.Start.Line - 1, functionBody.Tokens.First().Span.Start.Column - 1, functionBody.Tokens.Last().Span.End.Line - 1, functionBody.Tokens.Last().Span.End.Column - 1)
                });
            }
            default:
            {
                return null;
            }
        }
    }
}