Compute contains_subfunctions on the fly (#492)

middle_end/flambda/from_lambda/closure_conversion.ml

@@ -695,7 +695,6 @@ let close_one_function acc ~external_env ~by_closure_id decl
   let my_closure_id = closure_id in
   let my_depth = Depth_variable.create "my_depth" in
   let our_let_rec_ident = Function_decl.let_rec_ident decl in
-  let contains_closures = Function_decl.contains_closures decl in
   let compilation_unit = Compilation_unit.get_current_exn () in
   let code_id =
     Code_id.create ~name:(Closure_id.to_string closure_id) compilation_unit
@@ -773,6 +772,7 @@ let close_one_function acc ~external_env ~by_closure_id decl
         Kinded_parameter.create var (LC.value_kind kind))
       param_vars
   in
+  let acc = Acc.with_seen_a_function acc false in
   let acc, body =
     try body acc closure_env
     with Misc.Fatal_error -> begin
@@ -789,6 +789,7 @@ let close_one_function acc ~external_env ~by_closure_id decl
       raise Misc.Fatal_error
     end
   in
+  let contains_subfunctions = Acc.seen_a_function acc in
   let my_closure' = Simple.var my_closure in
   let acc, body =
     (* CR mshinwell: These Select_closure operations should maybe be inserted
@@ -853,7 +854,7 @@ let close_one_function acc ~external_env ~by_closure_id decl
        lifted during Closure (but will prevent inlining) but will likely have
        been lifted by our other check in [Inlining_cost] (thus preventing us
        seeing they were originally there). *)
-    if contains_closures
+    if contains_subfunctions
       && !Clflags.Flambda.Expert.fallback_inlining_heuristic
     then Never_inline
     else LC.inline_attribute (Function_decl.inline decl)
@@ -890,8 +891,9 @@ let close_one_function acc ~external_env ~by_closure_id decl
       ~cost_metrics
       ~inlining_arguments:Inlining_arguments.unknown
   in
-  Acc.add_code ~code_id ~code acc,
-  Closure_id.Map.add my_closure_id fun_decl by_closure_id
+  let acc = Acc.add_code ~code_id ~code acc in
+  let acc = Acc.with_seen_a_function acc true in
+  acc, Closure_id.Map.add my_closure_id fun_decl by_closure_id
 
 let close_functions acc external_env function_declarations =
   let compilation_unit = Compilation_unit.get_current_exn () in

middle_end/flambda/from_lambda/closure_conversion_aux.ml

@@ -192,19 +192,25 @@ module Acc = struct
     code : Flambda.Code.t Code_id.Map.t;
     free_names_of_current_function : Name_occurrences.t;
     cost_metrics : Flambda.Cost_metrics.t;
+    seen_a_function : bool;
   }
 
   let cost_metrics t = t.cost_metrics
   let increment_metrics metrics t =
     { t with cost_metrics = Flambda.Cost_metrics.(+) t.cost_metrics metrics }
   let with_cost_metrics cost_metrics t = { t with cost_metrics }
 
+  let seen_a_function t = t.seen_a_function
+  let with_seen_a_function t seen_a_function =
+    { t with seen_a_function; }
+
   let empty = {
     declared_symbols = [];
     shareable_constants = Flambda.Static_const.Map.empty;
     code = Code_id.Map.empty;
     free_names_of_current_function = Name_occurrences.empty;
     cost_metrics = Flambda.Cost_metrics.zero;
+    seen_a_function = false;
   }
 
   let declared_symbols t = t.declared_symbols
@@ -265,12 +271,11 @@ module Function_decls = struct
       loc : Lambda.scoped_location;
       stub : bool;
       recursive : Recursive.t;
-      contains_closures : bool;
     }
 
     let create ~let_rec_ident ~closure_id ~kind ~params ~return
         ~return_continuation ~exn_continuation ~body ~attr
-        ~loc ~free_idents_of_body ~stub recursive ~contains_closures =
+        ~loc ~free_idents_of_body ~stub recursive =
       let let_rec_ident =
         match let_rec_ident with
         | None -> Ident.create_local "unnamed_function"
@@ -289,7 +294,6 @@ module Function_decls = struct
         loc;
         stub;
         recursive;
-        contains_closures;
       }
 
     let let_rec_ident t = t.let_rec_ident
@@ -307,7 +311,6 @@ module Function_decls = struct
     let stub t = t.attr.stub
     let loc t = t.loc
     let recursive t = t.recursive
