lsp-bjolang/Program.fs

/// A language server for Bjolang. It answers diagnostics, from the compiler's
/// own frontend, when a file is opened or saved.
module BjolangLsp.Program

open System
open System.Reflection
open System.Text.Json.Nodes

/// Points the compiler at the Bjolang checkout this server was built against,
/// unless BJOLANG_HOME already names one. Done before anything touches
/// `Bjolang.Paths`, whose locations are computed once, when it is first used.
let private setBjolangHome () =
    match Environment.GetEnvironmentVariable "BJOLANG_HOME" with
    | null
    | "" ->
        Assembly.GetExecutingAssembly().GetCustomAttributes<AssemblyMetadataAttribute>()
        |> Seq.tryFind (fun a -> a.Key = "BjolangHome")
        |> Option.iter (fun a -> Environment.SetEnvironmentVariable("BJOLANG_HOME", a.Value))
    | _ -> ()

/// Makes the checks the installed compiler's: staleness judged against
/// `bin/compiler/Bjolang.dll` rather than this server's copy of the compiler,
/// and a stale import built by running that compiler. Built here, an import's
/// build record would name this server, and the command line would build it
/// again.
///
/// The variable is only read while the backend is installed, and is cleared
/// again so that the compiler it starts builds its own imports as usual.
let private useInstalledCompiler () =
    let installed = IO.Path.Combine(Bjolang.Paths.root, "bin", "compiler", "Bjolang.dll")

    if IO.File.Exists installed then
        Bjolang.Paths.compilerAssembly <- installed
        Environment.SetEnvironmentVariable("BJOLANG_OUT_OF_PROCESS_DEPS", "1")
        Bjolang.Build.installDependencyBackend ()
        Environment.SetEnvironmentVariable("BJOLANG_OUT_OF_PROCESS_DEPS", null)
    else
        Rpc.log $"no compiler at %s{installed}; run build_compiler.sh there. Files that import a stale module will fail to check."

let private uriToPath (uri: string) = Uri(uri).LocalPath

let private str (node: JsonNode) (path: string list) =
    path |> List.fold (fun (n: JsonNode) key -> if isNull n then null else n[key]) node
    |> fun n -> if isNull n then None else Some(n.GetValue<string>())

/// The text of each open file as the editor has it, saved or not. Symbols,
/// definitions, completions and checks are all read from it, so that they
/// follow the buffer.
let private texts = Collections.Generic.Dictionary<string, string>()

let private publish (uri: string) =
    let buffers =
        texts |> Seq.map (fun kv -> IO.Path.GetFullPath(uriToPath kv.Key), kv.Value) |> Map.ofSeq

    let found =
        try
            Check.file buffers (uriToPath uri)
        with ex ->
            Rpc.log $"check of %s{uri} failed: %O{ex}"
            []

    Rpc.notify "textDocument/publishDiagnostics" (Check.publishParams uri found)

/// When each changed file is next to be checked. A check waits for typing to
/// pause, as one per keystroke would queue up behind each other.
let private dueChecks = Collections.Generic.Dictionary<string, DateTime>()

let private pause = TimeSpan.FromMilliseconds 400.0

let private textOf (uri: string) =
    match texts.TryGetValue uri with
    | true, text -> text
    | _ -> IO.File.ReadAllText(uriToPath uri)

/// A compiler range as an LSP one. Lines are 1-based in the compiler and
/// 0-based in LSP; columns are 0-based in both.
let private range (r: Bjolang.Lexer.Range) : JsonNode =
    let position (p: Bjolang.Lexer.Position) =
        JsonObject(dict [ "line", JsonValue.Create(max 0 (p.Line - 1)) :> JsonNode; "character", JsonValue.Create p.Column ])
        :> JsonNode

    JsonObject(dict [ "start", position r.Start; "end", position r.End ])

let rec private documentSymbol (s: Symbols.Symbol) : JsonNode =
    JsonObject(
        dict
            [ "name", JsonValue.Create s.Name :> JsonNode
              "detail", JsonValue.Create s.Detail
              "kind", JsonValue.Create(int s.Kind)
              "range", range s.Range
              "selectionRange", range s.NameRange
              "children", JsonArray(s.Children |> List.map documentSymbol |> Array.ofList) ]
    )

