Evaluate partial application arguments once

changes/partial-application-args.fix.md

@@ -0,0 +1 @@
+- Evaluate the callee and arguments of a partial application `f(a, ...)` once, so an `await` in them no longer ends up in a non-async function.

compiler/core/lam_compile_const.ml

@@ -24,18 +24,13 @@
 
 module E = Js_exp_make
 
-(** return [val < 0] if not nested [Some (Some (Some None))]*)
-let rec is_some_none_aux (x : Lambda.structured_constant) acc =
-  match x with
-  | Const_some v -> is_some_none_aux v (acc + 1)
-  | Const_module_alias | Const_js_undefined _ -> acc
-  | _ -> -1
-
 let rec nested_some_none n none =
   if n = 0 then none else nested_some_none (n - 1) (E.optional_block none)
 
 let rec translate_some (x : Lambda.structured_constant) : J.expression =
-  let depth = is_some_none_aux x 0 in
+  (* Keep in sync with [Lambda.const_is_allocating]: a nested none is a new
+     optional block. *)
+  let depth = Lambda.const_some_none_depth x in
   if depth < 0 then E.optional_not_nest_block (translate x)
   else
     nested_some_none depth

compiler/core/lam_pass_lets_dce.ml

@@ -32,7 +32,7 @@ let lets_helper (count_var : Ident.t -> Lam_pass_count.used_info) lam : Lambda.t
       match (count_var v, l1) with
       | {times = 0; _}, _ -> simplif l2
       | {times = 1; captured = false}, _
