Fix #64: better beta-reduction, see issus #64 for more details

jscomp/lam_analysis.ml

@@ -306,10 +306,149 @@ and eq_primitive (p : Lambda.primitive) (p1 : Lambda.primitive) =
     try p = p1 with _ -> false
 
 
-let is_closed_by map lam = 
-  Lambda.IdentSet.for_all Ident.global 
-    (Lambda.IdentSet.diff (Lambda.free_variables lam) map )
+
+type stats = 
+  { 
+    mutable top : bool ; 
+    (* all appearances are in the top,  substitution is fine 
+       whether it is pure or not
+       {[
+         (fun x y          
+           ->  x + y + (f x )) (32) (console.log('hi'), 33)
+       ]}       
+       since in ocaml, the application order is intentionally undefined, 
+       note if [times] is not one, this field does not make sense       
+    *)    
+    mutable times : int ; 
+  }
+type env = 
+  { top  : bool ; 
+    loop : bool 
+  }
+
+let no_substitute = { top = false; loop = true }
+let fresh_env = {top = true; loop = false }
+let fresh_stats () = { top = true; times = 0 }
+
+let param_map_of_list lst = 
+  List.fold_left  (fun acc l -> Ident_map.add l (fresh_stats ()) acc) Ident_map.empty  lst 
+
+(** Sanity check, remove all varaibles in [local_set] in the last pass *)  
+
+let free_variables (export_idents : Ident_set.t ) (params : stats Ident_map.t ) lam = 
+  let fv = ref params in
+  let local_set = ref export_idents in
+
+  let local_add k =
+    local_set := Ident_set.add k !local_set in
+  let local_add_list ks = 
+    local_set :=  
+      List.fold_left (fun acc k -> Ident_set.add k acc) !local_set ks 
+  in    
+  let loop_use = 100 in
+  let map_use {top; loop} v = 
+    (* relies on [identifier] uniquely bound *)    
+    let times = if loop then loop_use else 1 in
+    if Ident_set.mem v !local_set then ()    
+    else begin match Ident_map.find v !fv with 
+      | exception Not_found
+        -> fv := Ident_map.add v { top ; times } !fv
+      | v ->
+        v.times <- v.times + times ; 
+        v.top <- v.top && top          
+    end
+  in
+  let new_env lam (env : env) = 
+    if env.top then 
+      if no_side_effects lam 
+      then env 
+      else { env with top = false}
+    else env      
+  in    
+  let rec iter (top : env) (lam : Lambda.lambda) =
+    match lam with 
+    | Lvar v -> map_use top v 
+    | Lconst _ -> ()
+    | Lapply(fn, args, _) ->
+      iter top  fn; 
+      let top = new_env fn top in
+      List.iter (iter top ) args  
+    | Lprim(_p, args) -> 
+      (* Check: can top be propoaged for all primitives *)
+      List.iter (iter top) args
+    | Lfunction(_kind, params, body) ->
+      local_add_list params;
+      iter no_substitute body 
+    | Llet(_let_kind, id, arg, body) ->
+      local_add id ;  
+      iter top  arg; iter no_substitute body
+    | Lletrec(decl, body) ->
+      local_set := List.fold_left (fun acc (id, _) -> 
+          Ident_set.add id acc) !local_set decl;        
+      List.iter (fun (_, exp) -> iter no_substitute exp) decl;
+      iter no_substitute body
+    | Lswitch(arg, sw) ->
+      iter top arg; 
+      let top = new_env arg top  in       
+      List.iter (fun (key, case) -> iter top case) sw.sw_consts;
+      List.iter (fun (key, case) -> iter top  case) sw.sw_blocks;
+
+      begin match sw.sw_failaction with 
+        | None -> ()
+        | Some x ->
+          let nconsts = List.length sw.sw_consts in
+          let nblocks = List.length sw.sw_blocks in
+
+          if nconsts < sw.sw_numconsts  && nblocks < sw.sw_numblocks then
+            iter no_substitute x
+          else
+            iter top x
+      end
+
+    | Lstringswitch (arg,cases,default) ->
+      iter top arg ;
+      let top = new_env arg top  in       
+      List.iter (fun (_,act) -> iter top  act) cases ;
+      begin match default with 
+      | None -> ()
+      | Some x -> iter top x 
+      end
+    | Lstaticraise (_,args) ->
+      List.iter (iter no_substitute ) args
+    | Lstaticcatch(e1, (_,vars), e2) ->
+      iter no_substitute  e1; 
+      local_add_list vars;       
+      iter no_substitute e2
+    | Ltrywith(e1, exn, e2) ->
+      iter top  e1; iter no_substitute  e2
+    | Lifthenelse(e1, e2, e3) ->
+      iter top e1; 
+      let top = new_env e1 top  in
+      iter top e2; iter top e3
+    | Lsequence(e1, e2) ->
+      iter top e1; iter no_substitute e2
+    | Lwhile(e1, e2) ->
+      iter no_substitute e1; iter no_substitute e2 (* in the loop, no substitution any way *)
+    | Lfor(v, e1, e2, dir, e3) ->
+      local_add v ; 
+      iter no_substitute e1; iter no_substitute e2; iter no_substitute e3
+    | Lassign(id, e) ->
+      map_use top  id ; 
+      iter top e
+    | Lsend (_k, met, obj, args, _) ->
+      iter no_substitute met ; 
+      iter no_substitute obj;
+      List.iter (iter no_substitute) args
+    | Levent (lam, evt) ->
+      iter top  lam
+    | Lifused (v, e) ->
+      iter no_substitute e in
+  iter fresh_env  lam ; !fv 
+
+
+let is_closed_by set lam = 
+  Ident_map.is_empty (free_variables set (Ident_map.empty ) lam   )
 
 
 let is_closed  lam = 
