Implement semantic tokens (for highlighting)

src/LanguageServer.jl

@@ -36,6 +36,7 @@ include("requests/actions.jl")
 include("requests/init.jl")
 include("requests/signatures.jl")
 include("requests/highlight.jl")
+include("requests/semantic.jl")
 include("utilities.jl")
 
 end

src/languageserverinstance.jl

@@ -344,6 +344,7 @@ function Base.run(server::LanguageServerInstance)
     msg_dispatcher[textDocument_prepareRename_request_type] = request_wrapper(textDocument_prepareRename_request, server)
     msg_dispatcher[textDocument_documentSymbol_request_type] = request_wrapper(textDocument_documentSymbol_request, server)
     msg_dispatcher[textDocument_documentHighlight_request_type] = request_wrapper(textDocument_documentHighlight_request, server)
+    msg_dispatcher[textDocument_semanticTokens_full_request_type] = request_wrapper(textDocument_semanticTokens_full_request, server)
     msg_dispatcher[julia_getModuleAt_request_type] = request_wrapper(julia_getModuleAt_request, server)
     msg_dispatcher[julia_getDocAt_request_type] = request_wrapper(julia_getDocAt_request, server)
     msg_dispatcher[textDocument_hover_request_type] = request_wrapper(textDocument_hover_request, server)

src/protocol/initialize.jl

@@ -77,6 +77,7 @@ end
     publishDiagnostics::Union{PublishDiagnosticsClientCapabilities,Missing}
     foldingRange::Union{FoldingRangeClientCapabilities,Missing}
     selectionRange::Union{SelectionRangeClientCapabilities,Missing}
+    semanticTokens::Union{SemanticTokensClientCapabilities,Missing}
 end
 
 @dict_readable struct WindowClientCapabilities <: Outbound
@@ -183,6 +184,7 @@ struct ServerCapabilities <: Outbound
     foldingRangeProvider::Union{Bool,FoldingRangeOptions,FoldingRangeRegistrationOptions,Missing}
     executeCommandProvider::Union{ExecuteCommandOptions,Missing}
     selectionRangeProvider::Union{Bool,SelectionRangeOptions,SelectionRangeRegistrationOptions,Missing}
+    semanticTokensProvider::Union{Bool,SemanticTokensOptions,SemanticTokensRegistrationOptions}
     workspaceSymbolProvider::Union{Bool,Missing}
     workspace::Union{WorkspaceOptions,Missing}
     experimental::Union{Any,Missing}

src/protocol/messagedefs.jl

