initial implementation of expanding variant type spreads

jscomp/ml/typedecl.ml

@@ -379,34 +379,77 @@ let transl_declaration ~typeRecordAsObject env sdecl id =
         let copy_tag_attr_from_decl attr =
           let tag_attrs = Ext_list.filter sdecl.ptype_attributes (fun ({txt}, _) -> txt = "tag" || txt = Ast_untagged_variants.untagged) in
           if tag_attrs = [] then attr else tag_attrs @ attr in
+        let constructors_from_variant_spreads = Hashtbl.create 10 in
         let make_cstr scstr =
           let name = Ident.create scstr.pcd_name.txt in
           let targs, tret_type, args, ret_type, _cstr_params =
             make_constructor env (Path.Pident id) params
                              scstr.pcd_args scstr.pcd_res
           in
-          let tcstr =
-            { cd_id = name;
-              cd_name = scstr.pcd_name;
-              cd_args = targs;
-              cd_res = tret_type;
-              cd_loc = scstr.pcd_loc;
-              cd_attributes = scstr.pcd_attributes |> copy_tag_attr_from_decl }
-          in
-          let cstr =
-            { Types.cd_id = name;
-              cd_args = args;
-              cd_res = ret_type;
-              cd_loc = scstr.pcd_loc;
-              cd_attributes = scstr.pcd_attributes |> copy_tag_attr_from_decl }
-          in
+          if String.starts_with scstr.pcd_name.txt ~prefix:"..." then (
+            (match args with
+            | Cstr_tuple [spread_variant] -> (
+              match Ctype.extract_concrete_typedecl env spread_variant with
+              | (_, _, {type_kind=Type_variant constructors}) -> (
+                constructors |> List.iter(fun (c: Types.constructor_declaration) ->
+                  Hashtbl.add constructors_from_variant_spreads c.cd_id.name c)
+              )
+              | _ -> ()
+            )
+            | _ -> ()); 
+              None)
+          else (
+          let tcstr, cstr = match Hashtbl.find_opt constructors_from_variant_spreads name.name with
+          | Some cstr ->
+            let tcstr =
+              {
+                cd_id = name;
+                cd_name = scstr.pcd_name;
+                cd_args =
+                  (match cstr.cd_args with
+                  | Cstr_tuple args ->
+                    Cstr_tuple
+                      (args
+                      |> List.map (fun texpr : Typedtree.core_type ->
+                             {
+                               ctyp_attributes = [];
+                               ctyp_loc = cstr.cd_loc;
+                               ctyp_env = env;
+                               ctyp_type = texpr;
+                               ctyp_desc = Ttyp_any;
+                             }))
+                  | Cstr_record _lbls -> assert false (* TODO: Translate *));
+                cd_res = tret_type;
+                cd_loc = scstr.pcd_loc;
+                cd_attributes = scstr.pcd_attributes |> copy_tag_attr_from_decl;
+              }
+            in
+            tcstr, cstr
+          | None ->
+            let tcstr =
+              { cd_id = name;
+                cd_name = scstr.pcd_name;
+                cd_args = targs;
+                cd_res = tret_type;
+                cd_loc = scstr.pcd_loc;
+                cd_attributes = scstr.pcd_attributes |> copy_tag_attr_from_decl }
+            in
+            let cstr =
+              { Types.cd_id = name;
+                cd_args = args;
+                cd_res = ret_type;
+                cd_loc = scstr.pcd_loc;
+                cd_attributes = scstr.pcd_attributes |> copy_tag_attr_from_decl }
+            in
             tcstr, cstr
+          in Some (tcstr, cstr)
+          )
         in
         let make_cstr scstr =
           Builtin_attributes.warning_scope scstr.pcd_attributes
             (fun () -> make_cstr scstr)
         in
-        let tcstrs, cstrs = List.split (List.map make_cstr scstrs) in
+        let tcstrs, cstrs = List.split (List.filter_map make_cstr scstrs) in
         let isUntaggedDef = Ast_untagged_variants.has_untagged sdecl.ptype_attributes in
         Ast_untagged_variants.check_well_formed ~env ~isUntaggedDef cstrs;
         Ttype_variant tcstrs, Type_variant cstrs, sdecl
@@ -1270,7 +1313,7 @@ let transl_type_decl env rec_flag sdecl_list =
         {sdecl with
          ptype_name; ptype_kind = Ptype_abstract; ptype_manifest = None})
       fixed_types
-    @ sdecl_list
+    @ (sdecl_list |> Variant_type_spread.expand_variant_spreads env)
   in
 
   (* Create identifiers. *)
@@ -1324,6 +1367,7 @@ let transl_type_decl env rec_flag sdecl_list =
     List.map2 transl_declaration sdecl_list (List.map id_slots id_list) in
   let decls =
     List.map (fun tdecl -> (tdecl.typ_id, tdecl.typ_type)) tdecls in
+  let sdecl_list = Variant_type_spread.expand_dummy_constructor_args sdecl_list decls in
   current_slot := None;
   (* Check for duplicates *)
   check_duplicates sdecl_list;

jscomp/ml/variant_type_spread.ml

