Better LSP hover interaction in stack-or-heap

src/analysis/stack_or_heap_enclosing.ml

@@ -10,20 +10,41 @@ type stack_or_heap =
 
 type stack_or_heap_enclosings = (Location.t * stack_or_heap) list
 
-let from_nodes ~pos ~path =
-  let aux (_env, node, _tail) =
+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 with
-    | Expression { exp_desc; _ } -> (
+    match (node, parent) with
+    | ( Pattern { pat_desc = Tpat_var _; _ },
+        Some
+          (Value_binding
+            { vb_expr = { exp_desc = Texp_function { alloc_mode; _ }; _ };
+              vb_loc;
+              _
+            }) ) ->
+      (* The location that most sensibly corresponds to the "allocation" is the entire
+         value binding. However, the LSP hover at this point will describe just the
+         pattern, so we don't override the location in the [lsp_compat] regime. *)
+      let loc = if lsp_compat then None else Some vb_loc in
+      ret ?loc (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
@@ -51,26 +72,30 @@ 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 =
+          (* The location of the "allocation" here is the entire expression, but the LSP
+             hover for a constructor reports information just for the constructor (not the
+             entire [Texp_construct] expression), so we override the location in the
+             [lsp_compat] regime. *)
+          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
@@ -90,4 +115,8 @@ let from_nodes ~pos ~path =
       | _ -> None)
     | _ -> None
   in
-  List.filter_map ~f:aux path
+  path
+  |> List.map ~f:(fun (_, node, _) -> node)
+  |> tails
+  |> List.filter_map ~f:(fun (node, ancestors) ->
+         aux node (List.nth_opt ancestors 0))

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/lsp_compat.ml

@@ -0,0 +1,53 @@
+(* Cursor on the constructor itself (we treat this case specially to improve LSP
+   compatibility) *)
+
+let f g x y =
+  let z = x + y in
+  Some (g z)
+ (* ^ *)
+;;
+
+let f g x y =
+  let z = x + y in
+  exclave_ Some (g z)
+          (* ^ *)
+;;
+
+let f g x y =
+  let z = Some (g x) in
+         (* ^ *)
+  y
+;;
+
+(* Pattern of a [let]-bound function (we treat this case specially to improve LSP
+   compatibility) *)
+
+let f g x y =
+ (* ^ *)
+  let z = x + y in
+  exclave_ Some (g z)
+and h g x y =
+ (* ^ *)
+  let z = x + y in
+  exclave_ Some (g z)
+;;
+
+let ignore (local_ _) = ()
+
+let () =
+  let f g x y =
+   (* ^ *)
+    let z = x + y in
+    exclave_ Some (g z)
+  and h g x y =
+   (* ^ *)
+    let z = x + y in
+    exclave_ Some (g z)
+  in
+  ignore f;
+  ignore h
+
+(* Ensure other [let]-bound patterns aren't treated this way *)
+
+let x = Some 5
+ (* ^ *)