Skip to content

Commit b1771ae

Browse files
committed
Merge pull request #67 from bloomberg/beta_reduce_refine
Fix #64: better beta-reduction, see issus #64 for more details
2 parents d3b6428 + 3ee879b commit b1771ae

32 files changed

+528
-195
lines changed

jscomp/lam_analysis.ml

Lines changed: 143 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -306,10 +306,149 @@ and eq_primitive (p : Lambda.primitive) (p1 : Lambda.primitive) =
306306
try p = p1 with _ -> false
307307

308308

309-
let is_closed_by map lam =
310-
Lambda.IdentSet.for_all Ident.global
311-
(Lambda.IdentSet.diff (Lambda.free_variables lam) map )
309+
310+
type stats =
311+
{
312+
mutable top : bool ;
313+
(* all appearances are in the top, substitution is fine
314+
whether it is pure or not
315+
{[
316+
(fun x y
317+
-> x + y + (f x )) (32) (console.log('hi'), 33)
318+
]}
319+
since in ocaml, the application order is intentionally undefined,
320+
note if [times] is not one, this field does not make sense
321+
*)
322+
mutable times : int ;
323+
}
324+
type env =
325+
{ top : bool ;
326+
loop : bool
327+
}
328+
329+
let no_substitute = { top = false; loop = true }
330+
let fresh_env = {top = true; loop = false }
331+
let fresh_stats () = { top = true; times = 0 }
332+
333+
let param_map_of_list lst =
334+
List.fold_left (fun acc l -> Ident_map.add l (fresh_stats ()) acc) Ident_map.empty lst
335+
336+
(** Sanity check, remove all varaibles in [local_set] in the last pass *)
337+
338+
let free_variables (export_idents : Ident_set.t ) (params : stats Ident_map.t ) lam =
339+
let fv = ref params in
340+
let local_set = ref export_idents in
341+
342+
let local_add k =
343+
local_set := Ident_set.add k !local_set in
344+
let local_add_list ks =
345+
local_set :=
346+
List.fold_left (fun acc k -> Ident_set.add k acc) !local_set ks
347+
in
348+
let loop_use = 100 in
349+
let map_use {top; loop} v =
350+
(* relies on [identifier] uniquely bound *)
351+
let times = if loop then loop_use else 1 in
352+
if Ident_set.mem v !local_set then ()
353+
else begin match Ident_map.find v !fv with
354+
| exception Not_found
355+
-> fv := Ident_map.add v { top ; times } !fv
356+
| v ->
357+
v.times <- v.times + times ;
358+
v.top <- v.top && top
359+
end
360+
in
361+
let new_env lam (env : env) =
362+
if env.top then
363+
if no_side_effects lam
364+
then env
365+
else { env with top = false}
366+
else env
367+
in
368+
let rec iter (top : env) (lam : Lambda.lambda) =
369+
match lam with
370+
| Lvar v -> map_use top v
371+
| Lconst _ -> ()
372+
| Lapply(fn, args, _) ->
373+
iter top fn;
374+
let top = new_env fn top in
375+
List.iter (iter top ) args
376+
| Lprim(_p, args) ->
377+
(* Check: can top be propoaged for all primitives *)
378+
List.iter (iter top) args
379+
| Lfunction(_kind, params, body) ->
380+
local_add_list params;
381+
iter no_substitute body
382+
| Llet(_let_kind, id, arg, body) ->
383+
local_add id ;
384+
iter top arg; iter no_substitute body
385+
| Lletrec(decl, body) ->
386+
local_set := List.fold_left (fun acc (id, _) ->
387+
Ident_set.add id acc) !local_set decl;
388+
List.iter (fun (_, exp) -> iter no_substitute exp) decl;
389+
iter no_substitute body
390+
| Lswitch(arg, sw) ->
391+
iter top arg;
392+
let top = new_env arg top in
393+
List.iter (fun (key, case) -> iter top case) sw.sw_consts;
394+
List.iter (fun (key, case) -> iter top case) sw.sw_blocks;
395+
396+
begin match sw.sw_failaction with
397+
| None -> ()
398+
| Some x ->
399+
let nconsts = List.length sw.sw_consts in
400+
let nblocks = List.length sw.sw_blocks in
401+
402+
if nconsts < sw.sw_numconsts && nblocks < sw.sw_numblocks then
403+
iter no_substitute x
404+
else
405+
iter top x
406+
end
407+
408+
| Lstringswitch (arg,cases,default) ->
409+
iter top arg ;
410+
let top = new_env arg top in
411+
List.iter (fun (_,act) -> iter top act) cases ;
412+
begin match default with
413+
| None -> ()
414+
| Some x -> iter top x
415+
end
416+
| Lstaticraise (_,args) ->
417+
List.iter (iter no_substitute ) args
418+
| Lstaticcatch(e1, (_,vars), e2) ->
419+
iter no_substitute e1;
420+
local_add_list vars;
421+
iter no_substitute e2
422+
| Ltrywith(e1, exn, e2) ->
423+
iter top e1; iter no_substitute e2
424+
| Lifthenelse(e1, e2, e3) ->
425+
iter top e1;
426+
let top = new_env e1 top in
427+
iter top e2; iter top e3
428+
| Lsequence(e1, e2) ->
429+
iter top e1; iter no_substitute e2
430+
| Lwhile(e1, e2) ->
431+
iter no_substitute e1; iter no_substitute e2 (* in the loop, no substitution any way *)
432+
| Lfor(v, e1, e2, dir, e3) ->
433+
local_add v ;
434+
iter no_substitute e1; iter no_substitute e2; iter no_substitute e3
435+
| Lassign(id, e) ->
436+
map_use top id ;
437+
iter top e
438+
| Lsend (_k, met, obj, args, _) ->
439+
iter no_substitute met ;
440+
iter no_substitute obj;
441+
List.iter (iter no_substitute) args
442+
| Levent (lam, evt) ->
443+
iter top lam
444+
| Lifused (v, e) ->
445+
iter no_substitute e in
446+
iter fresh_env lam ; !fv
447+
448+
449+
let is_closed_by set lam =
450+
Ident_map.is_empty (free_variables set (Ident_map.empty ) lam )
312451

