Evaluate a let before the callee of the call it feeds

changes/let-before-callee.fix.md

@@ -0,0 +1 @@
+- Fix `let x = e; f(x)` running `e` after it evaluates the callee `f`, as in `let x = c.n; next()(x)`.

compiler/core/lam_analysis.ml

@@ -192,6 +192,23 @@ let rec delay_safe ~assignable (lam : Lambda.t) : bool =
   | Lfor _ | Lfor_of _ | Lfor_await_of _ | Lassign _ | Lbreak | Lcontinue ->
     false
 
+let reads_callee (primitive : Lambda.primitive) =
+  match primitive with
+  | Pjs_call {ffi = {kind = Decl_val _ | Decl_new _}} -> true
+  | _ -> false
+
+type motion = Anywhere | Reads | Effects
+
+let motion ~assignable lam : motion =
+  if delay_safe ~assignable lam then Anywhere
+  else if no_side_effects lam then Reads
+  else Effects
+
+let conflict (a : motion) (b : motion) =
+  match (a, b) with
+  | Anywhere, _ | _, Anywhere | Reads, Reads -> false
+  | Effects, (Reads | Effects) | Reads, Effects -> true
+
 (* Size of a lambda that is too big to inline, e.g. one containing a loop *)
 let too_big = 1000
 

compiler/core/lam_analysis.mli

@@ -32,6 +32,22 @@ val delay_safe : assignable:Set_ident.t -> Lambda.t -> bool
     after other code, gives the same value. [assignable] are the variables
     some [Lassign] assigns. *)
 
+val reads_callee : Lambda.primitive -> bool
+(** The call of an external function or class, which reads it from its global
+    or module before the arguments run *)
+
+(** How the evaluation of a term may move with respect to other code *)
+type motion =
+  | Anywhere  (** no side effects, and depends on no state that may change *)
+  | Reads  (** no side effects, but reads state that may change *)
+  | Effects
+
+val motion : assignable:Set_ident.t -> Lambda.t -> motion
+
+val conflict : motion -> motion -> bool
+(** Whether code of the two motions can give other results when it runs in
+    the other order *)
+
 val size_upto : limit:int -> Lambda.t -> int
 (** The size of a lambda for inlining decisions if it is below [limit],
     [limit] otherwise. [limit] must be at most 1000. *)

compiler/core/lam_beta_reduce.ml

@@ -150,7 +150,7 @@ let propagate_beta_reduce (meta : Lam_stats.t) (params : Ident.t list)
         | Lprim {primitive = Psome | Psome_not_nest; args = [v]; _} ->
           Hash_ident.replace meta.ident_tbl param (Normal_optional v)
         | _ -> ());
-        Lam_util.refine_let ~kind:Strict param arg l)
+        Lam_util.refine_let ~assignable ~kind:Strict param arg l)
 
 let propagate_beta_reduce_with_map (meta : Lam_stats.t)
     (map : Lam_var_stats.stats Map_ident.t) params body args =
@@ -201,7 +201,7 @@ let propagate_beta_reduce_with_map (meta : Lam_stats.t)
         | Lprim {primitive = Psome | Psome_not_nest; args = [v]} ->
           Hash_ident.replace meta.ident_tbl param (Normal_optional v)
         | _ -> ());
-        Lam_util.refine_let ~kind:Strict param arg l)
+        Lam_util.refine_let ~assignable ~kind:Strict param arg l)
 
 let no_names_beta_reduce ~assignable params body args =
   match
@@ -212,4 +212,4 @@ let no_names_beta_reduce ~assignable params body args =
     (* Fold right, so that the first parameter's binding is outermost and the
        arguments evaluate in call order *)
     Ext_list.fold_right2 params args body (fun param arg l ->
-        Lam_util.refine_let ~kind:Strict param arg l)
+        Lam_util.refine_let ~assignable ~kind:Strict param arg l)

compiler/core/lam_beta_reduce_util.ml

