1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
(**************************************************************************)
(* *)
(* OCaml *)
(* *)
(* Pierre Chambart, OCamlPro *)
(* Mark Shinwell and Leo White, Jane Street Europe *)
(* *)
(* Copyright 2013--2016 OCamlPro SAS *)
(* Copyright 2014--2016 Jane Street Group LLC *)
(* *)
(* All rights reserved. This file is distributed under the terms of *)
(* the GNU Lesser General Public License version 2.1, with the *)
(* special exception on linking described in the file LICENSE. *)
(* *)
(**************************************************************************)
[@@@ocaml.warning "+a-4-9-30-40-41-42-66"]
open! Int_replace_polymorphic_compare
(** A variable in a closure can either be used by the closure itself
or by an inlined version of the function. *)
let remove_unused_closure_variables ~remove_direct_call_surrogates program =
let used_vars_within_closure, used_closure_ids =
let used = Var_within_closure.Tbl.create 13 in
let used_fun = Closure_id.Tbl.create 13 in
let aux_named (named : Flambda.named) =
match named with
| Project_closure { set_of_closures = _; closure_id } ->
Closure_id.Tbl.add used_fun closure_id ()
| Project_var { closure_id; var } ->
Var_within_closure.Tbl.add used var ();
Closure_id.Tbl.add used_fun closure_id ()
| Move_within_set_of_closures { closure = _; start_from; move_to } ->
Closure_id.Tbl.add used_fun start_from ();
Closure_id.Tbl.add used_fun move_to ()
| Symbol _ | Const _ | Set_of_closures _ | Prim _ | Expr _
| Allocated_const _ | Read_mutable _ | Read_symbol_field _ -> ()
in
Flambda_iterators.iter_named_of_program ~f:aux_named program;
used, used_fun
in
let aux_named _ (named : Flambda.named) : Flambda.named =
match named with
| Set_of_closures ({ function_decls; free_vars; _ } as set_of_closures) ->
let direct_call_surrogates =
if remove_direct_call_surrogates then Variable.Set.empty
else
Variable.Set.of_list
(Variable.Map.data set_of_closures.direct_call_surrogates)
in
let rec add_needed needed_funs remaining_funs free_vars_of_kept_funs =
let new_needed_funs, remaining_funs =
(* Keep a function if it is used either by the rest of the code,
(in used_closure_ids), or by any other kept function
(in free_vars_of_kept_funs) *)
Variable.Map.partition (fun fun_id _ ->
Variable.Set.mem fun_id free_vars_of_kept_funs
|| Closure_id.Tbl.mem used_closure_ids
(Closure_id.wrap fun_id)
|| Variable.Set.mem fun_id direct_call_surrogates)
remaining_funs
in
if Variable.Map.is_empty new_needed_funs then
(* If no new function is needed, we reached fixpoint *)
needed_funs, free_vars_of_kept_funs
else begin
let needed_funs =
Variable.Map.disjoint_union needed_funs new_needed_funs
in
let free_vars_of_kept_funs =
Variable.Map.fold (fun _ { Flambda. free_variables } acc ->
Variable.Set.union free_variables acc)
new_needed_funs
free_vars_of_kept_funs
in
add_needed needed_funs remaining_funs free_vars_of_kept_funs
end
in
let funs, free_vars_of_kept_funs =
add_needed Variable.Map.empty function_decls.funs Variable.Set.empty
in
let free_vars =
Variable.Map.filter (fun id _var ->
Variable.Set.mem id free_vars_of_kept_funs
|| Var_within_closure.Tbl.mem
used_vars_within_closure
(Var_within_closure.wrap id))
free_vars
in
let function_decls =
Flambda.update_function_declarations function_decls ~funs
in
let specialised_args =
(* Remove specialised args that are used by removed functions *)
let all_remaining_arguments =
Variable.Map.fold (fun _ { Flambda.params } set ->
Variable.Set.union set (Parameter.Set.vars params))
funs Variable.Set.empty
in
Variable.Map.filter (fun arg _ ->
Variable.Set.mem arg all_remaining_arguments)
set_of_closures.specialised_args
in
let free_vars =
Flambda_utils.clean_projections ~which_variables:free_vars
in
let direct_call_surrogates =
(* Remove direct call surrogates where either the existing function
or the surrogate has been eliminated. *)
Variable.Map.fold (fun existing surrogate surrogates ->
if not (Variable.Map.mem existing funs)
|| not (Variable.Map.mem surrogate funs)
then surrogates
else Variable.Map.add existing surrogate surrogates)
set_of_closures.direct_call_surrogates
Variable.Map.empty
in
let set_of_closures =
Flambda.create_set_of_closures ~function_decls
~free_vars ~specialised_args ~direct_call_surrogates
in
Set_of_closures set_of_closures
| e -> e
in
Flambda_iterators.map_named_of_program ~f:aux_named program
|
