From 556da72593c88b027a7124ae86da6cd638c7d679 Mon Sep 17 00:00:00 2001
From: PixieDust <111846546+PizieDust@users.noreply.github.com>
Date: Thu, 14 Nov 2024 15:19:18 +0100
Subject: [PATCH] Type search custom request (#1369)

This custom request allows clients to perform a type search at a specific position within a text document based on finding functions or types that match a specific query pattern. Result can include documentation in ocaml-doc or markdown format.
---
 .github/workflows/build-and-test.yml          |   2 +-
 CHANGES.md                                    |   3 +
 .../docs/ocamllsp/typeSearch-spec.md          |  55 ++++
 .../src/custom_requests/custom_request.ml     |   1 +
 .../src/custom_requests/custom_request.mli    |   1 +
 .../src/custom_requests/req_type_search.ml    | 119 +++++++++
 .../src/custom_requests/req_type_search.mli   |  22 ++
 ocaml-lsp-server/src/ocaml_lsp_server.ml      |   2 +
 ocaml-lsp-server/test/e2e-new/dune            |   1 +
 ocaml-lsp-server/test/e2e-new/start_stop.ml   |   3 +-
 ocaml-lsp-server/test/e2e-new/type_search.ml  | 234 ++++++++++++++++++
 11 files changed, 441 insertions(+), 2 deletions(-)
 create mode 100644 ocaml-lsp-server/docs/ocamllsp/typeSearch-spec.md
 create mode 100644 ocaml-lsp-server/src/custom_requests/req_type_search.ml
 create mode 100644 ocaml-lsp-server/src/custom_requests/req_type_search.mli
 create mode 100644 ocaml-lsp-server/test/e2e-new/type_search.ml

diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index 7416bc634..1889fd256 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -54,7 +54,7 @@ jobs:
 
       # Remove this pin once a compatible version of Merlin has been released
       # - name: Pin dev Merlin
-      #   run: opam pin https://github.com/ocaml/merlin.git#master
+      #   run: opam pin https://github.com/ocaml/merlin.git#main
 
       - name: Build and install dependencies
         run: opam install .
diff --git a/CHANGES.md b/CHANGES.md
index 341b3dc42..1626846d3 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,5 +1,8 @@
 # Unreleased
 
+- Add custom
+  [`ocamllsp/typeSearch`](/ocaml-lsp-server/docs/ocamllsp/typeSearch-spec.md) request (#1369)
+
 - Make MerlinJump code action configurable (#1376)
 
 ## Fixes
diff --git a/ocaml-lsp-server/docs/ocamllsp/typeSearch-spec.md b/ocaml-lsp-server/docs/ocamllsp/typeSearch-spec.md
new file mode 100644
index 000000000..2b34815db
--- /dev/null
+++ b/ocaml-lsp-server/docs/ocamllsp/typeSearch-spec.md
@@ -0,0 +1,55 @@
+# TypeSearch Request
+
+## Description
+
+This custom request allows clients to perform a type search at a specific position within a text document based on finding functions or types that match a specific query pattern.
+
+## Server capability
+
+- property name: `handleTypeSearch`
+- property type: `boolean`
+
+## Request
+
+```js
+export interface TypeSearchParams extends TexDocumentPositionParams
+{
+    query: string;
+    limit: int;
+    with_doc: bool;
+    doc_format: string;
+}
+```
+- method: `ocamllsp/typeSearch`
+- params:
+    - `TextDocumentPositionParams`: This is an existing interface that includes:
+            - `TextDocumentIdentifier`: Specifies the document uri for which the request is sent.
+            - `Position`: Specifies the cursor position.
+        More details can be found in the [TextDocumentPositionParams - LSP Specification](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocumentPositionParams).
+    - `query`: The search pattern.
+    - `limit`: The number of results to return
+    - `with_doc`: If to return documentation information or not
+
+## Response
+```json
+{
+    [
+        "name": string,
+        "typ": string,
+        "loc": Range,
+        "doc": {
+            "value": string,
+            "kind": string
+        },
+        "cost": int,
+        "constructible" : string
+    ]
+}
+ ```
+- name: The fully qualified name of this result.,
+- typ: The signature of this result,
+- loc: The location of the definition of this result in the source code.,
+- doc: Optional documentation associated with this result.,
+- cost: A numeric value representing the "cost" or distance between this result and the query.
+- constructible: A constructible form or template that can be used to invoke this result
+- A response with null result is returned if no entries are found.
diff --git a/ocaml-lsp-server/src/custom_requests/custom_request.ml b/ocaml-lsp-server/src/custom_requests/custom_request.ml
index ad68a560a..4f4398aa2 100644
--- a/ocaml-lsp-server/src/custom_requests/custom_request.ml
+++ b/ocaml-lsp-server/src/custom_requests/custom_request.ml
@@ -7,3 +7,4 @@ module Typed_holes = Req_typed_holes
 module Type_enclosing = Req_type_enclosing
 module Wrapping_ast_node = Req_wrapping_ast_node
 module Get_documentation = Req_get_documentation
+module Type_search = Req_type_search
diff --git a/ocaml-lsp-server/src/custom_requests/custom_request.mli b/ocaml-lsp-server/src/custom_requests/custom_request.mli
index 199e87d5e..b26880242 100644
--- a/ocaml-lsp-server/src/custom_requests/custom_request.mli
+++ b/ocaml-lsp-server/src/custom_requests/custom_request.mli
@@ -9,3 +9,4 @@ module Typed_holes = Req_typed_holes
 module Type_enclosing = Req_type_enclosing
 module Wrapping_ast_node = Req_wrapping_ast_node
 module Get_documentation = Req_get_documentation
+module Type_search = Req_type_search
diff --git a/ocaml-lsp-server/src/custom_requests/req_type_search.ml b/ocaml-lsp-server/src/custom_requests/req_type_search.ml
new file mode 100644
index 000000000..215280ec4
--- /dev/null
+++ b/ocaml-lsp-server/src/custom_requests/req_type_search.ml
@@ -0,0 +1,119 @@
+open Import
+module TextDocumentPositionParams = Lsp.Types.TextDocumentPositionParams
+
+let meth = "ocamllsp/typeSearch"
+let capability = "handleTypeSearch", `Bool true
+
+module TypeSearchParams = struct
+  type t =
+    { text_document : TextDocumentIdentifier.t
+    ; position : Position.t
+    ; limit : int
+    ; query : string
+    ; with_doc : bool
+    ; doc_format : MarkupKind.t option
+    }
+
+  let t_of_yojson json =
+    let open Yojson.Safe.Util in
+    let textDocumentPosition = Lsp.Types.TextDocumentPositionParams.t_of_yojson json in
+    let query = json |> member "query" |> to_string in
+    let limit = json |> member "limit" |> to_int in
+    let with_doc = json |> member "with_doc" |> to_bool in
+    let doc_format = json |> member "doc_format" |> to_option MarkupKind.t_of_yojson in
+    { position = textDocumentPosition.position
+    ; text_document = textDocumentPosition.textDocument
+    ; query
+    ; limit
+    ; with_doc
+    ; doc_format
+    }
+  ;;
+
+  let yojson_of_t { text_document; position; query; limit; with_doc; doc_format } =
+    let doc_format =
+      match doc_format with
+      | Some format -> [ "doc_format", MarkupKind.yojson_of_t format ]
+      | None -> []
+    in
+    `Assoc
+      (("textDocument", TextDocumentIdentifier.yojson_of_t text_document)
+       :: ("position", Position.yojson_of_t position)
+       :: ("limit", `Int limit)
+       :: ("with_doc", `Bool with_doc)
+       :: ("query", `String query)
+       :: doc_format)
+  ;;
+end
+
+module TypeSearch = struct
+  type t = string Query_protocol.type_search_result list
+
+  let doc_to_markupContent ~kind ~value =
+    let v =
+      match kind with
+      | MarkupKind.Markdown ->
+        (match Doc_to_md.translate value with
+         | Raw d -> d
+         | Markdown d -> d)
+      | MarkupKind.PlainText -> value
+    in
+    MarkupContent.create ~kind ~value:v
+  ;;
+
+  let yojson_of_t (t : t) doc_format =
+    let format =
+      match doc_format with
+      | Some format -> format
+      | None -> MarkupKind.PlainText
+    in
+    let yojson_of_type_search_result (res : string Query_protocol.type_search_result) =
+      `Assoc
+        [ "name", `String res.name
+        ; "typ", `String res.typ
+        ; "loc", Range.yojson_of_t (Range.of_loc res.loc)
+        ; ( "doc"
+          , match res.doc with
+            | Some value ->
+              doc_to_markupContent ~kind:format ~value |> MarkupContent.yojson_of_t
+            | None -> `Null )
+        ; "cost", `Int res.cost
+        ; "constructible", `String res.constructible
+        ]
+    in
+    `List (List.map ~f:yojson_of_type_search_result t)
+  ;;
+end
+
+type t = TypeSearch.t
+
+module Request_params = struct
+  type t = TypeSearchParams.t
+
+  let yojson_of_t t = TypeSearchParams.yojson_of_t t
+
+  let create text_document position limit query with_doc doc_format : t =
+    { text_document; position; limit; query; with_doc; doc_format }
+  ;;
+end
+
+let dispatch merlin position limit query with_doc doc_format =
+  Document.Merlin.with_pipeline_exn merlin (fun pipeline ->
+    let position = Position.logical position in
+    let query = Query_protocol.Type_search (query, position, limit, with_doc) in
+    let results = Query_commands.dispatch pipeline query in
+    TypeSearch.yojson_of_t results doc_format)
+;;
+
+let on_request ~params state =
+  Fiber.of_thunk (fun () ->
+    let params = (Option.value ~default:(`Assoc []) params :> Yojson.Safe.t) in
+    let TypeSearchParams.{ text_document; position; limit; query; with_doc; doc_format } =
+      TypeSearchParams.t_of_yojson params
+    in
+    let uri = text_document.uri in
+    let doc = Document_store.get state.State.store uri in
+    match Document.kind doc with
+    | `Other -> Fiber.return `Null
+    | `Merlin merlin -> dispatch merlin position limit query with_doc doc_format)
+;;
diff --git a/ocaml-lsp-server/src/custom_requests/req_type_search.mli b/ocaml-lsp-server/src/custom_requests/req_type_search.mli
new file mode 100644
index 000000000..d17e5a82d
--- /dev/null
+++ b/ocaml-lsp-server/src/custom_requests/req_type_search.mli
@@ -0,0 +1,22 @@
+open Import
+
+module Request_params : sig
+  type t
+
+  val yojson_of_t : t -> Json.t
+
+  val create
+    :  TextDocumentIdentifier.t
+    -> Position.t
+    -> int
+    -> string
+    -> bool
+    -> MarkupKind.t option
+    -> t
+end
+
+type t
+
+val meth : string
+val capability : string * [> `Bool of bool ]
+val on_request : params:Jsonrpc.Structured.t option -> State.t -> Json.t Fiber.t
diff --git a/ocaml-lsp-server/src/ocaml_lsp_server.ml b/ocaml-lsp-server/src/ocaml_lsp_server.ml
index 733b1903c..52b16c779 100644
--- a/ocaml-lsp-server/src/ocaml_lsp_server.ml
+++ b/ocaml-lsp-server/src/ocaml_lsp_server.ml
@@ -96,6 +96,7 @@ let initialize_info (client_capabilities : ClientCapabilities.t) : InitializeRes
               ; Req_type_enclosing.capability
               ; Req_get_documentation.capability
               ; Req_construct.capability
+              ; Req_type_search.capability
               ] )
         ]
     in
@@ -526,6 +527,7 @@ let on_request
        ; Req_type_enclosing.meth, Req_type_enclosing.on_request
        ; Req_get_documentation.meth, Req_get_documentation.on_request
        ; Req_wrapping_ast_node.meth, Req_wrapping_ast_node.on_request
+       ; Req_type_search.meth, Req_type_search.on_request
        ; Req_construct.meth, Req_construct.on_request
        ; ( Semantic_highlighting.Debug.meth_request_full
          , Semantic_highlighting.Debug.on_request_full )
diff --git a/ocaml-lsp-server/test/e2e-new/dune b/ocaml-lsp-server/test/e2e-new/dune
index 0ee7bbb23..e17a1b200 100644
--- a/ocaml-lsp-server/test/e2e-new/dune
+++ b/ocaml-lsp-server/test/e2e-new/dune
@@ -63,6 +63,7 @@
     test
     type_enclosing
     documentation
+    type_search
     with_pp
     with_ppx
     workspace_change_config))))
diff --git a/ocaml-lsp-server/test/e2e-new/start_stop.ml b/ocaml-lsp-server/test/e2e-new/start_stop.ml
index 31bfdd351..cc549955e 100644
--- a/ocaml-lsp-server/test/e2e-new/start_stop.ml
+++ b/ocaml-lsp-server/test/e2e-new/start_stop.ml
@@ -93,7 +93,8 @@ let%expect_test "start/stop" =
             "handleMerlinCallCompatible": true,
             "handleTypeEnclosing": true,
             "handleGetDocumentation": true,
-            "handleConstruct": true
+            "handleConstruct": true,
+            "handleTypeSearch": true
           }
         },
         "foldingRangeProvider": true,
diff --git a/ocaml-lsp-server/test/e2e-new/type_search.ml b/ocaml-lsp-server/test/e2e-new/type_search.ml
new file mode 100644
index 000000000..ece053402
--- /dev/null
+++ b/ocaml-lsp-server/test/e2e-new/type_search.ml
@@ -0,0 +1,234 @@
+open Test.Import
+module Req = Ocaml_lsp_server.Custom_request.Type_search
+
+module Util = struct
+  let call_search position query with_doc doc_format client =
+    let uri = DocumentUri.of_path "test.ml" in
+    let text_document = TextDocumentIdentifier.create ~uri in
+    let params =
+      Req.Request_params.create text_document position 3 query with_doc doc_format
+      |> Req.Request_params.yojson_of_t
+      |> Jsonrpc.Structured.t_of_yojson
+      |> Option.some
+    in
+    let req =
+      Lsp.Client_request.UnknownRequest { meth = "ocamllsp/typeSearch"; params }
+    in
+    Client.request client req
+  ;;
+
+  let test ~line ~character ~query source ~with_doc ?(doc_format = None) () =
+    let position = Position.create ~character ~line in
+    let request client =
+      let open Fiber.O in
+      let+ response = call_search position query with_doc doc_format client in
+      Test.print_result response
+    in
+    Helpers.test source request
+  ;;
+end
+
+let%expect_test "Polarity Search for a simple query that takes an int and returns a \
+                 string with documentation"
+  =
+  let source = "" in
+  let line = 1 in
+  let character = 0 in
+  let doc_format = Some MarkupKind.Markdown in
+  Util.test ~line ~character ~query:"-int +string" source ~with_doc:true ~doc_format ();
+  [%expect
+    {|
+    [
+      {
+        "name": "Int.to_string",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 29, "line": 152 },
+          "start": { "character": 0, "line": 152 }
+        },
+        "doc": {
+          "kind": "markdown",
+          "value": "`to_string x` is the written representation of `x` in decimal."
+        },
+        "cost": 4,
+        "constructible": "Int.to_string _"
+      },
+      {
+        "name": "string_of_int",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 33, "line": 740 },
+          "start": { "character": 0, "line": 740 }
+        },
+        "doc": {
+          "kind": "markdown",
+          "value": "Return the string representation of an integer, in decimal."
+        },
+        "cost": 4,
+        "constructible": "string_of_int _"
+      },
+      {
+        "name": "string_of_int",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 33, "line": 740 },
+          "start": { "character": 0, "line": 740 }
+        },
+        "doc": {
+          "kind": "markdown",
+          "value": "Return the string representation of an integer, in decimal."
+        },
+        "cost": 4,
+        "constructible": "string_of_int _"
+      }
+    ]
+    |}]
+;;
+
+let%expect_test "Polarity Search for a simple query that takes an int and returns a \
+                 string with no documentation"
+  =
+  let source = "" in
+  let line = 1 in
+  let character = 0 in
+  Util.test ~line ~character ~query:"-int +string" source ~with_doc:false ();
+  [%expect
+    {|
+    [
+      {
+        "name": "Int.to_string",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 29, "line": 152 },
+          "start": { "character": 0, "line": 152 }
+        },
+        "doc": null,
+        "cost": 4,
+        "constructible": "Int.to_string _"
+      },
+      {
+        "name": "string_of_int",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 33, "line": 740 },
+          "start": { "character": 0, "line": 740 }
+        },
+        "doc": null,
+        "cost": 4,
+        "constructible": "string_of_int _"
+      },
+      {
+        "name": "string_of_int",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 33, "line": 740 },
+          "start": { "character": 0, "line": 740 }
+        },
+        "doc": null,
+        "cost": 4,
+        "constructible": "string_of_int _"
+      }
+    ] |}]
+;;
+
+let%expect_test "Type Search for a simple query that takes an int and returns a string \
+                 with no documentation"
+  =
+  let source = "" in
+  let line = 1 in
+  let character = 0 in
+  Util.test ~line ~character ~query:"int -> string" source ~with_doc:false ();
+  [%expect
+    {|
+    [
+      {
+        "name": "Int.to_string",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 29, "line": 152 },
+          "start": { "character": 0, "line": 152 }
+        },
+        "doc": null,
+        "cost": 0,
+        "constructible": "Int.to_string _"
+      },
+      {
+        "name": "string_of_int",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 33, "line": 740 },
+          "start": { "character": 0, "line": 740 }
+        },
+        "doc": null,
+        "cost": 0,
+        "constructible": "string_of_int _"
+      },
+      {
+        "name": "string_of_int",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 33, "line": 740 },
+          "start": { "character": 0, "line": 740 }
+        },
+        "doc": null,
+        "cost": 0,
+        "constructible": "string_of_int _"
+      }
+    ] |}]
+;;
+
+let%expect_test "Type Search for a simple query that takes an int and returns a string \
+                 with documentation"
+  =
+  let source = "" in
+  let line = 1 in
+  let character = 0 in
+  Util.test ~line ~character ~query:"int -> string" source ~with_doc:true ();
+  [%expect
+    {|
+    [
+      {
+        "name": "Int.to_string",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 29, "line": 152 },
+          "start": { "character": 0, "line": 152 }
+        },
+        "doc": {
+          "kind": "plaintext",
+          "value": "[to_string x] is the written representation of [x] in decimal."
+        },
+        "cost": 0,
+        "constructible": "Int.to_string _"
+      },
+      {
+        "name": "string_of_int",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 33, "line": 740 },
+          "start": { "character": 0, "line": 740 }
+        },
+        "doc": {
+          "kind": "plaintext",
+          "value": "Return the string representation of an integer, in decimal."
+        },
+        "cost": 0,
+        "constructible": "string_of_int _"
+      },
+      {
+        "name": "string_of_int",
+        "typ": "int -> string",
+        "loc": {
+          "end": { "character": 33, "line": 740 },
+          "start": { "character": 0, "line": 740 }
+        },
+        "doc": {
+          "kind": "plaintext",
+          "value": "Return the string representation of an integer, in decimal."
+        },
+        "cost": 0,
+        "constructible": "string_of_int _"
+      }
+    ]
+    |}]
+;;