Add support for labeled tuple expressions

Signed-off-by: Nathan Rebours <[email protected]>

Author: Nathan Rebours

astlib/encoding_504.ml

@@ -1,49 +1,45 @@
 module Ext_name = struct
   let ptyp_labeled_tuple = "ppxlib.migration.ptyp_labeled_tuple_504"
+  let pexp_labeled_tuple = "ppxlib.migration.pexp_labeled_tuple_504"
 end
 
+let invalid_encoding ~loc name =
+  Location.raise_errorf ~loc "Invalid %s encoding" name
+
 module type AST = sig
   type payload
   type core_type
   type core_type_desc
+  type expression
+  type expression_desc
 
   module Construct : sig
     val ptyp_extension_desc : string Location.loc -> payload -> core_type_desc
     val ptyp_tuple : loc:Location.t -> core_type list -> core_type
     val ptyp_var : loc:Location.t -> string -> core_type
     val ptyp_any : loc:Location.t -> core_type
     val ptyp : core_type -> payload
+    val pexp_extension_desc : string Location.loc -> payload -> expression_desc
+    val pexp_tuple : loc:Location.t -> expression list -> expression
+
+    val pexp_variant :
+      loc:Location.t -> string -> expression option -> expression
+
+    val pstr_eval : loc:Location.t -> expression -> payload
   end
 
   module Destruct : sig
     val ptyp : payload -> core_type option
     val ptyp_tuple : core_type -> core_type list option
     val ptyp_var : core_type -> string option
     val ptyp_any : core_type -> unit option
+    val pstr_eval : payload -> expression option
+    val pexp_tuple : expression -> expression list option
+    val pexp_variant : expression -> (string * expression option) option
   end
 end
 
-module type S = sig
-  type payload
-  type core_type
-  type core_type_desc
-
-  val encode_ptyp_labeled_tuple :
-    loc:Location.t -> (string option * core_type) list -> core_type_desc
-
-  val decode_ptyp_labeled_tuple :
-    loc:Location.t -> payload -> (string option * core_type) list
-end
-
-module Make (X : AST) :
-  S
-    with type core_type = X.core_type
-     and type core_type_desc = X.core_type_desc
-     and type payload = X.payload = struct
-  type payload = X.payload
-  type core_type = X.core_type
-  type core_type_desc = X.core_type_desc
-
+module Make (X : AST) = struct
   let encode_ptyp_labeled_tuple ~loc args =
     let payload =
       let l =
@@ -80,9 +76,54 @@ module Make (X : AST) :
     in
     match res with
     | Some res -> res
-    | None ->
-        Location.raise_errorf ~loc "Invalid %s encoding"
-          Ext_name.ptyp_labeled_tuple
+    | None -> invalid_encoding ~loc Ext_name.ptyp_labeled_tuple
+
+  let encode_pexp_labeled_tuple ~loc args =
+    let payload =
+      let l =
+        List.map
+          (fun (label_opt, expr) ->
+            let label =
+              match label_opt with
+              | None -> X.Construct.pexp_variant ~loc "None" None
+              | Some s ->
+                  let string_as_variant =
+                    X.Construct.pexp_variant ~loc s None
+                  in
+                  X.Construct.pexp_variant ~loc "Some" (Some string_as_variant)
+            in
+            X.Construct.pexp_tuple ~loc [ label; expr ])
+          args
+      in
+      X.Construct.pexp_tuple ~loc l
+    in
+    X.Construct.pexp_extension_desc
+      { txt = Ext_name.pexp_labeled_tuple; loc }
+      (X.Construct.pstr_eval ~loc payload)
+
+  let decode_pexp_labeled_tuple ~loc payload =
+    let open Stdlib0.Option.Op in
+    let res =
+      let* exp = X.Destruct.pstr_eval payload in
+      let* exp_list = X.Destruct.pexp_tuple exp in
+      Stdlib0.Option.List.map exp_list ~f:(fun exp ->
+          let* exp_pair = X.Destruct.pexp_tuple exp in
+          match exp_pair with
+          | [ label; exp ] -> (
+              let* opt_variant = X.Destruct.pexp_variant label in
+              match opt_variant with
+              | "None", None -> Some (None, exp)
+              | "Some", Some exp' -> (
+                  let* label_variant = X.Destruct.pexp_variant exp' in
+                  match label_variant with
+                  | s, None -> Some (Some s, exp)
+                  | _, _ -> None)
+              | _ -> None)
+          | _ -> None)
+    in
+    match res with
+    | Some res -> res
+    | None -> invalid_encoding ~loc Ext_name.pexp_labeled_tuple
 end
 
 module Ast_503 = struct
