change transform_inline to be aware of checkpoints

To properly support osr across inlining the inlining pass
needs to know if there is a checkpoint after the function
call at the caller site.

This patch refactors compute_inlining_order to not just
return the order, but in fact an annotated graph that
contains all the information for the inliner to go forward.

At the core there is the `inlining_candidate` which contains
all information to apply one inlining. The inlinings are
performed recursively, all in one go.

At the end a new version is returned which can be added to
the first caller function.

Currently inlinings are skipped when changes to the osr
points would be needed.

Author: o-

transform_inline.ml

@@ -9,8 +9,19 @@ type inlining_input = {
   instrs : instructions;
 }
 
-let as_inlining_input (func : afunction) (version : version) : inlining_input =
-  { formals = Analysis.as_var_list func.formals; instrs = version.instrs }
+let as_inlining_input (func : afunction) : inlining_input =
+  { formals = Analysis.as_var_list func.formals; instrs = (active_version func).instrs }
+
+type inlining_candidate = {
+  pos : pc;
+  target : afunction;
+  ret : variable;
+  args : argument list;
+  osr : osr_frame_map list option;
+  next : inlining_site;
+}
+and inlining_site = inlining_candidate list
+
 
 (* FUNCTION INLINING *)
 (* Given a program, inline the functions in it to the maximum possible extent.
@@ -31,27 +42,6 @@ let inline ({main; functions} as orig_prog : program) : program option =
     VarSet.union (all_declared_vars instrs) (VarSet.of_list formals)
   in
 
-  (* Given a function and an array of instructions, generate a fresh version
-     for the function with these instructions. *)
-  let add_version (func : afunction) (instrs : instructions)  =
-    let label = fresh_version_label func "inlined_version" in
-    let version = {instrs = instrs; label = label; annotations = None} in
-    {func with body = version :: func.body}
-  in
-
-  (* Given `instructions` and the function identifier, extract the location,
-     return variable and arguments from the first callsite where a call is made
-     to that function *)
-  let extract_callsite instrs fun_ref : (pc * variable * argument list) =
-    let rec at pc =
-      match[@warning "-4"] instrs.(pc) with
-      | Call (x, (Simple (Constant (Fun_ref f))), es) ->
-        if f = fun_ref then (pc, x, es) else at (pc + 1)
-      | _ -> at (pc + 1)
-    in
-    at 0
-  in
-
   (* Replace variables in `callee` instructions so that they don't match the
      `caller` variables. The formals list of the callee is also updated
      accordingly *)
@@ -103,8 +93,7 @@ let inline ({main; functions} as orig_prog : program) : program option =
 
   (* Replace labels in `callee` instructions so that they don't match the
      `caller` labels *)
-  let replace_labels caller ({formals; instrs} as callee) =
-    let used_labels = extract_labels caller.instrs in
+  let replace_labels used_labels ({formals; instrs} as callee) =
     let callee_labels = LabelSet.elements (extract_labels callee.instrs) in
     let replacements = Edit.fresh_many used_labels callee_labels in
     let mapper instr =
@@ -116,7 +105,8 @@ let inline ({main; functions} as orig_prog : program) : program option =
       | Branch (e, l1, l2) -> Branch (e, replace l1, replace l2)
       | i -> i
     in
-    {callee with instrs = Array.map mapper instrs}
+    let new_labels = LabelSet.of_list (snd (List.split replacements)) in
+    {callee with instrs = Array.map mapper instrs}, LabelSet.union used_labels new_labels
   in
 
   (* Inserts the header for the inlined callee body. Assigns all the formals to
@@ -141,10 +131,10 @@ let inline ({main; functions} as orig_prog : program) : program option =
 
   (* Transforms the callee instructions to a form that can be substituted
      inside the caller.*)
-  let compose caller callee ret_var arguments : inlining_input =
-    let callee = callee
-                 |> replace_vars caller
-                 |> replace_labels caller
+  let compose caller callee used_labels ret_var arguments =
+    let callee, used_labels = callee
+                            |> replace_vars caller
+                            |> replace_labels used_labels
     in
     (* It is important to generate return label and result variable after
        replacing variables and labels. Otherwise there is a good chance that
@@ -154,108 +144,103 @@ let inline ({main; functions} as orig_prog : program) : program option =
        callee with `res_1` and `res_1` is already used for the result variable.
        This problem will not happen if we `replace_vars` before generating a
        fresh `res_var`. *)
-    let ret_lab =
-      fresh_label (Array.append callee.instrs caller.instrs) "inl"
-    in
+    let ret_lab = fresh used_labels "inl" in
+    let used_labels = LabelSet.add ret_lab used_labels in
+
     let res_var =
       fresh (VarSet.union (function_vars callee) (function_vars caller)) "res"
     in
-    callee
-    |> replace_returns res_var ret_lab
-    |> insert_prologue res_var arguments
-    |> insert_epilogue res_var ret_var ret_lab
-  in
-
-  (* Given the caller and callee identifiers, generate a fresh version of caller
-     with the callee inlined.*)
-  let inline_pair (caller_id : identifier)
-                  (callee_id : identifier)
-                  (prog : program) =
-    let caller = lookup_fun prog caller_id in
-    let caller' = as_inlining_input caller (List.hd caller.body) in
-    let callee = lookup_fun prog callee_id in
-    let callee' = as_inlining_input callee (List.hd callee.body) in
-    let (pc, ret_var, arguments) = extract_callsite caller'.instrs callee_id in
-    let callee'' = compose caller' callee' ret_var arguments in
-    let instrs, _ = subst caller'.instrs pc 1 callee''.instrs in
-    add_version caller instrs
+    let res = callee
+            |> replace_returns res_var ret_lab
+            |> insert_prologue res_var arguments
+            |> insert_epilogue res_var ret_var ret_lab in
+    (res, used_labels)
   in
 
-(* This function computes an order for inlining of the entire program.
-   Given a call graph starting from main, it descends as deep as possible into
-   the call chain and stops when it encounters an edge which leads to recursion.
-   Using this depth-first approach, it generates all caller-callee pairs with
-   the depth at which this pair was generated. This depth is used to sort the
-   list of these pairs such that the inlining happens in a bottom up fashion.
-   If a caller-callee pair is already encountered, then it is not included
-   again. This is because once the callee has been inlined in the caller at a
-   given callsite, then the callsite no longer exists. However, multiple
-   callsites with the same target will result in as many caller-callee pairs.
-   Note - The word reduced is used in the general sense of something being
-   simplified. If a function has been analyzed for inlining, its considered
-   reduced. In reality, inlining expands this function.
-   The accumulator `acc` is composed of the following components -
-   - `ord` - The list of (caller name, callee name, depth) triplets
-   - `caller` - The caller object
-   - `vis` - The set of visited function names
-   - `red` - The set of reduced (already analyzed for inlining) function names
-   - `dep` - The depth of this function from `main`
-*)
-  let compute_inlining_order (prog : program) : (identifier * identifier) list =
-    let rec inspect_instr ((ord, caller, vis, red, dep) as acc) instr =
-      match[@warning "-4"] instr with
+  (* This function computes an order for inlining of the entire program.
+     Given a call graph starting from main, it descends as deep as possible into
+     the call chain and stops when it encounters an edge which leads to recursion.
+  *)
+  let rec compute_inline_order func seen : inlining_site =
+    let instrs = (active_version func).instrs in
+    let inlinings = ref [] in
+    let visit_instr pc =
+      assert (pc+1 < Array.length instrs);
+      match[@warning "-4"] instrs.(pc) with
       | Call (x, (Simple (Constant (Fun_ref f))), es) ->
-        (* If `f` is already visited in this branch, then there is recursion,
-           so don't add this pair to accumulator.*)
-        if (VarSet.mem f vis) then acc
-        (* If `f` is already analyzed for inlining, then add it but don't go
-           inside it.*)
-        else if (VarSet.mem f red)
-        then ((caller.name, f, dep) :: ord, caller, vis, red, dep)
-        (* If `f` is neither visited in this branch, nor has it been analyzed
-           for inlining before, then descend into it and analyze it.*)
-        else
-          let callee = lookup_fun prog f in
-          let (callee_ord, _, _, callee_red, _) =
-            add_callees callee vis red (dep + 1)
-          in
-          let new_ord = (caller.name, f, dep) :: ord @ callee_ord in
-          let new_red = VarSet.union red callee_red in
-          (* Note that visited set is only maintained for a branch while
-             reduced set is carried across branches. *)
-          (new_ord, caller, vis, new_red, dep)
-      (* Any other instruction returns accumulator unchanged. *)
-      | _ -> acc
-    and add_callees caller vis red dep =
-      (* The `caller` is being visited and reduced. So add it to the two sets.*)
-      let vis = VarSet.add caller.name vis in
-      let red = VarSet.add caller.name red in
-      Array.fold_left
-        inspect_instr
-        ([], caller, vis, red, dep)
-        (List.hd caller.body).instrs
+        if LabelSet.mem f seen then () else begin
+          (* If the call is followed by a checkpoint we store it's
+           * osr map to be able to concatenate it to the inlinee's *)
+          let checkpoint = (match[@warning "-4"] instrs.(pc+1) with
+                            | Osr {varmap; frame_maps} -> Some []
+                            | _ -> None) in
+          let seen = LabelSet.add f seen in
+          let callee = lookup_fun orig_prog f in
+          let next = compute_inline_order callee seen in
+          let inlining = { pos = pc; target = callee; ret = x; args = es; osr = checkpoint; next } in
+          inlinings := inlining :: !inlinings
+        end
+      | _ -> ()
     in
-    (* Sort (caller, callee, depth) triplets in decreasing order of depth. *)
-    let comp (_, _, d1) (_, _, d2) = d2 - d1 in
-    let (ord, _, _, _, _) = add_callees main VarSet.empty VarSet.empty 1 in
-    List.map
-      (fun (caller_id, callee_id, _ ) -> (caller_id, callee_id))
-      (List.sort comp ord)
+    for pc = 0 to (Array.length instrs) - 2 do
+      visit_instr pc
+    done;
+    !inlinings
+  in
+
+  let needs_osr =
+    let is_osr = function[@warning "-4"]
+      | Osr _ -> true | _ -> false in
+    Array.exists is_osr
+  in
+
+  let rec apply_inlinings func inlinings =
+    let cur = as_inlining_input func in
+    if inlinings = []
+    then cur
+    else
+      let used_labels = ref (extract_labels cur.instrs) in
+      let get {target; next; pos; ret; args; osr} =
+        let apply next = if next = []
+                         then as_inlining_input target
+                         else apply_inlinings target next in
+        let callee = apply next in
+        match needs_osr callee.instrs, osr with
+        | false, _ ->
+            let inlinee, new_used = compose cur callee !used_labels ret args in
+            used_labels := new_used;
+            (pos, 1, inlinee.instrs)
+        | true, Some osr ->
+            (* TODO(osr): Here we need to update the inlinee osr points to create the
+             *            additional osr frame. *)
+            (pos, 0, [| |])
+        | true , None ->
+            (* The callee needs to osr but the caller does not have a safepoint
+             * after the call. We can't do anything. *)
+            (pos, 0, [| |])
+      in
+      let to_inline = List.map get inlinings in
+      let instrs, _ = Edit.subst_many cur.instrs to_inline in
+      { cur with instrs }
   in
 
-  (* Given a list of caller - callee pairs, inline the callee inside the caller
-     until the list is exhausted. Order matters, as only this order will result
-     in a maximally inlined program excluding recursive calls. The final program
-     will only have recursive calls left to be inlined. *)
-  let inline_with order prog =
-    List.fold_left
-      (fun p (caller, callee) -> replace_fun p (inline_pair caller callee p))
-      prog
-      order
+  let inline_at func =
+    let inline_order = compute_inline_order func LabelSet.empty in
+    (* If there are no caller-callee pairs to inline, return `None`, else return
+       the completely inlined program *)
+    if inline_order = [] then None
+    else
+      let result = apply_inlinings func inline_order in
+      let label = fresh_version_label func "inlined_version" in
+      Some { label; annotations = None; instrs = result.instrs }
   in
 
-  (* If there are no caller-callee pairs to inline, return `None`, else return
-     the completely inlined program *)
-  let inline_order = compute_inlining_order orig_prog in
-  if List.length inline_order = 0 then None
-  else Some (inline_with inline_order orig_prog)
+  (* Starting from main inline all the way down *)
+  match inline_at orig_prog.main with
+  | None -> None
+  | Some v ->
+    Some { orig_prog with
+      main = { orig_prog.main with
+        body = v :: orig_prog.main.body
+      }
+    }