/// The last symbols each file read to, kept for while it does not read: a
/// half-typed form would otherwise empty the outline on every keystroke.
let private lastSymbols = Collections.Generic.Dictionary<string, Symbols.Symbol list>()

let private symbolsOf (uri: string) : JsonNode =
    let symbols =
        match Symbols.ofText (uriToPath uri) (textOf uri) with
        | Some symbols ->
            lastSymbols[uri] <- symbols
            symbols
        | None ->
            match lastSymbols.TryGetValue uri with
            | true, symbols -> symbols
            | _ -> []

    JsonArray(symbols |> List.map documentSymbol |> Array.ofList)

/// The name at an LSP position, where it is, and its definition.
let private lookup (uri: string) (line: int) (character: int) =
    let path = uriToPath uri
    let text = textOf uri

    Symbols.nameAt path text (line + 1) character
    |> Option.bind (fun (name, at) -> Symbols.definition path text name |> Option.map (fun found -> at, found))

/// The local binding the name at an LSP position refers to, if it is one, and
/// where the name is.
let private localAt (uri: string) (line: int) (character: int) : (Scope.Binder * Bjolang.Lexer.Range) option =
    let path = uriToPath uri
    let text = textOf uri

    match Symbols.nameAt path text (line + 1) character, Symbols.forms path text with
    | Some(name, at), Some forms -> Scope.find forms (line + 1) character name |> Option.map (fun b -> b, at)
    | _ -> None

let private location (r: Bjolang.Lexer.Range) : JsonNode =
    JsonObject(dict [ "uri", JsonValue.Create(Uri(r.File).AbsoluteUri) :> JsonNode; "range", range r ])

let private definitionAt (uri: string) (line: int) (character: int) : JsonNode =
    match localAt uri line character with
    | Some(b, _) -> location b.Range
    | None ->
        match lookup uri line character with
        | None -> null
        | Some(_, found) -> location found.Symbol.NameRange

let private markdownContent (text: string) : JsonNode =
    JsonObject(dict [ "kind", JsonValue.Create "markdown" :> JsonNode; "value", JsonValue.Create text ])

/// What a hover says about `name`: of its definition, or of the builtin it
/// is.
let private describe (uri: string) (name: string) : string option =
    match Symbols.definition (uriToPath uri) (textOf uri) name with
    | Some found -> Some(Hover.text found)
    | None -> Hover.builtin name

let private hoverAt (uri: string) (line: int) (character: int) : JsonNode =
    let answer (text: string) (at: Bjolang.Lexer.Range) =
        JsonObject(dict [ "contents", markdownContent text; "range", range at ]) :> JsonNode

    match localAt uri line character with
    | Some(b, at) -> answer (Hover.local b (Symbols.sliceOf (textOf uri) b.Shown)) at
    | None ->
        match Symbols.nameAt (uriToPath uri) (textOf uri) (line + 1) character with
        | None -> null
        | Some(name, at) ->
            match describe uri name with
            | None -> null
            | Some text -> answer text at