@@ -92,11 +133,22 @@ module Ast_503 = struct
     let core_type ~loc ptyp_desc =
       { ptyp_desc; ptyp_loc = loc; ptyp_attributes = []; ptyp_loc_stack = [] }
 
+    let expression ~loc pexp_desc =
+      { pexp_desc; pexp_loc = loc; pexp_attributes = []; pexp_loc_stack = [] }
+
     let ptyp_extension_desc name payload = Ptyp_extension (name, payload)
     let ptyp_tuple ~loc typs = core_type ~loc (Ptyp_tuple typs)
     let ptyp_var ~loc s = core_type ~loc (Ptyp_var s)
     let ptyp_any ~loc = core_type ~loc Ptyp_any
     let ptyp typ = PTyp typ
+    let pexp_extension_desc name payload = Pexp_extension (name, payload)
+    let pexp_tuple ~loc l = expression ~loc (Pexp_tuple l)
+
+    let pexp_variant ~loc v exp_opt =
+      expression ~loc (Pexp_variant (v, exp_opt))
+
+    let pstr_eval ~loc expr =
+      PStr [ { pstr_desc = Pstr_eval (expr, []); pstr_loc = loc } ]
   end
 
   module Destruct = struct
@@ -111,6 +163,18 @@ module Ast_503 = struct
       | _ -> None
 
     let ptyp_any = function { ptyp_desc = Ptyp_any; _ } -> Some () | _ -> None
+
+    let pstr_eval = function
+      | PStr [ { pstr_desc = Pstr_eval (expr, []); _ } ] -> Some expr
+      | _ -> None
+
+    let pexp_tuple = function
+      | { pexp_desc = Pexp_tuple l; _ } -> Some l
+      | _ -> None
+
+    let pexp_variant = function
+      | { pexp_desc = Pexp_variant (s, e); _ } -> Some (s, e)
+      | _ -> None
   end
 end
 
@@ -121,11 +185,22 @@ module Ast_502 = struct
     let core_type ~loc ptyp_desc =
       { ptyp_desc; ptyp_loc = loc; ptyp_attributes = []; ptyp_loc_stack = [] }
 
+    let expression ~loc pexp_desc =
+      { pexp_desc; pexp_loc = loc; pexp_attributes = []; pexp_loc_stack = [] }
+
     let ptyp_extension_desc name payload = Ptyp_extension (name, payload)
     let ptyp_tuple ~loc typs = core_type ~loc (Ptyp_tuple typs)
     let ptyp_var ~loc s = core_type ~loc (Ptyp_var s)
     let ptyp_any ~loc = core_type ~loc Ptyp_any
     let ptyp typ = PTyp typ
+    let pexp_extension_desc name payload = Pexp_extension (name, payload)
+    let pexp_tuple ~loc l = expression ~loc (Pexp_tuple l)
+
+    let pexp_variant ~loc v exp_opt =
+      expression ~loc (Pexp_variant (v, exp_opt))
+
+    let pstr_eval ~loc expr =
+      PStr [ { pstr_desc = Pstr_eval (expr, []); pstr_loc = loc } ]
   end
 
   module Destruct = struct
@@ -140,6 +215,18 @@ module Ast_502 = struct
       | _ -> None
 
     let ptyp_any = function { ptyp_desc = Ptyp_any; _ } -> Some () | _ -> None
