From fca94c47c6d5cd72e8ed9c0877691ee6d561194c Mon Sep 17 00:00:00 2001
From: Stephen Dolan <sdolan@janestreet.com>
Date: Mon, 15 Nov 2021 13:39:32 +0000
Subject: [PATCH] Register allocations for Omitted parameter closures (#47)

---
 lambda/printlambda.ml                 | 5 +++--
 testsuite/tests/typing-local/alloc.ml | 3 +--
 typing/typecore.ml                    | 1 +
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/lambda/printlambda.ml b/lambda/printlambda.ml
index 20fac2321de..1446782ae1c 100644
--- a/lambda/printlambda.ml
+++ b/lambda/printlambda.ml
@@ -525,7 +525,7 @@ let rec lam ppf = function
         apply_tailcall_attribute ap.ap_tailcall
         apply_inlined_attribute ap.ap_inlined
         apply_specialised_attribute ap.ap_specialised
-  | Lfunction{kind; params; return; body; attr; ret_mode} ->
+  | Lfunction{kind; params; return; body; attr; mode; ret_mode} ->
       let pr_params ppf params =
         match kind with
         | Curried {nlocal} ->
@@ -543,7 +543,8 @@ let rec lam ppf = function
                 value_kind ppf k)
               params;
             fprintf ppf ")" in
-      fprintf ppf "@[<2>(function%a@ %a%a%a)@]" pr_params params
+      fprintf ppf "@[<2>(function%s%a@ %a%a%a)@]"
+        (alloc_kind mode) pr_params params
         function_attribute attr return_kind (ret_mode, return) lam body
   | (Llet _ | Lregion(Llet _)) as expr ->
       let kind = function
diff --git a/testsuite/tests/typing-local/alloc.ml b/testsuite/tests/typing-local/alloc.ml
index 472d245bf98..d69bad817fd 100644
--- a/testsuite/tests/typing-local/alloc.ml
+++ b/testsuite/tests/typing-local/alloc.ml
@@ -181,8 +181,7 @@ let currylocal1 n =
 let currylocal2 n =
   ignore_local (Sys.opaque_identity local_arg_fn ~a:n); ()
 let currylocal3 n =
-  (* FIXME broken, see Translcore.build_apply 
-  ignore_local (local_arg_fn ~b:n); *)
+  ignore_local (local_arg_fn ~b:n);
   ()
 
 let partprim1 n =
diff --git a/typing/typecore.ml b/typing/typecore.ml
index 5c1aaf56fae..f7b882199e8 100644
--- a/typing/typecore.ml
+++ b/typing/typecore.ml
@@ -2601,6 +2601,7 @@ let type_omitted_parameters expected_mode env ty_ret mode_ret args =
              let closed_args = new_closed_args @ closed_args in
              let open_args = [] in
              let mode_closure = Alloc_mode.join (mode_fun :: closed_args) in
+             register_allocation_mode mode_closure;
              let arg = Omitted { mode_closure; mode_arg; mode_ret } in
              let args = (lbl, arg) :: args in
              (ty_ret, mode_closure, open_args, closed_args, args))