Skip to content

Commit 77fd60f

Browse files
committed
Transform tail-recursive functions into recursive continuations
1 parent a77be10 commit 77fd60f

24 files changed

+511
-54
lines changed

middle_end/flambda2/from_lambda/closure_conversion.ml

Lines changed: 36 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1093,8 +1093,28 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
10931093
let params = Function_decl.params decl in
10941094
let return = Function_decl.return decl in
10951095
let return_continuation = Function_decl.return_continuation decl in
1096-
let recursive = Function_decl.recursive decl in
1096+
let acc, exn_continuation =
1097+
close_exn_continuation acc external_env
1098+
(Function_decl.exn_continuation decl)
1099+
in
1100+
assert (
1101+
match Exn_continuation.extra_args exn_continuation with
1102+
| [] -> true
1103+
| _ :: _ -> false);
10971104
let my_closure = Variable.create "my_closure" in
1105+
let recursive = Function_decl.recursive decl in
1106+
(* Mark function available for loopify only if it is a single recursive
1107+
function *)
1108+
let is_single_recursive_function =
1109+
match recursive, Function_decls.to_list function_declarations with
1110+
| Recursive, [_] -> true
1111+
| Recursive, ([] | _ :: _ :: _) -> false
1112+
| Non_recursive, _ -> false
1113+
in
1114+
let acc =
1115+
Acc.push_closure_info acc ~return_continuation ~exn_continuation ~my_closure
1116+
~is_purely_tailrec:is_single_recursive_function
1117+
in
10981118
let my_region = Function_decl.my_region decl in
10991119
let function_slot = Function_decl.function_slot decl in
11001120
let my_depth = Variable.create "my_depth" in
@@ -1280,14 +1300,6 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
12801300
Let_with_acc.create acc bound (Named.create_rec_info next_depth_expr) ~body
12811301
in
12821302
let cost_metrics = Acc.cost_metrics acc in
1283-
let acc, exn_continuation =
1284-
close_exn_continuation acc external_env
1285-
(Function_decl.exn_continuation decl)
1286-
in
1287-
assert (
1288-
match Exn_continuation.extra_args exn_continuation with
1289-
| [] -> true
1290-
| _ :: _ -> false);
12911303
let inline : Inline_attribute.t =
12921304
(* We make a decision based on [fallback_inlining_heuristic] here to try to
12931305
mimic Closure's behaviour as closely as possible, particularly when there
@@ -1321,6 +1333,7 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
13211333
|> Acc.remove_continuation_from_free_names
13221334
(Exn_continuation.exn_handler exn_continuation)
13231335
in
1336+
let closure_info, acc = Acc.pop_closure_info acc in
13241337
let params_arity = Bound_parameters.arity_with_subkinds params in
13251338
let is_tupled =
13261339
match Function_decl.kind decl with Curried _ -> false | Tupled -> true
@@ -1332,6 +1345,15 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
13321345
then Function_decl_inlining_decision_type.Stub
13331346
else Function_decl_inlining_decision_type.Not_yet_decided
13341347
in
1348+
let loopify : Loopify_attribute.t =
1349+
match Function_decl.loop decl with
1350+
| Always_loop -> Always_loopify
1351+
| Never_loop -> Never_loopify
1352+
| Default_loop ->
1353+
if closure_info.is_purely_tailrec
1354+
then Default_loopify_and_tailrec
1355+
else Default_loopify_and_not_tailrec
1356+
in
13351357
let code =
13361358
Code.create code_id ~params_and_body
13371359
~free_names_of_params_and_body:(Acc.free_names acc) ~params_arity
@@ -1351,7 +1373,7 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
13511373
~dbg ~is_tupled
13521374
~is_my_closure_used:
13531375
(Function_params_and_body.is_my_closure_used params_and_body)
1354-
~inlining_decision ~absolute_history ~relative_history
1376+
~inlining_decision ~absolute_history ~relative_history ~loopify
13551377
in
13561378
let approx =
13571379
let code = Code_or_metadata.create code in
@@ -1480,6 +1502,7 @@ let close_functions acc external_env ~current_region function_declarations =
14801502
~inlining_decision:Recursive
14811503
~absolute_history:(Inlining_history.Absolute.empty compilation_unit)
14821504
~relative_history:Inlining_history.Relative.empty
1505+
~loopify:Never_loopify
14831506
in
14841507
let code = Code_or_metadata.create_metadata_only metadata in
14851508
let approx =
@@ -2216,6 +2239,9 @@ let close_program (type mode) ~(mode : mode Flambda_features.mode)
22162239
defining_expr ~body)
22172240
(acc, body) (Acc.declared_symbols acc)
22182241
in
2242+
if Option.is_some (Acc.top_closure_info acc)
2243+
then
2244+
Misc.fatal_error "Information on nested closures should be empty at the end";
22192245
let get_code_metadata code_id =
22202246
Code_id.Map.find code_id (Acc.code acc) |> Code.code_metadata
22212247
in

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

