Fix locate for punned record pattern

src/analysis/context.ml

@@ -136,10 +136,14 @@ let inspect_expression ~cursor ~lid e : t =
   | Texp_constant _ -> Constant
   | _ -> Expr
 
-let inspect_browse_tree ?let_pun_behavior ~cursor lid browse : t option =
+let inspect_browse_tree ?let_pun_behavior ?record_pattern_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 ?let_pun_behavior cursor browse with
+  match
+    Mbrowse.enclosing ?let_pun_behavior ?record_pattern_pun_behavior cursor
+      browse
+  with
   | [] ->
     log ~title:"inspect_context" "no enclosing around: %a" Lexing.print_position
       cursor;

src/analysis/context.mli

@@ -62,6 +62,7 @@ val of_locate_context : Query_protocol.Locate_context.t -> t
 *)
 val inspect_browse_tree :
   ?let_pun_behavior:Mbrowse.Let_pun_behavior.t ->
+  ?record_pattern_pun_behavior:Mbrowse.Record_pattern_pun_behavior.t ->
   cursor:Std.Lexing.position ->
   Longident.t ->
   Mbrowse.t list ->

src/analysis/locate.ml

@@ -948,7 +948,8 @@ 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 ?let_pun_behavior ?namespaces ~pos lid browse is_label =
+let infer_namespace ?let_pun_behavior ?record_pattern_pun_behavior ?namespaces
+    ~pos lid browse is_label =
   match namespaces with
   | Some nss ->
     if not is_label then `Ok (nss :> Env_lookup.Namespace.inferred list)
@@ -961,7 +962,8 @@ let infer_namespace ?let_pun_behavior ?namespaces ~pos lid browse is_label =
       `Error `Missing_labels_namespace)
   | None -> (
     match
-      ( Context.inspect_browse_tree ?let_pun_behavior ~cursor:pos lid [ browse ],
+      ( Context.inspect_browse_tree ?let_pun_behavior
+          ?record_pattern_pun_behavior ~cursor:pos lid [ browse ],
         is_label )
     with
     | None, _ ->
@@ -976,7 +978,8 @@ let infer_namespace ?let_pun_behavior ?namespaces ~pos lid browse is_label =
       `Ok [ `Labels ])
 
 let from_string ~config ~env ~local_defs ~pos ?let_pun_behavior
-    ?(namespaces = Namespace_resolution.Inferred) path =
+    ?record_pattern_pun_behavior ?(namespaces = Namespace_resolution.Inferred)
+    path =
   File_switching.reset ();
   let browse = Mbrowse.of_typedtree local_defs in
   let lid = Type_utils.parse_longident path in
@@ -989,10 +992,11 @@ let from_string ~config ~env ~local_defs ~pos ?let_pun_behavior
         log ~title:"from_string" "overrode context: %s" (Context.to_string ctxt);
         `Ok (Env_lookup.Namespace.from_context ctxt)
       | Explicit namespaces ->
-        infer_namespace ?let_pun_behavior ~namespaces ~pos lid browse is_label
+        infer_namespace ?let_pun_behavior ?record_pattern_pun_behavior
+          ~namespaces ~pos lid browse is_label
       | Inferred ->
-        infer_namespace ?let_pun_behavior ?namespaces:None ~pos lid browse
-          is_label
+        infer_namespace ?let_pun_behavior ?record_pattern_pun_behavior
+          ?namespaces:None ~pos lid browse is_label
     in
     match namespaces with
     | `Error e -> e

src/analysis/locate.mli

@@ -84,6 +84,7 @@ val from_string :
   local_defs:Mtyper.typedtree ->
   pos:Lexing.position ->
   ?let_pun_behavior:Mbrowse.Let_pun_behavior.t ->
+  ?record_pattern_pun_behavior:Mbrowse.Record_pattern_pun_behavior.t ->
   ?namespaces:Namespace_resolution.t ->
   string ->
   [> `File_not_found of string

src/frontend/query_commands.ml

@@ -618,6 +618,9 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a = function
     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 record_pattern_pun_behavior =
+      Mbrowse.Record_pattern_pun_behavior.Prefer_label
+    in
     let env, _ = Mbrowse.leaf_node (Mtyper.node_at typer pos) in
     let path =
       match patho with
@@ -646,7 +649,7 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a = function
         in
         match
           Locate.from_string ~config ~env ~local_defs ~pos ?namespaces
-            ~let_pun_behavior path
+            ~let_pun_behavior ~record_pattern_pun_behavior path
         with
         | `Found { file; location; _ } ->
           Locate.log ~title:"result" "found: %s" file;

src/kernel/mbrowse.ml

@@ -35,6 +35,14 @@ type t = (Env.t * node) list
 
 module Let_pun_behavior = struct
   type t = Prefer_expression | Prefer_pattern
+
+  let default = Prefer_pattern
+end
+
+module Record_pattern_pun_behavior = struct
+  type t = Prefer_label | Prefer_pattern
+
+  let default = Prefer_pattern
 end
 
 let node_of_binary_part = Browse_raw.node_of_binary_part
@@ -106,60 +114,109 @@ let select_leafs pos root =
   !branches
 
 module Favorability = struct
-  type t = Neutral | Unfavored
+  type t = Unfavored | Neutral | Favored
 
   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.loc, Location_aux.compare_pos pos l2.loc)
-  with
-  (* Cursor inside both locations:
-       If one is unfavored, favor the other one.
-       Otherwise, favor the one closer to the end *)
-  | 0, 0 -> begin
-    match (l1.favorability, l2.favorability) with
-    | Unfavored, Neutral -> 1
-    | Neutral, Unfavored -> -1
-    | _ -> Lexing.compare_pos l1.loc.loc_end l2.loc.loc_end
-  end
-  (* Cursor inside one location: it has priority *)
-  | 0, _ -> t1_first
-  | _, 0 -> t2_first
-  (* Cursor outside locations: favor before *)
-  | n, m when n > 0 && m < 0 -> t1_first
-  | n, m when m > 0 && n < 0 -> t2_first
-  (* Cursor is after both, select the closest one *)
-  | _, _ -> 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 based_on_let_punning ~(let_pun_behavior : Let_pun_behavior.t) node =
     let is_punned =
-      Browse_raw.has_attr ~name:Builtin_attributes.merlin_let_punned node
+      Browse_raw.has_attr ~name:Builtin_attributes.merlin_punned_let 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
+    match (is_punned, node, let_pun_behavior) with
+    | true, Expression _, Prefer_expression -> Favored
+    | true, Expression _, Prefer_pattern -> Unfavored
+    | true, Pattern _, Prefer_expression -> Unfavored
+    | true, Pattern _, Prefer_pattern -> Favored
+    | _ -> Neutral
+
+  let based_on_record_pattern_punning
+      ~(record_pattern_pun_behavior : Record_pattern_pun_behavior.t) node =
+    let is_punned =
+      Browse_raw.has_attr ~name:Builtin_attributes.merlin_punned_record_pattern
+        node
     in
-    { loc = node_loc node; favorability }
-  in
-  compare_locations pos
-    (loc_with_favorability n1 loc1)
-    (loc_with_favorability n2 loc2)
+    match (is_punned, node, record_pattern_pun_behavior) with
+    | true, Pattern _, Prefer_label -> Unfavored
+    | true, Pattern _, Prefer_pattern -> Favored
+    | _ -> Neutral
+end
 
-let best_node ?let_pun_behavior pos = function
+module Node_comparison_result = struct
+  type t = Left | Neutral | Right
+
+  let left_or_neutral = function
+    | Left | Neutral -> true
+    | Right -> false
+
+  let invert = function
+    | Left -> Right
+    | Right -> Left
+    | Neutral -> Neutral
+
+  let from_favorabilities (left : Favorability.t) (right : Favorability.t) =
+    match (left, right) with
+    | Unfavored, (Neutral | Favored) | Neutral, Favored -> Right
+    | (Neutral | Favored), Unfavored | Favored, Neutral -> Left
+    | Unfavored, Unfavored | Neutral, Neutral | Favored, Favored -> Neutral
+
+  let rec combine_lazy_list = function
+    | [] -> Neutral
+    | (lazy Left) :: _ -> Left
+    | (lazy Right) :: _ -> Right
+    | (lazy Neutral) :: rest -> combine_lazy_list rest
+end
+
+let closer_to_end_of_file (l1 : Location.t) (l2 : Location.t) :
+    Node_comparison_result.t =
+  match Lexing.compare_pos l1.loc_end l2.loc_end with
+  | r when r < 0 -> Left
+  | r when r > 0 -> Right
+  | _ -> Neutral
+
+let compare_locations pos l1 l2 : Node_comparison_result.t =
+  match (Location_aux.compare_pos pos l1, Location_aux.compare_pos pos l2) with
+  (* Cursor inside both locations *)
+  | 0, 0 -> Neutral
+  (* Cursor inside one location: it has priority *)
+  | 0, _ -> Left
+  | _, 0 -> Right
+  (* Cursor between the two locations: favor the one before the cursor *)
+  | n, m when n > 0 && m < 0 -> Left
+  | n, m when m > 0 && n < 0 -> Right
+  (* Cursor is after both, select the one whose end is earlier in the file *)
+  | _, _ -> closer_to_end_of_file l1 l2 |> Node_comparison_result.invert
+
+let compare_nodes ~let_pun_behavior ~record_pattern_pun_behavior pos
+    (node1, loc1) (node2, loc2) =
+  (* Prioritization order:
+     1. Choose the node whose location encompasses [pos]
+     2. If node is part of a let-punned pattern, choose the one on the preferred side
+        (based on [let_pun_behavior])
+     3. Choose the one that isn't marked as ghost
+     4. Choose the node that is closer to the end of the file *)
+  [ lazy (compare_locations pos loc1 loc2);
+    lazy
+      (Node_comparison_result.from_favorabilities
+         (Favorability.based_on_let_punning ~let_pun_behavior node1)
+         (Favorability.based_on_let_punning ~let_pun_behavior node2));
+    lazy
+      (Node_comparison_result.from_favorabilities
+         (Favorability.based_on_record_pattern_punning
+            ~record_pattern_pun_behavior node1)
+         (Favorability.based_on_record_pattern_punning
+            ~record_pattern_pun_behavior node2));
+    lazy
+      (Node_comparison_result.from_favorabilities
+         (Favorability.based_on_ghostliness loc1)
+         (Favorability.based_on_ghostliness loc2));
+    lazy (closer_to_end_of_file loc1 loc2)
+  ]
+  |> Node_comparison_result.combine_lazy_list
+
+let best_node ~let_pun_behavior ~record_pattern_pun_behavior pos = function
   | [] -> []
   | init :: xs ->
     let f acc x =
@@ -168,22 +225,28 @@ let best_node ?let_pun_behavior pos = function
         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
+      match
+        compare_nodes ~let_pun_behavior ~record_pattern_pun_behavior pos
+          (leaf_with_loc acc) (leaf_with_loc x)
+      with
+      | Left | Neutral -> acc
+      | Right -> x
     in
     List.fold_left ~f ~init xs
 
-let enclosing ?let_pun_behavior pos roots =
-  match best_node ?let_pun_behavior pos roots with
+let enclosing ?(let_pun_behavior = Let_pun_behavior.default)
+    ?(record_pattern_pun_behavior = Record_pattern_pun_behavior.default) pos
+    roots =
+  match best_node ~let_pun_behavior ~record_pattern_pun_behavior pos roots with
   | [] -> []
-  | root -> best_node ?let_pun_behavior pos (select_leafs pos root)
+  | root ->
+    best_node ~let_pun_behavior ~record_pattern_pun_behavior pos
+      (select_leafs pos root)
 
-let deepest_before ?let_pun_behavior pos roots =
-  match enclosing ?let_pun_behavior pos roots with
+let deepest_before ?(let_pun_behavior = Let_pun_behavior.default)
+    ?(record_pattern_pun_behavior = Record_pattern_pun_behavior.default) pos
+    roots =
+  match enclosing ~let_pun_behavior ~record_pattern_pun_behavior pos roots with
   | [] -> []
   | root ->
     let rec aux path =
@@ -198,7 +261,9 @@ let deepest_before ?let_pun_behavior pos roots =
         then
           match acc with
           | Some (_, loc', node')
-            when compare_nodes pos (node', loc') (node, loc) <= 0 -> acc
+            when compare_nodes ~let_pun_behavior ~record_pattern_pun_behavior
+                   pos (node', loc') (node, loc)
+                 |> Node_comparison_result.left_or_neutral -> acc
           | Some _ | None -> Some (env, loc, node)
         else acc
       in

src/kernel/mbrowse.mli

@@ -35,6 +35,10 @@ module Let_pun_behavior : sig
   type t = Prefer_expression | Prefer_pattern
 end
 
+module Record_pattern_pun_behavior : sig
+  type t = Prefer_label | 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
@@ -52,14 +56,22 @@ val drop_leaf : t -> t option
  [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
+  ?let_pun_behavior:Let_pun_behavior.t ->
+  ?record_pattern_pun_behavior:Record_pattern_pun_behavior.t ->
+  Lexing.position ->
+  t list ->
+  t
 
 val select_open_node : t -> (Path.t * Longident.t * t) option
 
 (** [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
+  ?let_pun_behavior:Let_pun_behavior.t ->
+  ?record_pattern_pun_behavior:Record_pattern_pun_behavior.t ->
+  Lexing.position ->
+  t list ->
+  t
 
 val of_structure : Typedtree.structure -> t
 val of_signature : Typedtree.signature -> t

src/kernel/mtyper.ml

@@ -366,7 +366,8 @@ let get_index t = t.index
 
 let get_stamp t = t.stamp
 
-let node_at ?(skip_recovered = false) ?let_pun_behavior t pos_cursor =
+let node_at ?(skip_recovered = false) ?let_pun_behavior
+    ?record_pattern_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
@@ -376,7 +377,10 @@ let node_at ?(skip_recovered = false) ?let_pun_behavior t pos_cursor =
       -> select ancestors
     | l -> l
   in
-  match Mbrowse.deepest_before ?let_pun_behavior pos_cursor [ node ] with
+  match
+    Mbrowse.deepest_before ?let_pun_behavior ?record_pattern_pun_behavior
+      pos_cursor [ node ]
+  with
   | [] -> [ (get_env t, Browse_raw.Dummy) ]
   | path when skip_recovered -> select path
   | path ->

src/kernel/mtyper.mli

@@ -59,6 +59,7 @@ val get_cache_stat : result -> typer_cache_stats
 val node_at :
   ?skip_recovered:bool ->
   ?let_pun_behavior:Mbrowse.Let_pun_behavior.t ->
+  ?record_pattern_pun_behavior:Mbrowse.Record_pattern_pun_behavior.t ->
   result ->
   Lexing.position ->
   Mbrowse.t

src/ocaml/parsing/builtin_attributes.ml

@@ -1136,4 +1136,6 @@ let get_tracing_probe_payload (payload : Parsetree.payload) =
 
 (* Merlin specific *)
 
-let merlin_let_punned = "merlin.let-punned"
+let merlin_punned_let = "merlin.punned-let"
+
+let merlin_punned_record_pattern = "merlin.punned-record-pattern"