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
|
(**************************************************************************)
(* *)
(* 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
let pass_name = "remove-free-vars-equal-to-args"
let () = Pass_wrapper.register ~pass_name
let rewrite_one_function_decl ~(function_decl : Flambda.function_declaration)
~back_free_vars ~specialised_args =
let params_for_equal_free_vars =
List.fold_left (fun subst param ->
match Variable.Map.find param specialised_args with
| exception Not_found ->
(* param is not specialised *)
subst
| (spec_to : Flambda.specialised_to) ->
let outside_var = spec_to.var in
match Variable.Map.find outside_var back_free_vars with
| exception Not_found ->
(* No free variables equal to the param *)
subst
| set ->
(* Replace the free variables equal to a parameter *)
Variable.Set.fold (fun free_var subst ->
Variable.Map.add free_var param subst)
set subst)
Variable.Map.empty (Parameter.List.vars function_decl.params)
in
if Variable.Map.is_empty params_for_equal_free_vars then
function_decl
else
let body =
Flambda_utils.toplevel_substitution
params_for_equal_free_vars
function_decl.body
in
Flambda.update_function_declaration function_decl
~params:function_decl.params ~body:body
let rewrite_one_set_of_closures (set_of_closures : Flambda.set_of_closures) =
let back_free_vars =
Variable.Map.fold (fun var (outside_var : Flambda.specialised_to) map ->
let set =
match Variable.Map.find outside_var.var map with
| exception Not_found -> Variable.Set.singleton var
| set -> Variable.Set.add var set
in
Variable.Map.add outside_var.var set map)
set_of_closures.free_vars Variable.Map.empty
in
let done_something = ref false in
let funs =
Variable.Map.map (fun function_decl ->
let new_function_decl =
rewrite_one_function_decl ~function_decl ~back_free_vars
~specialised_args:set_of_closures.specialised_args
in
if not (new_function_decl == function_decl) then begin
done_something := true
end;
new_function_decl)
set_of_closures.function_decls.funs
in
if not !done_something then
None
else
let function_decls =
Flambda.update_function_declarations
set_of_closures.function_decls ~funs
in
let set_of_closures =
Flambda.create_set_of_closures
~function_decls
~free_vars:set_of_closures.free_vars
~specialised_args:set_of_closures.specialised_args
~direct_call_surrogates:set_of_closures.direct_call_surrogates
in
Some set_of_closures
let run ~ppf_dump set_of_closures =
Pass_wrapper.with_dump ~ppf_dump ~pass_name ~input:set_of_closures
~print_input:Flambda.print_set_of_closures
~print_output:Flambda.print_set_of_closures
~f:(fun () -> rewrite_one_set_of_closures set_of_closures)
|
