Fix locate command in presence of let punning

src/analysis/context.ml

@@ -132,12 +132,12 @@ let inspect_expression ~cursor ~lid e : t =
   | _ ->
     Expr
 
-let inspect_browse_tree ~cursor lid browse : t option =
+let inspect_browse_tree ?let_pun_behavior ~cursor lid browse : t option =
   log ~title:"inspect_context" "current node is: [%s]"
     (String.concat ~sep:"|" (
       List.map ~f:(Mbrowse.print ()) browse
     ));
-  match Mbrowse.enclosing cursor browse with
+  match Mbrowse.enclosing ?let_pun_behavior cursor browse with
   | [] ->
     log ~title:"inspect_context"
       "no enclosing around: %a" Lexing.print_position cursor;

src/analysis/context.mli

@@ -55,4 +55,5 @@ val to_string : t -> string
   breaking the context inference.
 *)
 val inspect_browse_tree :
+  ?let_pun_behavior:Mbrowse.Let_pun_behavior.t ->
   cursor:Std.Lexing.position -> Longident.t -> Mbrowse.t list -> t option

src/analysis/locate.ml

@@ -836,7 +836,7 @@ let from_path ~config ~env ~local_defs ~namespace path =
   | None -> `Not_in_env (Path.name path)
   | Some decl -> from_path ~config ~env ~local_defs ~decl path
 
-let infer_namespace ?namespaces ~pos lid browse is_label =
+let infer_namespace ?let_pun_behavior ?namespaces ~pos lid browse is_label =
   match namespaces with
   | Some nss ->
     if not is_label
@@ -850,7 +850,7 @@ let infer_namespace ?namespaces ~pos lid browse is_label =
       `Error `Missing_labels_namespace
     )
   | None ->
-    match Context.inspect_browse_tree ~cursor:pos lid [browse], is_label with
+    match Context.inspect_browse_tree ?let_pun_behavior ~cursor:pos lid [browse], is_label with
     | None, _ ->
       log ~title:"from_string" "already at origin, doing nothing" ;
       `Error `At_origin
@@ -864,13 +864,13 @@ let infer_namespace ?namespaces ~pos lid browse is_label =
         "dropping inferred context, it is not precise enough";
       `Ok [ `Labels ]
 
-let from_string ~config ~env ~local_defs ~pos ?namespaces path =
+let from_string ~config ~env ~local_defs ~pos ?let_pun_behavior ?namespaces path =
   File_switching.reset ();
   let browse = Mbrowse.of_typedtree local_defs in
   let lid = Type_utils.parse_longident path in
   let from_lid lid =
     let ident, is_label = Longident.keep_suffix lid in
-    match infer_namespace ?namespaces ~pos lid browse is_label with
+    match infer_namespace ?let_pun_behavior ?namespaces ~pos lid browse is_label with
     | `Error e -> e
     | `Ok nss ->
       log ~title:"from_string"

src/analysis/locate.mli

@@ -79,6 +79,7 @@ val from_string
   -> env:Env.t
   -> local_defs:Mtyper.typedtree
   -> pos:Lexing.position
+  -> ?let_pun_behavior:Mbrowse.Let_pun_behavior.t
   -> ?namespaces:Env_lookup.Namespace.inferred_basic list
   -> string
   -> [> `File_not_found of string

src/frontend/query_commands.ml

@@ -595,6 +595,7 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a =
     let typer = Mpipeline.typer_result pipeline in
     let local_defs = Mtyper.get_typedtree typer in
     let pos = Mpipeline.get_lexing_pos pipeline pos in
+    let let_pun_behavior = Mbrowse.Let_pun_behavior.Prefer_expression in
     let env, _ = Mbrowse.leaf_node (Mtyper.node_at typer pos) in
     let path =
       match patho with
@@ -614,7 +615,7 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a =
         traverse_aliases = true
       }
     in