313452

314453
let is_closed lam =
315-
Lambda.IdentSet.for_all Ident.global (Lambda.free_variables lam)
454+
Ident_map.is_empty (free_variables Ident_set.empty Ident_map.empty lam)

jscomp/lam_analysis.mli

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,31 @@ val eq_lambda : Lambda.lambda -> Lambda.lambda -> bool
3333
(** [is_closed_by map lam]
3434
return [true] if all unbound variables
3535
belongs to the given [map] *)
36-
val is_closed_by : Lambda.IdentSet.t -> Lambda.lambda -> bool
36+
val is_closed_by : (* Lambda. *) Ident_set.t -> Lambda.lambda -> bool
3737

3838
val is_closed : Lambda.lambda -> bool
3939

4040

41+
42+
43+
44+
type stats =
45+
{
46+
mutable top : bool ;
47+
(* all appearances are in the top, substitution is fine
48+
whether it is pure or not
49+
{[
50+
(fun x y
51+
-> x + y + (f x )) (32) (console.log('hi'), 33)
52+
]}
53+
since in ocaml, the application order is intentionally undefined,
54+
note if [times] is not one, this field does not make sense
55+
*)
56+
mutable times : int ;
57+
}
58+
59+
val param_map_of_list : Ident.t list -> stats Ident_map.t
60+
val free_variables : Ident_set.t -> stats Ident_map.t -> Lambda.lambda -> stats Ident_map.t
61+
4162
val small_inline_size : int
4263
val exit_inline_size : int

jscomp/lam_beta_reduce.ml

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,61 @@ let propogate_beta_reduce
228228
Lam_util.refine_let param arg l)
229229
rest_bindings new_body
230230