Lines changed: 127 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -375,6 +375,13 @@ module Acc = struct
375375
| Trackable_arguments of Env.value_approximation list
376376
| Untrackable
377377

378+
type closure_info =
379+
{ return_continuation : Continuation.t;
380+
exn_continuation : Exn_continuation.t;
381+
my_closure : Variable.t;
382+
is_purely_tailrec : bool
383+
}
384+
378385
type t =
379386
{ declared_symbols : (Symbol.t * Static_const.t) list;
380387
lifted_sets_of_closures :
@@ -389,7 +396,8 @@ module Acc = struct
389396
seen_a_function : bool;
390397
symbol_for_global : Ident.t -> Symbol.t;
391398
slot_offsets : Slot_offsets.t;
392-
regions_closed_early : Ident.Set.t
399+
regions_closed_early : Ident.Set.t;
400+
closure_infos : closure_info list
393401
}
394402

395403
let cost_metrics t = t.cost_metrics
@@ -414,7 +422,8 @@ module Acc = struct
414422
seen_a_function = false;
415423
symbol_for_global;
416424
slot_offsets;
417-
regions_closed_early = Ident.Set.empty
425+
regions_closed_early = Ident.Set.empty;
426+
closure_infos = []
418427
}
419428

420429
let declared_symbols t = t.declared_symbols
@@ -451,15 +460,47 @@ module Acc = struct
451460
let add_free_names free_names t =
452461
{ t with free_names = Name_occurrences.union free_names t.free_names }
453462

454-
let add_name_to_free_names ~name t =
463+
let add_free_names_and_check_my_closure_use free_names t =
464+
let t =
465+
match t.closure_infos with
466+
| [] -> t
467+
| closure_info :: closure_infos ->
468+
if closure_info.is_purely_tailrec
469+
&& Name_occurrences.mem_var free_names closure_info.my_closure
470+
then
471+
{ t with
472+
closure_infos =
473+
{ closure_info with is_purely_tailrec = false } :: closure_infos
474+
}
475+
else t
476+
in
477+
add_free_names free_names t
478+
479+
let add_name_to_free_names ~is_tail_call ~name t =
480+
let closure_infos =
481+
match is_tail_call, t.closure_infos with
482+
| true, closure_infos -> closure_infos
483+
| false, [] -> []
484+
| false, closure_info :: closure_infos ->
485+
if closure_info.is_purely_tailrec
486+
&& Name.equal (Name.var closure_info.my_closure) name
487+
then { closure_info with is_purely_tailrec = false } :: closure_infos
488+
else t.closure_infos
489+
in
455490
{ t with
491+
closure_infos;
456492
free_names = Name_occurrences.add_name t.free_names name Name_mode.normal
457493
}
458494

459-
let add_simple_to_free_names acc simple =
495+
let add_simple_to_free_names_maybe_tail_call ~is_tail_call acc simple =
460496
Simple.pattern_match simple
461497
~const:(fun _ -> acc)
462-
~name:(fun name ~coercion:_ -> add_name_to_free_names ~name acc)
498+
~name:(fun name ~coercion ->
499+
let acc = add_name_to_free_names ~is_tail_call ~name acc in
500+
add_free_names (Coercion.free_names coercion) acc)
501+
502+
let add_simple_to_free_names acc simple =
503+
add_simple_to_free_names_maybe_tail_call ~is_tail_call:false acc simple
463504

