Added function inlining; bugfix of nested_closure

src/riscv_generate.ml

@@ -908,7 +908,6 @@ let rec do_convert ssa (expr: Mcore.expr) =
 
           (* Let the pointer point to beginning of data, rather than the length section *)
           let beginning = new_temp T_bytes in
-          Basic_vec.push ssa (ExtArray { label; values });
           Basic_vec.push ssa (AssignLabel { rd = beginning; imm = label; });
           Basic_vec.push ssa (Addi { rd; rs = beginning; imm = 4 })
 
@@ -923,7 +922,6 @@ let rec do_convert ssa (expr: Mcore.expr) =
 
           (* Let the pointer point to beginning of data, rather than the length section *)
           let beginning = new_temp T_bytes in
-          Basic_vec.push ssa (ExtArray { label; values });
           Basic_vec.push ssa (AssignLabel { rd = beginning; imm = label; });
           Basic_vec.push ssa (Addi { rd; rs = beginning; imm = 4 })
 

src/riscv_opt.ml

@@ -26,9 +26,15 @@ let basic_blocks = Hashtbl.create 1024
 (** The exit block(s) for each function `fn`, i.e. whose final instruction is `return`. *)
 let exit_fn = Hashtbl.create 256
 
+(** The parameters of each function. *)
+let (params: (string, var list) Hashtbl.t) = Hashtbl.create 256
+
 (** Get the basic block with label `name`. *)
 let block_of name = Hashtbl.find basic_blocks name
 
+(** Get the body of a basic block. *)
+let body_of name = (block_of name).body |> Basic_vec.to_list
+
 (**
 Builds control flow graph.
 
@@ -219,7 +225,7 @@ let liveness_analysis fn =
     Hashtbl.add phiuses name !phiuse
   ) blocks;
 
-  (* Keep doing until reaches fixed point *)
+  (* Keep doing until fixed point is reached *)
   let rec iterate worklist =
     let last_item = Basic_vec.pop_opt worklist in
     match last_item with

src/riscv_opt_gather.ml

@@ -1,10 +1,14 @@
 (** Gathers all optimizations. *)
 open Riscv_opt
+open Riscv_ssa
 
 let opt ssa =
+  List.iter (fun top -> match top with
+  | FnDecl { fn; args } -> Hashtbl.add params fn args
+  | _ -> ()) ssa;
   iter_fn2 build_cfg ssa;
 
-  (* Peephole *)
+  Riscv_opt_inline.inline ssa;
   Riscv_opt_peephole.peephole ssa;
 
   let s = map_fn ssa_of_cfg ssa in

src/riscv_opt_inline.ml

@@ -0,0 +1,90 @@
+(**
+Inlines a function.
+
+Current inline condition:
+  1. The function has only one basic block;
+  2. The function is a leaf (i.e. not calling other functions)
+*)
+
+open Riscv_ssa
+open Riscv_opt
+
+(**
+Whether each function can get inlined.
+This cache is to avoid repeatedly scanning through the same function body.
+*)
+let inline_cache = Hashtbl.create 64
+
+let is_single fn = 
+  List.length (get_blocks fn) = 1
+
+let is_leaf fn =
+  let blocks = get_blocks fn in
+
+  let has_call = List.exists (fun block -> 
+    let body = body_of block in
+
+    List.exists (fun x -> match x with
+    | Call _ 
+    | CallIndirect _
+    | CallExtern _ -> true
+
+    | _ -> false) body
+  ) blocks in
+
+  not has_call
+
+let can_inline fn =
+  if Hashtbl.mem inline_cache fn then
+    Hashtbl.find inline_cache fn
+  else
+    let inlineable = is_leaf fn && is_single fn in
+    Hashtbl.add inline_cache fn inlineable;
+    inlineable
+
+(** Scan through body of `fn` and see if any function call can be inlined. *)
+let do_inline fn = 
+  let blocks = get_blocks fn in
+  List.iter (fun block ->
+    let body = body_of block in
+    let new_body = Basic_vec.empty () in
+
+    List.iter (fun x -> match x with
+    | Call { rd; fn; args } ->
+        (* Check if the function is inlineable *)
+        let able = can_inline fn in
+        if able then (
+          (* As `fn` is a single block, the content of it is the whole function body *)
+          let fn_body = body_of fn in
+
+          (* Record the correspondence between arguments and parameters *)
+          let subst = Hashtbl.create 16 in
+          List.iter2 (fun arg param ->
+            Hashtbl.add subst param arg
+          ) args (Hashtbl.find params fn);
+
+          (* Begin substitution *)
+          let replace var =
+            if Hashtbl.mem subst var then Hashtbl.find subst var
+            else var
+          in
+          let fn_new = List.map (fun x -> reg_map replace replace x) fn_body in
+
+          (* Convert `return` to an assignment operation *)
+          let fn_core = List.map
+            (fun x -> match x with
+            | Return return -> Assign { rd; rs = return }
+            | x -> x) fn_new
+          in
+
+          (* Add it to the original function body *)
+          Basic_vec.append new_body (Basic_vec.of_list fn_core))
+        else Basic_vec.push new_body x
+
+    | x -> Basic_vec.push new_body x) body;
+
+    (block_of block).body <- new_body
+  ) blocks
+
+let inline ssa = 
+  iter_fn do_inline ssa

src/riscv_opt_peephole.ml

@@ -64,7 +64,7 @@ let to_itype fn =
   let blocks = get_blocks fn in
 
   let convert block = 
-    let body = (block_of block).body |> Basic_vec.to_list in
+    let body = body_of block in
     List.map (fun x -> match x with
     | Add { rd; rs1; rs2 } ->
         if good rs1 then
@@ -119,7 +119,7 @@ let remove_dead_variable fn =
   let liveness = liveness_analysis fn in
 
   let remove block = 
-    let body = (block_of block).body |> Basic_vec.to_list in
+    let body = body_of block in
 
     (* Variables alive at the end of block *)
     let alive = Hashtbl.find liveness block in

src/riscv_ssa.ml

@@ -399,7 +399,7 @@ let rec reg_map fd fs t = match t with
 | AssignLabel { rd; imm } -> AssignLabel { rd = fd rd; imm; }
 | Assign { rd; rs } -> Assign { rd = fd rd; rs = fs rs }
 | Load { rd; rs; offset; byte } -> Load { rd = fd rd; rs = fs rs; offset; byte }
-| Store { rd; rs; offset; byte } -> Load { rd = fs rd; rs = fs rs; offset; byte }
+| Store { rd; rs; offset; byte } -> Store { rd = fs rd; rs = fs rs; offset; byte }
 | Jump label -> Jump label
 | Branch { cond; ifso; ifnot } -> Branch { cond = fs cond; ifso; ifnot }
 | Label label -> Label label

test/src/nested_closure/nested_closure.ans

@@ -2,3 +2,4 @@
 0
 0
 0
+()