+
+    let pstr_eval = function
+      | PStr [ { pstr_desc = Pstr_eval (expr, []); _ } ] -> Some expr
+      | _ -> None
+
+    let pexp_tuple = function
+      | { pexp_desc = Pexp_tuple l; _ } -> Some l
+      | _ -> None
+
+    let pexp_variant = function
+      | { pexp_desc = Pexp_variant (s, e); _ } -> Some (s, e)
+      | _ -> None
   end
 end
 

astlib/encoding_504.mli

@@ -1,5 +1,6 @@
 module Ext_name : sig
   val ptyp_labeled_tuple : string
+  val pexp_labeled_tuple : string
 end
 
 module To_503 : sig
@@ -10,6 +11,12 @@ module To_503 : sig
 
   val decode_ptyp_labeled_tuple :
     loc:Location.t -> payload -> (string option * core_type) list
+
+  val encode_pexp_labeled_tuple :
+    loc:Location.t -> (string option * expression) list -> expression_desc
+
+  val decode_pexp_labeled_tuple :
+    loc:Location.t -> payload -> (string option * expression) list
 end
 
 module To_502 : sig
@@ -20,4 +27,10 @@ module To_502 : sig
 
   val decode_ptyp_labeled_tuple :
     loc:Location.t -> payload -> (string option * core_type) list
+
+  val encode_pexp_labeled_tuple :
+    loc:Location.t -> (string option * expression) list -> expression_desc
+
+  val decode_pexp_labeled_tuple :
+    loc:Location.t -> payload -> (string option * expression) list
 end

astlib/migrate_503_504.ml

@@ -179,6 +179,11 @@ and copy_expression_desc_with_loc :
       Ast_504.Parsetree.Pexp_open (copy_open_declaration x0, copy_expression x1)
   | Ast_503.Parsetree.Pexp_letop x0 ->
       Ast_504.Parsetree.Pexp_letop (copy_letop x0)
+  | Ast_503.Parsetree.Pexp_extension ({ txt; loc }, payload)
+    when String.equal txt Encoding_504.Ext_name.pexp_labeled_tuple ->
+      let xs = Encoding_504.To_503.decode_pexp_labeled_tuple ~loc payload in
+      Ast_504.Parsetree.Pexp_tuple
+        (List.map (fun (lbl, exp) -> (lbl, copy_expression exp)) xs)
   | Ast_503.Parsetree.Pexp_extension x0 ->
       Ast_504.Parsetree.Pexp_extension (copy_extension x0)
   | Ast_503.Parsetree.Pexp_unreachable -> Ast_504.Parsetree.Pexp_unreachable

astlib/migrate_504_503.ml

@@ -55,14 +55,15 @@ and copy_expression :
        Ast_504.Parsetree.pexp_loc_stack;
        Ast_504.Parsetree.pexp_attributes;
      } ->
+  let loc = copy_location pexp_loc in
   {
-    Ast_503.Parsetree.pexp_desc = copy_expression_desc pexp_desc;
-    Ast_503.Parsetree.pexp_loc = copy_location pexp_loc;
+    Ast_503.Parsetree.pexp_desc = copy_expression_desc ~loc pexp_desc;
+    Ast_503.Parsetree.pexp_loc = loc;
     Ast_503.Parsetree.pexp_loc_stack = copy_location_stack pexp_loc_stack;
     Ast_503.Parsetree.pexp_attributes = copy_attributes pexp_attributes;
   }
 
-and copy_expression_desc :
+and copy_expression_desc ~loc :
     Ast_504.Parsetree.expression_desc -> Ast_503.Parsetree.expression_desc =
   function
   | Ast_504.Parsetree.Pexp_ident x0 ->
@@ -90,15 +91,12 @@ and copy_expression_desc :
   | Ast_504.Parsetree.Pexp_try (x0, x1) ->
       Ast_503.Parsetree.Pexp_try (copy_expression x0, List.map copy_case x1)
   | Ast_504.Parsetree.Pexp_tuple x0 ->