@@ -38,25 +38,12 @@
    3. an argument with side effects is not put in an operand that may not be
       evaluated, the right one of [&&] and [||] *)
 
-(* How the evaluation of an argument may move with respect to other code *)
-type motion =
-  | Anywhere  (** no side effects, and depends on no state that may change *)
-  | Reads  (** no side effects, but reads state that may change *)
-  | Effects
-
-let motion ~assignable lam : motion =
-  if Lam_analysis.delay_safe ~assignable lam then Anywhere
-  else if Lam_analysis.no_side_effects lam then Reads
-  else Effects
-
-(* Whether code of motions [a] and [b] can give other results when it runs in
-   the other order *)
-let conflict (a : motion) (b : motion) =
-  match (a, b) with
-  | Anywhere, _ | _, Anywhere | Reads, Reads -> false
-  | Effects, (Reads | Effects) | Reads, Effects -> true
-
-type value = {mutable used: bool; lambda: Lambda.t; rank: int; motion: motion}
+type value = {
+  mutable used: bool;
+  lambda: Lambda.t;
+  rank: int;
+  motion: Lam_analysis.motion;
+}
 
 (* Sequence the unused arguments before [body], in parameter order *)
 let seq_unused param_hash params body =
@@ -90,7 +77,7 @@ let simple_beta_reduce ~assignable params body args =
   (* The code the reduced body evaluates, latest first, with its rank in the
      call *)
   let evaluated = ref [] in
-  let evaluate ~conditional rank (motion : motion) =
+  let evaluate ~conditional rank (motion : Lam_analysis.motion) =
     match motion with
     | Effects when conditional -> raise_notrace Not_simple_apply
     | Anywhere | Reads | Effects -> evaluated := (rank, motion) :: !evaluated
@@ -102,7 +89,7 @@ let simple_beta_reduce ~assignable params body args =
       evaluate ~conditional exp.rank exp.motion;
       exp.lambda
     | None ->
-      evaluate ~conditional body_rank (motion ~assignable opt);
+      evaluate ~conditional body_rank (Lam_analysis.motion ~assignable opt);
       opt
   in
   (* Substitutes in the operands from left to right. The operands from
@@ -131,7 +118,7 @@ let simple_beta_reduce ~assignable params body args =
       | [] -> true
       | (rank, motion) :: later ->
         Ext_list.for_all later (fun (rank', motion') ->
-            rank <= rank' || not (conflict motion motion'))
+            rank <= rank' || not (Lam_analysis.conflict motion motion'))
         && in_order later
     in
     if not (in_order (unused @ List.rev !evaluated)) then
@@ -141,7 +128,12 @@ let simple_beta_reduce ~assignable params body args =
     List.iteri
       (fun rank (p, a) ->
         Hash_ident.add param_hash p
-          {lambda = a; used = false; rank; motion = motion ~assignable a})
+          {
+            lambda = a;
+            used = false;
+            rank;
+            motion = Lam_analysis.motion ~assignable a;
+          })
       (List.combine params args)
   in
   match (body : Lambda.t) with

compiler/core/lam_pass_deep_flatten.ml

@@ -100,12 +100,13 @@ let rec eliminate_tuple (id : Ident.t) (lam : Lambda.t) acc =
                     - also for function compilation, flattening should be done first
                     - [compile_group] and [compile] become mutually recursive function
                 *)
