Added Handler for Semantic Tokenization

src/PowerShellEditorServices/Server/PsesLanguageServer.cs

@@ -92,6 +92,7 @@ public async Task StartAsync()
                     .WithHandler<GetCommandHandler>()
                     .WithHandler<ShowHelpHandler>()
                     .WithHandler<ExpandAliasHandler>()
+                    .WithHandler<PsesSemanticTokens>()
                     .OnInitialize(
                         async (languageServer, request, cancellationToken) =>
                         {

src/PowerShellEditorServices/Services/TextDocument/Handlers/PsesSemanticTokens.cs

@@ -0,0 +1,145 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Management.Automation.Language;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+using Microsoft.PowerShell.EditorServices.Services;
+using Microsoft.PowerShell.EditorServices.Services.TextDocument;
+using Microsoft.PowerShell.EditorServices.Utility;
+using OmniSharp.Extensions.LanguageServer.Protocol;
+using OmniSharp.Extensions.LanguageServer.Protocol.Document.Proposals;
+using OmniSharp.Extensions.LanguageServer.Protocol.Models;
+using OmniSharp.Extensions.LanguageServer.Protocol.Models.Proposals;
+
+namespace Microsoft.PowerShell.EditorServices.Handlers
+{
+    //SemanticTokensHandler is labeled "Obsolete" because that is how Omnisharp marks proposed LSP features. Since we want this proposed feature, we disable this warning.
+#pragma warning disable 618
+    internal class PsesSemanticTokens : SemanticTokensHandler
+    {
+        private readonly ILogger _logger;
+        private readonly WorkspaceService _workspaceService;
+        static readonly SemanticTokensRegistrationOptions _registrationOptions = new SemanticTokensRegistrationOptions() {
+            DocumentSelector = LspUtils.PowerShellDocumentSelector,
+            Legend = new SemanticTokensLegend(),
+            DocumentProvider = new Supports<SemanticTokensDocumentProviderOptions>(isSupported: true,
+                new SemanticTokensDocumentProviderOptions {
+                    Edits = true
+                }),
+            RangeProvider = true
+        };
+
+        public PsesSemanticTokens(ILogger<PsesSemanticTokens> logger, WorkspaceService workspaceService) : base(_registrationOptions)
+        {
+            _logger = logger;
+            _workspaceService = workspaceService;
+        }
+
+        protected override Task Tokenize(SemanticTokensBuilder builder, ITextDocumentIdentifierParams identifier,
+            CancellationToken cancellationToken)
+        {
+            ScriptFile file = _workspaceService.GetFile(DocumentUri.GetFileSystemPath(identifier));
+            Token[] tokens = file.ScriptTokens;
+            foreach (var token in tokens){
+                PushToken(token, builder);
+            }
+            return Task.CompletedTask;
+        }
+
+        private static void PushToken(Token token, SemanticTokensBuilder builder)
+        {
+            if(token is StringExpandableToken stringExpandableToken)
+            {
+                // Try parsing tokens within the string
+                if (stringExpandableToken.NestedTokens != null)
+                {
+                    foreach (Token t in stringExpandableToken.NestedTokens)
+                    {
+                        PushToken(t, builder);
+                    }
+                    return;
+                }
+            }
+
+            //Tokens line and col numbers indexed starting from 1, expecting indexing from 0
+            int line = token.Extent.StartLineNumber - 1;
+            int index = token.Extent.StartColumnNumber - 1;
+
+            builder.Push(line: line, @char: index, length: token.Text.Length,
+                tokenType: MapSemanticToken(token), tokenModifiers: Array.Empty<string>());
+        }
+
+        private static SemanticTokenType MapSemanticToken(Token token)
+        {
+            // First check token flags
+            if ((token.TokenFlags & TokenFlags.Keyword) != 0)
+            {
+                return SemanticTokenType.Keyword;
+            }
+
+            if ((token.TokenFlags & TokenFlags.CommandName) != 0)
+            {
+                return SemanticTokenType.Function;
+            }
+
+            if (token.Kind != TokenKind.Generic && (token.TokenFlags &
+                (TokenFlags.BinaryOperator | TokenFlags.UnaryOperator | TokenFlags.AssignmentOperator)) != 0)
+            {
+                return SemanticTokenType.Operator;
+            }
+
+            if ((token.TokenFlags & TokenFlags.TypeName) != 0)
+            {
+                return SemanticTokenType.Type;
+            }
+
+            if ((token.TokenFlags & TokenFlags.MemberName) != 0)
+            {
+                return SemanticTokenType.Member;
+            }
+
+            // Only check token kind after checking flags
+            switch (token.Kind)
+            {
+                case TokenKind.Comment:
+                    return SemanticTokenType.Comment;
+
+                case TokenKind.Parameter:
+                case TokenKind.Generic when token is StringLiteralToken slt && slt.Text.StartsWith("--"):
+                    return SemanticTokenType.Parameter;
+
+                case TokenKind.Variable:
+                case TokenKind.SplattedVariable:
+                    return SemanticTokenType.Variable;
+
+                case TokenKind.StringExpandable:
+                case TokenKind.StringLiteral:
+                case TokenKind.HereStringExpandable:
+                case TokenKind.HereStringLiteral:
+                    return SemanticTokenType.String;
+
+                case TokenKind.Number:
+                    return SemanticTokenType.Number;
+
+                case TokenKind.Generic:
+                    return SemanticTokenType.Function;
+            }
+
+            // Default semantic token
+            return SemanticTokenType.Documentation;
+        }
+
+        protected override Task<SemanticTokensDocument> GetSemanticTokensDocument(
+            ITextDocumentIdentifierParams @params,
+            CancellationToken cancellationToken)
+        {
+            return Task.FromResult(new SemanticTokensDocument(GetRegistrationOptions().Legend));
+        }
+    }
+#pragma warning restore 618
+}

test/PowerShellEditorServices.Test.E2E/LanguageServerProtocolMessageTests.cs

@@ -19,9 +19,12 @@
 using OmniSharp.Extensions.LanguageServer.Protocol.Document;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using OmniSharp.Extensions.LanguageServer.Protocol.Workspace;
+using OmniSharp.Extensions.LanguageServer.Protocol.Models.Proposals;
 using Xunit;
 using Xunit.Abstractions;
 using Range = OmniSharp.Extensions.LanguageServer.Protocol.Models.Range;
+using OmniSharp.Extensions.LanguageServer.Protocol.Document.Proposals;
+using Microsoft.PowerShell.EditorServices.Utility;
 
 namespace PowerShellEditorServices.Test.E2E
 {
@@ -1136,5 +1139,36 @@ await PsesLanguageClient
 
             Assert.Equal("Get-ChildItem", expandAliasResult.Text);
         }
+
+#pragma warning disable 618
+         [Fact]
+        public async Task CanSendSemanticTokenRequest()
+        {
+
+            //var tokens = semanticTokensBuilder.Commit().GetSemanticTokens();
+            string scriptPath = NewTestFile(@"
+function add {}
+");
+            var result =
+                await PsesLanguageClient
+                    .SendRequest<SemanticTokensParams>(
+                        "textDocument/semanticTokens",
+                        new SemanticTokensParams
+                        {
+                            TextDocument = new TextDocumentIdentifier
+                            {
+                                Uri = new Uri(scriptPath)
+                            }
+                        })
+                    .Returning<SemanticTokens>(CancellationToken.None);
+
+            scriptPath = null;
+            // Information about how this data is generated can be found at
+            // https://github.com/microsoft/vscode-extension-samples/blob/5ae1f7787122812dcc84e37427ca90af5ee09f14/semantic-tokens-sample/vscode.proposed.d.ts#L71
+            var arr = new int[35]{0,0,2,0,0,1,0,8,2,0,0,9,3,0,0,0,4,1,0,0,0,1,1,0,0,0,1,2,0,0,1,1,0,0,0};
+            Assert.Equal(result.Data.ToArray(), arr);
+        }
+
+#pragma warning restore 618
     }
 }