Stop measuring lambda size once the inlining limit is reached

CHANGELOG.md

@@ -50,6 +50,7 @@
 
 - Compile dict spreads to object spread syntax instead of `Object.assign`. https://github.com/rescript-lang/rescript/pull/8747
 - Make compilation slightly faster and stop rewatch from copying files twice for modules with an interface. https://github.com/rescript-lang/rescript/pull/8769
+- Speed up inlining decisions for large functions. https://github.com/rescript-lang/rescript/pull/8771
 - Represent explicit expression braces as `Pexp_braces` in parsetree v1 and format `else` branches consistently with `if` branches. https://github.com/rescript-lang/rescript/pull/8678
 - Omit redundant braces around multi-statement switch case bodies when formatting. https://github.com/rescript-lang/rescript/pull/8677
 - Avoid running `rescript-schema-ppx` and `sury-ppx` on source files without an `@schema` annotation. https://github.com/rescript-lang/rescript/pull/8662

compiler/core/lam_analysis.ml

@@ -141,34 +141,57 @@ let rec no_side_effects (lam : Lambda.t) : bool =
   | Lapply _ -> false
 (* we need purity analysis .. *)
 
-(*
-     Estimate the size of lambda for better inlining
-     threshold is 1000 - so that we
-*)
-exception Too_big_to_inline
-
-let really_big () = raise_notrace Too_big_to_inline
+(* Size of a lambda that is too big to inline, e.g. one containing a loop *)
+let too_big = 1000
 
-(* let big_lambda = 1000 *)
+exception Limit_reached
 
-let rec size (lam : Lambda.t) =
-  try
+(*
+   Estimate the size of a lambda for inlining. Callers only compare it with
+   small thresholds, so stop counting once [limit] is reached: the result is
+   the size if it is below [limit], and [limit] otherwise.
+*)
+let size_upto ~limit (lam : Lambda.t) =
+  assert (limit <= too_big);
+  let total = ref 0 in
+  let add n =
+    total := !total + n;
+    if !total >= limit then raise_notrace Limit_reached
+  in
+  (* Constants count their leaves, and stop at the limit too: a large constant
+     list or record should not be walked completely. *)
+  let rec constant (c : Lambda.structured_constant) =
+    match c with
+    | Const_int _ | Const_assertfalse | Const_constructor _ | Const_char _
+    | Const_float _ | Const_bigint _ | Const_polyvar _ | Const_js_null
+    | Const_js_undefined _ | Const_module_alias | Const_js_true | Const_js_false
+    | Const_string _ ->
+      add 1
+    | Const_some s -> constant s
+    | Const_block (_, str) -> List.iter constant str
+  in
+  let rec size (lam : Lambda.t) =
     match lam with
-    | Lvar _ -> 1
-    | Lconst c -> size_constant c
-    | Llet (_, _, l1, l2) -> 1 + size l1 + size l2
-    | Lletrec _ -> really_big ()
+    | Lvar _ -> add 1
+    | Lconst c -> constant c
+    | Llet (_, _, l1, l2) ->
+      add 1;
+      size l1;
+      size l2
+    | Lletrec _ -> add too_big
     | Lprim
         {
           primitive = Pfield (_, Fld_module _);
           args = [(Lglobal_module _ | Lvar _)];
           _;
         } ->
-      1
+      add 1
     | Lprim {primitive = Praise | Pis_not_none; args = [l]; _} -> size l
-    | Lglobal_module _ -> 1
-    | Lprim {primitive = Praw_js_code _} -> really_big ()
-    | Lprim {args = ll; _} -> size_lams 1 ll
+    | Lglobal_module _ -> add 1
+    | Lprim {primitive = Praw_js_code _} -> add too_big
+    | Lprim {args = ll; _} ->
+      add 1;
+      List.iter size ll
     (* complicated
            1. inline this function
            2. ...
@@ -177,39 +200,37 @@ let rec size (lam : Lambda.t) =
        {var $$let=Make(funarg);
          return [0, $$let[5],... $$let[16]]}
     *)
-    | Lapply {ap_func; ap_args; _} -> size_lams (size ap_func) ap_args
-    (* | Lfunction(_, params, l) -> really_big () *)
+    | Lapply {ap_func; ap_args; _} ->
+      size ap_func;
+      List.iter size ap_args
     | Lfunction {body} -> size body