231+
let propogate_beta_reduce_with_map
232+
(meta : Lam_stats.meta) (map : Lam_analysis.stats Ident_map.t ) params body args =
233+
let rest_bindings, rev_new_params =
234+
List.fold_left2
235+
(fun (rest_bindings, acc) old_param (arg : Lambda.lambda) ->
236+
match arg with
237+
| Lconst _
238+
| Lvar _ -> rest_bindings , arg :: acc
239+
| Lprim (Pgetglobal ident, [])
240+
(* TODO: we can pass Global, but you also need keep track of it*)
241+
->
242+
let p = Ident.rename old_param in
243+
(p,arg) :: rest_bindings , (Lambda.Lvar p) :: acc
244+
245+
| _ ->
246+
if Lam_analysis.no_side_effects arg then
247+
begin match Ident_map.find old_param map with
248+
| exception Not_found -> assert false
249+
| {top = true ; times = 0 }
250+
| {top = true ; times = 1 }
251+
->
252+
rest_bindings, arg :: acc
253+
| _ ->
254+
let p = Ident.rename old_param in
255+
(p,arg) :: rest_bindings , (Lambda.Lvar p) :: acc
256+
end
257+
else
258+
let p = Ident.rename old_param in
259+
(p,arg) :: rest_bindings , (Lambda.Lvar p) :: acc
260+
) ([],[]) params args in
261+
let new_body = rewrite (Ext_hashtbl.of_list2 (List.rev params) (rev_new_params)) body in
262+
List.fold_right
263+
(fun (param, (arg : Lambda.lambda)) l ->
264+
let arg =
265+
match arg with
266+
| Lvar v ->
267+
begin
268+
match Hashtbl.find meta.ident_tbl v with
269+
| exception Not_found -> ()
270+
| ident_info ->
271+
Hashtbl.add meta.ident_tbl param ident_info
272+
end;
273+
arg
274+
| Lprim (Pgetglobal ident, []) ->
275+
(* It's not completeness, its to make it sound.. *)
276+
Lam_compile_global.query_lambda ident meta.env
277+
(* alias meta param ident (Module (Global ident)) Strict *)
278+
| Lprim (Pmakeblock (_, _, Immutable ) , ls) ->
279+
Hashtbl.replace meta.ident_tbl param
280+
(Lam_util.kind_of_lambda_block ls ); (** *)
281+
arg
282+
| _ -> arg in
283+
Lam_util.refine_let param arg l)
284+
rest_bindings new_body
285+
231286

232287

233288

jscomp/lam_beta_reduce.mli

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,9 @@ val propogate_beta_reduce :
4949
val refresh :
5050
Lambda.lambda ->
5151
Lambda.lambda
52+
53+
val propogate_beta_reduce_with_map :
54+
Lam_stats.meta ->
55+
Lam_analysis.stats Ident_map.t ->
56+
Ident_map.key list ->
57+
Lambda.lambda -> Lambda.lambda list -> Lambda.lambda

jscomp/lam_pass_collect.ml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ let collect_helper (meta : Lam_stats.meta) (lam : Lambda.lambda) =
9999
collect l
100100
| x ->
101101
collect x ;
102-
if Lambda.IdentSet.mem ident meta.export_idents then
102+
if Ident_set.mem ident meta.export_idents then
103103
annotate meta rec_flag ident (Lam_stats_util.get_arity meta x ) lam
104104

105105

jscomp/lam_pass_exits.ml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,8 @@ let count_helper (lam : Lambda.lambda) : (int, int ref) Hashtbl.t =
106106
| Levent(l, _) -> count l
107107
| Lifused(_, l) -> count l
108108

109-
and count_default sw = match sw.sw_failaction with
109+
and count_default sw =
110+
match sw.sw_failaction with
110111
| None -> ()
111112
| Some al ->
112113
let nconsts = List.length sw.sw_consts

jscomp/lam_pass_remove_alias.ml

Lines changed: 35 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ let simplify_alias
3939
since we aliased k, so it's safe to remove it?
4040
*)
4141
let v = simpl l in
42-
if Lambda.IdentSet.mem k meta.export_idents
42+
if Ident_set.mem k meta.export_idents
4343
then
4444
Llet(kind, k, g, v)
4545
(* in this case it is preserved, but will still be simplified
@@ -125,32 +125,43 @@ let simplify_alias
125125
simpl (Lam_beta_reduce.propogate_beta_reduce meta params body args)
126126
end
127127
else
128-
if lam_size < Lam_analysis.small_inline_size &&
129-
(Lam_analysis.is_closed_by meta.export_idents _m
130-
|| not (Lambda.IdentSet.mem v meta.export_idents))
131-
132-
then
133-
if rec_flag = Rec then
134-
begin
135-
(* Ext_log.dwarn __LOC__ "beta rec.. %s/%d@." v.name v.stamp ; *)
136-
Lam_beta_reduce.propogate_beta_reduce meta params body args
137-
end
128+
if lam_size < Lam_analysis.small_inline_size then
129+
let param_fresh_map = Lam_analysis.param_map_of_list params in
130+
let param_map =
131+
Lam_analysis.free_variables meta.export_idents param_fresh_map body in
132+
let old_count = List.length params in
133+
let new_count = Ident_map.cardinal param_map in
134+
if
135+
(
136+
not (Ident_set.mem v meta.export_idents)
137+
|| old_count = new_count
138+
)
139+
140+
then
141+
if rec_flag = Rec then
142+
begin
143+
(* Ext_log.dwarn __LOC__ "beta rec.. %s/%d@." v.name v.stamp ; *)
144+
(* Lam_beta_reduce.propogate_beta_reduce meta params body args *)
145+
Lam_beta_reduce.propogate_beta_reduce_with_map meta param_map params body args
146+
end
147+
else
148+
begin
149+
(* Ext_log.dwarn __LOC__ "beta nonrec..[%d] [%a] %s/%d@." *)
150+
(* (List.length args) *)
151+
(* Printlambda.lambda body *)
152+
(* v.name v.stamp ; *)
153+
simpl (Lam_beta_reduce.propogate_beta_reduce_with_map meta param_map params body args)
154+
155+
end
156+
else
157+
Lapply ( simpl l1, List.map simpl args, info)
138158
else
139159
begin
140-
(* Ext_log.dwarn __LOC__ "beta nonrec..[%d] [%a] %s/%d@." *)
141-
(* (List.length args) *)
142-
(* Printlambda.lambda body *)
143-
(* v.name v.stamp ; *)
144-
simpl (Lam_beta_reduce.propogate_beta_reduce meta params body args)
145-
160+
(* Ext_log.dwarn __LOC__ "%s/%d: %d @." *)
161+
(* v.name v.stamp lam_size *)
162+
(* ; *)
163+
Lapply ( simpl l1, List.map simpl args, info)
146164
end
147-
else
148-
begin
149-
(* Ext_log.dwarn __LOC__ "%s/%d: %d @." *)
150-
(* v.name v.stamp lam_size *)
151-
(* ; *)
152-
Lapply ( simpl l1, List.map simpl args, info)
153-
end
154165
else
155166
begin
156167
(* Ext_log.dwarn __LOC__ "%d vs %d @." (List.length args) (List.length params); *)

jscomp/lam_stats.ml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ type ident_info = {
8888
type meta = {
8989
env : Env.t;
9090
filename : string ;
91-
export_idents : Lambda.IdentSet.t ;
91+
export_idents : Ident_set.t ;
9292
exports : Ident.t list ;
9393

9494
alias_tbl : alias_tbl;

jscomp/lam_stats.mli

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ type ident_info = {
9090
type meta = {
9191
env : Env.t;
9292
filename : string ;
93-
export_idents : Lambda.IdentSet.t ;
93+
export_idents : Ident_set.t ;
9494
exports : Ident.t list ;
9595
alias_tbl : alias_tbl;
9696
exit_codes : int Hash_set.hashset;

jscomp/lam_stats_util.ml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,7 @@ let export_to_cmj
252252
let closed_lambda =
253253
if Lam_inline_util.maybe_functor x.name
254254
then
255-
if Lam_analysis.is_closed lambda
255+
if Lam_analysis.is_closed lambda (* TODO: seriealize more*)
256256
then Some lambda
257257
else None
258258
else None in

0 commit comments

Comments
 (0)