From 8a582b40dab12e60bf63a28bd7cf9aeb4831f81b Mon Sep 17 00:00:00 2001
From: Romain Tetley <romain.tetley@inria.fr>
Date: Tue, 10 Dec 2024 13:35:51 +0100
Subject: [PATCH] fix: give a more accurate representation of proof block

---
 client/src/protocol/types.ts            | 15 ++++++++++-----
 language-server/dm/document.ml          |  2 +-
 language-server/dm/documentManager.ml   |  8 +++++---
 language-server/protocol/proofState.ml  | 11 +++++++++--
 language-server/protocol/proofState.mli |  5 ++++-
 5 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/client/src/protocol/types.ts b/client/src/protocol/types.ts
index 74b555f59..77d2b10c6 100644
--- a/client/src/protocol/types.ts
+++ b/client/src/protocol/types.ts
@@ -149,17 +149,22 @@ export interface DocumentProofsRequest {
     textDocument: TextDocumentIdentifier;
 }
 
-type proof_step = {
+type ProofStatement = {
+    range: Range;
+    statement: string;
+}
+
+type ProofStep = {
     range: Range;
     tactic: string;
 };
 
-type proof_block = {
-    statement: string;
+type ProofBlock = {
+    statement: ProofStatement;
     range: Range;
-    proofs: proof_step[];
+    proofs: ProofStep[];
 };
 
 export interface DocumentProofsResponse {
-    proofs: proof_block[];
+    proofs: ProofBlock[];
 }
diff --git a/language-server/dm/document.ml b/language-server/dm/document.ml
index a71c81e98..a32cae5bf 100644
--- a/language-server/dm/document.ml
+++ b/language-server/dm/document.ml
@@ -176,7 +176,7 @@ let push_proof_step_in_outline document id (outline : outline) =
 let record_outline document id (ast : Synterp.vernac_control_entry) classif (outline : outline) =
   let open Vernacextend in
   match classif with
-  | VtProofStep _ -> push_proof_step_in_outline document id outline
+  | VtProofStep _ | VtQed _ -> push_proof_step_in_outline document id outline
   | VtStartProof (_, names) ->
     let vernac_gen_expr = ast.v.expr in
     let type_ = match vernac_gen_expr with
diff --git a/language-server/dm/documentManager.ml b/language-server/dm/documentManager.ml
index ab0d6af08..58d2377fc 100644
--- a/language-server/dm/documentManager.ml
+++ b/language-server/dm/documentManager.ml
@@ -328,10 +328,12 @@ let get_document_proofs st =
     | _ -> false
     in
   let mk_proof_block Document.{statement; proof; range } =
+    let statement = ProofState.mk_proof_statement statement range in
+    let last_step = List.hd proof in
+    let proof = List.rev proof in
+    let fst_step = List.hd proof in
+    let range = Range.create ~start:fst_step.range.start ~end_:last_step.range.end_ in
     let steps = List.map (fun Document.{tactic; range} -> ProofState.mk_proof_step tactic range) proof in
-    let step_ranges = List.map (fun (x: Document.proof_step) -> x.range) proof in
-    let range = List.fold_right RangeList.insert_or_merge_range step_ranges [range] in
-    let range = List.hd range in
     ProofState.mk_proof_block statement steps range
   in
   let proofs, _  = List.partition is_theorem outline in
diff --git a/language-server/protocol/proofState.ml b/language-server/protocol/proofState.ml
index 004e3222d..86549cbb2 100644
--- a/language-server/protocol/proofState.ml
+++ b/language-server/protocol/proofState.ml
@@ -17,13 +17,18 @@ open Lsp.Types
 
 open Ppx_yojson_conv_lib.Yojson_conv.Primitives
 
+type proof_statement = {
+  statement: string;
+  range: Range.t;
+} [@@deriving yojson]
+
 type proof_step = {
   tactic: string;
   range: Range.t;
 } [@@deriving yojson]
 
 type proof_block = {
-  statement: string;
+  statement: proof_statement;
   range: Range.t;
   steps: proof_step list;
 } [@@deriving yojson]
@@ -132,6 +137,8 @@ let get_proof ~previous diff_mode st =
       unfocusedGoals;
     }
 
+let mk_proof_statement statement range = {statement; range}
+
 let mk_proof_step tactic range = {tactic; range}
 
-let mk_proof_block statement steps range = {statement; steps; range}
\ No newline at end of file
+let mk_proof_block statement steps range = {steps; statement; range}
\ No newline at end of file
diff --git a/language-server/protocol/proofState.mli b/language-server/protocol/proofState.mli
index e1619027c..0acf7f232 100644
--- a/language-server/protocol/proofState.mli
+++ b/language-server/protocol/proofState.mli
@@ -13,6 +13,8 @@
 (**************************************************************************)
 open Settings
 
+type proof_statement [@@deriving yojson]
+
 type proof_step [@@deriving yojson]
 
 type proof_block [@@deriving yojson]
@@ -21,5 +23,6 @@ type t [@@deriving yojson]
 
 val get_proof : previous:Vernacstate.t option -> Goals.Diff.Mode.t -> Vernacstate.t -> t option
 
+val mk_proof_statement : string -> Lsp.Types.Range.t -> proof_statement
 val mk_proof_step : string -> Lsp.Types.Range.t -> proof_step
-val mk_proof_block : string -> proof_step list -> Lsp.Types.Range.t -> proof_block
\ No newline at end of file
+val mk_proof_block : proof_statement -> proof_step list -> Lsp.Types.Range.t -> proof_block
\ No newline at end of file