Fix SSA Init Bug and Introduce VirtASM Stage

.ocamlformat

@@ -0,0 +1,2 @@
+profile = janestreet
+version = 0.27.0

src/driver_util.ml

@@ -246,6 +246,8 @@ let riscv_gen (core : Mcore.t) =
   core
   |> Riscv_generate.ssa_of_mcore
   |> Riscv_opt_gather.opt
+  |> Riscv_virtasm_generate.virtasm_of_ssa
+  |> Riscv_reg_alloc.reg_alloc
   |> Riscv.generate
 
 let link_core ~(shrink_wasm : bool) ~(elim_unused_let : bool)

src/riscv_generate.ml

@@ -38,6 +38,9 @@ let variants = Hashtbl.create 64
 (** The enum type that each variant belongs to. *)
 let belong = Hashtbl.create 64
 
+(** Init function exprs**)
+let init_exprs = Basic_vec.empty ()
+
 (** Get offset of the `pos`-th field in the record type called `name`. *)
 (** Note that there are n+1 fields for offsets in a n-field variant; the last one is the total size. *)
 let offsetof ty pos = Hashtbl.find offset_table (ty, pos)
@@ -506,7 +509,7 @@ let update_types ({ defs; _ }: Mtype.defs) =
           let sizes = List.map (fun x -> sizeof x) types in
           let offsets = 0 :: Basic_lst.cumsum sizes in
 
-          (* Record the correspondence between name and index *)
+          (* Record the correspondence between nameand index *)
           let variant_name = Printf.sprintf "%s.%s" name tag.name_ in
           Hashtbl.add belong variant_name name;
           Basic_vec.push tag_offsets offsets
@@ -1313,6 +1316,11 @@ let generate_vtables () =
 Converts given `expr` into a list of SSA instructions,
 along with the variable in which the result of this expression is stored.
 *)
+let convert_expr_no_ret (expr: Mcore.expr) =
+  let ssa = Basic_vec.empty () in
+  let _ = do_convert ssa expr in
+  Basic_vec.map_into_list ssa (fun x -> x)
+
 let convert_expr (expr: Mcore.expr) =
   let ssa = Basic_vec.empty () in
   let return = do_convert ssa expr in
@@ -1448,7 +1456,6 @@ let convert_lambda (expr: Mcore.expr) =
 
   | w -> w
 
-
 (** Store captured variables for each closure in `captured` *)
 let process_closure ((fn: Mcore.fn), (name: Ident.t)) = 
   let free_ident = Mcore_util.free_vars ~exclude:(Ident.Set.singleton name) fn in
@@ -1484,6 +1491,11 @@ let convert_toplevel _start (top: Mcore.top_item) =
   in
 
   match top with
+  (* Init function *)
+  | Ctop_expr { expr; _ } ->
+    let expr = convert_expr_no_ret expr in
+    Basic_vec.append init_exprs @@ Basic_vec.of_list expr;
+    []
   | Ctop_fn { binder; func; export_info_; _ } ->
       let fn = Ident.to_string binder in
       let args = List.map var_of_param func.params in
@@ -1525,8 +1537,6 @@ let convert_toplevel _start (top: Mcore.top_item) =
       | Inline_code_text _
       | Inline_code_sexp _ -> failwith "RISC-V target does not support inline WASM"
       | Import _ -> failwith "riscv_ssa.ml: import should have been eliminated in link stage")
-
-  | _ -> failwith "TODO: riscv_ssa.ml: don't know this toplevel"
 
 let find_functions (top: Mcore.top_item) =
   match top with
@@ -1589,6 +1599,7 @@ let ssa_of_mcore (core: Mcore.t) =
 
   (* Add _start *)
   let unused = new_temp Mtype.T_unit in
+  Basic_vec.append _start init_exprs;
   Basic_vec.push _start (Call { rd = unused; fn = "main"; args = [] });
   Basic_vec.push _start (Return unused);
 

src/riscv_reg.ml

