Transform tail-recursive functions into recursive continuations

Author: Ekdohibs

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1089,8 +1089,28 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
   let params = Function_decl.params decl in
   let return = Function_decl.return decl in
   let return_continuation = Function_decl.return_continuation decl in
-  let recursive = Function_decl.recursive decl in
+  let acc, exn_continuation =
+    close_exn_continuation acc external_env
+      (Function_decl.exn_continuation decl)
+  in
+  assert (
+    match Exn_continuation.extra_args exn_continuation with
+    | [] -> true
+    | _ :: _ -> false);
   let my_closure = Variable.create "my_closure" in
+  let recursive = Function_decl.recursive decl in
+  (* Mark function available for loopify only if it is a single recursive
+     function *)
+  let is_single_recursive_function =
+    match recursive, Function_decls.to_list function_declarations with
+    | Recursive, [_] -> true
+    | Recursive, ([] | _ :: _ :: _) -> false
+    | Non_recursive, _ -> false
+  in
+  let acc =
+    Acc.push_closure_info acc ~return_continuation ~exn_continuation ~my_closure
+      ~is_purely_tailrec:is_single_recursive_function
+  in
   let my_region = Function_decl.my_region decl in
   let function_slot = Function_decl.function_slot decl in
   let my_depth = Variable.create "my_depth" in
@@ -1276,14 +1296,6 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
     Let_with_acc.create acc bound (Named.create_rec_info next_depth_expr) ~body
   in
   let cost_metrics = Acc.cost_metrics acc in
-  let acc, exn_continuation =
-    close_exn_continuation acc external_env
-      (Function_decl.exn_continuation decl)
-  in
-  assert (
-    match Exn_continuation.extra_args exn_continuation with
-    | [] -> true
-    | _ :: _ -> false);
   let inline : Inline_attribute.t =
     (* We make a decision based on [fallback_inlining_heuristic] here to try to
        mimic Closure's behaviour as closely as possible, particularly when there
@@ -1317,6 +1329,7 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
     |> Acc.remove_continuation_from_free_names
          (Exn_continuation.exn_handler exn_continuation)
   in
+  let closure_info, acc = Acc.pop_closure_info acc in
   let params_arity = Bound_parameters.arity_with_subkinds params in
   let is_tupled =
     match Function_decl.kind decl with Curried _ -> false | Tupled -> true
@@ -1328,6 +1341,15 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
     then Function_decl_inlining_decision_type.Stub
     else Function_decl_inlining_decision_type.Not_yet_decided
   in
+  let loopify : Loopify_attribute.t =
+    match Function_decl.loop decl with
+    | Always_loop -> Always_loopify
+    | Never_loop -> Never_loopify
+    | Default_loop ->
+      if closure_info.is_purely_tailrec
+      then Default_loopify_and_tailrec
+      else Default_loopify_and_not_tailrec
+  in
   let code =
     Code.create code_id ~params_and_body
       ~free_names_of_params_and_body:(Acc.free_names acc) ~params_arity
@@ -1345,7 +1367,7 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
       ~dbg ~is_tupled
       ~is_my_closure_used:
         (Function_params_and_body.is_my_closure_used params_and_body)
-      ~inlining_decision ~absolute_history ~relative_history
+      ~inlining_decision ~absolute_history ~relative_history ~loopify
   in
   let approx =
     let code = Code_or_metadata.create code in
@@ -1470,6 +1492,7 @@ let close_functions acc external_env ~current_region function_declarations =
             ~inlining_decision:Recursive
             ~absolute_history:(Inlining_history.Absolute.empty compilation_unit)
             ~relative_history:Inlining_history.Relative.empty
+            ~loopify:Never_loopify
         in
         let code = Code_or_metadata.create_metadata_only metadata in
         let approx =
@@ -2205,6 +2228,9 @@ let close_program (type mode) ~(mode : mode Flambda_features.mode)
           defining_expr ~body)
       (acc, body) (Acc.declared_symbols acc)
   in
+  if Option.is_some (Acc.top_closure_info acc)
+  then
+    Misc.fatal_error "Information on nested closures should be empty at the end";
   let get_code_metadata code_id =
     Code_id.Map.find code_id (Acc.code acc) |> Code.code_metadata
   in

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

@@ -375,6 +375,13 @@ module Acc = struct
     | Trackable_arguments of Env.value_approximation list
     | Untrackable
 
