feat: custom tag names via type[@mel.tag "foo"] t = .. (#1189)

* feat: custom tag names via `type[@mel.tag "foo"] t = ..`

* test: add cram test

* refactor: use inline record for J.Caml_block_tag

* chore: add changelog entry

Changes.md

@@ -23,6 +23,11 @@ Unreleased
 - Support OCaml 5.3 ([#1168](https://github.com/melange-re/melange/pull/1168))
 - Upgrade Stdlib to the OCaml 5.3 Stdlib
   ([#1182](https://github.com/melange-re/melange/pull/1182))
+- Support `[@mel.tag "the_tag"]` in variant definitions
+  ([#1189](https://github.com/melange-re/melange/pull/1189))
+  - combined with `[@mel.as ..]` in variant definitions and inline record
+    payloads, Melange can now express types for discriminated union object
+    shapes.
 
 4.0.1 2024-06-07
 ---------------

jscomp/core/j.ml

@@ -179,7 +179,7 @@ and expression_desc =
      [Caml_block_tag] can return [undefined],
      you have to use [E.tag] in a safe way
   *)
-  | Caml_block_tag of expression
+  | Caml_block_tag of { expr : expression; name : string }
   (* | Caml_block_set_tag of expression * expression *)
   (* | Caml_block_set_length of expression * expression *)
   (* It will just fetch tag, to make it safe, when creating it,

jscomp/core/js_analyzer.ml

@@ -109,7 +109,7 @@ let rec no_side_effect_expression_desc (x : J.expression_desc) =
   | Length { expr = e; _ }
   | Char_of_int e
   | Char_to_int e
-  | Caml_block_tag e
+  | Caml_block_tag { expr = e; _ }
   | Typeof e ->
       no_side_effect e
   | Bin { op; expr1 = a; expr2 = b } ->

jscomp/core/js_dump.ml

@@ -50,8 +50,6 @@ open Import
 
 *)
 
-let name_symbol = Js_op.Symbol_name
-
 module E = Js_exp_make
 module S = Js_stmt_make
 module L = Js_dump_lit
@@ -803,13 +801,19 @@ and expression_desc cxt ~(level : int) x : cxt =
       let objs =
         let tails =
           List.map_combine_array_append p.fields el
-            (if !Js_config.debug then [ (name_symbol, E.str p.name) ] else [])
+            (if !Js_config.debug then [ (Js_op.Symbol_name, E.str p.name) ]
+             else [])
             (fun i -> Js_op.Lit i)
         in
         let as_value =
           Lam_constant_convert.modifier ~name:p.name p.attributes
         in
-        ( Js_op.Lit L.tag,
+        let tag_name =
+          match Record_attributes_check.process_tag_name p.attributes with
+          | None -> L.tag
+          | Some s -> s
+        in
+        ( Js_op.Lit tag_name,
           {
             (match as_value.as_modifier with
             | Some modifier -> E.as_value modifier
@@ -830,15 +834,20 @@ and expression_desc cxt ~(level : int) x : cxt =
             el
           @
           if !Js_config.debug && not is_cons then
-            [ (name_symbol, E.str p.name) ]
+            [ (Symbol_name, E.str p.name) ]
           else []
         in
         if is_cons && p.num_nonconst = 1 then tails
         else
           let as_value =
             Lam_constant_convert.modifier ~name:p.name p.attributes
           in
-          ( Js_op.Lit L.tag,
+          let tag_name =
+            match Record_attributes_check.process_tag_name p.attributes with
+            | None -> L.tag
+            | Some s -> s
+          in
+          ( Js_op.Lit tag_name,
             {
               (match as_value.as_modifier with
               | Some modifier -> E.as_value modifier
@@ -858,11 +867,11 @@ and expression_desc cxt ~(level : int) x : cxt =
         _;
       } ->
       expression_desc cxt ~level (Array { items = el; mutable_flag })
-  | Caml_block_tag e ->
+  | Caml_block_tag { expr = e; name } ->
       group cxt 1 (fun () ->
           let cxt = expression ~level:15 cxt e in
           string cxt L.dot;
-          string cxt L.tag;
+          string cxt name;
           cxt)
   | Array_index { expr = e; index = p } | String_index { expr = e; index = p }
     ->

jscomp/core/js_exp_make.ml

@@ -847,7 +847,8 @@ let is_type_string ?loc ?comment (e : t) : t =
    call plain [dot]
 *)
 
-let tag ?loc ?comment e : t = make_expression ?loc ?comment (Caml_block_tag e)
+let tag ?loc ?comment ?(name = Js_dump_lit.tag) expr : t =
+  make_expression ?loc ?comment (Caml_block_tag { expr; name })
 
 (* according to the compiler, [Btype.hash_variant],
    it's reduced to 31 bits for hash

jscomp/core/js_exp_make.mli

@@ -295,7 +295,9 @@ val unit : t
 (** [unit] in ocaml will be compiled into [0]  in js *)
 
 val undefined : t
-val tag : ?loc:Location.t -> ?comment:string -> J.expression -> t
+
+val tag :
+  ?loc:Location.t -> ?comment:string -> ?name:string -> J.expression -> t
 
 (** Note that this is coupled with how we encode block, if we use the
     `Object.defineProperty(..)` since the array already hold the length,

jscomp/core/lam_compile.ml

@@ -163,9 +163,19 @@ let default_action ~saturated failaction =
 let get_const_name i (sw_names : Lambda.switch_names option) =
   match sw_names with None -> None | Some { consts; _ } -> Some consts.(i)
 
-let get_block_name i (sw_names : Lambda.switch_names option) =
+let get_block i (sw_names : Lambda.switch_names option) =
   match sw_names with None -> None | Some { blocks; _ } -> Some blocks.(i)
 
+let get_tag_name (sw_names : Lambda.switch_names option) =
+  match sw_names with
+  | None -> Js_dump_lit.tag
+  | Some { blocks; _ } -> (
+      match
+        Array.find_map ~f:(fun { Lambda.tag_name; _ } -> tag_name) blocks
+      with
+      | Some s -> s
+      | None -> Js_dump_lit.tag)
+
 let no_effects_const = lazy true
 (* let has_effects_const = lazy false *)
 
@@ -699,7 +709,13 @@ and compile_switch (switch_arg : Lam.t) (sw : Lam.lambda_switch)
     default_action ~saturated:sw_blocks_full sw_failaction
   in
   let get_const_name i = get_const_name i sw_names in
-  let get_block_name i = get_block_name i sw_names in
+  let get_block i = get_block i sw_names in
+  let get_block_name i =
+    match get_block i with
+    | Some { cstr_name; _ } -> Some cstr_name
+    | None -> None
+  in
+  let tag_name = get_tag_name sw_names in
   let compile_whole (cxt : Lam_compile_context.t) =
     match
       compile_lambda { cxt with continuation = NeedValue Not_tail } switch_arg
@@ -709,7 +725,7 @@ and compile_switch (switch_arg : Lam.t) (sw : Lam.lambda_switch)
         block
         @
         if sw_consts_full && sw_consts = [] then
-          compile_cases cxt (E.tag e) sw_blocks sw_blocks_default
+          compile_cases cxt (E.tag ~name:tag_name e) sw_blocks sw_blocks_default
             ~get_cstr_name:get_block_name
         else if sw_blocks_full && sw_blocks = [] then
           compile_cases cxt e sw_consts sw_num_default
@@ -722,8 +738,8 @@ and compile_switch (switch_arg : Lam.t) (sw : Lam.lambda_switch)
                  ~get_cstr_name:get_const_name)
               (* default still needed, could simplified*)
               ~else_:
-                (compile_cases cxt (E.tag e) sw_blocks sw_blocks_default
-                   ~get_cstr_name:get_block_name)
+                (compile_cases cxt (E.tag ~name:tag_name e) sw_blocks
+                   sw_blocks_default ~get_cstr_name:get_block_name)
           in
           match e.expression_desc with
           | J.Var _ -> [ dispatch e ]

jscomp/core/lam_print.ml

@@ -355,7 +355,7 @@ let lambda ppf v =
               fprintf ppf "@[<hv 1>case tag %i %S:@ %a@]" n
                 (match sw.sw_names with
                 | None -> ""
-                | Some x -> x.blocks.(n).name)
+                | Some x -> x.blocks.(n).cstr_name.name)
                 lam l)
             sw.sw_blocks;
           match sw.sw_failaction with

jscomp/core/matching_polyfill.cppo.ml

@@ -33,14 +33,19 @@ let names_from_construct_pattern
     let get_cstr_name (cstr: Types.constructor_declaration) =
       Lam_constant_convert.modifier ~name:(Ident.name cstr.cd_id) cstr.cd_attributes
     in
+    let get_block (cstr: Types.constructor_declaration) =
+      { Lambda.cstr_name = get_cstr_name cstr
+      ; tag_name = Record_attributes_check.process_tag_name cstr.cd_attributes
+      }
+    in
     let consts, blocks =
       List.fold_left
         ~f:(fun (consts, blocks)
                 (cstr : Types.constructor_declaration) ->
           if is_nullary_variant cstr.cd_args then
             (get_cstr_name cstr :: consts , blocks)
           else
-            (consts , get_cstr_name cstr :: blocks))
+            (consts , get_block cstr :: blocks))
         ~init:([], []) cstrs
     in
     Some

jscomp/core/record_attributes_check.cppo.ml

@@ -34,6 +34,41 @@ let namespace_error ~loc txt =
          `[@mel.*]' attributes. Use `[@mel.as]' instead."
   | _ -> ()
 
+
+let process_tag_name attrs =
+  let st = ref None in
+  List.iter attrs
+    ~f:(fun { Parsetree.attr_name = { txt; loc }; attr_payload; _ } ->
+      match txt with
+      | "mel.tag" ->
+          if !st = None then (
+            (match attr_payload with
+            | PStr
+                [
+                  {
+                    pstr_desc =
+                      Pstr_eval ({ pexp_desc = Pexp_constant const; _ }, _);
+                    _;
+                  };
+                ] -> (
+                namespace_error ~loc txt;
+                match const.pconst_desc with
+                | Pconst_string (s, _, _) -> st := Some s
+                | _ -> ())
+            | _ -> ());
+            if !st = None
+            then
+              Location.raise_errorf
+                ~loc
+                 "Variant tag annotation (`[@mel.tag \"..\"]') must be a string")
+          else
+            Location.raise_errorf
+              ~loc
+               "Duplicate `[@mel.tag \"..\"]' annotation"
+      | _ -> ());
+  !st
+
+
 let find_mel_as_name =
   let find_mel_as_name (attr : Parsetree.attribute) =
     match attr.attr_name with

jscomp/runtime/melange_mini_stdlib.ml

@@ -57,8 +57,6 @@ module Obj = struct
   external set_field : t -> int -> t -> unit = "%obj_set_field"
   external tag : t -> int = "caml_obj_tag"
 
-  (* The compiler ensures (|0) operation *)
-  external set_tag : t -> int -> unit = "TAG" [@@mel.set]
   external repr : 'a -> t = "%identity"
   external obj : t -> 'a = "%identity"
   external magic : 'a -> 'b = "%identity"

jscomp/stdlib/camlinternalLazy.ml

@@ -48,7 +48,6 @@ let is_val (type a ) (l : a lazy_t) : bool =
 
 let forward_with_closure (type a ) (blk : a concrete) (closure : unit -> a [@u]) : a =
   let result = closure () [@u] in
-  (* do set_field BEFORE set_tag *)
   blk.value <- result;
   blk.tag<- true;
   result

ppx/melange_ppx.cppo.ml

@@ -786,6 +786,12 @@ module Mapper = struct
               ~f:(fun (tdcl : type_declaration) ->
                 match tdcl.ptype_kind with
                 | Ptype_variant cstrs ->
+                    List.iter
+                      ~f:(function
+                          | ({ attr_name= { txt = "mel.tag"; _ }; _ } as attr) ->
+                            Mel_ast_invariant.mark_used_mel_attribute attr
+                          | _ -> ())
+                    tdcl.ptype_attributes;
                     List.iter
                       ~f:(fun
                           ({ pcd_attributes; _ } : constructor_declaration) ->

test/blackbox-tests/custom-tag.t

@@ -0,0 +1,102 @@
+Test `@@mel.tag` attribute for discriminated unions
+
+  $ . ./setup.sh
+  $ cat > x.ml <<EOF
+  > type foo = { a : string; b : string }
+  > type[@mel.tag "lol"] t = Foo of foo | Bar of { c : string; d : string }
+  > let x, y =
+  >   let x = Foo { a = "a"; b = "b" } in
+  >   let y = Bar { c = "c"; d = "d" } in
+  >   (x, y)
+  > let s = function Foo { a; b } -> (a, b) | Bar { c; d } -> (c, d)
+  > EOF
+
+  $ melc -ppx melppx x.ml
+  // Generated by Melange
+  'use strict';
+  
+  
+  function s(param) {
+    if (param.lol !== /* Foo */0) {
+      return [
+              param.c,
+              param.d
+            ];
+    }
+    const match = param._0;
+    return [
+            match.a,
+            match.b
+          ];
+  }
+  
+  const x = {
+    lol: /* Foo */0,
+    _0: {
+      a: "a",
+      b: "b"
+    }
+  };
+  
+  const y = {
+    lol: /* Bar */1,
+    c: "c",
+    d: "d"
+  };
+  
+  exports.x = x;
+  exports.y = y;
+  exports.s = s;
+  /* No side effect */
+
+`@mel.tag` + `[@mel.as "string"]`
+
+  $ cat > x.ml <<EOF
+  > type foo = { a : string; b : string }
+  > type[@mel.tag "lol"] t =
+  >   | Foo of foo
+  >   | Bar of { c : string; d : string } [@mel.as "bar"]
+  > let x, y =
+  >   let x = Foo { a = "a"; b = "b" } in
+  >   let y = Bar { c = "c"; d = "d" } in
+  >   (x, y)
+  > let s = function Foo { a; b } -> (a, b) | Bar { c; d } -> (c, d)
+  > EOF
+
+  $ melc -ppx melppx x.ml
+  // Generated by Melange
+  'use strict';
+  
+  
+  function s(param) {
+    if (param.lol !== /* Foo */0) {
+      return [
+              param.c,
+              param.d
+            ];
+    }
+    const match = param._0;
+    return [
+            match.a,
+            match.b
+          ];
+  }
+  
+  const x = {
+    lol: /* Foo */0,
+    _0: {
+      a: "a",
+      b: "b"
+    }
+  };
+  
+  const y = {
+    lol: /* Bar */"bar",
+    c: "c",
+    d: "d"
+  };
+  
+  exports.x = x;
+  exports.y = y;
+  exports.s = s;
+  /* No side effect */