-let lambda_of_groups ~(rev_bindings : Lam_group.t list) (result : Lambda.t) :
-    Lambda.t =
+let lambda_of_groups ~assignable ~(rev_bindings : Lam_group.t list)
+    (result : Lambda.t) : Lambda.t =
   Ext_list.fold_left rev_bindings result (fun acc x ->
       match x with
       | Nop l -> Lambda.seq l acc
-      | Single (kind, ident, lam) -> Lam_util.refine_let ~kind ident lam acc
+      | Single (kind, ident, lam) ->
+        Lam_util.refine_let ~assignable ~kind ident lam acc
       | Recursive bindings -> Lambda.letrec bindings acc)
 
 (* TODO:
@@ -155,6 +156,9 @@ let regroups_binding (str : Lambda.let_kind) (id : Ident.t) (arg : Lambda.t) =
     | _ -> false
 
 let deep_flatten (lam : Lambda.t) : Lambda.t =
+  (* See [Lam_pass_remove_alias.simplify_alias] *)
+  let assignable = Lambda_traverse.assigned_variables lam in
+  let lambda_of_groups = lambda_of_groups ~assignable in
   let rec flatten (acc : Lam_group.t list) (lam : Lambda.t) :
       Lambda.t * Lam_group.t list =
     match lam with
@@ -225,7 +229,8 @@ let deep_flatten (lam : Lambda.t) : Lambda.t =
   and aux (lam : Lambda.t) : Lambda.t =
     match lam with
     | Llet (str, id, arg, body) when not (regroups_binding str id arg) ->
-      Lam_util.refine_let ~original:lam ~kind:str id (aux arg) (aux body)
+      Lam_util.refine_let ~original:lam ~assignable ~kind:str id (aux arg)
+        (aux body)
     | Llet _ ->
       let res, groups = flatten [] lam in
       lambda_of_groups res ~rev_bindings:groups

compiler/core/lam_pass_lets_dce.ml

@@ -93,7 +93,9 @@ let lets_helper ~assignable (count_var : Ident.t -> Lam_pass_count.used_info)
           Hash_ident.add string_table v s;
           (* we need move [simplif lbody] later, since adding Hash does have side effect *)
           Lambda.let_ Alias v l1 (simplif lbody)
-        | _ -> Lam_util.refine_let ~original:lam ~kind v l1 (simplif lbody)
+        | _ ->
+          Lam_util.refine_let ~original:lam ~assignable ~kind v l1
+            (simplif lbody)
         (* TODO: check if it is correct rollback to [StrictOpt]? *))
     | Llet (((Strict | Variable) as kind), v, l1, l2) -> (
       if not (used v) then
@@ -107,7 +109,8 @@ let lets_helper ~assignable (count_var : Ident.t -> Lam_pass_count.used_info)
         | Strict, Lconst (Const_string s) ->
           Hash_ident.add string_table v s;
           Lambda.let_ Alias v l1 (simplif l2)
-        | _ -> Lam_util.refine_let ~original:lam ~kind v l1 (simplif l2))
+        | _ ->
+          Lam_util.refine_let ~original:lam ~assignable ~kind v l1 (simplif l2))
     | Lapply
         {ap_func = Lfunction ({params; body} as lfunction); ap_args = args; _}
       when Ext_list.same_length params args

compiler/core/lam_pass_sroa.ml

@@ -109,6 +109,9 @@ let replace ~block ~info ~initializers body =
     else
       let fields = fields_for_block block info field_count in
       let body = rewrite block fields uses body in
+      (* The fields that are written, of this block and of the ones replaced
+         in [body] *)
+      let assignable = lazy (Lambda_traverse.assigned_variables body) in
       let rec bind_fields index initializers body =
         match initializers with
         | [] -> body
@@ -121,7 +124,8 @@ let replace ~block ~info ~initializers body =
              still needs a mutable scalar binding. *)
           if not use.read then discard_value init body
           else if not use.written then
-            Lam_util.refine_let ~kind:Strict fields.(index) init body
+            Lam_util.refine_let ~assignable:(Lazy.force assignable) ~kind:Strict
+              fields.(index) init body
           else Lambda.let_ Variable fields.(index) init body
       in
       Some (bind_fields 0 initializers body)

compiler/core/lam_util.ml

@@ -41,6 +41,7 @@
    - (Return)      [let[k] x = e in x]               ⟶ e
    - (Prim)        [let[k] x = e in prim p x]        ⟶ prim p e   (p ≠ makeblock)
    - (Call)        [let[k] x = e in f x]             ⟶ f e        (x not captured in f)
