~jonsterling/forester-devel

LSP server based on Asai LSP v1 APPLIED

Kento Okura: 1
 LSP server based on Asai LSP

 17 files changed, 1232 insertions(+), 6 deletions(-)
Export patchset (mbox)
How do I use this?

Copy & paste the following snippet into your terminal to import this patchset into git:

curl -s https://lists.sr.ht/~jonsterling/forester-devel/patches/55701/mbox | git am -3
Learn more about email & git

[PATCH] LSP server based on Asai LSP Export this patch

Capabilities:
- Diagnostics (only parse errors at the momment)
- DocumentSync
- Hover
- Inlay hints for tree addresses
- Go to definition for imports. For transcludes, see
  https://todo.sr.ht/~jonsterling/forester/109
- Completion for addresses. For completion of unit paths, need to
  investigate how to use yuujinchou's get_visible function:
  https://ocaml.org/p/yuujinchou/5.2.0/doc/Yuujinchou/Scope/module-type-S/index.html#val-get_visible

New dependencies:
- ppx_yojson_conv, which depends on jane street's base...
- lsp, no additional dependencies
The LSP code was originally adapted from the Asai repo, which is
Apache-2.0 (with LLVM-exception) licensed. Thus, the files in
lib/language_server are dual-licensed with GPL and Apache licenses, and
the copyright is held by both the Forester and RedPRL developers.
---
 bin/forester/dune                   |   1 +
 bin/forester/main.ml                |  26 +-
 lib/compiler/Grammar.messages       |   2 +-
 lib/compiler/Parse.ml               |   3 +-
 lib/forest/Forester_forest.ml       |   1 +
 lib/frontend/Forest_reader.ml       |   2 +-
 lib/frontend/Import_graph.ml        |   2 +-
 lib/language_server/Analysis.ml     | 172 +++++++++++
 lib/language_server/Analysis.mli    |  15 +
 lib/language_server/Base.ml         |  18 ++
 lib/language_server/Completion.ml   |  78 +++++
 lib/language_server/Forester_lsp.ml | 244 +++++++++++++++
 lib/language_server/LspEio.ml       | 133 ++++++++
 lib/language_server/LspServer.ml    | 463 ++++++++++++++++++++++++++++
 lib/language_server/LspShims.ml     |  42 +++
 lib/language_server/dune            |  35 +++
 lib/prelude/Fun_util.ml             |   1 +
 17 files changed, 1232 insertions(+), 6 deletions(-)
 create mode 100644 lib/language_server/Analysis.ml
 create mode 100644 lib/language_server/Analysis.mli
 create mode 100644 lib/language_server/Base.ml
 create mode 100644 lib/language_server/Completion.ml
 create mode 100644 lib/language_server/Forester_lsp.ml
 create mode 100644 lib/language_server/LspEio.ml
 create mode 100644 lib/language_server/LspServer.ml
 create mode 100644 lib/language_server/LspShims.ml
 create mode 100644 lib/language_server/dune

diff --git a/bin/forester/dune b/bin/forester/dune
index 8c32178..8267c08 100644
--- a/bin/forester/dune
@@ -13,6 +13,7 @@
   forester.prelude
   forester.core
   forester.frontend
   forester.language_server
   forester.compiler
   cmdliner
   dune-build-info
diff --git a/bin/forester/main.ml b/bin/forester/main.ml
index 0391a18..1b8a3bf 100644
--- a/bin/forester/main.ml
@@ -242,6 +242,29 @@ let init_cmd ~env =
  let info = Cmd.info "init" ~version ~doc ~man in
  Cmd.v info Term.(const (init ~env) $ arg_dir)

let lsp ~env config =
  let config = Config.parse_forest_config_file config in
  Forester_lsp.start
    ~env
    ~config
    ~source: None

let lsp_cmd ~env =
  let man =
    [
      `S Manpage.s_description;
      `P "The $(tname) command starts the forester language server.";
    ]
  in
  let doc = "Start the LSP" in
  let info = Cmd.info "lsp" ~version ~doc ~man in
  Cmd.v
    info
    Term.(
      const (lsp ~env)
      $ arg_config
    )

let cmd ~env =
  let doc = "a tool for tending mathematical forests" in
  let man =
@@ -260,7 +283,8 @@ let cmd ~env =
      new_tree_cmd ~env;
      complete_cmd ~env;
      init_cmd ~env;
      query_cmd ~env
      query_cmd ~env;
      lsp_cmd ~env;
    ]

let () =
diff --git a/lib/compiler/Grammar.messages b/lib/compiler/Grammar.messages
index 794a4e1..f470287 100644
--- a/lib/compiler/Grammar.messages
+++ b/lib/compiler/Grammar.messages
@@ -4,10 +4,10 @@
main: TICK
main: OBJECT LBRACE LSQUARE TEXT RSQUARE WHITESPACE TICK
main: SUBTREE LBRACE TICK
main: IDENT TICK

Unexpected symbol: '

main: IDENT TICK
main: ALLOC XML_ELT_IDENT
main: CALL LBRACE RBRACE LBRACE XML_ELT_IDENT
main: CALL LBRACE RBRACE XML_ELT_IDENT
diff --git a/lib/compiler/Parse.ml b/lib/compiler/Parse.ml
index 162a432..b369bf1 100644
--- a/lib/compiler/Parse.ml
+++ b/lib/compiler/Parse.ml
@@ -121,8 +121,7 @@ let parse
          let end_position = lexbuf.lex_curr_p in
          if is_opening_delim token then
            let range = Range.of_lex_range (start_position, end_position) in
            Stack.push (token, range) delim_stack;
          ;
            Stack.push (token, range) delim_stack; ;
          if is_closing_delim token then
            begin
              match Stack.top_opt delim_stack with