-    begin match Locate.from_string ~config ~env ~local_defs ~pos path with
+    begin match Locate.from_string ~config ~env ~local_defs ~pos ~let_pun_behavior path with
     | `Found { file; location; _ } ->
       Locate.log ~title:"result"
         "found: %s"  file;

src/kernel/mbrowse.ml

@@ -33,6 +33,12 @@ open Browse_raw
 type node = Browse_raw.node
 type t = (Env.t * node) list
 
+module Let_pun_behavior = struct
+  type t =
+    | Prefer_expression
+    | Prefer_pattern
+end
+
 let node_of_binary_part = Browse_raw.node_of_binary_part
 
 let fold_node f env t acc =
@@ -71,7 +77,6 @@ let node_loc node = approximate_loc Browse_raw.node_real_loc node
 let node_merlin_loc node = approximate_loc Browse_raw.node_merlin_loc node
 
 let leaf_node = List.hd
-let leaf_loc t = node_loc (snd (leaf_node t))
 
 let drop_leaf t =
   match t with
@@ -112,20 +117,34 @@ let select_leafs pos root =
   (try traverse root with Exit -> ());
   !branches
 
-let compare_locations pos l1 l2 =
+module Favorability = struct
+  type t =
+    | Neutral
+    | Unfavored
+
+  let based_on_ghostliness (loc : Location.t) =
+    match loc.loc_ghost with
+    | true -> Unfavored
+    | false -> Neutral
+end
+type node_loc = { loc : Location.t; favorability : Favorability.t }
+
+let compare_locations pos (l1 : node_loc) (l2 : node_loc) =
   let t2_first = +1 in
   let t1_first = -1 in
   match
-    Location_aux.compare_pos pos l1,
-    Location_aux.compare_pos pos l2
+    Location_aux.compare_pos pos l1.loc,
+    Location_aux.compare_pos pos l2.loc
   with
   | 0, 0 ->
-    (* Cursor inside both locations: favor non-ghost closer to the end *)
-    begin match l1.Location.loc_ghost, l2.Location.loc_ghost with
-    | true, false -> 1
-    | false, true -> -1
+    (* Cursor inside both locations:
+       If one is unfavored, favor the other one.
+       Otherwise, favor the one closer to the end *)
+    begin match l1.favorability, l2.favorability with
+    | Unfavored, Neutral -> 1
+    | Neutral, Unfavored -> -1
     | _ ->
-        Lexing.compare_pos l1.Location.loc_end l2.Location.loc_end
+        Lexing.compare_pos l1.loc.loc_end l2.loc.loc_end
     end
   (* Cursor inside one location: it has priority *)
   | 0, _ -> t1_first
@@ -135,25 +154,45 @@ let compare_locations pos l1 l2 =
   | n, m when m > 0 && n < 0 -> t2_first
   (* Cursor is after both, select the closest one *)
   | _, _ ->
-      Lexing.compare_pos l2.Location.loc_end l1.Location.loc_end
+      Lexing.compare_pos l2.loc.loc_end l1.loc.loc_end
+
+let compare_nodes ?(let_pun_behavior = Let_pun_behavior.Prefer_pattern) pos (n1, loc1) (n2, loc2) =
+  let loc_with_favorability node (loc : Location.t) : node_loc =
+    let is_punned = Browse_raw.has_attr ~name:Builtin_attributes.merlin_let_punned node in
+    let favorability : Favorability.t =
+      match is_punned, node, let_pun_behavior with
+      | true, Expression _, Prefer_expression -> Neutral
+      | true, Expression _, Prefer_pattern -> Unfavored
+      | true, Pattern _, Prefer_expression -> Unfavored
+      | true, Pattern _, Prefer_pattern -> Neutral
+      | _ -> Favorability.based_on_ghostliness loc
+    in
+    { loc = node_loc node; favorability }
+  in
+  compare_locations pos (loc_with_favorability n1 loc1) (loc_with_favorability n2 loc2)
 
-let best_node pos = function
+let best_node ?let_pun_behavior pos = function
   | [] -> []
   | init :: xs ->
     let f acc x =
-      if compare_locations pos (leaf_loc acc) (leaf_loc x) <= 0
+      let leaf_with_loc leaf =
+        let _, node = leaf_node leaf in
+        let loc = node_loc node in
+        node, loc
+      in
+      if compare_nodes ?let_pun_behavior pos (leaf_with_loc acc) (leaf_with_loc x) <= 0
       then acc
       else x
     in
     List.fold_left ~f ~init xs
 
-let enclosing pos roots =
-  match best_node pos roots with
+let enclosing ?let_pun_behavior pos roots =
+  match best_node ?let_pun_behavior pos roots with
   | [] -> []
-  | root -> best_node pos (select_leafs pos root)
+  | root -> best_node ?let_pun_behavior pos (select_leafs pos root)
 
-let deepest_before pos roots =
-  match enclosing pos roots with
+let deepest_before ?let_pun_behavior pos roots =
+  match enclosing ?let_pun_behavior pos roots with
   | [] -> []
   | root ->
     let rec aux path =
@@ -165,7 +204,7 @@ let deepest_before pos roots =
            Location_aux.compare_pos pos loc = 0 ||
            Lexing.compare_pos loc.Location.loc_end loc0.Location.loc_end = 0
         then match acc with
-          | Some (_,loc',_) when compare_locations pos loc' loc <= 0 -> acc
+          | Some (_,loc',node') when compare_nodes pos (node', loc') (node, loc) <= 0 -> acc
           | Some _ | None -> Some (env,loc,node)
         else acc
       in

src/kernel/mbrowse.mli

@@ -31,6 +31,12 @@ open Std
 type node = Browse_raw.node
 type t = (Env.t * node) list
 
+module Let_pun_behavior : sig
+  type t =
+    | Prefer_expression
+    | Prefer_pattern
+end
+
 val fold_node : (Env.t -> Browse_raw.node -> 'a -> 'a) ->
                  Env.t -> Browse_raw.node -> 'a -> 'a
 val node_loc : Browse_raw.node -> Location.t
@@ -43,13 +49,18 @@ val drop_leaf : t -> t option
  * through:
  *    foo bar (baz :: tail) <cursor>
  * asking for node from cursor position will return context of "tail".
- * Returns the matching node and all its ancestors or the empty list. *)
-val deepest_before : Lexing.position -> t list -> t
+ * Returns the matching node and all its ancestors or the empty list.
+
+ [let_pun_behavior] dictates whether to prefer the expression or pattern node in a
+    punned let expression. The default is [Prefer_pattern] *)
+val deepest_before : ?let_pun_behavior:Let_pun_behavior.t -> Lexing.position -> t list -> t
 
 
 val select_open_node : t -> (Path.t * Longident.t * t) option
 
-val enclosing : Lexing.position -> t list -> t
+(** [let_pun_behavior] dictates whether to prefer the expression or pattern node in a
+    punned let expression. The default is [Prefer_pattern] *)
+val enclosing : ?let_pun_behavior:Let_pun_behavior.t -> Lexing.position -> t list -> t
 
 val of_structure : Typedtree.structure -> t
 val of_signature : Typedtree.signature -> t

src/kernel/mtyper.ml

@@ -284,7 +284,7 @@ let get_index t = t.index
 
 let get_stamp t = t.stamp
 
-let node_at ?(skip_recovered=false) t pos_cursor =
+let node_at ?(skip_recovered=false) ?let_pun_behavior t pos_cursor =
   let node = Mbrowse.of_typedtree (get_typedtree t) in
   log ~title:"node_at" "Node: %s" (Mbrowse.print () node);
   let rec select = function
@@ -294,7 +294,7 @@ let node_at ?(skip_recovered=false) t pos_cursor =
       when Mbrowse.is_recovered node' -> select ancestors
     | l -> l
   in
-  match Mbrowse.deepest_before pos_cursor [node] with
+  match Mbrowse.deepest_before ?let_pun_behavior pos_cursor [node] with
   | [] -> [get_env t, Browse_raw.Dummy]
   | path when skip_recovered -> select path
   | path ->

src/kernel/mtyper.mli

@@ -59,4 +59,4 @@ val get_cache_stat : result -> typer_cache_stats
  *      inside x definition.
  *)
 val node_at :
-  ?skip_recovered:bool -> result -> Lexing.position -> Mbrowse.t
+  ?skip_recovered:bool -> ?let_pun_behavior:Mbrowse.Let_pun_behavior.t -> result -> Lexing.position -> Mbrowse.t

src/ocaml/parsing/builtin_attributes.ml

@@ -1023,3 +1023,7 @@ let get_tracing_probe_payload (payload : Parsetree.payload) =
     | _ -> Error ()
   in
   Ok { name; name_loc; enabled_at_init; arg }
+
+(* Merlin specific *)
+
+let merlin_let_punned = "merlin.let-punned"

src/ocaml/parsing/builtin_attributes.mli

@@ -326,3 +326,10 @@ type tracing_probe =
 *)
 val get_tracing_probe_payload :
   Parsetree.payload -> (tracing_probe, unit) result
+
+(* Merlin specific *)
+
+(** The name of the attribute used to identify punned let expressions. When a let
+    expression is punned, an attribute with this name is added to the pattern and
+    expression nodes by the parser. *)
+val merlin_let_punned : string