Add support for Codelens requests

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [0.1.4] Next release
 
 ### Added
+- Support for LSP request `textDocument/codeLens` [#349](https://github.com/OCamlPro/superbol-studio-oss/pull/349)
 - Show display example of `NUMERIC-EDITED` data on hover [#337](https://github.com/OCamlPro/superbol-studio-oss/pull/337)
 - Support for dump and listing files, along with a task attribute for outputting the latter [#347](https://github.com/OCamlPro/superbol-studio-oss/pull/347)
 - Improved information shown on completion [#336](https://github.com/OCamlPro/superbol-studio-oss/pull/336)

src/lsp/cobol_lsp/lsp_capabilities.ml

@@ -47,6 +47,7 @@ let reply (_: ClientCapabilities.t) =
     ServerCapabilities.create_workspace ()
       ~workspaceFolders
   in
+  let codeLensProvider = CodeLensOptions.create () in
   ServerCapabilities.create ()
     ~textDocumentSync:(`TextDocumentSyncOptions sync)
     ~definitionProvider:(`Bool true)
@@ -59,3 +60,4 @@ let reply (_: ClientCapabilities.t) =
     ~completionProvider:completion_option
     ~workspace
     ~documentSymbolProvider:(`Bool true)
+    ~codeLensProvider

src/lsp/cobol_lsp/lsp_request.ml

@@ -586,6 +586,82 @@ let handle_document_symbol registry (params: DocumentSymbolParams.t) =
       Some (`DocumentSymbol symbols)
     end
 
+(** { Document Code Lens } *)
+
+module PosSet = Set.Make(struct
+    type t = Position.t
+    let compare (p1: t) (p2: t) =
+      let c = p2.line - p1.line in
+      if c <> 0 then c else p2.character - p1.character
+  end)
+
+let codelens_positions ~uri group =
+  let filename = Lsp.Uri.to_path uri in
+  let open struct
+    include Cobol_common.Visitor
+    include Cobol_data.Visitor
+    type context =
+      | ProcedureDiv
+      | DataDiv
+      | None
+  end in
+  let set_context context (old, acc) =
+    do_children_and_then (context, acc) (fun (_, acc) -> (old, acc))
+  in
+  let take_when_in context { loc; _ } (current, acc) =
+    if context <> current
+    then skip (current, acc)
+    else
+      let range = Lsp_position.range_of_srcloc_in ~filename loc in
+      skip (context, PosSet.add range.start acc)
+  in
+  Cobol_unit.Visitor.fold_unit_group
+    object (v)
+      inherit [_] Cobol_unit.Visitor.folder
+      method! fold_procedure _ = set_context ProcedureDiv
+      method! fold_data_definitions _ = set_context DataDiv
+      method! fold_paragraph' _ = skip
+      method! fold_procedure_name' = take_when_in ProcedureDiv
+      method! fold_qualname' = take_when_in DataDiv
+      method! fold_record_renaming { renaming_name; _ } =
+        take_when_in DataDiv renaming_name
+      method! fold_field_definition' { payload = field; _ } acc =
+        fold_field_definition v { field with field_redefines = None } acc
+        |> skip
+      method! fold_table_definition' { payload = table; _ } acc =
+        fold_table_definition v { table with table_redefines = None } acc
+        |> skip
+    end group (None, PosSet.empty)
+  |> snd
+
+let handle_codelens registry ({ textDocument; _ }: CodeLensParams.t) =
+  try_with_main_document_data registry textDocument
+    ~f:begin fun ~doc checked_doc ->
+      let uri = Lsp.Text_document.documentUri doc.textdoc in
+      let rootdir = Lsp_project.(string_of_rootdir @@ rootdir doc.project) in
+      let context = ReferenceContext.create ~includeDeclaration:true in
+      codelens_positions ~uri checked_doc.group
+      |> PosSet.to_seq
+      |> Seq.map begin fun position ->
+        let params =
+          ReferenceParams.create ~context ~position ~textDocument () in
+        let ref_count =
+          lookup_references_in_doc ~rootdir params checked_doc
+          |> Option.fold ~none:0 ~some:List.length in
+        let title = string_of_int ref_count
+                    ^ " reference"
+                    ^ if ref_count > 1 then "s" else "" in
+        let range = Range.create ~end_:position ~start:position in
+        let uri = DocumentUri.yojson_of_t textDocument.uri in
+        let command = Command.create () ~title
+            ~command:"superbol.editor.action.findReferences"
+            ~arguments:[uri; Position.yojson_of_t position] in
+        CodeLens.create ~command ~range ()
+      end
+      |> List.of_seq |> Option.some
+    end
+  |> Option.value ~default:[]
+
 (** {3 Generic handling} *)
 
 let shutdown: state -> unit = function
@@ -633,10 +709,11 @@ let on_request
         Ok (handle_shutdown registry, ShuttingDown)
     | DocumentSymbol  (* DocumentSymbolParams.t.t *) params ->
         Ok (handle_document_symbol registry params, state)
+    | TextDocumentCodeLens (* CodeLensParams.t.t *) params ->
+        Ok (handle_codelens registry params, state)
     | TextDocumentDeclaration  (* TextDocumentPositionParams.t.t *) _
     | TextDocumentTypeDefinition  (* TypeDefinitionParams.t.t *) _
     | TextDocumentImplementation  (* ImplementationParams.t.t *) _
-    | TextDocumentCodeLens  (* CodeLensParams.t.t *) _
     | TextDocumentCodeLensResolve  (* CodeLens.t.t *) _
     | TextDocumentPrepareCallHierarchy  (* CallHierarchyPrepareParams.t.t *) _
     | TextDocumentPrepareRename  (* PrepareRenameParams.t.t *) _
@@ -705,6 +782,7 @@ module INTERNAL = struct
   let lookup_references = handle_references
   let hover = handle_hover
   let completion = handle_completion
+  let codelens = handle_codelens
   let document_symbol = handle_document_symbol
   let formatting = handle_formatting
 end

src/lsp/cobol_lsp/lsp_request.mli

@@ -41,4 +41,8 @@ module INTERNAL: sig
     : Lsp_server.t
     -> Lsp.Types.DocumentSymbolParams.t
     -> [> `DocumentSymbol of Lsp.Types.DocumentSymbol.t list ] option
+  val codelens
+    : Lsp_server.t
+    -> Lsp.Types.CodeLensParams.t
+    -> Lsp.Types.CodeLens.t list
 end

src/lsp/cobol_typeck/typeck_procedure.ml

@@ -182,7 +182,7 @@ let references ~(data_definitions: Cobol_unit.Types.data_definitions) procedure
     { acc with diags = Proc_error err :: acc.diags }
   in
 
-  let visitor = object
+  let visitor = object (v)
     inherit [acc] Cobol_unit.Visitor.folder
 
     method! fold_qualname qn acc =                (* TODO: data_name' instead *)
@@ -212,10 +212,19 @@ let references ~(data_definitions: Cobol_unit.Types.data_definitions) procedure
           error acc @@ Ambiguous_data_name { given_qualname = qn &@ loc;
                                              matching_qualnames }
 
-    method! fold_procedure_section s ({ current_section; _ } as acc) =
-      Visitor.do_children_and_then
-        { acc with current_section = Some s }
-        (fun acc -> { acc with current_section })
+    method! fold_procedure_section ({ section_paragraphs; _ } as s)
+        ({ current_section; _ } as acc) =
+      { acc with current_section = Some s }
+      |> Cobol_common.Visitor.fold_list v section_paragraphs.list
+        ~fold:Cobol_unit.Visitor.fold_procedure_paragraph'
+      |> (fun acc -> { acc with current_section })
+      |> Cobol_common.Visitor.skip
+
+    method! fold_procedure_paragraph { paragraph; _ } acc =
+      acc
+      |> Cobol_ptree.Proc_division_visitor.fold_paragraph' v paragraph
+      |> Cobol_common.Visitor.skip
+
 
     method! fold_procedure_name' ({ loc; _ } as qn)
         ({ current_section = in_section; _ } as acc) =

src/vscode/superbol-vscode-platform/superbol_commands.ml

@@ -34,6 +34,22 @@ let command id handler =
   commands := command :: !commands;
   command
 
+let _editor_action_findReferences =
+  command "superbol.editor.action.findReferences" @@ Instance
+    begin fun _instance ~args ->
+      match args with
+      | [arg1; arg2] ->
+        let uri = Uri.parse (Ojs.string_of_js arg1) () in
+        let line = Ojs.get_prop_ascii arg2 "line" |> Ojs.int_of_js in
+        let character = Ojs.get_prop_ascii arg2 "character" |> Ojs.int_of_js in
+        let pos = Position.make ~line ~character in
+        let _ = Commands.executeCommand
+            ~command:"editor.action.findReferences"
+            ~args:[Uri.t_to_js uri; Position.t_to_js pos ]
+        in ()
+      | _ -> ()
+    end
+
 let _restart_language_server =
   command "superbol.server.restart" @@ Instance
     begin fun instance ~args:_ ->

test/lsp/lsp_codelens.ml

@@ -0,0 +1,199 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                        SuperBOL OSS Studio                             *)
+(*                                                                        *)
+(*  Copyright (c) 2022-2023 OCamlPro SAS                                  *)
+(*                                                                        *)
+(* All rights reserved.                                                   *)
+(* This source code is licensed under the GNU Affero General Public       *)
+(* License version 3 found in the LICENSE.md file in the root directory   *)
+(* of this source tree.                                                   *)
+(*                                                                        *)
+(**************************************************************************)
+
+open Lsp.Types
+open Lsp_testing
+
+let codelens doc : string -> unit =
+  let { end_with_postproc; projdir }, server = make_lsp_project () in
+  let server, prog = add_cobol_doc server ~projdir "prog.cob" doc in
+  let location_as_srcloc = new srcloc_resuscitator_cache in
+  let params = CodeLensParams.create () ~textDocument:prog in
+  LSP.Request.codelens server params |> List.rev
+  |> List.iter begin fun (codelens: CodeLens.t) ->
+    let location = Location.create ~range:codelens.range ~uri:prog.uri in
+    codelens.command
+    |> Option.iter begin fun (command: Command.t) ->
+      Pretty.out "%a%s@."
+        location_as_srcloc#pp location
+        command.title end
+  end;
+  end_with_postproc
+;;
+
+let%expect_test "codelens" =
+  let end_with_postproc = codelens {cobol|
+        IDENTIFICATION DIVISION.
+        PROGRAM-ID. prog.
+        DATA DIVISION.
+        WORKING-STORAGE SECTION.
+        01 AA.
+          02 BB PIC X.
+          02 BBprime REDEFINES BB PIC 9.
+          02 CC PIC X. 02 DD PIC X.
+          66 ABCD RENAMES BB THRU DD.
+        01 ZZ OCCURS 5 TIMES INDEXED BY INDEX1.
+          02 YY PIC X.
+          88 YYcond value "a".
+        PROCEDURE DIVISION.
+            MOVE aa TO aA.
+            STOP RUN.
+    |cobol} in
+  end_with_postproc [%expect.output];
+  [%expect {|
+    {"params":{"diagnostics":[],"uri":"file://__rootdir__/prog.cob"},"method":"textDocument/publishDiagnostics","jsonrpc":"2.0"}
+    __rootdir__/prog.cob:6.11:
+       3           PROGRAM-ID. prog.
+       4           DATA DIVISION.
+       5           WORKING-STORAGE SECTION.
+       6 >         01 AA.
+    ----              ^
+       7             02 BB PIC X.
+       8             02 BBprime REDEFINES BB PIC 9.
+    3 references
+    __rootdir__/prog.cob:7.13:
+       4           DATA DIVISION.
+       5           WORKING-STORAGE SECTION.
+       6           01 AA.
+       7 >           02 BB PIC X.
+    ----                ^
+       8             02 BBprime REDEFINES BB PIC 9.
+       9             02 CC PIC X. 02 DD PIC X.
+    3 references
+    __rootdir__/prog.cob:8.13:
+       5           WORKING-STORAGE SECTION.
+       6           01 AA.
+       7             02 BB PIC X.
+       8 >           02 BBprime REDEFINES BB PIC 9.
+    ----                ^
+       9             02 CC PIC X. 02 DD PIC X.
+      10             66 ABCD RENAMES BB THRU DD.
+    1 reference
+    __rootdir__/prog.cob:9.13:
+       6           01 AA.
+       7             02 BB PIC X.
+       8             02 BBprime REDEFINES BB PIC 9.
+       9 >           02 CC PIC X. 02 DD PIC X.
+    ----                ^
+      10             66 ABCD RENAMES BB THRU DD.
+      11           01 ZZ OCCURS 5 TIMES INDEXED BY INDEX1.
+    1 reference
+    __rootdir__/prog.cob:9.26:
+       6           01 AA.
+       7             02 BB PIC X.
+       8             02 BBprime REDEFINES BB PIC 9.
+       9 >           02 CC PIC X. 02 DD PIC X.
+    ----                             ^
+      10             66 ABCD RENAMES BB THRU DD.
+      11           01 ZZ OCCURS 5 TIMES INDEXED BY INDEX1.
+    2 references
+    __rootdir__/prog.cob:10.13:
+       7             02 BB PIC X.
+       8             02 BBprime REDEFINES BB PIC 9.
+       9             02 CC PIC X. 02 DD PIC X.
+      10 >           66 ABCD RENAMES BB THRU DD.
+    ----                ^
+      11           01 ZZ OCCURS 5 TIMES INDEXED BY INDEX1.
+      12             02 YY PIC X.
+    1 reference
+    __rootdir__/prog.cob:11.11:
+       8             02 BBprime REDEFINES BB PIC 9.
+       9             02 CC PIC X. 02 DD PIC X.
+      10             66 ABCD RENAMES BB THRU DD.
+      11 >         01 ZZ OCCURS 5 TIMES INDEXED BY INDEX1.
+    ----              ^
+      12             02 YY PIC X.
+      13             88 YYcond value "a".
+    1 reference
+    __rootdir__/prog.cob:11.40:
+       8             02 BBprime REDEFINES BB PIC 9.
+       9             02 CC PIC X. 02 DD PIC X.
+      10             66 ABCD RENAMES BB THRU DD.
+      11 >         01 ZZ OCCURS 5 TIMES INDEXED BY INDEX1.
+    ----                                           ^
+      12             02 YY PIC X.
+      13             88 YYcond value "a".
+    1 reference
+    __rootdir__/prog.cob:12.13:
+       9             02 CC PIC X. 02 DD PIC X.
+      10             66 ABCD RENAMES BB THRU DD.
+      11           01 ZZ OCCURS 5 TIMES INDEXED BY INDEX1.
+      12 >           02 YY PIC X.
+    ----                ^
+      13             88 YYcond value "a".
+      14           PROCEDURE DIVISION.
+    1 reference
+    __rootdir__/prog.cob:13.13:
+      10             66 ABCD RENAMES BB THRU DD.
+      11           01 ZZ OCCURS 5 TIMES INDEXED BY INDEX1.
+      12             02 YY PIC X.
+      13 >           88 YYcond value "a".
+    ----                ^
+      14           PROCEDURE DIVISION.
+      15               MOVE aa TO aA.
+    0 reference |}];;
+
+let%expect_test "codelens-procedure" =
+  let end_with_postproc = codelens {cobol|
+        IDENTIFICATION DIVISION.
+        PROGRAM-ID. prog.
+        PROCEDURE DIVISION.
+            AA SECTION.
+            BB.
+            PERFORM BB.
+            CC.
+            DD SECTION.
+            PERFORM AA.
+            PERFORM BB.
+            GO DD.
+            STOP RUN.
+    |cobol} in
+  end_with_postproc [%expect.output];
+  [%expect {|
+    {"params":{"diagnostics":[],"uri":"file://__rootdir__/prog.cob"},"method":"textDocument/publishDiagnostics","jsonrpc":"2.0"}
+    __rootdir__/prog.cob:5.12:
+       2           IDENTIFICATION DIVISION.
+       3           PROGRAM-ID. prog.
+       4           PROCEDURE DIVISION.
+       5 >             AA SECTION.
+    ----               ^
+       6               BB.
+       7               PERFORM BB.
+    2 references
+    __rootdir__/prog.cob:6.12:
+       3           PROGRAM-ID. prog.
+       4           PROCEDURE DIVISION.
+       5               AA SECTION.
+       6 >             BB.
+    ----               ^
+       7               PERFORM BB.
+       8               CC.
+    3 references
+    __rootdir__/prog.cob:8.12:
+       5               AA SECTION.
+       6               BB.
+       7               PERFORM BB.
+       8 >             CC.
+    ----               ^
+       9               DD SECTION.
+      10               PERFORM AA.
+    1 reference
+    __rootdir__/prog.cob:9.12:
+       6               BB.
+       7               PERFORM BB.
+       8               CC.
+       9 >             DD SECTION.
+    ----               ^
+      10               PERFORM AA.
+      11               PERFORM BB.
+    2 references |}];;