-      let args =
-        List.map
-          (function
-            | None, arg -> arg
-            | Some _l, (arg : Ast_504.Parsetree.expression) ->
-                migration_error arg.pexp_loc "labelled tuples")
-          x0
+      let exps =
+        List.map (fun (label, exp) -> (label, copy_expression exp)) x0
       in
-      Ast_503.Parsetree.Pexp_tuple (List.map copy_expression args)
+      if List.exists (function Some _, _ -> true | _ -> false) exps then
+        Encoding_504.To_503.encode_pexp_labeled_tuple ~loc exps
+      else Ast_503.Parsetree.Pexp_tuple (List.map snd exps)
   | Ast_504.Parsetree.Pexp_construct (x0, x1) ->
       Ast_503.Parsetree.Pexp_construct
         (copy_loc copy_Longident_t x0, Option.map copy_expression x1)

src/ast_builder.ml

@@ -270,6 +270,12 @@ module Default = struct
     in
     { ptyp_desc; ptyp_loc = loc; ptyp_attributes = []; ptyp_loc_stack = [] }
 
+  let pexp_labeled_tuple ~loc l =
+    let pexp_desc =
+      Astlib__.Encoding_504.To_502.encode_pexp_labeled_tuple ~loc l
+    in
+    { pexp_desc; pexp_loc = loc; pexp_attributes = []; pexp_loc_stack = [] }
+
   let pexp_tuple_opt ~loc l =
     match l with [] -> None | _ :: _ -> Some (pexp_tuple ~loc l)
 
@@ -562,6 +568,7 @@ end) : S = struct
   let ppat_tuple l = Default.ppat_tuple ~loc l
   let ptyp_tuple l = Default.ptyp_tuple ~loc l
   let ptyp_labeled_tuple l = Default.ptyp_labeled_tuple ~loc l
+  let pexp_labeled_tuple l = Default.pexp_labeled_tuple ~loc l
   let pexp_tuple_opt l = Default.pexp_tuple_opt ~loc l
   let ppat_tuple_opt l = Default.ppat_tuple_opt ~loc l
   let ptyp_poly vars ty = Default.ptyp_poly ~loc vars ty

src/ast_builder_intf.ml

@@ -170,6 +170,10 @@ module type Additional_helpers = sig
   val ptyp_labeled_tuple :
     ((string option * core_type) list -> core_type) with_loc
   (** Returns an encoded labeled tuple type as introduced in OCaml 5.4. *)
+
+  val pexp_labeled_tuple :
+    ((string option * expression) list -> expression) with_loc
+  (** Returns an encoded labeled tuple expression as introduced in OCaml 5.4. *)
 end
 
 module type Located = sig

src/ast_pattern.ml

@@ -292,3 +292,20 @@ let ptyp_labeled_tuple (T f0) =
           let k = f0 ctx loc x0 k in
           k
       | _ -> fail loc "labeled tuple")
+
+let pexp_labeled_tuple (T f0) =
+  T
+    (fun ctx _loc x k ->
+      let loc = x.pexp_loc in
+      let x = x.pexp_desc in
+      match x with
+      | Pexp_extension ({ txt; _ }, payload)
+        when String.equal txt Astlib__.Encoding_504.Ext_name.pexp_labeled_tuple
+        ->
+          let x0 =
+            Astlib__.Encoding_504.To_502.decode_pexp_labeled_tuple ~loc payload
+          in
+          ctx.matched <- ctx.matched + 1;
+          let k = f0 ctx loc x0 k in
+          k
+      | _ -> fail loc "labeled tuple")

src/ast_pattern.mli

@@ -227,3 +227,11 @@ val ptyp_labeled_tuple :
 
     It will fail on a regular tuple type and as a consequence, if it matches, at
     least one type in the tuple is guaranteed to be labeled. *)
+
+val pexp_labeled_tuple :
+  ((string option * expression) list, 'a, 'b) t -> (expression, 'a, 'b) t
+(** Match over an encoded OCaml 5.4 labeled tuple expression.
+
+    It will fail on a regular tuple expression and as a consequence, if it
+    matches, at least one expression in the tuple is guaranteed to be labeled.
+*)