464505
let remove_code_id_or_symbol_from_free_names code_id_or_symbol t =
465506
{ t with
@@ -538,6 +579,36 @@ module Acc = struct
538579
set_of_closures
539580
in
540581
{ t with slot_offsets }
582+
583+
let top_closure_info t =
584+
match t.closure_infos with
585+
| [] -> None
586+
| closure_info :: _ -> Some closure_info
587+
588+
let push_closure_info t ~return_continuation ~exn_continuation ~my_closure
589+
~is_purely_tailrec =
590+
{ t with
591+
closure_infos =
592+
{ return_continuation; exn_continuation; my_closure; is_purely_tailrec }
593+
:: t.closure_infos
594+
}
595+
596+
let pop_closure_info t =
597+
let closure_info, closure_infos =
598+
match t.closure_infos with
599+
| [] -> Misc.fatal_error "pop_closure_info called on empty stack"
600+
| closure_info :: closure_infos -> closure_info, closure_infos
601+
in
602+
let closure_infos =
603+
match closure_infos with
604+
| [] -> []
605+
| closure_info2 :: closure_infos2 ->
606+
if closure_info2.is_purely_tailrec
607+
&& Name_occurrences.mem_var t.free_names closure_info2.my_closure
608+
then { closure_info2 with is_purely_tailrec = false } :: closure_infos2
609+
else closure_infos
610+
in
611+
closure_info, { t with closure_infos }
541612
end
542613

543614
module Function_decls = struct
@@ -711,7 +782,40 @@ module Expr_with_acc = struct
711782
(Code_size.apply apply |> Cost_metrics.from_size)
712783
acc
713784
in
714-
let acc = Acc.add_free_names (Apply_expr.free_names apply) acc in
785+
let is_tail_call =
786+
match Acc.top_closure_info acc with
787+
| None -> false
788+
| Some { return_continuation; exn_continuation; _ } -> (
789+
(match Apply_expr.continuation apply with
790+
| Never_returns -> true
791+
| Return cont -> Continuation.equal cont return_continuation)
792+
&& Exn_continuation.equal
793+
(Apply_expr.exn_continuation apply)
794+
exn_continuation
795+
(* If the return and exn continuation match, the call is in tail
796+
position, but could still be an under- or over-application. By
797+
checking that it is a direct call, we are sure it has the correct
798+
arity. *)
799+
&&
800+
match Apply.call_kind apply with
801+
| Function { function_call = Direct _; _ } -> true
802+
| Function
803+
{ function_call = Indirect_unknown_arity | Indirect_known_arity _;
804+
_
805+
} ->
806+
false
807+
| Method _ -> false
808+
| C_call _ -> false)
809+
in
810+
let acc =
811+
Acc.add_simple_to_free_names_maybe_tail_call ~is_tail_call acc
812+
(Apply.callee apply)
813+
in
814+
let acc =
815+
Acc.add_free_names_and_check_my_closure_use
816+
(Apply_expr.free_names_except_callee apply)
817+
acc
818+
in
715819
let acc =
716820
match Apply_expr.continuation apply with
717821
| Never_returns -> acc
@@ -744,7 +848,11 @@ module Apply_cont_with_acc = struct
744848
let create acc ?trap_action ?args_approx cont ~args ~dbg =
745849
let apply_cont = Apply_cont.create ?trap_action cont ~args ~dbg in
746850
let acc = Acc.add_continuation_application ~cont args_approx acc in
747-
let acc = Acc.add_free_names (Apply_cont.free_names apply_cont) acc in
851+
let acc =
852+
Acc.add_free_names_and_check_my_closure_use
853+
(Apply_cont.free_names apply_cont)
854+
acc
855+
in
748856
acc, apply_cont
749857

750858
let goto acc cont =
@@ -799,7 +907,18 @@ module Let_with_acc = struct
799907
~code_id:(fun acc cid -> Acc.remove_code_id_from_free_names cid acc)
800908
in
801909
let let_expr = Let.create let_bound named ~body ~free_names_of_body in
802-
let acc = Acc.add_free_names (Named.free_names named) acc in
910+
let is_project_value_slot =
911+
match[@ocaml.warning "-4"] (named : Named.t) with
912+
| Prim (Unary (Project_value_slot _, _), _) -> true
913+
| _ -> false
914+
in
915+
let acc =
916+
if is_project_value_slot
917+
then Acc.add_free_names (Named.free_names named) acc
918+
else
919+
Acc.add_free_names_and_check_my_closure_use (Named.free_names named)
920+
acc
921+
in
803922
acc, Expr.create_let let_expr
804923
end
805924

middle_end/flambda2/from_lambda/closure_conversion_aux.mli

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,13 @@ end
179179

180180
(** Used to pipe some data through closure conversion *)
181181
module Acc : sig
182+
type closure_info = private
183+
{ return_continuation : Continuation.t;
184+
exn_continuation : Exn_continuation.t;
185+
my_closure : Variable.t;
186+
is_purely_tailrec : bool
187+
}
188+
182189
type t
183190

184191
val create :
@@ -252,6 +259,18 @@ module Acc : sig
252259

253260
val add_set_of_closures_offsets :
254261
is_phantom:bool -> t -> Set_of_closures.t -> t
262+
263+
val top_closure_info : t -> closure_info option
264+
265+
val push_closure_info :
266+
t ->
267+
return_continuation:Continuation.t ->
268+
exn_continuation:Exn_continuation.t ->
269+
my_closure:Variable.t ->
270+
is_purely_tailrec:bool ->
271+
t
272+
273+
val pop_closure_info : t -> closure_info * t
255274
end
256275

257276
(** Used to represent information about a set of function declarations during

middle_end/flambda2/parser/fexpr_to_flambda.ml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -778,6 +778,7 @@ let rec expr env (e : Fexpr.expr) : Flambda.Expr.t =
778778
in
779779
let code =
780780
(* CR mshinwell: [inlining_decision] should maybe be set properly *)
781+
(* CR ncourant: same for loopify *)
781782
Code.create code_id ~params_and_body ~free_names_of_params_and_body
782783
~newer_version_of ~params_arity ~num_trailing_local_params:0
783784
~result_arity ~result_types:Unknown
@@ -793,6 +794,7 @@ let rec expr env (e : Fexpr.expr) : Flambda.Expr.t =
793794
(Inlining_history.Absolute.empty
794795
(Compilation_unit.get_current_exn ()))
795796
~relative_history:Inlining_history.Relative.empty
797+
~loopify:Never_loopify
796798
in
797799
Flambda.Static_const_or_code.create_code code
798800
in

0 commit comments

Comments
 (0)