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,175 @@
+//
+// 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.Collections.Generic;
+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
+{
+    internal class PsesSemanticTokens : SemanticTokensHandler
+    {
+        private readonly ILogger _logger;
+        private readonly WorkspaceService _workspaceService;
+        private static readonly SemanticTokensRegistrationOptions s_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(s_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)
+        {
+
+            List<SemanticToken> semanticToken = ConvertToSemanticTokens(token);
+            foreach(SemanticToken sToken in semanticToken)
+            {
+                builder.Push(sToken.Line, sToken.Index, length: sToken.Text.Length,
+                    sToken.Type, tokenModifiers: sToken.TokenModifiers);
+            }
+        }
+
+        internal static List<SemanticToken> ConvertToSemanticTokens(Token token){
+            if(token is StringExpandableToken stringExpandableToken)
+            {
+                // Try parsing tokens within the string
+                if (stringExpandableToken.NestedTokens != null)
+                {
+                    List<SemanticToken> tokens = new List<SemanticToken>();
+                    foreach (Token t in stringExpandableToken.NestedTokens)
+                    {
+                        tokens.AddRange(ConvertToSemanticTokens(t));
+                    }
+                    return tokens;
+                }
+            }
+
+            SemanticTokenType mappedType =  MapSemanticTokenType(token);
+            if(mappedType == null){
+                return new List<SemanticToken>();
+            }
+
+            //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;
+            SemanticToken sToken = new SemanticToken(token.Text, mappedType,
+                line, index, Array.Empty<string>());
+            return new List<SemanticToken>{sToken};
+        }
+
+        private static SemanticTokenType MapSemanticTokenType(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;
+            }
+
+            return null;
+        }
+
+        protected override Task<SemanticTokensDocument> GetSemanticTokensDocument(
+            ITextDocumentIdentifierParams @params,
+            CancellationToken cancellationToken)
+        {
+            return Task.FromResult(new SemanticTokensDocument(GetRegistrationOptions().Legend));
+        }
+    }
+}
+
+class SemanticToken
+{
+    public SemanticToken(string text, SemanticTokenType type, int line, int index, IEnumerable<string> tokenModifiers)
+    {
+        Line = line;
+        Text = text;
+        Index = index;
+        Type = type;
+        TokenModifiers = tokenModifiers;
+    }
+    public string Text {get; set;}
+    public int Line {get; set;}
+    public int Index {get; set;}
+    public SemanticTokenType Type {get; set;}
+    public IEnumerable<string> TokenModifiers {get; set;}
+}

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);
         }
+
+        [Fact]
+        public async Task CanSendSemanticTokenRequest()
+        {
+            string scriptContent = "function";
+            string scriptPath = NewTestFile(scriptContent);
+
+            SemanticTokens result =
+                await PsesLanguageClient
+                    .SendRequest<SemanticTokensParams>(
+                        "textDocument/semanticTokens",
+                        new SemanticTokensParams
+                        {
+                            TextDocument = new TextDocumentIdentifier
+                            {
+                                Uri = new Uri(scriptPath)
+                            }
+                        })
+                    .Returning<SemanticTokens>(CancellationToken.None);
+
+            // More 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[10]
+                {
+                    //line, index, token length, token type, token modifiers
+                    0, 0, scriptContent.Length, 2, 0, //token 1 is function
+                    0, 9, 0, 0, 0 //token 2 is EOF
+                };
+
+            Assert.Equal(result.Data.ToArray(), arr);
+        }
     }
 }

test/PowerShellEditorServices.Test/Language/SemanticTokenTest.cs

@@ -0,0 +1,93 @@
+//
+// 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.Collections.Generic;
+using System.IO;
+using System.Management.Automation.Language;
+using System.Threading.Tasks;
+using Microsoft.PowerShell.EditorServices.Services.TextDocument;
+using Microsoft.PowerShell.EditorServices.Handlers;
+using OmniSharp.Extensions.LanguageServer.Protocol;
+using OmniSharp.Extensions.LanguageServer.Protocol.Document.Proposals;
+using OmniSharp.Extensions.LanguageServer.Protocol.Models.Proposals;
+using Xunit;
+
+namespace Microsoft.PowerShell.EditorServices.Test.Language
+{
+    public class SemanticTokenTest
+    {
+        [Fact]
+        public async Task MapsFunctionElements()
+        {
+            string text = @"
+function Get-Sum {
+    param( [int]$a, [int]$b )
+    return $a + $b
+}
+";
+            ScriptFile scriptFile = new ScriptFile(
+                // Use any absolute path. Even if it doesn't exist.
+                DocumentUri.FromFileSystemPath(Path.Combine(Path.GetTempPath(), "TestFile.ps1")),
+                text,
+                Version.Parse("5.0"));
+
+            foreach(Token t in scriptFile.ScriptTokens)
+            {
+                List<SemanticToken> mappedTokens = PsesSemanticTokens.ConvertToSemanticTokens(t);
+                switch(t.Text)
+                {
+                    case "function":
+                    case "param":
+                    case "return":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Keyword, mappedTokens[0].Type);
+                        break;
+                    case "Get-Sum":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Function, mappedTokens[0].Type);
+                        break;
+                    case "$a":
+                    case "$b":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Variable, mappedTokens[0].Type);
+                        break;
+                    case "[int]":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Type, mappedTokens[0].Type);
+                        break;
+                    case "+":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Operator, mappedTokens[0].Type);
+                        break;
+                }
+            }
+        }
+
+        [Fact]
+        public async Task TokenizesStringExpansion()
+        {
+            string text = "Write-Host \"$(Test-Property Get-Whatever) $(Get-Whatever)\"";
+            ScriptFile scriptFile = new ScriptFile(
+                // Use any absolute path. Even if it doesn't exist.
+                DocumentUri.FromFileSystemPath(Path.Combine(Path.GetTempPath(), "TestFile.ps1")),
+                text,
+                Version.Parse("5.0"));
+
+            Token commandToken = scriptFile.ScriptTokens[0];
+            List<SemanticToken> mappedTokens = PsesSemanticTokens.ConvertToSemanticTokens(commandToken);
+            Assert.Single(mappedTokens);
+            Assert.Equal(SemanticTokenType.Function, mappedTokens[0].Type);
+
+            Token stringExpandableToken = scriptFile.ScriptTokens[1];
+            mappedTokens = PsesSemanticTokens.ConvertToSemanticTokens(stringExpandableToken);
+            Assert.Collection(mappedTokens,
+                sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type),
+                sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type),
+                sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type)
+            );
+        }
+    }
+}