@@ -0,0 +1,199 @@
+(** Registers for RV64GC **)
+
+(* Define a label type*)
+type label_t = string
+
+(**
+Defines an immediate value type.
+
+Immediate values can either be an integer (`IntImm`) or a floating-point number (`FloatImm`).
+*)
+type imm_t =
+  | IntImm of int  (* Integer immediate value *)
+  | FloatImm of float  (* Floating-point immediate value *)
+
+(**
+Defines a slot type for both virtual and physical registers.
+
+This type encapsulates different kinds of registers, including general-purpose registers, floating-point registers,
+and specific slots representing values like `Unit` (no return value).
+*)
+type slot_t =
+  | Unit  (* Represents no return value, often used in function calls or returns *)
+  | Slot of int  
+  | FSlot of int
+  | Reg of reg_t 
+  | FReg of freg_t
+
+and ret_type = 
+  | IntRet
+  | FloatRet
+  | UnitRet
+
+and reg_t =
+  | Zero (* zero register *)
+  | Ra (* caller return address *)
+  | Sp (* caller (S0) stack pointer *)
+  | Gp (* global pointer *)
+  | Tp (* thread pointer *)
+  | T0 (* caller temporary register *)
+  | T1
+  | T2
+  | Fp (* callee stack bottom register *)
+  | S1 (* callee saved register *)
+  | A0 (* caller argument register *)
+  | A1
+  | A2
+  | A3
+  | A4
+  | A5
+  | A6
+  | A7
+  | S2 (* callee saved register *)
+  | S3
+  | S4
+  | S5
+  | S6
+  | S7
+  | S8
+  | S9
+  | S10
+  | S11
+  | T3 (* caller temporary register *)
+  | T4
+  | T5
+  | T6 (* caller swap register *)
+
+and freg_t =
+  | Ft0 (* caller floating-point temporary register *)
+  | Ft1
+  | Ft2
+  | Ft3
+  | Ft4
+  | Ft5
+  | Ft6
+  | Ft7
+  | Fs0 (* callee floating-point saved register *)
+  | Fs1
+  | Fa0 (* caller floating-point argument register *)
+  | Fa1
+  | Fa2
+  | Fa3
+  | Fa4
+  | Fa5
+  | Fa6
+  | Fa7
+  | Fs2 (* callee floating-point saved register *)
+  | Fs3
+  | Fs4
+  | Fs5
+  | Fs6
+  | Fs7
+  | Fs8
+  | Fs9
+  | Fs10
+  | Fs11
+  | Ft8
+  | Ft9
+  | Ft10
+  | Ft11 (* caller swap floating-point register *)
+
+(* Convert reg_t to string representation *)
+let reg_to_string r =
+  match r with
+  | Zero -> "zero"
+  | Ra -> "ra"
+  | Sp -> "sp"
+  | Gp -> "gp"
+  | Tp -> "tp"
+  | T0 -> "t0"
+  | T1 -> "t1"
+  | T2 -> "t2"
+  | Fp -> "s0"
+  | S1 -> "s1"
+  | A0 -> "a0"
+  | A1 -> "a1"
+  | A2 -> "a2"
+  | A3 -> "a3"
+  | A4 -> "a4"
+  | A5 -> "a5"
+  | A6 -> "a6"
+  | A7 -> "a7"
+  | S2 -> "s2"
+  | S3 -> "s3"
+  | S4 -> "s4"
+  | S5 -> "s5"
+  | S6 -> "s6"
+  | S7 -> "s7"
+  | S8 -> "s8"
+  | S9 -> "s9"
+  | S10 -> "s10"
+  | S11 -> "s11"
+  | T3 -> "t3"
+  | T4 -> "t4"
+  | T5 -> "t5"
+  | T6 -> "t6"
+;;
+
+(* Convert freg_t to string representation *)
+let freg_to_string fr =
+  match fr with
+  | Ft0 -> "ft0"
+  | Ft1 -> "ft1"
+  | Ft2 -> "ft2"
+  | Ft3 -> "ft3"
+  | Ft4 -> "ft4"
+  | Ft5 -> "ft5"
+  | Ft6 -> "ft6"
+  | Ft7 -> "ft7"
+  | Fs0 -> "fs0"
+  | Fs1 -> "fs1"
+  | Fa0 -> "fa0"
+  | Fa1 -> "fa1"
+  | Fa2 -> "fa2"
+  | Fa3 -> "fa3"
+  | Fa4 -> "fa4"
+  | Fa5 -> "fa5"
+  | Fa6 -> "fa6"
+  | Fa7 -> "fa7"
+  | Fs2 -> "fs2"
+  | Fs3 -> "fs3"
+  | Fs4 -> "fs4"
+  | Fs5 -> "fs5"
+  | Fs6 -> "fs6"
+  | Fs7 -> "fs7"
+  | Fs8 -> "fs8"
+  | Fs9 -> "fs9"
+  | Fs10 -> "fs10"
+  | Fs11 -> "fs11"
+  | Ft8 -> "ft8"
+  | Ft9 -> "ft9"
+  | Ft10 -> "ft10"
+  | Ft11 -> "ft11"
+;;
+
+let to_string (s : slot_t) : string =
+  match s with
+  | Slot i -> Printf.sprintf "%%%d" i
+  | FSlot i -> Printf.sprintf "%%f%d" i
+  | Reg r -> reg_to_string r
+  | FReg fr -> freg_to_string fr
+  | Unit -> "_"
+;;
+
+(** Counter of temporaries. *)
+let slot_cnt = ref 0
+
+let fslot_cnt = ref 0
+
+let new_slot () =
+  let i = !slot_cnt in
+  slot_cnt := i + 1;
+  Slot i
+;;
+
+let new_fslot () =
+  let i = !fslot_cnt in
+  fslot_cnt := i + 1;
+  FSlot i
+;;