-  Lambda.IdentSet.for_all Ident.global (Lambda.free_variables lam)  
+  Ident_map.is_empty (free_variables Ident_set.empty Ident_map.empty lam)  

jscomp/lam_analysis.mli

@@ -33,10 +33,31 @@ val eq_lambda : Lambda.lambda -> Lambda.lambda -> bool
 (** [is_closed_by map lam]
     return [true] if all unbound variables
     belongs to the given [map] *)
-val is_closed_by : Lambda.IdentSet.t -> Lambda.lambda -> bool
+val is_closed_by : (* Lambda. *) Ident_set.t -> Lambda.lambda -> bool
 
 val is_closed : Lambda.lambda -> bool
 
 
+
+
+
+type stats = 
+  { 
+    mutable top : bool ; 
+    (* all appearances are in the top,  substitution is fine 
+       whether it is pure or not
+       {[
+         (fun x y          
+           ->  x + y + (f x )) (32) (console.log('hi'), 33)
+       ]}       
+       since in ocaml, the application order is intentionally undefined, 
+       note if [times] is not one, this field does not make sense       
+    *)    
+    mutable times : int ; 
+  }
+
+val param_map_of_list : Ident.t list -> stats Ident_map.t
+val free_variables : Ident_set.t -> stats Ident_map.t -> Lambda.lambda -> stats Ident_map.t
+
 val small_inline_size : int 
 val exit_inline_size : int 

jscomp/lam_beta_reduce.ml

@@ -228,6 +228,61 @@ let propogate_beta_reduce
        Lam_util.refine_let param arg l) 
      rest_bindings new_body
 
