nub-lang-archive-2/compiler/NubLang.LSP/CompletionHandler.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 CompletionHandler(WorkspaceManager workspaceManager) : CompletionHandlerBase
{
    private readonly CompletionItem[] _definitionSnippets =
    [
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "func",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "func ${1:name}(${2:params})\n{\n    $0\n}",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "struct",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "struct ${1:name}\n{\n    $0\n}",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "module",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "module \"$0\"",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "import",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "import \"$0\"",
        }
    ];

    private readonly CompletionItem[] _statementSnippets =
    [
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "let",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "let ${1:name} = $0",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "if",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "if ${1:condition}\n{\n    $0\n}",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "else if",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "else if ${1:condition}\n{\n    $0\n}",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "else",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "else\n{\n    $0\n}",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "while",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "while ${1:condition}\n{\n    $0\n}",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "for",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "for ${1:name}, ${2:index} in ${3:array}\n{\n    $0\n}",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "return",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "return $0",
        },
        new()
        {
            Kind = CompletionItemKind.Keyword,
            Label = "defer",
            InsertTextFormat = InsertTextFormat.Snippet,
            InsertText = "defer $0",
        }
    ];

    protected override CompletionRegistrationOptions CreateRegistrationOptions(CompletionCapability capability, ClientCapabilities clientCapabilities)
    {
        return new CompletionRegistrationOptions();
    }

    public override Task<CompletionList> Handle(CompletionParams request, CancellationToken cancellationToken)
    {
        return Task.FromResult(HandleSync(request, cancellationToken));
    }

    private CompletionList HandleSync(CompletionParams request, CancellationToken cancellationToken)
    {
        var completions = new List<CompletionItem>();
        var position = request.Position;

        var uri = request.TextDocument.Uri;
        var compilationUnit = workspaceManager.GetCompilationUnit(uri);
        if (compilationUnit != null)
        {
            var function = compilationUnit.Functions.FirstOrDefault(x => x.Body != null && x.Body.ContainsPosition(position.Line, position.Character));
            if (function != null)
            {
                completions.AddRange(_statementSnippets);

                foreach (var prototype in compilationUnit.ImportedFunctions)
                {
                    var parameterStrings = new List<string>();
                    foreach (var (index, parameter) in prototype.Parameters.Index())
                    {
                        parameterStrings.AddRange($"${{{index + 1}:{parameter.Name}}}");
                    }

                    completions.Add(new CompletionItem
                    {
                        Kind = CompletionItemKind.Function,
                        Label = $"{prototype.Module}::{prototype.Name}",
                        InsertTextFormat = InsertTextFormat.Snippet,
                        InsertText = $"{prototype.Module}::{prototype.Name}({string.Join(", ", parameterStrings)})",
                    });
                }

                foreach (var parameter in function.Prototype.Parameters)
                {
                    completions.Add(new CompletionItem
                    {
                        Kind = CompletionItemKind.Variable,
                        Label = parameter.Name,
                        InsertText = parameter.Name,
                    });
                }

                var variables = function.Body!
                    .Descendants()
                    .OfType<VariableDeclarationNode>();

                foreach (var variable in variables)
                {
                    completions.Add(new CompletionItem
                    {
                        Kind = CompletionItemKind.Variable,
                        Label = variable.Name,
                        InsertText = variable.Name,
                    });
                }
            }
            else
            {
                completions.AddRange(_definitionSnippets);
            }
        }

        return new CompletionList(completions, false);
    }

    public override Task<CompletionItem> Handle(CompletionItem request, CancellationToken cancellationToken)
    {
        return Task.FromResult(new CompletionItem());
    }
}