src/riscv_reg_alloc.ml

@@ -0,0 +1,12 @@
+open Riscv_virtasm
+
+let reg_alloc (vprog : vprog_t) =
+  let vprog : vprog_t =
+    { blocks = VBlockMap.empty
+    ; funcs = VFuncMap.empty
+    ; consts = VSymbolMap.empty
+    ; loop_vars = VBlockMap.empty
+    }
+  in
+  vprog
+;;

src/riscv_ssa.ml

@@ -230,25 +230,31 @@ and t =
 let pointer_size = 8
 
 (* This is the size of their representations, not the actual size. *)
-let sizeof ty = match ty with
-| Mtype.T_bool -> 1
-| Mtype.T_byte -> 1
-| Mtype.T_bytes -> pointer_size
-| Mtype.T_char -> 2
-| Mtype.T_double -> 8
-| Mtype.T_float -> 4
-| Mtype.T_func _ -> pointer_size
-| Mtype.T_int -> 4
-| Mtype.T_int64 -> 8
-| Mtype.T_string -> pointer_size
-| Mtype.T_uint -> 4
-| Mtype.T_uint64 -> 8
-| Mtype.T_unit -> 0
-| Mtype.T_tuple _ -> pointer_size
-| Mtype.T_constr id -> pointer_size
-| Mtype.T_fixedarray _ -> pointer_size
-| _ -> failwith "riscv_ssa.ml: cannot calculate size"
-
+let rec sizeof ty =
+  match ty with
+  | Mtype.T_bool -> 1
+  | Mtype.T_byte -> 1
+  | Mtype.T_bytes -> pointer_size
+  | Mtype.T_char -> 2
+  | Mtype.T_double -> 8
+  | Mtype.T_float -> 4
+  | Mtype.T_func _ -> pointer_size
+  | Mtype.T_int -> 4
+  | Mtype.T_int64 -> 8
+  | Mtype.T_string -> pointer_size
+  | Mtype.T_uint -> 4
+  | Mtype.T_uint64 -> 8
+  | Mtype.T_unit -> 0 (* Unit type has no size *)
+  | Mtype.T_tuple _ -> pointer_size
+  | Mtype.T_constr id -> pointer_size
+  | Mtype.T_fixedarray _ -> pointer_size
+  | Mtype.T_trait _ -> pointer_size
+  (* | Mtype.T_optimized_option { elem } -> pointer_size  *)
+  (* | Mtype.T_any { name } -> pointer_size  *)
+  (* | Mtype.T_maybe_uninit x -> sizeof x *)(*Same size as the contained type *)
+  (* | Mtype.T_error_value_result { ok; err; id } -> sizeof ok + sizeof err + pointer_size *)
+  | _ -> failwith ("riscv_ssa.ml: cannot calculate size for type: "^ Mtype.to_string ty)
+;;
 
 (** Emits SSA form. We choose a less human-readable form to facilitate verifier. *)
 let to_string t =