/// The completions at a position. Each replaces the whole word before the
/// cursor. Only names that start with the word's first character are sent,
/// as the client filters the rest of the way itself, and every name at once
/// is two thousand of them.
let private completionAt (uri: string) (line: int) (character: int) : JsonNode =
    let text = textOf uri
    let lines = text.Replace("\r\n", "\n").Split('\n')
    let lineText = if line < lines.Length then lines[line] else ""
    let column = min character lineText.Length
    let start = Completion.wordStart lineText column
    let word = lineText.Substring(start, column - start)
    let own = IO.Path.GetFileName(uriToPath uri)

    let position (c: int) =
        JsonObject(dict [ "line", JsonValue.Create line :> JsonNode; "character", JsonValue.Create c ]) :> JsonNode

    let replace = JsonObject(dict [ "start", position start; "end", position column ])

    let rank (item: Completion.Item) =
        match item.From with
        | "local" -> 0
        | f when f = own -> 1
        | "builtin" -> 3
        | "special form" -> 4
        | _ -> 2

    // The locals in scope, innermost first and each name once: an inner
    // binding hides an outer one, and both hide a top-level name.
    let locals =
        match Symbols.forms (uriToPath uri) text with
        | Some forms ->
            Scope.at forms (line + 1) column
            |> List.rev
            |> List.distinctBy (fun b -> b.Name)
            // The word being typed is not yet a name in scope, unless it is
            // the binder itself, which is not worth offering either.
            |> List.filter (fun b -> not (b.Range.Start.Line = line + 1 && b.Range.Start.Column = start))
        | None -> []

    let localNames = locals |> List.map (fun b -> b.Name) |> Set.ofList

    let candidates =
        if Completion.inText lineText column then
            []
        else
            (locals |> List.map (fun b -> ({ Label = b.Name; Kind = Completion.ItemKind.Variable; From = "local" }: Completion.Item)))
            @ (Completion.items (uriToPath uri) text |> List.filter (fun item -> not (localNames.Contains item.Label)))

    let localDoc =
        locals
        |> List.map (fun b -> b.Name, Hover.local b (Symbols.sliceOf text b.Shown))
        |> Map.ofList

    let items =
        candidates
        |> List.filter (fun item -> word = "" || Char.ToLowerInvariant item.Label[0] = Char.ToLowerInvariant word[0])
        |> List.map (fun item ->
            let fields =
                [ "label", JsonValue.Create item.Label :> JsonNode
                  "kind", JsonValue.Create(int item.Kind)
                  "labelDetails", JsonObject(dict [ "description", JsonValue.Create item.From :> JsonNode ])
                  "sortText", JsonValue.Create($"%d{rank item}%s{item.Label}")
                  "textEdit",
                  JsonObject(dict [ "range", replace.DeepClone(); "newText", JsonValue.Create item.Label :> JsonNode ]) ]

            // A local's doc is at hand; anything else's is found when the
            // client asks for it.
            let more =
                match item.From, Map.tryFind item.Label localDoc with
                | "local", Some doc -> [ "documentation", markdownContent doc ]
                | _ -> [ "data", JsonObject(dict [ "uri", JsonValue.Create uri :> JsonNode; "name", JsonValue.Create item.Label ]) :> JsonNode ]

            JsonObject(dict (fields @ more)) :> JsonNode)

    JsonObject(dict [ "isIncomplete", JsonValue.Create false :> JsonNode; "items", JsonArray(Array.ofList items) ])

/// A completion item with its doc, which is the hover of its name.
let private resolve (item: JsonNode) : JsonNode =
    let data = item["data"]

    if not (isNull data) then
        let uri = data["uri"].GetValue<string>()
        let name = data["name"].GetValue<string>()

        match describe uri name with
        | Some text -> item["documentation"] <- markdownContent text
        | None -> ()

    item.DeepClone()

let private capabilities () : JsonNode =
    JsonNode.Parse
        """
        { "capabilities":
            { "textDocumentSync": { "openClose": true, "change": 1, "save": { "includeText": false } },
              "documentSymbolProvider": true,
              "definitionProvider": true,
              "hoverProvider": true,
              "completionProvider": { "resolveProvider": true, "triggerCharacters": [ "(" ] } },
          "serverInfo": { "name": "bjolang-lsp", "version": "0.4.0" } }
        """

