feat: added procedure codelens

OCamlPro · Aug 8, 2024 · 48c2947 · 48c2947
1 parent 7bf02aa
commit 48c2947
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 6 deletions.
diff --git a/src/lsp/cobol_lsp/lsp_request.ml b/src/lsp/cobol_lsp/lsp_request.ml
@@ -597,15 +597,33 @@ module PosSet = Set.Make(struct
 
 let codelens_positions ~uri group =
   let filename = Lsp.Uri.to_path uri in
+  let open struct
+    include Cobol_common.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
       inherit [_] Cobol_unit.Visitor.folder
-      method! fold_procedure _ = Cobol_common.Visitor.skip
-      method! fold_qualname' { loc; _ } acc =
-        let range = Lsp_position.range_of_srcloc_in ~filename loc in
-        PosSet.add range.start acc
-        |> Cobol_common.Visitor.do_children
-    end group PosSet.empty
+      method! fold_procedure _ = set_context ProcedureDiv
+      method! fold_paragraph' _ = skip
+      method! fold_data_definitions _ = set_context DataDiv
+      method! fold_procedure_name' = take_when_in ProcedureDiv
+      method! fold_qualname' = take_when_in DataDiv
+    end group (None, PosSet.empty)
+  |> snd
 
 let handle_codelens registry ({ textDocument; _ }: CodeLensParams.t) =
   try_with_main_document_data registry textDocument

diff --git a/test/lsp/lsp_codelens.ml b/test/lsp/lsp_codelens.ml
@@ -169,3 +169,58 @@ let%expect_test "codelens" =
       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.
+    4 references
+    __rootdir__/prog.cob:6.12:
+       3           PROGRAM-ID. prog.
+       4           PROCEDURE DIVISION.
+       5               AA SECTION.
+       6 >             BB.
+    ----               ^
+       7               PERFORM BB.
+       8               CC.
+    4 references
+    __rootdir__/prog.cob:8.12:
+       5               AA SECTION.
+       6               BB.
+       7               PERFORM BB.
+       8 >             CC.
+    ----               ^
+       9               DD SECTION.
+      10               PERFORM AA.
+    2 references
+    __rootdir__/prog.cob:9.12:
+       6               BB.
+       7               PERFORM BB.
+       8               CC.
+       9 >             DD SECTION.
+    ----               ^
+      10               PERFORM AA.
+      11               PERFORM BB.
+    4 references |}];;