Merge branch 'main' into emission-fix

README.md

@@ -7,7 +7,7 @@
 ![CI Status](https://github.com/ethanuppal/cs3110_compiler/actions/workflows/ci.yaml/badge.svg)
 
 > "x86 is simple trust me bro"  
-> Last updated: 2024-05-14 22:50:09.279091
+> Last updated: 2024-05-15 09:58:18.266097
 
 ```
 $ ./main -h

lib/backend/asm.ml

@@ -41,8 +41,16 @@ module Register = struct
 
   let compare = Stdlib.compare
 
-  (** Every register but RBX, RSP, RBP, and R12–R15. *)
-  let caller_saved = [ RAX; RCX; RDX; RSI; RDI; R8; R9; R10; R11 ]
+  (** Caller saved registers. rsp is excluded because it is used to manage the
+      stack and usually needs special care. *)
+  let caller_saved_data_registers =
+    [ RAX; RCX; RDX; RDI; RSI; R8; R9; R10; R11 ]
+
+  (** Caller saved registers. rbp is excluded because it is used to manage the
+      stack and usually needs special care. *)
+  let callee_saved_data_registers = [ RBX; R12; R13; R14; R15 ]
+
+  let parameter_passing_registers = [ RDI; RSI; RDX; RCX; R8; R9 ]
 end
 
 module Operand = struct

lib/backend/asm_emit.ml

@@ -7,8 +7,13 @@ let mangle name =
   "_x86istmb" ^ mangle_helper name
 
 let debug_print_symbol = mangle [ "std"; "debug_print" ]
+let stack_alignment = 16
 let var_size = 8
 
+let align_offset bytes =
+  let amount_over = bytes mod stack_alignment in
+  if amount_over = 0 then 0 else stack_alignment - amount_over
+
 let emit_var regalloc var =
   match Ir.VariableMap.find regalloc var with
   | Regalloc.Register reg -> Asm.Operand.Register reg
@@ -18,22 +23,40 @@ let emit_oper regalloc = function
   | Operand.Variable var -> emit_var regalloc var
   | Constant int -> Asm.Operand.Intermediate int
 
+(** [emit_save_registers text registers] emits instructions into [text] to save
+    [registers] on the stack. If the register contents are to be restored, they
+    must be restored using [emit_restore_registers]. *)
+let emit_save_registers text registers =
+  let open Asm.Instruction in
+  let push_instructions = List.map (fun reg -> Push (Register reg)) registers in
+  let added_size = var_size * List.length registers in
+  let extra_offset = align_offset added_size in
+  Asm.Section.add_all text push_instructions;
+  Asm.Section.add text (Sub (Register RSP, Intermediate extra_offset))
+
+(** [emit_save_registers text registers] emits instructions into [text] to pop
+    [registers] from the stack. The registers being restored must be the same as
+    those most recently pushed using [emit_save_registers]. *)
+let emit_restore_registers text registers =
+  let open Asm.Instruction in
+  let pop_instructions =
+    List.rev_map (fun reg -> Pop (Register reg)) registers
+  in
+  let added_size = var_size * List.length registers in
+  let extra_offset = align_offset added_size in
+  Asm.Section.add text (Add (Register RSP, Intermediate extra_offset));
+  Asm.Section.add_all text pop_instructions
+
 let emit_call text regalloc name args =
-  let to_save = Asm.Register.caller_saved in
-  let to_save =
-    if List.length to_save mod 2 = 0 then to_save
-    else List.hd to_save :: to_save
+  emit_save_registers text Asm.Register.caller_saved_data_registers;
+  let param_moves =
+    Util.zip_shortest args Asm.Register.parameter_passing_registers
+    |> List.map (fun (arg, reg) ->
+           Asm.Instruction.Mov (Register reg, emit_oper regalloc arg))
   in
-  let to_pass = [| Asm.Register.RDI; RSI; RDX; RCX; R8; R9 |] in
-  Asm.Section.add_all text
-    (List.map (fun r -> Asm.Instruction.Push (Register r)) to_save
-    @ List.mapi
-        (fun i arg ->
-          Asm.Instruction.Mov (Register to_pass.(i), emit_oper regalloc arg))
-        args
-    @ [ Asm.Instruction.Call (Label name) ]
-    @ (List.map (fun r -> Asm.Instruction.Pop (Register r)) to_save |> List.rev)
-    )
+  Asm.Section.add_all text param_moves;
+  Asm.Section.add text (Asm.Instruction.Call (Label name));
+  emit_restore_registers text Asm.Register.caller_saved_data_registers
 
 let emit_ir text regalloc = function
   | Ir.Assign (var, op) ->
@@ -62,6 +85,7 @@ let emit_ir text regalloc = function
           Asm.Section.add text (Mov (Register RAX, emit_oper regalloc op)))
         op_opt
       |> ignore;
+      emit_restore_registers text Asm.Register.callee_saved_data_registers;
       Asm.Section.add_all text
         [ Mov (Register RSP, Register RBP); Pop (Register RBP); Ret ]
 
@@ -102,17 +126,22 @@ let emit_cfg ~text cfg regalloc =
   in
   (* max_spill starts at zero but we need to start from rbp-8 *)
   let spill_bytes = var_size * (max_spill + 1) in
-  (* What we need to spill + however much is needed to round our total up to
-     16 *)
-  let total_bytes = spill_bytes + var_size in
-  let stack_bytes = spill_bytes + (total_bytes mod 16) in
+  let stack_bytes = spill_bytes + align_offset spill_bytes in
+
+  let entry = Cfg.entry_to cfg in
+  Asm.Section.add text
+    (Label
+       (Asm.Label.make ~is_global:true ~is_external:false
+          (mangle (Cfg.name_of cfg))));
+  (* no need to align here, we can assume as a callee that 8 bytes for the
+     return address was already pushed to the stack. *)
   Asm.Section.add_all text
     [
-      Label
-        (Asm.Label.make ~is_global:true ~is_external:false
-           (mangle (Cfg.name_of cfg)));
       Push (Register RBP);
       Mov (Register RBP, Register RSP);
       Sub (Register RSP, Intermediate stack_bytes);
     ];
+  (* restore is done at returns *)
+  emit_save_registers text Asm.Register.callee_saved_data_registers;
+  Asm.Section.add text (Jmp (Label (Basic_block.label_for entry)));
   Cfg.blocks_of cfg |> List.iter (emit_bb text cfg regalloc)

lib/backend/regalloc/regalloc.ml

@@ -19,6 +19,8 @@ type allocation =
 
 module BBAnalysis = Liveliness.BasicBlockAnalysis
 
+(* todo make this less awkward why are registers just listed here maybe pass em
+   in or smth *)
 let registers =
   let open Asm.Register in
   [ RAX; RBX; RCX; RDX; RSI; RDI; R8; R9; R10; R11; R12; R13; R14; R15 ]

lib/util/util.ml

@@ -90,6 +90,11 @@ let basename =
   in
   String.to_seq >> List.of_seq >> List.rev >> basename_aux
 
+let rec zip_shortest lst1 lst2 =
+  match (lst1, lst2) with
+  | h1 :: t1, h2 :: t2 -> (h1, h2) :: zip_shortest t1 t2
+  | _, _ -> []
+
 (** [pp_of string_of] is a pretty printer for a type with the string conversion
     function [string_of] that simply prints the result of [string_of] inline. *)
 let pp_of string_of fmt x = Format.fprintf fmt "%s" (string_of x)