Convert generated nested functions into multi-argument functions

lib/genletrec.ml

@@ -13,12 +13,23 @@ struct
   open Asttypes
   open Parsetree
 
+  (* Convert nested functions fun x -> fun y -> e
+     into multi-argument functions fun x y -> e *)
+  let rec squish_arity : expression -> expression = function
+    | { pexp_desc =
+          Pexp_function (ps, None,
+                         Pfunction_body ({pexp_desc =
+                                            Pexp_function (ps', None, b)} as d')) } as d ->
+       squish_arity { d with pexp_desc = Pexp_function (ps @ ps', None, b);
+                             pexp_attributes = d.pexp_attributes @ d'.pexp_attributes }
+    | e -> e
+
   let letrec : (string * expression) list -> expression -> expression =
     fun bindings e ->
       let binding (x, e) =
         { pvb_pat = Ast_helper.Pat.var (Location.mknoloc x);
           pvb_constraint = None;
-          pvb_expr = e;
+          pvb_expr = squish_arity e;
           pvb_attributes = [];
           pvb_loc = Location.none } in
       match bindings with

lib_test/arity.ml

@@ -0,0 +1,34 @@
+(*
+ * Copyright (c) 2025 Jeremy Yallop.
+ *
+ * This file is distributed under the terms of the MIT License.
+ * See the file LICENSE for details.
+ *)
+
+open Letrec
+
+module Sym = struct
+  type _ t =
+    Z : int t
+  | S : 'a t -> (int -> 'a) t
+  let rec eql : type a b. a t -> b t -> (a, b) eql option =
+    fun x y ->
+    match x, y with
+    | Z , Z -> Some Refl
+    | S x, S y -> begin match eql x y with
+                  | Some Refl -> Some Refl
+                  | None -> None
+                  end
+    | _ -> None
+end
+
+module R = Letrec.Make(Sym)
+
+let rec mkrhs : type a. (int code list -> int code) -> a Sym.t -> a code =
+  fun k -> function
+  | Z -> k []
+  | S n -> .< fun x -> .~(mkrhs (fun l -> k (.<x>. :: l)) n)>.
+
+let sumn n =
+  let rhs _ = mkrhs (List.fold_left (fun x y -> .< .~x + .~y >.) .<0>.) in
+  R.letrec {R.rhs} (fun {R.resolve} -> resolve n)

lib_test/tests.ml

@@ -31,4 +31,7 @@ let () = begin
 
   pr "(* Custom equality for indexes *)@\n";
   pr "%a@." print_code Custom_eq.evenp_oddp;
+
+  pr "(* Generating functions of varying arity *)@\n";
+  pr "%a@." print_code Arity.(sumn Sym.(S (S (S Z))));
 end