diff --git a/src/lsp/cobol_lsp/lsp_request.ml b/src/lsp/cobol_lsp/lsp_request.ml index cf7c2474c..6c2581490 100644 --- 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 index abde8836a..36e5ef10a 100644 --- 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 |}];;