ffi: allow @mel.unwrap + polyvariants without payload (#1239)

* ffi: allow `@mel.unwrap` + polyvariants without payload * remove test that no longer errors * test: add a cram test * test: show failure with @mel.as * support @mel.as * restrict it a bit * wip * ints' and strings' * fix regression on regular mel.unwrap * refactor: one common place for the logic * chore: add changelog entry * refactor: unify logic
melange-re · Dec 7, 2024 · 9016a73 · 9016a73
1 parent d0b3e3a
commit 9016a73
Show file tree

Hide file tree

Showing 10 changed files with 199 additions and 115 deletions.
diff --git a/Changes.md b/Changes.md
@@ -38,6 +38,8 @@ Unreleased
 - ppx,core: propagate internal FFI information via attributes instead of adding
   marshalled data in the native primitive name
   ([#1222](https://github.com/melange-re/melange/pull/1222))
+- melange-ppx: allow `@mel.unwrap` polyvariants not to have a payload
+  ([#1239](https://github.com/melange-re/melange/pull/1239))
 
 4.0.1 2024-06-07
 ---------------

diff --git a/jscomp/common/external_arg_spec.ml b/jscomp/common/external_arg_spec.ml
@@ -70,7 +70,7 @@ type attr =
   | Extern_unit
   | Nothing
   | Ignore
-  | Unwrap
+  | Unwrap of attr
 
 type 'a param = { arg_type : attr; arg_label : 'a }
 

diff --git a/jscomp/common/external_arg_spec.mli b/jscomp/common/external_arg_spec.mli
@@ -38,7 +38,7 @@ type attr =
   | Extern_unit
   | Nothing
   | Ignore
-  | Unwrap
+  | Unwrap of attr
 
 module Arg_label : sig
   type t = Arg_label | Arg_empty | Arg_optional

diff --git a/jscomp/core/js_of_lam_variant.ml b/jscomp/core/js_of_lam_variant.ml
@@ -130,4 +130,5 @@ let eval_as_int (arg : J.expression) (dispatches : (string * int) list) : E.t =
 let eval_as_unwrap (arg : J.expression) : E.t =
   match arg.expression_desc with
   | Caml_block { fields = [ { expression_desc = Number _; _ }; cb ]; _ } -> cb
+  | Str _ | Unicode _ -> arg
   | _ -> E.poly_var_value_access arg
diff --git a/jscomp/core/lam_compile_external_call.ml b/jscomp/core/lam_compile_external_call.ml
@@ -101,7 +101,7 @@ let append_list x xs =
 
      This would not work with [NonNullString]
 *)
-let ocaml_to_js_eff ~(arg_label : Melange_ffi.External_arg_spec.Arg_label.t)
+let rec ocaml_to_js_eff ~(arg_label : Melange_ffi.External_arg_spec.Arg_label.t)
     ~(arg_type : Melange_ffi.External_arg_spec.attr) (raw_arg : E.t) :
     arg_expression * E.t list =
   let arg =
@@ -131,11 +131,19 @@ let ocaml_to_js_eff ~(arg_label : Melange_ffi.External_arg_spec.Arg_label.t)
       *)
   | Int dispatches ->
       (Splice1 (Js_of_lam_variant.eval_as_int arg dispatches), [])
-  | Unwrap ->
-      let single_arg =
-        match arg_label with
-        | Arg_optional ->
-            (*
+  | Unwrap polyvar -> (
+      match (polyvar, raw_arg.expression_desc) with
+      | (Poly_var_string _ | Poly_var _ | Int _), Caml_block _ ->
+          Location.raise_errorf ?loc:raw_arg.loc
+            "`[@mel.as ..]' can only be used with `[@mel.unwrap]' variants \
+             without a payload."
+      | (Poly_var_string _ | Poly_var _ | Int _), _ ->
+          ocaml_to_js_eff ~arg_label ~arg_type:polyvar raw_arg
+      | Nothing, _ ->
+          let single_arg =
+            match arg_label with
+            | Arg_optional ->
+                (*
            If this is an optional arg (like `?arg`), we have to potentially do
            2 levels of unwrapping:
            - if ocaml arg is `None`, let js arg be `undefined` (no unwrapping)
@@ -144,10 +152,11 @@ let ocaml_to_js_eff ~(arg_label : Melange_ffi.External_arg_spec.Arg_label.t)
            - Here `Some x` is `x` due to the current encoding
            Lets inline here since it depends on the runtime encoding
         *)
-            Js_of_lam_option.option_unwrap raw_arg
-        | _ -> Js_of_lam_variant.eval_as_unwrap raw_arg
-      in
-      (Splice1 single_arg, [])
+                Js_of_lam_option.option_unwrap raw_arg
+            | _ -> Js_of_lam_variant.eval_as_unwrap raw_arg
+          in
+          (Splice1 single_arg, [])
+      | _, _ -> assert false)
   | Nothing -> (Splice1 arg, [])
 
 let empty_pair = ([], [])

diff --git a/ppx/ast_external_process.ml b/ppx/ast_external_process.ml
@@ -28,21 +28,54 @@ module External_ffi_types = Melange_ffi.External_ffi_types
 
 (* record pattern match complete checker *)
 
-let rec variant_can_unwrap_aux (row_fields : row_field list) : bool =
-  match row_fields with
-  | [] -> true
-  | { prf_desc = Rtag (_, false, [ _ ]); _ } :: rest ->
-      variant_can_unwrap_aux rest
-  | _ :: _ -> false
+let variant_unwrap =
+  let rec variant_can_unwrap_aux (row_fields : row_field list) : bool =
+    match row_fields with
+    | [] -> true
+    | { prf_desc = Rtag (_, true, []); _ } :: rest ->
+        variant_can_unwrap_aux rest
+    | { prf_desc = Rtag (_, false, [ _ ]); _ } :: rest ->
+        variant_can_unwrap_aux rest
+    | _ :: _ -> false
+  in
+  fun (row_fields : row_field list) : bool ->
+    match row_fields with
+    | [] -> false (* impossible syntax *)
+    | xs -> variant_can_unwrap_aux xs
 
-let variant_unwrap (row_fields : row_field list) : bool =
-  match row_fields with
-  | [] -> false (* impossible syntax *)
-  | xs -> variant_can_unwrap_aux xs
+let infer_mel_as ~loc row_fields ~allow_no_payload =
+  let mel_as_type =
+    (* No `@mel.string` / `@mel.int` present. Try to infer `@mel.as`, if
+       present, in polyvariants.
 
-(*
-  TODO: [nolabel] is only used once turn Nothing into Unit, refactor later
-*)
+       https://github.com/melange-re/melange/issues/578 *)
+    List.fold_left
+      ~f:(fun mel_as_type { prf_attributes; prf_loc; _ } ->
+        match List.filter ~f:Ast_attributes.is_mel_as prf_attributes with
+        | [] -> mel_as_type
+        | [ { attr_payload; attr_loc = loc; _ } ] -> (
+            match
+              ( mel_as_type,
+                Ast_payload.is_single_string attr_payload,
+                Ast_payload.is_single_int attr_payload )
+            with
+            | (`Nothing | `String), Some _, None -> `String
+            | (`Nothing | `Int), None, Some _ -> `Int
+            | (`Nothing | `String | `Int), None, None -> `Nothing
+            | `String, None, Some _ -> Error.err ~loc Expect_string_literal
+            | `Int, Some _, None -> Error.err ~loc Expect_int_literal
+            | _, Some _, Some _ -> assert false)
+        | _ :: _ -> Error.err ~loc:prf_loc Duplicated_mel_as)
+      ~init:`Nothing row_fields
+  in
+  match mel_as_type with
+  | `Nothing -> External_arg_spec.Nothing
+  | `String ->
+      Ast_polyvar.map_row_fields_into_strings row_fields ~loc ~allow_no_payload
+  | `Int ->
+      Ast_polyvar.map_row_fields_into_ints row_fields ~loc ~allow_no_payload
+
+(* TODO: [nolabel] is only used once turn Nothing into Unit, refactor later *)
 let spec_of_ptyp ~(nolabel : bool) (ptyp : core_type) : External_arg_spec.attr =
   let ptyp_desc = ptyp.ptyp_desc in
   match
@@ -52,22 +85,23 @@ let spec_of_ptyp ~(nolabel : bool) (ptyp : core_type) : External_arg_spec.attr =
   | `String -> (
       match ptyp_desc with
       | Ptyp_variant (row_fields, Closed, None) ->
-          Ast_polyvar.map_row_fields_into_strings ptyp.ptyp_loc row_fields
+          Ast_polyvar.map_row_fields_into_strings row_fields ~loc:ptyp.ptyp_loc
+            ~allow_no_payload:false
       | _ -> Error.err ~loc:ptyp.ptyp_loc Invalid_mel_string_type)
   | `Ignore -> Ignore
   | `Int -> (
       match ptyp_desc with
       | Ptyp_variant (row_fields, Closed, None) ->
-          let int_lists =
-            Ast_polyvar.map_row_fields_into_ints ptyp.ptyp_loc row_fields
-          in
-          Int int_lists
+          Ast_polyvar.map_row_fields_into_ints row_fields ~loc:ptyp.ptyp_loc
+            ~allow_no_payload:false
       | _ -> Error.err ~loc:ptyp.ptyp_loc Invalid_mel_int_type)
   | `Unwrap -> (
       match ptyp_desc with
       | Ptyp_variant (row_fields, Closed, _) when variant_unwrap row_fields ->
+          (* Unwrap attribute can only be attached to things like
+             `[a of a0 | b of b0]` *)
           Unwrap
-          (* Unwrap attribute can only be attached to things like `[a of a0 | b of b0]` *)
+            (infer_mel_as ~loc:ptyp.ptyp_loc row_fields ~allow_no_payload:true)
       | _ -> Error.err ~loc:ptyp.ptyp_loc Invalid_mel_unwrap_type)
   | `Uncurry opt_arity -> (
       let real_arity = Ast_core_type.get_uncurry_arity ptyp in
@@ -84,41 +118,8 @@ let spec_of_ptyp ~(nolabel : bool) (ptyp : core_type) : External_arg_spec.attr =
       match ptyp_desc with
       | Ptyp_constr ({ txt = Lident "unit"; _ }, []) ->
           if nolabel then Extern_unit else Nothing
-      | Ptyp_variant (row_fields, Closed, None) -> (
-          (* No `@mel.string` / `@mel.int` present. Try to infer `@mel.as`, if
-             present, in polyvariants.
-
-             https://github.com/melange-re/melange/issues/578 *)
-          let mel_as_type =
-            List.fold_left
-              ~f:(fun mel_as_type { prf_attributes; prf_loc; _ } ->
-                match
-                  List.filter ~f:Ast_attributes.is_mel_as prf_attributes
-                with
-                | [] -> mel_as_type
-                | [ { attr_payload; attr_loc = loc; _ } ] -> (
-                    match
-                      ( mel_as_type,
-                        Ast_payload.is_single_string attr_payload,
-                        Ast_payload.is_single_int attr_payload )
-                    with
-                    | (`Nothing | `String), Some _, None -> `String
-                    | (`Nothing | `Int), None, Some _ -> `Int
-                    | (`Nothing | `String | `Int), None, None -> `Nothing
-                    | `String, None, Some _ ->
-                        Error.err ~loc Expect_string_literal
-                    | `Int, Some _, None -> Error.err ~loc Expect_int_literal
-                    | _, Some _, Some _ -> assert false)
-                | _ :: _ -> Error.err ~loc:prf_loc Duplicated_mel_as)
-              ~init:`Nothing row_fields
-          in
-          match mel_as_type with
-          | `Nothing -> Nothing
-          | `String ->
-              Ast_polyvar.map_row_fields_into_strings ptyp.ptyp_loc row_fields
-          | `Int ->
-              Int
-                (Ast_polyvar.map_row_fields_into_ints ptyp.ptyp_loc row_fields))
+      | Ptyp_variant (row_fields, Closed, None) ->
+          infer_mel_as ~loc:ptyp.ptyp_loc row_fields ~allow_no_payload:false
       | _ -> Nothing)
 
 let const_payload_cst = function
@@ -527,7 +528,7 @@ let process_obj (loc : Location.t) (st : external_desc) (prim_name : string)
                           arg_types,
                           (* ignored in [arg_types], reserved in [result_types] *)
                           result_types )
-                    | Nothing | Unwrap ->
+                    | Nothing | Unwrap _ ->
                         let s = Melange_ffi.Lam_methname.translate name in
                         ( {
                             arg_label = External_arg_spec.Obj_label.obj s;
@@ -575,7 +576,7 @@ let process_obj (loc : Location.t) (st : external_desc) (prim_name : string)
                         ( External_arg_spec.empty_kind obj_arg_type,
                           param_type :: arg_types,
                           result_types )
-                    | Nothing | Unwrap ->
+                    | Nothing | Unwrap _ ->
                         let s = Melange_ffi.Lam_methname.translate name in
                         (* XXX(anmonteiro): it's unsafe to just read the type
                              of the labelled argument declaration, since it

diff --git a/ppx/ast_polyvar.ml b/ppx/ast_polyvar.ml
@@ -38,7 +38,8 @@ let is_enum_polyvar =
         Some row_fields
     | _ -> None
 
-let map_row_fields_into_ints ptyp_loc (row_fields : row_field list) =
+let map_row_fields_into_ints (row_fields : row_field list) ~loc
+    ~allow_no_payload =
   let _, acc =
     List.fold_left
       ~f:(fun (i, acc) rtag ->
@@ -52,53 +53,56 @@ let map_row_fields_into_ints ptyp_loc (row_fields : row_field list) =
               | None -> i
             in
             (i + 1, (txt, i) :: acc)
-        | _ -> Error.err ~loc:ptyp_loc Invalid_mel_int_type)
+        | Rtag ({ txt; _ }, _, _) when allow_no_payload ->
+            let i =
+              match
+                Ast_attributes.iter_process_mel_int_as rtag.prf_attributes
+              with
+              | Some i -> i
+              | None -> i
+            in
+            (i + 1, (txt, i) :: acc)
+        | _ -> Error.err ~loc Invalid_mel_int_type)
       ~init:(0, []) row_fields
   in
-  List.rev acc
+  Melange_ffi.External_arg_spec.Int (List.rev acc)
+
+let process_mel_as tag ~txt ~has_mel_as =
+  let name =
+    match Ast_attributes.iter_process_mel_string_as tag.prf_attributes with
+    | Some name ->
+        has_mel_as := true;
+        name
+    | None -> txt
+  in
+  (txt, name)
 
 (* It also check in-consistency of cases like
    {[ [`a  | `c of int ] ]} *)
-let map_row_fields_into_strings ptyp_loc (row_fields : row_field list) :
-    Melange_ffi.External_arg_spec.attr =
+let map_row_fields_into_strings (row_fields : row_field list) ~loc
+    ~allow_no_payload =
   let has_mel_as = ref false in
   let case, result =
     List.fold_right
       ~f:(fun tag (nullary, acc) ->
         match (nullary, tag.prf_desc) with
         | (`Nothing | `Null), Rtag ({ txt; _ }, true, []) ->
-            let name =
-              match
-                Ast_attributes.iter_process_mel_string_as tag.prf_attributes
-              with
-              | Some name ->
-                  has_mel_as := true;
-                  name
-              | None -> txt
-            in
-            (`Null, (txt, name) :: acc)
+            (`Null, process_mel_as tag ~txt ~has_mel_as :: acc)
+        | `NonNull, Rtag ({ txt; _ }, true, []) when allow_no_payload ->
+            (`Null, process_mel_as tag ~txt ~has_mel_as :: acc)
         | (`Nothing | `NonNull), Rtag ({ txt; _ }, false, [ _ ]) ->
-            let name =
-              match
-                Ast_attributes.iter_process_mel_string_as tag.prf_attributes
-              with
-              | Some name ->
-                  has_mel_as := true;
-                  name
-              | None -> txt
-            in
-            (`NonNull, (txt, name) :: acc)
-        | _ -> Error.err ~loc:ptyp_loc Invalid_mel_string_type)
+            (`NonNull, process_mel_as tag ~txt ~has_mel_as :: acc)
+        | _ -> Error.err ~loc Invalid_mel_string_type)
       row_fields ~init:(`Nothing, [])
   in
   match case with
-  | `Nothing -> Error.err ~loc:ptyp_loc Invalid_mel_string_type
+  | `Nothing -> Error.err ~loc Invalid_mel_string_type
   | `Null | `NonNull -> (
       let has_payload = case = `NonNull in
       let descr = if !has_mel_as then Some result else None in
       match (has_payload, descr) with
       | false, None ->
-          Mel_ast_invariant.warn ~loc:ptyp_loc Redundant_mel_string;
-          Nothing
+          Mel_ast_invariant.warn ~loc Redundant_mel_string;
+          Melange_ffi.External_arg_spec.Nothing
       | false, Some descr -> Poly_var_string { descr }
       | true, _ -> Poly_var { descr })
diff --git a/ppx/ast_polyvar.mli b/ppx/ast_polyvar.mli
@@ -27,8 +27,16 @@ open Import
 val is_enum_polyvar : type_declaration -> row_field list option
 
 val map_row_fields_into_ints :
-  Location.t -> row_field list -> (string * int) list
+  row_field list ->
+  loc:Location.t ->
+  (* allow `Foo [@mel.as "bar"] inside `@mel.unwrap` *)
+  allow_no_payload:bool ->
+  Melange_ffi.External_arg_spec.attr
 (** side effect: it will mark used attributes `mel.as`  *)
 
 val map_row_fields_into_strings :
-  Location.t -> row_field list -> Melange_ffi.External_arg_spec.attr
+  row_field list ->
+  loc:Location.t ->
+  (* allow `Foo [@mel.as "bar"] inside `@mel.unwrap` *)
+  allow_no_payload:bool ->
+  Melange_ffi.External_arg_spec.attr
diff --git a/test/blackbox-tests/ffi-error-debug.t b/test/blackbox-tests/ffi-error-debug.t
@@ -38,19 +38,3 @@
   Error: `[@mel.as ..]' must not be used with an optionally labelled
          polymorphic variant
   [2]
-
-Each [@mel.unwrap] variant constructor requires an argument
-
-  $ cat > x.ml <<EOF
-  > external err :
-  >   ?hi_should_error:([\`a of int | \`b] [@mel.unwrap]) ->
-  >   unit -> unit = "err"
-  > EOF
-  $ melc -ppx melppx x.ml
-  File "x.ml", line 2, characters 20-36:
-  2 |   ?hi_should_error:([`a of int | `b] [@mel.unwrap]) ->
-                          ^^^^^^^^^^^^^^^^
-  Error: Invalid type for `@mel.unwrap'. Type must be an inline variant
-         (closed), and each constructor must have an argument.
-  [2]
-