diff --git a/lib/forest/Forester_forest.ml b/lib/forest/Forester_forest.ml
index 90a33f6..e8001f2 100644
--- a/lib/forest/Forester_forest.ml
+++ b/lib/forest/Forester_forest.ml
@@ -4,6 +4,7 @@
 * SPDX-License-Identifier: GPL-3.0-or-later
 *)

module Forest_graph = Forest_graph
module Forest_graphs = Forest_graphs
module Forest = Forest
module Forest_util = Forest_util
diff --git a/lib/frontend/Forest_reader.ml b/lib/frontend/Forest_reader.ml
index b048794..756b0c8 100644
--- a/lib/frontend/Forest_reader.ml
+++ b/lib/frontend/Forest_reader.ml
@@ -32,7 +32,7 @@ module Job_runner = struct
  let eval { env; host } job =
    let@ () = Reporter.easy_run in
    match job with
    | Job.LaTeX_to_svg { hash; source; content } ->
    | Job.LaTeX_to_svg{ hash; source; content } ->
      let svg = Build_latex.latex_to_svg ~env source in
      let frontmatter = T.default_frontmatter ~iri: (Iri_scheme.hash_iri ~host hash) () in
      let mainmatter = content ~svg in
diff --git a/lib/frontend/Import_graph.ml b/lib/frontend/Import_graph.ml
index 93836aa..8c36891 100644
--- a/lib/frontend/Import_graph.ml
+++ b/lib/frontend/Import_graph.ml
@@ -37,7 +37,7 @@ let build (trees : Code.tree list) =
      analyse_tree roots addr code
    | Scope code | Namespace (_, code) | Group (_, code) | Math (_, code) | Let (_, _, code) | Fun (_, code) | Def (_, _, code) ->
      analyse_code roots code
    | Object { methods; _ } | Patch { methods; _ } ->
    | Object{ methods; _ } | Patch{ methods; _ } ->
      let@ _, code = List.iter @~ methods in
      analyse_code roots code
    | Dx_prop (rel, args) ->
diff --git a/lib/language_server/Analysis.ml b/lib/language_server/Analysis.ml
new file mode 100644
index 0000000..0170401
--- /dev/null
+++ b/lib/language_server/Analysis.ml
@@ -0,0 +1,172 @@
(*
 * SPDX-FileCopyrightText: 2024 The Forester Project Contributors AND The RedPRL Development Team
 *
 * SPDX-License-Identifier: GPL-3.0-or-later OR Apache-2.0 WITH LLVM-exception
 *
 *)

open Forester_prelude
open Forester_core
open Forester_compiler
open Forester_frontend

module L = Lsp.Types
module T = Types
module EP = Eio.Path
module G = Forester_forest.Forest_graphs.Make ()
module F = Forester_forest.Forest.Make(G)
module FU = Forester_forest.Forest_util.Make(F)

module Code_set = Set.Make(struct
  type t = Code.tree
  (* TODO: no polymorphic compare*)
  let compare = compare
end)

let parse_path path = Parse.parse_string @@ EP.load path