+let propogate_beta_reduce_with_map  
+    (meta : Lam_stats.meta) (map : Lam_analysis.stats Ident_map.t ) params body args =
+  let rest_bindings, rev_new_params  = 
+    List.fold_left2 
+      (fun (rest_bindings, acc) old_param (arg : Lambda.lambda) -> 
+         match arg with          
+         | Lconst _
+         | Lvar _  -> rest_bindings , arg :: acc 
+         | Lprim (Pgetglobal ident, [])
+           (* TODO: we can pass Global, but you also need keep track of it*)
+           ->
+           let p = Ident.rename old_param in 
+           (p,arg) :: rest_bindings , (Lambda.Lvar p) :: acc 
+
+         | _ -> 
+           if  Lam_analysis.no_side_effects arg then
+             begin match Ident_map.find old_param map with 
+               | exception Not_found -> assert false 
+               | {top = true ; times = 0 }
+               | {top = true ; times = 1 } 
+                 -> 
+                 rest_bindings, arg :: acc                
+               | _  ->  
+                 let p = Ident.rename old_param in 
+                 (p,arg) :: rest_bindings , (Lambda.Lvar p) :: acc 
+             end
+           else
+             let p = Ident.rename old_param in 
+             (p,arg) :: rest_bindings , (Lambda.Lvar p) :: acc 
+      )  ([],[]) params args in
+  let new_body = rewrite (Ext_hashtbl.of_list2 (List.rev params) (rev_new_params)) body in
+  List.fold_right
+    (fun (param, (arg : Lambda.lambda)) l -> 
+       let arg = 
+         match arg with 
+         | Lvar v -> 
+           begin 
+             match Hashtbl.find meta.ident_tbl v with 
+             | exception Not_found -> ()
+             | ident_info -> 
+               Hashtbl.add meta.ident_tbl param ident_info 
+           end;
+           arg 
+         | Lprim (Pgetglobal ident, []) -> 
+           (* It's not completeness, its to make it sound.. *)
+           Lam_compile_global.query_lambda ident meta.env 
+         (* alias meta param ident (Module (Global ident)) Strict *)
+         | Lprim (Pmakeblock (_, _, Immutable ) , ls) -> 
+           Hashtbl.replace meta.ident_tbl param 
+             (Lam_util.kind_of_lambda_block ls ); (** *)
+           arg
+         | _ -> arg in
+       Lam_util.refine_let param arg l) 
+     rest_bindings new_body
+
 
 
 

jscomp/lam_beta_reduce.mli

@@ -49,3 +49,9 @@ val propogate_beta_reduce :
 val refresh : 
   Lambda.lambda -> 
   Lambda.lambda 
+
+val propogate_beta_reduce_with_map : 
+  Lam_stats.meta ->
+  Lam_analysis.stats Ident_map.t ->
+  Ident_map.key list ->
+  Lambda.lambda -> Lambda.lambda list -> Lambda.lambda

jscomp/lam_pass_collect.ml

@@ -99,7 +99,7 @@ let collect_helper  (meta : Lam_stats.meta) (lam : Lambda.lambda)  =
       collect l
     | x -> 
         collect x ;
-        if Lambda.IdentSet.mem ident meta.export_idents then 
+        if Ident_set.mem ident meta.export_idents then 
           annotate meta rec_flag ident (Lam_stats_util.get_arity meta x ) lam
 
 

jscomp/lam_pass_exits.ml

@@ -106,7 +106,8 @@ let count_helper  (lam : Lambda.lambda) : (int, int ref) Hashtbl.t  =
     | Levent(l, _) -> count l
     | Lifused(_, l) -> count l 
 
-  and count_default sw = match sw.sw_failaction with
+  and count_default sw =
+    match sw.sw_failaction with
     | None -> ()
     | Some al ->
       let nconsts = List.length sw.sw_consts

jscomp/lam_pass_remove_alias.ml

@@ -39,7 +39,7 @@ let simplify_alias
           since we aliased k, so it's safe to remove it?
       *)
       let v = simpl l in
