Rewrite inlining pass

compiler/lib/code.ml

@@ -791,6 +791,12 @@ let fold_closures_outermost_first { start; blocks; _ } f accu =
   let accu = f None [] (start, []) None accu in
   visit blocks start f accu
 
+let rec last_instr l =
+  match l with
+  | [] | [ Event _ ] -> None
+  | [ i ] | [ i; Event _ ] -> Some i
+  | _ :: rem -> last_instr rem
+
 let eq p1 p2 =
   p1.start = p2.start
   && Addr.Map.equal

compiler/lib/code.mli

@@ -311,6 +311,9 @@ val traverse :
 val preorder_traverse :
   fold_blocs_poly -> (Addr.t -> 'c -> 'c) -> Addr.t -> block Addr.Map.t -> 'c -> 'c
 
+val last_instr : instr list -> instr option
+(** Last instruction of a block body, ignoring events *)
+
 val prepend : program -> instr list -> program
 
 val empty : program

compiler/lib/config.ml

@@ -135,7 +135,7 @@ module Param = struct
     p ~name:"switch_size" ~desc:"set the maximum number of case in a switch" (int 60)
 
   let inlining_limit =
-    p ~name:"inlining-limit" ~desc:"set the size limit for inlining" (int 200)
+    p ~name:"inlining-limit" ~desc:"set the size limit for inlining" (int 150)
 
   let tailcall_max_depth =
     p

compiler/lib/driver.ml

@@ -54,12 +54,12 @@ let deadcode p =
   let r, live_vars = deadcode' p in
   Deadcode.remove_empty_blocks ~live_vars r
 
-let inline p =
+let inline aggressive p =
   if Config.Flag.inline () && Config.Flag.deadcode ()
   then (
     let p, live_vars = deadcode' p in
     if debug () then Format.eprintf "Inlining...@.";
-    Inline.f p live_vars)
+    Inline.f ~aggressive p live_vars)
   else p
 
 let specialize_1 (p, info) =