-    let contains_closures t = t.contains_closures
   end
 
   type t = {

middle_end/flambda/from_lambda/closure_conversion_aux.mli

@@ -122,6 +122,9 @@ module Acc : sig
   val code : t -> Flambda.Code.t Code_id.Map.t
   val free_names_of_current_function : t -> Name_occurrences.t
 
+  val seen_a_function : t -> bool
+  val with_seen_a_function : t -> bool -> t
+
   val add_declared_symbol
      : symbol:Symbol.t
     -> constant:Flambda.Static_const.t
@@ -175,7 +178,6 @@ module Function_decls : sig
       -> free_idents_of_body:Ident.Set.t
       -> stub:bool
       -> Recursive.t
-      -> contains_closures:bool
       -> t
 
 
@@ -193,7 +195,6 @@ module Function_decls : sig
     val stub : t -> bool
     val loc : t -> Lambda.scoped_location
     val recursive : t -> Recursive.t
-    val contains_closures : t -> bool
 
     (* Like [all_free_idents], but for just one function. *)
     val free_idents : t -> Ident.Set.t

middle_end/flambda/from_lambda/lambda_to_flambda.ml

@@ -287,49 +287,6 @@ let name_for_function (func : Lambda.lfunction) =
   | Loc_known { loc; _ } ->
     Format.asprintf "anon-fn[%a]" Location.print_compact loc
 
-let contains_functions (lam : Lambda.lambda) =
-  let rec contains_functions_tail (lam : Lambda.lambda) k =
-  match lam with
-    | Lvar _  | Lconst _ -> k ()
-    | Lfunction _ | Lletrec _ -> true
-    | Lassign (_, lam) | Levent (lam, _) | Lifused (_, lam)
-      -> contains_functions_tail lam k
-    | Lapply { ap_func; ap_args; _}
-      -> contains_functions_list (ap_func::ap_args) k
-    | Llet (_, _, _, lam1, lam2) | Lstaticcatch (lam1, _, lam2)
-    | Ltrywith (lam1, _, lam2) | Lsequence (lam1, lam2)
-    | Lwhile (lam1, lam2)
-      -> contains_functions_list [lam1; lam2] k
-    | Lifthenelse (lam1, lam2, lam3) | Lfor (_, lam1, lam2, _, lam3)
-      -> contains_functions_list [lam1; lam2; lam3] k
-    | Lprim (_, lams, _) | Lstaticraise (_, lams)
-      -> contains_functions_list lams k
-    | Lsend (_, lam1, lam2, lams, _)
-      -> contains_functions_list (lam1::lam2::lams) k
-    | Lswitch (lam1, { sw_consts; sw_blocks; sw_failaction; _ }, _) ->
-      let lams1 = List.map snd sw_consts in
-      let lams2 = List.map snd sw_blocks in
-      let lams = match sw_failaction with
-        | None -> lam1::lams1 @ lams2
-        | Some lam2 -> lam1::lam2::lams1 @ lams2
-      in
-      contains_functions_list lams k
-    | Lstringswitch (lam1, branches, failaction, _) ->
-      let lams = List.map snd branches in
-      let lams = match failaction with
-        | None -> lam1::lams
-        | Some lam2 -> lam1::lam2::lams
-      in
-      contains_functions_list lams k
-  and contains_functions_list lams k =
-    match lams with
-    | [] -> k ()
-    | lam::lams ->
-        contains_functions_tail lam
-          (fun () -> contains_functions_list lams k)
-  in
-  contains_functions_tail lam (fun () -> false)
-
 let extra_args_for_exn_continuation env exn_handler =
   let more_extra_args =
     Env.extra_args_for_continuation_with_kinds env exn_handler
@@ -1420,20 +1377,14 @@ and cps_function env ~fid ~stub
     Closure_id.wrap (Compilation_unit.get_current_exn ())
       (Variable.create_with_same_name_as_ident fid)
   in
-  let contains_closures =
-    (* only useful with this flag *)
-    if !Clflags.Flambda.Expert.fallback_inlining_heuristic
-    then contains_functions body
-    else false
-  in
   let body = fun acc ccenv ->
     cps_tail acc new_env ccenv body body_cont body_exn_cont
   in
   Function_decl.create ~let_rec_ident:(Some fid)
     ~closure_id ~kind ~params ~return ~return_continuation:body_cont
     ~exn_continuation ~body
     ~attr ~loc ~free_idents_of_body ~stub
-    recursive ~contains_closures
+    recursive
 
 and cps_switch acc env ccenv (switch : L.lambda_switch) ~scrutinee
       (k : Continuation.t) (k_exn : Continuation.t)