@@ -10,6 +10,7 @@ const textDocument_prepareRename_request_type = JSONRPC.RequestType("textDocumen
 const textDocument_documentSymbol_request_type = JSONRPC.RequestType("textDocument/documentSymbol", DocumentSymbolParams, Union{Vector{DocumentSymbol}, Vector{SymbolInformation}, Nothing})
 const textDocument_documentHighlight_request_type = JSONRPC.RequestType("textDocument/documentHighlight", DocumentHighlightParams, Union{Vector{DocumentHighlight}, Nothing})
 const textDocument_hover_request_type = JSONRPC.RequestType("textDocument/hover", TextDocumentPositionParams, Union{Hover, Nothing})
+const textDocument_semanticTokens_full_request_type = JSONRPC.RequestType("textDocument/semanticTokens/full", SemanticTokensParams, Union{SemanticTokens,Nothing})
 const textDocument_didOpen_notification_type = JSONRPC.NotificationType("textDocument/didOpen", DidOpenTextDocumentParams)
 const textDocument_didClose_notification_type = JSONRPC.NotificationType("textDocument/didClose", DidCloseTextDocumentParams)
 const textDocument_didChange_notification_type = JSONRPC.NotificationType("textDocument/didChange", DidChangeTextDocumentParams)

src/protocol/protocol.jl

@@ -5,6 +5,7 @@ include("formatting.jl")
 include("hover.jl")
 include("goto.jl")
 include("highlight.jl")
+include("semantic.jl")
 include("signature.jl")
 include("symbols.jl")
 include("features.jl")

src/protocol/semantic.jl

@@ -0,0 +1,125 @@
+# https://microsoft.github.io/language-server-protocol/specifications/specification-3-17/#textDocument_semanticTokens
+const SemanticTokenKind = String
+const SemanticTokenKinds = (
+    Struct="struct",
+    TypeParameter="typeParameter",
+    Parameter="parameter",
+    Variable="variable",
+    Property="property",
+    Function="function",
+    Macro="macro",
+    Keyword="keyword",
+    Comment="comment",
+    String="string",
+    Number="number",
+    Regexp="regexp",
+    Operator="operator",
+)
+
+const SemanticTokenModifiersKind = String
+const SemanticTokenModifiersKinds = (
+    Declaration="declaration",
+    Definition="definition",
+    Modification="modification",
+    Documentation="documentation",
+    DefaultLibrary="defaultLibrary",
+)
+
+
+struct SemanticTokensLegend <: Outbound
+    tokenTypes::Vector{String} # The token types a server uses.
+
+    tokenModifiers::Vector{String} # The token modifiers a server uses.
+end
+
+const JuliaSemanticTokensLegend = SemanticTokensLegend(
+    collect(values(SemanticTokenKinds)),
+    collect(values(SemanticTokenModifiersKinds))
+)
+
+function semantic_token_encoding(token::String)::UInt32
+    for (i, type) in enumerate(JuliaSemanticTokensLegend.tokenTypes)
+        if token == type
+            return i - 1 # -1 to shift to 0-based indexing
+        end
+    end
+end
+
+@dict_readable struct SemanticTokensFullDelta <: Outbound
+    delta::Union{Bool,Missing}
+end
+
+@dict_readable struct SemanticTokensClientCapabilitiesRequests <: Outbound
+    range::Union{Bool,Missing}
+    full::Union{Bool,Missing,SemanticTokensFullDelta}
+
+end
+@dict_readable struct SemanticTokensClientCapabilities <: Outbound
+    dynamicRegistration::Union{Bool,Missing}
+    tokenTypes::Vector{String}
+    tokenModifiers::Vector{String}
+    formats::Vector{String}
+    overlappingTokenSupport::Union{Bool,Missing}
+    multilineTokenSupport::Union{Bool,Missing}
+end
+
+struct SemanticTokensOptions <: Outbound
+    legend::SemanticTokensLegend
+    range::Union{Bool,Missing}
+    full::Union{Bool,SemanticTokensFullDelta,Missing}
+end
+
+struct SemanticTokensRegistrationOptions <: Outbound
+    documentSelector::Union{DocumentSelector,Nothing}
+    # workDoneProgress::Union{Bool,Missing}
+end
+
+@dict_readable struct SemanticTokensParams <: Outbound
+    textDocument::TextDocumentIdentifier
+    # position::Position
+    workDoneToken::Union{Int,String,Missing} # ProgressToken
+    partialResultToken::Union{Int,String,Missing} # ProgressToken
+end
+
+struct SemanticTokens <: Outbound
+    resultId::Union{String,Missing}
+    data::Vector{UInt32}
+end
+
+SemanticTokens(data::Vector{UInt32}) = SemanticTokens(missing, data)
+
+
+struct SemanticTokensPartialResult <: Outbound
+    data::Vector{UInt32}
+end
+
+struct SemanticTokensDeltaParams <: Outbound
+    workDoneToken::Union{Int,String,Missing}
+    partialResultToken::Union{Int,String,Missing} # ProgressToken
+    textDocument::TextDocumentIdentifier
+    previousResultId::String
+end
+struct SemanticTokensEdit <: Outbound
+    start::UInt32
+    deleteCount::Int
+    data::Union{Vector{Int},Missing}
+end
+struct SemanticTokensDelta <: Outbound
+    resultId::Union{String,Missing}
+    edits::Vector{SemanticTokensEdit}
+end
+
+struct SemanticTokensDeltaPartialResult <: Outbound
+    edits::Vector{SemanticTokensEdit}
+end
+
+struct SemanticTokensRangeParams <: Outbound
+    workDoneToken::Union{Int,String,Missing}
+    partialResultToken::Union{Int,String,Missing} # ProgressToken
+    textDocument::TextDocumentIdentifier
+    range::Range
+end
+
+struct SemanticTokensWorkspaceClientCapabilities <: Outbound
+    refreshSupport::Union{Bool,Missing}
+end

src/requests/init.jl

@@ -30,6 +30,11 @@ function ServerCapabilities(client::ClientCapabilities)
         false,
         ExecuteCommandOptions(missing, collect(keys(LSActions))),
         true,
+        SemanticTokensOptions(
+            SemanticTokensLegend([values(SemanticTokenKinds)...],
+                [values(SemanticTokenModifiersKinds)...]),
+            false,
+            true),
         true,
         WorkspaceOptions(WorkspaceFoldersOptions(true, true)),
         missing

src/requests/semantic.jl

@@ -0,0 +1,160 @@
+const C = CSTParser
+
+# Struct-like version of a semantic token before being flattened into 5-number-pair
+struct NonFlattenedSemanticToken
+    # token line number
+    line::UInt32
+    # token start character within line
+    start::UInt32
+    # the length of the token.
+    length::UInt32
+    # will be looked up in SemanticTokensLegend.tokenTypes
+    tokenType::UInt32
+    # each set bit will be looked up in SemanticTokensLegend.tokenModifiers
+    tokenModifiers::UInt32
+end
+
+
+"""
+
+Map collection of tokens into SemanticTokens
+
+Note: currently uses absolute position
+"""
+function semantic_tokens(tokens)::SemanticTokens
+    # TODO implement relative position (track last token)
+    token_data_size = length(tokens) * 5
+    token_data = Vector{UInt32}(undef, token_data_size)
+    for (i_token, token::NonFlattenedSemanticToken) ∈ zip(1:5:token_data_size, tokens)
+        token_data[i_token:i_token+4] = [
+            token.line,
+            token.start,
+            token.length,
+            token.tokenType,
+            token.tokenModifiers
+        ]
+    end
+    SemanticTokens(token_data)
+end
+
+function textDocument_semanticTokens_full_request(params::SemanticTokensParams,
+    server::LanguageServerInstance, _)::Union{SemanticTokens,Nothing}
+    uri = params.textDocument.uri
+    d = getdocument(server, uri)
+
+    external_env = getenv(d, server)
+
+    repeated_tokens = collect(NonFlattenedSemanticToken, every_semantic_token(d, external_env))
+    sort!(repeated_tokens, lt=(l, r) -> begin
+        (l.line, l.start) < (r.line, r.start)
+    end)
+    return semantic_tokens(unique(repeated_tokens))
+end
+
+# TODO visit expressions correctly and collect tokens into a Vector, rather than a Set (see visit_every_expression() )
+TokenCollection = Set{NonFlattenedSemanticToken}
+mutable struct ExpressionVisitorState
+    collected_tokens::TokenCollection
+    # access to positioning (used with offset, see visit_every_expression() )
+    document::Document
+    # read-only
+    external_env::StaticLint.ExternalEnv
+end
+ExpressionVisitorState(args...) = ExpressionVisitorState(TokenCollection(), args...)
+
+"""
+Adds token to state.collected only if maybe_get_token_from_expr() parsed an actual token
+"""
+function maybe_collect_token_from_expr(ex::EXPR, state::ExpressionVisitorState, offset::Integer)
+    maybe_token = maybe_get_token_from_expr(ex, state, offset)
+    if maybe_token !== nothing
+        push!(state.collected_tokens, maybe_token)
+    end
+
+end
+
+function maybe_get_token_from_expr(ex::EXPR, state::ExpressionVisitorState, offset::Integer)::Union{Nothing,NonFlattenedSemanticToken}
+    kind = semantic_token_kind(ex, state.external_env)
+    if kind === nothing
+        return nothing
+    end
+    name = C.get_name(ex)
+    name_offset = 0
+    # get the offset of the name expr
+    if name !== nothing
+        found = false
+        for x in ex
+            if x == name
+                found = true
+                break
+            end
+            name_offset += x.fullspan
+        end
+        if !found
+            name_offset = -1
+        end
+    end
+    line, char = get_position_from_offset(state.document, offset)
+    return NonFlattenedSemanticToken(
+        line,
+        char,
+        ex.span,
+        semantic_token_encoding(kind),
+        0)
+end
+
+
+"""
+
+Visit each expression, collecting semantic-tokens into state
+
+Note: couldn't pack offset into ExpressionVisitorState and update, that's why it's a separate argument
+TODO: not sure about how to recurse an EXPR and its sub-expressions. For now, that'll be covered by collecting them into a Set
+"""
+function visit_every_expression(expr_in::EXPR, state::ExpressionVisitorState, offset=0)::Nothing
+    maybe_collect_token_from_expr(expr_in, state, offset)
+
+    # recurse into this expression's expressions
+    for e ∈ expr_in
+        maybe_collect_token_from_expr(e, state, offset)
+
+        visit_every_expression(e, state, offset)
+
+        offset += e.fullspan
+    end
+end
+
+function every_semantic_token(document::Document, external_env::StaticLint.ExternalEnv)
+    root_expr = getcst(document)
+    state = ExpressionVisitorState(document, external_env)
+    visit_every_expression(root_expr, state)
+    collect(state.collected_tokens)
+end
+
+
+"""
+Get the semantic token kind for `expr`, which is assumed to be an identifier
+
+See CSTParser.jl/src/interface.jl
+"""
+function semantic_token_kind(expr::EXPR, external_env::StaticLint.ExternalEnv)::Union{String,Nothing}
+    # TODO felipe use external_env
+
+    return if C.isidentifier(expr)
+        SemanticTokenKinds.Variable
+    elseif C.isoperator(expr)
+        SemanticTokenKinds.Operator
+    elseif C.isstringliteral(expr) || C.isstring(expr)
+        SemanticTokenKinds.String
+    elseif C.iskeyword(expr)
+        SemanticTokenKinds.Keyword
+    elseif C.defines_function(expr)
+        SemanticTokenKinds.Function
+    elseif C.defines_struct(expr)
+        SemanticTokenKinds.Struct
+    elseif C.defines_macro(expr)
+        SemanticTokenKinds.Macro
+    elseif C.isnumber(expr)
+        SemanticTokenKinds.Number
+    end
+end

test/requests/test_semantic.jl

@@ -0,0 +1,36 @@
+
+@testitem "simple token" begin
+    include("../test_shared_server.jl")
+
+    let _LS = LanguageServer,
+        _encoding = _LS.semantic_token_encoding
+
+        doc = settestdoc("""a=123""")
+        @test Int64.(token_full_test().data) == Int64.(_LS.SemanticTokens(
+            UInt32[0, 0, # first line, first column
+                1, # «a»
+                _encoding(_LS.SemanticTokenKinds.Variable), 0,
+                0, 1,
+                1, # «=»
+                _encoding(_LS.SemanticTokenKinds.Operator), 0,
+                0, 2,
+                3, # «123»
+                _encoding(_LS.SemanticTokenKinds.Number), 0,
+            ]).data)
+        doc = settestdoc("""const C = CSTParser""")
+        @test Int64.(token_full_test().data) == Int64.(_LS.SemanticTokens(
+            UInt32[0, 0,
+                5, # «const » TODO WIP
+                _encoding(_LS.SemanticTokenKinds.Keyword), 0,
+                0, 6,
+                1, # «C»
+                _encoding(_LS.SemanticTokenKinds.Variable), 0,
+                0, 8,
+                1, # «=»
+                _encoding(_LS.SemanticTokenKinds.Operator), 0,
+                0, 10,
+                9, # «CSTParser»
+                _encoding(_LS.SemanticTokenKinds.Variable), 0,
+            ]).data)
+    end
+end

test/test_shared_init_request.jl

@@ -40,6 +40,7 @@ init_request = LanguageServer.InitializeParams(
                 missing, # PublishDiagnosticsClientCapabilities(),
                 missing, # FoldingRangeClientCapabilities(),
                 missing, # SelectionRangeClientCapabilities()
+                missing, # SemanticTokensClientCapabilities()
             ),
             missing,
             missing

test/test_shared_server.jl

@@ -21,6 +21,7 @@ ref_test(line, char) = LanguageServer.textDocument_references_request(LanguageSe
 rename_test(line, char) = LanguageServer.textDocument_rename_request(LanguageServer.RenameParams(LanguageServer.TextDocumentIdentifier(uri"untitled:testdoc"), LanguageServer.Position(line, char), missing, "newname"), server, server.jr_endpoint)
 hover_test(line, char) = LanguageServer.textDocument_hover_request(LanguageServer.TextDocumentPositionParams(LanguageServer.TextDocumentIdentifier(uri"untitled:testdoc"), LanguageServer.Position(line, char)), server, server.jr_endpoint)
 range_formatting_test(line0, char0, line1, char1) = LanguageServer.textDocument_range_formatting_request(LanguageServer.DocumentRangeFormattingParams(LanguageServer.TextDocumentIdentifier(uri"untitled:testdoc"), LanguageServer.Range(LanguageServer.Position(line0, char0), LanguageServer.Position(line1, char1)), LanguageServer.FormattingOptions(4, true, missing, missing, missing)), server, server.jr_endpoint)
+token_full_test() = LanguageServer.textDocument_semanticTokens_full_request(LanguageServer.SemanticTokensParams(LanguageServer.TextDocumentIdentifier(uri"untitled:testdoc"), missing, missing), server, server.jr_endpoint)
 
 # TODO Replace this with a proper mock endpoint
 JSONRPC.send(::Nothing, ::Any, ::Any) = nothing