+                     both when e may run after what the call evaluates first
    - (Alias)       [let[k] x = e in body]            ⟶ let[Alias] x = e in body
                      when k ∈ {Strict, StrictOpt} and SafeAlias(e)
    - (Strict λ)    [let[Strict] x = fn in body]      ⟶ let[StrictOpt] x = fn in body
@@ -50,15 +51,21 @@
    evaluation strategy downstream, so we keep its predicate intentionally
    syntactic and narrow. *)
 (* [original] is the binding this one was taken apart from, if any. When
-   nothing is refined it is handed back untouched rather than rebuilt. *)
-let refine_let ?original ~kind param (arg : Lambda.t) (l : Lambda.t) : Lambda.t
-    =
+   nothing is refined it is handed back untouched rather than rebuilt.
+   [assignable] are the variables some [Lassign] assigns. *)
+let refine_let ?original ~assignable ~kind param (arg : Lambda.t) (l : Lambda.t)
+    : Lambda.t =
   (* A JSX element whose only argument is its tag keeps the tag bound: a
      function bound by [let] gets the binding's name, which React shows *)
   let keeps_binding = function
     | Lambda.Pmakeblock _ | Pjsx _ -> true
     | _ -> false
   in
+  (* JavaScript evaluates the callee of a call before the arguments, so
+     [f(e)] runs [e] after [f], while the binding ran it before *)
+  let may_follow (callee : Lam_analysis.motion) =
+    not (Lam_analysis.conflict callee (Lam_analysis.motion ~assignable arg))
+  in
   (* SafeAlias is the predicate that justifies the (Alias) rewrite
        let[k] x = e in body  -->  let[Alias] x = e in body
      for strict bindings.  Turning a binding into [Alias] authorises JS codegen
@@ -101,17 +108,23 @@ let refine_let ?original ~kind param (arg : Lambda.t) (l : Lambda.t) : Lambda.t
          cannot refer back to `x`, and we avoid generating a redundant local. *)
     arg
   | _, _, Lprim {primitive; args = [Lvar w]; loc; _}
-    when Ident.same w param && not (keeps_binding primitive) ->
+    when Ident.same w param
+         && (not (keeps_binding primitive))
+         && ((not (Lam_analysis.reads_callee primitive)) || may_follow Reads) ->
     (* When we immediately feed the binding into a primitive, like
          `{ let x = value; Array.length(x) }`, we inline the primitive call
-         with `value`. This only happens for primitives that are pure and do not
-         allocate new blocks, so evaluation order and side effects stay the same. *)
+         with `value`. Only the call of an external function does something
+         before its argument, which is to read the function, so evaluation
+         order and side effects stay the same. *)
     Lambda.prim ~primitive ~args:[arg] loc
   | _, _, Lapply {ap_func = fn; ap_args = [Lvar w]; ap_loc}
-    when Ident.same w param && not (Lam_hit.hit_variable param fn) ->
+    when Ident.same w param
+         && (not (Lam_hit.hit_variable param fn))
+         && may_follow (Lam_analysis.motion ~assignable fn) ->
     (* For a function call such as `{ let x = value; someFn(x) }`, we can
-         rewrite to `someFn(value)` as long as the callee does not capture `x`.
-         This removes the temporary binding while preserving the call semantics. *)
+         rewrite to `someFn(value)` as long as the callee does not capture `x`
+         and `value` may run after it: not `{ let x = c.n; next()(x) }`, where
+         `next()` may assign `c.n`. *)
     Lambda.apply fn [arg] ap_loc
   | (Strict | StrictOpt), arg, _ when is_safe_to_alias arg ->
     (* `Strict` and `StrictOpt` bindings both evaluate the RHS immediately

compiler/core/lam_util.mli

@@ -54,12 +54,14 @@ val alias_ident_or_global :
 
 val refine_let :
   ?original:Lambda.t ->
+  assignable:Set_ident.t ->
   kind:Lambda.let_kind ->
   Ident.t ->
   Lambda.t ->
   Lambda.t ->
   Lambda.t
 (** [original] is the binding being rebuilt, when there is one. It is returned
-    unchanged if no refinement applies. *)
+    unchanged if no refinement applies. [assignable] are the variables some
+    [Lassign] assigns. *)
 
 val not_function : Lambda.t -> bool