@@ -0,0 +1,134 @@
+let mk_constructor_comes_from_spread_attr () : Parsetree.attribute =
+  (Location.mknoloc "res.constructor_from_spread", PStr [])
+
+(* Spreads in variants are parsed as constructors named "...", with a single payload that's an identifier
+   pointing to the type that's spread. We need to expand those constructors as soon as we can, before type
+   checking. So, here we look for constructors named "...", look up their type, and add the constructors that
+   type itself has.
+*)
+let expand_variant_spreads (env : Env.t)
+    (sdecl_list : Parsetree.type_declaration list) =
+  sdecl_list
+  |> List.map (fun (sdecl : Parsetree.type_declaration) ->
+         match sdecl with
+         | {ptype_kind = Ptype_variant constructors} ->
+           {
+             sdecl with
+             ptype_kind =
+               Ptype_variant
+                 (constructors
+                 |> List.map (fun (c : Parsetree.constructor_declaration) ->
+                        match c with
+                        | {
+                         pcd_name = {txt = "..."};
+                         pcd_args =
+                           Pcstr_tuple
+                             [{ptyp_loc; ptyp_desc = Ptyp_constr (loc, [])}];
+                        } -> (
+                          try
+                            let _path, type_decl =
+                              Typetexp.find_type env ptyp_loc loc.txt
+                            in
+                            match type_decl with
+                            | {type_kind = Type_variant cstrs} ->
+                              let spread_constructor_name =
+                                "..."
+                                ^ (Longident.flatten loc.txt
+                                 |> String.concat ".")
+                              in
+                              {
+                                c with
+                                pcd_name =
+                                  {
+                                    c.pcd_name with
+                                    txt = spread_constructor_name;
+                                  };
+                              }
+                              :: (cstrs
+                                 |> List.map
+                                      (fun
+                                        (cstr : Types.constructor_declaration)
+                                        :
+                                        Parsetree.constructor_declaration
+                                      ->
+                                        {
+                                          (* This will mark this constructor as originating from a variant type spread.
+                                             We use that hint to fill in the real, typed constructor arguments (if any)
+                                             at a later stage when that information is available. *)
+                                          pcd_attributes =
+                                            mk_constructor_comes_from_spread_attr
+                                              ()
+                                            :: cstr.cd_attributes;
+                                          pcd_loc = cstr.cd_loc;
+                                          pcd_res = None;
+                                          (* It's important that we _don't_ fill in pcd_args here, since we have no way to produce
+                                             a valid set of args for the parsetree at this stage. Inserting dummies here instead
+                                             of later means that our dummies would end up being typechecked, and we don't want that.
+                                             We'll fill this in with dummy info later. *)
+                                          pcd_args = Pcstr_tuple [];
+                                          pcd_name =
+                                            Location.mkloc cstr.cd_id.name
+                                              cstr.cd_loc;
+                                        }))
+                            | _ -> [c]
+                          with _ ->
+                            (* Did not find type. Can't spread here, report as error that types need to be known before hand. *)
+                            [c])
+                        | _ -> [c])
+                 |> List.concat);
+           }
+         | _ -> sdecl)
+
+let constructor_is_from_spread (attrs : Parsetree.attributes) =
+  attrs
+  |> List.exists (fun (a : Parsetree.attribute) ->
+         match a with
+         | {txt = "res.constructor_from_spread"}, PStr [] -> true
+         | _ -> false)
+
+(* The type checker matches lengths of constructor arguments etc between the parsetree and 
+   typed constructor definitions at various places. However, it doesn't use the parsetree 
+   definition for anything but the loc after having done the initial type check. So, here
+   we add dummy constructor arguments that match the actual number of arguments that the 
+   type checker has told us each constructor has, so the various invariants check out. *)
+let expand_dummy_constructor_args (sdecl_list : Parsetree.type_declaration list)
+    (decls : (Ident.t * Types.type_declaration) list) =
+  List.map2
+    (fun sdecl (_, decl) ->
+      match (sdecl, decl) with
+      | ( {Parsetree.ptype_kind = Ptype_variant c1},
+          {Types.type_kind = Type_variant c2} ) ->
+        {
+          sdecl with
+          ptype_kind =
+            Ptype_variant
+              (c1
+              |> List.map (fun (c : Parsetree.constructor_declaration) ->
+                     if constructor_is_from_spread c.pcd_attributes then
+                       match
+                         c2
+                         |> List.find_opt
+                              (fun (cc : Types.constructor_declaration) ->
+                                Ident.name cc.cd_id = c.pcd_name.txt)
+                       with
+                       | None -> c
+                       | Some constructor -> (
+                         match constructor with
+                         | {cd_args = Cstr_tuple args} ->
+                           {
+                             c with
+                             pcd_args =
+                               Pcstr_tuple
+                                 (args
+                                 |> List.map (fun _t ->
+                                        {
+                                          Parsetree.ptyp_loc = c.pcd_loc;
+                                          ptyp_attributes = [];
+                                          ptyp_desc = Ptyp_any;
+                                        }));
+                           }
+                         | _ -> c)
+                     else c));
+        }
+      | _ -> sdecl)
+    sdecl_list decls

jscomp/test/VariantSpreads.res

@@ -0,0 +1,15 @@
+module S = {
+    type x = Foo | Bar
+    type s = Five(int) | Six
+}
+
+type a = One(bool, S.x) | Two
+
+type b = | ...a | Three | Four | ...S.s
+
+let b1: b = Two
+let b2: b = One(true, Bar)
+
+let c: b = Five(2)
+
+let ddd: b = Six