-    | Lswitch _ -> really_big ()
-    | Lstringswitch (_, _, _) -> really_big ()
+    | Lswitch _ -> add too_big
+    | Lstringswitch (_, _, _) -> add too_big
     | Lstaticraise (_i, ls) ->
-      Ext_list.fold_left ls 1 (fun acc x -> size x + acc)
-    | Lstaticcatch _ -> really_big ()
-    | Ltrywith _ -> really_big ()
-    | Lifthenelse (l1, l2, l3) -> 1 + size l1 + size l2 + size l3
-    | Lsequence (l1, l2) -> size l1 + size l2
-    | Lbreak | Lcontinue -> 1
-    | Lwhile _ -> really_big ()
-    | Lfor _ -> really_big ()
-    | Lfor_of _ | Lfor_await_of _ -> really_big ()
-    | Lassign (_, v) -> 1 + size v
+      add 1;
+      List.iter size ls
+    | Lstaticcatch _ -> add too_big
+    | Ltrywith _ -> add too_big
+    | Lifthenelse (l1, l2, l3) ->
+      add 1;
+      size l1;
+      size l2;
+      size l3
+    | Lsequence (l1, l2) ->
+      size l1;
+      size l2
+    | Lbreak | Lcontinue -> add 1
+    | Lwhile _ -> add too_big
+    | Lfor _ -> add too_big
+    | Lfor_of _ | Lfor_await_of _ -> add too_big
+    | Lassign (_, v) ->
+      add 1;
+      size v
     (* This is side effectful,  be careful *)
-  with Too_big_to_inline -> 1000
-
-and size_constant x =
-  match x with
-  | Const_int _ | Const_assertfalse | Const_constructor _ | Const_char _
-  | Const_float _ | Const_bigint _ | Const_polyvar _ | Const_js_null
-  | Const_js_undefined _ | Const_module_alias | Const_js_true | Const_js_false
-    ->
-    1
-  | Const_string _ -> 1
-  | Const_some s -> size_constant s
-  | Const_block (_, str) ->
-    Ext_list.fold_left str 0 (fun acc x -> acc + size_constant x)
-
-and size_lams acc (lams : Lambda.t list) =
-  Ext_list.fold_left lams acc (fun acc l -> acc + size l)
+  in
+  match size lam with
+  | () -> !total
+  | exception Limit_reached -> limit
 
 let args_all_const (args : Lambda.t list) =
   Ext_list.for_all args (fun x ->
@@ -241,13 +262,15 @@ let destruct_pattern (body : Lambda.t) params args =
   | Lswitch (Lvar v, switch) -> (
     match aux v params args with
     | Some (Lambda.Lconst _ as lam) ->
-      size (Lambda.switch lam switch) < small_inline_size
+      size_upto ~limit:small_inline_size (Lambda.switch lam switch)
+      < small_inline_size
     | Some _ | None -> false)
   | Lifthenelse (Lvar v, then_, else_) -> (
     (* -FIXME *)
     match aux v params args with
     | Some (Lconst _ as lam) ->
-      size (Lambda.if_ lam then_ else_) < small_inline_size
+      size_upto ~limit:small_inline_size (Lambda.if_ lam then_ else_)
+      < small_inline_size
     | Some _ | None -> false)
   | _ -> false
 
@@ -263,7 +286,7 @@ let ok_to_inline_fun_when_app (m : Lambda.lfunction) (args : Lambda.t list) =
   | Default_inline -> (
     match m with
     | {body; params} ->
-      let s = size body in
+      let s = size_upto ~limit:10 body in
       s < small_inline_size
       || destruct_pattern body params args
       || (args_all_const args && s < 10 && no_side_effects body))

compiler/core/lam_analysis.mli

@@ -27,7 +27,9 @@
 val no_side_effects : Lambda.t -> bool
 (** No side effect, but it might depend on data store *)
 
-val size : Lambda.t -> int
+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. *)
 
 val lfunction_can_be_inlined : Lambda.lfunction -> bool
 

compiler/core/lam_pass_exits.ml

@@ -173,7 +173,9 @@ let subst_helper (subst : subst_tbl) (query : int -> int) (lam : Lambda.t) :
         let ok_to_inline =
           i >= 0 && no_bounded_variables l2'
           &&