+  type closure_info =
+    { return_continuation : Continuation.t;
+      exn_continuation : Exn_continuation.t;
+      my_closure : Variable.t;
+      is_purely_tailrec : bool
+    }
+
   type t =
     { declared_symbols : (Symbol.t * Static_const.t) list;
       lifted_sets_of_closures :
@@ -389,7 +396,8 @@ module Acc = struct
       seen_a_function : bool;
       symbol_for_global : Ident.t -> Symbol.t;
       slot_offsets : Slot_offsets.t;
-      regions_closed_early : Ident.Set.t
+      regions_closed_early : Ident.Set.t;
+      closure_infos : closure_info list
     }
 
   let cost_metrics t = t.cost_metrics
@@ -414,7 +422,8 @@ module Acc = struct
       seen_a_function = false;
       symbol_for_global;
       slot_offsets;
-      regions_closed_early = Ident.Set.empty
+      regions_closed_early = Ident.Set.empty;
+      closure_infos = []
     }
 
   let declared_symbols t = t.declared_symbols
@@ -451,15 +460,47 @@ module Acc = struct
   let add_free_names free_names t =
     { t with free_names = Name_occurrences.union free_names t.free_names }
 
-  let add_name_to_free_names ~name t =
+  let add_free_names_and_check_my_closure_use free_names t =
+    let t =
+      match t.closure_infos with
+      | [] -> t
+      | closure_info :: closure_infos ->
+        if closure_info.is_purely_tailrec
+           && Name_occurrences.mem_var free_names closure_info.my_closure
+        then
+          { t with
+            closure_infos =
+              { closure_info with is_purely_tailrec = false } :: closure_infos
+          }
+        else t
+    in
+    add_free_names free_names t
+
+  let add_name_to_free_names ~is_tail_call ~name t =
+    let closure_infos =
+      match is_tail_call, t.closure_infos with
+      | true, closure_infos -> closure_infos
+      | false, [] -> []
+      | false, closure_info :: closure_infos ->
+        if closure_info.is_purely_tailrec
+           && Name.equal (Name.var closure_info.my_closure) name
+        then { closure_info with is_purely_tailrec = false } :: closure_infos
+        else t.closure_infos
+    in
     { t with
+      closure_infos;
       free_names = Name_occurrences.add_name t.free_names name Name_mode.normal
     }
 
-  let add_simple_to_free_names acc simple =
+  let add_simple_to_free_names_maybe_tail_call ~is_tail_call acc simple =
     Simple.pattern_match simple
       ~const:(fun _ -> acc)
