Complete the constructors of extensible variants

analysis/src/completion_back_end.ml

@@ -1783,7 +1783,22 @@ let rec complete_typed_value ?(type_arg_context : type_arg_context option)
       ]
   | Tvariant {env; constructors; variant_decl; variant_name} ->
     if Debug.verbose () then print_endline "[complete_typed_value]--> Tvariant";
+    (* An extension constructor is named by its module path, as in
+       [Sub.Flag]; the prefix may also be the start of its last part. *)
+    let matches_prefix name =
+      prefix = ""
+      || Utils.starts_with name prefix
+      ||
+      match String.rindex_opt name '.' with
+      | Some i ->
+        Utils.starts_with
+          (String.sub name (i + 1) (String.length name - i - 1))
+          prefix
+      | None -> false
+    in
     constructors
+    |> List.filter (fun (constructor : Constructor.t) ->
+        matches_prefix constructor.cname.txt)
     |> List.map (fun (constructor : Constructor.t) ->
         let num_args =
           match constructor.args with
@@ -1800,7 +1815,6 @@ let rec complete_typed_value ?(type_arg_context : type_arg_context option)
             (Constructor
                (constructor, variant_decl |> Shared.decl_to_string variant_name))
           ~env)
-    |> filter_items ~prefix
   | Tpolyvariant {env; constructors; type_expr} ->
     if Debug.verbose () then
       print_endline "[complete_typed_value]--> Tpolyvariant";
@@ -2091,8 +2105,8 @@ let rec complete_typed_value ?(type_arg_context : type_arg_context option)
 
 module String_set = Set.Make (String)
 
-let rec process_completable ~state ~debug ~full ~scope ~env ~pos ~for_hover
-    completable =
+let rec process_completable_in_scope ~state ~debug ~full ~scope ~env ~pos
+    ~for_hover completable =
   if debug then
     Printf.printf "Completable: %s\n" (Completable.to_string completable);
   let package = full.package in
@@ -2398,7 +2412,8 @@ let rec process_completable ~state ~debug ~full ~scope ~env ~pos ~for_hover
       match (fallback, items) with
       | Some fallback, (None | Some []) ->
         fallback
-        |> process_completable ~state ~debug ~full ~scope ~env ~pos ~for_hover
+        |> process_completable_in_scope ~state ~debug ~full ~scope ~env ~pos
+             ~for_hover
       | _, Some items -> items
       | None, None -> []
     in
@@ -2569,15 +2584,23 @@ let rec process_completable ~state ~debug ~full ~scope ~env ~pos ~for_hover
             typ_expr |> Type_utils.extract_type ~env:c.env ~package ~state
           with
           | Some (Tvariant v, _) ->
-            with_exhaustive_item c
-              ~cases:
-                (v.constructors
-                |> List.map (fun (constructor : Constructor.t) ->
-                    constructor.cname.txt
-                    ^
-                    match constructor.args with
-                    | Args [] -> ""
-                    | _ -> "(_)"))
+            let cases =
+              v.constructors
+              |> List.map (fun (constructor : Constructor.t) ->
+                  constructor.cname.txt
+                  ^
+                  match constructor.args with
+                  | Args [] -> ""
+                  | _ -> "(_)")
+            in
+            (* An open type can be extended elsewhere, so only a catch-all
+               case makes the switch exhaustive. *)
+            let cases =
+              match v.variant_decl.type_kind with
+              | Type_open -> cases @ ["_"]
+              | _ -> cases
+            in
+            with_exhaustive_item c ~cases
           | Some (Tpolyvariant v, _) ->
             with_exhaustive_item c
               ~cases:
@@ -2625,3 +2648,32 @@ let rec process_completable ~state ~debug ~full ~scope ~env ~pos ~for_hover
           ~env ~insert_text:"todo(\"${0:TODO}\")";
       ]
     else []
+
+(* The names of the modules around [pos] in [file], outermost first. *)
+let module_path_at ~(full : full) ~pos =
+  let rec in_items (items : Module.item list) =
+    match
+      items
+      |> List.find_opt (fun (item : Module.item) ->
+          Loc.has_pos ~pos item.loc
+          &&
+          match item.kind with
+          | Module _ -> true
+          | _ -> false)
+    with
+    | Some {name; kind = Module {type_}} -> name :: in_module type_
+    | _ -> []
+  and in_module (m : Module.t) =
+    match m with
+    | Structure {items} -> in_items items
+    | Constraint (implementation, _) -> in_module implementation
+    | Ident _ -> []
+  in
+  in_items full.file.structure.items
+
+let process_completable ~state ~debug ~full ~scope ~env ~pos ~for_hover
+    completable =
+  Type_utils.with_completion_scope ~file:full.file
+    ~module_path:(module_path_at ~full ~pos) ~pos ~scope (fun () ->
+      process_completable_in_scope ~state ~debug ~full ~scope ~env ~pos
+        ~for_hover completable)