let parse_from = function
  | `String s -> Parse.parse_string s
  | `Eio_path p -> parse_path p
  | `Uri(uri, cache: L.TextDocumentIdentifier.t * _) ->
    begin
      match Hashtbl.find_opt cache uri with
      | Some doc -> Parse.parse_string (Lsp.Text_document.text doc)
      | None ->
        Error (Reporter.fatalf Internal_error "Could not find %s in the internal document store. This is a bug!" (Lsp.Uri.to_path uri.uri))
    end
  | `Iri (env, iri) ->
    match F.get_article iri with
    | Some{ frontmatter = { source_path = Some str; _ }; _ } ->
      let p = EP.(env#fs / str) in
      parse_path p
    | _ ->
      Result.error @@
        Reporter.diagnosticf
          (Tree_not_found iri)
          "could not find tree %a "
          pp_iri
          iri

let dependencies (code : Code.t) host : iri Range.located list =
  let rec analyse_deps (node : Code.node Range.located) =
    match Range.(node.value) with
    | Import (_, dep) ->
      [Range.{ loc = node.loc; value = (Iri_scheme.user_iri ~host dep) }]
    | Subtree (_, code)
    | Scope code
    | Namespace (_, code)
    | Group (_, code)
    | Math (_, code)
    | Let (_, _, code)
    | Fun (_, code)
    | Def (_, _, code) ->
      List.concat_map analyse_deps code
    | Object{ methods; _ } | Patch{ methods; _ } ->
      let@ code = List.concat_map @~ methods in
      List.concat_map analyse_deps (snd code)
    | _ ->
      []
  in
  List.concat_map
    analyse_deps
    code

(* Does no IO*)
let get_dependencies (server : Base.server) code =
  let rec go c acc =
    let immediate_deps = dependencies c server.config.host in
    List.fold_left
      (
        fun acc' d ->
          match Hashtbl.find_opt server.resolver Range.(d.value) with
          | None ->
            Reporter.emitf ?loc: d.loc Resource_not_found "Could not find tree %a" pp_iri d.value;
            acc'
          | Some uri ->
            begin
              match Hashtbl.find_opt server.codes uri with
              | None ->
                Reporter.emitf ?loc: d.loc Resource_not_found "Could not find tree %s" @@ Lsp.Uri.to_path uri.uri;
                acc'
              | Some tree -> go tree.code (Code_set.add tree acc')
            end
      )
      acc
      immediate_deps
  in
  go code Code_set.empty

let check (server : Base.server) uri =
  let res = parse_from (`Uri (L.TextDocumentIdentifier.{ uri = uri }, server.documents)) in
  match res with
  | Ok code ->
    let tree = Code.{ source_path = None; addr = None; code } in
    let trans_deps = get_dependencies server code in
    let trees = trans_deps |> Code_set.to_list in
    let _units, _expanded_trees =
      Forest_reader.expand
        ~host: server.config.host
        (tree :: trees)
    in
    ()
  | Error diagnostic ->
    Reporter.emit_diagnostic diagnostic

let extract_addr (node : Code.node Range.located) : string option =
  match node.value with
  | Group (Braces, [{ value = Text addr; _ }])
  | Group (Parens, [{ value = Text addr; _ }])
  | Group (Squares, [{ value = Group (Squares, [{ value = Text addr; _ }]); _ }])
  | Import (_, addr) ->
    Some addr
  | Text _ | Verbatim _ | Math (_, _) | Ident _ | Hash_ident _ | Xml_ident _ | Subtree (_, _) | Let (_, _, _) | Open _ | Scope _ | Put (_, _) | Default (_, _) | Get _ | Fun (_, _) | Object _ | Patch _ | Call (_, _) | Def (_, _, _) | Decl_xmlns (_, _) | Alloc _ | Namespace (_, _) | _ -> None

(* If the code is already in the cache, this means that the editor already
   opened the file. In this case don't parse from disk, as we will be notified
   of changes to that file by the protocol*)

let _ =
  assert (
    extract_addr @@
      Range.{
        loc = None;
        value = Group (Parens, [{ value = Text "foo"; loc = None }])
      }
    = Some "foo"
  )

let rec flatten (tree : Code.t) : Code.t =
  tree
  |> List.concat_map @@
    fun (node : 'a Range.located) ->
      match node.value with
      | Code.Subtree (_, tree) -> flatten tree
      | Code.Scope tree -> flatten tree
      | _ -> [node]

let within ~range: (a, b) x = a <= x && x <= b

let is_at = fun
      ~(position : Lsp.Types.Position.t)
      (located : _ Range.located)
    ->
    match located.loc with
    | Some loc ->
      begin
        match Range.view loc with
        | `Range (start, end_) ->
          within ~range: (start.line_num, end_.line_num) (position.line + 1)
          && within
            ~range: ((start.offset - start.start_of_line), (end_.offset - end_.start_of_line - 1))
            position.character
        | _ -> false
      end
    | None -> false

let node_at ~(position : Lsp.Types.Position.t) (code : _) : _ option =
  let flattened = flatten code in
  List.find_opt (is_at ~position) flattened

let addr_at ~(position : Lsp.Types.Position.t) (code : _) : _ option =
  Option.bind (node_at ~position code) extract_addr
diff --git a/lib/language_server/Analysis.mli b/lib/language_server/Analysis.mli
new file mode 100644
index 0000000..75f9811
--- /dev/null
+++ b/lib/language_server/Analysis.mli
@@ -0,0 +1,15 @@
module G : Forester_forest.Forest_graphs.S
module F : Forester_forest.Forest.S
module L = Lsp.Types

val check : Base.server -> L.DocumentUri.t -> unit

val extract_addr :
  Forester_compiler.Code.node Forester_core.Range.located ->
  string option

val addr_at :
  position:Lsp.Types.Position.t ->
  Forester_compiler.Code.t ->
  string option

diff --git a/lib/language_server/Base.ml b/lib/language_server/Base.ml
new file mode 100644
index 0000000..f96baa4
--- /dev/null
+++ b/lib/language_server/Base.ml
@@ -0,0 +1,18 @@
open Forester_core
open Forester_frontend
open Forester_compiler

module L = Lsp.Types

type server = {
  env: Forest_reader.env;
  units: Expand.Env.t;
  config: Forester_frontend.Config.Forest_config.t;
  lsp_io: LspEio.io;
  should_shutdown: bool;
  source: string option;
  (* One hashtbl per phase? Annoying...*)
  resolver: (iri, L.TextDocumentIdentifier.t) Hashtbl.t;
  documents: (L.TextDocumentIdentifier.t, Lsp.Text_document.t) Hashtbl.t;
  codes: (L.TextDocumentIdentifier.t, Code.tree) Hashtbl.t;
}
diff --git a/lib/language_server/Completion.ml b/lib/language_server/Completion.ml
new file mode 100644
index 0000000..a01f714
--- /dev/null
+++ b/lib/language_server/Completion.ml
@@ -0,0 +1,78 @@
open Forester_compiler
module L = Lsp.Types

let kind
    : Syn.node -> L.CompletionItemKind.t option
  = function
  | Fun (_, _) -> Some Function
  | Text _ | Verbatim _ -> Some Text
  | Meta -> Some Field
  | Route_asset -> Some File
  | Var _ -> Some Variable
  | Prim _
  | Transclude
  | Embed_tex
  | Ref
  | Title
  | Parent
  | Taxon
  | Attribution (_, _)
  | Tag _
  | Date
  | Number ->
    Some Keyword
  | Group (_, _)
  | Math (_, _)
  | Link _
  | Subtree (_, _)
  | Sym _
  | Put (_, _, _)
  | Default (_, _, _)
  | Get _
  | Xml_tag (_, _, _)
  | TeX_cs _
  | Object _
  | Patch _
  | Call (_, _)
  | Query_polarity _
  | Query_mode _
  | Results_of_query
  | Query_rel _
  | Query_isect
  | Query_union
  | Query_compl
  | Query_isect_fam
  | Query_union_fam
  | Query_isect_fam_rel
  | Query_union_fam_rel
  | Query_builtin (_, _)
  | Dx_sequent (_, _)
  | Dx_query (_, _, _)
  | Dx_prop (_, _)
  | Dx_var _
  | Dx_const (_, _)
  | Dx_execute
  | Publish_results_of_query ->
    None

let insert_text path = String.concat "/" path

let make
    : Yuujinchou.Trie.path
    * (Resolver.P.data * unit) ->
    L.CompletionItem.t option
  = fun (path, (data, _)) ->
    match data with
    | Resolver.P.Term syn ->
      let kind = kind (List.hd syn).value in
      let insertText = insert_text path in
      Some
        (
          L.CompletionItem.create
            ?kind
            ~insertText
            ~label: (String.concat "/" path)
            ()
        )
    | Resolver.P.Xmlns _ ->
      None
diff --git a/lib/language_server/Forester_lsp.ml b/lib/language_server/Forester_lsp.ml
new file mode 100644
index 0000000..8e62a56
--- /dev/null
+++ b/lib/language_server/Forester_lsp.ml
@@ -0,0 +1,244 @@
(*
 * SPDX-FileCopyrightText: 2024 The Forester Project Contributors AND The RedPRL Development Team
 *
 * SPDX-License-Identifier: GPL-3.0-or-later OR Apache-2.0 WITH LLVM-exception
 *
 *)

open Forester_prelude
open Forester_core
open Forester_compiler
module Analysis = Analysis

module L = Lsp.Types
module RPC = Jsonrpc
module Base = Base

module Server = LspServer
open Server

let unwrap opt err =
  match opt with
  | Some opt -> opt
  | None -> raise @@ LspError err

let print_exn exn =
  let msg = Printexc.to_string exn
  and stack = Printexc.get_backtrace ()
  in
  Eio.traceln "%s\n%s" msg stack

(* [TODO: Reed M, 09/06/2022] Commands??? *)
let supported_code_actions = []
let supported_commands = []

let server_capabilities =
  let textDocumentSync =
    let opts =
      L.TextDocumentSyncOptions.create
        ~change: L.TextDocumentSyncKind.Full
        ~openClose: true
        ~save: (`SaveOptions (L.SaveOptions.create ~includeText: false ()))
        ()
    in
    `TextDocumentSyncOptions opts
  in
  let hoverProvider =
    let opts = L.HoverOptions.create () in
    `HoverOptions opts
  in
  let codeActionProvider =
    let opts = L.CodeActionOptions.create ~codeActionKinds: supported_code_actions () in
    `CodeActionOptions opts
  in
  let executeCommandProvider =
    L.ExecuteCommandOptions.create ~commands: supported_commands ()
  in
  let inlayHintProvider =
    let opts = L.InlayHintOptions.create () in
    `InlayHintOptions opts
  in
  let definitionProvider = `DefinitionOptions (L.DefinitionOptions.create ()) in
  let completionProvider =
    L.CompletionOptions.create
      ~triggerCharacters: ["\\"; "{"; "("; ]
      ~allCommitCharacters: ["}"; ")"; ]
      ()
  in
  let _semanticTokensProvider =
    let full = `Full (L.SemanticTokensOptions.create_full ~delta: true ()) in
    `SemanticTokensOptions (L.SemanticTokensOptions.create ~legend: Highlighting.legend ~full ())
  in
  (* [NOTE: Position Encodings]
     For various historical reasons, the spec states that we are _required_ to support UTF-16.
     This causes more trouble than it's worth, so we always select UTF-8 as our encoding, even
     if the client doesn't support it. *)
  let positionEncoding
    =
    L.PositionEncodingKind.UTF8
  in
  (* [FIME: Reed M, 09/06/2022] The current verison of the LSP library doesn't support 'positionEncoding' *)
  L.ServerCapabilities.create
    ~textDocumentSync
    ~hoverProvider
    ~codeActionProvider
    ~executeCommandProvider
    ~inlayHintProvider
    ~positionEncoding
    ~completionProvider
    ~definitionProvider
    (* ~semanticTokensProvider *)
    ()

let supports_utf8_encoding (init_params : L.InitializeParams.t) =
  let position_encodings =
    Option.value ~default: [] @@
    Option.bind init_params.capabilities.general @@
    fun gcap -> gcap.positionEncodings
  in
  List.mem L.PositionEncodingKind.UTF8 position_encodings

let get_root (init_params : L.InitializeParams.t) =
  match init_params.rootUri with
  | Some uri -> Some (L.DocumentUri.to_path uri)
  | None -> Option.join init_params.rootPath

module R = Lsp.Client_request

let build_once ~env (state : Base.server) () =
  let tree_dirs = (paths_of_dirs ~env state.config.trees) in
  Eio.traceln "Planting forest";
  let parsed_trees =
    Forester_frontend.Forester.parse_trees_in_dirs
      ~dev: true
      tree_dirs
  in
  parsed_trees
  |> List.iter
    (
      fun
          (Code.{ source_path; addr; _ } as code)
        ->
        match source_path with
        | Some p ->
          let uri = Lsp.Uri.of_path p in
          Hashtbl.add state.codes L.TextDocumentIdentifier.{ uri } code;
          begin
            match addr with
            | Some a ->
              Hashtbl.add
                state.resolver
                (Iri_scheme.user_iri ~host: state.config.host a)
                L.TextDocumentIdentifier.{ uri }
            | None -> ()
          end
        | None ->
          ()
    );
  try
    let articles, _ =
      Forester_frontend.Forest_reader.read_trees
        ~env
        ~host: state.config.host
        parsed_trees
    in
    let@ article = List.iter @~ articles in
    F.plant_resource @@ T.Article article
  with
    | _ -> ()

(** Perform the LSP initialization handshake.
    https://microsoft.github.io/language-server-protocol/specifications/specification-current/#initialize *)
let initialize () =
  let (id, req) =
    unwrap (Request.recv ()) @@
      HandshakeError "Initialization must begin with a request."
  in
  match req with
  | E (Initialize init_params as init_req) ->
    begin
      (* [HACK: Position Encodings]
         If the client doesn't support UTF-8, we shouldn't give up, as it might be using UTF-8 anyways...
         Therefore, we just produce a warning, and try to use UTF-8 regardless. *)
      if not (supports_utf8_encoding init_params) then
        Eio.traceln "Warning: client does not support UTF-8 encoding, which may lead to inconsistent positions.";
      let resp = L.InitializeResult.create ~capabilities: server_capabilities () in
      Request.respond id init_req resp;
      let notif =
        unwrap (Notification.recv ()) @@
          HandshakeError "Initialization must complete with an initialized notification."
      in
      match notif with
      | Initialized ->
        (* let root = get_root init_params in *)
        (* Eio.traceln "Root: %s" (Option.value root ~default: "<no-root>"); *)
        (* set_root root; *)
        ();
        Eio.traceln "Initialized!"
      | _ ->
        raise @@ LspError (HandshakeError "Initialization must complete with an initialized notification.")
    end
  | (E _) ->
    raise @@ LspError (HandshakeError "Initialization must begin with an initialize request.")

(** Perform the LSP shutdown sequence.
    See https://microsoft.github.io/language-server-protocol/specifications/specification-current/#exit *)
let shutdown () =
  let notif =
    unwrap (Notification.recv ()) @@
      ShutdownError "No requests can be recieved after a shutdown request."
  in
  match notif with
  | Exit ->
    ()
  | _ ->
    raise @@ LspError (ShutdownError "The only notification that can be recieved after a shutdown request is exit.")

(** {1 Main Event Loop} *)

let rec event_loop () =
  match recv () with
  | Some packet ->
    let _ =
      match packet with
      | RPC.Packet.Request req ->
        let resp = Request.handle req in
        send (RPC.Packet.Response resp)
      | RPC.Packet.Notification notif ->
        Notification.handle notif
      | _ ->
        Eio.traceln "Recieved unexpected packet type."
      | exception exn ->
        print_exn exn
    in
    if should_shutdown () then
      shutdown ()
    else
      event_loop ()
  | None ->
    Eio.traceln "Recieved an invalid message. Shutting down...@."

let start ~env ~source ~config =
  let lsp_io = LspEio.init env in
  let codes = Hashtbl.create 1000 in
  let resolver = Hashtbl.create 1000 in
  let init =
    Base.{
      env;
      lsp_io;
      config;
      source;
      codes;
      resolver;
      units = Expand.Env.empty;
      documents = Hashtbl.create 10;
      should_shutdown = false;
    }
  in
  build_once ~env init ();
  Server.run ~init @@
    fun () ->
      begin
        initialize ();
        event_loop ()
      end
diff --git a/lib/language_server/LspEio.ml b/lib/language_server/LspEio.ml
new file mode 100644
index 0000000..9efb610
--- /dev/null
+++ b/lib/language_server/LspEio.ml
@@ -0,0 +1,133 @@
(*
 * SPDX-FileCopyrightText: 2024 The Forester Project Contributors AND The RedPRL Development Team
 *
 * SPDX-License-Identifier: GPL-3.0-or-later OR Apache-2.0 WITH LLVM-exception
 *
 *)

open Eio
open Lsp.Import

module RPC = Jsonrpc

type io = {
  input: Buf_read.t;
  output: Eio_unix.sink_ty Eio.Resource.t;
}

(** See https://microsoft.github.io/language-server-protocol/specifications/specification-current/#headerPart *)
module Header = struct
  type t = {
    content_length: int;
    content_type: string
  }

  let empty =
    {
      content_length = -1;
      content_type = "application/vscode-jsonrpc; charset=utf-8"
    }

  let create ~(content_length : int) : t =
    { empty with content_length }

  let is_content_length key =
    String.equal (String.lowercase_ascii @@ String.trim key) "content-length"

  let is_content_type key =
    String.equal (String.lowercase_ascii @@ String.trim key) "content-type"

  (* NOTE: We should never really recieve an invalid header, as
     that would indicate a broken client implementation. Therefore,
     we just bail out when we see an invalid header, as there's
     no way we can really recover anyways. *)
  type header_error =
    | InvalidHeader of string
    | InvalidContentLength of string

  exception HeaderError of header_error

  (* If we do see any random header messages, we want to at least print out a decent error message. *)
  let () =
    Printexc.register_printer @@
      function
      | HeaderError (InvalidHeader err) -> Some (Format.asprintf "HeaderError: Invalid Header %s" err)
      | HeaderError (InvalidContentLength n) -> Some (Format.asprintf "HeaderError: Invalid Content Length '%s'" n)
      | _ -> None

  (* [TODO: Reed M, 09/06/2022] I could use some of the Buf_read parser module here, but this code works. *)
  let parse_header line headers =
    match String.split_on_char ~sep: ':' @@ String.trim line with
    | [key; value] when is_content_length key ->
      let content_length =
        match int_of_string_opt (String.trim value) with
        | Some n -> n
        | None -> raise (HeaderError (InvalidContentLength value))
      in
      { headers with content_length }
    | [key; value] when is_content_type key ->
      let content_type = String.trim value in
      { headers with content_type }
    | [_; _] ->
      (* We skip any unknown headers. *)
      headers
    | _ ->
      raise (HeaderError (InvalidHeader line))

  (** Read the header section of an LSP message. *)
  let read io =
    let rec loop headers =
      match Buf_read.line io.input with
      | "" -> headers
      | line -> loop (parse_header line headers)
    in
    let headers = loop empty in
    if headers.content_length < 0 then
      raise (HeaderError (InvalidContentLength (string_of_int headers.content_length)))
    else
      headers

  (** Write out the header section of an LSP message. *)
  let write io headers =
    let header_str =
      Format.asprintf
        "Content-Type: %s\r\nContent-Length: %d\r\n\r\n"
        headers.content_type
        headers.content_length
    in
    Flow.copy_string header_str io.output
end

module Message = struct
  let read io =
    try
      let header = Header.read io in
      let len = header.content_length in
      let json = Json.of_string @@ Buf_read.take len io.input in
      Some (RPC.Packet.t_of_yojson json)
    with
      | Sys_error _
      | End_of_file ->
        None

  let write io packet =
    let json = RPC.Packet.yojson_of_t packet in
    let data = Json.to_string json in
    let content_length = String.length data in
    let header = Header.create ~content_length in
    Header.write io header;
    Flow.copy_string data io.output
end

let init (env : Eio_unix.Stdenv.base) =
  {
    (* [TODO: Reed M, 09/06/2022] I should think about this buffer size... *)
    input = Buf_read.of_flow ~max_size: 1_000_000 @@ Eio.Stdenv.stdin env;
    output = Eio.Stdenv.stdout env
  }

let recv io =
  Message.read io

let send io packet =
  Message.write io packet
diff --git a/lib/language_server/LspServer.ml b/lib/language_server/LspServer.ml
new file mode 100644
index 0000000..7d83764
--- /dev/null
+++ b/lib/language_server/LspServer.ml
@@ -0,0 +1,463 @@
(*
 * SPDX-FileCopyrightText: 2024 The Forester Project Contributors AND The RedPRL Development Team
 *
 * SPDX-License-Identifier: GPL-3.0-or-later OR Apache-2.0 WITH LLVM-exception
 *
 *)

open Forester_prelude
open Forester_frontend
open Forester_compiler
open Forester_core

module T = Forester_core.Types
module EP = Eio.Path
module L = Lsp.Types
module RPC = Jsonrpc
module Broadcast = Lsp.Server_notification
module Lsp_Diagnostic = Lsp.Types.Diagnostic
module Lsp_Request = Lsp.Client_request
module Lsp_Notification = Lsp.Client_notification

module F = Analysis.F

module PT = Forester_forest.Plain_text_client.Make(F)(struct let route _ = "todo" end)

type eio_path = Eio.Fs.dir_ty EP.t

let path_of_dir ~env dir =
  EP.(Eio.Stdenv.fs env / dir)

let paths_of_dirs ~env =
  List.map (path_of_dir ~env)

type diagnostic = Reporter.Message.t Asai.Diagnostic.t

module State = Algaeff.State.Make(struct type t = Base.server end)

type lsp_error =
  | DecodeError of string
  | HandshakeError of string
  | ShutdownError of string
  | UnknownRequest of string
  | UnknownNotification of string

exception LspError of lsp_error

let () =
  Printexc.register_printer @@
    function
    | LspError (DecodeError err) ->
      Some (Format.asprintf "Lsp Error: Couldn't decode %s" err)
    | LspError (HandshakeError err) ->
      Some (Format.asprintf "Lsp Error: Invalid initialization handshake %s" err)
    | LspError (ShutdownError err) ->
      Some (Format.asprintf "Lsp Error: Invalid shutdown sequence %s" err)
    | LspError (UnknownRequest err) ->
      Some (Format.asprintf "Lsp Error: Unknown request %s" err)
    | LspError (UnknownNotification err) ->
      Some (Format.asprintf "Lsp Error: Unknown notification %s" err)
    | _ -> None

let recv () =
  let server = State.get () in
  LspEio.recv server.lsp_io

let send packet =
  let server = State.get () in
  LspEio.send server.lsp_io packet

let broadcast notif =
  let msg = Broadcast.to_jsonrpc notif in
  send (RPC.Packet.Notification msg)

let render_lsp_related_info (uri : L.DocumentUri.t) (message : Asai.Diagnostic.loctext) : L.DiagnosticRelatedInformation.t =
  let range = LspShims.Loc.lsp_range_of_range message.loc in
  let location = L.Location.create ~uri ~range in
  let message = Asai.Diagnostic.string_of_text message.value in
  L.DiagnosticRelatedInformation.create ~location ~message

let render_lsp_diagnostic (uri : L.DocumentUri.t) (diag : diagnostic) : Lsp_Diagnostic.t =
  let range = LspShims.Loc.lsp_range_of_range diag.explanation.loc in
  let severity = LspShims.Diagnostic.lsp_severity_of_severity @@ diag.severity in
  let code = `String (Reporter.Message.short_code diag.message) in
  let source = (State.get ()).source in
  let message = Asai.Diagnostic.string_of_text diag.explanation.value in
  let relatedInformation = Bwd.to_list @@ Bwd.map (render_lsp_related_info uri) diag.extra_remarks in
  Lsp_Diagnostic.create
    ~range
    ~severity
    ~code
    ?source
    ~message: (`String message)
    ~relatedInformation
    ()

let publish_diagnostics (uri : Lsp.Uri.t) (diagnostics : diagnostic list) =
  let diagnostics = List.map (render_lsp_diagnostic uri) diagnostics in
  let params = L.PublishDiagnosticsParams.create ~uri ~diagnostics () in
  broadcast (PublishDiagnostics params)

let should_shutdown () =
  let server = State.get () in
  server.should_shutdown

let initiate_shutdown () =
  State.modify @@ fun st -> { st with should_shutdown = true }

(* [TODO: Reed M, 12/12/2022] No code actions for now. *)
let code_action (_params : L.CodeActionParams.t) : L.CodeActionResult.t =
  None

let completion
    (params : L.CompletionParams.t)
  =
  match params with
  | {
    context;
    _;
  } ->
    let triggerCharacter =
      match context with
      | Some{ triggerCharacter; _ } ->
        triggerCharacter
      | None -> None
    in
    let server = State.get () in
    let addr_items () =
      server.codes
      |> Hashtbl.to_seq_values
      |> List.of_seq
      |> List.filter_map
        (
          fun (tree : Code.tree) ->
            let* addr = tree.addr in
            let* { frontmatter; mainmatter; _ } =
              (F.get_article @@ Iri_scheme.user_iri ~host: server.config.host addr)
            in
            let documentation =
              try
                let render = PT.string_of_content in
                let title = frontmatter.title in
                let taxon = frontmatter.taxon in
                let content =
                  Format.asprintf
                    {|%s
%s
%s
|}
                    (Option.fold ~none: "" ~some: (fun s -> Format.asprintf "# %s" (render s)) title)
                    (Option.fold ~none: "" ~some: (fun s -> Format.asprintf "taxon: %s" (render s)) taxon)
                    (render mainmatter)
                in
                Some (`String content)
              with
                | _ ->
                  Some (`String "computation of my value crashed")
            in
            let insertText =
              match triggerCharacter with
              | Some "{" -> addr ^ "}"
              | Some "(" -> addr ^ ")"
              | Some "[" -> addr ^ "]"
              | _ -> addr
            in
            Some (L.CompletionItem.create ?documentation ~label: addr ~insertText ())
        )
    in
    let trees = server.codes |> Hashtbl.to_seq_values |> List.of_seq in
    let scope_items () =
      let units, _expanded = Forest_reader.expand ~host: server.config.host trees in
      units
      |> Expand.Unit_map.to_list
      |> List.map snd
      |> List.concat_map
        (
          fun trie ->
            let open Yuujinchou in
            trie
            |> Trie.to_seq
            |> List.of_seq
            |> List.filter_map Completion.make
        )
    in
    let items =
      match triggerCharacter with
      | Some "(" -> addr_items ()
      | Some "{" -> addr_items ()
      | Some "\\" -> scope_items ()
      | _ -> []
    in
    Some
      (
        `CompletionList
          (L.CompletionList.create ~isIncomplete: false ~items ())
      )

let hover
    ({
      position;
      textDocument;
      _
    }: L.HoverParams.t)
    : L.Hover.t option
  =
  let server = State.get () in
  let host = server.config.host in
  let* tree = Hashtbl.find_opt server.codes { uri = textDocument.uri } in
  let* addr_at_cursor = Analysis.addr_at ~position tree.code in
  let iri_under_cursor = Iri_scheme.user_iri ~host addr_at_cursor in
  let content =
    match F.get_article iri_under_cursor with
    | None -> Format.asprintf "extracted iri %a" pp_iri iri_under_cursor
    | Some{ mainmatter; _ } ->
      PT.string_of_content mainmatter
  in
  Some
    (
      L.Hover.create
        ~contents: (
          `MarkupContent
            {
              kind = L.MarkupKind.Markdown;
              value = content
            }
        )
        ()
    )

module Highlighting = struct
  module Token_type = struct
    let legend : L.SemanticTokenTypes.t list =
      [
        Namespace;
        Type;
        Class;
        Enum;
        Interface;
        Struct;
        TypeParameter;
        Parameter;
        Variable;
        Property;
        EnumMember;
        Event;
        Function;
        Method;
        Macro;
        Keyword;
        Modifier;
        Comment;
        String;
        Number;
        Regexp;
        Operator;
        Decorator
      ]

    let token_types =
      List.map
        (
          fun s ->
            match L.SemanticTokenTypes.yojson_of_t s with
            | `String s -> s
            | _ -> assert false
        )
        legend
  end

  module Token_modifiers_set = struct
    let list = []
  end

  let legend = L.SemanticTokensLegend.create ~tokenTypes: Token_type.token_types ~tokenModifiers: Token_modifiers_set.list

  (* let compute_tokens doc = _ *)

  let on_request
      : L.SemanticTokensDeltaParams.t ->
      [`SemanticTokens of L.SemanticTokens.t
      | `SemanticTokensDelta of L.SemanticTokensDelta.t] option
    = fun
        {
          (* partialResultToken; *)
          (* previousResultId; *)
          (* workDoneToken; *)
          textDocument;
          _;
        }
      ->
      let server = State.get () in
      let doc = Hashtbl.find_opt server.documents { uri = textDocument.uri } in
      match doc with
      | None -> None
      | Some _doc ->
        (* let tokens = compute_tokens doc in *)
        Some (`SemanticTokens { L.SemanticTokens.resultId = None; data = Array.of_list [] })
end

let definitions
    (params : L.DefinitionParams.t)
    : L.Locations.t option
  =
  match params with
  | {
    position;
    textDocument;
    _;
  } ->
    let server = State.get () in
    let codes = server.codes in
    let* { code; _} = Hashtbl.find_opt codes { uri = textDocument.uri } in
    let* addr = Analysis.addr_at ~position code in
    let iri = Iri_scheme.user_iri ~host: server.config.host addr in
    let* uri = Hashtbl.find_opt server.resolver iri in
    let range = L.Range.create ~start: { character = 1; line = 0 } ~end_: { character = 1; line = 0 } in
    Some
      (`Location [L.Location.{ uri = uri.uri; range }])

let inlay_hint (params : L.InlayHintParams.t) : L.InlayHint.t list option =
  match params with
  | {
    textDocument;
    _;
  } ->
    let server = State.get () in
    match Hashtbl.find_opt server.codes { uri = textDocument.uri } with
    | None -> None
    | Some{ code; _ } ->
      List.filter_map
        (
          fun
              (Range.{ loc; _} as node)
            ->
            match Option.map Range.view loc with
            | Some (`Range (_, pos)) ->
              let* str = Analysis.extract_addr node in
              let iri = Iri_scheme.user_iri ~host: server.config.host str in
              let* { frontmatter; _} = F.get_article iri in
              let* title = frontmatter.title in
              let content = " " ^ PT.string_of_content title in
              Some
                (
                  L.InlayHint.create
                    ~position: (LspShims.Loc.lsp_pos_of_pos pos)
                    ~label: (`String content)
                    ()
                )
            | _ -> None
        )
        code
      |> Option.some

module Request = struct
  type 'resp t = 'resp Lsp.Client_request.t
  type packed = Lsp_Request.packed

  let dispatch : type resp. string -> resp t -> resp = fun mthd ->
      function
      | Initialize _ ->
        let err = "Server can only recieve a single initialization request." in
        raise @@ LspError (HandshakeError err)
      | Shutdown ->
        initiate_shutdown ()
      | CodeAction params ->
        code_action params
      | TextDocumentHover params ->
        hover params
      | TextDocumentCompletion params ->
        completion params
      | InlayHint params ->
        inlay_hint params
      | TextDocumentDefinition params ->
        definitions params
      | _ ->
        raise @@ LspError (UnknownRequest mthd)

  let handle (msg : RPC.Request.t) =
    Eio.traceln "Request: %s@." msg.method_;
    match Lsp_Request.of_jsonrpc msg with
    | Ok (E r) ->
      let resp = dispatch msg.method_ r in
      let json = Lsp_Request.yojson_of_result r resp in
      RPC.Response.ok msg.id json
    | Error err ->
      raise (LspError (DecodeError err))

  let recv () =
    Option.bind (recv ()) @@
      function
      | RPC.Packet.Request req ->
        begin
          match Lsp_Request.of_jsonrpc req with
          | Ok packed -> Some (req.id, packed)
          | Error err -> raise @@ LspError (DecodeError err)
        end
      | _ -> None

  let respond id req resp =
    let json = Lsp_Request.yojson_of_result req resp in
    send (RPC.Packet.Response (RPC.Response.ok id json))
end

module Notification = struct
  type t = Lsp.Client_notification.t

  let dispatch : string -> t -> unit = fun mthd ->
      let server = State.get () in
      function
      | TextDocumentDidOpen ({ textDocument = { uri; _ } } as params) ->
        let text_document = Lsp.Text_document.make ~position_encoding: `UTF16 params in
        Hashtbl.replace server.documents { uri } text_document;
        Reporter.lsp_run publish_diagnostics uri @@
          fun () ->
            Analysis.check server uri
      | DidSaveTextDocument{ textDocument; _; } ->
        begin
          match Hashtbl.find_opt server.documents textDocument with
          (* ocaml-lsp does *this* here: https://github.com/ocaml/ocaml-lsp/blob/8b47925eb44f907b8ec41a44c1b2a55447f1b439/ocaml-lsp-server/src/ocaml_lsp_server.ml#L757 *)
          | _ -> ()
        end
      | TextDocumentDidChange{ textDocument = { uri; _ }; contentChanges } ->
        begin
          match Hashtbl.find_opt server.documents { uri } with
          | Some doc ->
            let new_doc =
              Lsp.Text_document.apply_content_changes
                doc
                contentChanges
            in
            Hashtbl.replace server.documents { uri } new_doc;
            Reporter.lsp_run publish_diagnostics uri @@
              fun () ->
                Analysis.check server uri
          | None ->
            Reporter.lsp_run publish_diagnostics uri @@
              fun () ->
                Reporter.fatalf Internal_error "%s" "could not find document at %s" (uri |> Lsp.Uri.to_path)
        end
      | _ ->
        raise @@ LspError (UnknownNotification mthd)

  let handle (msg : RPC.Notification.t) =
    Eio.traceln "Request: %s@." msg.method_;
    match Lsp_Notification.of_jsonrpc msg with
    | Ok notif ->
      dispatch msg.method_ notif
    | Error err ->
      raise @@ LspError (DecodeError err)

  let recv () =
    Option.bind (recv ()) @@
      function
      | RPC.Packet.Notification msg ->
        begin
          match Lsp_Notification.of_jsonrpc msg with
          | Ok notif -> Some notif
          | Error err -> raise @@ LspError (DecodeError err)
        end
      | _ -> None
end

let run ~init k =
  State.run ~init k
diff --git a/lib/language_server/LspShims.ml b/lib/language_server/LspShims.ml
new file mode 100644
index 0000000..a50755c
--- /dev/null
+++ b/lib/language_server/LspShims.ml
@@ -0,0 +1,42 @@
(*
 * SPDX-FileCopyrightText: 2024 The Forester Project Contributors AND The RedPRL Development Team
 *
 * SPDX-License-Identifier: GPL-3.0-or-later OR Apache-2.0 WITH LLVM-exception
 *
 *)

module L = Lsp.Types

module Loc = struct
  let lsp_pos_of_pos (pos : Asai.Range.position) =
    L.Position.create
      ~line: (pos.line_num - 1)
      ~character: (pos.offset - pos.start_of_line)

  let lsp_range_of_range (r : Asai.Range.t option) =
    match r with
    | Some r ->
      let (start, stop) =
        match Asai.Range.view r with
        | `Range (start, stop) -> start, stop
        | `End_of_file pos -> pos, pos
      in
      L.Range.create
        ~start: (lsp_pos_of_pos start)
        ~end_: (lsp_pos_of_pos stop)
    | None ->
      (* When we have a message without a location,
         we set it's location to the start of the file,
         as we don't have any better choices. *)
      let start_of_file = L.Position.create ~line: 0 ~character: 0 in
      L.Range.create ~start: start_of_file ~end_: start_of_file
end

module Diagnostic = struct
  let lsp_severity_of_severity : Asai.Diagnostic.severity -> L.DiagnosticSeverity.t = function
    | Hint -> Hint
    | Info -> Information
    | Warning -> Warning
    | Error -> Error
    | Bug -> Error
end
diff --git a/lib/language_server/dune b/lib/language_server/dune
new file mode 100644
index 0000000..b0b2a9d
--- /dev/null
+++ b/lib/language_server/dune
@@ -0,0 +1,35 @@
;;; SPDX-FileCopyrightText: 2024 The Forester Project Contributors
;;; 
;;; SPDX-License-Identifier: GPL-3.0-or-later OR Apache-2.0 WITH LLVM-exception

(library
 (name Forester_lsp)
 (libraries
   algaeff
   repr
   unix
   forester.prelude
   forester.core
   forester.compiler
   forester.frontend
   forester.forest
   lsp
   asai
   eio
   eio_main
   eio.unix
   jsonrpc
   yojson
   bwd
   iri
   yuujinchou
   fmt
   )
 (preprocess
  (pps ppx_deriving.show ppx_repr ppx_yojson_conv))
 (public_name forester.language_server))

;(env
; (dev
;  (flags
;   (:standard -w -66-32-33-27-26))))
diff --git a/lib/prelude/Fun_util.ml b/lib/prelude/Fun_util.ml
index fa101e6..c23bd5f 100644
--- a/lib/prelude/Fun_util.ml
+++ b/lib/prelude/Fun_util.ml
@@ -6,3 +6,4 @@

let (let@) = ( @@ )
let ( @~ ) f x y = f y x
let (let*) = Option.bind
-- 
2.46.0