-      if Lambda.IdentSet.mem k meta.export_idents 
+      if Ident_set.mem k meta.export_idents 
       then 
         Llet(kind, k, g, v) 
         (* in this case it is preserved, but will still be simplified 
@@ -125,32 +125,43 @@ let simplify_alias
                 simpl (Lam_beta_reduce.propogate_beta_reduce meta params body args) 
               end
             else 
-            if lam_size < Lam_analysis.small_inline_size && 
-               (Lam_analysis.is_closed_by meta.export_idents _m
-                || not (Lambda.IdentSet.mem v meta.export_idents))
-
-            then 
-              if rec_flag = Rec then               
-                begin
-                  (* Ext_log.dwarn __LOC__ "beta rec.. %s/%d@." v.name v.stamp ; *)
-                  Lam_beta_reduce.propogate_beta_reduce meta params body args
-                end
+              if lam_size < Lam_analysis.small_inline_size then 
+                let param_fresh_map = Lam_analysis.param_map_of_list params in
+                let param_map = 
+                  Lam_analysis.free_variables meta.export_idents param_fresh_map body in
+                let old_count = List.length params in
+                let new_count = Ident_map.cardinal param_map in
+                if  
+                  (
+                    not (Ident_set.mem v meta.export_idents)
+                    || old_count = new_count
+                  )
+
+                then 
+                  if rec_flag = Rec then               
+                    begin
+                      (* Ext_log.dwarn __LOC__ "beta rec.. %s/%d@." v.name v.stamp ; *)
+                      (* Lam_beta_reduce.propogate_beta_reduce meta params body args *)
+                      Lam_beta_reduce.propogate_beta_reduce_with_map meta param_map params body args
+                    end
+                  else 
+                    begin
+                      (* Ext_log.dwarn __LOC__ "beta  nonrec..[%d] [%a]  %s/%d@."  *)
+                      (*   (List.length args)  *)
+                      (*   Printlambda.lambda body                      *)
+                      (*   v.name v.stamp ; *)
+                      simpl (Lam_beta_reduce.propogate_beta_reduce_with_map meta param_map params body args)
+
+                    end
+                else 
+                  Lapply ( simpl l1, List.map simpl args, info)
               else 
                 begin
-                  (* Ext_log.dwarn __LOC__ "beta  nonrec..[%d] [%a]  %s/%d@."  *)
-                  (*   (List.length args)  *)
-                  (*   Printlambda.lambda body                      *)
-                  (*   v.name v.stamp ; *)
-                  simpl (Lam_beta_reduce.propogate_beta_reduce meta params body args)
-
+                  (* Ext_log.dwarn __LOC__ "%s/%d: %d @."  *)
+                  (*   v.name v.stamp lam_size *)
+                  (* ;     *)
+                  Lapply ( simpl l1, List.map simpl args, info)
                 end
-            else 
-              begin
-                (* Ext_log.dwarn __LOC__ "%s/%d: %d @."  *)
-                (*   v.name v.stamp lam_size *)
-                (* ;     *)
-                Lapply ( simpl l1, List.map simpl args, info)
-              end
           else
             begin
               (* Ext_log.dwarn __LOC__ "%d vs %d @." (List.length args) (List.length params); *)

jscomp/lam_stats.ml

@@ -88,7 +88,7 @@ type ident_info = {
 type meta = {
   env : Env.t;
   filename : string ;
-  export_idents : Lambda.IdentSet.t ;
+  export_idents : Ident_set.t ;
   exports : Ident.t list ;
 
   alias_tbl : alias_tbl; 

jscomp/lam_stats.mli

@@ -90,7 +90,7 @@ type ident_info = {
 type meta = {
   env : Env.t;
   filename : string ;
-  export_idents : Lambda.IdentSet.t ;
+  export_idents : Ident_set.t ;
   exports : Ident.t list ;
   alias_tbl : alias_tbl; 
   exit_codes : int Hash_set.hashset;

jscomp/lam_stats_util.ml

@@ -252,7 +252,7 @@ let export_to_cmj
          let closed_lambda = 
            if Lam_inline_util.maybe_functor x.name 
            then
-             if Lam_analysis.is_closed lambda 
+             if Lam_analysis.is_closed lambda (* TODO: seriealize more*)
              then Some lambda
              else None
            else None in