compiler/ml/lambda_traverse.ml

@@ -245,11 +245,16 @@ let free_variables l =
 
 let assigned_variables l =
   let assigned = ref Set_ident.empty in
+  (* [visit] is the one callback of the walk, where [iter] would allocate a
+     closure for each node *)
   let rec walk l =
     (match l with
     | Lassign (id, _) -> assigned := Set_ident.add id !assigned
     | _ -> ());
-    iter walk l
+    ignore (shallow_exists visit l)
+  and visit l =
+    walk l;
+    false
   in
   walk l;
   !assigned

tests/tests/src/inline_arg_order_cross_module_lib.mjs

@@ -5,7 +5,12 @@ function minus3(_x, a, b) {
   return b - a | 0;
 }
 
+function callNext(next, x) {
+  return next()(x);
+}
+
 export {
   minus3,
+  callNext,
 }
 /* No side effect */

tests/tests/src/inline_arg_order_cross_module_lib.res

@@ -2,3 +2,5 @@
 
 // Inlined by inline_arg_order_cross_module_test.res
 let minus3 = (_x: int, a: int, b: int) => b - a
+
+let callNext = (next: unit => int => int, x: int) => next()(x)

tests/tests/src/inline_arg_order_cross_module_test.mjs

@@ -24,13 +24,25 @@ function bumpThenRead(x) {
   return s - a | 0;
 }
 
+function readThenCall() {
+  let s = 0;
+  let next = () => {
+    s = s + 1 | 0;
+    return y => y;
+  };
+  let x = s;
+  return next()(x);
+}
+
 Mocha.describe("Inline_arg_order_cross_module_test", () => {
-  Mocha.test("a read of a local ref before an argument that assigns it", () => Test_utils.eq("File \"inline_arg_order_cross_module_test.res\", line 30, characters 7-14", 1, readThenBump(0)));
-  Mocha.test("a read of a local ref after an argument that assigns it", () => Test_utils.eq("File \"inline_arg_order_cross_module_test.res\", line 34, characters 7-14", 9, bumpThenRead(0)));
+  Mocha.test("a read of a local ref before an argument that assigns it", () => Test_utils.eq("File \"inline_arg_order_cross_module_test.res\", line 38, characters 7-14", 1, readThenBump(0)));
+  Mocha.test("a read of a local ref after an argument that assigns it", () => Test_utils.eq("File \"inline_arg_order_cross_module_test.res\", line 42, characters 7-14", 9, bumpThenRead(0)));
+  Mocha.test("a read of a local ref before a callee that assigns it", () => Test_utils.eq("File \"inline_arg_order_cross_module_test.res\", line 46, characters 7-14", 0, readThenCall()));
 });
 
 export {
   readThenBump,
   bumpThenRead,
+  readThenCall,
 }
 /*  Not a pure module */

tests/tests/src/inline_arg_order_cross_module_test.res

@@ -25,6 +25,14 @@ let bumpThenRead = x => {
   Inline_arg_order_cross_module_lib.minus3(x, bumpS(), s.contents)
 }
 
+let readThenCall = () => {
+  let s = ref(0)
+  Inline_arg_order_cross_module_lib.callNext(() => {
+    s := s.contents + 1
+    y => y
+  }, s.contents)
+}
+
 describe(__MODULE__, () => {
   test("a read of a local ref before an argument that assigns it", () => {
     eq(__LOC__, 1, readThenBump(0))
@@ -33,4 +41,8 @@ describe(__MODULE__, () => {
   test("a read of a local ref after an argument that assigns it", () => {
     eq(__LOC__, 9, bumpThenRead(0))
   })
+
+  test("a read of a local ref before a callee that assigns it", () => {
+    eq(__LOC__, 0, readThenCall())
+  })
 })