C: Inlining support

Add a function Jib_optimize.inline which can inline functions. To make
this more efficient, we can make identifiers unique on a per-function
basis.

3 files changed, 168 insertions, 4 deletions

diff --git a/src/jib/jib_optimize.ml b/src/jib/jib_optimize.ml
index ea644bad..4e24315c 100644
--- a/src/jib/jib_optimize.ml
+++ b/src/jib/jib_optimize.ml
@@ -127,3 +127,162 @@ let flatten_cdef =
     CDEF_let (n, bindings, flatten_instrs instrs)
 
   | cdef -> cdef
+
+let unique_per_function_ids cdefs =
+  let unique_id i = function
+    | Name (id, ssa_num) -> Name (append_id id ("#u" ^ string_of_int i), ssa_num)
+    | name -> name
+  in
+  let rec unique_instrs i = function
+    | I_aux (I_decl (ctyp, id), aux) :: rest ->
+       I_aux (I_decl (ctyp, unique_id i id), aux) :: unique_instrs i (instrs_rename id (unique_id i id) rest)
+
+    | I_aux (I_init (ctyp, id, cval), aux) :: rest ->
+       I_aux (I_init (ctyp, unique_id i id, cval), aux) :: unique_instrs i (instrs_rename id (unique_id i id) rest)
+
+    | I_aux (I_block instrs, aux) :: rest ->
+       I_aux (I_block (unique_instrs i instrs), aux) :: unique_instrs i rest
+
+    | I_aux (I_try_block instrs, aux) :: rest ->
+       I_aux (I_try_block (unique_instrs i instrs), aux) :: unique_instrs i rest
+
+    | I_aux (I_if (cval, then_instrs, else_instrs, ctyp), aux) :: rest ->
+       I_aux (I_if (cval, unique_instrs i then_instrs, unique_instrs i else_instrs, ctyp), aux) :: unique_instrs i rest
+
+    | instr :: instrs -> instr :: unique_instrs i instrs
+    | [] -> []
+  in
+  let unique_cdef i = function
+    | CDEF_reg_dec (id, ctyp, instrs) -> CDEF_reg_dec (id, ctyp, unique_instrs i instrs)
+    | CDEF_type ctd -> CDEF_type ctd
+    | CDEF_let (n, bindings, instrs) -> CDEF_let (n, bindings, unique_instrs i instrs)
+    | CDEF_spec (id, ctyps, ctyp) -> CDEF_spec (id, ctyps, ctyp)
+    | CDEF_fundef (id, heap_return, args, instrs) -> CDEF_fundef (id, heap_return, args, unique_instrs i instrs)
+    | CDEF_startup (id, instrs) -> CDEF_startup (id, unique_instrs i instrs)
+    | CDEF_finish (id, instrs) -> CDEF_finish (id, unique_instrs i instrs)
+  in
+  List.mapi unique_cdef cdefs
+
+let rec frag_subst id subst = function
+  | F_id id' -> if Name.compare id id' = 0 then subst else F_id id'
+  | F_ref reg_id -> F_ref reg_id
+  | F_lit vl -> F_lit vl
+  | F_op (frag1, op, frag2) -> F_op (frag_subst id subst frag1, op, frag_subst id subst frag2)
+  | F_unary (op, frag) -> F_unary (op, frag_subst id subst frag)
+  | F_call (op, frags) -> F_call (op, List.map (frag_subst id subst) frags)
+  | F_field (frag, field) -> F_field (frag_subst id subst frag, field)
+  | F_raw str -> F_raw str
+  | F_poly frag -> F_poly (frag_subst id subst frag)
+
+let cval_subst id subst (frag, ctyp) = frag_subst id subst frag, ctyp
+
+let rec instrs_subst id subst =
+  function
+  | (I_aux (I_decl (_, id'), _) :: _) as instrs when Name.compare id id' = 0 ->
+     instrs
+
+  | I_aux (I_init (ctyp, id', cval), aux) :: rest when Name.compare id id' = 0 ->
+     I_aux (I_init (ctyp, id', cval_subst id subst cval), aux) :: rest
+
+  | (I_aux (I_reset (_, id'), _) :: _) as instrs when Name.compare id id' = 0 ->
+     instrs
+
+  | I_aux (I_reinit (ctyp, id', cval), aux) :: rest when Name.compare id id' = 0 ->
+     I_aux (I_reinit (ctyp, id', cval_subst id subst cval), aux) :: rest
+
+  | I_aux (instr, aux) :: instrs ->
+     let instrs = instrs_subst id subst instrs in
+     let instr = match instr with
+       | I_decl (ctyp, id') -> I_decl (ctyp, id')
+       | I_init (ctyp, id', cval) -> I_init (ctyp, id', cval_subst id subst cval)
+       | I_jump (cval, label) -> I_jump (cval_subst id subst cval, label)
+       | I_goto label -> I_goto label
+       | I_label label -> I_label label
+       | I_funcall (clexp, extern, fid, args) -> I_funcall (clexp, extern, fid, List.map (cval_subst id subst) args)
+       | I_copy (clexp, cval) -> I_copy (clexp, cval_subst id subst cval)
+       | I_clear (clexp, id') -> I_clear (clexp, id')
+       | I_undefined ctyp -> I_undefined ctyp
+       | I_match_failure -> I_match_failure
+       | I_end -> I_end
+       | I_if (cval, then_instrs, else_instrs, ctyp) ->
+          I_if (cval_subst id subst cval, instrs_subst id subst then_instrs, instrs_subst id subst else_instrs, ctyp)
+       | I_block instrs -> I_block (instrs_subst id subst instrs)
+       | I_try_block instrs -> I_try_block (instrs_subst id subst instrs)
+       | I_throw cval -> I_throw (cval_subst id subst cval)
+       | I_comment str -> I_comment str
+       | I_raw str -> I_raw str
+       | I_return cval -> I_return cval
+       | I_reset (ctyp, id') -> I_reset (ctyp, id')
+       | I_reinit (ctyp, id', cval) -> I_reinit (ctyp, id', cval_subst id subst cval)
+       | I_alias (clexp, cval) -> I_alias (clexp, cval_subst id subst cval)
+     in
+     I_aux (instr, aux) :: instrs
+
+  | [] -> []
+
+let rec clexp_subst id subst = function
+  | CL_id (id', ctyp) when Name.compare id id' = 0 ->
+     assert (ctyp_equal ctyp (clexp_ctyp subst));
+     subst
+  | CL_id (id', ctyp) -> CL_id (id', ctyp)
+  | CL_field (clexp, field) -> CL_field (clexp_subst id subst clexp, field)
+  | CL_addr clexp -> CL_addr (clexp_subst id subst clexp)
+  | CL_tuple (clexp, n) -> CL_tuple (clexp_subst id subst clexp, n)
+  | CL_void -> CL_void
+
+let rec find_function fid = function
+  | CDEF_fundef (fid', heap_return, args, body) :: _ when Id.compare fid fid' = 0 ->
+     Some (heap_return, args, body)
+
+  | cdef :: cdefs -> find_function fid cdefs
+
+  | [] -> None
+
+let inline cdefs should_inline instrs =
+  let inlines = ref (-1) in
+
+  let replace_return subst = function
+    | I_aux (I_funcall (clexp, extern, fid, args), aux) ->
+       I_aux (I_funcall (clexp_subst return subst clexp, extern, fid, args), aux)
+    | I_aux (I_copy (clexp, cval), aux) ->
+       I_aux (I_copy (clexp_subst return subst clexp, cval), aux)
+    | I_aux (I_alias (clexp, cval), aux) ->
+       I_aux (I_alias (clexp_subst return subst clexp, cval), aux)
+    | instr -> instr
+  in
+
+  let replace_end label = function
+    | I_aux (I_end, aux) -> I_aux (I_goto label, aux)
+    | instr -> instr
+  in
+
+  let rec inline_instr = function
+    | I_aux (I_funcall (clexp, false, function_id, args), aux) as instr when should_inline function_id ->
+       begin match find_function function_id cdefs with
+       | Some (None, ids, body) ->
+          incr inlines;
+          let inline_label = label "end_inline_" in
+          let body = List.fold_right2 instrs_subst (List.map name ids) (List.map fst args) body in
+          let body = List.map (map_instr (replace_end inline_label)) body in
+          let body = List.map (map_instr (replace_return clexp)) body in
+          I_aux (I_block (body @ [ilabel inline_label]), aux)
+       | Some (Some _, ids, body) ->
+          (* Some _ is only introduced by C backend, so we don't
+             expect it at this point. *)
+          raise (Reporting.err_general (snd aux) "Unexpected return method in IR")
+       | None -> instr
+       end
+    | instr -> instr
+  in
+
+  let rec go instrs =
+    if !inlines <> 0 then
+      begin
+        inlines := 0;
+        let instrs = List.map (map_instr inline_instr) instrs in
+        go instrs
+      end
+    else
+      instrs
+  in
+  go instrs
diff --git a/src/jib/jib_optimize.mli b/src/jib/jib_optimize.mli
index beffa81e..78759d08 100644
--- a/src/jib/jib_optimize.mli
+++ b/src/jib/jib_optimize.mli
@@ -61,3 +61,6 @@ val optimize_unit : instr list -> instr list
 val flatten_instrs : instr list -> instr list
 val flatten_cdef : cdef -> cdef
 
+val unique_per_function_ids : cdef list -> cdef list
+
+val inline : cdef list -> (Ast.id -> bool) -> instr list -> instr list
diff --git a/src/jib/jib_util.ml b/src/jib/jib_util.ml
index 78eca13b..22b983ff 100644
--- a/src/jib/jib_util.ml
+++ b/src/jib/jib_util.ml
@@ -81,7 +81,7 @@ let ifuncall ?loc:(l=Parse_ast.Unknown) clexp id cvals =
 
 let iextern ?loc:(l=Parse_ast.Unknown) clexp id cvals =
   I_aux (I_funcall (clexp, true, id, cvals), (instr_number (), l))
-  
+
 let icall ?loc:(l=Parse_ast.Unknown) clexp extern id cvals =
   I_aux (I_funcall (clexp, extern, id, cvals), (instr_number (), l))
 
@@ -108,11 +108,13 @@ let itry_block ?loc:(l=Parse_ast.Unknown) instrs =
 
 let ithrow ?loc:(l=Parse_ast.Unknown) cval =
   I_aux (I_throw cval, (instr_number (), l))
+
 let icomment ?loc:(l=Parse_ast.Unknown) str =
   I_aux (I_comment str, (instr_number (), l))
 
 let ilabel ?loc:(l=Parse_ast.Unknown) label =
   I_aux (I_label label, (instr_number (), l))
+
 let igoto ?loc:(l=Parse_ast.Unknown) label =
   I_aux (I_goto label, (instr_number (), l))
 
@@ -152,7 +154,7 @@ module NameMap = Map.Make(Name)
 let current_exception = Current_exception (-1)
 let have_exception = Have_exception (-1)
 let return = Return (-1)
-               
+
 let name id = Name (id, -1)
 
 let rec frag_rename from_id to_id = function
@@ -269,7 +271,7 @@ let string_of_name ?zencode:(zencode=true) =
      "return" ^ ssa_num n
   | Current_exception n ->
      "(*current_exception)" ^ ssa_num n
-    
+
 let rec string_of_fragment ?zencode:(zencode=true) = function
   | F_id id -> string_of_name ~zencode:zencode id
   | F_ref id -> "&" ^ string_of_name ~zencode:zencode id
@@ -503,7 +505,7 @@ let pp_name id =
   string (string_of_name ~zencode:false id)
 
 let pp_ctyp ctyp =
-  string (full_string_of_ctyp ctyp |> Util.yellow |> Util.clear)
+  string (string_of_ctyp ctyp |> Util.yellow |> Util.clear)
 
 let pp_keyword str =
   string ((str |> Util.red |> Util.clear) ^ " ")