Source file jsonrpc.ml

open Import
open Json.Conv
module Json = Json

module Id = struct
  type t =
    [ `String of string
    | `Int of int
    ]

  let yojson_of_t = function
    | `String s -> `String s
    | `Int i -> `Int i
  ;;

  let t_of_yojson = function
    | `String s -> `String s
    | `Int i -> `Int i
    | json -> Json.error "Id.t" json
  ;;

  let hash x = Hashtbl.hash x
  let equal = ( = )
end

module Constant = struct
  let jsonrpc = "jsonrpc"
  let jsonrpcv = "2.0"
  let id = "id"
  let method_ = "method"
  let params = "params"
  let result = "result"
  let error = "error"
end

let assert_jsonrpc_version fields =
  let jsonrpc = Json.field_exn fields Constant.jsonrpc Json.Conv.string_of_yojson in
  if not (String.equal jsonrpc Constant.jsonrpcv)
  then
    Json.error
      ("invalid packet: jsonrpc version doesn't match " ^ jsonrpc)
      (`Assoc fields)
;;

module Structured = struct
  type t =
    [ `Assoc of (string * Json.t) list
    | `List of Json.t list
    ]

  let t_of_yojson = function
    | `Assoc xs -> `Assoc xs
    | `List xs -> `List xs
    | json -> Json.error "invalid structured value" json
  ;;

  let yojson_of_t t = (t :> Json.t)
end

module Notification = struct
  type t =
    { method_ : string
    ; params : Structured.t option
    }

  let fields ~method_ ~params =
    let json =
      [ Constant.method_, `String method_; Constant.jsonrpc, `String Constant.jsonrpcv ]
    in
    match params with
    | None -> json
    | Some params -> (Constant.params, (params :> Json.t)) :: json
  ;;

  let yojson_of_t { method_; params } = `Assoc (fields ~method_ ~params)
  let create ?params ~method_ () = { params; method_ }
end

module Request = struct
  type t =
    { id : Id.t
    ; method_ : string
    ; params : Structured.t option
    }

  let yojson_of_t { id; method_; params } =
    let fields = Notification.fields ~method_ ~params in
    `Assoc ((Constant.id, Id.yojson_of_t id) :: fields)
  ;;

  let create ?params ~id ~method_ () = { params; id; method_ }
end

module Response = struct
  module Error = struct
    module Code = struct
      type t =
        | ParseError
        | InvalidRequest
        | MethodNotFound
        | InvalidParams
        | InternalError
        (* the codes below are LSP specific *)
        | ServerErrorStart
        | ServerErrorEnd
        | ServerNotInitialized
        | UnknownErrorCode
        | RequestFailed
        | ServerCancelled
        | ContentModified
        | RequestCancelled
        (* all other codes are custom *)
        | Other of int

      let of_int = function
        | -32700 -> ParseError
        | -32600 -> InvalidRequest
        | -32601 -> MethodNotFound
        | -32602 -> InvalidParams
        | -32603 -> InternalError
        | -32099 -> ServerErrorStart
        | -32000 -> ServerErrorEnd
        | -32002 -> ServerNotInitialized
        | -32001 -> UnknownErrorCode
        | -32800 -> RequestCancelled
        | -32801 -> ContentModified
        | -32802 -> ServerCancelled
        | -32803 -> RequestFailed
        | code -> Other code
      ;;

      let to_int = function
        | ParseError -> -32700
        | InvalidRequest -> -32600
        | MethodNotFound -> -32601
        | InvalidParams -> -32602
        | InternalError -> -32603
        | ServerErrorStart -> -32099
        | ServerErrorEnd -> -32000
        | ServerNotInitialized -> -32002
        | UnknownErrorCode -> -32001
        | RequestCancelled -> -32800
        | ContentModified -> -32801
        | ServerCancelled -> -32802
        | RequestFailed -> -32803
        | Other code -> code
      ;;

      let t_of_yojson json =
        match json with
        | `Int i -> of_int i
        | _ -> Json.error "invalid code" json
      ;;

      let yojson_of_t t = `Int (to_int t)
    end

    type t =
      { code : Code.t
      ; message : string
      ; data : Json.t option
      }

    let yojson_of_t { code; message; data } =
      let assoc = [ "code", Code.yojson_of_t code; "message", `String message ] in
      let assoc =
        match data with
        | None -> assoc
        | Some data -> ("data", data) :: assoc
      in
      `Assoc assoc
    ;;

    let t_of_yojson json =
      match json with
      | `Assoc fields ->
        let code = Json.field_exn fields "code" Code.t_of_yojson in
        let message = Json.field_exn fields "message" string_of_yojson in
        let data = Json.field fields "data" (fun x -> x) in
        { code; message; data }
      | _ -> Json.error "Jsonrpc.Response.t" json
    ;;

    exception E of t

    let raise t = raise (E t)
    let make ?data ~code ~message () = { data; code; message }

    let of_exn exn =
      let message = Printexc.to_string exn in
      make ~code:InternalError ~message ()
    ;;
  end

  type t =
    { id : Id.t
    ; result : (Json.t, Error.t) Result.t
    }

  let yojson_of_t { id; result } =
    let result =
      match result with
      | Ok json -> Constant.result, json
      | Error e -> Constant.error, Error.yojson_of_t e
    in
    `Assoc
      [ Constant.id, Id.yojson_of_t id
      ; Constant.jsonrpc, `String Constant.jsonrpcv
      ; result
      ]
  ;;

  let t_of_yojson json =
    match json with
    | `Assoc fields ->
      let id = Json.field_exn fields Constant.id Id.t_of_yojson in
      let jsonrpc = Json.field_exn fields Constant.jsonrpc Json.Conv.string_of_yojson in
      if jsonrpc <> Constant.jsonrpcv
      then Json.error "Invalid response" json
      else (
        match Json.field fields Constant.result (fun x -> x) with
        | Some res -> { id; result = Ok res }
        | None ->
          let result = Error (Json.field_exn fields Constant.error Error.t_of_yojson) in
          { id; result })
    | _ -> Json.error "Jsonrpc.Result.t" json
  ;;

  let make ~id ~result = { id; result }
  let ok id result = make ~id ~result:(Ok result)
  let error id error = make ~id ~result:(Error error)
end

module Packet = struct
  type t =
    | Notification of Notification.t
    | Request of Request.t
    | Response of Response.t
    | Batch_response of Response.t list
    | Batch_call of [ `Request of Request.t | `Notification of Notification.t ] list

  let yojson_of_t = function
    | Notification r -> Notification.yojson_of_t r
    | Request r -> Request.yojson_of_t r
    | Response r -> Response.yojson_of_t r
    | Batch_response r -> `List (List.map r ~f:Response.yojson_of_t)
    | Batch_call r ->
      `List
        (List.map r ~f:(function
           | `Request r -> Request.yojson_of_t r
           | `Notification r -> Notification.yojson_of_t r))
  ;;

  let t_of_fields (fields : (string * Json.t) list) =
    assert_jsonrpc_version fields;
    match Json.field fields Constant.id Id.t_of_yojson with
    | None ->
      let method_ = Json.field_exn fields Constant.method_ Json.Conv.string_of_yojson in
      let params = Json.field fields Constant.params Structured.t_of_yojson in
      Notification { Notification.params; method_ }
    | Some id ->
      (match Json.field fields Constant.method_ Json.Conv.string_of_yojson with
       | Some method_ ->
         let params = Json.field fields Constant.params Structured.t_of_yojson in
         Request { Request.method_; params; id }
       | None ->
         Response
           (match Json.field fields Constant.result (fun x -> x) with
            | Some result -> { Response.id; result = Ok result }
            | None ->
              let error =
                Json.field_exn fields Constant.error Response.Error.t_of_yojson
              in
              { id; result = Error error }))
  ;;

  let t_of_yojson_single json =
    match json with
    | `Assoc fields -> t_of_fields fields
    | _ -> Json.error "invalid packet" json
  ;;

  let t_of_yojson (json : Json.t) =
    match json with
    | `List [] -> Json.error "invalid packet" json
    | `List (x :: xs) ->
      (* we inspect the first element to see what we're dealing with *)
      let x =
        match x with
        | `Assoc fields -> t_of_fields fields
        | _ -> Json.error "invalid packet" json
      in
      (match
         match x with
         | Notification x -> `Call (`Notification x)
         | Request x -> `Call (`Request x)
         | Response r -> `Response r
         | _ -> Json.error "invalid packet" json
       with
       | `Call x ->
         Batch_call
           (x
            :: List.map xs ~f:(fun call ->
              let x = t_of_yojson_single call in
              match x with
              | Notification n -> `Notification n
              | Request n -> `Request n
              | _ -> Json.error "invalid packet" json))
       | `Response x ->
         Batch_response
           (x
            :: List.map xs ~f:(fun resp ->
              let resp = t_of_yojson_single resp in
              match resp with
              | Response n -> n
              | _ -> Json.error "invalid packet" json)))
    | _ -> t_of_yojson_single json
  ;;
end