open Ir
open Nemo
open Printer
open Parser

(* ---- Inline unit tests (no file I/O) ---- *)

let test_simple () =
  Printf.printf "Test Simple\n";
  let b1 = mk_block 1 [Alloca 100; Store (1, 100); Br 2] in
  let b2 = mk_block 2 [Load (200, 100); Ret (Some 200)] in
  let f = { blocks = [b1; b2]; entry = 1 } in
  Printf.printf "Before:\n";
  print_func f;
  let f' = mem2reg f in
  Printf.printf "\nAfter:\n";
  print_func f'

let test_phi () =
  Printf.printf "\nTest Phi\n";
  let b1 = mk_block 1 [Alloca 100; Store (1, 100); Br 2] in
  let b2 = mk_block 2 [Alloca 101; Store (2, 101); Br 3] in
  let b3 = mk_block 3 [Load (200, 100); Load (201, 101); Ret (Some 200)] in
  let f = { blocks = [b1; b2; b3]; entry = 1 } in
  Printf.printf "Before:\n";
  print_func f;
  let f' = mem2reg f in
  Printf.printf "\nAfter:\n";
  print_func f'

let test_branch () =
  Printf.printf "\nTest Branch with Phi\n";
  let b1 = mk_block 1 [Alloca 100; CondBr (1, 2, 3)] in
  let b2 = mk_block 2 [Store (10, 100); Br 4] in
  let b3 = mk_block 3 [Store (20, 100); Br 4] in
  let b4 = mk_block 4 [Load (200, 100); Ret (Some 200)] in
  let f = { blocks = [b1; b2; b3; b4]; entry = 1 } in
  Printf.printf "Before:\n";
  print_func f;
  let f' = mem2reg f in
  Printf.printf "\nAfter:\n";
  print_func f'

(* ---- File-based regression test ---- *)

(* Parse a .ir file, run mem2reg, and assert the printed SSA output equals the
   expected string. Exits with status 1 on mismatch so `make test` fails. *)
let string_of_func f =
  let buf = Buffer.create 256 in
  List.iter (fun b ->
    Buffer.add_string buf (Printf.sprintf "block %d:\n" b.id);
    List.iter (fun i ->
      Buffer.add_string buf (Printf.sprintf "  %s\n" (Printer.string_of_inst i))
    ) b.insts
  ) f.blocks;
  Buffer.contents buf

let test_file name expected () =
  Printf.printf "\nTest file: %s\n" name;
  let f = parse_func_from_file ("examples/" ^ name) in
  let f' = mem2reg f in
  let got = string_of_func f' in
  if got <> expected then begin
    Printf.eprintf "FAIL %s\n--- expected ---\n%s\n--- got ---\n%s\n" name expected got;
    exit 1
  end;
  print_string got

let expected_simple =
"block 1:\n  br 2\nblock 2:\n  move 200, 1\n  ret 200\n"

let expected_basic =
"block 1:\n  br 2\nblock 2:\n  move 200, 42\n  ret 200\n"

(* Example exercising the phi-handling path: an input phi node is filled in
   from the renamed value at each predecessor. *)
let expected_phi =
"block 1:\n  condbr 1, 2, 3\nblock 2:\n  br 4\nblock 3:\n  br 4\nblock 4:\n  phi 100 [(20, 3), (10, 2)]\n  move 200, 100\n  ret 200\n"

let () =
  test_simple ();
  test_phi ();
  test_branch ();
  test_file "simple.ir" expected_simple ();
  test_file "basic.ir" expected_basic ();
  test_file "phi.ir" expected_phi ()