Override context when locating (#91)

* Add context override to locate

* Create tests

* Add more context and tests

* Remove query_protocol.mli

* Cleanup signature of Locate.from_string

* Stop using array literals in test

* Declare arrays

* Rewrite test without using arrays

* Fix if syntax

* Change == to =

* Format files

Author: liam923

src/analysis/context.ml

@@ -35,8 +35,10 @@ type t =
     (* We attach the constructor description here so in the case of
        disambiguated constructors we actually directly look for the type
        path (cf. #486, #794). *)
+  | Unknown_constructor
   | Expr
   | Label of Types.label_description (* Similar to constructors. *)
+  | Unknown_label
   | Module_path
   | Module_type
   | Patt
@@ -46,15 +48,28 @@ type t =
 
 let to_string = function
   | Constructor (cd, _) -> Printf.sprintf "constructor %s" cd.cstr_name
+  | Unknown_constructor -> Printf.sprintf "unknown constructor"
   | Expr -> "expression"
   | Label lbl -> Printf.sprintf "record field %s" lbl.lbl_name
+  | Unknown_label -> Printf.sprintf "record field"
   | Module_path -> "module path"
   | Module_type -> "module type"
   | Patt -> "pattern"
   | Constant -> "constant"
   | Type -> "type"
   | Unknown -> "unknown"
 
+let of_locate_context : Query_protocol.Locate_context.t -> t = function
+  | Expr -> Expr
+  | Module_path -> Module_path
+  | Module_type -> Module_type
+  | Patt -> Patt
+  | Type -> Type
+  | Constant -> Constant
+  | Constructor -> Unknown_constructor
+  | Label -> Unknown_label
+  | Unknown -> Unknown
+
 (* Distinguish between "Mo[d]ule.something" and "Module.some[t]hing" *)
 let cursor_on_longident_end ~cursor:cursor_pos
     ~lid_loc:{ Asttypes.loc; txt = lid } name =

src/analysis/context.mli

@@ -31,8 +31,10 @@ type t =
     (* We attach the constructor description here so in the case of
        disambiguated constructors we actually directly look for the type
        path (cf. #486, #794). *)
+  | Unknown_constructor
   | Expr
   | Label of Types.label_description (* Similar to constructors. *)
+  | Unknown_label
   | Module_path
   | Module_type
   | Patt
@@ -42,6 +44,8 @@ type t =
 
 val to_string : t -> string
 
+val of_locate_context : Query_protocol.Locate_context.t -> t
+
 (**
   [inspect_browse_tree lid ~cursor mbrowse] tries to provide contextual
   information given the selected identifier, the position of the cursor and the

src/analysis/env_lookup.ml

@@ -21,6 +21,8 @@ module Namespace = struct
     | Type -> [ `Type; `Mod; `Modtype; `Constr; `Labels; `Vals ]
     | Module_type -> [ `Modtype; `Mod; `Type; `Constr; `Labels; `Vals ]
     | Expr | Constant -> [ `Vals; `Mod; `Modtype; `Constr; `Labels; `Type ]
+    | Unknown_constructor -> [ `Constr; `Vals; `Mod; `Modtype; `Labels; `Type ]
+    | Unknown_label -> [ `Labels; `Vals; `Mod; `Modtype; `Constr; `Type ]
     | Patt -> [ `Mod; `Modtype; `Type; `Constr; `Labels; `Vals ]
     | Unknown -> [ `Vals; `Type; `Constr; `Mod; `Modtype; `Labels ]
     | Label lbl -> [ `This_label lbl ]

src/analysis/locate.ml

@@ -43,6 +43,13 @@ type result =
     approximated : bool
   }
 
+module Namespace_resolution = struct
+  type t =
+    | From_context of Query_protocol.Locate_context.t
+    | Explicit of Env_lookup.Namespace.inferred_basic list
+    | Inferred
+end
+
 module File : sig
   type t = private
     | ML of string
@@ -897,16 +904,26 @@ let infer_namespace ?let_pun_behavior ?namespaces ~pos lid browse is_label =
         "dropping inferred context, it is not precise enough";
       `Ok [ `Labels ])
 
-let from_string ~config ~env ~local_defs ~pos ?let_pun_behavior ?namespaces path
-    =
+let from_string ~config ~env ~local_defs ~pos ?let_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
   let from_lid lid =
     let ident, is_label = Longident.keep_suffix lid in
-    match
-      infer_namespace ?let_pun_behavior ?namespaces ~pos lid browse is_label
-    with
+    let namespaces =
+      match namespaces with
+      | From_context ctxt ->
+        let ctxt = Context.of_locate_context ctxt in
+        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
+      | Inferred ->
+        infer_namespace ?let_pun_behavior ?namespaces:None ~pos lid browse
+          is_label
+    in
+    match namespaces with
     | `Error e -> e
     | `Ok nss ->
       log ~title:"from_string"

src/analysis/locate.mli

@@ -47,6 +47,16 @@ type result =
     approximated : bool
   }
 
+module Namespace_resolution : sig
+  type t =
+    | From_context of Query_protocol.Locate_context.t
+        (** Choose the namespaces based on a [Query_protocol.Locate_context.t] *)
+    | Explicit of Env_lookup.Namespace.inferred_basic list
+        (** Explicitly choose which namespaces to search in. The namespaces are prioritized
+        based on the list order (with the first element being highest priority) *)
+    | Inferred  (** Infer which namespaces to search in *)
+end
+
 val uid_of_result :
   traverse_aliases:bool -> Shape_reduce.result -> Shape.Uid.t option * bool
 
@@ -74,7 +84,7 @@ val from_string :
   local_defs:Mtyper.typedtree ->
   pos:Lexing.position ->
   ?let_pun_behavior:Mbrowse.Let_pun_behavior.t ->
-  ?namespaces:Env_lookup.Namespace.inferred_basic list ->
+  ?namespaces:Namespace_resolution.t ->
   string ->
   [> `File_not_found of string
   | `Found of result

src/commands/new_commands.ml

@@ -455,21 +455,34 @@ let all_commands =
     command "locate"
       ~spec:
         [ optional "-prefix" "<string> Prefix to complete"
-            (Marg.param "string" (fun txt (_, pos, kind) ->
-                 (Some txt, pos, kind)));
+            (Marg.param "string" (fun txt (_, pos, kind, ctx) ->
+                 (Some txt, pos, kind, ctx)));
           arg "-position" "<position> Position to complete"
-            (marg_position (fun pos (prefix, _pos, kind) -> (prefix, pos, kind)));
+            (marg_position (fun pos (prefix, _pos, kind, ctx) ->
+                 (prefix, pos, kind, ctx)));
           optional "-look-for"
             "<interface|implementation> Prefer opening interface or \
              implementation"
             (Marg.param "<interface|implementation>"
-               (fun kind (prefix, pos, _) ->
+               (fun kind (prefix, pos, _, ctx) ->
                  match kind with
-                 | "mli" | "interface" -> (prefix, pos, `MLI)
-                 | "ml" | "implementation" -> (prefix, pos, `ML)
+                 | "mli" | "interface" -> (prefix, pos, `MLI, ctx)
+                 | "ml" | "implementation" -> (prefix, pos, `ML, ctx)
                  | str ->
                    failwithf "expecting interface or implementation, got %S."
-                     str))
+                     str));
+          (let contexts =
+             let open Query_protocol.Locate_context in
+             all |> List.map ~f:to_string |> String.concat ~sep:"|"
+           in
+           optional "-context"
+             (Format.sprintf
+                "<%s> Which context to search for the identifier in" contexts)
+             (Marg.param (Format.sprintf "<%s>" contexts)
+                (fun ctx (prefix, pos, kind, _) ->
+                  match Query_protocol.Locate_context.of_string ctx with
+                  | Some ctx -> (prefix, pos, kind, Some ctx)
+                  | None -> failwithf "invalid context %s." ctx)))
         ]
       ~doc:
         "Finds the declaration of entity at the specified position, Or \
@@ -479,13 +492,13 @@ let all_commands =
          - `{'pos': position}` if the location is in the current buffer,\n\
          - `{'file': string, 'pos': position}` if definition is located in a \
          different file."
-      ~default:(None, `None, `MLI)
+      ~default:(None, `None, `MLI, None)
       begin
-        fun buffer (prefix, pos, lookfor) ->
+        fun buffer (prefix, pos, lookfor, context) ->
           match pos with
           | `None -> failwith "-position <pos> is mandatory"
           | #Msource.position as pos ->
-            run buffer (Query_protocol.Locate (prefix, lookfor, pos))
+            run buffer (Query_protocol.Locate (prefix, lookfor, pos, context))
       end;
     command "locate-type"
       ~spec:

src/commands/query_json.ml

@@ -105,7 +105,7 @@ let dump (type a) : a t -> json =
   | Syntax_document pos ->
     mk "syntax-document" [ ("position", mk_position pos) ]
   | Expand_ppx pos -> mk "ppx-expand" [ ("position", mk_position pos) ]
-  | Locate (prefix, look_for, pos) ->
+  | Locate (prefix, look_for, pos, context) ->
     mk "locate"
       [ ( "prefix",
           match prefix with
@@ -115,7 +115,12 @@ let dump (type a) : a t -> json =
           match look_for with
           | `ML -> `String "implementation"
           | `MLI -> `String "interface" );
-        ("position", mk_position pos)
+        ("position", mk_position pos);
+        ( "context",
+          match context with
+          | Some context ->
+            `String (Query_protocol.Locate_context.to_string context)
+          | None -> `Null )
       ]
   | Jump (target, pos) ->
     mk "jump" [ ("target", `String target); ("position", mk_position pos) ]

src/frontend/ocamlmerlin/old/old_IO.ml

@@ -170,10 +170,12 @@ let request_of_json context =
   | `String "document" :: `String path :: pos ->
     request (Query (Document (Some path, mandatory_position pos)))
   | `String "locate" :: (`String "" | `Null) :: `String choice :: pos ->
-    request (Query (Locate (None, ml_or_mli choice, mandatory_position pos)))
+    request
+      (Query (Locate (None, ml_or_mli choice, mandatory_position pos, None)))
   | `String "locate" :: `String path :: `String choice :: pos ->
     request
-      (Query (Locate (Some path, ml_or_mli choice, mandatory_position pos)))
+      (Query
+         (Locate (Some path, ml_or_mli choice, mandatory_position pos, None)))
   | `String "jump" :: `String target :: pos ->
     request (Query (Jump (target, mandatory_position pos)))
   | [ `String "outline" ] -> request (Query Outline)

src/frontend/query_commands.ml

@@ -614,7 +614,7 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a = function
         (Ppx_expand.get_ppxed_source ~ppxed_parsetree ~pos
            (Option.get ppx_kind_with_attr))
     | None -> `No_ppx)
-  | Locate (patho, ml_or_mli, pos) ->
+  | Locate (patho, ml_or_mli, pos, context) ->
     let typer = Mpipeline.typer_result pipeline in
     let local_defs = Mtyper.get_typedtree typer in
     let pos = Mpipeline.get_lexing_pos pipeline pos in
@@ -641,9 +641,13 @@ let dispatch pipeline (type a) : a Query_protocol.t -> a = function
           }
       in
       begin
+        let namespaces =
+          Option.map context ~f:(fun ctx ->
+              Locate.Namespace_resolution.From_context ctx)
+        in
         match
-          Locate.from_string ~config ~env ~local_defs ~pos ~let_pun_behavior
-            path
+          Locate.from_string ~config ~env ~local_defs ~pos ?namespaces
+            ~let_pun_behavior path
         with
         | `Found { file; location; _ } ->
           Locate.log ~title:"result" "found: %s" file;

src/frontend/query_protocol.ml

@@ -129,6 +129,54 @@ type _ _bool = bool
 type occurrences_status =
   [ `Not_requested | `Out_of_sync of string list | `No_def | `Included ]
 
+module Locate_context = struct
+  type t =
+    | Expr
+    | Module_path
+    | Module_type
+    | Patt
+    | Type
+    | Constant
+    | Constructor
+    | Label
+    | Unknown
+
+  let to_string = function
+    | Expr -> "expr"
+    | Module_path -> "module_path"
+    | Module_type -> "module_type"
+    | Patt -> "pattern"
+    | Type -> "type"
+    | Constant -> "constant"
+    | Constructor -> "constructor"
+    | Label -> "label"
+    | Unknown -> "unknown"
+
+  let of_string = function
+    | "expr" -> Some Expr
+    | "module_path" -> Some Module_path
+    | "module_type" -> Some Module_type
+    | "pattern" -> Some Patt
+    | "type" -> Some Type
+    | "constant" -> Some Constant
+    | "constructor" -> Some Constructor
+    | "label" -> Some Label
+    | "unknown" -> Some Unknown
+    | _ -> None
+
+  let all =
+    [ Expr;
+      Module_path;
+      Module_type;
+      Patt;
+      Type;
+      Constant;
+      Constructor;
+      Label;
+      Unknown
+    ]
+end
+
 type _ t =
   | Type_expr (* *) : string * Msource.position -> string t
   | Stack_or_heap_enclosing (* *) :
@@ -182,7 +230,10 @@ type _ t =
          | `At_origin ]
          t
   | Locate (* *) :
-      string option * [ `ML | `MLI ] * Msource.position
+      string option
+      * [ `ML | `MLI ]
+      * Msource.position
+      * Locate_context.t option
       -> [ `Found of string option * Lexing.position
          | `Invalid_context
          | `Builtin of string