Fix returns and add assembly minimization

Makefile

@@ -7,7 +7,7 @@ build:
 	@cp _build/install/default/bin/x86ISTMB ./main
 	@chmod u+x ./main
 	@chmod u+x .githooks/pre-commit
-	@make README
+	@make README 1> /dev/null
 
 .PHONY: protect
 protect:

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-15 22:27:21.599048
+> Last updated: 2024-05-15 23:49:52.345644
 
 ```
 $ ./main -h

lib/backend/asm.ml

@@ -78,6 +78,7 @@ module Label = struct
   }
 
   let make ~is_global ~is_external name = { is_global; is_external; name }
+  let name_of label = label.name
 
   let to_nasm label =
     match (label.is_global, label.is_external) with
@@ -92,7 +93,7 @@ module Instruction = struct
     | Mov of Operand.t * Operand.t
     | Add of Operand.t * Operand.t
     | Sub of Operand.t * Operand.t
-    | IMul of Operand.t
+    | IMul of Operand.t * Operand.t
     | Push of Operand.t
     | Pop of Operand.t
     | Call of Operand.t
@@ -111,7 +112,8 @@ module Instruction = struct
         "add " ^ Operand.to_nasm op1 ^ ", " ^ Operand.to_nasm op2
     | Sub (op1, op2) ->
         "sub " ^ Operand.to_nasm op1 ^ ", " ^ Operand.to_nasm op2
-    | IMul op -> "imul " ^ Operand.to_nasm op
+    | IMul (op1, op2) ->
+        "imul " ^ Operand.to_nasm op1 ^ ", " ^ Operand.to_nasm op2
     | Push op -> "push " ^ Operand.to_nasm op
     | Pop op -> "pop " ^ Operand.to_nasm op
     | Call op -> "call " ^ Operand.to_nasm op
@@ -146,6 +148,11 @@ module Section = struct
               | Label _ -> str
               | _ -> display_indent ^ str))
     |> String.concat "\n"
+
+  let get_instr section idx = BatDynArray.get section.contents idx
+  let set_instr section idx instr = BatDynArray.set section.contents idx instr
+  let rem_instr section idx = BatDynArray.delete section.contents idx
+  let length_of section = BatDynArray.length section.contents
 end
 
 module AssemblyFile = struct

lib/backend/asm.mli

@@ -64,6 +64,9 @@ module Label : sig
       if [is_external], but not both. *)
   val make : is_global:bool -> is_external:bool -> string -> t
 
+  (** [name_of label] is the name of [label]. *)
+  val name_of : t -> string
+
   (** [to_nasm label] is the NASM representation of [label]. *)
   val to_nasm : t -> string
 end
@@ -75,7 +78,7 @@ module Instruction : sig
     | Mov of Operand.t * Operand.t
     | Add of Operand.t * Operand.t
     | Sub of Operand.t * Operand.t
-    | IMul of Operand.t
+    | IMul of Operand.t * Operand.t
     | Push of Operand.t
     | Pop of Operand.t
     | Call of Operand.t
@@ -109,6 +112,21 @@ module Section : sig
 
   (** [to_nasm section] is the NASM representation of [section]. *)
   val to_nasm : t -> string
+
+  (** [get_instr section idx] is the instruction at index [idx] in [section]. *)
+  val get_instr : t -> int -> Instruction.t
+
+  (** [set_instr section idx instr] sets the instruction at index [idx] in
+      [section] to [instr]. *)
+  val set_instr : t -> int -> Instruction.t -> unit
+
+  (** [rem_instr section idx] deletes the instruction at index [idx] in
+      [section], invalidating all subsequent indices for [get_instr] and
+      [set_instr]. *)
+  val rem_instr : t -> int -> unit
+
+  (** [length_of section] is the number of instructions in [section]. *)
+  val length_of : t -> int
 end
 
 (** Contains functionality for creating assembly files. *)

lib/backend/asm_clean.ml

@@ -0,0 +1,51 @@
+let apply_clean_rules_single section i =
+  let cur = Asm.Section.get_instr section i in
+  match cur with
+  | Asm.Instruction.Mov (o1, o2) when o1 = o2 ->
+      Asm.Section.rem_instr section i;
+      true
+  | _ -> false
+
+let apply_clean_rules_pair section i =
+  let cur = Asm.Section.get_instr section i in
+  let next = Asm.Section.get_instr section (i + 1) in
+  let delete_pair () =
+    Asm.Section.rem_instr section (i + 1);
+    Asm.Section.rem_instr section i
+  in
+  match (cur, next) with
+  | Asm.Instruction.Jmp (Asm.Operand.Label label), Asm.Instruction.Label label2
+    when label = Asm.Label.name_of label2 ->
+      delete_pair ();
+      true
+  | (Push op1, Pop op2 | Pop op1, Push op2) when op1 = op2 ->
+      delete_pair ();
+      true
+  | (Add (r1, v1), Sub (r2, v2) | Sub (r1, v1), Add (r2, v2))
+    when r1 = r2 && v1 = v2 ->
+      delete_pair ();
+      true
+  | _ -> false
+
+let apply_clean_rules section i =
+  let first_try =
+    if i > 0 then apply_clean_rules_pair section (i - 1) else false
+  in
+  let second_try =
+    if i < Asm.Section.length_of section then apply_clean_rules_single section i
+    else false
+  in
+  first_try || second_try
+
+let clean_pass section =
+  let length = Asm.Section.length_of section in
+  let did_change = ref false in
+  for rev_i = 0 to length - 1 do
+    let i = length - rev_i - 1 in
+    did_change := apply_clean_rules section i || !did_change
+  done;
+  !did_change
+
+let clean section =
+  let rec clean_aux section = if clean_pass section then clean_aux section in
+  clean_aux section

lib/backend/asm_clean.mli

@@ -0,0 +1,3 @@
+(** [clean section] minimizes redundant assembly in [section] using naive rules
+    that are guaranteed to not affect the data. *)
+val clean : Asm.Section.t -> unit

lib/backend/asm_emit.ml

@@ -67,16 +67,27 @@ let emit_restore_registers text registers =
   Asm.Section.add text (Add (Register RSP, Intermediate extra_offset));
   Asm.Section.add_all text pop_instructions
 
-let emit_call text regalloc name args =
-  emit_save_registers text Asm.Register.caller_saved_data_registers;
+let emit_call text regalloc name args return_loc_opt =
+  let save_registers =
+    List.filter
+      (fun reg ->
+        match return_loc_opt with
+        | Some (Asm.Operand.Register return_reg) -> reg <> return_reg
+        | _ -> true)
+      Asm.Register.caller_saved_data_registers
+  in
+  emit_save_registers text save_registers;
   let param_moves =
     Util.zip_shortest args Asm.Register.parameter_registers
     |> List.map (fun (arg, reg) ->
            Asm.Instruction.Mov (Register reg, emit_oper regalloc arg))
   in
   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
+  (match return_loc_opt with
+  | Some return_loc -> Asm.Section.add text (Mov (return_loc, Register RAX))
+  | None -> ());
+  emit_restore_registers text save_registers
 
 let emit_ir text regalloc = function
   | Ir.Assign (var, op) ->
@@ -93,12 +104,17 @@ let emit_ir text regalloc = function
           Mov (emit_var regalloc var, emit_oper regalloc op);
           Sub (emit_var regalloc var, emit_oper regalloc op2);
         ]
+  | Mul (var, op, op2) ->
+      Asm.Section.add_all text
+        [
+          Mov (emit_var regalloc var, emit_oper regalloc op);
+          IMul (emit_var regalloc var, emit_oper regalloc op2);
+        ]
   | Ref _ -> failwith "ref not impl"
   | Deref _ -> failwith "deref not impl"
-  | DebugPrint op -> emit_call text regalloc debug_print_symbol [ op ]
+  | DebugPrint op -> emit_call text regalloc debug_print_symbol [ op ] None
   | Call (var, name, args) ->
-      emit_call text regalloc (mangle name) args;
-      Asm.Section.add text (Mov (emit_var regalloc var, Register RAX))
+      emit_call text regalloc (mangle name) args (Some (emit_var regalloc var))
   | GetParam var -> (
       let param_passing = ParameterPassingContext.make () in
       match ParameterPassingContext.get_next param_passing with

lib/backend/liveliness.ml

@@ -1,5 +1,4 @@
 open Util
-open Ir
 
 (** A set of IR variables. *)
 module VariableSet = struct
@@ -128,20 +127,8 @@ let apply_rules liveliness analysis cfg bb ir ir_idx ~is_final =
     instr_analysis.live_in <- VariableSet.remove var instr_analysis.live_in
   in
   bring_incoming ();
-  (match ir with
-  | DebugPrint op -> read_op op
-  | Assign (var, op) | Deref (var, op) | Ref (var, op) ->
-      write_var var;
-      read_op op
-  | Add (var, op1, op2) | Sub (var, op1, op2) | TestEqual (var, op1, op2) ->
-      write_var var;
-      read_op op1;
-      read_op op2
-  | Call (var, _, args) ->
-      write_var var;
-      List.iter read_op args
-  | GetParam var -> write_var var
-  | Return opt_op -> Option.map read_op opt_op |> ignore);
+  Option.map write_var (Ir.kill_of ir) |> ignore;
+  List.iter read_op (Ir.gen_of ir);
   check_for_changes ()
 
 (** [pass work_list liveliness cfg bb] performs a single pass of liveliness

lib/frontend/ir_gen.ml

@@ -27,6 +27,7 @@ let rec generate_expr ctx cfg block expr =
         match op with
         | Plus -> Ir.Add (result, lhs_result, rhs_result)
         | Minus -> Ir.Sub (result, lhs_result, rhs_result)
+        | Times -> Ir.Mul (result, lhs_result, rhs_result)
         | Equals -> Ir.TestEqual (result, lhs_result, rhs_result)
         | _ -> failwith "not implemented"
       in

lib/ir/ir.ml

@@ -2,6 +2,7 @@ type t =
   | Assign of Variable.t * Operand.t
   | Add of Variable.t * Operand.t * Operand.t
   | Sub of Variable.t * Operand.t * Operand.t
+  | Mul of Variable.t * Operand.t * Operand.t
   | Ref of Variable.t * Operand.t
   | Deref of Variable.t * Operand.t
   | TestEqual of Variable.t * Operand.t * Operand.t
@@ -14,13 +15,27 @@ let kill_of = function
   | Assign (var, _)
   | Add (var, _, _)
   | Sub (var, _, _)
+  | Mul (var, _, _)
   | Ref (var, _)
   | Deref (var, _)
   | TestEqual (var, _, _)
   | GetParam var
   | Call (var, _, _) -> Some var
   | DebugPrint _ | Return _ -> None
 
+let gen_of = function
+  | Assign (_, op)
+  | Ref (_, op)
+  | Deref (_, op)
+  | DebugPrint op
+  | Return (Some op) -> [ op ]
+  | Add (_, op1, op2)
+  | Sub (_, op1, op2)
+  | Mul (_, op1, op2)
+  | TestEqual (_, op1, op2) -> [ op1; op2 ]
+  | Call (_, _, ops) -> ops
+  | GetParam _ | Return None -> []
+
 let to_string =
   let open Printf in
   function
@@ -32,6 +47,9 @@ let to_string =
   | Sub (r, o1, o2) ->
       sprintf "%s = %s - %s" (Variable.to_string r) (Operand.to_string o1)
         (Operand.to_string o2)
+  | Mul (r, o1, o2) ->
+      sprintf "%s = %s * %s" (Variable.to_string r) (Operand.to_string o1)
+        (Operand.to_string o2)
   | Ref (r, o) ->
       sprintf "%s = &%s" (Variable.to_string r) (Operand.to_string o)
   | Deref (r, o) ->
@@ -43,10 +61,10 @@ let to_string =
   | Call (r, name, args) ->
       sprintf "%s = %s(%s)" (Variable.to_string r)
         (name |> String.concat "::")
-        (args |> List.map Operand.to_string |> String.concat ",")
+        (args |> List.map Operand.to_string |> String.concat ", ")
   | GetParam var -> sprintf "%s = <next parameter>" (Variable.to_string var)
   | Return op ->
-      sprintf "return %s"
+      sprintf "return%s"
         (match op with
         | Some op -> " " ^ Operand.to_string op
         | None -> "")

lib/ir/ir.mli

@@ -3,6 +3,7 @@ type t =
   | Assign of Variable.t * Operand.t
   | Add of Variable.t * Operand.t * Operand.t
   | Sub of Variable.t * Operand.t * Operand.t
+  | Mul of Variable.t * Operand.t * Operand.t
   | Ref of Variable.t * Operand.t
   | Deref of Variable.t * Operand.t
   | TestEqual of Variable.t * Operand.t * Operand.t
@@ -15,5 +16,8 @@ type t =
     otherwise. *)
 val kill_of : t -> Variable.t option
 
+(** [gen_of ir] is a list of operands read from in [ir]. *)
+val gen_of : t -> Operand.t list
+
 (** [to_string ir] is a string representation of the IRk instruction [ir]. *)
 val to_string : t -> string

lib/ir/ir_sim.ml

@@ -45,6 +45,10 @@ let rec run_cfg simulator cfgs cfg =
                  Context.insert simulator.context (Variable.to_string var)
                    (eval oper1 - eval oper2);
                  false
+             | Ir.Mul (var, oper1, oper2) ->
+                 Context.insert simulator.context (Variable.to_string var)
+                   (eval oper1 * eval oper2);
+                 false
              | Ir.TestEqual (var, oper1, oper2) ->
                  Context.insert simulator.context (Variable.to_string var)
                    (if eval oper1 = eval oper2 then 1 else 0);

lib/user/driver.ml

@@ -55,6 +55,7 @@ let compile paths flags build_dir_loc =
         in
         Asm_emit.emit_cfg ~text:text_section cfg regalloc)
       cfgs;
+    Asm_clean.clean text_section;
     let asm_file = Asm.AssemblyFile.make () in
     Asm.AssemblyFile.add asm_file text_section;
     let file_name_root =
@@ -110,7 +111,8 @@ let compile paths flags build_dir_loc =
         in
         if Sys.command nasm_command <> 0 then failwith "Failed to run NASM."
         else
-          Printf.printf "==> \x1B[32mGenerated \x1B[4m%s\x1B[m\n" asm_file_name)
+          Printf.printf "==> \x1B[32mGenerated \x1B[4m%s/%s\x1B[m\n" build_dir
+            asm_file_name)
       compiled_files;
 
     (* Run clang *)

test/e2e/complex_return.x

@@ -0,0 +1,9 @@
+// Computes ax^2 + bx + c
+func eval_quadratic(a: Int, b: Int, c: Int, x: Int) -> Int {
+    return a * x * x + b * x + c
+}
+
+func main() {
+    print eval_quadratic(0, 0, 0, 0)
+    print eval_quadratic(1, 1, 1, 1)
+}

test/e2e/return.x

@@ -0,0 +1,9 @@
+namespace math {
+    func inc(x: Int) -> Int {
+        return x + 1
+    }
+}
+
+func main() {
+    print math::inc(1)
+}

test/test_e2e.ml

@@ -26,5 +26,5 @@ let test_suite =
     Sys.readdir e2e_root |> Array.to_list
     |> List.map (fun filename ->
            let path = Util.merge_paths [ e2e_root; filename ] in
-           test_case filename `Slow (make_e2e_test path (Util.read_file path)))
+           test_case filename `Quick (make_e2e_test path (Util.read_file path)))
   )