-      ~name:(fun name ~coercion:_ -> add_name_to_free_names ~name acc)
+      ~name:(fun name ~coercion ->
+        let acc = add_name_to_free_names ~is_tail_call ~name acc in
+        add_free_names (Coercion.free_names coercion) acc)
+
+  let add_simple_to_free_names acc simple =
+    add_simple_to_free_names_maybe_tail_call ~is_tail_call:false acc simple
 
   let remove_code_id_or_symbol_from_free_names code_id_or_symbol t =
     { t with
@@ -538,6 +579,36 @@ module Acc = struct
         set_of_closures
     in
     { t with slot_offsets }
+
+  let top_closure_info t =
+    match t.closure_infos with
+    | [] -> None
+    | closure_info :: _ -> Some closure_info
+
+  let push_closure_info t ~return_continuation ~exn_continuation ~my_closure
+      ~is_purely_tailrec =
+    { t with
+      closure_infos =
+        { return_continuation; exn_continuation; my_closure; is_purely_tailrec }
+        :: t.closure_infos
+    }
+
+  let pop_closure_info t =
+    let closure_info, closure_infos =
+      match t.closure_infos with
+      | [] -> Misc.fatal_error "pop_closure_info called on empty stack"
+      | closure_info :: closure_infos -> closure_info, closure_infos
+    in
+    let closure_infos =
+      match closure_infos with
+      | [] -> []
+      | closure_info2 :: closure_infos2 ->
+        if closure_info2.is_purely_tailrec
+           && Name_occurrences.mem_var t.free_names closure_info2.my_closure
+        then { closure_info2 with is_purely_tailrec = false } :: closure_infos2
+        else closure_infos
+    in
+    closure_info, { t with closure_infos }
 end
 
 module Function_decls = struct
@@ -709,7 +780,40 @@ module Expr_with_acc = struct
         (Code_size.apply apply |> Cost_metrics.from_size)
         acc
     in
-    let acc = Acc.add_free_names (Apply_expr.free_names apply) acc in
+    let is_tail_call =
+      match Acc.top_closure_info acc with
+      | None -> false
+      | Some { return_continuation; exn_continuation; _ } -> (
+        (match Apply_expr.continuation apply with
+        | Never_returns -> true
+        | Return cont -> Continuation.equal cont return_continuation)
+        && Exn_continuation.equal
+             (Apply_expr.exn_continuation apply)
+             exn_continuation
+        (* If the return and exn continuation match, the call is in tail
+           position, but could still be an under- or over-application. By
+           checking that it is a direct call, we are sure it has the correct
+           arity. *)
+        &&
+        match Apply.call_kind apply with
+        | Function { function_call = Direct _; _ } -> true
+        | Function
+            { function_call = Indirect_unknown_arity | Indirect_known_arity _;
+              _
+            } ->
+          false
+        | Method _ -> false
+        | C_call _ -> false)
+    in
+    let acc =
+      Acc.add_simple_to_free_names_maybe_tail_call ~is_tail_call acc
+        (Apply.callee apply)
+    in
+    let acc =
+      Acc.add_free_names_and_check_my_closure_use
+        (Apply_expr.free_names_except_callee apply)
+        acc
+    in
     let acc =
       match Apply_expr.continuation apply with
       | Never_returns -> acc
@@ -742,7 +846,11 @@ module Apply_cont_with_acc = struct
   let create acc ?trap_action ?args_approx cont ~args ~dbg =
     let apply_cont = Apply_cont.create ?trap_action cont ~args ~dbg in
     let acc = Acc.add_continuation_application ~cont args_approx acc in
-    let acc = Acc.add_free_names (Apply_cont.free_names apply_cont) acc in
+    let acc =
+      Acc.add_free_names_and_check_my_closure_use
+        (Apply_cont.free_names apply_cont)
+        acc
+    in
     acc, apply_cont
 
   let goto acc cont =
@@ -797,7 +905,18 @@ module Let_with_acc = struct
           ~code_id:(fun acc cid -> Acc.remove_code_id_from_free_names cid acc)
       in
       let let_expr = Let.create let_bound named ~body ~free_names_of_body in
-      let acc = Acc.add_free_names (Named.free_names named) acc in
+      let is_project_value_slot =
+        match[@ocaml.warning "-4"] (named : Named.t) with
+        | Prim (Unary (Project_value_slot _, _), _) -> true
+        | _ -> false
+      in
+      let acc =
+        if is_project_value_slot
+        then Acc.add_free_names (Named.free_names named) acc
+        else
+          Acc.add_free_names_and_check_my_closure_use (Named.free_names named)
+            acc
+      in
       acc, Expr.create_let let_expr
 end
 

middle_end/flambda2/from_lambda/closure_conversion_aux.mli

@@ -179,6 +179,13 @@ end
 
 (** Used to pipe some data through closure conversion *)
 module Acc : sig
+  type closure_info = private
+    { return_continuation : Continuation.t;
+      exn_continuation : Exn_continuation.t;
+      my_closure : Variable.t;
+      is_purely_tailrec : bool
+    }
+
   type t
 
   val create :
@@ -252,6 +259,18 @@ module Acc : sig
 
   val add_set_of_closures_offsets :
     is_phantom:bool -> t -> Set_of_closures.t -> t
+
+  val top_closure_info : t -> closure_info option
+
+  val push_closure_info :
+    t ->
+    return_continuation:Continuation.t ->
+    exn_continuation:Exn_continuation.t ->
+    my_closure:Variable.t ->
+    is_purely_tailrec:bool ->
+    t
+
+  val pop_closure_info : t -> closure_info * t
 end
 
 (** Used to represent information about a set of function declarations during

middle_end/flambda2/parser/fexpr_to_flambda.ml

@@ -778,6 +778,7 @@ let rec expr env (e : Fexpr.expr) : Flambda.Expr.t =
         in
         let code =
           (* CR mshinwell: [inlining_decision] should maybe be set properly *)
+          (* CR ncourant: same for loopify *)
           Code.create code_id ~params_and_body ~free_names_of_params_and_body
             ~newer_version_of ~params_arity ~num_trailing_local_params:0
             ~result_arity ~result_types:Unknown
@@ -793,6 +794,7 @@ let rec expr env (e : Fexpr.expr) : Flambda.Expr.t =
               (Inlining_history.Absolute.empty
                  (Compilation_unit.get_current_exn ()))
             ~relative_history:Inlining_history.Relative.empty
+            ~loopify:Never_loopify
         in
         Flambda.Static_const_or_code.create_code code
     in