@@ -159,36 +159,38 @@ let rec loop max name round i (p : 'a) : 'a =
 
 (* o1 *)
 
-let o1 : 'a -> 'a =
+let round opt : 'a -> 'a =
   print
   +> tailcall
   +> flow
   +> specialize
   +> eval
-  +> inline (* inlining may reveal new tailcall opt *)
+  +> inline opt (* inlining may reveal new tailcall opt *)
   +> deadcode
   +> tailcall
   +> phi
   +> flow
   +> specialize
   +> eval
-  +> inline
+  +> inline opt
   +> deadcode
   +> print
   +> flow
   +> specialize
   +> eval
-  +> inline
+  +> inline opt
   +> deadcode
   +> phi
 
+let o1 = round false
+
 (* o2 *)
 
-let o2 = loop 10 "o1" o1 1 +> print
+let o2 = loop 10 "o1" (round true) 1 +> print
 
 (* o3 *)
 
-let o3 = loop 10 "o1" o1 1 +> print
+let o3 = loop 10 "o1" (round true) 1 +> print
 
 let generate
     ~exported_runtime
@@ -438,7 +440,9 @@ let check_js js =
     Javascript.IdentSet.fold
       (fun x acc ->
         match x with
-        | V _ -> assert false
+        | V y ->
+            Format.eprintf "ZZZ %a@." Code.Var.print y;
+            assert false
         | S { name = Utf8 x; _ } -> StringSet.add x acc)
       free
       StringSet.empty

compiler/lib/duplicate.ml

@@ -19,29 +19,65 @@
 open! Stdlib
 open Code
 
-let closure p ~bound_vars ~f ~params ~cont:(pc, args) =
-  let s =
-    Subst.from_map
-      (Var.Set.fold (fun x s -> Var.Map.add x (Var.fork x) s) bound_vars Var.Map.empty)
-  in
-  let free_pc, m =
+let bound_variables { blocks; _ } ~f ~params ~cont:(pc, _) =
+  let bound_vars = ref Var.Map.empty in
+  let add_var x = bound_vars := Var.Map.add x (Var.fork x) !bound_vars in
+  List.iter ~f:add_var (f :: params);
+  let rec traverse blocks pc =
     Code.traverse
-      { fold = Code.fold_children }
-      (fun pc (pc', m) -> pc' + 1, Addr.Map.add pc pc' m)
+      { fold = fold_children }
+      (fun pc _ ->
+        let block = Addr.Map.find pc blocks in
+        Freevars.iter_block_bound_vars add_var block;
+        List.iter
+          ~f:(fun i ->
+            match i with
+            | Let (_, Closure (params, (pc', _), _)) ->
+                List.iter ~f:add_var params;
+                traverse blocks pc'
+            | _ -> ())
+          block.body)
       pc
-      p.blocks
-      (p.free_pc, Addr.Map.empty)
+      blocks
+      ()
+  in
+  traverse blocks pc;
+  !bound_vars
+
+let rec blocks_to_rename p pc lst =
+  Code.traverse
+    { fold = Code.fold_children }
+    (fun pc lst ->
+      let block = Addr.Map.find pc p.blocks in
+      List.fold_left
+        ~f:(fun lst i ->
+          match i with
+          | Let (_, Closure (_, (pc', _), _)) -> blocks_to_rename p pc' lst
+          | _ -> lst)
+        ~init:(pc :: lst)
+        block.body)
+    pc
+    p.blocks
+    lst
+
+let closure p ~f ~params ~cont =
+  let s = Subst.from_map (bound_variables p ~f ~params ~cont) in
+  let pc, args = cont in
+  let blocks = blocks_to_rename p pc [] in
+  let free_pc, m =
+    List.fold_left
+      ~f:(fun (pc', m) pc -> pc' + 1, Addr.Map.add pc pc' m)
+      ~init:(p.free_pc, Addr.Map.empty)
+      blocks
   in
   let blocks =
-    Code.traverse
-      { fold = Code.fold_children }
-      (fun pc blocks ->
+    List.fold_left
+      ~f:(fun blocks pc ->
         let b = Addr.Map.find pc blocks in
         let b = Subst.Including_Binders.And_Continuations.block m s b in
         Addr.Map.add (Addr.Map.find pc m) b blocks)
-      pc
-      p.blocks
-      p.blocks
+      ~init:p.blocks
+      blocks
   in
   let p = { p with blocks; free_pc } in
   p, s f, List.map ~f:s params, (Addr.Map.find pc m, List.map ~f:s args)

compiler/lib/duplicate.mli

@@ -18,14 +18,11 @@
 
 val closure :
      Code.program
-  -> bound_vars:Code.Var.Set.t
   -> f:Code.Var.t
   -> params:Code.Var.t list
   -> cont:int * Code.Var.t list
   -> Code.program * Code.Var.t * Code.Var.t list * (int * Code.Var.t list)
-(** Given a program and a closure [f] -- defined by its name, parameters, and its
-    continuation --, return a program in which the body of [f] has been updated with fresh
-    variable names to replace elements of [bound_vars]. Also returns the new name of [f]
-    (fresh if [f] is in [bound_vars]), and the similarly substituted parameter list and
-    continuation.
-  *)
+(** Given a program and a closure [f] -- defined by its name,
+    parameters, and its continuation --, return a program with a copy
+    of [f]. Also returns the new name of [f], and the similarly
+    substituted parameter list and continuation. *)

compiler/lib/effects.ml

@@ -142,13 +142,6 @@ let dominance_frontier g idom =
     g.preds;
   frontiers
 
-(* Last instruction of a block, ignoring events *)
-let rec last_instr l =
-  match l with
-  | [] -> None
-  | [ i ] | [ i; Event _ ] -> Some i
-  | _ :: rem -> last_instr rem
-
 (* Split a block, separating the last instruction from the preceeding
    ones, ignoring events *)
 let block_split_last xs =
@@ -210,7 +203,7 @@ let compute_needed_transformations ~cfg ~idom ~cps_needed ~blocks ~start =
       let block = Addr.Map.find pc blocks in
       (match block.branch with
       | Branch (dst, _) -> (
-          match last_instr block.body with
+          match Code.last_instr block.body with
           | Some (Let (x, e))
             when effect_primitive_or_application e && Var.Set.mem x cps_needed ->
               (* The block after a function application that needs to