implement restricted constructor location

gated behind [-lsp-compat] flag

Signed-off-by: David Vulakh <[email protected]>

Author: dvulakh

src/analysis/stack_or_heap_enclosing.ml

@@ -10,19 +10,24 @@ type stack_or_heap =
 
 type stack_or_heap_enclosings = (Location.t * stack_or_heap) list
 
-let from_nodes ~pos ~path =
+let from_nodes ~lsp_compat ~pos ~path =
   let[@tail_mod_cons] rec tails = function
     | hd :: tl -> (hd, tl) :: tails tl
     | [] -> []
   in
+  let cursor_is_inside ({ loc_start; loc_end; _ } : Location.t) =
+    Lexing.compare_pos pos loc_start >= 0 && Lexing.compare_pos pos loc_end <= 0
+  in
   let aux node parent =
     let open Browse_raw in
-    let ret mode_result = Some (Mbrowse.node_loc node, mode_result) in
-    let ret_alloc alloc_mode = ret (Alloc_mode alloc_mode) in
-    let ret_no_alloc reason = ret (No_alloc { reason }) in
-    let ret_maybe_alloc reason = function
-      | Some alloc_mode -> ret_alloc alloc_mode
-      | None -> ret_no_alloc reason
+    let ret ?(loc = Mbrowse.node_loc node) mode_result =
+      Some (loc, mode_result)
+    in
+    let ret_alloc ?loc alloc_mode = ret ?loc (Alloc_mode alloc_mode) in
+    let ret_no_alloc ?loc reason = ret ?loc (No_alloc { reason }) in
+    let ret_maybe_alloc ?loc reason = function
+      | Some alloc_mode -> ret_alloc ?loc alloc_mode
+      | None -> ret_no_alloc ?loc reason
     in
     match (node, parent) with
     | ( Pattern { pat_desc = Tpat_var _; _ },
@@ -32,7 +37,7 @@ let from_nodes ~pos ~path =
       ) -> ret (Alloc_mode alloc_mode.mode)
     | Expression { exp_desc; _ }, _ -> (
       match exp_desc with
-      | Texp_function { alloc_mode; body; _ } ->
+      | Texp_function { alloc_mode; body; _ } -> (
         let body_loc =
           (* A function expression is often in a non-obvious way the nearest enclosing
              allocating expression. To avoid confusion, we only consider a function
@@ -60,26 +65,26 @@ let from_nodes ~pos ~path =
                 }
             | [] -> None)
         in
-        let cursor_is_inside_function_body =
-          match body_loc with
-          | None -> false
-          | Some { loc_start; loc_end; loc_ghost = _ } ->
-            Lexing.compare_pos pos loc_start >= 0
-            && Lexing.compare_pos pos loc_end <= 0
-        in
-        if cursor_is_inside_function_body then None
-        else ret (Alloc_mode alloc_mode.mode)
+        match body_loc with
+        | Some loc when cursor_is_inside loc -> None
+        | _ -> ret (Alloc_mode alloc_mode.mode))
       | Texp_array (_, _, _, alloc_mode) -> ret (Alloc_mode alloc_mode.mode)