[<EntryPoint>]
let main _ =
    setBjolangHome ()
    useInstalledCompiler ()
    let mutable shutdownAsked = false

    // Messages are read on a thread of their own and handled on this one.
    // Compiler work has to stay on one thread, as `Session` swaps the
    // compiler's global state around each piece of it; reading apart from it
    // is what lets this one notice that typing has paused.
    let inbox = new Collections.Concurrent.BlockingCollection<JsonNode option>()

    let reader =
        Threading.Thread(
            (fun () ->
                let rec go () =
                    match Rpc.read () with
                    | Some m ->
                        inbox.Add(Some m)
                        go ()
                    | None -> inbox.Add None

                try
                    go ()
                with ex ->
                    Rpc.log $"reading failed: %O{ex}"
                    inbox.Add None),
            IsBackground = true
        )

    reader.Start()

    /// The next message, running every check that falls due while waiting
    /// for it.
    let rec next () : JsonNode option =
        if dueChecks.Count = 0 then
            inbox.Take()
        else
            let soonest = dueChecks.Values |> Seq.min
            let untilDue = soonest - DateTime.UtcNow
            let wait = if untilDue < TimeSpan.Zero then TimeSpan.Zero else untilDue
            let mutable message: JsonNode option = None

            if inbox.TryTake(&message, wait) then
                message
            else
                let due =
                    dueChecks |> Seq.filter (fun kv -> kv.Value <= DateTime.UtcNow) |> Seq.map (fun kv -> kv.Key) |> List.ofSeq

                for uri in due do
                    dueChecks.Remove uri |> ignore
                    publish uri

                next ()

    let rec serve () =
        match next () with
        | None -> if shutdownAsked then 0 else 1
        | Some message ->
            let id = message["id"]
            let methodName = str message [ "method" ] |> Option.defaultValue ""
            let uri () = str message [ "params"; "textDocument"; "uri" ]

            match methodName with
            | "exit" -> if shutdownAsked then 0 else 1
            | _ ->
                try
                    match methodName with
                    | "initialize" -> Rpc.respond id (capabilities ())
                    | "shutdown" ->
                        shutdownAsked <- true
                        Rpc.respond id null
                    | "textDocument/didOpen" ->
                        uri ()
                        |> Option.iter (fun u ->
                            str message [ "params"; "textDocument"; "text" ] |> Option.iter (fun t -> texts[u] <- t)
                            publish u)
                    // Full sync: the one change is the whole text.
                    | "textDocument/didChange" ->
                        match uri (), message["params"]["contentChanges"] with
                        | Some u, (:? JsonArray as changes) when changes.Count > 0 ->
                            let last = changes[changes.Count - 1]
                            texts[u] <- last["text"].GetValue<string>()
                            dueChecks[u] <- DateTime.UtcNow + pause
                        | _ -> ()
                    | "textDocument/didSave" ->
                        uri ()
                        |> Option.iter (fun u ->
                            dueChecks.Remove u |> ignore
                            publish u)
                    // A closed file's diagnostics are the client's to forget,
                    // and an empty list is how it is told to.
                    | "textDocument/didClose" ->
                        uri ()
                        |> Option.iter (fun u ->
                            texts.Remove u |> ignore
                            dueChecks.Remove u |> ignore
                            lastSymbols.Remove u |> ignore
                            Rpc.notify "textDocument/publishDiagnostics" (Check.publishParams u []))
                    | "textDocument/documentSymbol" ->
                        Rpc.respond id (uri () |> Option.map symbolsOf |> Option.defaultValue (JsonArray()))
                    | "textDocument/definition"
                    | "textDocument/hover" ->
                        let position = message["params"]["position"]
                        let line = position["line"].GetValue<int>()
                        let character = position["character"].GetValue<int>()
                        let answer = if methodName = "textDocument/hover" then hoverAt else definitionAt

                        Rpc.respond id (uri () |> Option.map (fun u -> answer u line character) |> Option.defaultValue null)
                    | "textDocument/completion" ->
                        let position = message["params"]["position"]
                        let line = position["line"].GetValue<int>()
                        let character = position["character"].GetValue<int>()

                        Rpc.respond id (uri () |> Option.map (fun u -> completionAt u line character) |> Option.defaultValue null)
                    | "completionItem/resolve" -> Rpc.respond id (resolve message["params"])
                    // Not part of LSP: the expansion of the macro call at a
                    // position, from the file's last check. A check that is
                    // due runs first, so that the expansion is of the text as
                    // it stands.
                    | "bjolang/expandMacro" ->
                        let p = message["params"]
                        let position = p["position"]
                        let flag (key: string) (otherwise: bool) =
                            match p[key] with
                            | null -> otherwise
                            | v -> v.GetValue<bool>()

                        let answer =
                            uri ()
                            |> Option.bind (fun u ->
                                if dueChecks.Remove u then publish u

                                Expansion.at
                                    (uriToPath u)
                                    (position["line"].GetValue<int>() + 1)
                                    (position["character"].GetValue<int>())
                                    (flag "full" true)
                                    (flag "hygienic" false))
                            |> Option.map (fun (macro, text) ->
                                JsonObject(dict [ "macro", JsonValue.Create macro :> JsonNode; "text", JsonValue.Create text ])
                                :> JsonNode)
                            |> Option.defaultValue null

                        Rpc.respond id answer
                    | _ when not (isNull id) -> Rpc.respondError id -32601 $"%s{methodName} is not supported"
                    // Any other notification, `initialized` included, needs no
                    // answer.
                    | _ -> ()
                with ex ->
                    Rpc.log $"%s{methodName} failed: %O{ex}"

                    if not (isNull id) then
                        Rpc.respondError id -32603 ex.Message

                serve ()

    serve ()