analysis/src/completion_expressions.ml

@@ -157,6 +157,7 @@ let rec traverse_expr (exp : Parsetree.expression) ~expr_path ~pos
           Completable.NVariantPayload
             {
               constructor_name = Longident.last txt;
+              constructor_path = Longident.flatten txt;
               item_num = 0;
               source_arity = 1;
             };
@@ -170,6 +171,7 @@ let rec traverse_expr (exp : Parsetree.expression) ~expr_path ~pos
              Completable.NVariantPayload
                {
                  constructor_name = Longident.last txt;
+                 constructor_path = Longident.flatten txt;
                  item_num;
                  source_arity = List.length args;
                };
@@ -180,6 +182,7 @@ let rec traverse_expr (exp : Parsetree.expression) ~expr_path ~pos
              Completable.NVariantPayload
                {
                  constructor_name = Longident.last txt;
+                 constructor_path = Longident.flatten txt;
                  item_num = item_num + 1;
                  source_arity = List.length args;
                };

analysis/src/completion_front_end.ml

@@ -504,6 +504,7 @@ let completion_with_parser1 ~debug ~offset ~pos_cursor ~kind_file
                    item_num = index;
                    source_arity = List.length patterns;
                    constructor_name = Longident.last txt;
+                   constructor_path = Longident.flatten txt;
                  }
               :: pattern_path)
             ?context_path)
@@ -744,6 +745,12 @@ let completion_with_parser1 ~debug ~offset ~pos_cursor ~kind_file
     | Pstr_open {popen_lid} ->
       scope := !scope |> Scope.add_open ~lid:popen_lid.txt
     | Pstr_primitive vd -> scope_value_description vd
+    | Pstr_typext {ptyext_constructors} ->
+      ptyext_constructors
+      |> List.iter (fun (ext : Parsetree.extension_constructor) ->
+          scope :=
+            !scope
+            |> Scope.add_constructor ~name:ext.pext_name.txt ~loc:ext.pext_loc)
     | Pstr_value (rec_flag, bindings) ->
       if rec_flag = Recursive then bindings |> List.iter scope_value_binding;
       bindings |> List.iter (fun vb -> iterator.value_binding iterator vb);

analysis/src/completion_patterns.ml

@@ -212,6 +212,7 @@ and traverse_pattern (pat : Parsetree.pattern) ~pattern_path ~loc_has_cursor
           Completable.NVariantPayload
             {
               constructor_name = Longident.last txt;
+              constructor_path = Longident.flatten txt;
               item_num = 0;
               source_arity = 1;
             };
@@ -226,6 +227,7 @@ and traverse_pattern (pat : Parsetree.pattern) ~pattern_path ~loc_has_cursor
              Completable.NVariantPayload
                {
                  constructor_name = Longident.last txt;
+                 constructor_path = Longident.flatten txt;
                  item_num;
                  source_arity = List.length patterns;
                };
@@ -236,6 +238,7 @@ and traverse_pattern (pat : Parsetree.pattern) ~pattern_path ~loc_has_cursor
              Completable.NVariantPayload
                {
                  constructor_name = Longident.last txt;
+                 constructor_path = Longident.flatten txt;
                  item_num = item_num + 1;
                  source_arity = List.length patterns;
                };

analysis/src/process_cmt.ml

@@ -35,6 +35,54 @@ let map_record_field {Types.ld_id; ld_type; ld_attributes; ld_optional} =
     deprecated = Process_attributes.find_deprecated_attribute ld_attributes;
   }
 