-      | Texp_construct (_, { cstr_repr; _ }, args, maybe_alloc_mode) -> (
+      | Texp_construct
+          ({ loc; txt = _lident }, { cstr_repr; _ }, args, maybe_alloc_mode)
+        -> (
+        let loc =
+          if lsp_compat && cursor_is_inside loc then Some loc else None
+        in
         match maybe_alloc_mode with
-        | Some alloc_mode -> ret (Alloc_mode alloc_mode.mode)
+        | Some alloc_mode -> ret ?loc (Alloc_mode alloc_mode.mode)
         | None -> (
           match args with
-          | [] -> ret_no_alloc "constructor without arguments"
+          | [] -> ret_no_alloc ?loc "constructor without arguments"
           | _ :: _ -> (
             match cstr_repr with
-            | Variant_unboxed -> ret_no_alloc "unboxed constructor"
-            | Variant_extensible | Variant_boxed _ -> ret Unexpected_no_alloc)))
+            | Variant_unboxed -> ret_no_alloc ?loc "unboxed constructor"
+            | Variant_extensible | Variant_boxed _ ->
+              ret ?loc Unexpected_no_alloc)))
       | Texp_record { representation; alloc_mode = maybe_alloc_mode; _ } -> (
         match (maybe_alloc_mode, representation) with
         | _, Record_inlined _ -> None

src/analysis/stack_or_heap_enclosing.mli

@@ -27,6 +27,7 @@ type stack_or_heap =
 type stack_or_heap_enclosings = (Location.t * stack_or_heap) list
 
 val from_nodes :
+  lsp_compat:bool ->
   pos:Lexing.position ->
   path:(Env.t * Browse_raw.node * Query_protocol.is_tail_position) list ->
   stack_or_heap_enclosings

src/commands/new_commands.ml

@@ -708,6 +708,8 @@ let all_commands =
          of expressions known not to allocate, give \"unknown (does your code \
          contain a type error?)\". As suggested by the message, this should \
          only occur if the input does not typecheck.\n\n\
+         `-lsp-compat` can be used to change the locations reported for better \
+         LSP hover interaction.\n\n\
          `-index` can be used to print only one \"stack-or-heap\".\n\n\
          The result is returned as a list of:\n\
          ```javascript\n\
@@ -719,21 +721,31 @@ let all_commands =
          ```"
       ~spec:
         [ arg "-position" "<position> Position to complete"
-            (marg_position (fun pos (expr, cursor, _pos, index) ->
-                 (expr, cursor, pos, index)));
+            (marg_position (fun pos (expr, cursor, _pos, lsp_compat, index) ->
+                 (expr, cursor, pos, lsp_compat, index)));
+          optional "-lsp-compat"
+            "<bool> Report ranges that are less accurate but work better with \
+             LSP hover"
+            (Marg.param "bool"
+               (fun lsp_compat (expr, cursor, pos, _lsp_compat, index) ->
+                 match bool_of_string lsp_compat with
+                 | lsp_compat -> (expr, cursor, pos, lsp_compat, index)
+                 | exception _ -> failwith "lsp_compat should be a bool"));
           optional "-index" "<int> Only print type of <index>'th result"
-            (Marg.param "int" (fun index (expr, cursor, pos, _index) ->
+            (Marg.param "int"
+               (fun index (expr, cursor, pos, lsp_compat, _index) ->
                  match int_of_string index with
-                 | index -> (expr, cursor, pos, Some index)
+                 | index -> (expr, cursor, pos, lsp_compat, Some index)
                  | exception _ -> failwith "index should be an integer"))
         ]
-      ~default:("", -1, `None, None)
+      ~default:("", -1, `None, false, None)
       begin
-        fun buffer (_, _, pos, index) ->
+        fun buffer (_, _, pos, lsp_compat, index) ->
           match pos with
           | `None -> failwith "-position <pos> is mandatory"
           | #Msource.position as pos ->
-            run buffer (Query_protocol.Stack_or_heap_enclosing (pos, index))
+            run buffer
+              (Query_protocol.Stack_or_heap_enclosing (pos, lsp_compat, index))
       end;
     command "type-enclosing"
       ~doc:

src/commands/query_json.ml

@@ -56,12 +56,13 @@ let dump (type a) : a t -> json =
   | Type_expr (expr, pos) ->
     mk "type-expression"
       [ ("expression", `String expr); ("position", mk_position pos) ]
-  | Stack_or_heap_enclosing (pos, index) ->
+  | Stack_or_heap_enclosing (pos, lsp_compat, index) ->
     mk "stack-or-heap-enclosing"
       [ ( "index",
           match index with
           | None -> `String "all"
           | Some n -> `Int n );
+        ("lsp-compat", `Bool lsp_compat);
         ("position", mk_position pos)
       ]
   | Type_enclosing (opt_cursor, pos, index) ->

src/frontend/query_commands.ml

@@ -264,7 +264,7 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a = function
     let context = Context.Expr in
     ignore (Type_utils.type_in_env ~verbosity ~context env ppf source : bool);
     to_string ()
-  | Stack_or_heap_enclosing (pos, index) ->
+  | Stack_or_heap_enclosing (pos, lsp_compat, index) ->
     let typer = Mpipeline.typer_result pipeline in
 
     (* Optimise allocations only on programs that have type-checked. *)
@@ -292,7 +292,7 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a = function
       | browse -> Browse_misc.annotate_tail_calls browse
     in
 
-    let result = Stack_or_heap_enclosing.from_nodes ~pos ~path in
+    let result = Stack_or_heap_enclosing.from_nodes ~lsp_compat ~pos ~path in
 
     let all_results =
       List.mapi result ~f:(fun i (loc, text) ->

src/frontend/query_protocol.ml

@@ -180,7 +180,7 @@ end
 type _ t =
   | Type_expr (* *) : string * Msource.position -> string t
   | Stack_or_heap_enclosing (* *) :
-      Msource.position * int option
+      Msource.position * bool * int option
       -> (Location.t * [ `String of string | `Index of int ]) list t
   | Type_enclosing (* *) :
       (string * int) option * Msource.position * int option

tests/test-dirs/stack-or-heap.t/run.t

@@ -43,12 +43,13 @@ escape characters in string literals, so we use the revert-newlines script.
   >   file=$1
   >   position=$2
   >   index=$3
+  >   trailing=$4
   >   line=$(echo "$position" | cut -d ':' -f 1)
   >   col=$(echo "$position" | cut -d ':' -f 2)
   >   highlight_range "$file" $line $(expr $col - 1) $line $col
   >   merlin=$(
   >     $MERLIN single stack-or-heap-enclosing -position "$position" -verbosity "$verbosity" \
-  >       -filename "$file" < "$file" | revert-newlines
+  >       -filename "$file" $trailing < "$file" | revert-newlines
   >   )
   >   echo
   >   if [ "$(echo "$merlin" | jq ".value[$index]")" != null ]
@@ -71,7 +72,7 @@ escape characters in string literals, so we use the revert-newlines script.
   >   do
   >     for i in $(seq 0 $until)
   >     do
-  >       run "$orig_file.tmp.ml" $lc $i
+  >       run "$orig_file.tmp.ml" $lc $i "$3"
   >     done
   >   done
   >   rm "$orig_file.tmp.ml"
@@ -184,6 +185,113 @@ escape characters in string literals, so we use the revert-newlines script.
 
   "not an allocation (unboxed constructor)"
 
+
+  $ run_annotated_file constructors.ml 1 "-lsp-compat true"
+  |  Some (g z)
+  |        ^
+  
+  |  Some (g z)
+  |  ^^^^^^^^^^
+  
+  "heap"
+  
+  |  exclave_ Some (g z)
+  |                 ^
+  
+  |  exclave_ Some (g z)
+  |           ^^^^^^^^^^
+  
+  "stack"
+  
+  |  let z = Some (g x) in
+  |                ^
+  
+  |  let z = Some (g x) in
+  |          ^^^^^^^^^^
+  
+  "stack"
+  
+  |  Some (g z)
+  |    ^
+  
+  |  Some (g z)
+  |  ^^^^
+  
+  "heap"
+  
+  |  exclave_ Some (g z)
+  |             ^
+  
+  |  exclave_ Some (g z)
+  |           ^^^^
+  
+  "stack"
+  
+  |  let z = Some (g x) in
+  |            ^
+  
+  |  let z = Some (g x) in
+  |          ^^^^
+  
+  "stack"
+  
+  |  None
+  |    ^
+  
+  |  None
+  |  ^^^^
+  
+  "not an allocation (constructor without arguments)"
+  
+  |  exclave_ None
+  |             ^
+  
+  |  exclave_ None
+  |           ^^^^
+  
+  "not an allocation (constructor without arguments)"
+  
+  |  f (Some x);
+  |          ^
+  
+  |  f (Some x);
+  |    ^^^^^^^^
+  
+  "stack"
+  
+  |  f (local_ Some x);
+  |                 ^
+  
+  |  f (local_ Some x);
+  |    ^^^^^^^^^^^^^^^
+  
+  "stack"
+  
+  |  f (Some x)
+  |          ^
+  
+  |  f (Some x)
+  |    ^^^^^^^^
+  
+  "heap"
+  
+  |let g x = f (Some x) [@nontail]
+  |                  ^
+  
+  |let g x = f (Some x) [@nontail]
+  |            ^^^^^^^^
+  
+  "stack"
+  
+  |  Box (g z)
+  |       ^
+  
+  |  Box (g z)
+  |  ^^^^^^^^^
+  
+  "not an allocation (unboxed constructor)"
+  
+
 (II) Variants
 
   $ run_annotated_file variants.ml