show missed optimization / currying issue on OCaml 5.2+

Changes.md

@@ -114,6 +114,9 @@ Unreleased
     `[@mel.send]` with `[@mel.this]`, including when used with `@mel.variadic`.
 - ppx: deprecate `[@mel.send.pipe]`
   ([#1321](https://github.com/melange-re/melange/pull/1321))
+- core: fix missed optimization on OCaml versions 5.2 and above, caused by
+  [ocaml/ocaml#12236](https://github.com/ocaml/ocaml/pull/12236) generating
+  multiple function nodes for `fun a -> fun b -> ...` in the Lambda IR.
 
 4.0.1 2024-06-07
 ---------------

jscomp/core/lam_convert.cppo.ml

@@ -536,6 +536,20 @@ let nat_of_string_exn =
     let acc = int_of_string_aux s 0 0 (String.length s) in
     if acc < 0 then invalid_arg s else acc
 
+let convert_lfunction_params_and_body params body =
+  let just_params = List.map ~f:fst params in
+  let new_map, body =
+    rename_optional_parameters Ident.Map.empty just_params body
+  in
+  let params =
+    if Ident.Map.is_empty new_map then just_params
+    else
+      List.map
+        ~f:(fun x -> Ident.Map.find_default new_map x x)
+        just_params
+  in
+  params, body
+
 let convert (exports : Ident.Set.t) (lam : Lambda.lambda) :
     Lam.t * Lam_module_ident.Hash_set.t =
   let alias_tbl = Ident.Hash.create 64 in
@@ -751,20 +765,19 @@ let convert (exports : Ident.Set.t) (lam : Lambda.lambda) :
             ap_inlined;
             ap_status = App_na;
           }
-    | Lfunction { params; body; attr; _ } ->
-        let just_params = List.map ~f:fst params in
-        let body = convert_aux ~dynamic_import body in
-        let new_map, body =
-          rename_optional_parameters Ident.Map.empty just_params body
-        in
-        let params =
-          if Ident.Map.is_empty new_map then just_params
-          else
-            List.map
-              ~f:(fun x -> Ident.Map.find_default new_map x x)
-              just_params
-        in
-        Lam.function_ ~attr ~arity:(List.length params) ~params ~body
+    | Lfunction { params; body = l; attr = attr1; _ } -> (
+        (* because of ocaml/ocaml#12236, `fun a -> fun b -> ..` becomes 2
+           `Lfunction` nodes in the AST on OCaml 5.2 and up. *)
+        match convert_aux ~dynamic_import l with
+        | Lfunction { arity = arity'; params = params'; body; attr = attr2 }
+          when List.length params + List.length params' <= Lambda.max_arity() ->
+
+          let params, body = convert_lfunction_params_and_body params body in
+          let arity = (List.length params) + arity' in
+          Lam.function_ ~arity ~params:(params @ params') ~body ~attr:attr2
+        | body ->
+          let params, body = convert_lfunction_params_and_body params body in
+          Lam.function_ ~attr:attr1 ~arity:(List.length params) ~params ~body)
     | Llet (kind, _value_kind, id, e, body) (*FIXME*) ->
         convert_let kind id e body
     | Lmutlet (_value_kind, id, e, body) (*FIXME*) -> convert_mutlet id e body

test/blackbox-tests/syntactic-function-arity.t

@@ -0,0 +1,58 @@
+
+  $ . ./setup.sh
+  $ cat > x.ml <<EOF
+  > let sum = function a -> function b -> a + b
+  > let sum2 = fun a -> fun b -> a + b
+  > let sum3 = fun (a, b) -> fun c -> a + b + c
+  > let sum4 = fun a -> fun b -> fun c -> a + b + c
+  > EOF
+
+  $ melc -drawlambda x.ml
+  (setglobal X!
+    (let
+      (sum/267 =
+         (function a/269[int] (function b/270[int] : int (+ a/269 b/270)))
+       sum2/271 =
+         (function a/273[int] (function b/274[int] : int (+ a/273 b/274)))
+       sum3/275 =
+         (function param/280
+           (let
+             (b/278 =a (field_imm :[]/1 param/280)
+              a/277 =a (field_imm :[]/0 param/280))
+             (function c/279[int] : int (+ (+ a/277 b/278) c/279))))
+       sum4/281 =
+         (function a/283[int]
+           (function b/284[int]
+             (function c/285[int] : int (+ (+ a/283 b/284) c/285)))))
+      (makeblock 0module/exports sum/267 sum2/271 sum3/275 sum4/281)))
+  // Generated by Melange
+  'use strict';
+  
+  
+  function sum(a, b) {
+    return a + b | 0;
+  }
+  
+  function sum2(a, b) {
+    return a + b | 0;
+  }
+  
+  function sum3(param) {
+    const b = param[1];
+    const a = param[0];
+    return function (c) {
+      return (a + b | 0) + c | 0;
+    };
+  }
+  
+  function sum4(a, b, c) {
+    return (a + b | 0) + c | 0;
+  }
+  
+  module.exports = {
+    sum,
+    sum2,
+    sum3,
+    sum4,
+  }
+  /* No side effect */