+(* The constructors of [type t += ...] are registered as constructors of an
+   open type named [t], so that completion can offer them for values of
+   type [t]. They are not part of any type declaration, hence the stand-in
+   declaration that carries the kind and, as its manifest, the extended
+   type, which completion compares with the type it completes. *)
+let add_extension_constructor ~(env : Env.t) ~id ~name
+    (ext : Types.extension_constructor) =
+  let type_decl : Types.type_declaration =
+    {
+      type_params = [];
+      type_arity = 0;
+      type_kind = Type_open;
+      type_private = Public;
+      type_manifest = Some (Ctype.newconstr ext.ext_type_path []);
+      type_variance = [];
+      type_newtype_level = None;
+      type_loc = Location.none;
+      type_attributes = [];
+      type_immediate = false;
+      type_representation = Boxed;
+      type_shape = Unstated;
+      type_inlined_types = [];
+    }
+  in
+  let stamp = Ident.binding_time id in
+  let item =
+    {
+      Constructor.stamp;
+      cname = name;
+      args =
+        (match ext.ext_args with
+        | Cstr_tuple args ->
+          Args (args |> List.map (fun t -> (t, Location.none)))
+        | Cstr_record fields ->
+          InlineRecord (fields |> List.map map_record_field));
+      res = ext.ext_ret_type;
+      type_decl = (Path.last ext.ext_type_path, type_decl);
+      docstring = attrs_to_docstring ext.ext_attributes;
+      deprecated =
+        Process_attributes.find_deprecated_attribute ext.ext_attributes;
+    }
+  in
+  let declared =
+    Process_attributes.new_declared ~item ~extent:ext.ext_loc ~name ~stamp
+      ~module_path:env.module_path true ext.ext_attributes
+  in
+  Stamps.add_constructor env.stamps stamp declared
+
 let rec for_type_signature_item ~(env : Shared_types.Env.t)
     ~(exported : Exported.t) (item : Types.signature_item) =
   match item with
@@ -148,7 +196,7 @@ let rec for_type_signature_item ~(env : Shared_types.Env.t)
     let name = Ident.name ident in
     let declared =
       add_declared ~extent:md_loc
-        ~item:(for_type_module ~name ~env md_type)
+        ~item:(for_type_module ~name ~env:(env |> Env.add_module ~name) md_type)
         ~name:(Location.mkloc name md_loc)
         ~stamp:(Ident.binding_time ident) ~env md_attributes
         (Exported.add exported Exported.Module)
@@ -165,6 +213,11 @@ let rec for_type_signature_item ~(env : Shared_types.Env.t)
         loc = declared.extent_loc;
       };
     ]
+  | Sig_typext (id, ext, _) ->
+    add_extension_constructor ~env ~id
+      ~name:(Location.mkloc (Ident.name id) ext.ext_loc)
+      ext;
+    []
   | _ -> []
 
 and for_type_signature ~name ~env signature =
@@ -328,6 +381,12 @@ let for_type_declaration ~env ~(exported : Exported.t)
     loc = declared.extent_loc;
   }
 
+let add_extension_constructors ~(env : Env.t) {Typedtree.tyext_constructors} =
+  tyext_constructors
+  |> List.iter
+       (fun ({ext_id; ext_name; ext_type} : Typedtree.extension_constructor) ->
+         add_extension_constructor ~env ~id:ext_id ~name:ext_name ext_type)
+
 let rec for_signature_item ~env ~(exported : Exported.t)
     (item : Typedtree.signature_item) =
   match item.sig_desc with
@@ -401,6 +460,9 @@ let rec for_signature_item ~env ~(exported : Exported.t)
         incl_type []
     in
     top_level
+  | Tsig_typext type_extension ->
+    add_extension_constructors ~env type_extension;
+    []
   (* TODO: process other things here *)
   | _ -> []
 
@@ -645,6 +707,9 @@ let rec for_structure_item ~(env : Shared_types.Env.t) ~(exported : Exported.t)
           | _ -> Types.Trec_next
         in
         decl |> for_type_declaration ~env ~exported ~rec_status)
+  | Tstr_typext type_extension ->
+    add_extension_constructors ~env type_extension;
+    []
   | _ -> []
 
 and for_module ~env mod_desc module_name =

analysis/src/shared_types.ml

@@ -378,6 +378,9 @@ and completion_type =
   | Tvariant of {
       env: Query_env.t;
       constructors: Constructor.t list;
+      extension_envs: (Constructor.t * Query_env.t) list;
+          (** The environments of the files declaring extension
+              constructors, where their payloads are typed *)
       variant_decl: Types.type_declaration;
       variant_name: string;
     }
@@ -606,6 +609,8 @@ module Completable = struct
     | NRecordBody of {seen_fields: string list}
     | NVariantPayload of {
         constructor_name: string;
+        constructor_path: string list;
+            (** The constructor as written, with its module path *)
         item_num: int;
         source_arity: int;
       }