-          let lam_size = Lam_analysis.size l2' in
+          let lam_size =
+            Lam_analysis.size_upto ~limit:Lam_analysis.exit_inline_size l2'
+          in
           (i_occur <= 2 && lam_size < Lam_analysis.exit_inline_size)
           || lam_size < 5
         in

compiler/core/lam_stats_export.ml

@@ -77,7 +77,10 @@ let values_of_export (meta : Lam_stats.t) (export_map : Lambda.t Map_ident.t) :
                 optlam
               else None
             | _ ->
-              let lam_size = Lam_analysis.size lambda in
+              let lam_size =
+                Lam_analysis.size_upto ~limit:Lam_analysis.small_inline_size
+                  lambda
+              in
               (* TODO:
                  1. global need re-assocate when do the beta reduction
                  2. [lambda_exports] is not precise

tests/ounit_tests/ounit_lam_size_tests.ml

@@ -0,0 +1,62 @@
+open OUnit
+
+let ( =~ ) = assert_equal ~printer:string_of_int
+
+let x = Ident.create "x"
+
+(* [Lstaticraise] counts 1 plus its arguments and a variable counts 1, so this
+   lambda has exactly size [n + 1]. *)
+let raise_of_vars n = Lambda.staticraise 0 (List.init n (fun _ -> Lambda.var x))
+
+(* A loop is too big to inline, whatever its size. *)
+let loop = Lambda.while_ (Lambda.var x) (Lambda.var x)
+
+let size_upto = Lam_analysis.size_upto
+
+let int i = Lambda.const_int i
+
+(* A constant list of [n] elements, nested [n] levels deep *)
+let constant_list n =
+  let rec go i acc =
+    if i < 0 then acc
+    else go (i - 1) (Lambda.Const_block (Blk_tuple, [int i; acc]))
+  in
+  go (n - 1) (int 0)
+
+let suites =
+  __FILE__
+  >::: [
+         ( "returns the size below the limit and the limit otherwise"
+         >:: fun _ ->
+           (* the limits used by the inlining heuristics *)
+           List.iter
+             (fun limit ->
+               for n = 0 to limit + 3 do
+                 size_upto ~limit (raise_of_vars n) =~ min (n + 1) limit
+               done)
+             [Lam_analysis.small_inline_size; Lam_analysis.exit_inline_size; 10]
+         );
+         ( "adds up nested lambdas" >:: fun _ ->
+           (* 1 + (1 + 2) + (1 + 1) = 6 *)
+           let lam = Lambda.let_ Strict x (raise_of_vars 2) (raise_of_vars 1) in
+           size_upto ~limit:10 lam =~ 6;
+           size_upto ~limit:7 lam =~ 6;
+           size_upto ~limit:6 lam =~ 6;
+           size_upto ~limit:5 lam =~ 5 );
+         ( "counts a construct that is too big to inline as the limit"
+         >:: fun _ ->
+           size_upto ~limit:10 loop =~ 10;
+           size_upto ~limit:5
+             (Lambda.staticraise 0 [Lambda.var x; loop; Lambda.var x])
+           =~ 5;
+           (* the old unbounded size was 1000 for these *)
+           size_upto ~limit:1000 loop =~ 1000 );
+         ( "counts the leaves of constants" >:: fun _ ->
+           let tuple = Lambda.Const_block (Blk_tuple, [int 1; int 2; int 3]) in
+           size_upto ~limit:10 (Lambda.const tuple) =~ 3;
+           size_upto ~limit:10 (Lambda.const (Lambda.Const_some (int 1))) =~ 1;
+           size_upto ~limit:2 (Lambda.const tuple) =~ 2 );
+         ( "stops inside large constants" >:: fun _ ->
+           (* would need a million nested calls without the cutoff *)
+           size_upto ~limit:5 (Lambda.const (constant_list 1_000_000)) =~ 5 );
+       ]

tests/ounit_tests/ounit_tests_main.ml

@@ -22,6 +22,7 @@ let suites =
       Ounit_lambda_traverse_tests.suites;
       Ounit_deep_flatten_tests.suites;
       Ounit_exits_tests.suites;
+      Ounit_lam_size_tests.suites;
       Ounit_sroa_tests.suites;
       Ounit_ast_mapper0_tests.suites;
       Ounit_constructor_arguments_tests.suites;