-      | {times = 1; captured = true}, (Lconst _ | Lvar _)
+      | {times = 1; captured = true}, Lvar _
       | ( _,
           ( Lconst
               ( Const_int _ | Const_assertfalse | Const_constructor _
@@ -47,6 +47,12 @@ let lets_helper (count_var : Ident.t -> Lam_pass_count.used_info) lam : Lambda.t
         ->
         Hash_ident.add subst v (simplif l1);
         simplif l2
+      (* A captured constant is moved into the closure only when it allocates
+         nothing: a block would be a new one on each call. *)
+      | {times = 1; captured = true}, Lconst c
+        when not (Lambda.const_is_allocating c) ->
+        Hash_ident.add subst v (simplif l1);
+        simplif l2
       | _, Lconst (Const_string s) ->
         (* only "" added for later inlining *)
         Hash_ident.add string_table v s;

compiler/ml/lambda.ml

@@ -580,10 +580,22 @@ let str_of_field_info (x : field_dbg_info) : string option =
 let is_immutable_block (info : tag_info) =
   mutable_flag_of_tag_info info = Immutable
 
-(* A constant that has to be built at run time rather than shared. *)
+(* The number of [Some] around a none, as in [Some (Some None)], or [-1] if
+   [c] is not a nested none. *)
+let rec const_some_none_depth (c : structured_constant) : int =
+  match c with
+  | Const_some v ->
+    let depth = const_some_none_depth v in
+    if depth < 0 then depth else depth + 1
+  | Const_module_alias | Const_js_undefined _ -> 0
+  | _ -> -1
+
+(* A constant that has to be built at run time rather than shared. A [Some]
+   around a (nested) none is emitted as a new optional block, see
+   [Lam_compile_const.translate_some]. *)
 let rec const_is_allocating (c : structured_constant) : bool =
   match c with
-  | Const_some t -> const_is_allocating t
+  | Const_some t -> const_some_none_depth t >= 0 || const_is_allocating t
   | Const_block _ -> true
   | Const_js_null | Const_js_undefined _ | Const_js_true | Const_js_false
   | Const_int _ | Const_assertfalse | Const_constructor _ | Const_char _

compiler/ml/lambda.mli

@@ -435,6 +435,8 @@ val str_of_field_info : field_dbg_info -> string option
 
 val is_immutable_block : tag_info -> bool
 
+val const_some_none_depth : structured_constant -> int
+
 val const_is_allocating : structured_constant -> bool
 
 (* Constructors. [t] is private, so every term outside this module is

compiler/ml/translcore.ml

@@ -1550,17 +1550,37 @@ and transl_cases cases = List.map transl_case cases
 
 and transl_apply ?(uncurried_partial_application = None) lam sargs loc =
   let lapply ap_func ap_args = apply ap_func ap_args loc in
+  (* A partial application creates a function. Its callee and arguments are
+     bound outside that function, so they are evaluated once, where the
+     application is, and an await in them stays in the enclosing function. *)
+  let rec is_module_path (lam : Lambda.t) =
+    match lam with
+    | Lvar _ | Lglobal_module _ -> true
+    | Lprim {primitive = Pfield (_, Fld_module _); args = [m]} ->
+      is_module_path m
+    | _ -> false
+  in
+  let protect defs name lam =
+    match lam with
+    (* A variable, a module's value or a constant that allocates nothing reads
+       the same on each call, so it stays in place and keeps the function a
+       known one that can be inlined. A closure or a constant block does not:
+       each evaluation is a new one. Nor does an external value, which can be
+       reassigned or be a getter. *)
+    | Lconst c when not (Lambda.const_is_allocating c) -> lam
+    | _ when is_module_path lam -> lam
+    | _ ->
+      let id = Ident.create name in
+      defs := (id, lam) :: !defs;
+      var id
+  in
+  let bind_protected defs body =
+    List.fold_left (fun body (id, lam) -> let_ Strict id lam body) body !defs
+  in
   let rec build_apply lam args = function
     | (None, optional) :: l ->
       let defs = ref [] in
-      let protect name lam =
-        match lam with
-        | Lvar _ | Lconst _ -> lam
-        | _ ->
-          let id = Ident.create name in
-          defs := (id, lam) :: !defs;
-          var id
-      in
+      let protect = protect defs in
       let args, args' =
         if List.for_all (fun (_, opt) -> opt) args then ([], args)
         else (args, [])
@@ -1577,17 +1597,19 @@ and transl_apply ?(uncurried_partial_application = None) lam sargs loc =
           function_ ~loc ~attr:default_function_attribute ~params:[id_arg]
             ~body:lam
       in
-      List.fold_left (fun body (id, lam) -> let_ Strict id lam body) body !defs
+      bind_protected defs body
     | (Some arg, optional) :: l -> build_apply lam ((arg, optional) :: args) l
     | [] -> lapply lam (List.rev_map fst args)
   in
   match uncurried_partial_application with
   | Some arity ->
     let extra_arity = arity - List.length sargs in
     let none_ids = ref [] in
+    let defs = ref [] in
+    let lam = protect defs "func" lam in
     let args =
       Ext_list.filter_map sargs (function
-        | _, Some e -> Some (transl_exp e)
+        | _, Some e -> Some (protect defs "arg" (transl_exp e))
         | _, None ->
           let id_arg = Ident.create "none" in
           none_ids := id_arg :: !none_ids;
@@ -1599,9 +1621,10 @@ and transl_apply ?(uncurried_partial_application = None) lam sargs loc =
     let extra_args = Ext_list.map extra_ids (fun id -> var id) in
     let ap_args = args @ extra_args in
     let l0 = apply lam ap_args loc in
-    function_ ~loc ~attr:default_function_attribute
-      ~params:(List.rev_append !none_ids extra_ids)
-      ~body:l0
+    bind_protected defs
+      (function_ ~loc ~attr:default_function_attribute
+         ~params:(List.rev_append !none_ids extra_ids)
+         ~body:l0)
   | _ ->
     (build_apply lam []
        (List.map

compiler/ml/typecore.ml

@@ -2002,11 +2002,12 @@ let rec is_nonexpansive exp =
   | Texp_apply {funct = {exp_desc}; args; _} when is_component_identity exp_desc
     ->
     List.for_all is_nonexpansive_opt (List.map snd args)
-  | Texp_apply {partial = true; _} ->
-    (* ReScript partial applications (`foo(args, ...)`) lower to wrapper
-       functions in codegen, so creating the partial itself is nonexpansive
-       like an explicit lambda. *)
-    true
+  | Texp_apply {partial = true; funct; args} ->
+    (* A partial application (`foo(args, ...)`) evaluates its callee and
+       supplied arguments once, when it is created, and the resulting function
+       shares them, so it is nonexpansive only when they are. *)
+    is_nonexpansive funct
+    && List.for_all is_nonexpansive_opt (List.map snd args)
   | Texp_apply {funct = e; args = (_, None) :: el} ->
     is_nonexpansive e && List.for_all is_nonexpansive_opt (List.map snd el)
   | Texp_match (e, cases, [], _) ->

tests/belt_tests/src/bs_array_test.mjs

@@ -50,13 +50,13 @@ Mocha.describe("Bs_array_test", () => {
         1
       ], 2);
     });
-    let f = extra => Belt_Array.getExn([
+    let arg = [
       0,
       1
-    ], extra);
+    ];
     Test_utils.ok("File \"bs_array_test.res\", line 28, characters 6-13", Primitive_object.equal([
-      f(0),
-      f(1)
+      Belt_Array.getExn(arg, 0),
+      Belt_Array.getExn(arg, 1)
     ], [
       0,
       1

tests/build_tests/super_errors/expected/partial_application_value_restriction.res.expected

@@ -0,0 +1,16 @@
+
+  We've found a bug for you!
+  /.../fixtures/partial_application_value_restriction.res:6:30-33
+
+  4 │ let h = first([], ...)
+  5 │ let ints: array<int> = h(0)
+  6 │ let strings: array<string> = h(0)
+  7 │ 
+
+  This has type: array<int>
+  But it's expected to have type: array<string>
+
+  The incompatible parts:
+    int vs string
+
+  You can convert int to string with Int.toString.

tests/build_tests/super_errors/fixtures/partial_application_value_restriction.res

@@ -0,0 +1,6 @@
+let first = (x, _y) => x
+
+// `[]` is evaluated once, so every call of `h` returns the same array
+let h = first([], ...)
+let ints: array<int> = h(0)
+let strings: array<string> = h(0)

tests/ounit_tests/ounit_lambda_constant_tests.ml

@@ -331,4 +331,19 @@ let suites =
              with
              | Lambda.Lconst (Lambda.Const_int length) -> 3l =~ length
              | _ -> OUnit.assert_failure "expected a folded string length" );
+           ( "a nested None is an allocating constant" >:: fun _ ->
+             let none = Lambda.Const_js_undefined {is_unit = false} in
+             let some x = Lambda.Const_some x in
+             let int = Lambda.Const_int 1l in
+             List.iter
+               (fun (constant, expected) ->
+                 OUnit.assert_equal ~printer:string_of_bool expected
+                   (Lambda.const_is_allocating constant))
+               [
+                 (some none, true);
+                 (some (some none), true);
+                 (some int, false);
+                 (some (some int), false);
+                 (none, false);
+               ] );
          ]

tests/ounit_tests/ounit_partial_application_tests.ml

@@ -0,0 +1,111 @@
+let ( >:: ), ( >::: ) = OUnit.(( >:: ), ( >::: ))
+
+(* The raw Lambda bound to [name] in [source]. *)
+let binding_lambda ~name source =
+  let parsed =
+    Res_driver.parse_implementation_from_source
+      ~display_filename:"PartialApplication.res" ~source
+  in
+  OUnit.assert_bool "source parses" (not parsed.invalid);
+  let typed, _, _ =
+    Typemod.type_structure (Env.initial_safe_string ()) parsed.parsetree
+  in
+  let implementation =
+    Translmod.transl_implementation "PartialApplication"
+      (typed, Typedtree.Tcoerce_none)
+  in
+  let found = ref None in
+  let rec find (lam : Lambda.t) =
+    (match lam with
+    | Llet (_, id, arg, _) when Ident.name id = name -> found := Some arg
+    | _ -> ());
+    Lambda_traverse.iter find lam
+  in
+  find implementation.lambda;
+  match !found with
+  | Some lam -> lam
+  | None -> OUnit.assert_failure ("no binding for " ^ name)
+
+let unexpected lam =
+  OUnit.assert_failure ("unexpected Lambda: " ^ Printlambda.lambda_to_string lam)
+
+let is_call_of name (lam : Lambda.t) =
+  match lam with
+  | Lapply {ap_func = Lvar f} -> Ident.name f = name
+  | _ -> false
+
+let suites =
+  __FILE__
+  >::: [
+         ( "callee and arguments are bound outside the function, callee first"
+         >:: fun _ ->
+           let lam =
+             binding_lambda ~name:"h"
+               "let add = (a, _b, _c) => a\n\
+                let getAdd = () => add\n\
+                let next = () => 1\n\
+                let h = getAdd()(next(), ...)\n"
+           in
+           match lam with
+           | Llet
+               ( Strict,
+                 func,
+                 callee,
+                 Llet
+                   ( Strict,
+                     arg,
+                     argument,
+                     Lfunction
+                       {
+                         params = [extra1; extra2];
+                         body =
+                           Lapply
+                             {
+                               ap_func = Lvar f;
+                               ap_args = [Lvar a; Lvar e1; Lvar e2];
+                             };
+                       } ) )
+             when is_call_of "getAdd" callee && is_call_of "next" argument ->
+             OUnit.assert_bool "the function calls the bound callee"
+               (Ident.same f func);
+             OUnit.assert_bool "the function passes the bound argument"
+               (Ident.same a arg);
+             OUnit.assert_bool "the remaining parameters are passed on"
+               (Ident.same e1 extra1 && Ident.same e2 extra2)
+           | lam -> unexpected lam );
+         ( "variables and constants stay in the function" >:: fun _ ->
+           let lam =
+             binding_lambda ~name:"h"
+               "let add = (a, _b, _c) => a\nlet x = 1\nlet h = add(x, 2, ...)\n"
+           in
+           match lam with
+           | Lfunction
+               {
+                 body =
+                   Lapply
+                     {
+                       ap_func = Lvar f;
+                       ap_args = [Lvar x; Lconst (Const_int 2l); Lvar _];
+                     };
+               }
+             when Ident.name f = "add" && Ident.name x = "x" ->
+             ()
+           | lam -> unexpected lam );
+         ( "an allocating constant is bound outside the function" >:: fun _ ->
+           let lam =
+             binding_lambda ~name:"h"
+               "let first = (x, _y) => x\nlet h = first((1, 2), ...)\n"
+           in
+           match lam with
+           | Llet
+               ( Strict,
+                 arg,
+                 Lconst (Const_block _),
+                 Lfunction
+                   {
+                     body = Lapply {ap_func = Lvar f; ap_args = [Lvar a; Lvar _]};
+                   } )
+             when Ident.name f = "first" && Ident.same a arg ->
+             ()
+           | lam -> unexpected lam );
+       ]

tests/ounit_tests/ounit_tests_main.ml

@@ -15,6 +15,7 @@ let suites =
       Ounit_rec_check_tests.suites;
       Ounit_lambda_constant_tests.suites;
       Ounit_lambda_traverse_tests.suites;
+      Ounit_partial_application_tests.suites;
       Ounit_deep_flatten_tests.suites;
       Ounit_exits_tests.suites;
       Ounit_lam_size_tests.suites;

tests/tests/src/PartialApplicationNoRuntimeCurry.mjs

@@ -18,18 +18,20 @@ function add5(extra) {
   return 5 + extra | 0;
 }
 
+let arg = hook;
+
 function addHookPartial(extra) {
-  addHook(hook, extra);
+  addHook(arg, extra);
 }
 
-addHook(hook, _x => {});
+addHookPartial(_x => {});
 
-addHook(hook, _x => {});
+addHookPartial(_x => {});
 
 export {
   f,
   add$1 as add,
   add5,
   addHookPartial,
 }
-/*  Not a pure module */
+/* arg Not a pure module */

tests/tests/src/earger_curry_test.mjs

@@ -72,7 +72,9 @@ function add5(a0, a1, a2, a3, a4) {
 }
 
 function f(x) {
-  return (extra, extra$1) => add5(x, (v.contents = v.contents + 1 | 0, 1), (v.contents = v.contents + 1 | 0, 2), extra, extra$1);
+  v.contents = v.contents + 1 | 0;
+  v.contents = v.contents + 1 | 0;
+  return (extra, extra$1) => add5(x, 1, 2, extra, extra$1);
 }
 
 function g(x) {