Merge pull request #211 from tlaplus/lsp-rename-step-v2

Small fixes in LSP step rename command.

Author: kape1395

lsp/doc/proof_step_rename.gif

@@ -2,20 +2,49 @@
 open Tlapm_lib
 module StepMap = Map.Make (String)
 
-exception Found_the_step of string * int * Range.t list
+module StepInfo = struct
+  type t = { name : string; label_offset : int; step_ranges : Range.t list }
+
+  let make level suffix step_loc =
+    let prefix_str = Printf.sprintf "<%d>" level in
+    let step_str = Printf.sprintf "<%d>%s" level suffix in
+    let range = Range.of_locus_must step_loc in
+    let range = Range.of_len range (String.length step_str) in
+    {
+      name = step_str;
+      label_offset = String.length prefix_str;
+      step_ranges = [ range ];
+    }
+
+  let with_range si range = { si with step_ranges = range :: si.step_ranges }
+
+  let matching_range si range =
+    List.find_opt
+      (fun r -> Range.intersect range r || Range.touches range r)
+      si.step_ranges
+
+  let step_label si =
+    String.sub si.name si.label_offset (String.length si.name - si.label_offset)
+
+  let step_label_range si range = Range.crop_line_prefix range si.label_offset
+end
+
+exception Found_the_step of StepInfo.t
 
 class step_rename_visitor (at_loc : Range.t) =
   object (self : 'self)
     inherit Module.Visit.map as m_super
     inherit [unit] Proof.Visit.iter as p_super
-    val mutable step_map : (int * Range.t list) StepMap.t = StepMap.empty
+    val mutable step_map : StepInfo.t StepMap.t = StepMap.empty
 
-    (* Leaf at the `Proof.Visit.iter as p_super`.
+    (* Leaf at the `Module.Visit.map as m_super`.
        We will look at `orig_pf`, it is closer to the initial source code. *)
-    method! theorem cx name sq naxs _pf orig_pf summ =
-      let pf = orig_pf in
-      let ctx = Expr.Visit.empty () in
-      self#proof ctx pf;
+    method! theorem cx name sq naxs pf orig_pf summ =
+      (if Range.lines_intersect at_loc (Range.of_wrapped_must orig_pf) then
+         (* Only look for steps in proofs that intersects with the user selection. *)
+         let pf = orig_pf in
+         let ctx = Expr.Visit.empty () in
+         self#proof ctx pf);
       m_super#theorem cx name sq naxs pf orig_pf summ
 
     (* After each proof drop new step names that were defined within the proof.
@@ -33,14 +62,13 @@ class step_rename_visitor (at_loc : Range.t) =
       | _ -> ());
       (* If we are removing some step from the context, all its references
          are collected already. Check maybe we found what we were looking for. *)
-      let matches range = Range.intersect range at_loc in
       step_map <-
         StepMap.filter
-          (fun step_str (label_offset, step_ranges) ->
-            if not (StepMap.mem step_str step_map_before) then (
-              if List.exists matches step_ranges then
-                raise (Found_the_step (step_str, label_offset, step_ranges));
-              false)
+          (fun step_str step_info ->
+            if not (StepMap.mem step_str step_map_before) then
+              match StepInfo.matching_range step_info at_loc with
+              | None -> false
+              | Some _ -> raise (Found_the_step step_info)
             else true)
           step_map
 
@@ -63,32 +91,26 @@ class step_rename_visitor (at_loc : Range.t) =
     method add_step_opt stepno_opt (step_loc : Loc.locus) =
       match stepno_opt with
       | Proof.T.Named (level, suffix, false) ->
-          let prefix_str = Printf.sprintf "<%d>" level in
-          let step_str = Printf.sprintf "<%d>%s" level suffix in
-          let range = Range.of_locus_must step_loc in
-          let range = Range.of_len range (String.length step_str) in
-          let step_info = (String.length prefix_str, [ range ]) in
-          step_map <- StepMap.add step_str step_info step_map
+          let si = StepInfo.make level suffix step_loc in
+          step_map <- StepMap.add si.name si step_map
       | Proof.T.Named (_level, _suffix, true) -> ()
       | Proof.T.Unnamed (_, _) -> ()
 
     (* Private helper. *)
     method add_step_range_opt (step_name : string) (locus : Loc.locus) =
       match StepMap.find_opt step_name step_map with
       | None -> ()
-      | Some (label_offset, ranges) ->
-          let new_info = (label_offset, Range.of_locus_must locus :: ranges) in
-          step_map <- StepMap.add step_name new_info step_map
+      | Some si ->
+          let si = StepInfo.with_range si (Range.of_locus_must locus) in
+          step_map <- StepMap.add step_name si step_map
   end
 
-let find_ranges (at_loc : Range.t) (mule : Module.T.mule) :
-    (string * int * Range.t list) option =
+let find_ranges (at_loc : Range.t) (mule : Module.T.mule) : StepInfo.t option =
   let visitor = new step_rename_visitor at_loc in
   try
     let _ = visitor#tla_module_root mule in
     None
-  with Found_the_step (step_str, label_offset, step_ranges) ->
-    Some (step_str, label_offset, step_ranges)
+  with Found_the_step si -> Some si
 
 let%test_module "poc renames" =
   (module struct
@@ -111,47 +133,47 @@ let%test_module "poc renames" =
     let%test_unit "find <1>1 by step name" =
       match find_ranges (Range.of_ints ~lf:3 ~cf:6 ~lt:3 ~ct:6) mule with
       | None -> assert false
-      | Some (st_name, st_lb_o, st_ranges) ->
-          assert (st_name = "<1>1");
-          assert (st_lb_o = 3);
-          assert (List.length st_ranges = 3);
-          ()
+      | Some si ->
+          assert (si.name = "<1>1");
+          assert (si.label_offset = 3);
+          assert (List.length si.step_ranges = 3)
 
     let%test_unit "find <1>1 by step use" =
       match find_ranges (Range.of_ints ~lf:5 ~cf:18 ~lt:5 ~ct:18) mule with
       | None -> assert false
-      | Some (st_name, st_lb_o, st_ranges) ->
-          assert (st_name = "<1>1");
-          assert (st_lb_o = 3);
-          assert (List.length st_ranges = 3);
-          ()
+      | Some si ->
+          assert (si.name = "<1>1");
+          assert (si.label_offset = 3);
+          assert (List.length si.step_ranges = 3)
 
     let%test_unit "find <1>2" =
       match find_ranges (Range.of_ints ~lf:4 ~cf:6 ~lt:4 ~ct:6) mule with
       | None -> assert false
-      | Some (st_name, st_lb_o, st_ranges) ->
-          assert (st_name = "<1>2");
-          assert (st_lb_o = 3);
-          assert (List.length st_ranges = 2);
-          ()
+      | Some si ->
+          assert (si.name = "<1>2");
+          assert (si.label_offset = 3);
+          assert (List.length si.step_ranges = 2)
+
+    let%test_unit "find <1>2 - start col" =
+      match find_ranges (Range.of_ints ~lf:4 ~cf:5 ~lt:4 ~ct:5) mule with
+      | None -> assert false
+      | Some si -> assert (si.name = "<1>2")
+
+    let%test_unit "find <1>2 - end col" =
+      match find_ranges (Range.of_ints ~lf:4 ~cf:9 ~lt:4 ~ct:9) mule with
+      | None -> assert false
+      | Some si -> assert (si.name = "<1>2")
 
     let%test_unit "find <1>1 -- qed" =
       match find_ranges (Range.of_ints ~lf:5 ~cf:6 ~lt:5 ~ct:6) mule with
       | None -> assert false
-      | Some (st_name, st_lb_o, st_ranges) ->
-          assert (st_name = "<1>q");
-          assert (st_lb_o = 3);
-          assert (List.length st_ranges = 1);
-          ()
+      | Some si ->
+          assert (si.name = "<1>q");
+          assert (si.label_offset = 3);
+          assert (List.length si.step_ranges = 1)
 
     let%test_unit "not found" =
       match find_ranges (Range.of_ints ~lf:2 ~cf:6 ~lt:2 ~ct:6) mule with
       | None -> ()
-      | Some (_st_name, _st_lb_o, _st_ranges) -> assert false
+      | Some _si -> assert false
   end)
-
-let step_label step_name label_offset =
-  String.sub step_name label_offset (String.length step_name - label_offset)
-
-let step_label_range range label_offset =
-  Range.crop_line_prefix range label_offset

lsp/doc/proof_step_renumber.gif

@@ -1,6 +1,17 @@
-val find_ranges :
-  Range.t -> Tlapm_lib.Module.T.mule -> (string * int * Range.t list) option
-(** Find ranges related to the step at the cursor position. *)
+module StepInfo : sig
+  type t = { name : string; label_offset : int; step_ranges : Range.t list }
+
+  val matching_range : t -> Range.t -> Range.t option
+  (** Returns a range where the step was mentioned and matches the specified
+      range. *)
+
+  val step_label : t -> string
+  (** Returns the label part of the step. *)
 
-val step_label : string -> int -> string
-val step_label_range : Range.t -> int -> Range.t
+  val step_label_range : t -> Range.t -> Range.t
+  (** Returns a label part range of a step. The range must be one of the step
+      ranges. *)
+end
+
+val find_ranges : Range.t -> Tlapm_lib.Module.T.mule -> StepInfo.t option
+(** Find ranges related to the step at the cursor position. *)

lsp/lib/analysis/step_rename.ml

@@ -68,8 +68,7 @@ let of_lsp_range (range : LspT.Range.t) =
   R (f, t)
 
 let of_lsp_position (pos : LspT.Position.t) =
-  let p = (pos.line + 1, pos.character + 1) in
-  R (p, p)
+  R ((pos.line + 1, pos.character + 1), (pos.line + 1, pos.character))
 
 let of_string_opt s =
   match String.split_on_char ':' s with
@@ -117,6 +116,12 @@ let of_all = R ((0, 0), (0, 0))
 (** [before p r] means range [r] is before point [p]. *)
 let before p r = Position.less (till r) p
 
+(** [touches a b] is true, if ranges [a] and [b] touches with their endpoints.
+    Consider here that the ranges here are with the end character inclusive.
+    I.e. "..-x:10" touches with "x:11-.." *)
+let touches (R ((alf, acf), (alt, act))) (R ((blf, bcf), (blt, bct))) =
+  (alt = blf && act + 1 = bcf) || (blt = alf && bct + 1 = acf)
+
 (** [intersect a b] is true, if ranges [a] and [b] overlaps. *)
 let intersect a b =
   Position.leq (from a) (till b) && Position.leq (from b) (till a)

lsp/lib/analysis/step_rename.mli

@@ -33,6 +33,7 @@ val string_of_range : t -> string
 val string_of_pos : Position.t -> string
 val compare : t -> t -> int
 val before : Position.t -> t -> bool
+val touches : t -> t -> bool
 val intersect : t -> t -> bool
 val lines_intersect : t -> t -> bool
 val line_covered : t -> Position.t -> bool

lsp/lib/range.ml

@@ -237,13 +237,14 @@ module Make (CB : Callbacks) = struct
     let uri = params.textDocument.uri in
     let pos = Range.of_lsp_position params.position in
     CB.with_docs cb_state @@ fun cb_st docs ->
-    let f _vsn mule = Analysis.Step_rename.find_ranges pos mule in
+    let open Analysis.Step_rename in
+    let f _vsn mule = find_ranges pos mule in
     let docs, res = Docs.on_parsed_mule_latest docs uri f in
     let cb_st =
       match res with
       | None -> reply_ok jsonrpc_req `Null cb_st
-      | Some (step_name, label_offset, step_ranges) -> (
-          match List.find_opt (Range.intersect pos) step_ranges with
+      | Some si -> (
+          match StepInfo.matching_range si pos with
           | None -> reply_ok jsonrpc_req `Null cb_st
           | Some step_range ->
               reply_ok jsonrpc_req
@@ -252,12 +253,8 @@ module Make (CB : Callbacks) = struct
                      ( "range",
                        LspT.Range.yojson_of_t
                          (Range.as_lsp_range
-                            (Analysis.Step_rename.step_label_range step_range
-                               label_offset)) );
-                     ( "placeholder",
-                       `String
-                         (Analysis.Step_rename.step_label step_name label_offset)
-                     );
+                            (StepInfo.step_label_range si step_range)) );
+                     ("placeholder", `String (StepInfo.step_label si));
                    ])
                 cb_st)
     in
@@ -267,27 +264,27 @@ module Make (CB : Callbacks) = struct
       (params : LspT.RenameParams.t) cb_state =
     let uri = params.textDocument.uri in
     let pos = Range.of_lsp_position params.position in
-    let new_text = params.newName in
+    let newText = params.newName in
     CB.with_docs cb_state @@ fun cb_st docs ->
-    let f _vsn mule = Analysis.Step_rename.find_ranges pos mule in
+    let open Analysis.Step_rename in
+    let f _vsn mule = find_ranges pos mule in
     let docs, res = Docs.on_parsed_mule_latest docs uri f in
     let cb_st =
       match res with
       | None -> reply_ok jsonrpc_req `Null cb_st
-      | Some (_step_name, label_offset, step_ranges) -> (
-          match List.find_opt (Range.intersect pos) step_ranges with
+      | Some si -> (
+          match StepInfo.matching_range si pos with
           | None -> reply_ok jsonrpc_req `Null cb_st
           | Some _step_range ->
               let edits =
                 List.map
                   (fun step_range ->
                     let range =
                       Range.as_lsp_range
-                        (Analysis.Step_rename.step_label_range step_range
-                           label_offset)
+                        (StepInfo.step_label_range si step_range)
                     in
-                    LspT.TextEdit.create ~newText:new_text ~range)
-                  step_ranges
+                    LspT.TextEdit.create ~newText ~range)
+                  si.step_ranges
               in
               let workspace_edit =
                 LspT.WorkspaceEdit.create ~changes:[ (uri, edits) ] ()